wehub-resource-sync
1035 lines
38 KiB
Python
1035 lines
38 KiB
Python
import json
|
|
import os
|
|
import re
|
|
import shutil
|
|
import signal
|
|
import socket
|
|
import subprocess
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
import pexpect
|
|
from datasets import load_dataset
|
|
from git import Repo
|
|
|
|
DATASET_NAME = os.environ.get("DATASET_NAME", "Contextbench/ContextBench")
|
|
DATASET_SPLIT = os.environ.get("DATASET_SPLIT", "train")
|
|
|
|
|
|
# Parse integer environment variables with validation and defaults.
|
|
def _get_int_env(name: str, default: int, min_value: int = 1) -> int:
|
|
raw = os.environ.get(name, "").strip()
|
|
if not raw:
|
|
return default
|
|
try:
|
|
value = int(raw)
|
|
except ValueError:
|
|
print(f"⚠️ Invalid {name}={raw!r}; falling back to {default}")
|
|
return default
|
|
if value < min_value:
|
|
print(f"⚠️ {name} must be >= {min_value}; falling back to {default}")
|
|
return default
|
|
return value
|
|
|
|
|
|
LEANN_TOP_K = _get_int_env("LEANN_TOP_K", 5, min_value=1)
|
|
USAGE_LOG_FILE = os.environ.get("USAGE_LOG_FILE", "task_session_usage.log").strip()
|
|
TASK_METRICS_FILE = os.environ.get("TASK_METRICS_FILE", "task_metrics.jsonl").strip()
|
|
PREFETCH_STRICT = os.environ.get("PREFETCH_STRICT", "1").strip() != "0"
|
|
|
|
|
|
# Get uncommitted changes compared to the current HEAD commit.
|
|
def get_git_diff(repo_dir: Path) -> str:
|
|
try:
|
|
repo = Repo(repo_dir)
|
|
return repo.git.diff(repo.head.commit)
|
|
except Exception as e:
|
|
print(f"⚠️ Diff Error: {e}")
|
|
return ""
|
|
|
|
|
|
# Prepare repository state for one ContextBench task instance.
|
|
def setup_task_environment(
|
|
instance_id: str, repo_url: str, work_root: Path, task: Optional[dict] = None
|
|
) -> Path:
|
|
print(f"🚀 [1/4] Preparing environment: {instance_id}")
|
|
if not task:
|
|
ds = load_dataset(DATASET_NAME, split=DATASET_SPLIT)
|
|
task = next((x for x in ds if x["instance_id"] == instance_id), None)
|
|
if not task:
|
|
raise RuntimeError(f"Instance {instance_id} not found in dataset.")
|
|
|
|
target_dir = work_root / instance_id
|
|
if not target_dir.exists():
|
|
Repo.clone_from(repo_url, target_dir)
|
|
|
|
repo = Repo(target_dir)
|
|
repo.git.reset("--hard")
|
|
# Preserve LEANN indexes if present.
|
|
repo.git.clean("-fdx", "-e", ".leann/")
|
|
repo.git.checkout(task["base_commit"])
|
|
(target_dir / "PROBLEM.md").write_text(task["problem_statement"], encoding="utf-8")
|
|
return target_dir
|
|
|
|
|
|
# Pre-download all task repositories before model execution starts.
|
|
def prefetch_task_repositories(tasks: list[dict], work_root: Path) -> None:
|
|
if not tasks:
|
|
print("📦 [Prefetch] No tasks to download.")
|
|
return
|
|
|
|
print(f"📦 [Prefetch] Downloading repositories for {len(tasks)} tasks...")
|
|
failures: list[str] = []
|
|
|
|
for i, task in enumerate(tasks, start=1):
|
|
instance_id = task.get("instance_id", "").strip()
|
|
repo_url = task.get("repo_url", "").strip()
|
|
target_dir = work_root / instance_id
|
|
|
|
if not instance_id or not repo_url:
|
|
failures.append(f"task[{i}] missing instance_id/repo_url")
|
|
print(f" [{i}/{len(tasks)}] ❌ Invalid task metadata; skipping.")
|
|
continue
|
|
|
|
if target_dir.exists():
|
|
# A previous interrupted clone can leave a broken .git dir.
|
|
# Treat only repos with a valid HEAD as reusable.
|
|
try:
|
|
repo = Repo(target_dir)
|
|
repo.git.rev_parse("HEAD")
|
|
print(f" [{i}/{len(tasks)}] ✅ Already present: {instance_id}")
|
|
continue
|
|
except Exception:
|
|
print(f" [{i}/{len(tasks)}] ♻️ Found incomplete repo, re-cloning: {instance_id}")
|
|
shutil.rmtree(target_dir, ignore_errors=True)
|
|
|
|
cloned = False
|
|
last_error: Optional[Exception] = None
|
|
for attempt in range(1, 3 + 1):
|
|
try:
|
|
print(f" [{i}/{len(tasks)}] ⬇️ Cloning: {instance_id} (attempt {attempt}/3)")
|
|
Repo.clone_from(repo_url, target_dir)
|
|
print(f" [{i}/{len(tasks)}] ✅ Cloned: {instance_id}")
|
|
cloned = True
|
|
break
|
|
except Exception as e:
|
|
last_error = e
|
|
print(f" [{i}/{len(tasks)}] ❌ Clone attempt failed: {instance_id} -> {e}")
|
|
shutil.rmtree(target_dir, ignore_errors=True)
|
|
if attempt < 3:
|
|
sleep_s = 4 * attempt
|
|
print(f" [{i}/{len(tasks)}] ⏳ Retrying in {sleep_s}s...")
|
|
time.sleep(sleep_s)
|
|
|
|
if not cloned:
|
|
failures.append(f"{instance_id} ({last_error})")
|
|
|
|
if failures:
|
|
sample = ", ".join(failures[:3])
|
|
extra = f" (+{len(failures) - 3} more)" if len(failures) > 3 else ""
|
|
message = f"Repository prefetch failed for {len(failures)} task(s): {sample}{extra}"
|
|
if PREFETCH_STRICT:
|
|
raise RuntimeError(message)
|
|
print(f"⚠️ {message}")
|
|
print("⚠️ PREFETCH_STRICT=0; continuing with available repositories.")
|
|
print("⚠️ [Prefetch] Completed with partial failures.")
|
|
return
|
|
|
|
print("✅ [Prefetch] All repositories are ready.")
|
|
|
|
|
|
# Wait until a local TCP port starts accepting connections.
|
|
def wait_for_port(port: int, timeout: int = 10):
|
|
start = time.time()
|
|
while time.time() - start < timeout:
|
|
try:
|
|
with socket.create_connection(("127.0.0.1", port), timeout=0.2):
|
|
return
|
|
except OSError:
|
|
time.sleep(0.1)
|
|
print(f"⚠️ Warning: Proxy port {port} did not open in {timeout}s")
|
|
|
|
|
|
# Resolve mitmproxy CA cert path for Python HTTPS clients behind proxy.
|
|
def _resolve_mitm_ca_cert() -> Optional[Path]:
|
|
env_cert = os.environ.get("MITM_CA_CERT_PATH", "").strip()
|
|
candidates = []
|
|
if env_cert:
|
|
candidates.append(Path(env_cert).expanduser())
|
|
candidates.extend(
|
|
[
|
|
Path.home() / ".mitmproxy" / "mitmproxy-ca-cert.pem",
|
|
Path.home() / ".mitmproxy" / "mitmproxy-ca-cert.cer",
|
|
]
|
|
)
|
|
for cert_path in candidates:
|
|
if cert_path.exists():
|
|
return cert_path
|
|
return None
|
|
|
|
|
|
def _resolve_mitmdump_bin() -> str:
|
|
env_bin = os.environ.get("MITMDUMP_BIN", "").strip()
|
|
if env_bin:
|
|
return env_bin
|
|
|
|
# Prefer a dedicated local mitmproxy venv to avoid dependency conflicts
|
|
# with the main project environment.
|
|
project_root = Path(__file__).resolve().parents[1]
|
|
local_bin = project_root / ".mitmproxy-venv" / "bin" / "mitmdump"
|
|
if local_bin.exists():
|
|
return str(local_bin)
|
|
|
|
return "mitmdump"
|
|
|
|
|
|
# Start mitmproxy traffic recorder for the current task.
|
|
def start_mitmproxy(instance_id: str, mitm_script_path: str) -> subprocess.Popen:
|
|
print("🕵️ [2/4] Starting Traffic Recorder...")
|
|
env = os.environ.copy()
|
|
env["TASK_INSTANCE"] = instance_id
|
|
mitmdump_bin = _resolve_mitmdump_bin()
|
|
cmd = [mitmdump_bin, "--no-http2", "-s", mitm_script_path, "-q"]
|
|
try:
|
|
process = subprocess.Popen(cmd, env=env, preexec_fn=os.setsid)
|
|
except FileNotFoundError as exc:
|
|
raise RuntimeError(
|
|
"mitmdump not found. Install mitmproxy in a dedicated env with:\n"
|
|
" python3 -m venv .mitmproxy-venv\n"
|
|
" .mitmproxy-venv/bin/python -m pip install mitmproxy\n"
|
|
"or set MITMDUMP_BIN to an existing mitmdump path."
|
|
) from exc
|
|
wait_for_port(8080)
|
|
return process
|
|
|
|
|
|
# Check whether a server name exists in dict/list MCP server containers.
|
|
def _server_in_container(servers_obj, server_name: str) -> bool:
|
|
if isinstance(servers_obj, dict):
|
|
return server_name in servers_obj
|
|
if isinstance(servers_obj, list):
|
|
for item in servers_obj:
|
|
if isinstance(item, dict) and item.get("name") == server_name:
|
|
return True
|
|
return False
|
|
|
|
|
|
# Recursively scan a config tree for a specific MCP server entry.
|
|
def _has_mcp_server_config(node, server_name: str, parent_key: str = "") -> bool:
|
|
if isinstance(node, dict):
|
|
for key, value in node.items():
|
|
key_l = key.lower()
|
|
if key_l in {"mcpservers", "mcp_servers"} and _server_in_container(value, server_name):
|
|
return True
|
|
if (
|
|
parent_key.lower() == "mcp"
|
|
and key_l == "servers"
|
|
and _server_in_container(value, server_name)
|
|
):
|
|
return True
|
|
if _has_mcp_server_config(value, server_name, key):
|
|
return True
|
|
elif isinstance(node, list):
|
|
for item in node:
|
|
if _has_mcp_server_config(item, server_name, parent_key):
|
|
return True
|
|
return False
|
|
|
|
|
|
# Resolve Claude settings path that defines the required MCP server.
|
|
def resolve_claude_mcp_config(server_name: str) -> Optional[Path]:
|
|
env_path = os.environ.get("CLAUDE_MCP_CONFIG_PATH", "").strip()
|
|
candidates = []
|
|
if env_path:
|
|
candidates.append(Path(env_path).expanduser())
|
|
candidates.extend(
|
|
[
|
|
Path.home() / ".claude" / "settings.json",
|
|
Path.home() / ".config" / "claude" / "settings.json",
|
|
Path.home() / ".config" / "claude-code" / "settings.json",
|
|
Path.home() / ".claude.json",
|
|
]
|
|
)
|
|
|
|
seen = set()
|
|
for cfg_path in candidates:
|
|
cfg_path = cfg_path.resolve()
|
|
if str(cfg_path) in seen:
|
|
continue
|
|
seen.add(str(cfg_path))
|
|
if not cfg_path.exists():
|
|
continue
|
|
try:
|
|
data = json.loads(cfg_path.read_text(encoding="utf-8"))
|
|
except Exception:
|
|
continue
|
|
if _has_mcp_server_config(data, server_name):
|
|
return cfg_path
|
|
return None
|
|
|
|
|
|
def _normalize_mcp_servers(servers_obj) -> dict[str, dict]:
|
|
servers: dict[str, dict] = {}
|
|
if isinstance(servers_obj, dict):
|
|
for name, config in servers_obj.items():
|
|
if isinstance(name, str) and isinstance(config, dict):
|
|
servers[name] = config
|
|
elif isinstance(servers_obj, list):
|
|
for item in servers_obj:
|
|
if not isinstance(item, dict):
|
|
continue
|
|
name = item.get("name")
|
|
if not isinstance(name, str) or not name:
|
|
continue
|
|
config = dict(item)
|
|
config.pop("name", None)
|
|
servers[name] = config
|
|
return servers
|
|
|
|
|
|
def _extract_mcp_servers(node, parent_key: str = "") -> dict[str, dict]:
|
|
servers: dict[str, dict] = {}
|
|
if isinstance(node, dict):
|
|
for key, value in node.items():
|
|
key_l = key.lower()
|
|
if key_l in {"mcpservers", "mcp_servers"}:
|
|
servers.update(_normalize_mcp_servers(value))
|
|
continue
|
|
if parent_key.lower() == "mcp" and key_l == "servers":
|
|
servers.update(_normalize_mcp_servers(value))
|
|
continue
|
|
servers.update(_extract_mcp_servers(value, key))
|
|
elif isinstance(node, list):
|
|
for item in node:
|
|
servers.update(_extract_mcp_servers(item, parent_key))
|
|
return servers
|
|
|
|
|
|
def build_strict_mcp_config_without_server(server_name: str) -> Optional[str]:
|
|
cfg_path = resolve_claude_mcp_config(server_name)
|
|
if not cfg_path:
|
|
return None
|
|
try:
|
|
data = json.loads(cfg_path.read_text(encoding="utf-8"))
|
|
except Exception:
|
|
return None
|
|
servers = _extract_mcp_servers(data)
|
|
servers.pop(server_name, None)
|
|
return json.dumps({"mcpServers": servers}, separators=(",", ":"))
|
|
|
|
|
|
# Decide whether LEANN/MCP integration is available for this task repo.
|
|
def resolve_leann_integration(target_dir: Path) -> dict[str, str]:
|
|
leann_enabled = os.environ.get("LEANN_ENABLED", "1") != "0"
|
|
use_mcp = os.environ.get("LEANN_USE_MCP", "1") != "0"
|
|
mcp_server_name = os.environ.get("LEANN_MCP_SERVER", "leann-server")
|
|
leann_index_exists = (target_dir / ".leann").exists()
|
|
|
|
mode = "none"
|
|
if leann_enabled and leann_index_exists:
|
|
if use_mcp:
|
|
mode = "mcp"
|
|
print(
|
|
f" -> 🔍 LEANN MCP enabled (server: {mcp_server_name}, "
|
|
f"forced top_k={LEANN_TOP_K})"
|
|
)
|
|
else:
|
|
print(" -> ⚠️ LEANN_USE_MCP=0 but CLI mode is disabled; continuing without LEANN")
|
|
elif leann_enabled:
|
|
print(" -> ⚠️ LEANN enabled but no .leann index found; continuing without LEANN")
|
|
|
|
print(mode)
|
|
return {
|
|
"mode": mode,
|
|
"mcp_server_name": mcp_server_name,
|
|
}
|
|
|
|
|
|
# Build the initial Claude prompt from PROBLEM.md plus optional LEANN hints.
|
|
def build_initial_prompt(
|
|
target_dir: Path, leann_info: dict[str, str], instance_id: str = ""
|
|
) -> str:
|
|
problem_file = target_dir / "PROBLEM.md"
|
|
try:
|
|
problem_text = problem_file.read_text(encoding="utf-8").strip()
|
|
if leann_info.get("mode") == "mcp":
|
|
leann_info.get("mcp_server_name", "leann-server")
|
|
mcp_hint = (
|
|
f"Your LEANN index name is: {instance_id}\n"
|
|
"Use LEANN MCP as a cost-aware semantic entry-point router. "
|
|
"Before broad exploration, call leann_search once. "
|
|
"Use top_k=5 by default, top_k=10 only for clearly multi-hop tasks spanning multiple subsystems; never use top_k>10. "
|
|
"Use show_metadata=false.\n"
|
|
"After LEANN, open the top 1-3 likely implementation/source files first, preferring source files over tests/docs/examples/generated files. "
|
|
"Do not open many retrieved files just because they were returned. "
|
|
"After identifying the likely fix location, use targeted Grep to find existing callers, API/server entry points, tests, and config files that may need to be updated. "
|
|
"This targeted Grep is preferred over additional LEANN once concrete symbols, functions, file paths, config keys, or error strings are known.\n"
|
|
"Run at most one additional leann_search only if no plausible implementation entry point is found or a required subsystem is clearly missing. "
|
|
"The second query must be more specific, using concrete identifiers/literals found so far. "
|
|
"Use show_metadata=false unless metadata is necessary to disambiguate files. "
|
|
"Never run more than two leann_search calls total.\n"
|
|
)
|
|
return f"{mcp_hint}\n\n{problem_text}"
|
|
return problem_text
|
|
except Exception as e:
|
|
raise RuntimeError(f"Failed to read problem statement from {problem_file}: {e}")
|
|
|
|
|
|
# Run Claude CLI autonomously in the prepared repository.
|
|
def run_claude_autonomous(
|
|
target_dir: Path,
|
|
model: str,
|
|
leann_info: Optional[dict[str, str]] = None,
|
|
instance_id: str = "",
|
|
):
|
|
print("🤖 [3/4] Launching Claude Code")
|
|
if (target_dir / ".claude").exists():
|
|
shutil.rmtree(target_dir / ".claude")
|
|
|
|
if leann_info is None:
|
|
leann_info = resolve_leann_integration(target_dir)
|
|
server_name = leann_info.get("mcp_server_name", "leann-server")
|
|
if leann_info.get("mode") == "mcp":
|
|
cfg_path = resolve_claude_mcp_config(server_name)
|
|
if not cfg_path:
|
|
raise RuntimeError(
|
|
f"LEANN_USE_MCP=1 but MCP server '{server_name}' was not found in Claude config. "
|
|
"Set CLAUDE_MCP_CONFIG_PATH or configure this server in Claude settings."
|
|
)
|
|
print(f" -> ✅ Claude MCP config found for '{server_name}': {cfg_path}")
|
|
initial_prompt = build_initial_prompt(target_dir, leann_info, instance_id=instance_id)
|
|
|
|
env = os.environ.copy()
|
|
env.update(
|
|
{
|
|
"HTTP_PROXY": "http://127.0.0.1:8080",
|
|
"HTTPS_PROXY": "http://127.0.0.1:8080",
|
|
"NODE_TLS_REJECT_UNAUTHORIZED": "0",
|
|
}
|
|
)
|
|
|
|
claude_args = ["-p", initial_prompt, "--dangerously-skip-permissions"]
|
|
if model:
|
|
claude_args.extend(["--model", model])
|
|
if leann_info.get("mode") != "mcp":
|
|
filtered_mcp_config = build_strict_mcp_config_without_server(server_name)
|
|
if filtered_mcp_config is not None:
|
|
claude_args.extend(
|
|
[
|
|
"--strict-mcp-config",
|
|
"--mcp-config",
|
|
filtered_mcp_config,
|
|
]
|
|
)
|
|
print(
|
|
f" -> 🚫 Non-LEANN mode; removed MCP server '{server_name}' from strict MCP config"
|
|
)
|
|
|
|
child = pexpect.spawn(
|
|
"claude", claude_args, cwd=target_dir, env=env, encoding="utf-8", timeout=1200
|
|
)
|
|
|
|
try:
|
|
while True:
|
|
index = child.expect(
|
|
[r"\[y/n\]", r"Allow execution", pexpect.EOF, pexpect.TIMEOUT], timeout=5
|
|
)
|
|
if index in [0, 1]:
|
|
child.sendline("y")
|
|
elif index == 2:
|
|
break
|
|
except Exception as e:
|
|
print(f"❌ Interaction error: {e}")
|
|
finally:
|
|
if child.isalive():
|
|
child.terminate(force=True)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Trajectory extraction: parse mitmproxy trace → ContextBench traj_data
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _parse_sse_tool_uses(text: str) -> list[dict]:
|
|
"""Reconstruct tool_use calls from a streamed SSE response body."""
|
|
tool_blocks: dict[int, dict] = {} # index -> {name, partial_json}
|
|
completed: list[dict] = []
|
|
|
|
for line in text.split("\n"):
|
|
if not line.startswith("data: "):
|
|
continue
|
|
data_str = line[6:]
|
|
try:
|
|
event = json.loads(data_str)
|
|
except Exception:
|
|
continue
|
|
|
|
event_type = event.get("type", "")
|
|
|
|
if event_type == "content_block_start":
|
|
idx = event.get("index", 0)
|
|
block = event.get("content_block", {})
|
|
if block.get("type") == "tool_use":
|
|
tool_blocks[idx] = {"name": block.get("name", ""), "partial_json": ""}
|
|
|
|
elif event_type == "content_block_delta":
|
|
idx = event.get("index", 0)
|
|
delta = event.get("delta", {})
|
|
if idx in tool_blocks and delta.get("type") == "input_json_delta":
|
|
tool_blocks[idx]["partial_json"] += delta.get("partial_json", "")
|
|
|
|
elif event_type == "content_block_stop":
|
|
idx = event.get("index", 0)
|
|
if idx in tool_blocks:
|
|
tool = tool_blocks.pop(idx)
|
|
try:
|
|
input_data = json.loads(tool["partial_json"]) if tool["partial_json"] else {}
|
|
except Exception:
|
|
input_data = {}
|
|
completed.append({"name": tool["name"], "input": input_data})
|
|
|
|
return completed
|
|
|
|
|
|
def _parse_json_tool_uses(text: str) -> list[dict]:
|
|
"""Extract tool_use calls from a non-streamed JSON response body."""
|
|
try:
|
|
resp = json.loads(text)
|
|
except Exception:
|
|
return []
|
|
completed = []
|
|
for block in resp.get("content", []):
|
|
if block.get("type") == "tool_use":
|
|
completed.append({"name": block.get("name", ""), "input": block.get("input", {})})
|
|
return completed
|
|
|
|
|
|
def _make_relative(file_path: str, target_dir: Path) -> str:
|
|
"""Convert an absolute file path to one relative to target_dir."""
|
|
try:
|
|
return str(Path(file_path).relative_to(target_dir.resolve()))
|
|
except ValueError:
|
|
return file_path
|
|
|
|
|
|
def _is_leann_search_tool(name: str) -> bool:
|
|
return bool(name) and (
|
|
name == "leann_search" or name.endswith("__leann_search") or "leann_search" in name
|
|
)
|
|
|
|
|
|
def _extract_trace_artifacts(
|
|
trace_path: Path,
|
|
target_dir: Path,
|
|
since_timestamp: Optional[float] = None,
|
|
) -> dict:
|
|
"""
|
|
Parse the mitmproxy JSONL trace once and derive both ContextBench trajectory data
|
|
and LEANN usage statistics.
|
|
|
|
Each Anthropic /messages response that contains tool_use calls becomes one pred_step.
|
|
Read calls with line offsets are mapped to line spans; other file-touching tool calls
|
|
(Glob, Grep, Bash cat/head) contribute to pred_files only.
|
|
"""
|
|
pred_steps: list[dict] = []
|
|
leann_search_calls = 0
|
|
|
|
if not trace_path.exists():
|
|
print(f"⚠️ Trace file not found: {trace_path}")
|
|
return {
|
|
"traj_data": _empty_traj_data(),
|
|
"leann_tool_used": False,
|
|
"leann_search_calls": 0,
|
|
}
|
|
|
|
with open(trace_path, encoding="utf-8") as f:
|
|
for raw_line in f:
|
|
try:
|
|
entry = json.loads(raw_line)
|
|
except Exception:
|
|
continue
|
|
|
|
ts = entry.get("timestamp")
|
|
if (
|
|
since_timestamp is not None
|
|
and isinstance(ts, (int, float))
|
|
and ts < since_timestamp
|
|
):
|
|
continue
|
|
|
|
url = entry.get("request", {}).get("url", "")
|
|
if "anthropic.com" not in url or "/messages" not in url:
|
|
continue
|
|
|
|
response_text = entry.get("response", {}).get("text", "") or ""
|
|
if not response_text:
|
|
continue
|
|
|
|
# Support both streaming (SSE) and non-streaming responses.
|
|
if "data: " in response_text:
|
|
tool_uses = _parse_sse_tool_uses(response_text)
|
|
else:
|
|
tool_uses = _parse_json_tool_uses(response_text)
|
|
|
|
if not tool_uses:
|
|
continue
|
|
|
|
step_files: list[str] = []
|
|
step_spans: dict[str, list[dict]] = {}
|
|
|
|
for tool in tool_uses:
|
|
name = tool.get("name", "")
|
|
inp = tool.get("input", {}) or {}
|
|
|
|
if _is_leann_search_tool(name):
|
|
leann_search_calls += 1
|
|
|
|
if name == "Read":
|
|
raw_path = inp.get("file_path", "")
|
|
if not raw_path:
|
|
continue
|
|
rel = _make_relative(raw_path, target_dir)
|
|
if rel not in step_files:
|
|
step_files.append(rel)
|
|
offset = inp.get("offset")
|
|
limit = inp.get("limit")
|
|
if offset is not None:
|
|
start = int(offset)
|
|
end = start + int(limit) if limit is not None else start + 2000
|
|
step_spans.setdefault(rel, []).append({"start": start, "end": end})
|
|
|
|
elif name in ("Grep", "Glob"):
|
|
# These don't return exact spans; record any explicit path argument.
|
|
raw_path = inp.get("path", "")
|
|
if raw_path and not raw_path.endswith(("*", "/")):
|
|
rel = _make_relative(raw_path, target_dir)
|
|
if rel not in step_files:
|
|
step_files.append(rel)
|
|
|
|
elif name == "Bash":
|
|
# Heuristic: detect `cat`, `head`, `tail` calls on files.
|
|
cmd = inp.get("command", "") or ""
|
|
for m in re.finditer(
|
|
r"(?:cat|head|tail|sed|awk)\s+(?:-[^\s]+\s+)*([^\s|><;]+)", cmd
|
|
):
|
|
candidate = m.group(1)
|
|
if "/" in candidate or candidate.endswith(
|
|
(".py", ".go", ".ts", ".js", ".java", ".c", ".cpp", ".rs")
|
|
):
|
|
rel = _make_relative(candidate, target_dir)
|
|
if rel not in step_files:
|
|
step_files.append(rel)
|
|
|
|
if step_files:
|
|
pred_steps.append(
|
|
{
|
|
"files": step_files,
|
|
"spans": step_spans,
|
|
"symbols": {},
|
|
}
|
|
)
|
|
|
|
# Aggregate across all steps.
|
|
all_files: list[str] = []
|
|
all_spans: dict[str, list[dict]] = {}
|
|
for step in pred_steps:
|
|
for f in step["files"]:
|
|
if f not in all_files:
|
|
all_files.append(f)
|
|
for f, spans in step["spans"].items():
|
|
all_spans.setdefault(f, []).extend(spans)
|
|
|
|
return {
|
|
"traj_data": {
|
|
"pred_steps": pred_steps,
|
|
"pred_files": all_files,
|
|
"pred_spans": all_spans,
|
|
"pred_symbols": {},
|
|
},
|
|
"leann_tool_used": leann_search_calls > 0,
|
|
"leann_search_calls": leann_search_calls,
|
|
}
|
|
|
|
|
|
def extract_trajectory_from_traces(
|
|
trace_path: Path,
|
|
target_dir: Path,
|
|
since_timestamp: Optional[float] = None,
|
|
) -> dict:
|
|
return _extract_trace_artifacts(
|
|
trace_path,
|
|
target_dir,
|
|
since_timestamp=since_timestamp,
|
|
)["traj_data"]
|
|
|
|
|
|
def detect_messages_api_error(
|
|
trace_path: Path,
|
|
since_timestamp: Optional[float] = None,
|
|
) -> Optional[str]:
|
|
"""Return a concise API error message if /v1/messages requests failed."""
|
|
if not trace_path.exists():
|
|
return None
|
|
try:
|
|
with open(trace_path, encoding="utf-8") as f:
|
|
for raw_line in f:
|
|
try:
|
|
entry = json.loads(raw_line)
|
|
except Exception:
|
|
continue
|
|
ts = entry.get("timestamp")
|
|
if (
|
|
since_timestamp is not None
|
|
and isinstance(ts, (int, float))
|
|
and ts < since_timestamp
|
|
):
|
|
continue
|
|
|
|
url = entry.get("request", {}).get("url", "")
|
|
if "/v1/messages" not in url:
|
|
continue
|
|
response = entry.get("response", {}) or {}
|
|
status_code = response.get("status_code")
|
|
if not isinstance(status_code, int) or status_code < 400:
|
|
continue
|
|
text = response.get("text", "") or ""
|
|
err_type = ""
|
|
err_message = ""
|
|
try:
|
|
payload = json.loads(text)
|
|
error_obj = payload.get("error", {}) if isinstance(payload, dict) else {}
|
|
if isinstance(error_obj, dict):
|
|
err_type = str(error_obj.get("type", "") or "")
|
|
err_message = str(error_obj.get("message", "") or "")
|
|
except Exception:
|
|
pass
|
|
detail = f"/v1/messages returned HTTP {status_code}"
|
|
if err_type:
|
|
detail += f" ({err_type})"
|
|
if err_message:
|
|
detail += f": {err_message}"
|
|
return detail
|
|
except Exception:
|
|
return None
|
|
return None
|
|
|
|
|
|
def _empty_traj_data() -> dict:
|
|
return {"pred_steps": [], "pred_files": [], "pred_spans": {}, "pred_symbols": {}}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Usage tracking (mirrored from swebench runner)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _blank_usage() -> dict:
|
|
return {
|
|
"input_tokens": 0,
|
|
"output_tokens": 0,
|
|
"cache_creation_tokens": 0,
|
|
"cache_read_tokens": 0,
|
|
"total_tokens": 0,
|
|
"cost_usd": 0.0,
|
|
}
|
|
|
|
|
|
def _normalize_usage(raw: Optional[dict]) -> dict:
|
|
base = _blank_usage()
|
|
if not isinstance(raw, dict):
|
|
return base
|
|
base["input_tokens"] = int(raw.get("inputTokens", 0) or 0)
|
|
base["output_tokens"] = int(raw.get("outputTokens", 0) or 0)
|
|
base["cache_creation_tokens"] = int(raw.get("cacheCreationTokens", 0) or 0)
|
|
base["cache_read_tokens"] = int(raw.get("cacheReadTokens", 0) or 0)
|
|
base["total_tokens"] = int(raw.get("totalTokens", 0) or 0)
|
|
base["cost_usd"] = float(raw.get("totalCost", 0) or 0.0)
|
|
return base
|
|
|
|
|
|
def _usage_delta(before: dict, after: dict) -> dict:
|
|
return {
|
|
"input_tokens": after["input_tokens"] - before["input_tokens"],
|
|
"output_tokens": after["output_tokens"] - before["output_tokens"],
|
|
"cache_creation_tokens": after["cache_creation_tokens"] - before["cache_creation_tokens"],
|
|
"cache_read_tokens": after["cache_read_tokens"] - before["cache_read_tokens"],
|
|
"total_tokens": after["total_tokens"] - before["total_tokens"],
|
|
"cost_usd": round(after["cost_usd"] - before["cost_usd"], 6),
|
|
}
|
|
|
|
|
|
def _has_positive_delta(usage: dict) -> bool:
|
|
return (
|
|
usage["total_tokens"] > 0
|
|
or usage["input_tokens"] > 0
|
|
or usage["output_tokens"] > 0
|
|
or usage["cache_creation_tokens"] > 0
|
|
or usage["cache_read_tokens"] > 0
|
|
or usage["cost_usd"] > 0
|
|
)
|
|
|
|
|
|
def _format_usage(usage: dict) -> str:
|
|
return (
|
|
f"input: {usage['input_tokens']:,}, output: {usage['output_tokens']:,}, "
|
|
f"cache_read: {usage['cache_read_tokens']:,}, cache_create: {usage['cache_creation_tokens']:,} "
|
|
f"total: {usage['total_tokens']:,}, cost: ${usage['cost_usd']:.4f}"
|
|
)
|
|
|
|
|
|
def append_usage_log(line: str) -> None:
|
|
if not USAGE_LOG_FILE:
|
|
return
|
|
log_path = Path(USAGE_LOG_FILE).expanduser()
|
|
try:
|
|
if log_path.parent and str(log_path.parent) != ".":
|
|
log_path.parent.mkdir(parents=True, exist_ok=True)
|
|
with log_path.open("a", encoding="utf-8") as f:
|
|
f.write(f"{line}\n")
|
|
except Exception as e:
|
|
print(f"⚠️ Failed to write usage log {log_path}: {e}")
|
|
|
|
|
|
def append_task_metrics(entry: dict) -> None:
|
|
if not TASK_METRICS_FILE:
|
|
return
|
|
metrics_path = Path(TASK_METRICS_FILE).expanduser()
|
|
try:
|
|
if metrics_path.parent and str(metrics_path.parent) != ".":
|
|
metrics_path.parent.mkdir(parents=True, exist_ok=True)
|
|
with metrics_path.open("a", encoding="utf-8") as f:
|
|
f.write(json.dumps(entry) + "\n")
|
|
except Exception as e:
|
|
print(f"⚠️ Failed to write task metrics {metrics_path}: {e}")
|
|
|
|
|
|
def get_ccusage_sessions() -> Optional[dict[str, dict]]:
|
|
try:
|
|
result = subprocess.run(
|
|
["npx", "ccusage", "session", "--json", "--offline"],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=30,
|
|
)
|
|
if result.returncode != 0:
|
|
stderr = (result.stderr or "").strip()
|
|
raise RuntimeError(stderr or "ccusage returned non-zero exit code")
|
|
data = json.loads(result.stdout or "{}")
|
|
sessions = data.get("sessions", [])
|
|
if not isinstance(sessions, list):
|
|
return {}
|
|
by_session: dict[str, dict] = {}
|
|
for row in sessions:
|
|
if not isinstance(row, dict):
|
|
continue
|
|
session_id = str(row.get("sessionId", "")).strip()
|
|
if not session_id:
|
|
continue
|
|
by_session[session_id] = _normalize_usage(row)
|
|
return by_session
|
|
except Exception as e:
|
|
print(f"⚠️ Failed to get ccusage sessions: {e}")
|
|
return None
|
|
|
|
|
|
def compute_usage_diff(
|
|
before_sessions: Optional[dict[str, dict]],
|
|
after_sessions: Optional[dict[str, dict]],
|
|
instance_id: str,
|
|
) -> Optional[dict]:
|
|
if before_sessions is None or after_sessions is None:
|
|
return None
|
|
|
|
changed_sessions: dict[str, dict] = {}
|
|
all_session_ids = set(before_sessions.keys()) | set(after_sessions.keys())
|
|
for session_id in all_session_ids:
|
|
before = before_sessions.get(session_id, _blank_usage())
|
|
after = after_sessions.get(session_id, _blank_usage())
|
|
delta = _usage_delta(before, after)
|
|
if _has_positive_delta(delta):
|
|
changed_sessions[session_id] = delta
|
|
|
|
if not changed_sessions:
|
|
return None
|
|
|
|
instance_lower = instance_id.lower()
|
|
subagent_ids = [
|
|
sid for sid in changed_sessions if sid.lower() == "subagents" or "subagent" in sid.lower()
|
|
]
|
|
non_subagent_items = [
|
|
(sid, usage) for sid, usage in changed_sessions.items() if sid not in subagent_ids
|
|
]
|
|
matched_non_subagent = [
|
|
(sid, usage)
|
|
for sid, usage in non_subagent_items
|
|
if sid.lower() in instance_lower or instance_lower in sid.lower()
|
|
]
|
|
primary_candidates = matched_non_subagent or non_subagent_items
|
|
if not primary_candidates:
|
|
return None
|
|
|
|
primary_session_id, session_usage = max(
|
|
primary_candidates, key=lambda x: (x[1]["total_tokens"], x[1]["cost_usd"])
|
|
)
|
|
return {
|
|
"session_id": primary_session_id,
|
|
"input_tokens": session_usage["input_tokens"],
|
|
"output_tokens": session_usage["output_tokens"],
|
|
"cache_creation_tokens": session_usage["cache_creation_tokens"],
|
|
"cache_read_tokens": session_usage["cache_read_tokens"],
|
|
"total_tokens": session_usage["total_tokens"],
|
|
"cost_usd": session_usage["cost_usd"],
|
|
"session_usage": session_usage,
|
|
}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Main task runner
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def run_single_task(
|
|
instance_id: str,
|
|
repo_url: str,
|
|
work_root: Path,
|
|
mitm_script_path: str,
|
|
trace_dir: Path,
|
|
model: str = "",
|
|
task: Optional[dict] = None,
|
|
) -> tuple:
|
|
"""
|
|
Run one ContextBench instance end-to-end.
|
|
|
|
Returns: (patch, latency_seconds, agent_seconds, traj_data, usage)
|
|
- patch : git diff string (the model's code changes)
|
|
- latency_seconds : end-to-end wall-clock seconds
|
|
- agent_seconds : Claude agent runtime only
|
|
- traj_data : ContextBench trajectory dict (pred_steps, pred_files, …)
|
|
- usage : token/cost usage dict, or None
|
|
"""
|
|
print(f"\n🌊 Starting Claude Pipeline for {instance_id}")
|
|
work_root = Path(work_root)
|
|
trace_dir = Path(trace_dir)
|
|
mitm_proc = None
|
|
start_time = time.time()
|
|
usage_before: Optional[dict[str, dict]] = None
|
|
usage: Optional[dict] = None
|
|
status = "failed"
|
|
error_message: Optional[str] = None
|
|
claude_start_ts: Optional[float] = None
|
|
claude_end_ts: Optional[float] = None
|
|
leann_mode: Optional[str] = None
|
|
leann_tool_used = False
|
|
leann_search_calls = 0
|
|
|
|
try:
|
|
target_dir = setup_task_environment(instance_id, repo_url, work_root, task=task)
|
|
leann_info = resolve_leann_integration(target_dir)
|
|
leann_mode = leann_info.get("mode")
|
|
mitm_proc = start_mitmproxy(instance_id, str(mitm_script_path))
|
|
|
|
usage_before = get_ccusage_sessions()
|
|
claude_start_ts = time.time()
|
|
run_claude_autonomous(target_dir, model, leann_info=leann_info, instance_id=instance_id)
|
|
claude_end_ts = time.time()
|
|
usage_after = get_ccusage_sessions()
|
|
|
|
trace_path = trace_dir / f"{instance_id}_trace.jsonl"
|
|
api_error = detect_messages_api_error(trace_path, since_timestamp=claude_start_ts)
|
|
if api_error:
|
|
raise RuntimeError(
|
|
f"Claude API request failed for {instance_id}: {api_error}. "
|
|
"Set MODEL/CLAUDE_MODEL to an available model for your account."
|
|
)
|
|
|
|
usage = compute_usage_diff(usage_before, usage_after, instance_id=instance_id)
|
|
if usage:
|
|
task_usage_line = (
|
|
f"💰 Task session usage ({usage['session_id']}) — "
|
|
f"{_format_usage(usage['session_usage'])}"
|
|
)
|
|
print(task_usage_line)
|
|
|
|
print(f"📍 [4/4] Extracting trajectory from trace: {trace_path.name}")
|
|
trace_artifacts = _extract_trace_artifacts(
|
|
trace_path,
|
|
target_dir,
|
|
since_timestamp=claude_start_ts,
|
|
)
|
|
traj_data = trace_artifacts["traj_data"]
|
|
leann_tool_used = bool(trace_artifacts["leann_tool_used"])
|
|
leann_search_calls = int(trace_artifacts["leann_search_calls"])
|
|
print(
|
|
f" -> {len(traj_data['pred_steps'])} steps, "
|
|
f"{len(traj_data['pred_files'])} unique files, "
|
|
f"leann_calls={leann_search_calls}"
|
|
)
|
|
|
|
patch = get_git_diff(target_dir)
|
|
elapsed = time.time() - start_time
|
|
agent_elapsed = (
|
|
max(0.0, claude_end_ts - claude_start_ts)
|
|
if claude_start_ts is not None and claude_end_ts is not None
|
|
else None
|
|
)
|
|
print(f"⏱️ Task completed in {elapsed:.1f}s ({elapsed / 60:.1f}min)")
|
|
status = "succeeded"
|
|
return patch, elapsed, agent_elapsed, traj_data, usage
|
|
except Exception as e:
|
|
error_message = str(e)
|
|
raise
|
|
|
|
finally:
|
|
elapsed = time.time() - start_time
|
|
agent_elapsed = None
|
|
if claude_start_ts is not None:
|
|
agent_end = claude_end_ts if claude_end_ts is not None else time.time()
|
|
agent_elapsed = max(0.0, agent_end - claude_start_ts)
|
|
if usage is None and usage_before is not None:
|
|
usage_after = get_ccusage_sessions()
|
|
usage = compute_usage_diff(usage_before, usage_after, instance_id=instance_id)
|
|
|
|
metrics_entry: dict = {
|
|
"instance_id": instance_id,
|
|
"repo_url": repo_url,
|
|
"model": model or "cli-default",
|
|
"status": status,
|
|
"latency_seconds": round(elapsed, 1),
|
|
"timestamp_unix": int(time.time()),
|
|
"token_usage": usage,
|
|
"leann_mode": leann_mode,
|
|
"leann_tool_used": leann_tool_used,
|
|
"leann_search_calls": leann_search_calls,
|
|
}
|
|
if agent_elapsed is not None:
|
|
metrics_entry["agent_seconds"] = round(agent_elapsed, 1)
|
|
if error_message:
|
|
metrics_entry["error"] = error_message
|
|
append_task_metrics(metrics_entry)
|
|
|
|
metric_line = (
|
|
f"📊 Task metrics ({instance_id}) — status={status}, "
|
|
f"agent={agent_elapsed:.1f}s, latency={elapsed:.1f}s"
|
|
if agent_elapsed is not None
|
|
else f"📊 Task metrics ({instance_id}) — status={status}, latency={elapsed:.1f}s"
|
|
)
|
|
if usage:
|
|
metric_line += (
|
|
f", session={usage['session_id']}, {_format_usage(usage['session_usage'])}"
|
|
)
|
|
else:
|
|
metric_line += ", token usage unavailable"
|
|
metric_line += (
|
|
f", leann_mode={leann_mode}, "
|
|
f"leann_tool_used={str(leann_tool_used).lower()}, "
|
|
f"leann_search_calls={leann_search_calls}"
|
|
)
|
|
if error_message:
|
|
metric_line += f", error={error_message}"
|
|
print(metric_line)
|
|
append_usage_log(metric_line)
|
|
|
|
if mitm_proc:
|
|
os.killpg(os.getpgid(mitm_proc.pid), signal.SIGTERM)
|