@salesforce/afv-skills 1.67.0 → 1.68.0

package/skills/agentforce-offline-experimentation-run/scripts/sfap_replay.py

@@ -1,811 +1,487 @@
 #!/usr/bin/env python3
-"""
-SFAP-based replay pipeline for offline experimentation.
-Reads STDM data, replays prompts via SFAP REST API with structured tool calling,
-scores results, and generates a CSV report.
-"""
-import json, re, csv, os, sys, subprocess, time
-
-# ---------------------------------------------------------------------------
-# Config
-# ---------------------------------------------------------------------------
-SFAP_BASE_URL  = os.environ.get("SFAP_BASE_URL", "https://test.api.salesforce.com/ai/gpt/v1")
-SFAP_TOKEN_FILE = os.environ.get("SFAP_TOKEN_FILE", "")
-CORE_TENANT_ID = os.environ.get("CORE_TENANT_ID", "")
-REPLAY_MODEL   = os.environ.get("REPLAY_MODEL", "llmgateway__OpenAIGPT41")
-JUDGE_MODEL    = os.environ.get("JUDGE_MODEL", REPLAY_MODEL)
-OUTPUT_CSV     = os.environ.get("OUTPUT_CSV", "replay_results_sfap.csv")
-ORG_ID         = os.environ.get("ORG_ID", "")
-AGENT_API_NAME = os.environ.get("AGENT_API_NAME", "")
-OUTPUT_JSON    = os.environ.get("OUTPUT_JSON", "replay_results.json")
-STDM_DATA_FILE = os.environ.get("STDM_DATA_FILE", "/tmp/details_result.json")
-LLM_STEPS_FILES = os.environ.get("LLM_STEPS_FILES", "/tmp/llm_steps_1.json,/tmp/llm_steps_2.json").split(",")
-
-# ---------------------------------------------------------------------------
-# Tool extraction is dynamic-only: definitions come from STDM step inputs
-# (gen_ai.request.tools) at runtime. No hardcoded tool schemas.
-# ---------------------------------------------------------------------------
-
-
-def _is_routing_tools(tools):
-    """True when ALL tools in the list are routing functions (go_to_* / __end_session_action__).
-    Action steps may include a single go_to_* for cross-topic transfer alongside domain tools,
-    so checking 'any' would mis-classify them."""
-    if not tools:
-        return False
-    for t in tools:
-        name = t.get('function', {}).get('name', '')
-        if not (name.startswith('go_to_') or name == '__end_session_action__'):
-            return False
-    return True
-
-
-def _extract_tool_call_from_step_output(step):
-    """Extract the first tool invocation from a step's gen_ai.output.messages."""
-    step_output = step.get('output', '') or ''
-    try:
-        out_json = json.loads(step_output) if isinstance(step_output, str) else step_output
-        out_msgs = out_json.get('gen_ai.output.messages', [])
-        for om in out_msgs:
-            for inv in om.get('tool_invocations', []):
-                func = inv.get('function', {})
-                if func.get('name'):
-                    args = func.get('arguments', '')
-                    if isinstance(args, dict):
-                        args = json.dumps(args)
-                    return func['name'], args
-    except Exception:
-        pass
-    return None, None
-
-
-def sanitize_tools_for_sfap(tools):
-    """Clean STDM tool schemas to be compatible with SFAP's OpenAI-format endpoint.
-    Removes additionalProperties and empty required arrays that cause 400 errors."""
-    sanitized = []
-    for tool in tools:
-        tool = json.loads(json.dumps(tool))  # deep copy
-        func = tool.get('function', {})
-        params = func.get('parameters', {})
-        params.pop('additionalProperties', None)
-        if 'required' in params and not params['required']:
-            del params['required']
-        sanitized.append(tool)
-    return sanitized
-
-
-def extract_tools_from_llm_responses(llm_step_map):
-    """Scan all LLM responses to discover the actual tool names used by the agent."""
-    seen = set()
-    for step in llm_step_map.values():
-        resp_text = step.get('llm_response', '') or ''
-        if not resp_text:
-            continue
-        try:
-            resp = json.loads(resp_text)
-            for inv in resp.get('toolInvocations', []):
-                name = inv.get('function', {}).get('name', '')
-                if name:
-                    seen.add(name)
-        except:
-            pass
-    return seen
-
-
-def read_token():
-    if SFAP_TOKEN_FILE and os.path.exists(SFAP_TOKEN_FILE):
-        with open(SFAP_TOKEN_FILE) as f:
-            return f.read().strip()
-    raise RuntimeError("SFAP_TOKEN_FILE not set or file not found")
-
-
-def parse_prompt_to_messages(prompt_text):
-    """Parse role-based prompt text into OpenAI-style messages array.
-
-    Keeps system messages in their original position — production keeps
-    subagent instructions at the end of the conversation, not merged to the top.
-    Merges consecutive same-role messages to keep alternation valid.
-    Appends a user "Continue." if conversation ends on assistant."""
-    messages = []
-    current_role = None
-    current_content = []
-
-    for line in prompt_text.split('\n'):
-        stripped = line.strip()
-        if stripped in ('role: system', 'role: user', 'role: assistant'):
-            if current_role and current_content:
-                content = '\n'.join(current_content).strip()
-                if content:
-                    messages.append({"role": current_role, "content": content})
-            current_role = stripped.replace('role: ', '')
-            current_content = []
-        else:
-            if current_role:
-                current_content.append(line)
-
-    if current_role and current_content:
-        content = '\n'.join(current_content).strip()
-        if content:
-            messages.append({"role": current_role, "content": content})
-
-    # Merge consecutive same-role messages (keeps system at original position)
-    merged = []
-    for m in messages:
-        if merged and merged[-1]['role'] == m['role']:
-            merged[-1]['content'] = f"{merged[-1]['content']}\n{m['content']}".strip()
-        else:
-            merged.append(m)
-
-    if merged and merged[-1]['role'] == 'assistant':
-        merged.append({"role": "user", "content": "Continue."})
-    return merged
-
-
-def convert_stdm_messages_for_sfap(stdm_messages):
-    """Convert STDM gen_ai.input.messages to SFAP-compatible OpenAI format.
-
-    Matches the temporal version's approach (src/replay/traces.py:_messages):
-    - Folds tool-role messages into the previous message as text
-    - Converts assistant tool_invocations to text (keeps content)
-    - Merges consecutive same-role messages
-    - Keeps system messages in their original position (does NOT merge trailing
-      system into the first — production keeps subagent instructions at the end)
-    - Appends a user "Continue." if conversation ends on assistant
-    """
-    raw = []
-    for m in stdm_messages:
-        role = m.get('role', '')
-        content = m.get('content', '') or ''
-        tool_invocations = m.get('tool_invocations', [])
-        if role == 'assistant' and tool_invocations:
-            parts = []
-            if content.strip():
-                parts.append(content)
-            for inv in tool_invocations:
-                func = inv.get('function', {})
-                name = func.get('name', '')
-                args = func.get('arguments', '')
-                if isinstance(args, dict):
-                    args = json.dumps(args)
-                parts.append(f"[tool_call] {name}({args})")
-            raw.append({"role": "assistant", "content": "\n".join(parts)})
-        else:
-            raw.append({"role": role, "content": content})
-
-    out = []
-    for m in raw:
-        role, content = m["role"], m["content"]
-        if role == "tool":
-            text = f"[tool result] {content}".strip()
-            if out:
-                out[-1]["content"] = f"{out[-1]['content']}\n{text}".strip()
-            else:
-                out.append({"role": "assistant", "content": text})
-            continue
-        if not content.strip():
-            continue
-        if out and out[-1]["role"] == role:
-            out[-1]["content"] = f"{out[-1]['content']}\n{content}".strip()
-        else:
-            out.append({"role": role, "content": content})
-
-    if out and out[-1]["role"] == "assistant":
-        out.append({"role": "user", "content": "Continue."})
-    return out
-
-
-def call_sfap(messages, model, tools=None, max_tokens=2048):
-    """Call SFAP chat generations endpoint."""
-    token = read_token()
-    url = f"{SFAP_BASE_URL}/chat/generations"
-
-    payload = {
-        "model": model,
-        "messages": messages,
-        "max_tokens": max_tokens,
-    }
-    if tools:
-        payload["tools"] = tools
-        payload["tool_choice"] = "required"
-
-    payload_json = json.dumps(payload)
-
-    cmd = [
-        "curl", "-sS", "-w", "\n__HTTP_STATUS:%{http_code}",
-        url,
-        "-H", f"Authorization: Bearer {token}",
-        "-H", "Content-Type: application/json",
-        "-H", "x-sfdc-app-context: EinsteinGPT",
-        "-H", "x-client-feature-id: AgentBacktesting",
-        "-H", f"x-sfdc-core-tenant-id: {CORE_TENANT_ID}",
-        "-d", payload_json,
-    ]
-
-    result = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
-    output = result.stdout
-
-    # Split response body and HTTP status
-    parts = output.rsplit('\n__HTTP_STATUS:', 1)
-    body = parts[0] if parts else output
-    status = int(parts[1].strip()) if len(parts) > 1 else 0
-
-    if status != 200:
-        return {"error": f"HTTP {status}", "body": body[:500]}
-
-    return json.loads(body)
-
-
-def extract_tool_call(sfap_response):
-    """Extract tool invocation from SFAP response."""
-    try:
-        gens = sfap_response.get("generation_details", {}).get("generations", [])
-        if not gens:
-            return None, None, None
+"""Replay production LLM steps (from STDM) on a candidate model and score them against production.
 
-        gen = gens[0]
-        content = gen.get("content", "")
-        invocations = gen.get("tool_invocations") or []
+Every LLM_STEP of every turn is one unit (btlib/units.py): the request exactly as production
+sent it, and production's own answer as the golden. The candidate model answers the same
+request through SFAP and is scored step by step (btlib/scoring.py).
 
-        if invocations:
-            inv = invocations[0]
-            func = inv.get("function", {})
-            return func.get("name"), func.get("arguments", ""), content
-        return None, None, content
-    except Exception as e:
-        return None, None, str(e)
+    python3 scripts/bt_fetch.py --org <alias> --agent <api name> --start <date> --end <date> \\
+        --out-dir runs/fetch
+    bash scripts/resolve-token.sh --org <alias> > /tmp/token_info.json
+    python3 scripts/sfap_replay.py --details runs/fetch/details.json \\
+        --token-info /tmp/token_info.json --model llmgateway__EinsteinSalesforceCerebro \\
+        --out-dir runs/cerebro
+    python3 scripts/build_report.py runs/cerebro
 
+--model trace replays each step on the model that served it in production. That self-control
+should agree almost fully; what it misses is replay noise, not a model difference.
 
-def extract_original_tool_call(llm_response_text):
-    """Extract tool invocation from the original STDM LLM response."""
-    if not llm_response_text:
-        return None, None, llm_response_text
+--patch FILE applies edit ops to the requests first (references/patching.md): the replay
+then scores the edited agent without a deploy. Untouched steps are sent unchanged.
 
-    try:
-        resp = json.loads(llm_response_text)
-        invocations = resp.get("toolInvocations", [])
-        if invocations:
-            func = invocations[0].get("function", {})
-            return func.get("name"), func.get("arguments", ""), resp.get("content", "")
-        return None, None, resp.get("content", llm_response_text)
-    except:
-        return None, None, llm_response_text
-
-
-def judge_text(actual, expected, question, model):
-    """LLM-as-judge for text responses via SFAP."""
-    prompt = f"""You are a scoring judge. Compare the predicted response against the reference response.
-
-Output ONLY valid JSON with "reasoning" and "score" fields.
-
-**Scoring Rubric**
-1. score=1 -- the predicted response covers key information of the reference response OR it takes a different but reasonable action
-2. score=0.66 -- the predicted response covers a good amount of information OR it takes a somewhat reasonable action
-3. score=0.33 -- the predicted response covers little information OR it takes a less reasonable action
-4. score=0 -- the predicted response covers no information OR it takes a completely unreasonable action
-
-**Guidelines**
-1. The reference answer is a reference, not strict ground truth. Be tolerant to different but reasonable answers.
-2. It is fine if the predicted response sounds like an intermediate step rather than a direct answer.
-3. The use case context is Salesforce Agentforce assisting customers.
-4. Output valid JSON only.
-
-[question]: {question}
-[expected text]: {expected}
-[predicted text]: {actual}
-
-Output:
-"""
-    messages = [{"role": "user", "content": prompt}]
-    resp = call_sfap(messages, model, max_tokens=512)
-    if "error" in resp:
-        return None, f"Judge error: {resp['error']}"
+An SFAP token lasts about 2 h. With --token-org <alias>, a refused token is minted again
+with resolve-token.sh during the run; with --cache-dir, a rerun sends only what is unanswered.
 
-    gens = resp.get("generation_details", {}).get("generations", [])
-    if not gens:
-        return None, "No judge response"
+Outputs in --out-dir: units.jsonl, results.jsonl, summary.json, replay_results_sfap.csv,
+and patch_coverage.json with --patch. They contain user utterances: keep them out of git.
+"""
 
-    text = gens[0].get("content", "")
+import argparse
+import hashlib
+import json
+import os
+import sys
+import time
+from collections import Counter
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+from btlib import SCHEMA_UNIT, agentsrc  # noqa: E402
+from btlib import cache as cachelib  # noqa: E402
+from btlib import judge as judgelib  # noqa: E402
+from btlib import patch as patchlib  # noqa: E402
+from btlib import results, scoring, stdm  # noqa: E402
+from btlib import splits as splitlib  # noqa: E402
+from btlib import tokens as tokenlib  # noqa: E402
+from btlib import units as unitlib  # noqa: E402
+from btlib import versions as versionlib  # noqa: E402
+from btlib.redact import redact  # noqa: E402
+from btlib.sfap import DEFAULT_TIMEOUT, SfapClient, build_payload, request_sha  # noqa: E402
+
+AGENT_CONFIG_KEYS = {"router_nodes", "equivalent_tools", "arg_defaults", "description"}
+
+
+def parse_args(argv=None):
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    src = p.add_mutually_exclusive_group(required=True)
+    src.add_argument("--details", action="extend", nargs="+", metavar="FILE",
+                     help="bt_fetch.py's details.json, or getMultipleConversationDetails output "
+                          "(`sf apex run --json` file or plain JSON); one or more files, repeatable")
+    src.add_argument("--from-units", metavar="FILE", help="units.jsonl of an earlier run")
+    p.add_argument("--out-dir", required=True, help="run dir for all outputs (created; must not hold a run)")
+    p.add_argument("--model", default=os.environ.get("REPLAY_MODEL"),
+                   help="candidate model alias, or 'trace' for each step's production model (env REPLAY_MODEL)")
+    p.add_argument("--judge-model", default=os.environ.get("JUDGE_MODEL") or judgelib.DEFAULT_JUDGE_MODEL,
+                   help=f"model for text / free-text-argument equivalence (env JUDGE_MODEL; default "
+                        f"{judgelib.DEFAULT_JUDGE_MODEL}); must differ from the candidate")
+    p.add_argument("--token-info", metavar="FILE",
+                   help="resolve-token.sh output; otherwise env SFAP_BASE_URL, SFAP_TOKEN_FILE, CORE_TENANT_ID")
+    p.add_argument("--token-org", metavar="ALIAS",
+                   help="when SFAP refuses the token (HTTP 401: it expired), mint a new one with "
+                        "resolve-token.sh --org ALIAS into the token file and send the request again, once")
+    p.add_argument("--steps", choices=("router", "subagent", "all"), default="all",
+                   help="which LLM steps to replay (default all)")
+    p.add_argument("--router-nodes", help="comma-separated router step names (default: agent config, else "
+                                          + ",".join(unitlib.DEFAULT_ROUTER_NODES) + ")")
+    p.add_argument("--agent-config", metavar="FILE",
+                   help='JSON {"router_nodes": [...], "equivalent_tools": [[tool, tool], ...], '
+                        '"arg_defaults": {tool: {param: value}}}')
+    p.add_argument("--org-id", default=os.environ.get("ORG_ID"),
+                   help="recorded in the outputs (env ORG_ID; default: from the core tenant id)")
+    p.add_argument("--agent", default=os.environ.get("AGENT_API_NAME"),
+                   help="agent API name recorded in the outputs (env AGENT_API_NAME)")
+    p.add_argument("--k", type=int, default=1, help="repeats per unit (seed = --seed + repeat)")
+    p.add_argument("--seed", type=int, default=0)
+    p.add_argument("--limit", type=int, help="replay only the first N units")
+    p.add_argument("--concurrency", type=int, default=4)
+    p.add_argument("--timeout", type=float, default=DEFAULT_TIMEOUT,
+                   help="seconds per SFAP call; above the provider's own limit (~180 s) by default")
+    p.add_argument("--max-retries", type=int, default=3, help="retries after the first attempt; a timeout is "
+                   "retried once")
+    p.add_argument("--cache-dir", metavar="DIR",
+                   help="answer a payload already answered from DIR instead of SFAP, and store every new answer "
+                        "there (judge calls too); the payload as sent is the key, model and seed included")
+    p.add_argument("--agent-versions", metavar="DIR",
+                   help="the retrieved .agent versions (a bundle dir, a dir of bundle dirs or of .agent files): "
+                        "attributes each unit to the versions that render it; writes versions.json")
+    p.add_argument("--agent-version", metavar="LABEL",
+                   help="replay only the sessions this version can have served (needs --agent-versions)")
+    p.add_argument("--val-fraction", type=float, metavar="F",
+                   help="label each session test or val: a share of about F goes to val, by a hash of the session ID")
+    p.add_argument("--split-salt", default=splitlib.DEFAULT_SALT, help="salt of that hash (default %(default)s)")
+    p.add_argument("--split", choices=("all",) + splitlib.SPLITS, default="all",
+                   help="replay only this split (needs --val-fraction)")
+    p.add_argument("--exclude-sessions", metavar="FILE",
+                   help="skip these sessions: a file that names them, e.g. one ID per line, a "
+                        "units.jsonl, a sessions or details JSON, a replay CSV (references/splits.md)")
+    p.add_argument("--patch", metavar="FILE",
+                   help="edit ops (bt.patch/1) applied to the requests before the replay; "
+                        "writes patch_coverage.json and stops if an op matches nothing")
+    p.add_argument("--dry-run", action="store_true",
+                   help="build units.jsonl (and patch_coverage.json) only; no SFAP calls")
+    p.add_argument("--overwrite", action="store_true", help="replace an existing run in --out-dir")
+    args = p.parse_args(argv)
+    if not args.dry_run and not args.model:
+        p.error("--model is required (an alias, or 'trace')")
+    if args.agent_version and not args.agent_versions:
+        p.error("--agent-version needs --agent-versions")
+    if args.val_fraction is not None and not 0 < args.val_fraction < 1:
+        p.error("--val-fraction must be between 0 and 1")
+    if args.split != "all" and args.val_fraction is None:
+        p.error("--split needs --val-fraction (and the --split-salt the units were labelled with)")
+    if args.k < 1 or args.concurrency < 1:
+        p.error("--k and --concurrency must be >= 1")
+    return args
+
+
+def load_agent_config(path):
+    if not path:
+        return {}
+    with open(path) as f:
+        cfg = json.load(f)
+    unknown = set(cfg) - AGENT_CONFIG_KEYS
+    if unknown:
+        raise SystemExit(f"--agent-config: unknown keys {sorted(unknown)}")
+    groups = cfg.get("equivalent_tools") or []
+    if not all(isinstance(g, list) and all(isinstance(n, str) for n in g) for g in groups):
+        raise SystemExit("--agent-config: equivalent_tools must be a list of lists of tool names")
+    defaults = cfg.get("arg_defaults") or {}
+    if not (isinstance(defaults, dict) and all(isinstance(d, dict) for d in defaults.values())):
+        raise SystemExit("--agent-config: arg_defaults must map each tool name to {param: value}")
+    return cfg
+
+
+def sfap_config(token_info):
+    info = {}
+    if token_info:
+        with open(token_info) as f:
+            info = json.load(f)
+        if info.get("error"):
+            raise SystemExit(f"--token-info reports an error: {info['error']} {info.get('hint', '')}")
+    return (info.get("baseUrl") or os.environ.get("SFAP_BASE_URL"),
+            info.get("tokenFile") or os.environ.get("SFAP_TOKEN_FILE"),
+            info.get("coreTenantId") or os.environ.get("CORE_TENANT_ID"))
+
+
+def load_units(args, router_nodes, org_id):
+    """(every unit of the inputs, skipped steps Counter, input paths, duplicate sessions)."""
+    if args.from_units:
+        units = results.read_jsonl(args.from_units)
+        wrong = {u.get("schema") for u in units} - {SCHEMA_UNIT}
+        if wrong:
+            raise SystemExit(f"{args.from_units}: not a units file (schema {sorted(map(str, wrong))})")
+        for u in units:  # the router set of this run decides the kind
+            u["step_kind"] = "router" if u["node"] in router_nodes else "subagent"
+        return units, Counter(), [args.from_units], 0
+    convos, dupes = stdm.load_conversations(args.details)
+    units, skipped = unitlib.build_units(convos, org=org_id, agent=args.agent, router_nodes=router_nodes)
+    return units, skipped, args.details, dupes
+
+
+def attribute_versions(units, args):
+    """(session version sets, versions.json report) with --agent-versions, else (None, None)."""
+    if not args.agent_versions:
+        return None, None
+    try:
+        found = agentsrc.load_versions(args.agent_versions)
+    except (OSError, agentsrc.AgentSourceError) as e:
+        raise SystemExit(f"--agent-versions {args.agent_versions}: {e}")
+    labels = [v.label for v in found]
+    if args.agent_version and args.agent_version not in labels:
+        raise SystemExit(f"--agent-version {args.agent_version}: not among the versions in "
+                         f"{args.agent_versions} ({', '.join(labels)})")
+    session_sets, report = versionlib.attribute(units, found)
+    report["source"] = os.path.basename(os.path.normpath(args.agent_versions))
+    s = report["sessions"]
+    print(f"versions: {report['units_rendered']} of {report['units']} units render from a version in "
+          f"{report['source']} (" + ", ".join(f"{v['label']} {v['units_rendered']}" for v in report["versions"])
+          + f"); sessions {s['total']}: one version {s['one_version']}, several {s['several_versions']}, "
+          f"conflict {s['conflict']}, unknown {s['unknown']}")
+    missing = report["units_not_rendered_by_node"]
+    if missing:
+        where = "; ".join(f"{node} " + ", ".join(f"{check} {n}" for check, n in c.items()) for node, c in missing.items())
+        print(f"WARNING: {report['units'] - report['units_rendered']} units render from none of the versions "
+              f"(node, then the check the closest version failed: {where}); retrieve every version the "
+              "traces may come from")
+    return session_sets, report
+
+
+def label_splits(units, args):
+    """Label every unit test or val (--val-fraction); the split config, else None."""
+    if args.val_fraction is None:
+        return None
+    try:
+        counts = splitlib.label(units, args.val_fraction, args.split_salt)
+    except splitlib.SplitError as e:
+        raise SystemExit(f"--val-fraction: {e}")
+    print(f"split (salt {args.split_salt}, val fraction {args.val_fraction}): "
+          + ", ".join(f"{s} {n} sessions {k} units" for s, (n, k) in counts.items()))
+    return {"val_fraction": args.val_fraction, "salt": args.split_salt, "selected": args.split,
+            "sessions": {s: n for s, (n, _) in counts.items()}}
+
+
+def exclude_sessions(units, skipped, args):
+    """(the units of sessions not listed in --exclude-sessions, its config or None)."""
+    if not args.exclude_sessions:
+        return units, None
+    try:
+        excluded = splitlib.read_session_ids(args.exclude_sessions)
+    except (OSError, splitlib.SplitError) as e:
+        raise SystemExit(f"--exclude-sessions {args.exclude_sessions}: {e}")
+    present = {u["session_id"] for u in units} & excluded
+    keep = [u for u in units if u["session_id"] not in excluded]
+    skipped["excluded_session"] += len(units) - len(keep)
+    print(f"--exclude-sessions: {len(present)} of its {len(excluded)} sessions are in the input; "
+          f"{len(units) - len(keep)} units skipped")
+    return keep, {"file": os.path.basename(args.exclude_sessions), "sessions": len(excluded),
+                  "in_input": len(present)}
+
+
+def select_units(units, skipped, args, session_sets, report):
+    """The units to replay: --steps, then --agent-version, then --split, then --limit."""
+    if args.steps != "all":
+        keep = [u for u in units if u["step_kind"] == args.steps]
+        skipped["not_selected"] += len(units) - len(keep)
+        units = keep
+    if args.agent_version:
+        units, dropped = versionlib.select(units, args.agent_version, session_sets)
+        skipped.update(dropped)
+        unrendered = sum(1 for u in units if not u["compatible_versions"])
+        report.update(selected=args.agent_version, kept=len(units), dropped=dict(dropped),
+                      kept_not_rendered=unrendered)
+        print(f"--agent-version {args.agent_version}: {len(units)} units kept; dropped: "
+              + (", ".join(f"{k} {v}" for k, v in sorted(dropped.items())) or "none"))
+        if unrendered:
+            print(f"WARNING: {unrendered} kept units do not render from {args.agent_version} exactly; "
+                  "a patch may not apply to them (see patch_coverage.json)")
+    if args.split != "all":
+        keep = [u for u in units if u["split"] == args.split]
+        skipped["other_split"] += len(units) - len(keep)
+        units = keep
+    if args.limit:
+        units = units[:args.limit]
+    return units
+
+
+RUN_FILES = ("summary.json", "results.jsonl", "units.jsonl", "replay_results_sfap.csv", "patch_coverage.json",
+             "versions.json", "compare.json", "report.html")
+
+
+def prepare_out_dir(out_dir, overwrite):
+    if os.path.exists(os.path.join(out_dir, "results.jsonl")) and not overwrite:
+        raise SystemExit(f"{out_dir} already holds a run (results.jsonl); use a new --out-dir or --overwrite")
+    if not os.path.isdir(out_dir):
+        os.makedirs(out_dir)
+        os.chmod(out_dir, 0o700)  # user utterances inside
+    elif overwrite:  # the old run's files go first: a re-run that stops early must not pass for it
+        for name in RUN_FILES:
+            if os.path.isfile(os.path.join(out_dir, name)):
+                os.remove(os.path.join(out_dir, name))
+
+
+def warn_loose_mode(path, what, expected):
     try:
-        # Extract JSON from response
-        json_match = re.search(r'\{[^{}]*"score"[^{}]*\}', text, re.DOTALL)
-        if json_match:
-            parsed = json.loads(json_match.group())
-            return float(parsed.get("score", 0)), parsed.get("reasoning", "")
-    except:
+        if os.stat(path).st_mode & 0o077:
+            print(f"WARNING: {what} is readable by other users (expected mode {expected})")
+    except OSError:
         pass
-    return None, f"Could not parse judge response: {text[:200]}"
 
 
-def judge_args(actual_args, expected_args, func_name, question, model):
-    """LLM-as-judge for function arguments via SFAP."""
-    prompt = f"""You are a scoring judge. Compare function call arguments.
-
-Output ONLY valid JSON with "reasoning" and "score" fields.
-
-**Scoring Rubric**
-1. score=1 -- the arguments are semantically equivalent (same values, possibly different format)
-2. score=0.66 -- most arguments match, minor differences in non-critical fields
-3. score=0.33 -- some arguments match but key fields differ
-4. score=0 -- arguments are completely different or wrong
-
-**Guidelines**
-1. Compare semantic meaning, not exact string match.
-2. Order of JSON keys does not matter.
-3. The use case context is Salesforce Agentforce.
-4. Output valid JSON only.
-
-[user question]: {question}
-[function name]: {func_name}
-[expected arguments]: {expected_args}
-[predicted arguments]: {actual_args}
-
-Output:
-"""
-    messages = [{"role": "user", "content": prompt}]
-    resp = call_sfap(messages, model, max_tokens=512)
-    if "error" in resp:
-        return None, f"Judge error: {resp['error']}"
-
-    gens = resp.get("generation_details", {}).get("generations", [])
-    if not gens:
-        return None, "No judge response"
-
-    text = gens[0].get("content", "")
+def apply_patch(units, path, out_dir):
+    """(requests, patch config): the (messages, tools, patch info) to send per unit."""
+    if not path:
+        return [(u["request"]["messages"], u["request"]["tools"], None) for u in units], None
     try:
-        json_match = re.search(r'\{[^{}]*"score"[^{}]*\}', text, re.DOTALL)
-        if json_match:
-            parsed = json.loads(json_match.group())
-            return float(parsed.get("score", 0)), parsed.get("reasoning", "")
-    except:
-        pass
-    return None, f"Could not parse judge response: {text[:200]}"
-
-
-# ---------------------------------------------------------------------------
-# Enriched-output helpers (for HTML report)
-# ---------------------------------------------------------------------------
-TOOL_GROUPS = {
-    'record_lookup': {'IdentifyRecordByName', 'GetRecordDetails', 'QueryRecords'},
-    'record_modify': {'UpdateRecordFields', 'ExtractRecordFieldsAndValuesFromUserInput'},
-    'activity': {'GetActivitiesTimeline', 'GetActivityDetails'},
-    'routing_discovery': {'go_to_product_discovery'},
-    'routing_cart': {'go_to_cart_and_checkout'},
-    'routing_help': {'go_to_store_help'},
-    'routing_off': {'go_to_off_topic'},
-    'shopper_search': {'SearchProducts', 'GetProductDetails', 'GetRecommendations'},
-    'shopper_cart': {'AddToCart', 'ViewCart', 'UpdateCart', 'Checkout'},
-}
-
-
-def names_match_lenient(golden, predicted):
-    if not golden or not predicted:
-        return False
-    if golden.lower() == predicted.lower():
-        return True
-    for group in TOOL_GROUPS.values():
-        if golden in group and predicted in group:
-            return True
-    return False
-
-
-def build_structured_answer(tool_name, tool_args, text_content):
-    answer = {"content": text_content or ""}
-    if tool_name:
-        args = tool_args
-        if isinstance(args, str):
+        ops = patchlib.load_patch(path)
+        patched, coverage = patchlib.patch_units(units, ops)
+    except (OSError, patchlib.PatchError) as e:
+        raise SystemExit(f"--patch {path}: {e}")
+    requests, changed = [], 0
+    for u, (messages, tools, applied) in zip(units, patched):
+        base = request_sha(u["request"]["messages"], u["request"]["tools"])
+        changed += request_sha(messages, tools) != base
+        requests.append((messages, tools, {"base_request_sha": base, "ops": applied}))
+    report = patchlib.coverage_report(os.path.basename(path), coverage, changed, len(units))
+    results.write_json(os.path.join(out_dir, "patch_coverage.json"), report)
+    for c in coverage:
+        print(f"  patch {c['id']:<16} {c['target']:<40} matched {c['matched']}/{c['units_with_target']} "
+              f"{c['status']}" + (f": {c['message']}" if c["message"] else ""))
+    print(f"patch: {changed} of {len(units)} requests changed; coverage in {out_dir}/patch_coverage.json")
+    if report["errors"]:
+        raise SystemExit(f"--patch: {len(report['errors'])} op(s) failed, nothing replayed "
+                         f"(see patch_coverage.json): " + "; ".join(report["errors"]))
+    with open(path, "rb") as f:
+        digest = hashlib.sha256(f.read()).hexdigest()[:16]
+    return requests, {"file": os.path.basename(path), "sha256": digest, "ops": [op.id for op in ops],
+                      "units_changed": changed}
+
+
+def replay(units, requests, args, client, judge, equivalence, out_dir, arg_defaults=None):
+    jobs = []
+    for u, (messages, tools, patch) in zip(units, requests):
+        sha = request_sha(messages, tools)
+        model = u["request"]["model"] if args.model == "trace" else args.model
+        for rep in range(args.k):
+            jobs.append((u, messages, tools, patch, rep, model, args.seed + rep, sha))
+
+    def work(job):
+        u, messages, tools, patch, rep, model, seed, sha = job
+        call = client.generate(build_payload(model, messages, tools, seed=seed))
+        # scored against the unit as captured: a patch edits texts, never tool names or schemas
+        scores = (scoring.score_unit(u, call["response"], judge, equivalence, arg_defaults)
+                  if call["status"] == "ok" else None)
+        return results.make_record(u, rep, model, seed, sha, call, scores, patch=patch)
+
+    records, interrupted, collected = [], False, set()
+    path = os.path.join(out_dir, "results.jsonl")
+    with results.open_private(path) as f:
+        with ThreadPoolExecutor(max_workers=args.concurrency) as pool:
+            futures = [pool.submit(work, j) for j in jobs]
             try:
-                args = json.loads(args)
+                for i, fut in enumerate(as_completed(futures), 1):
+                    collected.add(fut)
+                    rec = fut.result()
+                    records.append(rec)
+                    results.append_jsonl(f, rec)  # partial results survive a crash
+                    print(f"  [{i}/{len(jobs)}] {rec['node'][:28]:<28} r{rec['repeat']} {rec['status']:<11} "
+                          f"{rec['bucket'] or rec['error_kind'] or ''}", flush=True)
+            except KeyboardInterrupt:
+                interrupted = True
+                for fut in futures:
+                    fut.cancel()
+                print("\nInterrupted: waiting for calls in flight, then writing partial results.")
             except Exception:
-                args = {"raw": args} if args else {}
-        answer["tool_calls"] = [{"name": tool_name, "arguments": args or {}}]
-    return answer
-
-
-def classify_golden_type(tool_name):
-    if not tool_name:
-        return "action"
-    lower = tool_name.lower()
-    if any(p in lower for p in ('route_to', 'go_to', 'transfer_to')):
-        return "topic"
-    return "action"
-
-
-def build_report_json(results, csv_rows):
-    rows = []
-    recs = {}
-    for i, (r, cr) in enumerate(zip(results, csv_rows)):
-        golden = build_structured_answer(
-            r.get('original_tool_name'), r.get('original_tool_args'), r.get('original_text'))
-        predicted = build_structured_answer(
-            r.get('replayed_tool_name'), r.get('replayed_tool_args'), r.get('replayed_text'))
-
-        g_name = r.get('original_tool_name')
-        p_name = r.get('replayed_tool_name')
-        name_match = (g_name == p_name) if (g_name and p_name) else (not g_name and not p_name)
-        lenient = names_match_lenient(g_name, p_name) or name_match
-
-        args_score = cr.get('function_args_score')
-        args_is_num = isinstance(args_score, (int, float))
-        args_match = args_is_num and args_score >= 0.8
-        args_lenient = args_is_num and args_score >= 0.5
-
-        tool_alt = ""
-        scoring_err = r.get('replay_error') or ""
-        if g_name and not p_name:
-            scoring_err = f"Expected tool call ({g_name}), got text response"
-        elif not g_name and p_name:
-            scoring_err = f"Expected text, got tool call ({p_name})"
-        elif g_name and p_name and not name_match and lenient:
-            tool_alt = f"Used {p_name} instead of {g_name} (functionally similar)"
-
-        rows.append({
-            'SessionRowIndex': i,
-            'AiAgentSessionId': r['session_id'],
-            'GoldenType': r.get('golden_type', classify_golden_type(g_name)),
-            'Utterance': r['user_utterance'],
-            'GoldenAnswer': json.dumps(golden),
-            'PredictedAnswer': json.dumps(predicted),
-            'FunctionNameAgreement': name_match,
-            'FunctionNameAgreementLenient': lenient,
-            'PassStrict': name_match and args_match,
-            'FunctionArgsAgreement': args_match,
-            'FunctionArgsAgreementLenient': args_lenient,
-            'FunctionArgsScore': args_score if args_is_num else None,
-            'FunctionArgsScoreReason': r.get('args_reasoning', ''),
-            'ToolAltReason': tool_alt,
-            'SchemaViolationNotes': '',
-            'ScoringError': scoring_err,
-            'BaselineModel': cr.get('original_model', ''),
-            'ReplayModel': cr.get('replay_model', REPLAY_MODEL),
-            'OrgId': ORG_ID,
-            'AgentApiName': AGENT_API_NAME,
-        })
-
-        recs[str(i)] = {
-            'messages': r.get('replay_messages', []),
-            'tools': [t['function']['name'] for t in (r.get('replay_tools') or [])],
-            'session_id': r['session_id'],
-        }
-
-    return {
-        'metadata': {
-            'baseline_model': 'openai/gpt-4.1',
-            'replay_model': REPLAY_MODEL,
-            'org_id': ORG_ID,
-            'agent_api_name': AGENT_API_NAME,
-            'timestamp': time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime()),
-            'total_rows': len(rows),
-        },
-        'rows': rows,
-        'recs': recs,
-    }
-
-
-# ---------------------------------------------------------------------------
-# Main pipeline
-# ---------------------------------------------------------------------------
-def main():
-    print("=" * 60)
-    print("SFAP Replay Pipeline")
-    print("=" * 60)
-
-    # 1. Load STDM conversation data
-    print(f"\n[Phase 1] Loading STDM data from {STDM_DATA_FILE}...")
-    with open(STDM_DATA_FILE) as f:
-        logs = json.load(f)['result']['logs']
-    idx = logs.find('DEBUG|STDM_RESULT:')
-    convos = json.loads(logs[idx + len('DEBUG|STDM_RESULT:'):].split('\n')[0].strip())
-    print(f"  Loaded {len(convos)} conversations")
-
-    # Load LLM step details
-    all_llm_steps = []
-    for f_path in LLM_STEPS_FILES:
-        f_path = f_path.strip()
-        try:
-            all_llm_steps.extend(json.load(open(f_path)))
-        except:
-            pass
-    llm_step_map = {s['step_id']: s for s in all_llm_steps}
-    print(f"  Loaded {len(all_llm_steps)} LLM step details")
-
-    # 2. Build replay targets — both router (topic) and action steps per turn
-    print("\n[Phase 1.5] Building replay targets...")
-    replay_targets = []
-
-    for convo in convos:
-        session_id = convo['session_id']
-        for turn_idx, turn in enumerate(convo.get('turns', [])):
-            user_msg = None
-            output_msg = None
-            for msg in turn.get('messages', []):
-                if msg.get('message_type') == 'Input' and msg.get('text'):
-                    user_msg = msg['text'].strip()
-                if msg.get('message_type') == 'Output' and msg.get('text'):
-                    output_msg = msg['text'].strip()
-
-            if not user_msg:
-                continue
-
-            topic = turn.get('topic')
-
-            # Collect all LLM steps with tools, separating router vs action
-            router_step = None
-            action_step = None
-
-            for step in turn.get('steps', []):
-                if step.get('step_type') != 'LLM_STEP' or not step.get('input'):
-                    continue
-                try:
-                    step_input = json.loads(step['input'])
-                except Exception:
-                    continue
-
-                tools_from_input = step_input.get('gen_ai.request.tools', [])
-                if not tools_from_input:
-                    continue
-
-                if _is_routing_tools(tools_from_input):
-                    if not router_step:
-                        router_step = (step, step_input, tools_from_input)
-                else:
-                    if not action_step:
-                        action_step = (step, step_input, tools_from_input)
-
-            # --- Emit router target (golden_type = "topic") ---
-            if router_step:
-                r_step, r_input, r_tools = router_step
-                r_msgs = r_input.get('gen_ai.input.messages', [])
-                r_tool_name, r_tool_args = _extract_tool_call_from_step_output(r_step)
-                replay_targets.append({
-                    'session_id': session_id,
-                    'turn_index': turn_idx,
-                    'user_utterance': user_msg,
-                    'topic': topic,
-                    'golden_type': 'topic',
-                    'original_tool_name': r_tool_name,
-                    'original_tool_args': r_tool_args,
-                    'original_text': output_msg,
-                    'prompt': (llm_step_map.get(r_step.get('step_id', ''), {})
-                               .get('prompt', '')),
-                    'step_name': r_step.get('name', ''),
-                    'step_tools': r_tools,
-                    'step_messages': r_msgs,
-                })
-
-            # --- Emit action target (golden_type = "action") ---
-            if action_step:
-                a_step, a_input, a_tools = action_step
-                a_msgs = a_input.get('gen_ai.input.messages', [])
-                a_tool_name, a_tool_args = _extract_tool_call_from_step_output(a_step)
-
-                # Fallback: fold ACTION_STEP tool call onto this target
-                if not a_tool_name:
-                    for step in turn.get('steps', []):
-                        if step.get('step_type') == 'ACTION_STEP':
-                            try:
-                                si = json.loads(step.get('input', '{}') or '{}')
-                                name = si.get('actionName', '').split('.')[-1]
-                                if name:
-                                    a_tool_name = name
-                                    args = si.get('actionInput', '')
-                                    a_tool_args = json.dumps(args) if isinstance(args, dict) else (args or '')
-                                    break
-                            except Exception:
-                                pass
-
-                # Fallback: check llm_step_map
-                if not a_tool_name:
-                    step_id = a_step.get('step_id')
-                    if step_id in llm_step_map:
-                        resp = llm_step_map[step_id].get('llm_response', '') or ''
-                        name, args, _ = extract_original_tool_call(resp)
-                        if name:
-                            a_tool_name = name
-                            a_tool_args = args
-
-                replay_targets.append({
-                    'session_id': session_id,
-                    'turn_index': turn_idx,
-                    'user_utterance': user_msg,
-                    'topic': topic,
-                    'golden_type': 'action',
-                    'original_tool_name': a_tool_name,
-                    'original_tool_args': a_tool_args,
-                    'original_text': output_msg,
-                    'prompt': (llm_step_map.get(a_step.get('step_id', ''), {})
-                               .get('prompt', '')),
-                    'step_name': a_step.get('name', ''),
-                    'step_tools': a_tools,
-                    'step_messages': a_msgs,
-                })
-
-    topic_count = sum(1 for t in replay_targets if t['golden_type'] == 'topic')
-    action_count = sum(1 for t in replay_targets if t['golden_type'] == 'action')
-    print(f"  Built {len(replay_targets)} replay targets ({topic_count} router, {action_count} action)")
-    for t in replay_targets:
-        print(f"    Turn {t['turn_index']} [{t['golden_type']:6s}]: tool={t['original_tool_name']} user=\"{t['user_utterance'][:55]}\"")
-
-    # 3. Replay via SFAP
-    print(f"\n[Phase 2] Replaying via SFAP (model: {REPLAY_MODEL})...")
-    results = []
-
-    for i, target in enumerate(replay_targets):
-        gtype = target.get('golden_type', '?')
-        print(f"\n  [{i+1}/{len(replay_targets)}] Turn {target['turn_index']} [{gtype}]: \"{target['user_utterance'][:50]}\"")
-
-        # Parse messages from prompt text, or fall back to structured STDM messages
-        messages = parse_prompt_to_messages(target['prompt'])
-        if not messages and target.get('step_messages'):
-            messages = convert_stdm_messages_for_sfap(target['step_messages'])
-        if not messages:
-            print(f"    SKIP: Could not parse prompt into messages")
-            continue
-
-        # Use tool definitions from STDM step input (dynamic-only, no hardcoded fallback)
-        tools = target.get('step_tools')
-        if tools:
-            tools = sanitize_tools_for_sfap(tools)
-
-        print(f"    Messages: {len(messages)}, Tools: {'yes' if tools else 'no'}")
-
-        # Call SFAP
-        try:
-            sfap_resp = call_sfap(messages, REPLAY_MODEL, tools=tools, max_tokens=2048)
-        except Exception as e:
-            print(f"    ERROR: {e}")
-            results.append({**target, 'replayed_tool_name': None, 'replayed_tool_args': None,
-                           'replayed_text': str(e), 'replay_error': str(e),
-                           'replay_messages': messages, 'replay_tools': tools})
-            continue
-
-        if "error" in sfap_resp:
-            print(f"    ERROR: {sfap_resp['error']}")
-            results.append({**target, 'replayed_tool_name': None, 'replayed_tool_args': None,
-                           'replayed_text': sfap_resp.get('body', ''), 'replay_error': sfap_resp['error'],
-                           'replay_messages': messages, 'replay_tools': tools})
-            continue
-
-        # Extract tool call from SFAP response
-        rep_tool_name, rep_tool_args, rep_text = extract_tool_call(sfap_resp)
-
-        print(f"    Original: tool={target['original_tool_name']}, args={str(target['original_tool_args'])[:80]}")
-        print(f"    Replayed: tool={rep_tool_name}, args={str(rep_tool_args)[:80]}")
-        if rep_text:
-            print(f"    Text: {str(rep_text)[:100]}")
-
-        results.append({
-            **target,
-            'replayed_tool_name': rep_tool_name,
-            'replayed_tool_args': rep_tool_args,
-            'replayed_text': rep_text,
-            'replay_error': None,
-            'replay_messages': messages,
-            'replay_tools': tools,
-        })
-
-        # Small delay to avoid rate limiting
-        time.sleep(0.5)
-
-    # 4. Score
-    print(f"\n[Phase 3] Scoring...")
-    csv_rows = []
-
-    for r in results:
-        row = {
-            'session_id': r['session_id'],
-            'utterance': r['user_utterance'],
-            'golden_type': r.get('golden_type', ''),
-            'original_model': 'openai/gpt-4.1',
-            'replay_model': REPLAY_MODEL,
-            'actual_topic': r.get('topic', ''),
-            'replayed_topic': r.get('topic', ''),
-            'topic_match': 1 if r.get('topic') else '',
-            'actual_action': r.get('original_tool_name', ''),
-            'replayed_action': r.get('replayed_tool_name', ''),
-            'function_match_score': '',
-            'function_args_score': '',
-            'actual_outcome': r.get('original_text', ''),
-            'replayed_outcome': r.get('replayed_text', ''),
-            'text_resp_score': '',
-        }
-
-        # Function name match
-        if r.get('original_tool_name') and r.get('replayed_tool_name'):
-            row['function_match_score'] = 1 if r['original_tool_name'] == r['replayed_tool_name'] else 0
-            print(f"  Turn {r['turn_index']} [{r.get('golden_type', '?')}]: func match = {row['function_match_score']} "
-                  f"({r['original_tool_name']} vs {r['replayed_tool_name']})")
-
-            # Judge arguments on exact OR lenient name match
-            lenient = names_match_lenient(r['original_tool_name'], r['replayed_tool_name'])
-            if (row['function_match_score'] == 1 or lenient) and r.get('original_tool_args') and r.get('replayed_tool_args'):
-                label = "exact" if row['function_match_score'] == 1 else "lenient"
-                print(f"    Judging arguments ({label} name match)...")
-                score, reasoning = judge_args(
-                    r['replayed_tool_args'], r['original_tool_args'],
-                    r['original_tool_name'], r['user_utterance'], JUDGE_MODEL
-                )
-                row['function_args_score'] = score if score is not None else ''
-                r['args_reasoning'] = reasoning or ''
-                print(f"    Args score: {score} - {str(reasoning)[:100]}")
-                time.sleep(0.5)
-        elif r.get('original_tool_name') and not r.get('replayed_tool_name'):
-            row['function_match_score'] = 0
-            print(f"  Turn {r['turn_index']}: func match = 0 (expected {r['original_tool_name']}, got text)")
-        elif not r.get('original_tool_name') and r.get('replayed_tool_name'):
-            row['function_match_score'] = 0
-            print(f"  Turn {r['turn_index']}: func match = 0 (expected text, got {r['replayed_tool_name']})")
-
-        # Text response scoring — for turns where replayed model returned text
-        if r.get('original_text') and r.get('replayed_text'):
-            print(f"  Turn {r['turn_index']}: Judging text response...")
-            score, reasoning = judge_text(
-                r['replayed_text'], r['original_text'], r['user_utterance'], JUDGE_MODEL
-            )
-            row['text_resp_score'] = score if score is not None else ''
-            r['text_reasoning'] = reasoning or ''
-            print(f"    Text score: {score} - {str(reasoning)[:100]}")
-            time.sleep(0.5)
-
-        csv_rows.append(row)
-
-    # 5. Write CSV
-    print(f"\n[Phase 4] Writing CSV to {OUTPUT_CSV}...")
-    fieldnames = [
-        'session_id', 'utterance', 'golden_type', 'original_model', 'replay_model',
-        'actual_topic', 'replayed_topic', 'topic_match',
-        'actual_action', 'replayed_action', 'function_match_score',
-        'function_args_score',
-        'actual_outcome', 'replayed_outcome', 'text_resp_score'
-    ]
-    with open(OUTPUT_CSV, 'w', newline='') as f:
-        writer = csv.DictWriter(f, fieldnames=fieldnames)
-        writer.writeheader()
-        writer.writerows(csv_rows)
-
-    # 5b. Write enriched JSON for HTML report
-    print(f"  Writing JSON to {OUTPUT_JSON}...")
-    report_data = build_report_json(results, csv_rows)
-    with open(OUTPUT_JSON, 'w') as f:
-        json.dump(report_data, f, indent=2, default=str)
-
-    # 6. Summary
-    print(f"\n{'=' * 60}")
-    print(f"REPLAY SUMMARY")
-    print(f"{'=' * 60}")
-    print(f"Sessions: {len(set(r['session_id'] for r in results))}")
-    print(f"Total turns replayed: {len(results)}")
-    print(f"Replay model: {REPLAY_MODEL}")
-
-    func_matches = [r['function_match_score'] for r in csv_rows if r['function_match_score'] != '']
-    func_args = [r['function_args_score'] for r in csv_rows if r['function_args_score'] != '']
-    text_scores = [r['text_resp_score'] for r in csv_rows if r['text_resp_score'] != '']
-    topic_matches = [r['topic_match'] for r in csv_rows if r['topic_match'] != '']
-
-    if topic_matches:
-        print(f"\nTopic Agreement:           {sum(topic_matches)/len(topic_matches)*100:.2f}% ({len(topic_matches)} turns)")
-    if func_matches:
-        print(f"Function Name Accuracy:    {sum(func_matches)/len(func_matches):.4f} ({len(func_matches)} turns)")
-    if func_args:
-        print(f"Function Argument Accuracy: {sum(func_args)/len(func_args):.4f} ({len(func_args)} turns)")
-    if text_scores:
-        print(f"Text Response Accuracy:    {sum(text_scores)/len(text_scores):.4f} ({len(text_scores)} turns)")
-
-    # Breakdown by golden_type
-    for gtype in ('topic', 'action'):
-        gt_rows = [r for r in csv_rows if r.get('golden_type') == gtype]
-        if not gt_rows:
-            continue
-        gt_func = [r['function_match_score'] for r in gt_rows if r['function_match_score'] != '']
-        gt_args = [r['function_args_score'] for r in gt_rows if r['function_args_score'] != '']
-        label = "Router" if gtype == "topic" else "Action"
-        print(f"\n  [{label} steps: {len(gt_rows)} targets]")
-        if gt_func:
-            print(f"    Function Name Accuracy:    {sum(gt_func)/len(gt_func):.4f} ({len(gt_func)} turns)")
-        if gt_args:
-            print(f"    Function Argument Accuracy: {sum(gt_args)/len(gt_args):.4f} ({len(gt_args)} turns)")
-
-    print(f"\nCSV saved to: {OUTPUT_CSV}")
-
-
-if __name__ == '__main__':
-    main()
+                for fut in futures:  # a bug: stop spending calls; results.jsonl keeps what finished
+                    fut.cancel()
+                raise
+        # The pool has shut down, so calls that were in flight at an interrupt are done: keep them.
+        for fut in futures:
+            if fut not in collected and not fut.cancelled() and fut.exception() is None:
+                rec = fut.result()
+                records.append(rec)
+                results.append_jsonl(f, rec)
+    order = {u["unit_id"]: i for i, u in enumerate(units)}
+    records.sort(key=lambda r: (order[r["unit_id"]], r["repeat"]))
+    results.write_jsonl(path, records)
+    return records, interrupted
+
+
+def print_summary(summary, out_dir):
+    c, o = summary["counts"], summary["overall"]
+    skipped = ", ".join(f"{k} {v}" for k, v in c["skipped_steps"].items()) or "none"
+    print("\n" + "=" * 64)
+    print(f"units {c['units']} x k={summary['config'].get('k')} | scored {c['scored']} | "
+          f"infra errors {c['infra_error']} ({100 * c['infra_rate']:.1f}%) | skipped steps: {skipped}")
+    if c["infra_by_kind"]:
+        print("  infra errors by kind: " + ", ".join(f"{k} {v}" for k, v in c["infra_by_kind"].items()))
+    if o["n"]:
+        print(f"agreement {o['agreement']:.3f}   strict {o['strict']:.3f}   ({o['match']}/{o['n']} match)")
+        for kind, agg in summary["by_kind"].items():
+            print(f"  {kind:<16} n={agg['n']:<4} agreement {agg['agreement']:.3f}")
+        for split, agg in summary["by_split"].items():
+            print(f"  split {split:<10} n={agg['n']:<4} agreement {agg['agreement']:.3f}")
+        t = summary["topic"]
+        if t["n"]:
+            print(f"  topic (router go_to) {t['match']}/{t['n']}")
+        print("buckets: " + ", ".join(f"{b} {n}" for b, n in summary["by_bucket"].items()))
+    cache = summary["config"].get("cache")
+    if cache:
+        print(f"cache {cache['dir']}: {cache['hits']} answers reused, {cache['calls']} calls made (replay and judge)"
+              + (f"; WARNING: {cache['write_errors']} answers could not be stored" if cache["write_errors"] else ""))
+    refresh = summary["config"].get("token_refresh")
+    if refresh and (refresh["refreshes"] or refresh["failures"]):
+        print(f"SFAP token (--token-org {refresh['org']}): {refresh['refreshes']} new tokens minted"
+              + (f"; {refresh['failures']} refreshes failed (see the WARNING above)" if refresh["failures"] else ""))
+    unauthorized = summary["counts"]["infra_by_kind"].get("http_401")
+    if unauthorized and not refresh:
+        print(f"{unauthorized} calls were refused with HTTP 401: the SFAP token has probably expired (it lasts "
+              "about 2 h). Run resolve-token.sh again before a rerun. Pass --token-org <alias> to refresh it "
+              "during the run, and --cache-dir so that a rerun sends only what is unanswered")
+    if c.get("timed_out_units"):
+        print(f"{c['timed_out_units']} units timed out on every repeat: the model may not finish them (a runaway "
+              "answer). Here they are infra errors, left out of the scores; bt_compare scores them 0 (no answer) "
+              "against a run that answered them on another request or model")
+    print(f"outputs: {out_dir}/ units.jsonl results.jsonl summary.json replay_results_sfap.csv")
+    print(f"report:  python3 {os.path.relpath(os.path.join(os.path.dirname(__file__), 'build_report.py'))} {out_dir}")
+
+
+def main(argv=None):
+    args = parse_args(argv)
+    cfg = load_agent_config(args.agent_config)
+    if args.router_nodes:
+        router_nodes = [n.strip() for n in args.router_nodes.split(",") if n.strip()]
+    else:
+        router_nodes = cfg.get("router_nodes") or list(unitlib.DEFAULT_ROUTER_NODES)
+    base_url, token_file, tenant = sfap_config(args.token_info)
+    org_id = args.org_id or (tenant.rstrip("/").rsplit("/", 1)[-1] if tenant else None)
+
+    prepare_out_dir(args.out_dir, args.overwrite)
+    units, skipped, inputs, dupes = load_units(args, set(router_nodes), org_id)
+    units, exclude_info = exclude_sessions(units, skipped, args)  # another run's sessions: not input at all
+    session_sets, versions_report = attribute_versions(units, args)  # on every unit: a session's other steps count
+    split_info = label_splits(units, args)
+    units = select_units(units, skipped, args, session_sets, versions_report)
+    if versions_report is not None:
+        results.write_json(os.path.join(args.out_dir, "versions.json"), versions_report)
+    results.write_jsonl(os.path.join(args.out_dir, "units.jsonl"), units)
+    kinds = Counter(results.unit_kind(u) for u in units)
+    print(f"{len(units)} units ({', '.join(f'{k} {v}' for k, v in sorted(kinds.items()))}); "
+          f"skipped steps: {dict(skipped) or 'none'}" + (f"; duplicate sessions ignored: {dupes}" if dupes else ""))
+    if not units:
+        raise SystemExit("no units to replay (the skipped steps above say why)")
+    requests, patch_info = apply_patch(units, args.patch, args.out_dir)
+    if args.dry_run:
+        print(f"dry run: wrote {args.out_dir}/units.jsonl" + (" and patch_coverage.json" if args.patch else ""))
+        return 0
+
+    candidates = sorted({u["request"]["model"] for u in units}) if args.model == "trace" else [args.model]
+    try:
+        judgelib.check_judge_model(args.judge_model, candidates)
+        refresher = (tokenlib.TokenRefresher(token_file, args.token_org, base_url, tenant)
+                     if args.token_org and token_file else None)
+        client = SfapClient(base_url, token_file, tenant, timeout=args.timeout, max_retries=args.max_retries,
+                            on_unauthorized=refresher)
+        if args.cache_dir:
+            client = cachelib.CachingClient(client, args.cache_dir)
+    except (ValueError, OSError) as e:
+        raise SystemExit(f"config error: {e}")
+    family = judgelib.same_family(args.judge_model, candidates)
+    if family:
+        print(f"WARNING: judge {args.judge_model} is the same model family as {', '.join(family)}; "
+              "its text verdicts may favour that model")
+    warn_loose_mode(token_file, "the SFAP token file", "0600")
+    if args.cache_dir:
+        warn_loose_mode(args.cache_dir, f"--cache-dir {args.cache_dir}", "0700")
+    judge = judgelib.Judge(client, args.judge_model)
+    equivalence = scoring.Equivalence(cfg.get("equivalent_tools"))
+
+    config = {"model": args.model, "candidate_models": candidates, "judge_model": args.judge_model,
+              "judge_same_family": family, "steps": args.steps, "router_nodes": sorted(router_nodes),
+              "equivalent_tools": cfg.get("equivalent_tools") or [], "arg_defaults": cfg.get("arg_defaults") or {},
+              "k": args.k, "seed": args.seed,
+              "limit": args.limit, "timeout": args.timeout, "max_retries": args.max_retries,
+              "concurrency": args.concurrency, "inputs": [os.path.basename(p) for p in inputs],
+              "org": org_id, "agent": args.agent, "sfap_base_url": base_url, "patch": patch_info,
+              "versions": None if versions_report is None else {
+                  "source": versions_report["source"], "labels": [v["label"] for v in versions_report["versions"]],
+                  "selected": args.agent_version},
+              "split": split_info, "exclude_sessions": exclude_info,
+              "cache": {"dir": os.path.basename(os.path.abspath(args.cache_dir))} if args.cache_dir else None,
+              "token_refresh": {"org": args.token_org} if refresher else None,
+              "code": results.code_fingerprint(),
+              "started": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())}
+    print(f"replaying {len(units)} units x {args.k} on {args.model} (judge {args.judge_model}, "
+          f"concurrency {args.concurrency})")
+    records, interrupted = replay(units, requests, args, client, judge, equivalence, args.out_dir,
+                                  cfg.get("arg_defaults"))
+    config["finished"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
+    config["interrupted"] = interrupted
+    if args.cache_dir:
+        config["cache"].update(client.stats())
+    if refresher:
+        config["token_refresh"].update(refresher.stats())
+    summary = results.summarize(records, config, skipped, unit_count=len(units))
+    results.write_json(os.path.join(args.out_dir, "summary.json"), summary)
+    results.write_csv(os.path.join(args.out_dir, "replay_results_sfap.csv"), records)
+    print_summary(summary, args.out_dir)
+    if interrupted:
+        return 130
+    return 0 if summary["counts"]["scored"] else 1
+
+
+if __name__ == "__main__":
+    try:
+        sys.exit(main())
+    except stdm.StdmError as e:
+        sys.exit(f"STDM input error: {redact(str(e))}")