初始化仓库：AI 接口自动化测试平台
纳入 FastAPI 后端、Vue 管理端、MCP 桥接与文档；通过 .gitignore 排除本地数据库与构建产物。 Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Cursor
@@ -0,0 +1,953 @@
|
||||
import asyncio
|
||||
import json
|
||||
import re
|
||||
import time
|
||||
from collections import defaultdict, deque
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
|
||||
def render_template_value(value: Any, variables: dict[str, Any]) -> Any:
|
||||
if isinstance(value, str):
|
||||
matches = re.findall(r"\{\{(.*?)\}\}", value)
|
||||
rendered = value
|
||||
for raw_key in matches:
|
||||
key = raw_key.strip()
|
||||
replacement = variables.get(key, "")
|
||||
rendered = rendered.replace(f"{{{{{raw_key}}}}}", str(replacement))
|
||||
return rendered
|
||||
if isinstance(value, dict):
|
||||
return {k: render_template_value(v, variables) for k, v in value.items()}
|
||||
if isinstance(value, list):
|
||||
return [render_template_value(v, variables) for v in value]
|
||||
return value
|
||||
|
||||
|
||||
def build_execution_order(nodes: list[dict[str, Any]], edges: list[dict[str, Any]]) -> list[str]:
|
||||
node_ids = [n["id"] for n in nodes]
|
||||
indegree = {node_id: 0 for node_id in node_ids}
|
||||
graph = defaultdict(list)
|
||||
for edge in edges:
|
||||
source = edge.get("source")
|
||||
target = edge.get("target")
|
||||
if source in indegree and target in indegree:
|
||||
graph[source].append(target)
|
||||
indegree[target] += 1
|
||||
|
||||
queue = deque([node_id for node_id in node_ids if indegree[node_id] == 0])
|
||||
order = []
|
||||
while queue:
|
||||
cur = queue.popleft()
|
||||
order.append(cur)
|
||||
for nxt in graph[cur]:
|
||||
indegree[nxt] -= 1
|
||||
if indegree[nxt] == 0:
|
||||
queue.append(nxt)
|
||||
|
||||
if len(order) != len(node_ids):
|
||||
return node_ids
|
||||
return order
|
||||
|
||||
|
||||
LOOP_BODY_BRANCHES = {"", "body", "loop", "in", "next", "continue"}
|
||||
LOOP_EXIT_BRANCHES = {"", "done", "break", "out", "exit", "default"}
|
||||
|
||||
|
||||
def _safe_eval(left: Any, op: str, right: Any) -> bool:
|
||||
if op == "==":
|
||||
return left == right
|
||||
if op == "!=":
|
||||
return left != right
|
||||
if op == ">":
|
||||
return left > right
|
||||
if op == "<":
|
||||
return left < right
|
||||
if op == ">=":
|
||||
return left >= right
|
||||
if op == "<=":
|
||||
return left <= right
|
||||
return False
|
||||
|
||||
|
||||
async def execute_workflow(
|
||||
definition: dict[str, Any],
|
||||
api_map: dict[int, dict[str, Any]],
|
||||
mock_map: dict[str, dict[str, Any]],
|
||||
base_url: str,
|
||||
fail_fast: bool = True,
|
||||
default_headers: dict[str, Any] | None = None,
|
||||
on_node_start: Callable[[dict[str, Any], int, int], Any] | None = None,
|
||||
on_node_finish: Callable[[dict[str, Any], dict[str, Any], int, int], Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
nodes = definition.get("nodes", [])
|
||||
edges = definition.get("edges", [])
|
||||
variables = dict(definition.get("variables", {}))
|
||||
node_map = {node["id"]: node for node in nodes}
|
||||
order = build_execution_order(nodes, edges)
|
||||
order_index = {node_id: idx for idx, node_id in enumerate(order)}
|
||||
results: list[dict[str, Any]] = []
|
||||
halted = False
|
||||
graph = _build_edge_graph(edges)
|
||||
pending = deque(_sorted_start_nodes(nodes, edges, order_index))
|
||||
queued = set(pending)
|
||||
executed = set()
|
||||
total_nodes = len(order)
|
||||
default_headers = dict(default_headers or {})
|
||||
|
||||
async with httpx.AsyncClient(timeout=20.0) as client:
|
||||
while pending:
|
||||
node_id = pending.popleft()
|
||||
queued.discard(node_id)
|
||||
if node_id in executed or node_id not in node_map:
|
||||
continue
|
||||
node = node_map[node_id]
|
||||
node_type = str(node.get("data", {}).get("type", "http")).lower()
|
||||
await _invoke_callback(on_node_start, node, len(executed) + 1, total_nodes)
|
||||
|
||||
if node_type == "loop":
|
||||
loop_results, result = await _execute_loop_node(
|
||||
node=node,
|
||||
node_map=node_map,
|
||||
graph=graph,
|
||||
edges=edges,
|
||||
order_index=order_index,
|
||||
variables=variables,
|
||||
api_map=api_map,
|
||||
mock_map=mock_map,
|
||||
base_url=base_url,
|
||||
client=client,
|
||||
default_headers=default_headers,
|
||||
fail_fast=fail_fast,
|
||||
on_node_start=on_node_start,
|
||||
on_node_finish=on_node_finish,
|
||||
executed_count=len(executed),
|
||||
total_nodes=total_nodes,
|
||||
)
|
||||
results.extend(loop_results)
|
||||
else:
|
||||
result = await execute_single_node(
|
||||
node, variables, api_map, mock_map, base_url, client, default_headers=default_headers
|
||||
)
|
||||
results.append(result)
|
||||
|
||||
executed.add(node_id)
|
||||
_persist_node_last_run(node, result)
|
||||
await _invoke_callback(on_node_finish, node, result, len(executed), total_nodes)
|
||||
|
||||
if result["status"] == "failed" and fail_fast:
|
||||
halted = True
|
||||
break
|
||||
|
||||
allowed = _allowed_branches(result)
|
||||
next_nodes = []
|
||||
for target, branch in graph.get(node_id, []):
|
||||
if target in executed or target in queued:
|
||||
continue
|
||||
if node_type == "loop":
|
||||
if branch not in LOOP_EXIT_BRANCHES:
|
||||
continue
|
||||
if branch in allowed:
|
||||
next_nodes.append(target)
|
||||
next_nodes.sort(key=lambda item: order_index.get(item, 10**9))
|
||||
for target in next_nodes:
|
||||
pending.append(target)
|
||||
queued.add(target)
|
||||
|
||||
if halted:
|
||||
for node_id in order:
|
||||
if node_id in executed:
|
||||
continue
|
||||
node_type = node_map[node_id].get("data", {}).get("type", "http")
|
||||
skipped_result = {
|
||||
"node_id": node_id,
|
||||
"node_type": node_type,
|
||||
"status": "skipped",
|
||||
"output": {},
|
||||
"error": "halted by previous failure",
|
||||
}
|
||||
results.append(skipped_result)
|
||||
_persist_node_last_run(node_map[node_id], skipped_result)
|
||||
await _invoke_callback(on_node_finish, node_map[node_id], skipped_result, len(executed), total_nodes)
|
||||
|
||||
status = "failed" if any(item["status"] == "failed" for item in results) else "success"
|
||||
return {"status": status, "variables": variables, "results": results, "raw": json.dumps(results)}
|
||||
|
||||
|
||||
async def _invoke_callback(callback: Callable[..., Any] | None, *args: Any) -> None:
|
||||
if callback is None:
|
||||
return
|
||||
result = callback(*args)
|
||||
if asyncio.iscoroutine(result):
|
||||
await result
|
||||
|
||||
|
||||
def _persist_node_last_run(node: dict[str, Any], result: dict[str, Any]) -> None:
|
||||
if not isinstance(node, dict):
|
||||
return
|
||||
data = node.setdefault("data", {})
|
||||
if not isinstance(data, dict):
|
||||
return
|
||||
data["last_run"] = {
|
||||
"status": result.get("status"),
|
||||
"error": result.get("error"),
|
||||
"output": result.get("output", {}),
|
||||
"ts": int(time.time()),
|
||||
}
|
||||
|
||||
|
||||
def _edge_branch(edge: dict[str, Any]) -> str:
|
||||
data = edge.get("data")
|
||||
if isinstance(data, dict):
|
||||
value = str(data.get("branch", "")).strip().lower()
|
||||
if value:
|
||||
return value
|
||||
label = str(edge.get("label", "")).strip().lower()
|
||||
return label
|
||||
|
||||
|
||||
def _build_edge_graph(edges: list[dict[str, Any]]) -> dict[str, list[tuple[str, str]]]:
|
||||
graph: dict[str, list[tuple[str, str]]] = defaultdict(list)
|
||||
for edge in edges:
|
||||
source = edge.get("source")
|
||||
target = edge.get("target")
|
||||
if not source or not target:
|
||||
continue
|
||||
graph[source].append((target, _edge_branch(edge)))
|
||||
return graph
|
||||
|
||||
|
||||
def _sorted_start_nodes(
|
||||
nodes: list[dict[str, Any]], edges: list[dict[str, Any]], order_index: dict[str, int]
|
||||
) -> list[str]:
|
||||
indegree = {node["id"]: 0 for node in nodes if "id" in node}
|
||||
for edge in edges:
|
||||
target = edge.get("target")
|
||||
if target in indegree:
|
||||
indegree[target] += 1
|
||||
starts = [node_id for node_id, degree in indegree.items() if degree == 0]
|
||||
starts.sort(key=lambda item: order_index.get(item, 10**9))
|
||||
return starts
|
||||
|
||||
|
||||
def _allowed_branches(result: dict[str, Any]) -> set[str]:
|
||||
status = result.get("status")
|
||||
node_type = result.get("node_type")
|
||||
output = result.get("output", {})
|
||||
if status == "failed":
|
||||
return {"", "always", "default", "fail", "failed", "failure", "on_fail"}
|
||||
if node_type == "condition":
|
||||
if bool(output.get("result")):
|
||||
return {"", "always", "default", "true", "on_true", "success"}
|
||||
return {"", "always", "default", "false", "on_false"}
|
||||
if node_type == "loop":
|
||||
return set(LOOP_EXIT_BRANCHES) | {"always", "success", "on_success"}
|
||||
return {"", "always", "default", "success", "on_success"}
|
||||
|
||||
|
||||
def _resolve_field_path(source: Any, field: str) -> Any:
|
||||
extracted = source
|
||||
for part in str(field or "").split("."):
|
||||
if not part:
|
||||
continue
|
||||
if isinstance(extracted, dict):
|
||||
extracted = extracted.get(part)
|
||||
else:
|
||||
return None
|
||||
return extracted
|
||||
|
||||
|
||||
def _resolve_condition_operand(node_data: dict[str, Any], side: str, variables: dict[str, Any]) -> Any:
|
||||
mode_key = f"{side}_mode"
|
||||
mode = str(node_data.get(mode_key) or node_data.get("mode", "template")).strip().lower()
|
||||
if mode == "json_path":
|
||||
source_var = str(node_data.get(f"{side}_source_var") or node_data.get("source_var") or "").strip()
|
||||
field = str(node_data.get(f"{side}_field") or node_data.get("field") or "").strip()
|
||||
if not source_var:
|
||||
return None
|
||||
return _resolve_field_path(variables.get(source_var), field)
|
||||
|
||||
raw = node_data.get(side, "")
|
||||
return render_template_value(raw, variables)
|
||||
|
||||
|
||||
def _evaluate_condition(node_data: dict[str, Any], variables: dict[str, Any]) -> bool:
|
||||
left = _resolve_condition_operand(node_data, "left", variables)
|
||||
right = _resolve_condition_operand(node_data, "right", variables)
|
||||
op = str(node_data.get("op", "=="))
|
||||
return _safe_eval(left, op, right)
|
||||
|
||||
|
||||
def _collect_loop_body_node_ids(loop_id: str, graph: dict[str, list[tuple[str, str]]]) -> set[str]:
|
||||
entries = [target for target, branch in graph.get(loop_id, []) if branch in LOOP_BODY_BRANCHES]
|
||||
if not entries:
|
||||
return set()
|
||||
|
||||
body: set[str] = set()
|
||||
queue = deque(entries)
|
||||
while queue:
|
||||
current = queue.popleft()
|
||||
if current == loop_id or current in body:
|
||||
continue
|
||||
body.add(current)
|
||||
for target, _branch in graph.get(current, []):
|
||||
if target != loop_id and target not in body:
|
||||
queue.append(target)
|
||||
return body
|
||||
|
||||
|
||||
async def _execute_node_subgraph(
|
||||
*,
|
||||
node_ids: set[str],
|
||||
entry_node_ids: list[str],
|
||||
node_map: dict[str, dict[str, Any]],
|
||||
graph: dict[str, list[tuple[str, str]]],
|
||||
edges: list[dict[str, Any]],
|
||||
variables: dict[str, Any],
|
||||
api_map: dict[int, dict[str, Any]],
|
||||
mock_map: dict[str, dict[str, Any]],
|
||||
base_url: str,
|
||||
client: httpx.AsyncClient,
|
||||
default_headers: dict[str, Any],
|
||||
fail_fast: bool,
|
||||
iteration: int,
|
||||
on_node_start: Callable[[dict[str, Any], int, int], Any] | None,
|
||||
on_node_finish: Callable[[dict[str, Any], dict[str, Any], int, int], Any] | None,
|
||||
executed_count: int,
|
||||
total_nodes: int,
|
||||
) -> tuple[list[dict[str, Any]], bool]:
|
||||
if not node_ids:
|
||||
return [], False
|
||||
|
||||
subgraph_nodes = [node_map[node_id] for node_id in node_ids if node_id in node_map]
|
||||
subgraph_edges = [
|
||||
edge
|
||||
for edge in edges
|
||||
if edge.get("source") in node_ids and edge.get("target") in node_ids
|
||||
]
|
||||
order = build_execution_order(subgraph_nodes, subgraph_edges)
|
||||
order_index = {node_id: idx for idx, node_id in enumerate(order)}
|
||||
|
||||
indegree = {node_id: 0 for node_id in node_ids}
|
||||
for edge in subgraph_edges:
|
||||
target = edge.get("target")
|
||||
if target in indegree:
|
||||
indegree[target] += 1
|
||||
|
||||
starts = [node_id for node_id in entry_node_ids if node_id in node_ids]
|
||||
if not starts:
|
||||
starts = [node_id for node_id, degree in indegree.items() if degree == 0]
|
||||
starts.sort(key=lambda item: order_index.get(item, 10**9))
|
||||
|
||||
pending = deque(starts)
|
||||
queued = set(starts)
|
||||
local_executed: set[str] = set()
|
||||
results: list[dict[str, Any]] = []
|
||||
halted = False
|
||||
|
||||
while pending:
|
||||
node_id = pending.popleft()
|
||||
queued.discard(node_id)
|
||||
if node_id in local_executed or node_id not in node_map:
|
||||
continue
|
||||
|
||||
node = node_map[node_id]
|
||||
await _invoke_callback(on_node_start, node, executed_count + len(local_executed) + 1, total_nodes)
|
||||
result = await execute_single_node(
|
||||
node, variables, api_map, mock_map, base_url, client, default_headers=default_headers
|
||||
)
|
||||
if isinstance(result.get("output"), dict):
|
||||
result["output"]["loop_iteration"] = iteration
|
||||
results.append(result)
|
||||
local_executed.add(node_id)
|
||||
_persist_node_last_run(node, result)
|
||||
await _invoke_callback(on_node_finish, node, result, executed_count + len(local_executed), total_nodes)
|
||||
|
||||
if result["status"] == "failed" and fail_fast:
|
||||
halted = True
|
||||
break
|
||||
|
||||
allowed = _allowed_branches(result)
|
||||
next_nodes = []
|
||||
for target, branch in graph.get(node_id, []):
|
||||
if target not in node_ids or target in local_executed or target in queued:
|
||||
continue
|
||||
if branch in allowed:
|
||||
next_nodes.append(target)
|
||||
next_nodes.sort(key=lambda item: order_index.get(item, 10**9))
|
||||
for target in next_nodes:
|
||||
pending.append(target)
|
||||
queued.add(target)
|
||||
|
||||
return results, halted
|
||||
|
||||
|
||||
async def _execute_loop_node(
|
||||
*,
|
||||
node: dict[str, Any],
|
||||
node_map: dict[str, dict[str, Any]],
|
||||
graph: dict[str, list[tuple[str, str]]],
|
||||
edges: list[dict[str, Any]],
|
||||
order_index: dict[str, int],
|
||||
variables: dict[str, Any],
|
||||
api_map: dict[int, dict[str, Any]],
|
||||
mock_map: dict[str, dict[str, Any]],
|
||||
base_url: str,
|
||||
client: httpx.AsyncClient,
|
||||
default_headers: dict[str, Any],
|
||||
fail_fast: bool,
|
||||
on_node_start: Callable[[dict[str, Any], int, int], Any] | None,
|
||||
on_node_finish: Callable[[dict[str, Any], dict[str, Any], int, int], Any] | None,
|
||||
executed_count: int,
|
||||
total_nodes: int,
|
||||
) -> tuple[list[dict[str, Any]], dict[str, Any]]:
|
||||
node_id = str(node.get("id", ""))
|
||||
node_data = node.get("data", {}) if isinstance(node.get("data"), dict) else {}
|
||||
max_iterations = max(1, min(int(node_data.get("max_iterations", 10) or 10), 1000))
|
||||
iteration_var = str(node_data.get("iteration_var") or "loop_index").strip() or "loop_index"
|
||||
|
||||
body_node_ids = _collect_loop_body_node_ids(node_id, graph)
|
||||
entry_node_ids = [
|
||||
target for target, branch in graph.get(node_id, []) if branch in LOOP_BODY_BRANCHES and target in body_node_ids
|
||||
]
|
||||
entry_node_ids.sort(key=lambda item: order_index.get(item, 10**9))
|
||||
|
||||
if not body_node_ids:
|
||||
return [], {
|
||||
"node_id": node_id,
|
||||
"node_type": "loop",
|
||||
"status": "failed",
|
||||
"output": {"iterations": 0, "reason": "missing loop body branch"},
|
||||
"error": "循环节点需要连接 body 分支到循环体",
|
||||
}
|
||||
|
||||
all_results: list[dict[str, Any]] = []
|
||||
iterations_run = 0
|
||||
continue_loop = True
|
||||
last_error: str | None = None
|
||||
|
||||
while continue_loop and iterations_run < max_iterations:
|
||||
variables[iteration_var] = iterations_run
|
||||
body_results, halted = await _execute_node_subgraph(
|
||||
node_ids=body_node_ids,
|
||||
entry_node_ids=entry_node_ids,
|
||||
node_map=node_map,
|
||||
graph=graph,
|
||||
edges=edges,
|
||||
variables=variables,
|
||||
api_map=api_map,
|
||||
mock_map=mock_map,
|
||||
base_url=base_url,
|
||||
client=client,
|
||||
default_headers=default_headers,
|
||||
fail_fast=fail_fast,
|
||||
iteration=iterations_run,
|
||||
on_node_start=on_node_start,
|
||||
on_node_finish=on_node_finish,
|
||||
executed_count=executed_count + len(all_results),
|
||||
total_nodes=total_nodes,
|
||||
)
|
||||
all_results.extend(body_results)
|
||||
iterations_run += 1
|
||||
|
||||
if halted or any(item.get("status") == "failed" for item in body_results):
|
||||
last_error = next(
|
||||
(str(item.get("error") or "loop body failed") for item in body_results if item.get("status") == "failed"),
|
||||
"loop body failed",
|
||||
)
|
||||
break
|
||||
|
||||
while_cfg = {
|
||||
"left_mode": node_data.get("while_left_mode") or node_data.get("left_mode"),
|
||||
"left": node_data.get("while_left", ""),
|
||||
"left_source_var": node_data.get("while_left_source_var") or node_data.get("left_source_var"),
|
||||
"left_field": node_data.get("while_left_field") or node_data.get("left_field"),
|
||||
"op": node_data.get("while_op", node_data.get("op", "==")),
|
||||
"right_mode": node_data.get("while_right_mode") or node_data.get("right_mode"),
|
||||
"right": node_data.get("while_right", ""),
|
||||
"right_source_var": node_data.get("while_right_source_var"),
|
||||
"right_field": node_data.get("while_right_field"),
|
||||
}
|
||||
has_while_rule = any(
|
||||
str(while_cfg.get(key) or "").strip()
|
||||
for key in ("left", "left_source_var", "left_field", "right", "right_source_var", "right_field")
|
||||
)
|
||||
if not has_while_rule:
|
||||
continue_loop = False
|
||||
else:
|
||||
continue_loop = _evaluate_condition(while_cfg, variables)
|
||||
|
||||
save_as = str(node_data.get("save_as") or f"{node_id}_loop").strip()
|
||||
summary = {
|
||||
"iterations": iterations_run,
|
||||
"max_iterations": max_iterations,
|
||||
"continue_loop": continue_loop,
|
||||
"body_nodes": sorted(body_node_ids),
|
||||
}
|
||||
variables[save_as] = summary
|
||||
|
||||
status = "failed" if last_error else "success"
|
||||
return all_results, {
|
||||
"node_id": node_id,
|
||||
"node_type": "loop",
|
||||
"status": status,
|
||||
"output": summary,
|
||||
"error": last_error,
|
||||
}
|
||||
|
||||
|
||||
async def execute_single_node(
|
||||
node: dict[str, Any],
|
||||
variables: dict[str, Any],
|
||||
api_map: dict[int, dict[str, Any]],
|
||||
mock_map: dict[str, dict[str, Any]],
|
||||
base_url: str,
|
||||
client: httpx.AsyncClient | None = None,
|
||||
default_headers: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
node_id = node.get("id", "")
|
||||
node_data = node.get("data", {})
|
||||
node_type = node_data.get("type", "http")
|
||||
own_client = client is None
|
||||
if own_client:
|
||||
client = httpx.AsyncClient(timeout=20.0)
|
||||
assert client is not None
|
||||
try:
|
||||
# Legacy compatibility: 旧的 mock 节点已废弃 → 视为 http
|
||||
if node_type == "mock":
|
||||
node_type = "http"
|
||||
node_data["type"] = "http"
|
||||
|
||||
if node_type == "http":
|
||||
return await _execute_http_node(
|
||||
node_id=node_id,
|
||||
node_data=node_data,
|
||||
variables=variables,
|
||||
api_map=api_map,
|
||||
mock_map=mock_map,
|
||||
base_url=base_url,
|
||||
client=client,
|
||||
default_headers=default_headers or {},
|
||||
)
|
||||
|
||||
if node_type == "condition":
|
||||
ok = _evaluate_condition(node_data, variables)
|
||||
variables[node_data.get("save_as", f"{node_id}_result")] = ok
|
||||
return {
|
||||
"node_id": node_id,
|
||||
"node_type": node_type,
|
||||
"status": "success",
|
||||
"output": {
|
||||
"result": ok,
|
||||
"left": _resolve_condition_operand(node_data, "left", variables),
|
||||
"right": _resolve_condition_operand(node_data, "right", variables),
|
||||
"op": node_data.get("op", "=="),
|
||||
},
|
||||
"error": None,
|
||||
}
|
||||
|
||||
if node_type == "loop":
|
||||
return {
|
||||
"node_id": node_id,
|
||||
"node_type": node_type,
|
||||
"status": "failed",
|
||||
"output": {},
|
||||
"error": "loop node must be executed by execute_workflow",
|
||||
}
|
||||
|
||||
if node_type in {"start", "end"}:
|
||||
return {
|
||||
"node_id": node_id,
|
||||
"node_type": node_type,
|
||||
"status": "success",
|
||||
"output": {"pass": True},
|
||||
"error": None,
|
||||
}
|
||||
|
||||
if node_type == "extract":
|
||||
source_key = node_data.get("source_var", "")
|
||||
field = node_data.get("field", "")
|
||||
save_as = node_data.get("save_as", "")
|
||||
source = variables.get(source_key, {})
|
||||
extracted = source
|
||||
for part in field.split("."):
|
||||
if not part:
|
||||
continue
|
||||
if isinstance(extracted, dict):
|
||||
extracted = extracted.get(part)
|
||||
else:
|
||||
extracted = None
|
||||
break
|
||||
variables[save_as] = extracted
|
||||
return {
|
||||
"node_id": node_id,
|
||||
"node_type": node_type,
|
||||
"status": "success",
|
||||
"output": {"value": extracted},
|
||||
"error": None,
|
||||
}
|
||||
|
||||
if node_type == "transform":
|
||||
template = node_data.get("template", {})
|
||||
transformed = render_template_value(template, variables)
|
||||
save_as = node_data.get("save_as", f"{node_id}_output")
|
||||
variables[save_as] = transformed
|
||||
return {
|
||||
"node_id": node_id,
|
||||
"node_type": node_type,
|
||||
"status": "success",
|
||||
"output": {"value": transformed},
|
||||
"error": None,
|
||||
}
|
||||
|
||||
if node_type == "delay":
|
||||
seconds = float(node_data.get("seconds", 1))
|
||||
await asyncio.sleep(seconds)
|
||||
return {
|
||||
"node_id": node_id,
|
||||
"node_type": node_type,
|
||||
"status": "success",
|
||||
"output": {"waited": seconds},
|
||||
"error": None,
|
||||
}
|
||||
|
||||
if node_type == "ai":
|
||||
prompt = node_data.get("prompt", "")
|
||||
response = {"message": "AI node reserved", "prompt": prompt}
|
||||
save_as = node_data.get("save_as", f"{node_id}_ai")
|
||||
variables[save_as] = response
|
||||
return {
|
||||
"node_id": node_id,
|
||||
"node_type": node_type,
|
||||
"status": "success",
|
||||
"output": response,
|
||||
"error": None,
|
||||
}
|
||||
|
||||
raise ValueError(f"unsupported node type: {node_type}")
|
||||
except Exception as exc:
|
||||
return {
|
||||
"node_id": node_id,
|
||||
"node_type": node_type,
|
||||
"status": "failed",
|
||||
"output": {},
|
||||
"error": str(exc),
|
||||
}
|
||||
finally:
|
||||
if own_client:
|
||||
await client.aclose()
|
||||
|
||||
|
||||
def _merge_str_dict(a: dict[str, Any] | None, b: dict[str, Any] | None) -> dict[str, Any]:
|
||||
out = dict(a or {})
|
||||
for k, v in (b or {}).items():
|
||||
out[k] = v
|
||||
return out
|
||||
|
||||
|
||||
def _mock_payload_for_overlay(
|
||||
node_data: dict[str, Any], mock_map: dict[str, dict[str, Any]], variables: dict[str, Any]
|
||||
) -> dict[str, Any]:
|
||||
raw: Any = node_data.get("mock_data")
|
||||
if raw is None:
|
||||
raw = node_data.get("mock")
|
||||
if isinstance(raw, str) and raw in mock_map:
|
||||
raw = mock_map.get(raw) or {}
|
||||
name = node_data.get("mock_name")
|
||||
if (not raw or raw == {}) and isinstance(name, str) and name.strip() and name in mock_map:
|
||||
raw = mock_map.get(name) or {}
|
||||
rendered = render_template_value(raw if isinstance(raw, dict) else {}, variables)
|
||||
return rendered if isinstance(rendered, dict) else {}
|
||||
|
||||
|
||||
def _looks_like_response_mock(m: dict[str, Any]) -> bool:
|
||||
if not m:
|
||||
return False
|
||||
if "status_code" in m:
|
||||
return True
|
||||
if m.get("text") is not None or m.get("json") is not None:
|
||||
if not any(k in m for k in ("method", "url", "query", "path_params")):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _has_request_overlay_keys(m: dict[str, Any]) -> bool:
|
||||
return any(k in m for k in ("method", "url", "headers", "query", "body", "path_params"))
|
||||
|
||||
|
||||
def _apply_request_overlay(
|
||||
method: str,
|
||||
url: str,
|
||||
headers: dict[str, Any],
|
||||
params: dict[str, Any],
|
||||
json_body: dict[str, Any],
|
||||
path_params: dict[str, Any],
|
||||
overlay: dict[str, Any],
|
||||
) -> tuple[str, str, dict[str, Any], dict[str, Any], dict[str, Any], dict[str, Any]]:
|
||||
if not overlay:
|
||||
return method, url, headers, params, json_body, path_params
|
||||
if overlay.get("method"):
|
||||
method = str(overlay["method"]).upper()
|
||||
if "url" in overlay and overlay["url"] is not None:
|
||||
url = str(overlay["url"])
|
||||
if isinstance(overlay.get("headers"), dict):
|
||||
headers = _merge_str_dict(headers, overlay["headers"])
|
||||
if isinstance(overlay.get("query"), dict):
|
||||
params = dict(overlay["query"])
|
||||
if isinstance(overlay.get("body"), dict):
|
||||
json_body = dict(overlay["body"])
|
||||
if isinstance(overlay.get("path_params"), dict):
|
||||
path_params = dict(overlay["path_params"])
|
||||
return method, url, headers, params, json_body, path_params
|
||||
|
||||
|
||||
async def _execute_http_node(
|
||||
node_id: str,
|
||||
node_data: dict[str, Any],
|
||||
variables: dict[str, Any],
|
||||
api_map: dict[int, dict[str, Any]],
|
||||
mock_map: dict[str, dict[str, Any]],
|
||||
base_url: str,
|
||||
client: httpx.AsyncClient,
|
||||
default_headers: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
request_cfg = _resolve_http_request(node_data, api_map)
|
||||
mode = str(node_data.get("mode", "http")).lower()
|
||||
overlay_cfg = _mock_payload_for_overlay(node_data, mock_map, variables)
|
||||
|
||||
use_request_mock = bool(node_data.get("use_mock"))
|
||||
legacy_mode_mock = mode == "mock"
|
||||
# 旧 mode=mock：若 overlay 中含请求字段则作为请求 mock；仅有响应字段时不覆盖请求（仍发真实 HTTP）
|
||||
if legacy_mode_mock:
|
||||
use_request_mock = use_request_mock or _has_request_overlay_keys(overlay_cfg)
|
||||
|
||||
method = str(request_cfg.get("method", "GET")).upper()
|
||||
url = render_template_value(request_cfg.get("url", ""), variables)
|
||||
|
||||
dh = render_template_value(dict(default_headers or {}), variables) or {}
|
||||
base_h = render_template_value(request_cfg.get("headers", {}), variables) or {}
|
||||
headers = _merge_str_dict(dh, base_h)
|
||||
params = render_template_value(request_cfg.get("query", {}), variables) or {}
|
||||
json_body = render_template_value(request_cfg.get("body", {}), variables) or {}
|
||||
path_params = render_template_value(request_cfg.get("path_params", {}), variables) or {}
|
||||
|
||||
mock_request_applied = False
|
||||
if use_request_mock and overlay_cfg and _has_request_overlay_keys(overlay_cfg):
|
||||
method, url, headers, params, json_body, path_params = _apply_request_overlay(
|
||||
method, url, headers, params, json_body, path_params, overlay_cfg
|
||||
)
|
||||
mock_request_applied = True
|
||||
|
||||
if base_url and url and not url.startswith("http"):
|
||||
url = f"{base_url.rstrip('/')}/{url.lstrip('/')}"
|
||||
|
||||
for key, value in path_params.items():
|
||||
url = url.replace(f"{{{key}}}", str(value))
|
||||
|
||||
request_payload = {
|
||||
"method": method,
|
||||
"url": url,
|
||||
"headers": headers,
|
||||
"query": params,
|
||||
"body": json_body if method in {"POST", "PUT", "PATCH", "DELETE"} else None,
|
||||
"path_params": path_params,
|
||||
"mock_request_applied": mock_request_applied,
|
||||
}
|
||||
|
||||
response_payload: dict[str, Any] = {}
|
||||
error_message: str | None = None
|
||||
mock_response_used = False
|
||||
started = time.perf_counter()
|
||||
|
||||
try:
|
||||
response = await client.request(
|
||||
method=method,
|
||||
url=url,
|
||||
headers=headers,
|
||||
params=params,
|
||||
json=json_body if method in {"POST", "PUT", "PATCH", "DELETE"} else None,
|
||||
)
|
||||
response_payload = _build_response_payload(response)
|
||||
except Exception as exc:
|
||||
error_message = str(exc)
|
||||
if mode == "auto" and overlay_cfg and _looks_like_response_mock(overlay_cfg):
|
||||
response_payload = _normalize_mock_response(overlay_cfg)
|
||||
response_payload["mock_fallback_reason"] = error_message
|
||||
mock_response_used = True
|
||||
error_message = None
|
||||
|
||||
duration_ms = int((time.perf_counter() - started) * 1000)
|
||||
assertions = _evaluate_http_assertions(node_data.get("expect"), response_payload, error_message)
|
||||
status = "failed" if (error_message or any(item["passed"] is False for item in assertions)) else "success"
|
||||
|
||||
output = {
|
||||
"request": request_payload,
|
||||
"response": response_payload,
|
||||
"mock_used": mock_request_applied or mock_response_used,
|
||||
"duration_ms": duration_ms,
|
||||
"assertions": assertions,
|
||||
}
|
||||
|
||||
save_as = node_data.get("save_as")
|
||||
if save_as:
|
||||
variables[save_as] = response_payload
|
||||
|
||||
if status == "failed" and error_message is None:
|
||||
error_message = "; ".join(item.get("reason", "assertion failed") for item in assertions if item["passed"] is False)
|
||||
|
||||
return {
|
||||
"node_id": node_id,
|
||||
"node_type": "http",
|
||||
"status": status,
|
||||
"output": output,
|
||||
"error": error_message,
|
||||
}
|
||||
|
||||
|
||||
def _resolve_http_request(node_data: dict[str, Any], api_map: dict[int, dict[str, Any]]) -> dict[str, Any]:
|
||||
"""Merge api_map definition (if api_id given) with inline node fields. Inline overrides."""
|
||||
|
||||
base: dict[str, Any] = {}
|
||||
api_id_value = node_data.get("api_id")
|
||||
if api_id_value is not None:
|
||||
try:
|
||||
api_id = int(api_id_value)
|
||||
except (TypeError, ValueError):
|
||||
api_id = None
|
||||
if api_id is not None and api_id in api_map:
|
||||
api_def = api_map[api_id]
|
||||
base = {
|
||||
"method": api_def.get("method", "GET"),
|
||||
"url": api_def.get("url", ""),
|
||||
"headers": dict(api_def.get("headers", {}) or {}),
|
||||
"query": dict(api_def.get("query", {}) or {}),
|
||||
"body": dict(api_def.get("body", {}) or {}),
|
||||
"path_params": dict(api_def.get("path_params", {}) or {}),
|
||||
}
|
||||
|
||||
for field in ("method", "url"):
|
||||
if node_data.get(field):
|
||||
base[field] = node_data[field]
|
||||
for field in ("headers", "query", "body", "path_params"):
|
||||
if isinstance(node_data.get(field), dict):
|
||||
merged = dict(base.get(field, {}) or {})
|
||||
merged.update(node_data[field])
|
||||
base[field] = merged
|
||||
|
||||
base.setdefault("method", "GET")
|
||||
base.setdefault("url", "")
|
||||
base.setdefault("headers", {})
|
||||
base.setdefault("query", {})
|
||||
base.setdefault("body", {})
|
||||
base.setdefault("path_params", {})
|
||||
return base
|
||||
|
||||
|
||||
def _build_response_payload(response: httpx.Response) -> dict[str, Any]:
|
||||
payload: dict[str, Any] = {
|
||||
"status_code": response.status_code,
|
||||
"headers": dict(response.headers),
|
||||
"text": response.text,
|
||||
}
|
||||
try:
|
||||
payload["json"] = response.json()
|
||||
except Exception:
|
||||
payload["json"] = None
|
||||
return payload
|
||||
|
||||
|
||||
def _normalize_mock_response(mock_data: dict[str, Any]) -> dict[str, Any]:
|
||||
if not isinstance(mock_data, dict):
|
||||
return {"status_code": 200, "headers": {}, "text": "", "json": None}
|
||||
body = mock_data.get("json")
|
||||
if body is None and "data" in mock_data:
|
||||
body = mock_data.get("data")
|
||||
text = mock_data.get("text")
|
||||
if text is None and body is not None:
|
||||
try:
|
||||
text = json.dumps(body, ensure_ascii=False)
|
||||
except Exception:
|
||||
text = str(body)
|
||||
return {
|
||||
"status_code": int(mock_data.get("status_code", 200)),
|
||||
"headers": dict(mock_data.get("headers", {}) or {}),
|
||||
"text": text or "",
|
||||
"json": body,
|
||||
"is_mock": True,
|
||||
}
|
||||
|
||||
|
||||
def _evaluate_http_assertions(
|
||||
expect: Any, response_payload: dict[str, Any], error_message: str | None
|
||||
) -> list[dict[str, Any]]:
|
||||
if not isinstance(expect, dict):
|
||||
return []
|
||||
if error_message:
|
||||
return [{"name": "request", "passed": False, "reason": error_message}]
|
||||
|
||||
results: list[dict[str, Any]] = []
|
||||
status_in = expect.get("status_in") or expect.get("status_codes")
|
||||
if isinstance(status_in, list) and status_in:
|
||||
actual = response_payload.get("status_code")
|
||||
passed = actual in status_in
|
||||
results.append(
|
||||
{
|
||||
"name": "status_in",
|
||||
"passed": passed,
|
||||
"expected": status_in,
|
||||
"actual": actual,
|
||||
"reason": None if passed else f"status_code {actual} not in {status_in}",
|
||||
}
|
||||
)
|
||||
|
||||
json_contains = expect.get("json_contains")
|
||||
if isinstance(json_contains, dict):
|
||||
body = response_payload.get("json") or {}
|
||||
passed, reason = _dict_contains(body, json_contains)
|
||||
results.append(
|
||||
{
|
||||
"name": "json_contains",
|
||||
"passed": passed,
|
||||
"expected": json_contains,
|
||||
"actual": body,
|
||||
"reason": None if passed else reason,
|
||||
}
|
||||
)
|
||||
|
||||
text_includes = expect.get("text_includes")
|
||||
if isinstance(text_includes, str):
|
||||
text = response_payload.get("text") or ""
|
||||
passed = text_includes in text
|
||||
results.append(
|
||||
{
|
||||
"name": "text_includes",
|
||||
"passed": passed,
|
||||
"expected": text_includes,
|
||||
"actual": text[:200],
|
||||
"reason": None if passed else f"response text does not contain '{text_includes}'",
|
||||
}
|
||||
)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def _dict_contains(actual: Any, expected: Any) -> tuple[bool, str | None]:
|
||||
if isinstance(expected, dict):
|
||||
if not isinstance(actual, dict):
|
||||
return False, "expected object"
|
||||
for key, value in expected.items():
|
||||
if key not in actual:
|
||||
return False, f"missing key: {key}"
|
||||
ok, reason = _dict_contains(actual[key], value)
|
||||
if not ok:
|
||||
return False, f"{key}: {reason}"
|
||||
return True, None
|
||||
if isinstance(expected, list):
|
||||
if not isinstance(actual, list):
|
||||
return False, "expected list"
|
||||
for item in expected:
|
||||
if item not in actual:
|
||||
return False, f"missing list item: {item}"
|
||||
return True, None
|
||||
return (actual == expected, None if actual == expected else f"expected {expected}, got {actual}")
|
||||
@@ -0,0 +1,126 @@
|
||||
"""Build General / Grafana Loki explore links from workflow run context."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any
|
||||
from urllib.parse import quote
|
||||
|
||||
|
||||
def _parse_run_time(value: str | None) -> datetime:
|
||||
if not value:
|
||||
return datetime.now(timezone.utc)
|
||||
text = str(value).strip()
|
||||
if text.endswith("Z"):
|
||||
text = text[:-1] + "+00:00"
|
||||
try:
|
||||
parsed = datetime.fromisoformat(text)
|
||||
except ValueError:
|
||||
return datetime.now(timezone.utc)
|
||||
if parsed.tzinfo is None:
|
||||
parsed = parsed.replace(tzinfo=timezone.utc)
|
||||
return parsed.astimezone(timezone.utc)
|
||||
|
||||
|
||||
def _url_path_hint(url: str) -> str:
|
||||
text = str(url or "").strip()
|
||||
if not text:
|
||||
return ""
|
||||
if "://" in text:
|
||||
try:
|
||||
from urllib.parse import urlparse
|
||||
|
||||
parsed = urlparse(text)
|
||||
return parsed.path or text
|
||||
except Exception:
|
||||
return text
|
||||
return text.split("?")[0]
|
||||
|
||||
|
||||
def _extract_http_request(node_result: dict[str, Any] | None) -> dict[str, str]:
|
||||
if not isinstance(node_result, dict):
|
||||
return {}
|
||||
output = node_result.get("output") or {}
|
||||
if not isinstance(output, dict):
|
||||
return {}
|
||||
request = output.get("request") or {}
|
||||
if not isinstance(request, dict):
|
||||
return {}
|
||||
method = str(request.get("method") or "GET").upper()
|
||||
url = str(request.get("url") or "")
|
||||
return {"method": method, "url": url, "path": _url_path_hint(url)}
|
||||
|
||||
|
||||
def _build_logql(method: str, url_path: str) -> str:
|
||||
service_label = os.getenv("LOKI_LABEL_SELECTOR", "").strip()
|
||||
if service_label:
|
||||
base = service_label
|
||||
else:
|
||||
service_key = os.getenv("LOKI_SERVICE_LABEL", "app").strip() or "app"
|
||||
service_value = os.getenv("LOKI_SERVICE_VALUE", "").strip() or "ai-auto-test"
|
||||
base = f'{{{service_key}="{service_value}"}}'
|
||||
|
||||
filters: list[str] = [base]
|
||||
if url_path:
|
||||
escaped = url_path.replace('"', '\\"')
|
||||
filters.append(f'|= "{escaped}"')
|
||||
if method and method != "GET":
|
||||
filters.append(f'|= "{method}"')
|
||||
return " ".join(filters)
|
||||
|
||||
|
||||
def _build_grafana_explore_url(logql: str, time_from: datetime, time_to: datetime) -> str:
|
||||
base = (os.getenv("GENERAL_LOKI_EXPLORE_URL") or os.getenv("LOKI_EXPLORE_URL") or "").strip().rstrip("/")
|
||||
if not base:
|
||||
return ""
|
||||
|
||||
datasource = os.getenv("LOKI_DATASOURCE", "Loki").strip() or "Loki"
|
||||
org_id = os.getenv("LOKI_ORG_ID", "1").strip() or "1"
|
||||
from_ms = int(time_from.timestamp() * 1000)
|
||||
to_ms = int(time_to.timestamp() * 1000)
|
||||
left = json.dumps(
|
||||
[str(from_ms), str(to_ms), datasource, {"expr": logql, "refId": "A"}],
|
||||
ensure_ascii=False,
|
||||
)
|
||||
return f"{base}?orgId={quote(org_id)}&left={quote(left)}"
|
||||
|
||||
|
||||
def build_run_loki_link(
|
||||
*,
|
||||
run_created_at: str | None,
|
||||
duration_ms: int = 0,
|
||||
node_result: dict[str, Any] | None = None,
|
||||
padding_seconds: int | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Return explore URL + LogQL for a single HTTP node execution."""
|
||||
|
||||
padding = int(padding_seconds or os.getenv("LOKI_TIME_PADDING_SECONDS", "120") or 120)
|
||||
started = _parse_run_time(run_created_at)
|
||||
ended = started + timedelta(milliseconds=max(duration_ms, 0))
|
||||
time_from = started - timedelta(seconds=padding)
|
||||
time_to = ended + timedelta(seconds=padding)
|
||||
|
||||
http = _extract_http_request(node_result)
|
||||
method = http.get("method", "")
|
||||
url = http.get("url", "")
|
||||
path = http.get("path", "")
|
||||
logql = _build_logql(method, path)
|
||||
explore_url = _build_grafana_explore_url(logql, time_from, time_to)
|
||||
enabled = bool(explore_url)
|
||||
|
||||
hint = "已生成 General Loki 查询链接，将按执行时间窗口与接口路径过滤日志。"
|
||||
if not enabled:
|
||||
hint = "未配置 GENERAL_LOKI_EXPLORE_URL / LOKI_EXPLORE_URL，请在环境变量中设置 General 探索页地址。"
|
||||
|
||||
return {
|
||||
"enabled": enabled,
|
||||
"explore_url": explore_url,
|
||||
"logql": logql,
|
||||
"time_from": time_from.isoformat(),
|
||||
"time_to": time_to.isoformat(),
|
||||
"method": method,
|
||||
"url": url,
|
||||
"hint": hint,
|
||||
}
|
||||
@@ -0,0 +1,220 @@
|
||||
import asyncio
|
||||
import contextlib
|
||||
import os
|
||||
import posixpath
|
||||
import uuid
|
||||
from collections.abc import AsyncIterator
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import paramiko
|
||||
|
||||
from ..schemas import SshProfileOut, SshScriptOut
|
||||
|
||||
PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
|
||||
SSH_SCRIPTS_DIR = PROJECT_ROOT / "data" / "ssh-scripts"
|
||||
|
||||
|
||||
def resolve_script_body(script: SshScriptOut) -> str:
|
||||
if script.source_type == "file":
|
||||
if not script.storage_path:
|
||||
raise ValueError("脚本文件路径为空")
|
||||
file_path = (SSH_SCRIPTS_DIR / script.storage_path).resolve()
|
||||
scripts_root = SSH_SCRIPTS_DIR.resolve()
|
||||
if scripts_root not in file_path.parents and file_path != scripts_root:
|
||||
raise ValueError("脚本文件路径非法")
|
||||
if not file_path.is_file():
|
||||
raise ValueError("脚本文件不存在")
|
||||
return file_path.read_text(encoding="utf-8")
|
||||
body = (script.content or "").strip()
|
||||
if not body:
|
||||
raise ValueError("脚本内容为空")
|
||||
return body
|
||||
|
||||
|
||||
def _build_connect_kwargs(profile: SshProfileOut, password: str) -> dict[str, Any]:
|
||||
if not password.strip():
|
||||
raise ValueError("执行前必须输入密码")
|
||||
|
||||
connect_kwargs: dict[str, Any] = {
|
||||
"hostname": profile.host,
|
||||
"port": int(profile.port or 22),
|
||||
"username": profile.username,
|
||||
"timeout": 15,
|
||||
"banner_timeout": 15,
|
||||
"auth_timeout": 15,
|
||||
}
|
||||
if profile.auth_type == "password":
|
||||
connect_kwargs.update({"password": password, "look_for_keys": False, "allow_agent": False})
|
||||
elif profile.auth_type == "key" and profile.key_path:
|
||||
connect_kwargs.update({"key_filename": profile.key_path, "look_for_keys": False, "allow_agent": False})
|
||||
else:
|
||||
connect_kwargs.update({"look_for_keys": True, "allow_agent": True})
|
||||
return connect_kwargs
|
||||
|
||||
|
||||
def _connect_ssh(profile: SshProfileOut, password: str) -> paramiko.SSHClient:
|
||||
if profile.auth_type == "key" and profile.key_path and not os.path.exists(profile.key_path):
|
||||
raise ValueError("配置的 key_path 不存在")
|
||||
|
||||
client = paramiko.SSHClient()
|
||||
client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
||||
client.connect(**_build_connect_kwargs(profile, password))
|
||||
return client
|
||||
|
||||
|
||||
def _upload_and_build_command(client: paramiko.SSHClient, script_body: str, script_name: str) -> str:
|
||||
run_id = uuid.uuid4().hex[:12]
|
||||
remote_dir = posixpath.join("/tmp", f"qip-script-{run_id}")
|
||||
safe_name = "".join(ch if ch.isalnum() or ch in "._-" else "_" for ch in script_name) or "script.sh"
|
||||
if not safe_name.endswith(".sh"):
|
||||
safe_name = f"{safe_name}.sh"
|
||||
remote_path = posixpath.join(remote_dir, safe_name)
|
||||
|
||||
stdin, stdout, stderr = client.exec_command(f"mkdir -p {remote_dir} && chmod 700 {remote_dir}")
|
||||
exit_status = stdout.channel.recv_exit_status()
|
||||
if exit_status != 0:
|
||||
err = (stderr.read() or b"").decode("utf-8", errors="replace").strip()
|
||||
raise RuntimeError(err or "无法在远端创建临时目录")
|
||||
|
||||
with client.open_sftp() as sftp:
|
||||
with sftp.file(remote_path, "w") as remote_file:
|
||||
remote_file.write(script_body)
|
||||
sftp.chmod(remote_path, 0o700)
|
||||
|
||||
return f"bash {remote_path}"
|
||||
|
||||
|
||||
def _cleanup_remote_session(stdout: Any, stderr: Any, client: paramiko.SSHClient | None) -> None:
|
||||
for stream in (stdout, stderr):
|
||||
channel = getattr(stream, "channel", None)
|
||||
if channel is not None:
|
||||
with contextlib.suppress(Exception):
|
||||
channel.close()
|
||||
if client is not None:
|
||||
with contextlib.suppress(Exception):
|
||||
client.close()
|
||||
|
||||
|
||||
async def stream_script_run(
|
||||
profile: SshProfileOut,
|
||||
script: SshScriptOut,
|
||||
password: str,
|
||||
cancel_event: asyncio.Event | None = None,
|
||||
) -> AsyncIterator[dict[str, Any]]:
|
||||
cancel_event = cancel_event or asyncio.Event()
|
||||
script_body = resolve_script_body(script)
|
||||
client = await asyncio.to_thread(_connect_ssh, profile, password)
|
||||
stdout = None
|
||||
stderr = None
|
||||
|
||||
try:
|
||||
remote_command = await asyncio.to_thread(_upload_and_build_command, client, script_body, script.name)
|
||||
yield {"type": "meta", "command": remote_command, "profile": profile.name, "script": script.name}
|
||||
|
||||
stdin, stdout, stderr = await asyncio.to_thread(client.exec_command, remote_command)
|
||||
if stdin:
|
||||
with contextlib.suppress(Exception):
|
||||
stdin.close()
|
||||
|
||||
while True:
|
||||
if cancel_event.is_set():
|
||||
yield {"type": "stopped", "detail": "用户手动停止", "ok": False, "exit_status": -1}
|
||||
return
|
||||
|
||||
emitted = False
|
||||
if stdout.channel.recv_ready():
|
||||
chunk = await asyncio.to_thread(stdout.channel.recv, 4096)
|
||||
if chunk:
|
||||
emitted = True
|
||||
yield {"type": "stdout", "data": chunk.decode("utf-8", errors="replace")}
|
||||
if stderr.channel.recv_stderr_ready():
|
||||
chunk = await asyncio.to_thread(stderr.channel.recv_stderr, 4096)
|
||||
if chunk:
|
||||
emitted = True
|
||||
yield {"type": "stderr", "data": chunk.decode("utf-8", errors="replace")}
|
||||
if (
|
||||
stdout.channel.exit_status_ready()
|
||||
and not stdout.channel.recv_ready()
|
||||
and not stderr.channel.recv_stderr_ready()
|
||||
):
|
||||
break
|
||||
if not emitted:
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
if cancel_event.is_set():
|
||||
yield {"type": "stopped", "detail": "用户手动停止", "ok": False, "exit_status": -1}
|
||||
return
|
||||
|
||||
exit_status = await asyncio.to_thread(stdout.channel.recv_exit_status)
|
||||
yield {"type": "done", "exit_status": int(exit_status), "ok": exit_status == 0}
|
||||
finally:
|
||||
_cleanup_remote_session(stdout, stderr, client)
|
||||
|
||||
|
||||
async def run_script_collect(
|
||||
profile: SshProfileOut,
|
||||
script: SshScriptOut,
|
||||
password: str,
|
||||
timeout_seconds: float = 120.0,
|
||||
max_output_chars: int = 200_000,
|
||||
) -> dict[str, Any]:
|
||||
timeout_value = max(5.0, min(float(timeout_seconds or 120.0), 600.0))
|
||||
cancel_event = asyncio.Event()
|
||||
stdout_parts: list[str] = []
|
||||
stderr_parts: list[str] = []
|
||||
meta: dict[str, Any] = {}
|
||||
result: dict[str, Any] = {
|
||||
"ok": False,
|
||||
"exit_status": -1,
|
||||
"stdout": "",
|
||||
"stderr": "",
|
||||
"command": "",
|
||||
"timed_out": False,
|
||||
"truncated": False,
|
||||
}
|
||||
|
||||
async def consume() -> None:
|
||||
nonlocal result
|
||||
async for event in stream_script_run(profile, script, password, cancel_event):
|
||||
event_type = str(event.get("type", ""))
|
||||
if event_type == "meta":
|
||||
meta.update(event)
|
||||
result["command"] = str(event.get("command", ""))
|
||||
continue
|
||||
if event_type in {"stdout", "stderr"}:
|
||||
chunk = str(event.get("data", ""))
|
||||
bucket = stdout_parts if event_type == "stdout" else stderr_parts
|
||||
bucket.append(chunk)
|
||||
total_len = sum(len(item) for item in stdout_parts) + sum(len(item) for item in stderr_parts)
|
||||
if total_len > max_output_chars:
|
||||
result["truncated"] = True
|
||||
cancel_event.set()
|
||||
return
|
||||
continue
|
||||
if event_type == "done":
|
||||
result["ok"] = bool(event.get("ok"))
|
||||
result["exit_status"] = int(event.get("exit_status", 0))
|
||||
return
|
||||
if event_type == "stopped":
|
||||
result["ok"] = False
|
||||
result["exit_status"] = int(event.get("exit_status", -1))
|
||||
result["stopped"] = True
|
||||
return
|
||||
if event_type == "error":
|
||||
raise RuntimeError(str(event.get("detail", "ssh script run failed")))
|
||||
|
||||
try:
|
||||
await asyncio.wait_for(consume(), timeout=timeout_value)
|
||||
except asyncio.TimeoutError:
|
||||
cancel_event.set()
|
||||
result["timed_out"] = True
|
||||
result["ok"] = False
|
||||
result["exit_status"] = -1
|
||||
finally:
|
||||
result["stdout"] = "".join(stdout_parts)
|
||||
result["stderr"] = "".join(stderr_parts)
|
||||
if meta:
|
||||
result["profile"] = meta.get("profile")
|
||||
result["script"] = meta.get("script")
|
||||
return result
|
||||
Reference in New Issue
Block a user