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 []
|