#!/usr/bin/env python3 """Forenbench evaluation client. Python 3.9+, standard library only. Set DF_BENCH_KEY in the environment. Never put the key in a URL or source file. """ import argparse import json import os import sys import time import threading import subprocess import urllib.error import urllib.parse import urllib.request import uuid from pathlib import Path class NoRedirect(urllib.request.HTTPRedirectHandler): def redirect_request(self, req, fp, code, msg, headers, newurl): return None # Do not forward credentials through redirects. class ApiError(ValueError): def __init__(self, status, message, retry_after=None): self.status = status self.retry_after = int(retry_after) if retry_after and str(retry_after).isdigit() else None suffix = ('；请等待 ' + str(self.retry_after) + ' 秒后使用相同 request_id 重试') if self.retry_after else '' super().__init__('HTTP ' + str(status) + '：' + message + suffix) class Client: def __init__(self, base, key): self.base = base.rstrip('/') self.parsed = urllib.parse.urlsplit(self.base) p = self.parsed if p.scheme not in ('https', 'http') or not p.hostname or p.username or p.password or p.query or p.fragment: raise ValueError('API Base 须为不含凭证或查询参数的 HTTP(S) 地址。') if p.scheme == 'http' and p.hostname not in ('localhost', '127.0.0.1', '::1'): raise ValueError('远程 API 必须使用 HTTPS；HTTP 仅允许本机调试。') if not key or not key.startswith('dfb_') or len(key) != 68: raise ValueError('请设置环境变量 DF_BENCH_KEY 为本轮评测 Key。') self.key = key self.pending_usage = [] self._usage_lock = threading.RLock() self.opener = urllib.request.build_opener(NoRedirect()) def request(self, path, data=None, binary=False): url = self.base + path parsed = urllib.parse.urlsplit(url) if (parsed.scheme, parsed.hostname, parsed.port) != (self.parsed.scheme, self.parsed.hostname, self.parsed.port): raise ValueError('拒绝向其他站点发送评测 Key。') body = None if data is None else json.dumps(data, ensure_ascii=False).encode('utf-8') request = urllib.request.Request(url, data=body, headers={ 'Authorization': 'Bearer ' + self.key, 'Accept': 'application/json', 'Content-Type': 'application/json', 'User-Agent': 'Forenbench-Client/1.4', }) try: with self.opener.open(request, timeout=60) as response: # Evidence downloads are bounded to the platform's 50 MiB file limit. payload = response.read(52 * 1024 * 1024 + 1) if len(payload) > 52 * 1024 * 1024: raise ValueError('响应超过 52 MiB，请联系组织方确认检材。') return payload if binary else json.loads(payload) except urllib.error.HTTPError as error: payload = error.read(8192) try: message = json.loads(payload).get('message', '请求未完成') except (ValueError, AttributeError): message = '请求未完成，响应不是 JSON' retry = error.headers.get('Retry-After') raise ApiError(error.code, message, retry) from None def heartbeat(self, stage=None, message=None, question_id=None): """Report real progress; do not include answers, credentials or private paths.""" data = {key: value for key, value in { 'stage': stage, 'message': message, 'question_id': str(question_id) if question_id else None, }.items() if value is not None} return self.request('/heartbeat', data) def record_model_call(self, model, event_id, provider=None, input_tokens=None, output_tokens=None, duration_ms=None, question_id=None, status='success', cached_input_tokens=None): """Call after one actual provider invocation; reuse event_id if delivery is retried. Pass provider-response token counts. Missing counts stay None, never estimates. Persist event_id before sending so a network retry cannot double-count usage. """ data = {'event_id': str(uuid.UUID(str(event_id))), 'model': model, 'status': status, 'provider': provider, 'input_tokens': input_tokens, 'output_tokens': output_tokens, 'cached_input_tokens': cached_input_tokens, 'duration_ms': duration_ms, 'question_id': str(question_id) if question_id else None} return self.request('/usage', data) def flush_usage(self): """Retry buffered usage with its original event UUID; never replay model calls.""" with self._usage_lock: for payload in list(self.pending_usage): try: self.request('/usage', payload) except (ValueError, OSError, urllib.error.URLError): break self.pending_usage.remove(payload) return len(self.pending_usage) def invoke_model(self, invoke, *, model, provider=None, read_usage=None, question_id=None, event_id=None): """Wrap one user-owned model invocation; the platform never receives prompts/keys. read_usage(response) returns the provider's actual token fields, or omit it. Delivery errors buffer metadata only and never retry/mask the model request. Flush pending_usage before finishing the evaluation. """ event_id = str(uuid.UUID(str(event_id))) if event_id else str(uuid.uuid4()) started = time.perf_counter() response = None status = 'error' try: response = invoke() status = 'success' return response finally: duration_ms = max(0, min(86400000, round((time.perf_counter() - started) * 1000))) counts = {} if status == 'success' and read_usage: try: counts = read_usage(response) or {} except Exception: counts = {} # A usage-parser failure must not destroy a valid model response. payload = {'event_id': event_id, 'model': model, 'provider': provider, 'status': status, 'duration_ms': duration_ms, 'question_id': str(question_id) if question_id else None} for field in ('input_tokens', 'output_tokens', 'cached_input_tokens'): value = counts.get(field) if isinstance(counts, dict) else None payload[field] = value if isinstance(value, int) and not isinstance(value, bool) and 0 <= value <= 100000000 else None if payload['input_tokens'] is None or (payload['cached_input_tokens'] or 0) > payload['input_tokens']: payload['cached_input_tokens'] = None with self._usage_lock: self.pending_usage.append(payload) if self.flush_usage(): print('模型调用用量投递尚未完成；记录保留在 Client.pending_usage，结束前请 flush_usage() 重试。', file=sys.stderr) def run_tool(client, command, stage='analyzing', message='本地工具正在运行', question_id=None): """Run a local tool and report liveness only while that child is alive. Does not submit answers, count model calls, finish the run, or forward the evaluation key into the tool environment. Tool output stays on the terminal. """ if command and command[0] == '--': command = command[1:] if not command: raise ValueError('run-tool 后须使用 -- 指定本地工具命令。') if len(message) > 120: raise ValueError('进展说明不能超过 120 字。') heartbeat_enabled = True next_report_at = 0 def ping(): nonlocal heartbeat_enabled, next_report_at if not heartbeat_enabled or time.monotonic() < next_report_at: return try: client.heartbeat(stage=stage, message=message, question_id=question_id) except (ValueError, OSError, urllib.error.URLError) as error: if any(str(error).startswith('HTTP ' + str(code) + '：') for code in (401, 403, 409)): heartbeat_enabled = False if isinstance(error, ApiError) and error.status == 429: next_report_at = time.monotonic() + max(30, error.retry_after or 30) print('进度暂未同步；本地工具继续运行，请检查工作台与连接。', file=sys.stderr) environment = dict(os.environ) environment.pop('DF_BENCH_KEY', None) process = subprocess.Popen(command, env=environment) try: ping() while True: try: return process.wait(timeout=30) except subprocess.TimeoutExpired: if process.poll() is None: ping() finally: if process.poll() is None: process.terminate() try: process.wait(timeout=5) except subprocess.TimeoutExpired: process.kill() process.wait() def main(argv=None): parser = argparse.ArgumentParser(description='连接智能体取证评测；Key 从 DF_BENCH_KEY 环境变量读取。') parser.add_argument('--base', required=True, help='评测工作台显示的 API Base') commands = parser.add_subparsers(dest='command', required=True) commands.add_parser('connect', help='连接并获取所选题目；开始计时') ping = commands.add_parser('ping', help='上报当前阶段与简短进展，不消耗答题次数') tool = commands.add_parser('run-tool', help='运行本地工具，存活期间每 30 秒心跳；退出时停止，不会结束评测') for reporter in (ping, tool): reporter.add_argument('--stage', choices=['preparing', 'analyzing', 'submitting', 'blocked', 'finishing'], default='analyzing') reporter.add_argument('--message', default='本地工具正在运行' if reporter is tool else '正在分析当前任务', help='120 字内，不含答案、Key 或私有路径') reporter.add_argument('--question-id', type=uuid.UUID) tool.add_argument('tool_command', nargs=argparse.REMAINDER, help='-- 本地工具命令及参数') commands.add_parser('results', help='获取当前结果') commands.add_parser('finish', help='结束评测，之后不能继续提交') question = commands.add_parser('question', help='读取题目') question.add_argument('id', type=uuid.UUID) download = commands.add_parser('download', help='下载题目检材') download.add_argument('id', type=uuid.UUID) download.add_argument('--out', required=True, help='指定新文件路径；不会覆盖已有文件') submit = commands.add_parser('submit', help='提交答案并消耗次数') submit.add_argument('id', type=uuid.UUID) answer = submit.add_mutually_exclusive_group(required=True) answer.add_argument('--answer', help='单选/判断选项键或填空答案') answer.add_argument('--choices', nargs='+', help='多选选项键，例如 --choices A C') submit.add_argument('--evidence', default='', help='简短证据与操作说明') submit.add_argument('--request-id', type=uuid.UUID, help='新答案不填即可生成；网络失败重试时沿用同一 UUID') usage = commands.add_parser('usage', help='上报一次真实模型调用；未知 Token 请省略；结束前上报') usage.add_argument('--event-id', type=uuid.UUID, help='一条模型调用对应一个 UUID，投递重试沿用同一 ID') usage.add_argument('--model', required=True) usage.add_argument('--provider') usage.add_argument('--input-tokens', type=int) usage.add_argument('--output-tokens', type=int) usage.add_argument('--cached-input-tokens', type=int) usage.add_argument('--duration-ms', type=int) usage.add_argument('--question-id', type=uuid.UUID) usage.add_argument('--status', choices=['success', 'error'], default='success') args = parser.parse_args(argv) try: client = Client(args.base, os.environ.get('DF_BENCH_KEY', '')) if args.command == 'connect': result = client.request('/session') elif args.command == 'ping': result = client.heartbeat(stage=args.stage, message=args.message, question_id=args.question_id) elif args.command == 'run-tool': return run_tool(client, args.tool_command, args.stage, args.message, args.question_id) elif args.command == 'results': result = client.request('/results') elif args.command == 'finish': result = client.request('/finish', {}) elif args.command == 'usage': event_id = str(args.event_id or uuid.uuid4()) print('event_id=' + event_id + '（投递失败时请沿用此 ID 重试）', file=sys.stderr) result = client.record_model_call(model=args.model, provider=args.provider, input_tokens=args.input_tokens, output_tokens=args.output_tokens, cached_input_tokens=args.cached_input_tokens, duration_ms=args.duration_ms, event_id=event_id, question_id=args.question_id, status=args.status) elif args.command == 'question': result = client.request('/questions/' + str(args.id)) elif args.command == 'download': destination = Path(args.out) if destination.exists(): raise ValueError('目标文件已存在，请选择其他文件名。') payload = client.request('/questions/' + str(args.id) + '/artifact', binary=True) with destination.open('xb') as file: file.write(payload) result = {'saved_to': str(destination), 'size_bytes': len(payload)} else: request_id = str(args.request_id or uuid.uuid4()) print('request_id=' + request_id + '（失败重试时请沿用）', file=sys.stderr) result = client.request('/questions/' + str(args.id) + '/answers', { 'request_id': request_id, 'answer': args.choices if args.choices is not None else args.answer, 'evidence': args.evidence, }) print(json.dumps(result, ensure_ascii=False, indent=2)) return 0 except KeyboardInterrupt: print('本地工具已停止，评测与已有答案保留。', file=sys.stderr) return 130 except (ValueError, OSError, urllib.error.URLError) as error: print(str(error), file=sys.stderr) return 1 if __name__ == '__main__': sys.exit(main())