79 lines
2.7 KiB
Python
79 lines
2.7 KiB
Python
|
# -*- coding: utf-8 -*-
|
|||
|
|
"""
|
|||
|
|
tests/_common.py — Agent-C 验证脚本共享工具
|
|||
|
|
============================================
|
|||
|
|
|
|||
|
|
- load_llm_env(): 从工作区根 .env 读取 LLM_API_KEY 等配置（只注入内存/子进程环境，
|
|||
|
|
绝不打印、绝不写入任何文件 —— 硬约束）。
|
|||
|
|
- patch_child_env(): 猴子补丁 orchestrator.build_child_env，把 LLM_API_KEY 注入
|
|||
|
|
pi 子进程环境（models.json 里 apiKey 写的是环境变量名 "LLM_API_KEY"）。
|
|||
|
|
- count_node_processes(): 统计本机 node.exe 进程数（攻击测试 d 的残留检查用）。
|
|||
|
|
"""
|
|||
|
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import os
|
|||
|
|
import subprocess
|
|||
|
|
import sys
|
|||
|
|
from pathlib import Path
|
|||
|
|
|
|||
|
|
POC_ROOT = Path(__file__).resolve().parent.parent
|
|||
|
|
sys.path.insert(0, str(POC_ROOT))
|
|||
|
|
|
|||
|
|
WORKSPACE_ROOT = POC_ROOT.parent.parent
|
|||
|
|
ENV_FILE = WORKSPACE_ROOT / ".env"
|
|||
|
|
|
|||
|
|
|
|||
|
|
def load_llm_env() -> dict:
|
|||
|
|
"""解析 .env，返回 {LLM_API_KEY, LLM_BASE_URL, LLM_MODEL}（缺失键为空串）。"""
|
|||
|
|
out = {"LLM_API_KEY": "", "LLM_BASE_URL": "", "LLM_MODEL": ""}
|
|||
|
|
if not ENV_FILE.exists():
|
|||
|
|
return out
|
|||
|
|
for line in ENV_FILE.read_text(encoding="utf-8").splitlines():
|
|||
|
|
line = line.strip()
|
|||
|
|
if not line or line.startswith("#") or "=" not in line:
|
|||
|
|
continue
|
|||
|
|
k, _, v = line.partition("=")
|
|||
|
|
k = k.strip()
|
|||
|
|
if k in out:
|
|||
|
|
out[k] = v.strip().strip('"').strip("'")
|
|||
|
|
return out
|
|||
|
|
|
|||
|
|
|
|||
|
|
def patch_child_env(api_key: str) -> None:
|
|||
|
|
"""让 orchestrator 拉起的 pi 子进程额外拿到 LLM_API_KEY（其余环境清洗逻辑不变）。"""
|
|||
|
|
import orchestrator
|
|||
|
|
|
|||
|
|
orig = orchestrator.build_child_env
|
|||
|
|
|
|||
|
|
def patched(extra=None):
|
|||
|
|
merged = dict(extra or {})
|
|||
|
|
if api_key:
|
|||
|
|
merged["LLM_API_KEY"] = api_key
|
|||
|
|
return orig(merged)
|
|||
|
|
|
|||
|
|
orchestrator.build_child_env = patched
|
|||
|
|
|
|||
|
|
|
|||
|
|
def count_node_processes() -> int:
|
|||
|
|
"""当前 node.exe 进程数（Windows tasklist）。"""
|
|||
|
|
r = subprocess.run(
|
|||
|
|
["tasklist", "/FI", "IMAGENAME eq node.exe", "/NH"],
|
|||
|
|
capture_output=True, text=True, timeout=30,
|
|||
|
|
)
|
|||
|
|
return sum(1 for line in r.stdout.splitlines() if line.lower().startswith("node.exe"))
|
|||
|
|
|
|||
|
|
|
|||
|
|
def find_pi_node_processes() -> list:
|
|||
|
|
"""命令行里含 pi-coding-agent 的 node 进程（wmic 不可用时退化为全量 node 列表）。"""
|
|||
|
|
try:
|
|||
|
|
r = subprocess.run(
|
|||
|
|
["powershell", "-NoProfile", "-Command",
|
|||
|
|
"Get-CimInstance Win32_Process -Filter \"Name='node.exe'\" | "
|
|||
|
|
"Where-Object { $_.CommandLine -like '*pi-coding-agent*' } | "
|
|||
|
|
"Select-Object -ExpandProperty ProcessId"],
|
|||
|
|
capture_output=True, text=True, timeout=30,
|
|||
|
|
)
|
|||
|
|
return [int(x) for x in r.stdout.split() if x.strip().isdigit()]
|
|||
|
|
except Exception:
|
|||
|
|
return []
|