diff --git a/.flocks/flockshub/index.json b/.flocks/flockshub/index.json index 3c07841ed..4d3560e97 100644 --- a/.flocks/flockshub/index.json +++ b/.flocks/flockshub/index.json @@ -14640,7 +14640,7 @@ "name": "SOC Workspace WebUI", "description": "SOC workspace pages for posture, overview, and alert investigation.", "descriptionCn": "SOC 工作区页面，包含态势、SOC 总览和告警调查。", - "version": "1.1.4", + "version": "1.1.5", "category": "workflow-automation", "tags": [ "siem", @@ -14708,7 +14708,7 @@ "nameCn": "HTTP研判工作流", "description": "Downstream alert triage workflow that writes triage results to SOC DB by default with optional JSONL output.", "descriptionCn": "下游告警研判工作流，默认写入 SOC DB，并保留 JSONL 输出配置。", - "version": "1.0.0", + "version": "1.1.1", "category": "workflow-automation", "tags": [ "siem", diff --git a/.flocks/flockshub/plugins/webuis/soc_ui/access/soc_alerts_operations.py b/.flocks/flockshub/plugins/webuis/soc_ui/access/soc_alerts_operations.py index d8fd8b430..3718859a0 100644 --- a/.flocks/flockshub/plugins/webuis/soc_ui/access/soc_alerts_operations.py +++ b/.flocks/flockshub/plugins/webuis/soc_ui/access/soc_alerts_operations.py @@ -78,7 +78,7 @@ "is_duplicate", "_syslog_meta", ) -DRIVER_FIELDS = frozenset(DISPLAY_FIELDS) +DRIVER_FIELDS = frozenset((*DISPLAY_FIELDS, "triage_attack_verdict", "triage_attack_success")) FILTER_FIELDS = frozenset( { "_source_type", @@ -89,7 +89,8 @@ "threat_name", "threat_type", "threat_phase", - "threat_result", + "triage_attack_verdict", + "triage_attack_success", "rsp_status_code", "sip", "dport", @@ -104,7 +105,7 @@ {"key": "threat_name", "label": "Threat Name"}, {"key": "threat_type", "label": "Threat Type"}, {"key": "threat_phase", "label": "Attack Stage"}, - {"key": "threat_result", "label": "Attack Result"}, + {"key": "triage_attack_success", "label": "Attack Result"}, {"key": "direction", "label": "Direction"}, {"key": "sip", "label": "Source IP"}, {"key": "sport", "label": "Source Port"}, @@ -159,6 +160,7 @@ def resolve(self, *, page_id: str, slot_id: str, contract_id: str, contract_vers "dateColumn": settings["date_column"], "eventTimeColumn": settings["event_time_column"], }, + predicate_value_resolver=_filter_value, capabilities=frozenset({"query"}), ) @@ -236,6 +238,8 @@ def _incident_from_row(row: InternalDataRow) -> dict[str, Any]: table_cells = _table_cells(record) triage_report = _text(record.get("triage_report")) report_title = _first_text(record, "report_title") or _report_title_from_markdown(triage_report) + triage_attack_verdict = _triage_attack_verdict(record) + triage_attack_success = _triage_attack_success(record) return { "id": record_id, @@ -243,6 +247,8 @@ def _incident_from_row(row: InternalDataRow) -> dict[str, Any]: "observedAt": observed_at, "rawAlerts": 1, "priority": "P1" if verdict == "success" else "P2", + "triageAttackVerdict": triage_attack_verdict, + "triageAttackSuccess": triage_attack_success, "reportTitle": report_title, "reason": threat_msg, "owner": "", @@ -265,7 +271,7 @@ def _incident_from_row(row: InternalDataRow) -> dict[str, Any]: "business": _text(record.get("asset_group_name")), }, "conclusion": { - "verdict": _verdict_label(verdict), + "verdict": triage_attack_verdict, "summary": threat_msg or threat_name, "recommendation": "", }, @@ -288,6 +294,8 @@ def _table_cells(record: dict[str, Any]) -> dict[str, dict[str, str]]: cells[key] = {"value": value} if "time" not in cells and (observed := _observed_at(record)): cells["time"] = {"value": observed} + cells["triage_attack_verdict"] = {"value": _triage_attack_verdict(record)} + cells["triage_attack_success"] = {"value": _triage_attack_success(record)} return cells @@ -423,30 +431,51 @@ def _request_method(record: dict[str, Any]) -> str: def _verdict_bucket(record: dict[str, Any]) -> str: - if record.get("attack_success") is True: - return "success" - raw = " ".join( - _text(record.get(key)).lower() - for key in ("attack_verdict", "threat_result") - ) - if any(marker in raw for marker in ("success", "attack_success", "succeeded")): + verdict = _triage_attack_verdict(record) + if verdict == "non_attack": + return "benign" + if verdict == "attack": + attack_success = _triage_attack_success(record) + if attack_success == "success": + return "success" + if attack_success == "failed": + return "failed" + return "attack" + return "unknown" + + +def _triage_attack_success(record: dict[str, Any]) -> str: + raw_verdict = _text(record.get("triage_attack_verdict")).strip().lower() + if _triage_attack_verdict(record) != "attack": + return "unknown" + value = _text(record.get("triage_attack_success")).strip().lower() + if value in {"success", "failed", "unknown"}: + return value + legacy_value = record.get("triage_attack_success") + if legacy_value is True: return "success" - if any(marker in raw for marker in ("failed", "blocked", "attack_failed")): + if legacy_value is False and raw_verdict == "attack_failed": return "failed" - if any(marker in raw for marker in ("benign", "normal", "safe")): - return "benign" - if "attack" in raw: + return "unknown" + + +def _triage_attack_verdict(record: dict[str, Any]) -> str: + value = _text(record.get("triage_attack_verdict")).strip().lower() + if value in {"attack", "non_attack", "unknown"}: + return value + if value in {"attack_success", "attack_failed"}: return "attack" + if value == "benign": + return "non_attack" return "unknown" -def _verdict_label(bucket: str) -> str: - return { - "success": "success", - "failed": "failed", - "benign": "benign", - "attack": "attack", - }.get(bucket, "unknown") +def _filter_value(record: dict[str, Any], field: str) -> Any: + if field == "triage_attack_verdict": + return _triage_attack_verdict(record) + if field == "triage_attack_success": + return _triage_attack_success(record) + return record.get(field) def _source_file_label(paths: tuple[Path, ...]) -> str: diff --git a/.flocks/flockshub/plugins/webuis/soc_ui/manifest.json b/.flocks/flockshub/plugins/webuis/soc_ui/manifest.json index 21a8d74fa..25c7d5cdd 100644 --- a/.flocks/flockshub/plugins/webuis/soc_ui/manifest.json +++ b/.flocks/flockshub/plugins/webuis/soc_ui/manifest.json @@ -5,7 +5,7 @@ "name": "SOC Workspace WebUI", "description": "SOC workspace pages for posture, overview, and alert investigation.", "descriptionCn": "SOC 工作区页面，包含态势、SOC 总览和告警调查。", - "version": "1.1.4", + "version": "1.1.5", "author": "Flocks Team", "license": "MIT", "homepage": "", diff --git a/.flocks/flockshub/plugins/webuis/soc_ui/soc_alerts/src/filterValues.ts b/.flocks/flockshub/plugins/webuis/soc_ui/soc_alerts/src/filterValues.ts index 9ae0a4f8c..50def4b59 100644 --- a/.flocks/flockshub/plugins/webuis/soc_ui/soc_alerts/src/filterValues.ts +++ b/.flocks/flockshub/plugins/webuis/soc_ui/soc_alerts/src/filterValues.ts @@ -49,18 +49,14 @@ const FILTER_VALUE_TEXT: Record> = { exploit: '利用', exploitation: '利用', }, - threat_result: { - attack_success: '攻击成功', - success: '成功', - succeeded: '成功', - attack_failed: '攻击失败', - failed: '失败', - blocked: '已阻断', - detected: '已检测', - attack: '攻击行为', - benign: '安全', - safe: '安全', - normal: '正常', + triage_attack_success: { + success: '攻击成功', + failed: '攻击失败', + unknown: '未知', + }, + triage_attack_verdict: { + attack: '攻击', + non_attack: '非攻击', unknown: '未知', }, }; diff --git a/.flocks/flockshub/plugins/webuis/soc_ui/soc_alerts/src/index.tsx b/.flocks/flockshub/plugins/webuis/soc_ui/soc_alerts/src/index.tsx index c681d4aa9..6ba6232a5 100644 --- a/.flocks/flockshub/plugins/webuis/soc_ui/soc_alerts/src/index.tsx +++ b/.flocks/flockshub/plugins/webuis/soc_ui/soc_alerts/src/index.tsx @@ -3,7 +3,7 @@ import { api } from '@flocks/webui-contract-sdk'; import { filterOptionText, matchesFilterOptionSearch } from './filterValues'; type Tone = 'red' | 'orange' | 'blue' | 'green' | 'purple' | 'slate'; -type FilterKey = '_source_type' | 'net_type' | 'direction' | 'threat_severity' | 'threat_level' | 'threat_name' | 'threat_type' | 'threat_phase' | 'threat_result' | 'rsp_status_code' | 'sip' | 'dport' | 'dip' | 'req_host' | 'threat_rule_id'; +type FilterKey = '_source_type' | 'net_type' | 'direction' | 'threat_severity' | 'threat_level' | 'threat_name' | 'threat_type' | 'threat_phase' | 'triage_attack_verdict' | 'triage_attack_success' | 'rsp_status_code' | 'sip' | 'dport' | 'dip' | 'req_host' | 'threat_rule_id'; type TimeRangeKey = '15m' | '1h' | '2h' | '24h' | 'today' | '7d' | '30d'; type TimeFilterMode = 'relative' | 'custom'; type TimePanelTab = 'auto' | 'custom'; @@ -50,6 +50,8 @@ interface IncidentCluster { rawAlerts?: number; confidence?: number; priority?: 'P1' | 'P2'; + triageAttackVerdict?: 'attack' | 'non_attack' | 'unknown'; + triageAttackSuccess?: 'success' | 'failed' | 'unknown'; reportTitle?: string; reason?: string; owner?: string; @@ -264,6 +266,8 @@ const EN_TEXT: Record = { '展开趋势图': 'Expand timeline', '攻击成功': 'Attack Success', '攻击失败': 'Attack Failed', + '攻击': 'Attack', + '非攻击': 'Non-attack', '入站': 'Inbound', '出站': 'Outbound', '横向': 'Lateral', @@ -395,7 +399,8 @@ const BASE_FILTER_CONFIGS: FilterConfig[] = [ const MORE_FILTER_CONFIGS: FilterConfig[] = [ { key: 'threat_type', label: '威胁类型' }, { key: 'threat_phase', label: '攻击阶段' }, - { key: 'threat_result', label: '攻击结果' }, + { key: 'triage_attack_verdict', label: '攻击行为' }, + { key: 'triage_attack_success', label: '攻击结果' }, { key: 'rsp_status_code', label: '响应状态' }, { key: 'sip', label: '源地址' }, { key: 'dport', label: '目标端口' }, @@ -415,7 +420,8 @@ const DEFAULT_FILTER_VALUES: Record = { threat_name: [], threat_type: [], threat_phase: [], - threat_result: [], + triage_attack_verdict: [], + triage_attack_success: [], rsp_status_code: [], sip: [], dport: [], @@ -923,15 +929,9 @@ function cellValue(incident: IncidentCluster, key: string, fallback = '') { return textValue(incident.tableCells?.[key]?.value, fallback); } -function rawAttackResultValue(incident: IncidentCluster) { - return cellValue(incident, 'attach_result') || cellValue(incident, 'attack_result') || cellValue(incident, 'threat_result'); -} - function verdictBucket(incident: IncidentCluster): 'success' | 'failed' | 'unknown' { - const rawResult = rawAttackResultValue(incident).toLowerCase(); - const verdict = incident.conclusion?.verdict || ''; - if (rawResult === 'success' || rawResult === 'succeeded' || verdict.includes('成功')) return 'success'; - if (rawResult === 'failed' || rawResult === 'blocked' || verdict.includes('失败')) return 'failed'; + if (incident.triageAttackSuccess === 'success') return 'success'; + if (incident.triageAttackSuccess === 'failed') return 'failed'; return 'unknown'; } @@ -942,6 +942,12 @@ function attackResultLabel(incident: IncidentCluster, tr: Translate) { return tr('未知'); } +function attackVerdictLabel(incident: IncidentCluster, tr: Translate) { + if (incident.triageAttackVerdict === 'attack') return tr('攻击'); + if (incident.triageAttackVerdict === 'non_attack') return tr('非攻击'); + return tr('未知'); +} + function attackResultTone(incident: IncidentCluster): Tone { const bucket = verdictBucket(incident); if (bucket === 'success') return 'red'; @@ -949,12 +955,6 @@ function attackResultTone(incident: IncidentCluster): Tone { return 'slate'; } -function severityTone(incident: IncidentCluster): Tone { - if (incident.priority === 'P1' || incident.conclusion?.verdict?.includes('成功')) return 'red'; - if (incident.conclusion?.verdict?.includes('失败')) return 'green'; - return 'orange'; -} - function dateFromIncident(incident: IncidentCluster) { const value = incident.observedAt || cellValue(incident, 'time'); if (!value) return null; @@ -1021,6 +1021,8 @@ function niceAxisMax(value: number) { } function readFilterValue(incident: IncidentCluster, key: FilterKey) { + if (key === 'triage_attack_verdict') return incident.triageAttackVerdict || 'unknown'; + if (key === 'triage_attack_success') return incident.triageAttackSuccess || 'unknown'; return cellValue(incident, key); } @@ -1948,7 +1950,7 @@ function IncidentInlineDetail({ incident, onClose, tr }: { incident: IncidentClu const [copyDone, setCopyDone] = useState(false); const report = parseTaggedReport(incident.triageReport); const markdownReport = report ? null : parseMarkdownReport(incident.triageReport); - const attackJudgement = incident.conclusion?.verdict || tr('待确认'); + const attackJudgement = attackVerdictLabel(incident, tr); const attackResult = attackResultLabel(incident, tr); const attackResultBucket = verdictBucket(incident); const srcAddress = incident.srcIp || cellValue(incident, 'sip', '-'); @@ -2065,7 +2067,7 @@ function IncidentInlineDetail({ incident, onClose, tr }: { incident: IncidentClu
- {attackJudgement} + {attackResult} {observedAt} / {ruleId} diff --git a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/handlers.py b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/handlers.py index fb0fe6bc7..b85305ca4 100644 --- a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/handlers.py +++ b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/handlers.py @@ -7,7 +7,7 @@ import time from collections import Counter, OrderedDict from dataclasses import dataclass -from datetime import datetime, timedelta +from datetime import datetime, timedelta, timezone from pathlib import Path from threading import RLock @@ -21,7 +21,7 @@ FACTS_TABLE = "soc_dashboard_alert_facts" ACTIVITY_TABLE = "soc_dashboard_activity" META_TABLE = "soc_dashboard_meta" -SCHEMA_VERSION = "3" +SCHEMA_VERSION = "4" ACTIVITY_DEFAULT_LIMIT = 20 ACTIVITY_MAX_LIMIT = 50 ACTIVITY_WINDOW_MS = 3000 @@ -32,6 +32,37 @@ WORKFLOW_DB = Path.home() / ".flocks" / "data" / "workflow.db" WORKFLOW_SNAPSHOT_TABLE = "soc_dashboard_workflow_stats_samples" +TASK_DB = Path.home() / ".flocks" / "data" / "tasks.db" +USAGE_DB = Path.home() / ".flocks" / "data" / "flocks.db" +SOC_PINNED_WORKFLOW_NAMES = { + "stream_alert_denoise": "告警降噪工作流", + "stream_alert_triage": "告警研判工作流", +} +TRIAGE_WORKFLOW_IDS = { + "stream_alert_triage", + "onesec_kafka_investigation", + "tdp_alert_triage", + "sec_alert_unified_ops", +} +WORKFLOW_DISPLAY_NAMES = { + **SOC_PINNED_WORKFLOW_NAMES, + "onesec_kafka_investigation": "OneSEC Kafka 告警研判工作流", + "tdp_alert_triage": "TDP 告警研判工作流", + "sec_alert_unified_ops": "统一告警运营工作流", +} +WORKFLOW_RUNNING_STATUSES = {"running", "queued", "pending"} +WORKFLOW_SUCCESS_STATUSES = {"success", "completed"} +WORKFLOW_TRIGGER_CONFIG_KINDS = { + "workflow_kafka_config", + "workflow_poller_config", + "workflow_syslog_config", +} +WORKFLOW_DEFAULT_ACTIVE_TIMEOUT_SECONDS = 7200 +WORKFLOW_MIN_ACTIVE_TIMEOUT_SECONDS = 60 +UUID_RE = re.compile( + r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$", + re.IGNORECASE, +) _workflow_stats_cache: OrderedDict = OrderedDict() _CACHE_TTL: float = 30.0 @@ -42,6 +73,8 @@ _stats_response_cache: OrderedDict = OrderedDict() _STATS_RESPONSE_CACHE_TTL = 300.0 _STATS_RESPONSE_CACHE_MAX = 32 +_token_usage_cache = {"updatedAt": 0.0, "mtimeNs": 0, "value": None} +_TOKEN_USAGE_CACHE_TTL = 30.0 _cache_lock = RLock() _schema_lock = RLock() _schema_ready: set = set() @@ -79,10 +112,10 @@ class _RecordSource: "triage_persisted_at", "triage_status", "triage_source", - "verdict", + "triage_attack_verdict", "risk_level", "triage_ms", - "attack_success", + "triage_attack_success", ) @@ -104,8 +137,31 @@ def _fact_expressions(prefix): kill_chain = _json_value(prefix, "kill_chain_phase") direction = _json_value(prefix, "direction") traffic_direction = _json_value(prefix, "traffic_direction") - result = _json_value(prefix, "threat_result") - verdict = _json_value(prefix, "attack_verdict") + raw_result = _json_value(prefix, "threat_result") + raw_verdict = _json_value(prefix, "attack_verdict") + triage_attack_verdict = _json_value(prefix, "triage_attack_verdict") + triage_attack_success = _json_value(prefix, "triage_attack_success") + normalized_triage_verdict = ( + f"CASE WHEN LOWER(COALESCE({triage_attack_verdict}, '')) " + "IN ('attack', 'non_attack', 'unknown') " + f"THEN LOWER({triage_attack_verdict}) " + f"WHEN LOWER(COALESCE({triage_attack_verdict}, '')) " + "IN ('attack_success', 'attack_failed') THEN 'attack' " + f"WHEN LOWER(COALESCE({triage_attack_verdict}, '')) = 'benign' " + "THEN 'non_attack' ELSE 'unknown' END" + ) + normalized_triage_success = ( + f"CASE WHEN {normalized_triage_verdict} != 'attack' THEN 'unknown' " + f"WHEN LOWER(COALESCE({triage_attack_success}, '')) " + "IN ('success', 'failed', 'unknown') " + f"THEN LOWER({triage_attack_success}) " + f"WHEN LOWER(COALESCE({triage_attack_verdict}, '')) = 'attack_success' " + "THEN 'success' " + f"WHEN LOWER(COALESCE({triage_attack_verdict}, '')) = 'attack_failed' " + "THEN 'failed' " + f"WHEN LOWER(COALESCE({triage_attack_success}, '')) IN ('1', 'true') " + "THEN 'success' ELSE 'unknown' END" + ) protocol = _json_value(prefix, "net_type") app_protocol = _json_value(prefix, "net_app_proto") protocol_fallback = _json_value(prefix, "protocol") @@ -121,7 +177,6 @@ def _fact_expressions(prefix): triage_source = _json_value(prefix, "triage_source") triage_report = _json_value(prefix, "triage_report") triage_ms = _json_value(prefix, "triage_ms") - attack_success = _json_value(prefix, "attack_success") has_triage = ( "CASE WHEN " f"NULLIF({triage_status}, '') IS NOT NULL " @@ -140,7 +195,9 @@ def _fact_expressions(prefix): f"COALESCE({prefix}.is_duplicate, 0)", f"COALESCE(NULLIF({phase}, ''), NULLIF({attack_phase}, ''), NULLIF({kill_chain}, ''), 'unknown')", f"COALESCE(NULLIF({direction}, ''), NULLIF({traffic_direction}, ''), 'unknown')", - f"COALESCE(NULLIF({result}, ''), NULLIF({verdict}, ''), 'unknown')", + f"CASE WHEN {has_triage} = 1 " + f"THEN {normalized_triage_success} " + f"ELSE COALESCE(NULLIF({raw_result}, ''), NULLIF({raw_verdict}, ''), 'unknown') END", f"COALESCE(NULLIF({protocol}, ''), NULLIF({app_protocol}, ''), " f"NULLIF({protocol_fallback}, ''), 'unknown')", f"COALESCE(NULLIF({severity}, ''), 'unknown')", @@ -151,10 +208,10 @@ def _fact_expressions(prefix): f"COALESCE({triage_persisted_at}, '')", f"COALESCE({triage_status}, '')", f"COALESCE({triage_source}, '')", - f"COALESCE({verdict}, 'unknown')", + normalized_triage_verdict, f"COALESCE(NULLIF({risk_level}, ''), 'unknown')", f"COALESCE(CAST({triage_ms} AS INTEGER), 0)", - f"CASE WHEN {attack_success} IN (1, '1', 'true') THEN 1 ELSE 0 END", + normalized_triage_success, ) @@ -374,10 +431,10 @@ def _ensure_sqlite_schema(): triage_persisted_at TEXT, triage_status TEXT, triage_source TEXT, - verdict TEXT, + triage_attack_verdict TEXT, risk_level TEXT, triage_ms INTEGER NOT NULL DEFAULT 0, - attack_success INTEGER NOT NULL DEFAULT 0 + triage_attack_success TEXT NOT NULL DEFAULT 'unknown' ) """ ) @@ -477,6 +534,89 @@ def _safe_json_object(value): return parsed if isinstance(parsed, dict) else {} +def _usage_iso(dt): + return dt.astimezone(timezone.utc).isoformat() + + +def _read_token_usage(): + empty = { + "totalTokens": 0, + "todayTokens": 0, + "todayRequests": 0, + "dailySeries": [], + "dailyLabels": [], + "source": "usage_records", + } + if not USAGE_DB.is_file(): + return empty + + try: + mtime_ns = USAGE_DB.stat().st_mtime_ns + except Exception: + mtime_ns = 0 + now_monotonic = time.monotonic() + with _cache_lock: + cached_value = _token_usage_cache.get("value") + if ( + cached_value is not None + and _token_usage_cache.get("mtimeNs") == mtime_ns + and now_monotonic - float(_token_usage_cache.get("updatedAt") or 0) < _TOKEN_USAGE_CACHE_TTL + ): + return cached_value + + now_local = datetime.now().astimezone() + today_start = now_local.replace(hour=0, minute=0, second=0, microsecond=0) + tomorrow_start = today_start + timedelta(days=1) + first_day = today_start - timedelta(days=6) + labels = [(first_day + timedelta(days=index)).strftime("%m/%d") for index in range(7)] + series_by_date = { + (first_day + timedelta(days=index)).date().isoformat(): 0 + for index in range(7) + } + + try: + with sqlite3.connect(f"file:{USAGE_DB}?mode=ro", uri=True, timeout=1.0) as conn: + conn.execute("PRAGMA query_only = ON") + if not _table_exists(conn, "usage_records"): + return {**empty, "dailyLabels": labels, "dailySeries": [0] * 7} + total_tokens = _safe_int( + conn.execute( + "SELECT COALESCE(SUM(total_tokens), 0) FROM usage_records" + ).fetchone()[0] + ) + today_row = conn.execute( + "SELECT COALESCE(SUM(total_tokens), 0), COUNT(*) " + "FROM usage_records WHERE created_at >= ? AND created_at < ?", + (_usage_iso(today_start), _usage_iso(tomorrow_start)), + ).fetchone() + series_rows = conn.execute( + "SELECT date(created_at, 'localtime') AS usage_day, " + "COALESCE(SUM(total_tokens), 0) AS token_count " + "FROM usage_records WHERE created_at >= ? AND created_at < ? " + "GROUP BY usage_day", + (_usage_iso(first_day), _usage_iso(tomorrow_start)), + ).fetchall() + except Exception: + return {**empty, "dailyLabels": labels, "dailySeries": [0] * 7} + + for usage_day, total in series_rows: + key = str(usage_day or "") + if key in series_by_date: + series_by_date[key] += max(_safe_int(total), 0) + + result = { + **empty, + "totalTokens": max(total_tokens, 0), + "todayTokens": max(_safe_int(today_row[0]), 0) if today_row else 0, + "todayRequests": max(_safe_int(today_row[1]), 0) if today_row else 0, + "dailySeries": list(series_by_date.values()), + "dailyLabels": labels, + } + with _cache_lock: + _token_usage_cache.update({"updatedAt": now_monotonic, "mtimeNs": mtime_ns, "value": result}) + return result + + def _workflow_stats_sample_deltas(conn, workflow_name, start_time=0, end_time=0): stats_exists = conn.execute( "SELECT 1 FROM sqlite_master WHERE type='table' AND name='workflow_stats'" @@ -890,11 +1030,27 @@ def _get_workflow_recent_events( ) -> list: if not WORKFLOW_DB.is_file(): return [] - query = ( - "SELECT id, status, started_at, output_results, input_params " - "FROM workflow_executions WHERE workflow_id = ?" - ) + workflow_stage = "triage" if workflow_name in TRIAGE_WORKFLOW_IDS else "denoise" query_params = [workflow_name] + try: + with sqlite3.connect(WORKFLOW_DB) as conn: + execution_columns = { + row[1] + for row in conn.execute("PRAGMA table_info(workflow_executions)").fetchall() + } + except Exception: + return [] + latest_select = ", ".join( + [ + "id", + "status", + "started_at", + _workflow_execution_column_expr(execution_columns, "output_results", "'{}'"), + _workflow_execution_column_expr(execution_columns, "input_params", "'{}'"), + _workflow_execution_column_expr(execution_columns, "payload", "'{}'"), + ] + ) + query = f"SELECT {latest_select} FROM workflow_executions WHERE workflow_id = ?" if start_time > 0 and end_time > 0: query += " AND started_at >= ? AND started_at <= ?" query_params.extend((int(start_time * 1000), int(end_time * 1000))) @@ -902,28 +1058,59 @@ def _get_workflow_recent_events( query_params.append(max(1, min(_safe_int(limit), 10))) try: with sqlite3.connect(WORKFLOW_DB) as conn: + conn.row_factory = sqlite3.Row rows = conn.execute(query, query_params).fetchall() except Exception: return [] events = [] - for execution_id, status, started_at, output_text, input_text in rows: + for row in rows: + execution_id = row["id"] + status = row["status"] + started_at = row["started_at"] + output_text = row["output_results"] + input_text = row["input_params"] + payload_text = row["payload"] metrics = _workflow_execution_metrics(output_text, input_text) preview = metrics["preview"] raw_count = metrics["rawCount"] unique_count = metrics["uniqueCount"] - threat_name = str( - preview.get("threat_name") or f"降噪批次 · 原始 {raw_count} 条" + threat_name = "" + if workflow_stage == "triage": + threat_name = _workflow_latest_alert_name(workflow_name, output_text, input_text) + if not threat_name: + threat_name = str( + preview.get("threat_name") + or preview.get("_threat_type") + or preview.get("threat_type") + or f"降噪批次 · 原始 {raw_count} 条" + ) + normalized_status = str(status or "").lower() + event_status = ( + "completed" + if normalized_status in {"success", "completed"} + else "running" + if normalized_status in {"running", "queued", "pending"} + else "failed" + ) + session_id, message_id = _workflow_link_context( + { + "payload": payload_text, + "input_params": input_text, + } ) events.append( { "eventId": f"workflow-execution:{execution_id}", - "stage": "denoise", - "status": "completed" if str(status).lower() == "success" else "failed", + "stage": workflow_stage, + "status": event_status, "occurredAt": datetime.fromtimestamp( _safe_int(started_at) / 1000 ).astimezone().isoformat(timespec="seconds"), "triggerSource": "workflow_execution", + "workflowId": workflow_name, + "sessionId": session_id, + "messageId": message_id, "sampleCount": max(unique_count, 1), "alert": { "id": str(preview.get("id") or execution_id), @@ -972,23 +1159,683 @@ def _get_workflow_recent_events( } RESULT_LABELS = { - "success": "攻击成功", - "succeeded": "攻击成功", - "failed": "攻击失败", - "blocked": "已阻断", "attack_success": "攻击成功", "attack": "攻击行为", "attack_failed": "攻击失败", - "benign": "良性", + "non_attack": "非攻击", "unknown": "待确认", } +def _triage_outcome_key(record): + verdict = _triage_attack_verdict_value(record) + result = _triage_attack_success_value(record) + if verdict == "attack": + if result == "success": + return "attack_success" + if result == "failed": + return "attack_failed" + return "attack" + if verdict == "non_attack": + return "non_attack" + return "unknown" + + +def _triage_attack_verdict_value(record): + verdict = _norm(record.get("triage_attack_verdict") or "unknown") + if verdict in {"attack", "non_attack", "unknown"}: + return verdict + if verdict in {"attack_success", "attack_failed"}: + return "attack" + if verdict == "benign": + return "non_attack" + return "unknown" + + +def _triage_attack_success_value(record): + raw_verdict = _norm(record.get("triage_attack_verdict") or "unknown") + verdict = _triage_attack_verdict_value(record) + result = _norm(record.get("triage_attack_success") or "unknown") + if verdict == "attack" and result in {"success", "failed", "unknown"}: + return result + if verdict != "attack": + return "unknown" + if raw_verdict == "attack_success" or record.get("triage_attack_success") is True: + return "success" + if raw_verdict == "attack_failed": + return "failed" + return "unknown" + + async def get_activity(ctx, request): params = dict(request.query_params) return await asyncio.to_thread(_get_activity, params) +async def get_task_center(ctx, request): + params = dict(request.query_params) + include_mock = _truthy(params.get("mockActivity")) or _truthy(params.get("mockTaskCenter")) + return await asyncio.to_thread(_get_task_center, include_mock) + + +def _table_exists(conn, table_name): + return bool( + conn.execute( + "SELECT 1 FROM sqlite_master WHERE type='table' AND name=?", + (table_name,), + ).fetchone() + ) + + +def _task_center_empty(): + return { + "generatedAt": datetime.now().astimezone().isoformat(timespec="seconds"), + "sessionCount": 0, + "scheduledTasks": [], + "workflows": [], + "sourceStatus": { + "tasksDb": str(TASK_DB), + "workflowDb": str(WORKFLOW_DB), + "tasksAvailable": TASK_DB.is_file(), + "workflowAvailable": WORKFLOW_DB.is_file(), + }, + } + + +def _truthy(value): + return str(value or "").strip().lower() in {"1", "true", "yes", "on"} + + +def _task_trigger_value(trigger, key): + if not isinstance(trigger, dict): + return "" + return trigger.get(key) or trigger.get(key[0].lower() + key[1:]) or "" + + +def _today_bounds(): + now_local = datetime.now().astimezone() + today_start = now_local.replace(hour=0, minute=0, second=0, microsecond=0) + return today_start, today_start + timedelta(days=1) + + +def _task_center_task_rows(limit=12): + if not TASK_DB.is_file(): + return 0, [], 0, 0 + try: + with sqlite3.connect(f"file:{TASK_DB}?mode=ro", uri=True, timeout=1.0) as conn: + conn.row_factory = sqlite3.Row + conn.execute("PRAGMA query_only = ON") + if not ( + _table_exists(conn, "task_schedulers") + and _table_exists(conn, "task_executions") + ): + return 0, [], 0, 0 + today_start, tomorrow_start = _today_bounds() + today_start_iso = today_start.isoformat(timespec="seconds") + tomorrow_start_iso = tomorrow_start.isoformat(timespec="seconds") + session_count = _safe_int( + conn.execute( + "SELECT COUNT(DISTINCT session_id) FROM task_executions " + "WHERE session_id IS NOT NULL AND session_id <> ''" + ).fetchone()[0] + ) + scheduler_rows = conn.execute( + "SELECT id, title, mode, status, trigger, execution_mode, workflow_id, updated_at " + "FROM task_schedulers WHERE status <> 'archived' " + "ORDER BY updated_at DESC" + ).fetchall() + tasks = [] + for scheduler in scheduler_rows: + summary = conn.execute( + "SELECT COUNT(*) AS execution_count, " + "SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END) AS success_count, " + "SUM(CASE WHEN status IN ('pending', 'queued', 'running') THEN 1 ELSE 0 END) AS active_count, " + "SUM(CASE WHEN " + "julianday(COALESCE(completed_at, updated_at, started_at, queued_at, created_at)) >= julianday(?) " + "AND julianday(COALESCE(completed_at, updated_at, started_at, queued_at, created_at)) < julianday(?) " + "THEN 1 ELSE 0 END) AS today_execution_count " + "FROM task_executions WHERE scheduler_id = ?", + (today_start_iso, tomorrow_start_iso, scheduler["id"]), + ).fetchone() + latest = conn.execute( + "SELECT status, queued_at, started_at, completed_at, updated_at " + "FROM task_executions WHERE scheduler_id = ? " + "ORDER BY julianday(COALESCE(completed_at, updated_at, started_at, queued_at, created_at)) DESC " + "LIMIT 1", + (scheduler["id"],), + ).fetchone() + trigger = _safe_json_object(scheduler["trigger"]) + execution_count = max(_safe_int(summary["execution_count"]), 0) + success_count = max(_safe_int(summary["success_count"]), 0) + active_count = max(_safe_int(summary["active_count"]), 0) + today_execution_count = max(_safe_int(summary["today_execution_count"]), 0) + last_run_at = "" + if latest: + last_run_at = latest["completed_at"] or latest["updated_at"] or latest["started_at"] or latest["queued_at"] or "" + tasks.append( + { + "id": scheduler["id"], + "name": scheduler["title"] or scheduler["id"], + "mode": scheduler["mode"] or "once", + "status": scheduler["status"] or "active", + "executionMode": scheduler["execution_mode"] or "agent", + "workflowId": scheduler["workflow_id"] or "", + "executionCount": execution_count, + "todayExecutionCount": today_execution_count, + "successCount": success_count, + "successRate": _ratio(success_count, execution_count), + "activeCount": active_count, + "lastStatus": latest["status"] if latest else "", + "lastRunAt": last_run_at, + "nextRunAt": _task_trigger_value(trigger, "nextRun"), + "cron": _task_trigger_value(trigger, "cron"), + "cronDescription": _task_trigger_value(trigger, "cronDescription"), + } + ) + tasks.sort( + key=lambda item: ( + item["activeCount"] > 0, + item["lastRunAt"] or "", + item["executionCount"], + ), + reverse=True, + ) + return ( + session_count, + tasks[:limit], + sum(task["executionCount"] for task in tasks), + sum(task["todayExecutionCount"] for task in tasks), + ) + except Exception: + return 0, [], 0, 0 + + +def _workflow_manifest_name_map(): + roots = [ + Path.home() / ".flocks" / "plugins" / "workflows", + Path.home() / ".flocks" / "workspace" / "workflows", + Path(__file__).resolve().parents[4] / "workflows", + ] + names = {} + for root in roots: + if not root.is_dir(): + continue + for manifest_path in root.glob("*/manifest.json"): + workflow_id = manifest_path.parent.name + try: + manifest = _safe_json_object(manifest_path.read_text(encoding="utf-8")) + except Exception: + manifest = {} + name = ( + manifest.get("nameCn") + or manifest.get("name") + or manifest.get("title") + or workflow_id + ) + name_i18n = manifest.get("nameI18n") + if isinstance(name_i18n, dict): + name = name_i18n.get("zh-CN") or name_i18n.get("zh") or name + names.setdefault(workflow_id, str(name)) + return names + + +def _workflow_config_name_map(conn): + if not _table_exists(conn, "workflow_configs"): + return {} + names = {} + for row in conn.execute( + "SELECT workflow_id, config FROM workflow_configs WHERE kind = ?", + ("workflow.integration-config",), + ).fetchall(): + workflow_id = row["workflow_id"] if isinstance(row, sqlite3.Row) else row[0] + config_raw = row["config"] if isinstance(row, sqlite3.Row) else row[1] + if not workflow_id: + continue + config = _safe_json_object(config_raw) + workflow = config.get("workflow") + if not isinstance(workflow, dict): + continue + name = workflow.get("name") or workflow.get("title") or workflow.get("id") + if name: + names.setdefault(str(workflow_id), str(name)) + return names + + +def _truthy_config_value(value): + if isinstance(value, bool): + return value + if isinstance(value, (int, float)): + return value != 0 + text = str(value or "").strip().lower() + if not text: + return False + return text in {"1", "true", "yes", "on", "enabled"} + + +def _workflow_trigger_state(conn, workflow_id): + state = { + "hasConfig": False, + "enabled": True, + "timeoutSeconds": WORKFLOW_DEFAULT_ACTIVE_TIMEOUT_SECONDS, + } + if not _table_exists(conn, "workflow_configs"): + return state + + placeholders = ",".join("?" for _ in WORKFLOW_TRIGGER_CONFIG_KINDS) + rows = conn.execute( + "SELECT kind, config FROM workflow_configs WHERE workflow_id = ? " + f"AND (kind IN ({placeholders}) OR kind LIKE ? OR kind LIKE ? OR kind LIKE ?)", + ( + workflow_id, + *WORKFLOW_TRIGGER_CONFIG_KINDS, + "workflow_kafka_config/%", + "workflow_poller_config/%", + "workflow_syslog_config/%", + ), + ).fetchall() + if not rows: + return state + + enabled = False + timeout_seconds = WORKFLOW_DEFAULT_ACTIVE_TIMEOUT_SECONDS + for row in rows: + config = _safe_json_object(row["config"] if isinstance(row, sqlite3.Row) else row[1]) + if "enabled" not in config or _truthy_config_value(config.get("enabled")): + enabled = True + timeout_seconds = max( + timeout_seconds, + _safe_int(config.get("timeoutSeconds") or config.get("timeout_seconds")), + ) + state["hasConfig"] = True + state["enabled"] = enabled + state["timeoutSeconds"] = max(timeout_seconds, WORKFLOW_MIN_ACTIVE_TIMEOUT_SECONDS) + return state + + +def _workflow_effective_status(latest, active_count, trigger_state): + if not latest: + return "" + status = str(latest["status"] or "").lower() + if status not in WORKFLOW_RUNNING_STATUSES: + return status + if active_count > 0: + return status + if trigger_state.get("hasConfig") and not trigger_state.get("enabled"): + return "disabled" + return "stale" + + +def _workflow_execution_column_expr(columns, column_name, fallback): + return column_name if column_name in columns else f"{fallback} AS {column_name}" + + +def _workflow_node_count(workflow_id): + roots = [ + Path.home() / ".flocks" / "plugins" / "workflows", + Path.home() / ".flocks" / "workspace" / "workflows", + Path(__file__).resolve().parents[4] / "workflows", + ] + for root in roots: + workflow_path = root / str(workflow_id) / "workflow.json" + if not workflow_path.is_file(): + continue + try: + workflow_json = _safe_json_object(workflow_path.read_text(encoding="utf-8")) + nodes = workflow_json.get("nodes") + if isinstance(nodes, list): + return max(len(nodes), 0) + except Exception: + continue + return 0 + + +def _first_text(*values): + for value in values: + text = str(value or "").strip() + if text and text.lower() not in {"unknown", "none", "null", "--"}: + return text + return "" + + +def _alert_name_from_record(record): + if not isinstance(record, dict): + return "" + return _first_text( + record.get("threat_name"), + record.get("_threat_type"), + record.get("threat_type"), + record.get("alert_name"), + record.get("name"), + record.get("report_title"), + record.get("title"), + ) + + +def _workflow_latest_alert_name(workflow_id, output_text="", input_text=""): + output = _safe_json_object(output_text) + inputs = _safe_json_object(input_text) + preview = _workflow_alert_preview(output) or _workflow_input_preview(inputs) + name = _alert_name_from_record(preview) + if name: + return name + for key in ( + "enriched_alerts_with_triage", + "triage_results", + "unique_alerts", + "enriched_alerts", + ): + value = output.get(key) + if isinstance(value, dict): + value = value.get("preview") + if isinstance(value, list): + for item in value: + name = _alert_name_from_record(item) + if name: + return name + name = _first_text( + output.get("top_report_title"), + output.get("report_title"), + output.get("title"), + ) + if name: + return name + if workflow_id in TRIAGE_WORKFLOW_IDS: + input_date = _first_text(inputs.get("input_date")) + return f"研判批次 {input_date}" if input_date else "研判批次" + if workflow_id == "stream_alert_denoise": + return "降噪批次" + return "" + + +def _workflow_progress(status, current_step_index, completed_steps, total_steps): + normalized = str(status or "").lower() + current = max(_safe_int(current_step_index), _safe_int(completed_steps), 0) + total = max(_safe_int(total_steps), current, 1) + if normalized == "disabled": + return 0, "已关闭" + if normalized == "stale": + return 0, "已停止" + if normalized in WORKFLOW_SUCCESS_STATUSES: + return 1, "已完成" + if normalized in {"error", "failed", "timeout"}: + percent = _ratio(current or total, total) + return percent, f"失败于 {current or total}/{total} 步" + if normalized == "cancelled": + percent = _ratio(current or total, total) + return percent, f"已取消 {current or total}/{total} 步" + if normalized in WORKFLOW_RUNNING_STATUSES: + visible_current = current if current > 0 else 1 + return max(_ratio(visible_current, total), 0.03), f"第 {visible_current}/{total} 步" + return 0, "待执行" + + +def _workflow_link_context(latest): + payload = _safe_json_object(latest["payload"] if latest and "payload" in latest.keys() else "") + input_params = _safe_json_object(latest["input_params"] if latest and "input_params" in latest.keys() else "") + session_id = _first_text( + payload.get("sessionId"), + payload.get("sessionID"), + payload.get("session_id"), + input_params.get("sessionId"), + input_params.get("sessionID"), + input_params.get("session_id"), + ) + message_id = _first_text( + payload.get("messageId"), + payload.get("messageID"), + payload.get("message_id"), + input_params.get("messageId"), + input_params.get("messageID"), + input_params.get("message_id"), + ) + return session_id, message_id + + +def _task_center_workflow_rows(limit=12, include_mock=False): + if not WORKFLOW_DB.is_file(): + return [], 0, 0 + try: + with sqlite3.connect(f"file:{WORKFLOW_DB}?mode=ro", uri=True, timeout=1.0) as conn: + conn.row_factory = sqlite3.Row + conn.execute("PRAGMA query_only = ON") + if not ( + _table_exists(conn, "workflow_stats") + or _table_exists(conn, "workflow_executions") + ): + return [], 0, 0 + today_start, tomorrow_start = _today_bounds() + today_start_ms = int(today_start.timestamp() * 1000) + tomorrow_start_ms = int(tomorrow_start.timestamp() * 1000) + workflow_ids = set() + if _table_exists(conn, "workflow_stats"): + workflow_ids.update( + row[0] + for row in conn.execute("SELECT workflow_id FROM workflow_stats").fetchall() + if row[0] + ) + if _table_exists(conn, "workflow_executions"): + workflow_ids.update( + row[0] + for row in conn.execute( + "SELECT DISTINCT workflow_id FROM workflow_executions" + ).fetchall() + if row[0] + ) + names = { + **_workflow_manifest_name_map(), + **_workflow_config_name_map(conn), + **WORKFLOW_DISPLAY_NAMES, + } + if include_mock: + workflow_ids.update(SOC_PINNED_WORKFLOW_NAMES) + stats_columns = ( + { + row[1] + for row in conn.execute("PRAGMA table_info(workflow_stats)").fetchall() + } + if _table_exists(conn, "workflow_stats") + else set() + ) + stats_success_expr = "success_count" if "success_count" in stats_columns else "0" + stats_error_expr = "error_count" if "error_count" in stats_columns else "0" + stats_updated_expr = "updated_at" if "updated_at" in stats_columns else "0" + execution_columns = ( + { + row[1] + for row in conn.execute("PRAGMA table_info(workflow_executions)").fetchall() + } + if _table_exists(conn, "workflow_executions") + else set() + ) + finished_at_expr = "finished_at" if "finished_at" in execution_columns else "NULL" + updated_at_expr = "updated_at" if "updated_at" in execution_columns else "NULL" + execution_time_expr = f"COALESCE({finished_at_expr}, {updated_at_expr}, started_at, 0)" + active_time_expr = f"COALESCE({updated_at_expr}, started_at, 0)" + latest_select = ", ".join( + [ + "id", + "status", + "started_at", + _workflow_execution_column_expr(execution_columns, "finished_at", "NULL"), + _workflow_execution_column_expr(execution_columns, "updated_at", "0"), + _workflow_execution_column_expr(execution_columns, "current_phase", "''"), + _workflow_execution_column_expr(execution_columns, "current_step_index", "0"), + _workflow_execution_column_expr(execution_columns, "step_count", "0"), + _workflow_execution_column_expr(execution_columns, "input_params", "'{}'"), + _workflow_execution_column_expr(execution_columns, "output_results", "'{}'"), + _workflow_execution_column_expr(execution_columns, "payload", "'{}'"), + _workflow_execution_column_expr(execution_columns, "error_message", "''"), + ] + ) + workflows = [] + now_ms = int(time.time() * 1000) + for workflow_id in workflow_ids: + workflow_name = names.get(workflow_id, workflow_id) + if UUID_RE.match(str(workflow_id)) and workflow_name == workflow_id: + continue + trigger_state = _workflow_trigger_state(conn, workflow_id) + stats = None + if _table_exists(conn, "workflow_stats"): + stats = conn.execute( + f"SELECT call_count, {stats_success_expr}, {stats_error_expr}, {stats_updated_expr} " + "FROM workflow_stats WHERE workflow_id = ?", + (workflow_id,), + ).fetchone() + latest = None + exec_summary = None + if _table_exists(conn, "workflow_executions"): + exec_summary = conn.execute( + "SELECT COUNT(*) AS execution_count, " + "SUM(CASE WHEN status IN ('success', 'completed') THEN 1 ELSE 0 END) AS success_count, " + f"SUM(CASE WHEN {execution_time_expr} >= ? " + f"AND {execution_time_expr} < ? THEN 1 ELSE 0 END) " + "AS today_execution_count " + "FROM workflow_executions WHERE workflow_id = ?", + (today_start_ms, tomorrow_start_ms, workflow_id), + ).fetchone() + latest = conn.execute( + f"SELECT {latest_select} " + "FROM workflow_executions WHERE workflow_id = ? " + f"ORDER BY {execution_time_expr} DESC LIMIT 1", + (workflow_id,), + ).fetchone() + active_count = 0 + if _table_exists(conn, "workflow_executions"): + if trigger_state["hasConfig"] and trigger_state["enabled"]: + active_since_ms = now_ms - trigger_state["timeoutSeconds"] * 1000 + active_summary = conn.execute( + "SELECT COUNT(*) AS active_count " + "FROM workflow_executions WHERE workflow_id = ? " + "AND status IN ('running', 'queued', 'pending') " + f"AND {active_time_expr} >= ?", + (workflow_id, active_since_ms), + ).fetchone() + elif not trigger_state["hasConfig"] and "updated_at" in execution_columns: + active_since_ms = now_ms - WORKFLOW_DEFAULT_ACTIVE_TIMEOUT_SECONDS * 1000 + active_summary = conn.execute( + "SELECT COUNT(*) AS active_count " + "FROM workflow_executions WHERE workflow_id = ? " + "AND status IN ('running', 'queued', 'pending') " + f"AND {active_time_expr} >= ?", + (workflow_id, active_since_ms), + ).fetchone() + elif not trigger_state["hasConfig"]: + active_summary = conn.execute( + "SELECT COUNT(*) AS active_count " + "FROM workflow_executions WHERE workflow_id = ? " + "AND status IN ('running', 'queued', 'pending')", + (workflow_id,), + ).fetchone() + else: + active_summary = None + active_count = max( + _safe_int(active_summary["active_count"] if active_summary else 0), + 0, + ) + execution_count = max( + _safe_int(stats["call_count"] if stats else 0), + _safe_int(exec_summary["execution_count"] if exec_summary else 0), + ) + success_count = max( + _safe_int(stats["success_count"] if stats else 0), + _safe_int(exec_summary["success_count"] if exec_summary else 0), + ) + today_execution_count = max( + _safe_int(exec_summary["today_execution_count"] if exec_summary else 0), + 0, + ) + last_run_at = 0 + if latest: + last_run_at = ( + _safe_int(latest["finished_at"]) + or _safe_int(latest["updated_at"]) + or _safe_int(latest["started_at"]) + ) + workflow_total_steps = _workflow_node_count(workflow_id) + latest_alert_name = "" + progress_percent = 0 + progress_label = "待执行" + current_phase = "" + session_id = "" + message_id = "" + effective_status = _workflow_effective_status(latest, active_count, trigger_state) + if latest: + workflow_total_steps = max(workflow_total_steps, _safe_int(latest["step_count"])) + latest_alert_name = _workflow_latest_alert_name( + workflow_id, + latest["output_results"], + latest["input_params"], + ) + progress_percent, progress_label = _workflow_progress( + effective_status, + latest["current_step_index"], + latest["step_count"], + workflow_total_steps, + ) + current_phase = str(latest["current_phase"] or "") + session_id, message_id = _workflow_link_context(latest) + workflows.append( + { + "id": workflow_id, + "name": workflow_name, + "executionCount": execution_count, + "todayExecutionCount": today_execution_count, + "successCount": success_count, + "successRate": _ratio(success_count, execution_count), + "activeCount": active_count, + "lastStatus": effective_status, + "lastRunAt": last_run_at, + "latestExecutionHash": str(latest["id"] if latest else ""), + "latestAlertName": latest_alert_name, + "progressPercent": progress_percent, + "progressLabel": progress_label, + "currentPhase": current_phase, + "sessionId": session_id, + "messageId": message_id, + } + ) + soc_order = { + "stream_alert_denoise": 3, + "stream_alert_triage": 2, + "onesec_kafka_investigation": 1, + "tdp_alert_triage": 1, + "sec_alert_unified_ops": 1, + } + workflows.sort( + key=lambda item: ( + item["activeCount"] > 0, + item["lastRunAt"], + item["executionCount"], + soc_order.get(item["id"], 0), + ), + reverse=True, + ) + return ( + workflows[:limit], + sum(workflow["executionCount"] for workflow in workflows), + sum(workflow["todayExecutionCount"] for workflow in workflows), + ) + except Exception: + return [], 0, 0 + + +def _get_task_center(include_mock=False): + session_count, tasks, scheduled_execution_count, scheduled_today_execution_count = _task_center_task_rows() + workflows, workflow_execution_count, workflow_today_execution_count = _task_center_workflow_rows( + include_mock=include_mock, + ) + return { + **_task_center_empty(), + "sessionCount": session_count, + "scheduledTasks": tasks, + "scheduledExecutionCount": scheduled_execution_count, + "scheduledTodayExecutionCount": scheduled_today_execution_count, + "workflowExecutionCount": workflow_execution_count, + "workflowTodayExecutionCount": workflow_today_execution_count, + "workflows": workflows, + } + + def _get_activity(params): _ensure_sqlite_schema() _maybe_prune_activity() @@ -1004,11 +1851,10 @@ def _get_activity(params): start_time, end_time, ) - workflow_events = _get_workflow_recent_events( - "stream_alert_denoise", - start_time, - end_time, - ) + workflow_events = [ + *_get_workflow_recent_events("stream_alert_denoise", start_time, end_time), + *_get_workflow_recent_events("stream_alert_triage", start_time, end_time), + ] raw_cursor = str(params.get("cursor") or "").strip() bootstrap = str(params.get("bootstrap") or "").strip().lower() == "latest" limit = max(1, min(_safe_int(params.get("limit") or ACTIVITY_DEFAULT_LIMIT), ACTIVITY_MAX_LIMIT)) @@ -1030,6 +1876,18 @@ def _get_activity(params): with sqlite3.connect(db_path) as conn: conn.row_factory = sqlite3.Row conn.execute("PRAGMA query_only = ON") + if not ( + _table_exists(conn, DEFAULT_SQLITE_TABLE) + and _table_exists(conn, ACTIVITY_TABLE) + ): + return _activity_response( + [], + 0, + 0, + cursor_reset=cursor_reset, + workflow_stats=workflow_stats, + workflow_events=workflow_events, + ) latest_row_id, latest_activity_id = _activity_latest_cursor(conn, settings) if bootstrap or cursor is None: @@ -1121,6 +1979,7 @@ def _activity_response( "cursorReset": cursor_reset, "workflowStats": workflow_stats or {"callCount": None, "latestStartedAt": None}, "workflowEvents": workflow_events or [], + "tokenUsage": _read_token_usage(), } @@ -1388,13 +2247,16 @@ def _activity_event(row): "dedupKey": _first_activity_text(record, "dedup_key"), } else: - verdict = _norm(record.get("attack_verdict") or "unknown") + verdict = _triage_attack_verdict_value(record) + attack_success = _triage_attack_success_value(record) + outcome = _triage_outcome_key(record) event["result"] = { "triageStatus": triage_status or "completed", "triageSource": str(record.get("triage_source") or "").strip().lower() or "triaged", "durationMs": _safe_int(record.get("triage_ms")), "verdict": verdict, - "verdictLabel": RESULT_LABELS.get(verdict, "待确认"), + "attackSuccess": attack_success, + "verdictLabel": RESULT_LABELS.get(outcome, "待确认"), "threatSeverity": _first_activity_text(record, "threat_severity"), "riskLevel": _first_activity_text(record, "risk_level"), "reportTitle": _first_activity_text(record, "report_title"), @@ -1505,6 +2367,7 @@ def _get_stats(params): if cached is not None: return { **cached, + "tokenUsage": _read_token_usage(), "generatedAt": datetime.now().isoformat(timespec="seconds"), "latencyMs": round((time.time() - started) * 1000), "cacheHit": True, @@ -1584,6 +2447,7 @@ def _get_stats(params): "allowedDatabases": [ _display_path(DEFAULT_SQLITE_DB), _display_path(WORKFLOW_DB), + _display_path(USAGE_DB), ], }, "workflowStatsDb": _display_path(WORKFLOW_DB), @@ -1620,7 +2484,7 @@ def _get_stats(params): {"key": "attack_success", "label": "攻击成功", "value": triage["attackSuccess"], "color": "#ff4d6d"}, {"key": "attack", "label": "攻击行为", "value": triage["attack"], "color": "#ffb020"}, {"key": "attack_failed", "label": "攻击失败", "value": triage["attackFailed"], "color": "#2ee6a6"}, - {"key": "benign", "label": "良性", "value": triage["benign"], "color": "#58a6ff"}, + {"key": "non_attack", "label": "非攻击", "value": triage["benign"], "color": "#58a6ff"}, {"key": "unknown", "label": "未知", "value": triage["unknown"], "color": "#9b8cff"}, ], "topThreatTypes": _counter_items( @@ -1632,6 +2496,7 @@ def _get_stats(params): 8, ), "riskLevels": _counter_items(triage["riskCounter"], 5), + "tokenUsage": _read_token_usage(), "timeline": { "labels": denoise.get("_timelineLabels") or _series_labels(max(len(denoise["seriesRaw"]), len(triage["seriesTotal"]))), @@ -2034,9 +2899,7 @@ def _sqlite_timeline(conn, settings, where_clause, query_params, dates, start_ti f"COALESCE(SUM(has_triage), 0) AS triage_count, " f"COALESCE(SUM(CASE WHEN has_triage = 1 " f"AND LOWER(triage_status) NOT IN ('failed', 'error') " - f"AND (LOWER(verdict) IN ('attack_success', 'attack', 'attack_failed') " - f"OR (attack_success = 1 AND LOWER(verdict) NOT IN " - f"('attack_success', 'attack', 'attack_failed', 'benign'))) " + f"AND LOWER(triage_attack_verdict) = 'attack' " f"THEN 1 ELSE 0 END), 0) AS attack_count " f"FROM {settings['facts_table']} WHERE {where_clause} AND event_time IS NOT NULL " f"GROUP BY bucket_index ORDER BY bucket_index", @@ -2152,7 +3015,6 @@ def _read_sqlite_triage(paths): "OR LOWER(triage_status) = 'follower_reused')" ) failed_condition = "LOWER(triage_status) IN ('failed', 'error')" - resolved_condition = f"NOT ({failed_condition})" new_triage_condition = ( f"NOT {cache_condition} AND NOT {follower_condition} AND NOT ({failed_condition})" ) @@ -2166,17 +3028,18 @@ def _read_sqlite_triage(paths): f"THEN 1 ELSE 0 END), 0), " f"COALESCE(SUM(CASE WHEN {new_triage_condition} " f"THEN 1 ELSE 0 END), 0), " - f"COALESCE(SUM(CASE WHEN {resolved_condition} AND " - f"(LOWER(verdict) = 'attack_success' OR (attack_success = 1 AND LOWER(verdict) NOT IN " - f"('attack_success', 'attack', 'attack_failed', 'benign'))) THEN 1 ELSE 0 END), 0), " - f"COALESCE(SUM(CASE WHEN {resolved_condition} AND LOWER(verdict) = 'attack' " + f"COALESCE(SUM(CASE WHEN LOWER(triage_attack_verdict) = 'attack' " + f"AND LOWER(triage_attack_success) = 'success' THEN 1 ELSE 0 END), 0), " + f"COALESCE(SUM(CASE WHEN LOWER(triage_attack_verdict) = 'attack' " + f"AND LOWER(triage_attack_success) NOT IN ('success', 'failed') " f"THEN 1 ELSE 0 END), 0), " - f"COALESCE(SUM(CASE WHEN {resolved_condition} AND LOWER(verdict) = 'attack_failed' " + f"COALESCE(SUM(CASE WHEN LOWER(triage_attack_verdict) = 'attack' " + f"AND LOWER(triage_attack_success) = 'failed' " f"THEN 1 ELSE 0 END), 0), " - f"COALESCE(SUM(CASE WHEN {resolved_condition} AND LOWER(verdict) = 'benign' " + f"COALESCE(SUM(CASE WHEN LOWER(triage_attack_verdict) = 'non_attack' " f"THEN 1 ELSE 0 END), 0), " - f"COALESCE(SUM(CASE WHEN {resolved_condition} AND LOWER(verdict) NOT IN " - f"('attack_success', 'attack', 'attack_failed', 'benign') AND attack_success <> 1 " + f"COALESCE(SUM(CASE WHEN LOWER(triage_attack_verdict) " + f"NOT IN ('attack', 'non_attack') " f"THEN 1 ELSE 0 END), 0), MIN(event_time), MAX(event_time), " f"COALESCE(ROUND(AVG(CASE WHEN triage_ms > 0 THEN triage_ms END)), 0) " f"FROM {settings['facts_table']} WHERE {triage_where}", @@ -2321,9 +3184,8 @@ def _read_triage(paths): total_records += 1 file_total += 1 - verdict = _norm(obj.get("attack_verdict") or "unknown") - if verdict not in {"attack_success", "attack", "attack_failed", "benign", "unknown"}: - verdict = "unknown" + verdict = _triage_attack_verdict_value(obj) + outcome = _triage_outcome_key(obj) source = _norm(obj.get("_source_type") or obj.get("source_type") or obj.get("device_type")) source_counter[source] += 1 threat_type_counter[_norm(obj.get("_threat_type") or obj.get("threat_type"))] += 1 @@ -2332,7 +3194,7 @@ def _read_triage(paths): if triage_ms > 0: triage_ms_total += triage_ms triage_ms_count += 1 - _update_profile_counters(obj, profile_counters) + _update_profile_counters(obj, profile_counters, triage_result=True) event_start, event_end = _merge_record_time(event_start, event_end, obj) triage_source = _norm(obj.get("triage_source")) triage_status = _norm(obj.get("triage_status")) @@ -2340,10 +3202,8 @@ def _read_triage(paths): triage_failed = triage_status in {"failed", "error"} if not triage_failed: - if verdict == "unknown" and obj.get("attack_success") is True: - verdict = "attack_success" - verdict_counter[verdict] += 1 - if verdict in {"attack_success", "attack", "attack_failed"}: + verdict_counter[outcome] += 1 + if verdict == "attack": file_attack += 1 else: fallback_failed += 1 @@ -2371,7 +3231,7 @@ def _read_triage(paths): attack = verdict_counter["attack"] attack_failed = verdict_counter["attack_failed"] attack_total = attack_success + attack + attack_failed - benign = verdict_counter["benign"] + benign = verdict_counter["non_attack"] unknown = verdict_counter["unknown"] series_total = _expand_series(series_total, total_records, seed=23) @@ -2422,10 +3282,15 @@ def _new_profile_counters(): } -def _update_profile_counters(obj, counters): +def _update_profile_counters(obj, counters, *, triage_result=False): counters["phaseCounter"][_norm(obj.get("threat_phase") or obj.get("attack_phase") or obj.get("kill_chain_phase"))] += 1 counters["directionCounter"][_norm(obj.get("direction") or obj.get("traffic_direction"))] += 1 - counters["resultCounter"][_norm(obj.get("threat_result") or obj.get("attack_verdict"))] += 1 + result = ( + _triage_attack_success_value(obj) + if triage_result + else obj.get("threat_result") or obj.get("attack_verdict") + ) + counters["resultCounter"][_norm(result)] += 1 counters["protocolCounter"][_norm(obj.get("net_type") or obj.get("net_app_proto") or obj.get("protocol"))] += 1 counters["severityCounter"][_norm(obj.get("threat_severity"))] += 1 counters["responseCounter"][_norm(obj.get("rsp_status_code") or obj.get("status_code"))] += 1 @@ -2476,7 +3341,7 @@ def _build_closed_loop(triage): total = triage["totalRecords"] auto_closed = triage["attackFailed"] + triage["benign"] manual = triage["unknown"] - pending = triage["triageFailed"] + triage["unknown"] + pending = triage["unknown"] resolved = max(total - pending, 0) return { "autoClosed": auto_closed, diff --git a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/routes.yaml b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/routes.yaml index 8f85fbabb..020897392 100644 --- a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/routes.yaml +++ b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/routes.yaml @@ -9,3 +9,8 @@ routes: handler: handlers.get_activity timeoutMs: 5000 description: Incremental alert denoise and triage activity + - method: GET + path: /task-center + handler: handlers.get_task_center + timeoutMs: 5000 + description: Task center scheduler and workflow execution statistics diff --git a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx index 9ebc10b18..84640a3ad 100644 --- a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx +++ b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx @@ -74,6 +74,7 @@ const EMPTY_STATS = { }, sources: [], closedLoop: { autoClosed: 0, resolved: 0, manualDecision: 0, pending: 0, resolutionRate: 0 }, + tokenUsage: { totalTokens: 0, todayTokens: 0, todayRequests: 0, dailySeries: [], dailyLabels: [], source: '' }, verdicts: [], attackProfile: [], topThreatTypes: [], @@ -87,7 +88,18 @@ const EVENT_RAIL_TASK_LIMIT = 10; const ACTIVITY_POLL_MS = 3000; const ACTIVITY_REPLAY_WINDOW_MS = 10 * 60 * 1000; const ACTIVITY_SEEN_KEY = 'soc-dashboard-seen-activity-v1'; +const EVENT_RAIL_DEFAULT_WIDTH = 330; +const EVENT_RAIL_MIN_WIDTH = 280; +const EVENT_RAIL_MAX_WIDTH = 560; +const EVENT_RAIL_COMPACT_WIDTH = 292; const DEFAULT_TIME_RANGE = '7d'; +const DEFAULT_COMMAND_TITLE = 'Flocks AI 智能告警态势中心'; +const CUSTOM_COMMAND_TITLE_KEY = 'soc-dashboard-custom-title-v1'; +const CUSTOM_COMMAND_TITLE_CHANGED_EVENT = 'soc-dashboard:title-changed'; +const SOC_MOCK_ACTIVITY_KEY = 'soc-dashboard-mock-activity-v1'; +const SOC_MOCK_TASK_CENTER_KEY = 'soc-dashboard-mock-task-center-v1'; +const SOC_MOCK_DASHBOARD_KEY = 'soc-dashboard-mock-v1'; +const SOC_MOCK_TRUE_VALUES = ['1', 'true', 'yes', 'on']; const TIME_RANGE_OPTIONS = [ { value: '15m', label: '最近15分钟' }, { value: '2h', label: '最近2小时' }, @@ -113,6 +125,44 @@ const REFRESH_INTERVAL_MS = { '1h': 3600000, }; +function readCustomCommandTitle() { + if (typeof window === 'undefined') return ''; + try { + return window.localStorage.getItem(CUSTOM_COMMAND_TITLE_KEY)?.trim() || ''; + } catch { + return ''; + } +} + +function isMockSwitchEnabled(value) { + return SOC_MOCK_TRUE_VALUES.includes(String(value || '').trim().toLowerCase()); +} + +function readMockDashboardEnabled() { + if (typeof window === 'undefined') return false; + try { + const params = new URLSearchParams(window.location.search || ''); + for (const key of ['mockActivity', 'mockTaskCenter', 'mockDashboard']) { + const queryValue = params.get(key); + if (queryValue !== null) return isMockSwitchEnabled(queryValue); + } + return [ + SOC_MOCK_ACTIVITY_KEY, + SOC_MOCK_TASK_CENTER_KEY, + SOC_MOCK_DASHBOARD_KEY, + ].some((key) => isMockSwitchEnabled(window.localStorage.getItem(key))); + } catch { + return false; + } +} + +function defaultEventRailWidth() { + if (typeof window === 'undefined') return EVENT_RAIL_DEFAULT_WIDTH; + if (window.innerWidth <= 1120) return EVENT_RAIL_MIN_WIDTH; + if (window.innerWidth <= 1360) return EVENT_RAIL_COMPACT_WIDTH; + return EVENT_RAIL_DEFAULT_WIDTH; +} + function emptyActivityBatch() { return { mode: 'normal', @@ -142,6 +192,291 @@ function createActivityState() { }; } +function createMockActivityEvent(overrides) { + const now = new Date(); + return { + eventId: overrides.eventId || `mock-${overrides.stage}-${overrides.alert?.id || Math.random().toString(36).slice(2)}`, + stage: overrides.stage, + status: overrides.status || 'running', + occurredAt: overrides.occurredAt || now.toISOString(), + triggerSource: overrides.triggerSource || 'mock', + playbackMode: overrides.playbackMode || 'normal', + playbackStartedAt: overrides.playbackStartedAt || Date.now() - 4200, + sampleCount: overrides.sampleCount || 1, + alert: { + id: overrides.alert?.id || `mock-alert-${overrides.stage}`, + threatName: overrides.alert?.threatName || '模拟告警', + sourceType: overrides.alert?.sourceType || 'mock', + srcIp: overrides.alert?.srcIp || '10.24.8.16', + dstIp: overrides.alert?.dstIp || '172.16.32.20', + requestUri: overrides.alert?.requestUri || '/api/admin/export', + }, + result: { + isDuplicate: false, + clusterId: 'MOCK-03', + riskLevel: 'high', + verdictLabel: '待确认', + durationMs: 18000, + rawCount: 8, + normalizedCount: 8, + reducedCount: 3, + uniqueCount: 5, + filterRemovedCount: 2, + duplicateCount: 1, + reductionRate: 0.375, + ...(overrides.result || {}), + }, + hiddenFromQueue: Boolean(overrides.hiddenFromQueue), + }; +} + +function createMockActivityState() { + const now = Date.now(); + const denoiseCurrent = createMockActivityEvent({ + stage: 'denoise', + eventId: 'mock-denoise-current', + playbackMode: 'burst', + playbackStartedAt: now - 3600, + sampleCount: 6, + alert: { + id: 'mock-alert-login-burst', + threatName: '异常登录爆发（Mock）', + sourceType: 'skyeye', + srcIp: '10.23.18.44', + dstIp: '172.16.8.21', + requestUri: '/login', + }, + result: { + clusterId: 'MOCK-LOGIN-07', + rawCount: 18, + normalizedCount: 18, + reducedCount: 11, + uniqueCount: 7, + duplicateCount: 7, + filterRemovedCount: 0, + reductionRate: 0.6111, + }, + }); + const triageCurrent = createMockActivityEvent({ + stage: 'triage', + eventId: 'mock-triage-current', + playbackStartedAt: now - 6200, + alert: { + id: 'mock-alert-rce', + threatName: '远程命令执行攻击（Mock）', + sourceType: 'tdp', + srcIp: '203.0.113.41', + dstIp: '10.12.4.18', + requestUri: '/cgi-bin/luci/;stok=/locale', + }, + result: { + riskLevel: 'high', + verdictLabel: '待确认', + triageSource: 'llm', + durationMs: 22000, + threatSeverity: 'high', + }, + }); + const triageWaiting = createMockActivityEvent({ + stage: 'triage', + status: 'queued', + eventId: 'mock-triage-waiting', + occurredAt: new Date(now - 18000).toISOString(), + playbackStartedAt: now - 18000, + alert: { + id: 'mock-alert-sql', + threatName: 'SQL 注入探测（Mock）', + sourceType: 'onesec', + srcIp: '198.51.100.12', + dstIp: '10.12.4.32', + requestUri: '/search?q=1%27', + }, + result: { + riskLevel: 'medium', + verdictLabel: '待确认', + durationMs: 0, + }, + }); + const denoiseWaiting = createMockActivityEvent({ + stage: 'denoise', + status: 'queued', + eventId: 'mock-denoise-waiting', + occurredAt: new Date(now - 26000).toISOString(), + playbackStartedAt: now - 26000, + sampleCount: 4, + alert: { + id: 'mock-alert-scan', + threatName: '端口扫描聚类（Mock）', + sourceType: 'qingteng', + srcIp: '192.0.2.88', + dstIp: '10.12.5.10', + requestUri: 'TCP/22,80,443', + }, + result: { + clusterId: 'MOCK-SCAN-02', + rawCount: 12, + normalizedCount: 12, + reducedCount: 8, + uniqueCount: 4, + reductionRate: 0.6667, + }, + }); + return { + ...createActivityState(), + connection: 'online', + mode: 'burst', + denoise: { current: denoiseCurrent, queue: [denoiseWaiting], last: null }, + triage: { current: triageCurrent, queue: [triageWaiting], last: null }, + recent: [denoiseCurrent, triageCurrent, triageWaiting, denoiseWaiting], + batch: { + mode: 'burst', + windowMs: ACTIVITY_POLL_MS, + receivedCount: 22, + duplicateCount: 8, + uniqueCount: 14, + clusterCount: 5, + triageUpdatedCount: 3, + sampledCount: 10, + suppressedCount: 2, + ratePerSecond: 7.3, + }, + batchUpdatedAt: now, + generatedAt: new Date(now).toISOString(), + mock: true, + }; +} + +function activityHasVisibleEvents(activity) { + return Boolean( + activity?.denoise?.current + || activity?.denoise?.last + || activity?.denoise?.queue?.length + || activity?.triage?.current + || activity?.triage?.last + || activity?.triage?.queue?.length + || activity?.recent?.length + ); +} + +function createTaskCenterState() { + return { + connection: 'initializing', + generatedAt: '', + sessionCount: 0, + scheduledExecutionCount: 0, + scheduledTodayExecutionCount: 0, + workflowExecutionCount: 0, + workflowTodayExecutionCount: 0, + scheduledTasks: [], + workflows: [], + error: '', + }; +} + +function createMockTaskCenterState() { + const now = Date.now(); + const startedAt = now - 7 * 60 * 1000; + const nextRunAt = new Date(now + 18 * 60 * 1000).toISOString(); + const lastRunAt = new Date(now - 11 * 60 * 1000).toISOString(); + return { + ...createTaskCenterState(), + connection: 'online', + generatedAt: new Date(now).toISOString(), + sessionCount: 6, + scheduledExecutionCount: 18, + scheduledTodayExecutionCount: 5, + workflowExecutionCount: 42, + workflowTodayExecutionCount: 9, + scheduledTasks: [ + { + id: 'mock-soc-scheduler-patrol', + name: 'SOC 告警自动巡检（Mock）', + mode: 'cron', + status: 'active', + executionMode: 'workflow', + workflowId: 'stream_alert_denoise', + executionCount: 12, + todayExecutionCount: 4, + successCount: 10, + successRate: 0.8333, + activeCount: 1, + lastStatus: 'running', + lastRunAt, + nextRunAt, + cron: '*/15 * * * *', + cronDescription: '每 15 分钟', + }, + { + id: 'mock-soc-scheduler-triage', + name: '高危告警智能研判（Mock）', + mode: 'cron', + status: 'active', + executionMode: 'workflow', + workflowId: 'stream_alert_triage', + executionCount: 6, + todayExecutionCount: 1, + successCount: 5, + successRate: 0.8333, + activeCount: 0, + lastStatus: 'completed', + lastRunAt: new Date(now - 42 * 60 * 1000).toISOString(), + nextRunAt: new Date(now + 36 * 60 * 1000).toISOString(), + cron: '*/30 * * * *', + cronDescription: '每 30 分钟', + }, + ], + workflows: [ + { + id: 'stream_alert_denoise', + name: '告警降噪工作流（Mock）', + executionCount: 24, + todayExecutionCount: 6, + successCount: 21, + successRate: 0.875, + activeCount: 1, + lastStatus: 'running', + lastRunAt: startedAt, + latestExecutionHash: 'mock-denoise-run-001', + latestAlertName: '异常登录爆发（Mock）', + progressPercent: 0.58, + progressLabel: '第 4/7 步', + currentPhase: '聚类降噪', + sessionId: '', + messageId: '', + }, + { + id: 'stream_alert_triage', + name: '告警研判工作流（Mock）', + executionCount: 18, + todayExecutionCount: 3, + successCount: 15, + successRate: 0.8333, + activeCount: 1, + lastStatus: 'running', + lastRunAt: now - 4 * 60 * 1000, + latestExecutionHash: 'mock-triage-run-002', + latestAlertName: '远程命令执行攻击（Mock）', + progressPercent: 0.67, + progressLabel: '第 2/3 步', + currentPhase: '证据汇总', + sessionId: '', + messageId: '', + }, + ], + mock: true, + }; +} + +function taskCenterHasVisibleRows(taskCenter) { + return Boolean( + taskCenter?.scheduledTasks?.length + || taskCenter?.workflows?.length + || taskCenter?.sessionCount + || taskCenter?.scheduledExecutionCount + || taskCenter?.workflowExecutionCount + ); +} + function activityDuration(event) { if (!event) return 0; if (event.stage === 'denoise') { @@ -154,6 +489,11 @@ function activityDuration(event) { return 30000; } +function isRunningWorkflowEvent(event) { + return event?.triggerSource === 'workflow_execution' + && ['running', 'queued', 'pending'].includes(String(event?.status || '').toLowerCase()); +} + function normalizeActivityBatch(raw) { const batch = { ...emptyActivityBatch(), ...(raw || {}) }; for (const key of ['windowMs', 'receivedCount', 'duplicateCount', 'uniqueCount', 'clusterCount', 'triageUpdatedCount', 'sampledCount', 'suppressedCount', 'ratePerSecond']) { @@ -205,8 +545,17 @@ function enqueueActivity(previous, events, generatedAt, recentEvents, rawBatch) ? { ...event, playbackMode: batch.mode, batch } : event; const lane = next[enriched.stage]; - const known = [lane.current, lane.last, ...lane.queue].some((item) => item?.eventId === enriched.eventId); - if (!known) lane.queue.push(enriched); + if (lane.current?.eventId === enriched.eventId) { + lane.current = { ...lane.current, ...enriched, playbackStartedAt: lane.current.playbackStartedAt }; + continue; + } + const queuedIndex = lane.queue.findIndex((item) => item?.eventId === enriched.eventId); + if (queuedIndex >= 0) { + lane.queue[queuedIndex] = { ...lane.queue[queuedIndex], ...enriched }; + continue; + } + const knownLast = lane.last?.eventId === enriched.eventId; + if (!knownLast || isRunningWorkflowEvent(enriched)) lane.queue.push(enriched); } for (const kind of ['denoise', 'triage']) { const lane = next[kind]; @@ -365,6 +714,7 @@ function mergeStats(raw) { triage: { ...EMPTY_STATS.triage, ...((raw || {}).triage || {}) }, pipeline: { ...EMPTY_STATS.pipeline, ...((raw || {}).pipeline || {}) }, closedLoop: { ...EMPTY_STATS.closedLoop, ...((raw || {}).closedLoop || {}) }, + tokenUsage: { ...EMPTY_STATS.tokenUsage, ...((raw || {}).tokenUsage || {}) }, dateRange: { ...EMPTY_STATS.dateRange, ...((raw || {}).dateRange || {}) }, eventRange: { ...EMPTY_STATS.eventRange, ...((raw || {}).eventRange || {}) }, timeline: { ...EMPTY_STATS.timeline, ...((raw || {}).timeline || {}) }, @@ -418,6 +768,12 @@ function compactNumber(value) { return fullNumber(n); } +function formatTokenVolume(value) { + const n = Math.max(Number(value || 0), 0); + if (n >= 1000000000) return `${(n / 1000000000).toFixed(2)}B`; + return `${(n / 1000000).toFixed(2)}M`; +} + function AnimatedNumber({ value, format, tag = 'span', className, duration = 900 }) { const { useEffect, useRef, useState } = getReact(); const target = Number(value || 0); @@ -913,6 +1269,22 @@ function Sparkline({ values, color }) { ]); } +function TokenSparkline({ values, color }) { + const nums = (values || []).map((v) => Math.max(Number(v || 0), 0)); + const series = nums.length > 1 ? nums : [0, nums[0] || 0]; + const max = Math.max(...series, 1); + const min = Math.min(...series); + const span = Math.max(max - min, 1); + const points = series.map((value, index) => { + const x = series.length === 1 ? 0 : (index / (series.length - 1)) * 300; + const y = 80 - ((value - min) / span) * 52; + return `${x},${y}`; + }).join(' '); + return h('svg', { className: 'token-chart', viewBox: '0 0 300 92', role: 'img' }, [ + h('polyline', { key: 'line', className: 'token-chart-line', points, fill: 'none', stroke: color || '#ff674d', strokeWidth: 4, strokeLinecap: 'round', strokeLinejoin: 'round' }), + ]); +} + function polarPoint(cx, cy, radius, angle) { const radians = (angle - 90) * Math.PI / 180; return { @@ -1240,7 +1612,7 @@ function TimeRefreshPopover({ value, refreshValue, open, onToggle, onApply, onCl ]); } -function CommandHeader({ timeFilter, refreshKey, timeMenuOpen, setTimeMenuOpen, applyTimeRefresh, stats, loading, refresh, activity }) { +function CommandHeader({ title, timeFilter, refreshKey, timeMenuOpen, setTimeMenuOpen, applyTimeRefresh, stats, loading, refresh, activity }) { const active = activity.denoise.current || activity.triage.current; const loadActive = activity.mode !== 'normal' && activity.batch?.receivedCount > 0; const status = activity.connection === 'error' @@ -1252,7 +1624,7 @@ function CommandHeader({ timeFilter, refreshKey, timeMenuOpen, setTimeMenuOpen, h('div', { className: 'command-brand', key: 'brand' }, [ h('div', { className: 'command-logo', key: 'logo' }, 'AI'), h('div', { key: 'copy' }, [ - h('strong', { key: 'title' }, 'Flocks AI 智能告警态势中心'), + h('strong', { key: 'title', title }, title), h('span', { key: 'sub' }, '告警汇聚 · 智能降噪 · 自动研判 · 风险聚合'), ]), ]), @@ -1417,15 +1789,15 @@ function CommandGraph({ stats, activity }) { h('em', { key: 'text' }, triageActive ? '处理中' : '安全事件'), ]), ]), - h('div', { title: 'AI 研判结论为良性的事件数量', key: 'benign' }, [ + h('div', { title: 'AI 研判结论为非攻击或良性的事件数量', key: 'benign' }, [ h(AnimatedNumber, { tag: 'b', value: stats.triage.benign, key: 'value' }), h('span', { className: 'outcome-label', key: 'label' }, [ h('i', { key: 'ai' }, 'AI判定'), h('em', { key: 'text' }, '非安全事件'), ]), ]), - h('div', { title: 'AI 判定需要进入人工复核的事件数量', key: 'manual' }, [ - h(AnimatedNumber, { tag: 'b', value: stats.closedLoop.pending, key: 'value' }), + h('div', { title: 'AI 研判结论为未知、需要进入人工复核的事件数量', key: 'manual' }, [ + h(AnimatedNumber, { tag: 'b', value: stats.triage.unknown, key: 'value' }), h('span', { className: 'outcome-label', key: 'label' }, [ h('i', { key: 'ai' }, 'AI判定'), h('em', { key: 'text' }, '待人工复核'), @@ -1455,12 +1827,34 @@ function CommandMetric({ label, value, format, sub, values, color }) { ]); } +function TokenUsageMetric({ tokenUsage }) { + return h('div', { className: 'command-metric token-usage-metric', style: { '--metric-color': '#ff674d' } }, [ + h('span', { className: 'token-title', key: 'label' }, 'Token 消耗'), + h('div', { className: 'token-summary-row', key: 'summary' }, [ + h(AnimatedNumber, { + tag: 'b', + className: 'token-value', + value: tokenUsage.todayTokens, + format: formatTokenVolume, + duration: 1200, + key: 'value', + }), + h('small', { className: 'token-sub', key: 'sub' }, [ + h('span', { className: 'token-sub-line', key: 'total' }, `累计 ${formatTokenVolume(tokenUsage.totalTokens)}`), + h('span', { className: 'token-sub-line', key: 'today' }, `今日调用 ${compactNumber(tokenUsage.todayRequests)} 次`), + ]), + ]), + h(TokenSparkline, { values: tokenUsage.dailySeries, color: '#ff674d', key: 'chart' }), + ]); +} + function CommandMetrics({ stats }) { + const tokenUsage = stats.tokenUsage || EMPTY_STATS.tokenUsage; return h('section', { className: 'command-metrics' }, [ h(CommandMetric, { label: '原始告警量', value: stats.denoise.totalRaw, sub: `${compactNumber(stats.denoise.totalUnique)} 条进入研判`, values: stats.timeline.denoiseRaw, color: '#2e72ff', key: 'raw' }), h(CommandMetric, { label: '安全事件量', value: stats.triage.attackTotal, sub: `${compactNumber(stats.triage.attackSuccess)} 条攻击成功`, values: stats.timeline.triageAttack, color: '#23ca8e', key: 'events' }), h(CommandMetric, { label: '降噪率', value: stats.denoise.duplicateRate * 100, format: (value) => `${trim(value)}%`, sub: `${compactNumber(stats.denoise.duplicates)} 条告警已过滤/收敛`, values: stats.timeline.denoiseUnique, color: '#21d8a3', key: 'rate' }), - h(CommandMetric, { label: '平均研判时间', value: stats.triage.avgTriageMs, format: (value) => formatDurationMs(value), sub: `${compactNumber(stats.triage.totalRecords)} 条已完成研判`, values: stats.timeline.triageTotal, color: '#ff674d', key: 'mtta' }), + h(TokenUsageMetric, { tokenUsage, key: 'tokens' }), ]); } @@ -1657,7 +2051,253 @@ function EventQueueProgress({ event }) { ]); } -function CommandEventRail({ activity, timeFilter, collapsed, onToggle }) { +function taskCenterPercent(value) { + return `${Math.round(Math.max(Math.min(Number(value || 0), 1), 0) * 100)}%`; +} + +function taskCenterTimeLabel(value) { + if (!value) return '暂无记录'; + const raw = Number(value); + const date = Number.isFinite(raw) && raw > 0 ? new Date(raw) : new Date(value); + if (Number.isNaN(date.getTime())) return '暂无记录'; + return date.toLocaleString('zh-CN', { + month: '2-digit', + day: '2-digit', + hour: '2-digit', + minute: '2-digit', + hour12: false, + }); +} + +function taskCenterStatusLabel(status) { + const value = String(status || '').toLowerCase(); + if (['running', 'queued', 'pending'].includes(value)) return '执行中'; + if (['completed', 'success'].includes(value)) return '成功'; + if (['failed', 'error', 'timeout'].includes(value)) return '失败'; + if (['disabled', 'stopped'].includes(value)) return '已关闭'; + if (value === 'stale') return '已停止'; + if (value === 'cancelled') return '取消'; + return '待执行'; +} + +function taskCenterHashLabel(value) { + const text = String(value || '').trim(); + if (!text) return '--'; + if (text.length <= 12) return text; + return `${text.slice(0, 6)}...${text.slice(-4)}`; +} + +function taskCenterHashValue(value) { + const text = String(value || '').trim(); + return text || '--'; +} + +function taskCenterWorkflowName(item) { + const name = String(item?.name || item?.id || '').trim(); + const id = String(item?.id || '').trim(); + if (/^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i.test(name)) { + return `动态工作流 ${taskCenterHashLabel(name)}`; + } + return name || id || '未命名工作流'; +} + +function taskCenterProgressValue(item) { + return Math.max(Math.min(Number(item?.progressPercent || 0), 1), 0); +} + +function taskCenterProgressLabel(item) { + return String(item?.progressLabel || '').trim() || '待执行'; +} + +function openTaskCenterConversation(item) { + const sessionId = String(item?.sessionId || item?.sessionID || '').trim(); + if (!sessionId || typeof window === 'undefined') return false; + const messageId = String(item?.messageId || item?.messageID || '').trim(); + const params = new URLSearchParams({ session: sessionId }); + if (messageId) params.set('focusMessage', messageId); + window.location.href = `/sessions?${params.toString()}`; + return true; +} + +function openConversationFromEvent(event) { + const sessionId = String(event?.sessionId || event?.sessionID || '').trim(); + if (!sessionId || typeof window === 'undefined') return false; + const messageId = String(event?.messageId || event?.messageID || '').trim(); + const params = new URLSearchParams({ session: sessionId }); + if (messageId) params.set('focusMessage', messageId); + window.location.href = `/sessions?${params.toString()}`; + return true; +} + +function TaskCenterSummary({ taskCenter }) { + const scheduledTasks = taskCenter.scheduledTasks || []; + const workflows = taskCenter.workflows || []; + const activeCount = [ + ...scheduledTasks, + ...workflows, + ].filter((item) => Number(item.activeCount || 0) > 0).length; + const summaryMetric = (key, label, value, sub, className = '') => h('div', { className, key }, [ + h('span', { key: 'label' }, label), + h(AnimatedNumber, { tag: 'b', value: value || 0, duration: 800, key: 'value' }), + sub ? h('small', { key: 'sub' }, sub) : null, + ]); + return h('div', { className: 'task-center-summary' }, [ + summaryMetric('sessions', '会话次数', taskCenter.sessionCount, ''), + summaryMetric('active', '执行中', activeCount, '', activeCount ? 'active' : ''), + summaryMetric('scheduledRuns', '定时执行', taskCenter.scheduledExecutionCount, `今日 ${taskCenter.scheduledTodayExecutionCount || 0}`), + summaryMetric('workflowRuns', '工作流执行', taskCenter.workflowExecutionCount, `今日 ${taskCenter.workflowTodayExecutionCount || 0}`), + ]); +} + +function TaskCenterItem({ item, kind }) { + const successRate = Math.max(Math.min(Number(item.successRate || 0), 1), 0); + const progressValue = taskCenterProgressValue(item); + const progressLabel = taskCenterProgressLabel(item); + const active = Number(item.activeCount || 0) > 0; + const status = taskCenterStatusLabel(item.lastStatus); + const statusClass = String(item.lastStatus || '').toLowerCase(); + const latestTime = taskCenterTimeLabel(item.lastRunAt); + const latestExecutionHash = taskCenterHashValue(item.latestExecutionHash); + const itemName = kind === 'workflow' ? taskCenterWorkflowName(item) : item.name || item.id; + const alertName = String(item.latestAlertName || '').trim(); + const hasConversation = kind === 'workflow' && Boolean(String(item.sessionId || item.sessionID || '').trim()); + const sub = kind === 'scheduled' + ? item.nextRunAt + ? `下次 ${taskCenterTimeLabel(item.nextRunAt)}` + : item.cronDescription || item.cron || taskCenterTimeLabel(item.lastRunAt) + : `最近执行 ${latestTime}`; + const stats = kind === 'workflow' + ? [ + h('span', { key: 'total' }, ['执行 ', h(AnimatedNumber, { tag: 'b', value: item.executionCount || 0, duration: 700, key: 'value' })]), + h('span', { key: 'today' }, ['今日 ', h(AnimatedNumber, { tag: 'b', value: item.todayExecutionCount || 0, duration: 700, key: 'value' })]), + h('span', { key: 'progress' }, ['进度 ', h('b', { key: 'value' }, progressLabel)]), + h('span', { key: 'rate' }, ['成功率 ', h('b', { key: 'value' }, taskCenterPercent(successRate))]), + ] + : [ + h('span', { key: 'total' }, ['执行 ', h(AnimatedNumber, { tag: 'b', value: item.executionCount || 0, duration: 700, key: 'value' })]), + h('span', { key: 'today' }, ['今日 ', h(AnimatedNumber, { tag: 'b', value: item.todayExecutionCount || 0, duration: 700, key: 'value' })]), + h('span', { key: 'success' }, ['成功 ', h(AnimatedNumber, { tag: 'b', value: item.successCount || 0, duration: 700, key: 'value' })]), + h('span', { key: 'rate' }, ['成功率 ', h('b', { key: 'value' }, taskCenterPercent(successRate))]), + ]; + const handleOpen = () => { + if (hasConversation) openTaskCenterConversation(item); + }; + const handleKeyDown = (event) => { + if (!hasConversation) return; + if (event.key === 'Enter' || event.key === ' ') { + event.preventDefault(); + openTaskCenterConversation(item); + } + }; + return h('article', { + className: cx('task-center-item', active && 'active', hasConversation && 'clickable'), + role: hasConversation ? 'button' : undefined, + tabIndex: hasConversation ? 0 : undefined, + title: hasConversation ? '打开对应对话' : undefined, + onClick: handleOpen, + onKeyDown: handleKeyDown, + }, [ + h('div', { className: 'task-center-item-head', key: 'head' }, [ + h('strong', { title: item.name || item.id, key: 'name' }, itemName), + h('span', { className: cx('task-center-status', active && 'active', statusClass), key: 'status' }, active ? '执行中' : status), + ]), + h('div', { className: 'task-center-item-sub', title: sub, key: 'sub' }, sub), + kind === 'workflow' ? h('div', { + className: cx('task-center-alert', alertName && 'has-alert'), + title: alertName || '暂无告警名称', + key: 'alert', + }, [ + h('span', { key: 'label' }, '研判告警'), + h('b', { key: 'value' }, alertName || '暂无告警名称'), + ]) : null, + kind === 'workflow' ? h('div', { className: 'task-center-hash', title: latestExecutionHash, key: 'hash' }, [ + h('span', { key: 'label' }, '执行哈希'), + h('code', { key: 'value' }, latestExecutionHash), + h('span', { key: 'link-label' }, '更多信息'), + h('code', { className: cx('task-center-jump', hasConversation && 'enabled'), key: 'link' }, hasConversation ? '查看对话' : '暂无关联对话'), + ]) : null, + h('div', { className: cx('task-center-stats', kind === 'workflow' && 'workflow-stats'), key: 'stats' }, stats), + h('div', { + className: cx('task-center-rate', kind === 'workflow' && 'progress-rate'), + style: { '--task-center-rate': kind === 'workflow' ? progressValue : successRate }, + key: 'rateBar', + }, [ + h('i', { key: 'fill' }), + ]), + ]); +} + +function TaskCenterSection({ title, count, items, kind, emptyText, expanded, onToggle, collapsed, onCollapseToggle }) { + const hasOverflow = items.length > 3; + const visibleItems = expanded || !hasOverflow ? items : items.slice(0, 3); + return h('section', { className: 'task-center-section' }, [ + h('div', { className: 'task-center-section-title', key: 'title' }, [ + h('button', { + className: 'task-center-section-toggle', + type: 'button', + 'aria-expanded': !collapsed, + onClick: onCollapseToggle, + key: 'toggle', + }, [ + h('i', { key: 'chevron' }, collapsed ? '›' : '⌄'), + h('strong', { key: 'label' }, title), + ]), + h('span', { key: 'count' }, collapsed ? `${count} 项` : expanded || !hasOverflow ? `${count} 项` : `显示 3/${count}`), + ]), + collapsed ? null : h('div', { className: 'task-center-section-list', key: 'list' }, visibleItems.length + ? visibleItems.map((item) => h(TaskCenterItem, { item, kind, key: `${kind}-${item.id}` })) + : h('div', { className: 'event-rail-empty' }, emptyText)), + !collapsed && hasOverflow ? h('button', { + className: 'task-center-expand', + type: 'button', + onClick: onToggle, + key: 'expand', + }, expanded ? '收起' : `展开全部 ${count} 项`) : null, + ]); +} + +function CommandTaskCenterPanel({ taskCenter }) { + const { useState } = getReact(); + const [scheduledExpanded, setScheduledExpanded] = useState(false); + const [workflowExpanded, setWorkflowExpanded] = useState(false); + const [scheduledCollapsed, setScheduledCollapsed] = useState(false); + const [workflowCollapsed, setWorkflowCollapsed] = useState(false); + const scheduledTasks = taskCenter.scheduledTasks || []; + const workflows = taskCenter.workflows || []; + return h('div', { className: 'task-center-panel', key: 'taskCenterPanel' }, [ + taskCenter.connection === 'error' ? h('div', { + className: 'task-center-inline-warn', + key: 'warn', + }, `任务中心连接异常：${taskCenter.error || '正在重试'}`) : null, + h(TaskCenterSummary, { taskCenter, key: 'summary' }), + h(TaskCenterSection, { + title: '定时执行', + count: scheduledTasks.length, + items: scheduledTasks, + kind: 'scheduled', + emptyText: '暂无定时任务执行记录', + expanded: scheduledExpanded, + onToggle: () => setScheduledExpanded((current) => !current), + collapsed: scheduledCollapsed, + onCollapseToggle: () => setScheduledCollapsed((current) => !current), + key: 'scheduled', + }), + h(TaskCenterSection, { + title: '工作流执行', + count: workflows.length, + items: workflows, + kind: 'workflow', + emptyText: '暂无工作流执行记录', + expanded: workflowExpanded, + onToggle: () => setWorkflowExpanded((current) => !current), + collapsed: workflowCollapsed, + onCollapseToggle: () => setWorkflowCollapsed((current) => !current), + key: 'workflows', + }), + ]); +} + +function CommandAiTaskPanel({ activity, timeFilter }) { const tasks = buildEventQueueTasks(activity, timeFilter); const filterTransitionKey = [timeFilter.mode, timeFilter.range, timeFilter.start, timeFilter.end].join('|'); const visibleTasks = useAnimatedTaskWindow( @@ -1674,11 +2314,7 @@ function CommandEventRail({ activity, timeFilter, collapsed, onToggle }) { : counts.processing ? `AI 正在并行处理 ${counts.processing} 个任务` : counts.waiting ? '最新 10 条待处理任务' : '等待新的降噪或研判任务'; - const content = collapsed ? [] : [ - h('div', { className: 'event-rail-head', key: 'head' }, [ - h('div', { key: 'title' }, [h('strong', { key: 'label' }, 'AI处理任务'), h('span', { key: 'sub' }, '最新 10 条待处理任务')]), - h(AnimatedNumber, { tag: 'b', value: queueCount, duration: 600, key: 'count' }), - ]), + return [ h('div', { className: cx('event-update-banner', activity.connection === 'error' && 'warn'), key: 'banner' }, banner), h('div', { className: 'event-rail-list', key: 'list' }, visibleTasks.length ? visibleTasks.map((task) => { const event = task.event; @@ -1691,13 +2327,31 @@ function CommandEventRail({ activity, timeFilter, collapsed, onToggle }) { ? '处理中' : '等待处理'; const detail = event?.triggerSource === 'workflow_execution' - ? event.result?.isDuplicate ? '重复告警已收敛' : '降噪处理完成' + ? task.stage === 'triage' + ? task.state === 'processing' ? '研判工作流处理中' : '研判工作流待处理' + : event.result?.isDuplicate ? '重复告警已收敛' : '降噪处理完成' : task.state === 'processing' ? task.stage === 'triage' ? '证据关联与结论生成中' : '特征提取与相似聚类中' : '等待 AI 处理'; + const hasConversation = Boolean(String(event?.sessionId || event?.sessionID || '').trim()); + const handleOpen = () => { + if (hasConversation) openConversationFromEvent(event); + }; + const handleKeyDown = (keyboardEvent) => { + if (!hasConversation) return; + if (keyboardEvent.key === 'Enter' || keyboardEvent.key === ' ') { + keyboardEvent.preventDefault(); + openConversationFromEvent(event); + } + }; return h('article', { - className: cx('event-rail-item', `state-${task.state}`, `kind-${task.stage}`, `motion-${task.motion || 'stable'}`), + className: cx('event-rail-item', `state-${task.state}`, `kind-${task.stage}`, `motion-${task.motion || 'stable'}`, hasConversation && 'clickable'), key: task.key, + role: hasConversation ? 'button' : undefined, + tabIndex: hasConversation ? 0 : undefined, + title: hasConversation ? '打开对应对话' : undefined, + onClick: handleOpen, + onKeyDown: handleKeyDown, }, [ h('div', { className: 'event-rail-meta', key: 'meta' }, [ h('span', { className: cx('event-queue-kind', `kind-${task.stage}`), key: 'kind' }, stageLabel), @@ -1706,39 +2360,162 @@ function CommandEventRail({ activity, timeFilter, collapsed, onToggle }) { ]), h('strong', { title, key: 'title' }, title), h('span', { title: eventEndpoint(event), key: 'endpoint' }, eventEndpoint(event)), - h('small', { key: 'result' }, detail), + h('small', { key: 'result' }, hasConversation ? `${detail} · 查看对话` : detail), task.state === 'processing' ? h(EventQueueProgress, { event, key: 'progress' }) : null, ]); }) : h('div', { className: 'event-rail-empty' }, '等待新的降噪或研判任务')), ]; +} + +function CommandEventRail({ activity, timeFilter, taskCenter, view, onViewChange, collapsed, onToggle, railWidth, onResizeStart, onResizeKeyDown }) { + const tasks = buildEventQueueTasks(activity, timeFilter); + const queueCount = tasks.filter((task) => task.state !== 'completed').length; + const taskCenterCount = Number(taskCenter.sessionCount || 0); + const content = collapsed ? [] : [ + h('div', { className: 'event-rail-head rail-view-head', key: 'head' }, [ + h('div', { className: 'rail-view-tabs', role: 'tablist', key: 'tabs' }, [ + h('button', { + className: cx(view === 'aiTasks' && 'active'), + type: 'button', + role: 'tab', + 'aria-selected': view === 'aiTasks', + onClick: () => onViewChange('aiTasks'), + key: 'aiTasks', + }, 'AI处理任务'), + h('button', { + className: cx(view === 'taskCenter' && 'active'), + type: 'button', + role: 'tab', + 'aria-selected': view === 'taskCenter', + onClick: () => onViewChange('taskCenter'), + key: 'taskCenter', + }, '任务中心'), + ]), + h(AnimatedNumber, { tag: 'b', value: view === 'aiTasks' ? queueCount : taskCenterCount, duration: 600, key: 'count' }), + ]), + view === 'taskCenter' + ? h(CommandTaskCenterPanel, { taskCenter, key: 'taskCenterContent' }) + : h(CommandAiTaskPanel, { activity, timeFilter, key: 'aiTaskContent' }), + ]; return h('aside', { className: cx('command-event-rail', collapsed && 'collapsed') }, [ + collapsed ? null : h('div', { + className: 'event-rail-resize', + role: 'separator', + tabIndex: 0, + 'aria-label': '调整任务面板宽度', + 'aria-orientation': 'vertical', + 'aria-valuemin': EVENT_RAIL_MIN_WIDTH, + 'aria-valuemax': EVENT_RAIL_MAX_WIDTH, + 'aria-valuenow': railWidth, + onPointerDown: onResizeStart, + onKeyDown: onResizeKeyDown, + key: 'resize', + }), h('button', { className: 'event-rail-toggle', type: 'button', - title: collapsed ? '展开 AI处理任务' : '向右折叠', - 'aria-label': collapsed ? '展开 AI处理任务' : '向右折叠 AI处理任务', + title: collapsed ? '展开任务面板' : '向右折叠', + 'aria-label': collapsed ? '展开任务面板' : '向右折叠任务面板', onClick: onToggle, key: 'toggle', - }, collapsed ? h('span', { key: 'label' }, ['任', '务', '面', '板'].map((text) => h('i', { key: text }, text))) : '›'), + }, collapsed ? '展开' : '收起'), ...content, ]); } export default function Page() { - const { useCallback, useEffect, useRef, useState } = getReact(); + const { useCallback, useEffect, useMemo, useRef, useState } = getReact(); const [timeFilter, setTimeFilter] = useState(() => createRelativeTimeFilter()); const [refreshKey, setRefreshKey] = useState('off'); const [timeMenuOpen, setTimeMenuOpen] = useState(false); const [eventRailCollapsed, setEventRailCollapsed] = useState(false); + const [eventRailWidth, setEventRailWidth] = useState(defaultEventRailWidth); + const [rightRailView, setRightRailView] = useState('aiTasks'); + const [customCommandTitle, setCustomCommandTitle] = useState(readCustomCommandTitle); + const [mockDashboardEnabled, setMockDashboardEnabled] = useState(readMockDashboardEnabled); const [stats, setStats] = useState(EMPTY_STATS); const [loading, setLoading] = useState(true); const [error, setError] = useState(''); const [activity, setActivity] = useState(createActivityState); + const [taskCenter, setTaskCenter] = useState(createTaskCenterState); const activityCursor = useRef(''); const workflowProgressByFilter = useRef(new Map()); const statsRequestId = useRef(0); const statsPending = useRef(0); + const clampEventRailWidth = useCallback((value) => { + const viewportLimit = typeof window === 'undefined' + ? EVENT_RAIL_MAX_WIDTH + : Math.max(EVENT_RAIL_MIN_WIDTH, Math.min(EVENT_RAIL_MAX_WIDTH, window.innerWidth - 760)); + return Math.max(EVENT_RAIL_MIN_WIDTH, Math.min(viewportLimit, Math.round(value))); + }, []); + + const startEventRailResize = useCallback((event) => { + if (eventRailCollapsed) return; + const startX = event.clientX; + const startWidth = eventRailWidth; + const previousCursor = document.body.style.cursor; + const previousUserSelect = document.body.style.userSelect; + document.body.style.cursor = 'col-resize'; + document.body.style.userSelect = 'none'; + + const move = (moveEvent) => { + setEventRailWidth(clampEventRailWidth(startWidth + startX - moveEvent.clientX)); + }; + const stop = () => { + document.body.style.cursor = previousCursor; + document.body.style.userSelect = previousUserSelect; + window.removeEventListener('pointermove', move); + window.removeEventListener('pointerup', stop); + window.removeEventListener('pointercancel', stop); + }; + + window.addEventListener('pointermove', move); + window.addEventListener('pointerup', stop); + window.addEventListener('pointercancel', stop); + event.preventDefault(); + }, [clampEventRailWidth, eventRailCollapsed, eventRailWidth]); + + const adjustEventRailWidth = useCallback((event) => { + if (event.key !== 'ArrowLeft' && event.key !== 'ArrowRight') return; + const delta = event.shiftKey ? 48 : 16; + setEventRailWidth((current) => clampEventRailWidth(current + (event.key === 'ArrowLeft' ? delta : -delta))); + event.preventDefault(); + }, [clampEventRailWidth]); + + useEffect(() => { + const handleResize = () => setEventRailWidth((current) => clampEventRailWidth(current)); + window.addEventListener('resize', handleResize); + return () => window.removeEventListener('resize', handleResize); + }, [clampEventRailWidth]); + + useEffect(() => { + const refreshTitle = (event) => { + const nextTitle = event?.detail && typeof event.detail.title === 'string' + ? event.detail.title + : readCustomCommandTitle(); + setCustomCommandTitle(nextTitle.trim()); + }; + const handleStorage = (event) => { + if (event.key === CUSTOM_COMMAND_TITLE_KEY) { + setCustomCommandTitle((event.newValue || '').trim()); + } else if ( + event.key === null + || event.key === SOC_MOCK_ACTIVITY_KEY + || event.key === SOC_MOCK_TASK_CENTER_KEY + || event.key === SOC_MOCK_DASHBOARD_KEY + ) { + setMockDashboardEnabled(readMockDashboardEnabled()); + } + }; + window.addEventListener(CUSTOM_COMMAND_TITLE_CHANGED_EVENT, refreshTitle); + window.addEventListener('storage', handleStorage); + return () => { + window.removeEventListener(CUSTOM_COMMAND_TITLE_CHANGED_EVENT, refreshTitle); + window.removeEventListener('storage', handleStorage); + }; + }, []); + const loadStats = useCallback(async (filter, options = {}) => { if (options.skipIfBusy && statsPending.current > 0) return; const requestId = ++statsRequestId.current; @@ -1794,6 +2571,54 @@ export default function Page() { setTimeMenuOpen(false); }, []); + useEffect(() => { + let stopped = false; + let timer = 0; + + const schedule = (delay) => { + if (!stopped) timer = window.setTimeout(() => void poll(), delay); + }; + + const poll = async () => { + if (stopped) return; + try { + const params = mockDashboardEnabled ? { mockActivity: '1' } : {}; + const response = await getApi().page.get('/task-center', { params }); + const payload = response.data || {}; + if (!stopped) { + setTaskCenter({ + ...createTaskCenterState(), + ...payload, + scheduledTasks: Array.isArray(payload.scheduledTasks) ? payload.scheduledTasks : [], + workflows: Array.isArray(payload.workflows) ? payload.workflows : [], + sessionCount: Math.max(Number(payload.sessionCount || 0), 0), + scheduledExecutionCount: Math.max(Number(payload.scheduledExecutionCount || 0), 0), + scheduledTodayExecutionCount: Math.max(Number(payload.scheduledTodayExecutionCount || 0), 0), + workflowExecutionCount: Math.max(Number(payload.workflowExecutionCount || 0), 0), + workflowTodayExecutionCount: Math.max(Number(payload.workflowTodayExecutionCount || 0), 0), + connection: 'online', + error: '', + }); + } + } catch (taskCenterError) { + if (!stopped) { + setTaskCenter((previous) => ({ + ...previous, + connection: 'error', + error: taskCenterError instanceof Error ? taskCenterError.message : 'task center api failed', + })); + } + } + schedule(ACTIVITY_POLL_MS); + }; + + void poll(); + return () => { + stopped = true; + window.clearTimeout(timer); + }; + }, [mockDashboardEnabled]); + useEffect(() => { let stopped = false; let timer = 0; @@ -1833,11 +2658,20 @@ export default function Page() { if (payload.error) throw new Error(payload.error); activityCursor.current = payload.cursor || activityCursor.current; if (!stopped) { + if (payload.tokenUsage) { + setStats((previous) => mergeStats({ ...previous, tokenUsage: payload.tokenUsage })); + } const rawIncomingEvents = bootstrap ? recentUnseenActivity(payload.recentEvents) : (payload.events || []); const incomingEvents = rawIncomingEvents.filter((event) => event?.stage !== 'denoise'); const workflowEvents = Array.isArray(payload.workflowEvents) ? payload.workflowEvents : []; + for (const workflowEvent of workflowEvents) { + const hasConversation = Boolean(String(workflowEvent?.sessionId || workflowEvent?.sessionID || '').trim()); + if (hasConversation) { + incomingEvents.push(workflowEvent); + } + } const rawCallCount = payload.workflowStats?.callCount; const hasWorkflowCount = rawCallCount !== null && rawCallCount !== undefined @@ -1924,20 +2758,22 @@ export default function Page() { useEffect(() => { const event = activity.denoise.current; if (!event) return undefined; + if (isRunningWorkflowEvent(event)) return undefined; const id = window.setTimeout(() => { setActivity((previous) => completeActivity(previous, 'denoise', event)); }, activityDuration(event)); return () => window.clearTimeout(id); - }, [activity.denoise.current?.eventId]); + }, [activity.denoise.current?.eventId, activity.denoise.current?.status]); useEffect(() => { const event = activity.triage.current; if (!event) return undefined; + if (isRunningWorkflowEvent(event)) return undefined; const id = window.setTimeout(() => { setActivity((previous) => completeActivity(previous, 'triage', event)); }, activityDuration(event)); return () => window.clearTimeout(id); - }, [activity.triage.current?.eventId]); + }, [activity.triage.current?.eventId, activity.triage.current?.status]); useEffect(() => { if (!activity.batchUpdatedAt) return undefined; @@ -1951,28 +2787,72 @@ export default function Page() { return () => window.clearTimeout(id); }, [activity.batchUpdatedAt]); - const activityBusy = Boolean( - activity.denoise.current - || activity.triage.current - || activity.denoise.queue.length - || activity.triage.queue.length - || activity.batch?.receivedCount - || activity.batch?.triageUpdatedCount + const displayActivity = useMemo( + () => ( + !activityHasVisibleEvents(activity) && mockDashboardEnabled + ? createMockActivityState() + : activity + ), + [activity, mockDashboardEnabled], + ); + const displayTaskCenter = useMemo( + () => ( + !taskCenterHasVisibleRows(taskCenter) && mockDashboardEnabled + ? createMockTaskCenterState() + : taskCenter + ), + [mockDashboardEnabled, taskCenter], + ); + const displayActivityBusy = Boolean( + displayActivity.denoise.current + || displayActivity.triage.current + || displayActivity.denoise.queue.length + || displayActivity.triage.queue.length + || displayActivity.batch?.receivedCount + || displayActivity.batch?.triageUpdatedCount ); return h('div', { - className: cx('adtd-root command-root', activityBusy && 'command-is-processing', eventRailCollapsed && 'event-rail-is-collapsed'), + className: cx('adtd-root command-root', displayActivityBusy && 'command-is-processing', eventRailCollapsed && 'event-rail-is-collapsed'), 'data-animations': 'on', + style: { '--event-rail-width': `${eventRailWidth}px` }, }, [ h('style', { key: 'style' }, CSS), - h(CommandHeader, { key: 'header', timeFilter, refreshKey, timeMenuOpen, setTimeMenuOpen, applyTimeRefresh, stats, loading, refresh, activity }), + h(CommandHeader, { + key: 'header', + title: customCommandTitle || DEFAULT_COMMAND_TITLE, + timeFilter, + refreshKey, + timeMenuOpen, + setTimeMenuOpen, + applyTimeRefresh, + stats, + loading, + refresh, + activity: displayActivity, + }), error ? h('div', { className: 'error-banner', key: 'error' }, `统计接口异常：${error}`) : null, - h('main', { className: cx('command-shell', eventRailCollapsed && 'event-rail-collapsed'), key: 'main' }, [ + h('main', { + className: cx('command-shell', eventRailCollapsed && 'event-rail-collapsed'), + key: 'main', + }, [ h('div', { className: 'command-main', key: 'workspace' }, [ - h(CommandGraph, { key: 'graph', stats, activity }), + h(CommandGraph, { key: 'graph', stats, activity: displayActivity }), h(CommandMetrics, { key: 'metrics', stats }), ]), - h(CommandEventRail, { key: 'events', activity, timeFilter, collapsed: eventRailCollapsed, onToggle: () => setEventRailCollapsed((current) => !current) }), + h(CommandEventRail, { + key: 'events', + activity: displayActivity, + timeFilter, + taskCenter: displayTaskCenter, + view: rightRailView, + onViewChange: setRightRailView, + collapsed: eventRailCollapsed, + onToggle: () => setEventRailCollapsed((current) => !current), + railWidth: eventRailWidth, + onResizeStart: startEventRailResize, + onResizeKeyDown: adjustEventRailWidth, + }), ]), ]); } @@ -3310,7 +4190,7 @@ const CSS = ` .command-root:before { content: ""; position: absolute; - inset: 68px 330px 142px 0; + inset: 68px var(--event-rail-width, 330px) 142px 0; pointer-events: none; opacity: .16; background-image: @@ -3514,7 +4394,7 @@ const CSS = ` position: relative; z-index: 2; display: grid; - grid-template-columns: minmax(0, 1fr) 330px; + grid-template-columns: minmax(0, 1fr) var(--event-rail-width, 330px); height: calc(100vh - 68px); min-height: 652px; transition: grid-template-columns .24s ease; @@ -4010,6 +4890,62 @@ const CSS = ` .command-metric .sparkline { width: 100%; height: 43px; margin-top: 6px; opacity: .75; } .command-metric .spark-grid { opacity: 0; } .command-metric .spark-line { stroke-width: 4; filter: drop-shadow(0 0 5px var(--metric-color)); } +.token-usage-metric { + display: grid; + grid-template-rows: auto auto minmax(0, 1fr); + row-gap: 7px; + padding-bottom: 8px; +} +.token-title { + display: block; + color: #b6bdbb; + font-size: 12px; +} +.token-summary-row { + display: grid; + grid-template-columns: minmax(96px, 48%) minmax(0, 1fr); + align-items: center; + column-gap: 8px; + min-width: 0; +} +.token-value { + display: block; + min-width: 0; + min-height: 30px; + overflow: hidden; + color: #f1f4f3; + font-size: 25px; + font-weight: 700; + line-height: 1; + text-overflow: ellipsis; + white-space: nowrap; +} +.token-sub { + display: block; + min-width: 0; + overflow: hidden; + color: var(--metric-color); + font-size: 9px; + line-height: 13.5px; +} +.token-chart { + display: block; + width: 100%; + height: 42px; + margin: 0; + overflow: visible; +} +.token-chart-line { + filter: drop-shadow(0 0 5px var(--metric-color)); + animation: commandMetricGlow 2.8s ease-in-out infinite; +} +.token-sub-line { + display: block; + min-width: 0; + overflow: hidden; + text-overflow: ellipsis; + white-space: nowrap; +} .command-event-rail { position: relative; z-index: 6; @@ -4026,56 +4962,86 @@ const CSS = ` background: transparent; overflow: visible; } +.event-rail-resize { + position: absolute; + z-index: 14; + top: 0; + bottom: 0; + left: -5px; + width: 10px; + cursor: col-resize; + outline: none; + touch-action: none; +} +.event-rail-resize:before { + content: ""; + position: absolute; + top: 12px; + bottom: 12px; + left: 4px; + width: 2px; + border-radius: 999px; + background: rgba(43,231,255,.16); + opacity: 0; + transition: opacity .16s ease, background .16s ease, box-shadow .16s ease; +} +.event-rail-resize:hover:before, +.event-rail-resize:focus-visible:before { + background: rgba(43,231,255,.58); + box-shadow: 0 0 12px rgba(43,231,255,.32); + opacity: 1; +} .event-rail-toggle { position: absolute; z-index: 12; top: 50%; - left: -14px; - width: 28px; - height: 44px; - border: 0; - color: #6ee8ff; - background: transparent; + left: -25px; + display: flex; + align-items: center; + justify-content: center; + width: 26px; + height: 64px; + border: 1px solid rgba(43,231,255,.34); + border-right: 0; + border-radius: 7px 0 0 7px; + color: #8df0ff; + background: linear-gradient(180deg, rgba(8,34,53,.96), rgba(5,22,38,.96)); transform: translateY(-50%); - opacity: .72; + opacity: .9; padding: 0; - font-size: 28px; - line-height: 1; - text-shadow: 0 0 9px rgba(43,231,255,.48); + font-size: 12px; + font-weight: 700; + line-height: 1.15; + text-orientation: upright; + writing-mode: vertical-rl; + box-shadow: -4px 0 14px rgba(0,0,0,.28), 0 0 12px rgba(43,231,255,.1); cursor: pointer; - transition: color .16s ease, opacity .16s ease, transform .16s ease; + transition: border-color .1s ease, color .1s ease, background .1s ease, box-shadow .1s ease, transform .1s ease; } -.event-rail-toggle:hover { +.event-rail-toggle:hover, +.event-rail-toggle:focus-visible { + border-color: rgba(70,240,255,.72); color: #fff; - opacity: 1; - transform: translate(2px, -50%); + background: rgba(7,39,62,.98); + box-shadow: -5px 0 16px rgba(0,0,0,.32), 0 0 16px rgba(43,231,255,.24); + outline: none; + transform: translate(-1px, -50%); +} +.event-rail-toggle:active { + transform: translate(0, -50%) scale(.98); } .command-event-rail.collapsed .event-rail-toggle { - left: -40px; - width: 40px; - height: 104px; - border-radius: 12px 0 0 12px; - color: #f4f7f9; - background: rgba(65, 73, 82, .96); - box-shadow: -5px 0 16px rgba(0,0,0,.22); - opacity: .96; - text-shadow: none; -} -.command-event-rail.collapsed .event-rail-toggle:hover { + left: -26px; + width: 26px; + height: 64px; + color: #dffcff; + background: linear-gradient(180deg, rgba(8,34,53,.96), rgba(5,22,38,.96)); +} +.command-event-rail.collapsed .event-rail-toggle:hover, +.command-event-rail.collapsed .event-rail-toggle:focus-visible { color: #fff; - background: rgba(78, 88, 98, .98); - transform: translate(-2px, -50%); -} -.command-event-rail.collapsed .event-rail-toggle span { - display: grid; - place-items: center; - gap: 2px; -} -.command-event-rail.collapsed .event-rail-toggle i { - font-style: normal; - font-size: 14px; - font-weight: 700; - line-height: 1.12; + background: rgba(11,52,78,.98); + transform: translate(-1px, -50%); } .event-rail-head { display: flex; @@ -4096,6 +5062,38 @@ const CSS = ` text-align: center; font-size: 11px; } +.rail-view-head { padding: 0 14px; } +.rail-view-head .rail-view-tabs { + display: flex; + flex-direction: row; + align-items: center; + gap: 4px; + min-width: 0; + border: 1px solid rgba(88,166,255,.18); + border-radius: 6px; + background: rgba(3,14,26,.72); + padding: 3px; +} +.rail-view-tabs button { + height: 28px; + min-width: 0; + border: 0; + border-radius: 4px; + color: #7e8f99; + background: transparent; + padding: 0 9px; + font: inherit; + font-size: 11px; + font-weight: 650; + white-space: nowrap; + cursor: pointer; +} +.rail-view-tabs button:hover { color: #d9f7ff; background: rgba(43,231,255,.08); } +.rail-view-tabs button.active { + color: #061725; + background: linear-gradient(135deg, #35d4ff, #40e1bd); + box-shadow: 0 0 13px rgba(43,231,255,.14); +} .event-update-banner { margin: 10px 14px 4px; padding: 10px 12px; @@ -4107,6 +5105,322 @@ const CSS = ` } .event-update-banner:before { content: "ⓘ"; margin-right: 7px; color: #6ba4fb; } .event-update-banner.warn { border-color: rgba(255,174,52,.42); color: #f1c67d; background: rgba(139,88,22,.2); } +.task-center-panel { + min-height: 0; + padding: 8px 14px 18px; + overflow: auto; + scrollbar-width: thin; + scrollbar-color: #3a403e transparent; +} +.task-center-inline-warn { + margin-bottom: 10px; + border: 1px solid rgba(255,174,52,.42); + border-radius: 6px; + color: #f1c67d; + background: rgba(139,88,22,.2); + padding: 9px 11px; + font-size: 11px; +} +.task-center-summary { + display: grid; + grid-template-columns: repeat(2, minmax(0, 1fr)); + gap: 8px; + margin-bottom: 12px; +} +.task-center-summary > div { + min-width: 0; + min-height: 64px; + display: flex; + flex-direction: column; + justify-content: center; + gap: 5px; + border: 1px solid rgba(43,231,255,.14); + border-radius: 6px; + background: rgba(10,33,51,.46); + padding: 10px 11px; +} +.task-center-summary > div.active { + border-color: rgba(46,230,166,.32); + background: rgba(20,75,57,.32); + animation: commandEventActive 1.8s ease-in-out infinite; +} +.task-center-summary span { + overflow: hidden; + color: #7f8e99; + font-size: 10px; + text-overflow: ellipsis; + white-space: nowrap; +} +.task-center-summary b { + color: #e9edeb; + font-size: 18px; + font-variant-numeric: tabular-nums; +} +.task-center-summary small { + overflow: hidden; + color: #49e4c8; + font-size: 10px; + font-weight: 800; + text-overflow: ellipsis; + white-space: nowrap; +} +.task-center-section { margin-top: 12px; } +.task-center-section-title { + display: flex; + align-items: center; + justify-content: space-between; + min-height: 28px; + border-bottom: 1px solid rgba(43,231,255,.12); + margin-bottom: 3px; +} +.task-center-section-toggle { + display: inline-flex; + align-items: center; + min-width: 0; + gap: 6px; + border: 0; + color: #dce1df; + background: transparent; + padding: 0; + font: inherit; + cursor: pointer; +} +.task-center-section-toggle:hover, +.task-center-section-toggle:focus-visible { color: #66e6d6; outline: none; } +.task-center-section-toggle i { + display: grid; + place-items: center; + width: 14px; + height: 14px; + color: #53e9c4; + font-style: normal; + font-size: 15px; + line-height: 1; +} +.task-center-section-title strong { color: currentColor; font-size: 12px; } +.task-center-section-title span { color: #66716d; font-size: 10px; } +.task-center-section-list { display: grid; gap: 0; } +.task-center-item { + position: relative; + min-height: 118px; + display: flex; + flex-direction: column; + gap: 8px; + border-bottom: 1px solid rgba(255,255,255,.09); + padding: 13px 2px 13px 12px; +} +.task-center-item:before { + content: ""; + position: absolute; + top: 16px; + bottom: 16px; + left: 1px; + width: 1px; + background: rgba(105,129,143,.25); +} +.task-center-item:after { + content: ""; + position: absolute; + top: 20px; + left: -2px; + width: 7px; + height: 7px; + border-radius: 50%; + background: #2be7ff; + box-shadow: 0 0 9px rgba(43,231,255,.54); +} +.task-center-item.active { + background: linear-gradient(90deg, rgba(43,231,255,.06), transparent 64%); + animation: commandEventActive 1.8s ease-in-out infinite; +} +.task-center-item.clickable { + cursor: pointer; +} +.task-center-item.clickable:hover { + background: linear-gradient(90deg, rgba(43,231,255,.085), rgba(43,231,255,.025) 58%, transparent); +} +.task-center-item.clickable:focus-visible { + outline: 1px solid rgba(85,232,255,.68); + outline-offset: -2px; +} +.task-center-item.active:after { + background: #46e4af; + box-shadow: 0 0 10px rgba(70,228,175,.78); +} +.task-center-item-head, +.task-center-stats { + display: flex; + align-items: center; + min-width: 0; +} +.task-center-item-head { gap: 8px; } +.task-center-item-head strong { + flex: 1 1 auto; + min-width: 0; + overflow: hidden; + color: #dce1df; + font-size: 12px; + font-weight: 650; + text-overflow: ellipsis; + white-space: nowrap; +} +.task-center-status { + flex: 0 0 auto; + max-width: 74px; + overflow: hidden; + border-radius: 4px; + color: #83aef1; + background: rgba(52,86,143,.5); + padding: 4px 7px; + font-size: 10px; + text-overflow: ellipsis; + white-space: nowrap; +} +.task-center-status.active, +.task-center-status.running { color: #57e1b5; background: rgba(23,111,83,.48); } +.task-center-status.failed, +.task-center-status.error, +.task-center-status.timeout { color: #ff9a76; background: rgba(133,57,33,.42); } +.task-center-status.disabled, +.task-center-status.stopped, +.task-center-status.stale { color: #9aa7a3; background: rgba(73,86,95,.42); } +.task-center-item-sub { + min-width: 0; + overflow: hidden; + color: #7b8581; + font-size: 10px; + text-overflow: ellipsis; + white-space: nowrap; +} +.task-center-alert { + display: grid; + grid-template-columns: auto minmax(0, 1fr); + align-items: center; + gap: 7px; + min-width: 0; + color: #74827d; + font-size: 10px; +} +.task-center-alert span { + white-space: nowrap; +} +.task-center-alert b { + min-width: 0; + overflow: hidden; + color: #dce1df; + font-size: 11px; + font-weight: 700; + text-overflow: ellipsis; + white-space: nowrap; +} +.task-center-alert.has-alert b { + color: #f0f6f4; +} +.task-center-hash { + position: relative; + display: grid; + grid-template-columns: auto minmax(0, 1fr); + align-items: center; + gap: 6px; + min-width: 0; + overflow: hidden; + border: 1px solid rgba(43,231,255,.16); + border-radius: 5px; + background: rgba(4,22,36,.64); + padding: 5px 6px; +} +.task-center-hash:after { + content: ""; + position: absolute; + top: 0; + bottom: 0; + left: -38%; + width: 32%; + background: linear-gradient(90deg, transparent, rgba(43,231,255,.18), transparent); + animation: commandBannerSweep 3.6s ease-in-out infinite; +} +.task-center-hash span { + color: #66716d; + font-size: 9px; + line-height: 1.2; + white-space: nowrap; +} +.task-center-hash code { + display: block; + min-width: 0; + overflow: hidden; + color: #77e8ff; + font: 700 10px/1.2 ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, monospace; + text-overflow: ellipsis; + white-space: nowrap; + user-select: text; +} +.task-center-hash .task-center-jump { + color: #7b8581; + user-select: none; +} +.task-center-hash .task-center-jump.enabled { + color: #57e1b5; +} +.task-center-stats { + display: grid; + grid-template-columns: repeat(4, minmax(0, 1fr)); + gap: 6px; + color: #7f8e99; + font-size: 10px; +} +.task-center-stats.workflow-stats { + grid-template-columns: .7fr .7fr 1.05fr .85fr; + gap: 6px; +} +.task-center-stats span { + min-width: 0; + overflow: hidden; + text-overflow: ellipsis; + white-space: nowrap; +} +.task-center-stats b { + color: #4dcfa7; + font-weight: 700; + font-variant-numeric: tabular-nums; +} +.task-center-rate { + position: relative; + height: 3px; + overflow: hidden; + border-radius: 3px; + background: rgba(127,111,255,.14); + box-shadow: inset 0 0 0 1px rgba(155,140,255,.11); +} +.task-center-rate i { + position: absolute; + inset: 0; + border-radius: inherit; + background: linear-gradient(90deg, #38d39f 0%, #55e8ff 100%); + box-shadow: 0 0 8px rgba(43,231,255,.55); + transform: scaleX(var(--task-center-rate, 0)); + transform-origin: left center; + transition: transform .5s ease; +} +.task-center-rate.progress-rate { + background: rgba(43,231,255,.12); +} +.task-center-expand { + width: 100%; + height: 30px; + margin-top: 8px; + border: 1px solid rgba(43,231,255,.2); + border-radius: 5px; + color: #73e7ff; + background: rgba(10,31,51,.52); + font: inherit; + font-size: 11px; + cursor: pointer; +} +.task-center-expand:hover { + border-color: rgba(43,231,255,.46); + background: rgba(20,57,82,.62); +} .event-rail-list { min-height: 0; padding: 4px 16px 18px 20px; @@ -4149,6 +5463,16 @@ const CSS = ` .event-rail-item.state-processing { background: linear-gradient(90deg, rgba(35,214,157,.07), transparent 72%); } +.event-rail-item.clickable { + cursor: pointer; +} +.event-rail-item.clickable:hover { + background: linear-gradient(90deg, rgba(155,140,255,.12), rgba(43,231,255,.04) 60%, transparent); +} +.event-rail-item.clickable:focus-visible { + outline: 1px solid rgba(155,140,255,.72); + outline-offset: -2px; +} .event-rail-meta { display: flex; align-items: center; width: 100%; gap: 7px; } .event-rail-meta time { margin-left: auto; color: #646d69; font-size: 10px; } .event-queue-kind, @@ -4237,8 +5561,8 @@ const CSS = ` @keyframes laneResult { to { opacity: 1; } } @media (max-width: 1360px) { .command-header { grid-template-columns: minmax(330px, 1fr) auto minmax(500px, 1fr); gap: 12px; padding: 0 14px; } - .command-shell { grid-template-columns: minmax(0, 1fr) 292px; } - .command-root:before { right: 292px; } + .command-shell { grid-template-columns: minmax(0, 1fr) var(--event-rail-width, 292px); } + .command-root:before { right: var(--event-rail-width, 292px); } .command-core { width: 288px; height: 288px; min-width: 288px; min-height: 288px; } .command-time-trigger { max-width: 320px; } .command-clock { min-width: 92px; } @@ -4249,7 +5573,7 @@ const CSS = ` .command-root { padding: 0; } .command-header { grid-template-columns: 330px 1fr; } .command-live { display: none; } - .command-shell { grid-template-columns: minmax(850px, 1fr) 280px; } + .command-shell { grid-template-columns: minmax(850px, 1fr) var(--event-rail-width, 280px); } .command-core { width: 270px; height: 270px; min-width: 270px; min-height: 270px; } .agent-badge { min-width: 96px; padding: 7px 8px; } .command-lanes { right: 12%; left: 12%; } diff --git a/.flocks/flockshub/plugins/webuis/soc_ui/soc_overview/api/handlers.py b/.flocks/flockshub/plugins/webuis/soc_ui/soc_overview/api/handlers.py index 6cc44b672..af5a17696 100644 --- a/.flocks/flockshub/plugins/webuis/soc_ui/soc_overview/api/handlers.py +++ b/.flocks/flockshub/plugins/webuis/soc_ui/soc_overview/api/handlers.py @@ -133,11 +133,51 @@ def _get_workflow_call_count(workflow_name: str, date: str = None) -> int: "attack_success": "攻击成功", "attack": "攻击行为", "attack_failed": "攻击失败", - "benign": "良性", + "non_attack": "非攻击", "unknown": "待确认", } +def _triage_outcome_key(record): + verdict = _triage_attack_verdict_value(record) + result = _triage_attack_success_value(record) + if verdict == "attack": + if result == "success": + return "attack_success" + if result == "failed": + return "attack_failed" + return "attack" + if verdict == "non_attack": + return "non_attack" + return "unknown" + + +def _triage_attack_verdict_value(record): + verdict = _norm(record.get("triage_attack_verdict") or "unknown") + if verdict in {"attack", "non_attack", "unknown"}: + return verdict + if verdict in {"attack_success", "attack_failed"}: + return "attack" + if verdict == "benign": + return "non_attack" + return "unknown" + + +def _triage_attack_success_value(record): + raw_verdict = _norm(record.get("triage_attack_verdict") or "unknown") + verdict = _triage_attack_verdict_value(record) + result = _norm(record.get("triage_attack_success") or "unknown") + if verdict == "attack" and result in {"success", "failed", "unknown"}: + return result + if verdict != "attack": + return "unknown" + if raw_verdict == "attack_success" or record.get("triage_attack_success") is True: + return "success" + if raw_verdict == "attack_failed": + return "failed" + return "unknown" + + def get_stats(ctx, request): start_time, end_time = _normalize_time_range( request.query_params.get("startTime"), @@ -226,7 +266,7 @@ def get_stats(ctx, request): {"key": "attack_success", "label": "攻击成功", "value": triage["attackSuccess"], "color": "#ff4d6d"}, {"key": "attack", "label": "攻击行为", "value": triage["attack"], "color": "#ffb020"}, {"key": "attack_failed", "label": "攻击失败", "value": triage["attackFailed"], "color": "#2ee6a6"}, - {"key": "benign", "label": "良性", "value": triage["benign"], "color": "#58a6ff"}, + {"key": "non_attack", "label": "非攻击", "value": triage["benign"], "color": "#58a6ff"}, {"key": "unknown", "label": "未知", "value": triage["unknown"], "color": "#9b8cff"}, ], "topThreats": _counter_items(triage["threatCounter"] or denoise["threatCounter"], 14), @@ -676,7 +716,6 @@ def _read_triage(paths): fallback_cache = 0 fallback_failed = 0 fallback_followers = 0 - extra_success = 0 series_total = [] series_attack = [] @@ -701,20 +740,17 @@ def _read_triage(paths): total_records += 1 file_total += 1 - verdict = _norm(obj.get("attack_verdict") or "unknown") - if verdict not in {"attack_success", "attack", "attack_failed", "benign", "unknown"}: - verdict = "unknown" - verdict_counter[verdict] += 1 - if obj.get("attack_success") is True and verdict != "attack_success": - extra_success += 1 - if verdict in {"attack_success", "attack", "attack_failed"}: + verdict = _triage_attack_verdict_value(obj) + outcome = _triage_outcome_key(obj) + verdict_counter[outcome] += 1 + if verdict == "attack": file_attack += 1 source = _norm(obj.get("_source_type") or obj.get("source_type") or obj.get("device_type")) source_counter[source] += 1 threat_counter[_norm(obj.get("threat_name"))] += 1 risk_counter[_norm(obj.get("risk_level"))] += 1 - _update_profile_counters(obj, profile_counters) + _update_profile_counters(obj, profile_counters, triage_result=True) event_start, event_end = _merge_record_time(event_start, event_end, obj) triage_source = _norm(obj.get("triage_source")) triage_status = _norm(obj.get("triage_status")) @@ -741,11 +777,11 @@ def _read_triage(paths): triage_failed = fallback_failed followers_reused = fallback_followers - attack_success = verdict_counter["attack_success"] + extra_success + attack_success = verdict_counter["attack_success"] attack = verdict_counter["attack"] attack_failed = verdict_counter["attack_failed"] attack_total = attack_success + attack + attack_failed - benign = verdict_counter["benign"] + benign = verdict_counter["non_attack"] unknown = verdict_counter["unknown"] series_total = _expand_series(series_total, total_records, seed=23) @@ -795,10 +831,15 @@ def _new_profile_counters(): } -def _update_profile_counters(obj, counters): +def _update_profile_counters(obj, counters, *, triage_result=False): counters["phaseCounter"][_norm(obj.get("threat_phase") or obj.get("attack_phase") or obj.get("kill_chain_phase"))] += 1 counters["directionCounter"][_norm(obj.get("direction") or obj.get("traffic_direction"))] += 1 - counters["resultCounter"][_norm(obj.get("threat_result") or obj.get("attack_verdict"))] += 1 + result = ( + _triage_attack_success_value(obj) + if triage_result + else obj.get("threat_result") or obj.get("attack_verdict") + ) + counters["resultCounter"][_norm(result)] += 1 counters["protocolCounter"][_norm(obj.get("net_type") or obj.get("net_app_proto") or obj.get("protocol"))] += 1 counters["severityCounter"][_norm(obj.get("threat_severity"))] += 1 counters["responseCounter"][_norm(obj.get("rsp_status_code") or obj.get("status_code"))] += 1 diff --git a/.flocks/flockshub/plugins/webuis/soc_ui/soc_overview/src/index.tsx b/.flocks/flockshub/plugins/webuis/soc_ui/soc_overview/src/index.tsx index cbf4d3600..2291ef2aa 100644 --- a/.flocks/flockshub/plugins/webuis/soc_ui/soc_overview/src/index.tsx +++ b/.flocks/flockshub/plugins/webuis/soc_ui/soc_overview/src/index.tsx @@ -65,8 +65,10 @@ interface Stats { attackSuccess?: number; attack?: number; attackFailed?: number; + benign?: number; unknown?: number; }; + verdicts?: CounterItem[]; topThreats?: CounterItem[]; fieldStats?: FieldStats; } @@ -101,11 +103,20 @@ const EMPTY_STATS: Required = { eventRange: { label: '', start: '', end: '' }, generatedAt: '', denoise: { totalRaw: 0, totalUnique: 0, duplicates: 0, duplicateRate: 0 }, - triage: { totalRecords: 0, attackSuccess: 0, attack: 0, attackFailed: 0, unknown: 0 }, + triage: { totalRecords: 0, attackSuccess: 0, attack: 0, attackFailed: 0, benign: 0, unknown: 0 }, + verdicts: [], topThreats: [], fieldStats: EMPTY_FIELD_STATS, }; +const VERDICT_CLASSES = [ + { key: 'attack_success', label: '攻击成功' }, + { key: 'attack', label: '攻击' }, + { key: 'attack_failed', label: '攻击失败' }, + { key: 'non_attack', label: '非攻击' }, + { key: 'unknown', label: '待人工审核' }, +] as const; + const PHASE_LABELS: Record = { exploit: '漏洞利用', recon: '侦察探测', @@ -182,6 +193,7 @@ function mergeStats(raw: Stats | undefined): Required { eventRange: { ...EMPTY_STATS.eventRange, ...(value.eventRange || {}) }, denoise: { ...EMPTY_STATS.denoise, ...(value.denoise || {}) }, triage: { ...EMPTY_STATS.triage, ...(value.triage || {}) }, + verdicts: list(value.verdicts), topThreats: list(value.topThreats), fieldStats: mergeFieldStats(value.fieldStats), }; @@ -274,10 +286,6 @@ function truncate(value: string, max = 34) { return value.length > max ? `${value.slice(0, max)}...` : value; } -function count(items: CounterItem[], key: string) { - return items.find((item) => item.key === key || item.label === key)?.value || 0; -} - export default function SocOverviewPage() { const [stats, setStats] = useState(EMPTY_STATS); const [loading, setLoading] = useState(true); @@ -349,10 +357,13 @@ export default function SocOverviewPage() { { label: '目标端口', value: stats.fieldStats.uniqueDestinationPorts, hint: '字段 dport' }, ], [stats]); - const resultTotal = Math.max(stats.denoise.totalUnique, 1); - const success = stats.triage.attackSuccess || count(stats.fieldStats.threatResults, 'success'); - const failed = stats.triage.attackFailed || count(stats.fieldStats.threatResults, 'failed'); - const unknown = Math.max(0, stats.denoise.totalUnique - success - failed) || count(stats.fieldStats.threatResults, 'unknown'); + const verdicts = VERDICT_CLASSES.map((verdictClass) => { + const backendVerdict = stats.verdicts.find((item) => item.key === verdictClass.key); + const value = Number(backendVerdict?.value || 0); + return { ...verdictClass, value: Number.isFinite(value) && value > 0 ? value : 0 }; + }); + const resultTotal = verdicts.reduce((total, item) => total + item.value, 0); + const resultDenominator = Math.max(resultTotal, 1); return (
@@ -391,19 +402,26 @@ export default function SocOverviewPage() {

告警研判结果

-

按 threat_result 与研判结论聚合。

+

按模型研判的是否攻击与攻击结果字段聚合。

{formatNumber(resultTotal)} 条有效告警
- - - + {verdicts.map((item) => ( + + ))}
- 攻击成功 {formatNumber(success)} - 待确认 {formatNumber(unknown)} - 攻击失败 {formatNumber(failed)} + {verdicts.map((item) => ( + + + {item.label} {formatNumber(item.value)} + + ))}
@@ -777,15 +795,19 @@ const CSS = ` .section-head p { margin: 6px 0 0; color: #667085; font-size: 13px; } .section-head > b { color: #475467; font-size: 13px; } .result-bar { display: flex; height: 14px; overflow: hidden; margin-top: 18px; border-radius: 999px; background: #edf2f7; } -.result-bar .success { background: #ef4444; } -.result-bar .unknown { background: #f59e0b; } -.result-bar .failed { background: #cbd5e1; } +.result-bar .verdict-attack_success { background: #ff4d6d; } +.result-bar .verdict-attack { background: #ffb020; } +.result-bar .verdict-attack_failed { background: #2ee6a6; } +.result-bar .verdict-non_attack { background: #58a6ff; } +.result-bar .verdict-unknown { background: #9b8cff; } .result-legend { display: flex; flex-wrap: wrap; gap: 18px; margin-top: 12px; color: #667085; font-size: 13px; } .result-legend span { display: inline-flex; align-items: center; gap: 7px; } .result-legend i { width: 9px; height: 9px; border-radius: 50%; } -.result-legend .success { background: #ef4444; } -.result-legend .unknown { background: #f59e0b; } -.result-legend .failed { background: #cbd5e1; } +.result-legend .verdict-attack_success { background: #ff4d6d; } +.result-legend .verdict-attack { background: #ffb020; } +.result-legend .verdict-attack_failed { background: #2ee6a6; } +.result-legend .verdict-non_attack { background: #58a6ff; } +.result-legend .verdict-unknown { background: #9b8cff; } .panel-grid { display: grid; grid-template-columns: repeat(2, minmax(0, 1fr)); gap: 12px; margin-top: 14px; } .panel-card { min-height: 288px; padding: 0 16px 16px; } .panel-head { height: 48px; display: flex; align-items: center; border-bottom: 1px solid #edf2f7; margin-bottom: 14px; } diff --git a/.flocks/flockshub/plugins/webuis/soc_ui/workspace.json b/.flocks/flockshub/plugins/webuis/soc_ui/workspace.json index b17ffc05a..3bf045c38 100644 --- a/.flocks/flockshub/plugins/webuis/soc_ui/workspace.json +++ b/.flocks/flockshub/plugins/webuis/soc_ui/workspace.json @@ -1,6 +1,6 @@ { "id": "soc_ui", - "version": "1.1.4", + "version": "1.1.5", "title": "SOC 工作区", "titleEn": "SOC Workspace", "icon": "ShieldCheck", diff --git a/.flocks/flockshub/plugins/workflows/stream_alert_denoise/workflow.json b/.flocks/flockshub/plugins/workflows/stream_alert_denoise/workflow.json index 12f97ed24..ea1b4193e 100644 --- a/.flocks/flockshub/plugins/workflows/stream_alert_denoise/workflow.json +++ b/.flocks/flockshub/plugins/workflows/stream_alert_denoise/workflow.json @@ -15,13 +15,13 @@ "id": "normalize", "type": "python", "description": "Normalize TDP and Skyeye alerts into a unified schema. Per-alert type detection via field signatures; falls back to batch_hint. Carries _syslog_meta and _source_type to downstream nodes.", - "code": "\nimport uuid\n\nHTTP_METHODS = ['GET', 'POST', 'PUT', 'DELETE', 'HEAD', 'OPTIONS', 'PATCH', 'TRACE']\n\nTDP_FIELD_MAP = {\n 'customer_uuid': 'customer_uuid',\n 'device_id': 'device_id',\n 'id': 'id',\n 'time': 'time',\n 'direction': 'direction',\n 'sip': 'net_real_src_ip',\n 'dip': 'net_dest_ip',\n 'sport': 'net_src_port',\n 'dport': 'net_dest_port',\n 'net_type': 'net_type',\n 'net_app_proto': 'net_app_proto',\n 'req_http_url': 'net_http_url',\n 'req_user_agent':'net_http_reqs_user_agent',\n 'req_host': 'net_http_reqs_host',\n 'req_line': 'net_http_reqs_line',\n 'req_header': 'net_http_reqs_header',\n 'req_body': 'net_http_reqs_body',\n 'req_cookie': 'net_http_reqs_cookie',\n 'req_body_len': 'net_http_reqs_content_length',\n 'rsp_status_code': 'net_http_status',\n 'rsp_line': 'net_http_resp_line',\n 'rsp_header': 'net_http_resp_header',\n 'rsp_body': 'net_http_resp_body',\n 'rsp_body_len': 'net_http_resp_content_length',\n 'net_bytes_toclient': 'net_bytes_toclient',\n 'net_bytes_toserver': 'net_bytes_toserver',\n 'threat_rule_id': 'threat_suuid',\n 'threat_name': 'threat_name',\n 'threat_msg': 'threat_msg',\n 'threat_ioc': 'threat_ioc',\n 'threat_level': 'threat_level',\n 'threat_severity': 'threat_severity',\n 'threat_phase': 'threat_phase',\n 'threat_type': 'threat_type',\n 'threat_result': 'threat_result',\n 'threat_confidence': 'threat_confidence',\n 'connection_established': 'established',\n 'asset_group_name': 'dest_assets_group_name',\n 'asset_name': 'dest_assets_latestName',\n}\n\nSKYEYE_FIELD_MAP = {\n 'id': 'none',\n 'time': 'time',\n 'direction': 'none',\n 'sip': 'sip',\n 'dip': 'dip',\n 'sport': 'sport',\n 'dport': 'dport',\n 'net_type': 'none',\n 'net_app_proto': 'none',\n 'req_http_url': 'uri',\n 'req_user_agent':'agent',\n 'req_host': 'host',\n 'req_line': 'none',\n 'req_header': 'req_header',\n 'req_body': 'req_body',\n 'req_cookie': 'none',\n 'req_body_len': 'none',\n 'rsp_status_code': 'rsp_status',\n 'rsp_line': 'none',\n 'rsp_header': 'rsp_header',\n 'rsp_body': 'rsp_body',\n 'rsp_body_len': 'rsp_body_len',\n 'threat_rule_id': 'rule_id',\n 'threat_name': 'vuln_name',\n 'threat_msg': 'vuln_desc',\n 'threat_ioc': 'none',\n 'threat_level': 'none',\n 'threat_severity': 'severity',\n 'threat_phase': 'none',\n 'threat_type': 'vuln_type',\n 'threat_tactic_id': 'attck_tactic',\n 'threat_technique_id': 'attck_tech',\n 'threat_result': 'attack_result',\n 'threat_confidence': 'confidence',\n 'connection_established': 'established',\n 'real_attack': 'attack_flag',\n}\n\ndef flatten_dict(d, prefix=''):\n res = {}\n for k, v in d.items():\n if isinstance(v, dict):\n res.update(flatten_dict(v, f'{prefix}{k}_'))\n else:\n res[f'{prefix}{k}'] = v\n return res\n\ndef make_uuid(norm):\n return str(uuid.uuid3(uuid.NAMESPACE_DNS, ''.join(str(v) for v in norm.values())))\n\ndef detect_alert_type(alert, batch_hint):\n if isinstance(alert.get('net'), dict):\n return 'tdp'\n if any(k in alert for k in ('behave_uuid', 'flow_id')):\n return 'tdp'\n if any(k in alert for k in ('net_real_src_ip', 'net_http_url', 'threat_suuid')):\n return 'tdp'\n if any(k in alert for k in ('uri', 'vuln_name', 'attack_result', 'attack_flag')):\n return 'skyeye'\n return batch_hint\n\ndef normalize_single(alert, source_type):\n flat = flatten_dict(alert)\n field_map = TDP_FIELD_MAP if source_type == 'tdp' else SKYEYE_FIELD_MAP\n norm = {}\n for std_key, raw_key in field_map.items():\n norm[std_key] = flat.get(raw_key, 'none') if raw_key != 'none' else 'none'\n if norm.get('id') in ('none', None, ''):\n norm['id'] = make_uuid(norm)\n if norm.get('net_type') in ('none', None, ''):\n method = flat.get('method', 'none')\n norm['net_type'] = 'http' if method in HTTP_METHODS else ('none' if method == 'none' else 'other')\n norm['_source_type'] = source_type\n # Carry syslog metadata if present\n if '_syslog_meta' in alert:\n norm['_syslog_meta'] = alert['_syslog_meta']\n return norm\n\nraw_alerts = inputs.get('raw_alerts', [])\nstats = dict(inputs.get('stats', {}))\nbatch_hint = str(inputs.get('source_log_type', 'tdp') or 'tdp').lower()\nif batch_hint not in ('tdp', 'skyeye'):\n batch_hint = 'tdp'\n\ntype_counts = {'tdp': 0, 'skyeye': 0}\nnormalized = []\nfor alert in raw_alerts:\n src_type = detect_alert_type(alert, batch_hint)\n type_counts[src_type] = type_counts.get(src_type, 0) + 1\n normalized.append(normalize_single(alert, src_type))\n\nstats['normalized_count'] = len(normalized)\nstats['normalize_type_counts'] = type_counts\nprint(f'[normalize] {len(raw_alerts)} alerts -> {len(normalized)} normalized '\n f'(tdp={type_counts.get(\"tdp\",0)}, skyeye={type_counts.get(\"skyeye\",0)}, '\n f'batch_hint={batch_hint!r})')\n\noutputs['normalized_alerts'] = normalized\noutputs['stats'] = stats\nfor k in ['input_mode', 'source_log_type', 'filter_enabled', 'dedup_enabled',\n 'dedup_threshold', 'strict_fields', 'lsh_fields', 'max_field_len', 'max_dedup_keys']:\n outputs[k] = inputs.get(k)\n" + "code": "\nimport uuid\n\nHTTP_METHODS = ['GET', 'POST', 'PUT', 'DELETE', 'HEAD', 'OPTIONS', 'PATCH', 'TRACE']\n\nTDP_FIELD_MAP = {\n 'customer_uuid': 'customer_uuid',\n 'device_id': 'device_id',\n 'id': 'id',\n 'time': 'time',\n 'direction': 'direction',\n 'sip': 'net_real_src_ip',\n 'dip': 'net_dest_ip',\n 'sport': 'net_src_port',\n 'dport': 'net_dest_port',\n 'net_type': 'net_type',\n 'net_app_proto': 'net_app_proto',\n 'req_http_url': 'net_http_url',\n 'req_user_agent':'net_http_reqs_user_agent',\n 'req_host': 'net_http_reqs_host',\n 'req_line': 'net_http_reqs_line',\n 'req_header': 'net_http_reqs_header',\n 'req_body': 'net_http_reqs_body',\n 'req_cookie': 'net_http_reqs_cookie',\n 'req_body_len': 'net_http_reqs_content_length',\n 'rsp_status_code': 'net_http_status',\n 'rsp_line': 'net_http_resp_line',\n 'rsp_header': 'net_http_resp_header',\n 'rsp_body': 'net_http_resp_body',\n 'rsp_body_len': 'net_http_resp_content_length',\n 'net_bytes_toclient': 'net_bytes_toclient',\n 'net_bytes_toserver': 'net_bytes_toserver',\n 'threat_rule_id': 'threat_suuid',\n 'threat_name': 'threat_name',\n 'threat_msg': 'threat_msg',\n 'threat_ioc': 'threat_ioc',\n 'threat_level': 'threat_level',\n 'threat_severity': 'threat_severity',\n 'threat_phase': 'threat_phase',\n 'threat_type': 'threat_type',\n 'threat_result': 'threat_result',\n 'threat_confidence': 'threat_confidence',\n 'connection_established': 'established',\n 'asset_group_name': 'dest_assets_group_name',\n 'asset_name': 'dest_assets_latestName',\n}\n\nSKYEYE_FIELD_MAP = {\n 'id': 'none',\n 'time': 'time',\n 'direction': 'none',\n 'sip': 'sip',\n 'dip': 'dip',\n 'sport': 'sport',\n 'dport': 'dport',\n 'net_type': 'none',\n 'net_app_proto': 'none',\n 'req_http_url': 'uri',\n 'req_user_agent':'agent',\n 'req_host': 'host',\n 'req_line': 'none',\n 'req_header': 'req_header',\n 'req_body': 'req_body',\n 'req_cookie': 'none',\n 'req_body_len': 'none',\n 'rsp_status_code': 'rsp_status',\n 'rsp_line': 'none',\n 'rsp_header': 'rsp_header',\n 'rsp_body': 'rsp_body',\n 'rsp_body_len': 'rsp_body_len',\n 'threat_rule_id': 'rule_id',\n 'threat_name': 'vuln_name',\n 'threat_msg': 'vuln_desc',\n 'threat_ioc': 'none',\n 'threat_level': 'none',\n 'threat_severity': 'severity',\n 'threat_phase': 'none',\n 'threat_type': 'vuln_type',\n 'threat_tactic_id': 'attck_tactic',\n 'threat_technique_id': 'attck_tech',\n 'threat_result': 'attack_result',\n 'threat_confidence': 'confidence',\n 'connection_established': 'established',\n 'real_attack': 'attack_flag',\n}\n\ndef flatten_dict(d, prefix=''):\n res = {}\n for k, v in d.items():\n if isinstance(v, dict):\n res.update(flatten_dict(v, f'{prefix}{k}_'))\n else:\n res[f'{prefix}{k}'] = v\n return res\n\ndef make_uuid(norm):\n return str(uuid.uuid3(uuid.NAMESPACE_DNS, ''.join(str(v) for v in norm.values())))\n\ndef detect_alert_type(alert, batch_hint):\n if isinstance(alert.get('net'), dict):\n return 'tdp'\n if any(k in alert for k in ('behave_uuid', 'flow_id')):\n return 'tdp'\n if any(k in alert for k in ('net_real_src_ip', 'net_http_url', 'threat_suuid')):\n return 'tdp'\n if any(k in alert for k in ('uri', 'vuln_name', 'attack_result', 'attack_flag')):\n return 'skyeye'\n return batch_hint\n\ndef normalize_single(alert, source_type):\n flat = flatten_dict(alert)\n field_map = TDP_FIELD_MAP if source_type == 'tdp' else SKYEYE_FIELD_MAP\n norm = {}\n for std_key, raw_key in field_map.items():\n norm[std_key] = flat.get(raw_key, 'none') if raw_key != 'none' else 'none'\n if norm.get('id') in ('none', None, ''):\n norm['id'] = make_uuid(norm)\n if norm.get('net_type') in ('none', None, ''):\n method = flat.get('method', 'none')\n norm['net_type'] = 'http' if method in HTTP_METHODS else ('none' if method == 'none' else 'other')\n norm['_source_type'] = source_type\n # Carry syslog metadata if present\n if '_syslog_meta' in alert:\n norm['_syslog_meta'] = alert['_syslog_meta']\n return norm\n\nraw_alerts = inputs.get('raw_alerts', [])\nstats = dict(inputs.get('stats', {}))\nbatch_hint = str(inputs.get('source_log_type', 'tdp') or 'tdp').lower()\nif batch_hint not in ('tdp', 'skyeye'):\n batch_hint = 'tdp'\n\ntype_counts = {'tdp': 0, 'skyeye': 0}\nnormalized = []\nfor alert in raw_alerts:\n src_type = detect_alert_type(alert, batch_hint)\n type_counts[src_type] = type_counts.get(src_type, 0) + 1\n normalized.append(normalize_single(alert, src_type))\n\nstats['normalized_count'] = len(normalized)\nstats['normalize_type_counts'] = type_counts\nprint(f'[normalize] {len(raw_alerts)} alerts -> {len(normalized)} normalized '\n f'(tdp={type_counts.get(\"tdp\",0)}, skyeye={type_counts.get(\"skyeye\",0)}, '\n f'batch_hint={batch_hint!r})')\n\noutputs['normalized_alerts'] = normalized\noutputs['stats'] = stats\noutputs['input_mode'] = inputs.get('input_mode')\noutputs['source_log_type'] = inputs.get('source_log_type')\noutputs['filter_enabled'] = inputs.get('filter_enabled')\noutputs['dedup_enabled'] = inputs.get('dedup_enabled')\noutputs['dedup_threshold'] = inputs.get('dedup_threshold')\noutputs['strict_fields'] = inputs.get('strict_fields')\noutputs['lsh_fields'] = inputs.get('lsh_fields')\noutputs['max_field_len'] = inputs.get('max_field_len')\noutputs['max_dedup_keys'] = inputs.get('max_dedup_keys')\n" }, { "id": "filter_logs", "type": "python", "description": "Filter: classify into 9 process_types, keep non-scan HTTP alerts (direction in/out/lateral). Adds _process_type and _threat_type fields. When filter_enabled=False, all alerts pass through.", - "code": "\nnormalized_alerts = inputs.get('normalized_alerts', [])\nfilter_enabled = inputs.get('filter_enabled', True)\nbatch_hint = str(inputs.get('source_log_type', 'tdp') or 'tdp').lower()\nif batch_hint not in ('tdp', 'skyeye'):\n batch_hint = 'tdp'\nstats = dict(inputs.get('stats', {}))\n\ndef is_scan_alert(threat_name):\n tnl = str(threat_name or '').lower()\n return ('扫描' in tnl) and ('webshell' not in tnl)\n\ndef get_threat_type(alert):\n src = alert.get('_source_type') or batch_hint\n if src == 'skyeye':\n return str(alert.get('threat_type', 'general') or 'general')\n return str(alert.get('threat_name', 'general') or 'general')\n\ndef is_http(alert):\n for field in ('application_layer_protocol', 'net_type', 'net_app_proto'):\n val = str(alert.get(field, '') or '').lower()\n if val and val != 'none' and 'http' in val:\n return True\n return False\n\ndef get_process_type(alert):\n src = alert.get('_source_type') or batch_hint\n threat_name = alert.get('threat_name', '')\n direction = str(alert.get('direction', '') or '').lower()\n scan = is_scan_alert(threat_name)\n http = is_http(alert)\n if src == 'skyeye':\n return 'alert_scan_direction_in' if scan else 'alert_not_scan_http_direction_in'\n if scan:\n return f'alert_scan_direction_{direction}' if direction in ('in', 'out', 'lateral') else 'alert_scan_direction_in'\n if http:\n return f'alert_not_scan_http_direction_{direction}' if direction in ('in', 'out', 'lateral') else 'alert_not_scan_http_direction_in'\n return f'alert_not_scan_not_http_direction_{direction}' if direction in ('in', 'out', 'lateral') else 'alert_not_process'\n\nNEED_ANALYSIS = {\n 'alert_not_scan_http_direction_in',\n 'alert_not_scan_http_direction_out',\n 'alert_not_scan_http_direction_lateral',\n}\n\nfiltered = []\nprocess_type_counts = {}\nfor alert in normalized_alerts:\n alert = dict(alert)\n if filter_enabled:\n ptype = get_process_type(alert)\n need = ptype in NEED_ANALYSIS\n threat_type = get_threat_type(alert)\n else:\n ptype = 'filter_disabled'\n need = True\n threat_type = get_threat_type(alert)\n process_type_counts[ptype] = process_type_counts.get(ptype, 0) + 1\n alert['_process_type'] = ptype\n alert['_threat_type'] = threat_type\n if need:\n filtered.append(alert)\n\nprint(f'[filter] input={len(normalized_alerts)}, kept={len(filtered)}')\nprint(f'[filter] process_type_counts={process_type_counts}')\n\nstats['after_filter_count'] = len(filtered)\nstats['filter_removed_count'] = len(normalized_alerts) - len(filtered)\nstats['filter_process_type_counts'] = process_type_counts\n\noutputs['filtered_alerts'] = filtered\noutputs['stats'] = stats\nfor k in ['input_mode', 'dedup_enabled', 'dedup_threshold', 'strict_fields',\n 'lsh_fields', 'max_field_len', 'max_dedup_keys']:\n outputs[k] = inputs.get(k)\n" + "code": "\nnormalized_alerts = inputs.get('normalized_alerts', [])\nfilter_enabled = inputs.get('filter_enabled', True)\nbatch_hint = str(inputs.get('source_log_type', 'tdp') or 'tdp').lower()\nif batch_hint not in ('tdp', 'skyeye'):\n batch_hint = 'tdp'\nstats = dict(inputs.get('stats', {}))\n\ndef is_scan_alert(threat_name):\n tnl = str(threat_name or '').lower()\n return ('扫描' in tnl) and ('webshell' not in tnl)\n\ndef get_threat_type(alert):\n src = alert.get('_source_type') or batch_hint\n if src == 'skyeye':\n return str(alert.get('threat_type', 'general') or 'general')\n return str(alert.get('threat_name', 'general') or 'general')\n\ndef is_http(alert):\n for field in ('application_layer_protocol', 'net_type', 'net_app_proto'):\n val = str(alert.get(field, '') or '').lower()\n if val and val != 'none' and 'http' in val:\n return True\n return False\n\ndef get_process_type(alert):\n src = alert.get('_source_type') or batch_hint\n threat_name = alert.get('threat_name', '')\n direction = str(alert.get('direction', '') or '').lower()\n scan = is_scan_alert(threat_name)\n http = is_http(alert)\n if src == 'skyeye':\n return 'alert_scan_direction_in' if scan else 'alert_not_scan_http_direction_in'\n if scan:\n return f'alert_scan_direction_{direction}' if direction in ('in', 'out', 'lateral') else 'alert_scan_direction_in'\n if http:\n return f'alert_not_scan_http_direction_{direction}' if direction in ('in', 'out', 'lateral') else 'alert_not_scan_http_direction_in'\n return f'alert_not_scan_not_http_direction_{direction}' if direction in ('in', 'out', 'lateral') else 'alert_not_process'\n\nNEED_ANALYSIS = {\n 'alert_not_scan_http_direction_in',\n 'alert_not_scan_http_direction_out',\n 'alert_not_scan_http_direction_lateral',\n}\n\nfiltered = []\nprocess_type_counts = {}\nfor alert in normalized_alerts:\n alert = dict(alert)\n if filter_enabled:\n ptype = get_process_type(alert)\n need = ptype in NEED_ANALYSIS\n threat_type = get_threat_type(alert)\n else:\n ptype = 'filter_disabled'\n need = True\n threat_type = get_threat_type(alert)\n process_type_counts[ptype] = process_type_counts.get(ptype, 0) + 1\n alert['_process_type'] = ptype\n alert['_threat_type'] = threat_type\n if need:\n filtered.append(alert)\n\nprint(f'[filter] input={len(normalized_alerts)}, kept={len(filtered)}')\nprint(f'[filter] process_type_counts={process_type_counts}')\n\nstats['after_filter_count'] = len(filtered)\nstats['filter_removed_count'] = len(normalized_alerts) - len(filtered)\nstats['filter_process_type_counts'] = process_type_counts\n\noutputs['filtered_alerts'] = filtered\noutputs['stats'] = stats\noutputs['input_mode'] = inputs.get('input_mode')\noutputs['dedup_enabled'] = inputs.get('dedup_enabled')\noutputs['dedup_threshold'] = inputs.get('dedup_threshold')\noutputs['strict_fields'] = inputs.get('strict_fields')\noutputs['lsh_fields'] = inputs.get('lsh_fields')\noutputs['max_field_len'] = inputs.get('max_field_len')\noutputs['max_dedup_keys'] = inputs.get('max_dedup_keys')\n" }, { "id": "dedup_and_write", @@ -34,17 +34,54 @@ { "from": "receive_alert", "to": "normalize", - "order": 0 + "order": 0, + "mapping": { + "raw_alerts": "raw_alerts", + "input_mode": "input_mode", + "source_log_type": "source_log_type", + "filter_enabled": "filter_enabled", + "dedup_enabled": "dedup_enabled", + "dedup_threshold": "dedup_threshold", + "strict_fields": "strict_fields", + "lsh_fields": "lsh_fields", + "max_field_len": "max_field_len", + "max_dedup_keys": "max_dedup_keys", + "stats": "stats" + } }, { "from": "normalize", "to": "filter_logs", - "order": 0 + "order": 0, + "mapping": { + "normalized_alerts": "normalized_alerts", + "input_mode": "input_mode", + "source_log_type": "source_log_type", + "filter_enabled": "filter_enabled", + "dedup_enabled": "dedup_enabled", + "dedup_threshold": "dedup_threshold", + "strict_fields": "strict_fields", + "lsh_fields": "lsh_fields", + "max_field_len": "max_field_len", + "max_dedup_keys": "max_dedup_keys", + "stats": "stats" + } }, { "from": "filter_logs", "to": "dedup_and_write", - "order": 0 + "order": 0, + "mapping": { + "filtered_alerts": "filtered_alerts", + "input_mode": "input_mode", + "dedup_enabled": "dedup_enabled", + "dedup_threshold": "dedup_threshold", + "strict_fields": "strict_fields", + "lsh_fields": "lsh_fields", + "max_field_len": "max_field_len", + "max_dedup_keys": "max_dedup_keys", + "stats": "stats" + } } ], "metadata": { @@ -80,6 +117,10 @@ "threat_type": "web攻击" } ] + }, + "runtime": { + "strict_edge_mapping": true, + "dataflow_mode": "vertex_cache" } }, "triggers": [ diff --git a/.flocks/flockshub/plugins/workflows/stream_alert_denoise/workflow.md b/.flocks/flockshub/plugins/workflows/stream_alert_denoise/workflow.md index 4217ff517..8e831cf97 100644 --- a/.flocks/flockshub/plugins/workflows/stream_alert_denoise/workflow.md +++ b/.flocks/flockshub/plugins/workflows/stream_alert_denoise/workflow.md @@ -30,6 +30,8 @@ receive_alert -> normalize -> filter_logs -> dedup_and_write ``` +节点边使用显式字段映射并启用严格映射模式。每一跳只传递当前阶段的告警列表、统计和后续必需配置，不透传完整上游 payload。 + | 顺序 | 节点 | 作用 | | --- | --- | --- | | 1 | `receive_alert` | 接收输入，判断输入模式和来源类型。 | diff --git a/.flocks/flockshub/plugins/workflows/stream_alert_triage/guide.md b/.flocks/flockshub/plugins/workflows/stream_alert_triage/guide.md index 108483378..627266ef6 100644 --- a/.flocks/flockshub/plugins/workflows/stream_alert_triage/guide.md +++ b/.flocks/flockshub/plugins/workflows/stream_alert_triage/guide.md @@ -1,370 +1,211 @@ -# stream_alert_triage 配置引导 +# stream_alert_triage 配置指南 -这个文件是 `stream_alert_triage` 的工作流专属 `guide.md`。Rex 处理这个工作流的配置、输入、并发、缓存、验证或查配置快捷入口时，必须先读取本文全文，再把 `workflow.md`、`workflow.json` 和 `workflow_config_manage(action="get" 或 "status", workflow_id="stream_alert_triage")` 的结果作为支撑上下文。 +本文是 `stream_alert_triage` 的工作流专属配置指南。配置、验证或查询该工作流时，先读取本文，再结合 `workflow.md`、`workflow.json` 和后端运行态配置。 -`workflow-config-guide` skill 只提供交互协议；本文才是本工作流配置细节、默认选项、提问顺序和验证方式的来源。 +## 0. 配置库访问约束 -Rex 引导用户时必须遵守： - -1. 根据用户点击的入口或自然语言需求，自动定位本文相关章节。 -2. 一次只问一个最关键问题。 -3. 每个选择都必须允许自定义/补充输入；没有补充则填 `none`。 -4. 涉及输入来源、并发、缓存上限、持久化输出或发布模板变更时，先展示计划和 diff，再用 question 工具确认。 -5. 查配置只能只读，不得修改文件、触发 LLM 研判、发布 API、启动监听或清理缓存。 - -## 0. 后端配置库访问约束 - -本节优先级高于通用会话提示中的后端 API token 或 curl 示例。处理本工作流的发布、定时触发、API 接入、输出策略或查配置时，必须按本文执行： - -- 配置库读取/写入必须使用内置工具 `workflow_config_manage`，不要读取 `server_api_token` 或 `service_api_token`，也不要手工 curl 本机后端配置接口。 -- 查配置使用 `workflow_config_manage(action="get", workflow_id="stream_alert_triage")` 或 `workflow_config_manage(action="status", workflow_id="stream_alert_triage")`。 -- 查定时触发配置使用 `workflow_config_manage(action="get", workflow_id="stream_alert_triage", config_type="poller")` 或 `workflow_config_manage(action="status", workflow_id="stream_alert_triage", config_type="poller")`。 -- 修改配置前先使用 `workflow_config_manage(action="diff", workflow_id="stream_alert_triage", config={...})` 展示差异并用 question 工具确认；确认后才使用 `workflow_config_manage(action="put", workflow_id="stream_alert_triage", config={...})`。 -- 修改定时触发配置前先使用 `workflow_config_manage(action="diff", workflow_id="stream_alert_triage", config_type="poller", config={...})` 展示差异并用 question 工具确认；确认后才使用 `workflow_config_manage(action="put", workflow_id="stream_alert_triage", config_type="poller", config={...})`。 -- 如果后端配置库没有模板，只能使用 `workflow_config_manage(action="sync", workflow_id="stream_alert_triage")`，让后端从工作流目录 `config.json` 迁移或生成模板。 -- `config.json` 只能作为模板来源或兜底迁移来源，不是直接写入目标，也不能证明配置已生效。 -- 需要启动或停止 API 服务、定时触发或其它运行态能力时，必须使用对应运行态接口；不要通过修改模板字段冒充运行态状态。 -- 如果 `workflow_config_manage` 不可用、返回未授权、拒绝访问、连接失败或后端不可达，必须停止配置流程，明确说明本次未应用、未发布、未启动；如已生成目标配置，只能保存草稿到 outputs，不要继续读取 token 或改写 `config.json`。 +- 查询工作流配置使用 `workflow_config_manage(action="get", workflow_id="stream_alert_triage")` 或 `status`。 +- 查询 poller 使用 `workflow_config_manage(action="get", workflow_id="stream_alert_triage", config_type="poller")`。 +- 修改前必须先调用 `workflow_config_manage(action="diff", workflow_id="stream_alert_triage", config_type="poller", config={...})` 展示差异并取得确认。 +- 用户确认后才调用 `workflow_config_manage(action="put", workflow_id="stream_alert_triage", config_type="poller", config={...})`。 +- 如果后端没有模板，可使用 `workflow_config_manage(action="sync", workflow_id="stream_alert_triage")` 从工作流目录迁移模板。 +- 不要读取 `server_api_token` 或 `service_api_token`，不要手工调用 `/api/workflow/stream_alert_triage/poller-config`。 +- 后端不可用时只能把目标配置保存为 outputs 下的草稿，并明确说明未应用、未发布、未启动。 +- 查询配置是只读操作，不得触发 LLM、写缓存、写 SOC DB、写 JSONL 或修改游标。 ## 1. 工作流定位 - 工作流 ID：`stream_alert_triage` -- 工作流名称：`stream_alert_denoise` 的下游批量研判 Pipeline。 -- 主要用途：读取 `stream_alert_denoise` 写出的 `dedup_result_NNN.jsonl`，按 `dedup_key` 做 leader/follower 分组，只对每组 leader 执行研判，followers 复用 leader 结果。 -- 当前状态：`meta.json` 标记为 `active`。 -- 当前发布状态：`workflow.json` 中 `triggers` 为空；工作流目录有 `config.json` 用于声明默认持久化策略。默认按手动运行或 API run 输入来引导。 - -本工作流适合： +- 上游：`stream_alert_denoise` +- 入口：`load_dedup_file` +- 流程：`load_dedup_file -> concurrent_triage -> commit_cursor -> summarize` +- 默认输出：`soc_db` +- 默认批次：最多 10 条、32 MiB +- 默认调度模板：5 分钟，`noOverlap=true` -- 对降噪后的 HTTP 告警做攻击研判。 -- 复用 `stream_alert_denoise` 的 `dedup_key` 降低 LLM 调用。 -- 按日期重放某天全部去重结果。 -- 默认只把明确 `is_duplicate=false`、包含 `dedup_key` 且批内首次出现的研判告警写入 `~/.flocks/data/soc.db`，并保留 JSONL 可选输出，供归档或下游工作流消费。 +本工作流只处理上游去重后的 HTTP 告警，不负责原始告警接入、字段归一化、LSH 去重或跨日期积压补偿。 -本工作流不适合： +## 2. 引导顺序 -- 直接接收原始 TDP / SkyEye 告警。 -- 执行上游过滤、字段归一化或 LSH 去重。 -- 调用或嵌入 `tdp_alert_triage` 子工作流。 -- 生成每条告警一个独立 markdown 报告文件。 +一次只确认一个关键问题，推荐顺序： -## 2. AI 引导方式 +1. 使用每日自动增量模式，还是显式文件重放？ +2. 是否保持 `batch_max_records=10` 和 `batch_max_bytes=33554432`？ +3. 是否保持 `concurrency=1` 和 `max_triage_cache_size=100000`？ +4. 输出使用 `soc_db`、`jsonl`、`both` 还是 `none`？ +5. 是否启用每 5 分钟一次的定时触发？ +6. 展示计划、配置 diff 和副作用后，再确认应用或只保存草稿。 -如果用户点击输入入口，优先确认要读哪些 `dedup_result_NNN.jsonl`；如果用户点击规则入口，优先确认并发和缓存策略；如果用户点击样例入口，优先做文件格式检查，不要直接触发 LLM 研判。 - -推荐提问顺序： - -1. 你要读取上游本次输出文件、单个文件、某个日期，还是默认读取今天？ -2. 是否显式设置 `concurrency=1`，还是确认提高到 2 到 5？ -3. 是否保持 `max_triage_cache_size=100000` 和 `triage_output_mode=soc_db`，还是改为 JSONL / both / none？ -4. 是否开启定时触发；默认建议每 3 分钟执行一次，并显式使用 `concurrency=1`。 -5. 是否只做轻量文件检查，还是确认执行真实研判？ -6. 是否保存配置草稿、应用发布模板，或暂不修改？ - -如果用户只问“查一下现在怎么配的”，不要提问，直接按第 9 节只读检查。 +如果用户只要求“查配置”，直接执行第 9 节的只读检查，不提问、不修改。 ## 3. 输入模式 -工作流代码支持四种输入定位方式，解析优先级固定为： - -1. `input_paths`: 显式 JSONL 文件路径列表，推荐直接使用 `stream_alert_denoise.outputs.output_paths`。 -2. `input_path`: 单个 JSONL 文件路径，通常来自 `stream_alert_denoise.outputs.output_path`。 -3. `input_date`: `YYYY-MM-DD`，自动读取该日 `stream_alert_denoise` 输出目录下全部 `dedup_result_*.jsonl`。 -4. 全部不传：默认读取今天目录下全部 `dedup_result_*.jsonl`。 +### 自动目录模式 -上游默认输出目录： +不传 `input_path` 和 `input_paths`。未显式设置 `input_date` 时，每次运行动态计算当天日期并扫描： ```text ~/.flocks/workspace/workflows/stream_alert_denoise//dedup_result_NNN.jsonl ``` -输入模式建议： - -| 模式 | 适用场景 | 推荐输入 | -| --- | --- | --- | -| 上游本次输出 | 刚跑完 `stream_alert_denoise`，需要立即研判本批首见告警 | `input_paths = denoise.outputs.output_paths` | -| 单文件重放 | 只检查某个文件或某个序号文件 | `input_path = ".../dedup_result_001.jsonl"` | -| 按日期重放 | 对某天所有去重结果统一研判 | `input_date = "YYYY-MM-DD"` | -| 今日默认 | 调试或日常手动执行 | 不传 `input_*`，但先确认今天目录存在文件 | - -默认推荐：上游本次输出文件。如果用户没有上游结果，再推荐 `input_date`。 - -互斥关系： - -- `input_paths` 和 `input_path` 可以同时传，但会合并并按顺序去重。 -- 只要显式路径存在，就不会再按 `input_date` 自动发现。 -- 传入不存在的路径会被跳过，不会报错中止，但 `load_stats` 会显示实际读取为 0。 - -## 4. 来源形态 - -真实来源是 `stream_alert_denoise` 的 JSONL 输出。每个文件形态： - -- 第一行：`{"_type":"file_header", ...}`，`load_dedup_file` 会跳过。 -- 后续每行：一条 JSON 告警。 -- 关键字段：`dedup_key`、`is_duplicate`、`_lsh_cluster_id`、`_source_type`、`_process_type`、`sip`、`dip`、`req_http_url`、`req_body`、`rsp_body`、`threat_name`。 - -`load_dedup_file` 输出： +自动模式读取并在成功后更新： -| 字段 | 说明 | -| --- | --- | -| `enriched_alerts` | 从 JSONL 读出的告警列表 | -| `loaded_files` | 实际读取到的文件路径 | -| `load_stats` | 文件数、记录数、跳过 header 数、坏行数 | -| `concurrency` | 下游外层并发参数 | -| `max_triage_cache_size` | 下游研判缓存上限 | -| `input_date` | 实际日期字符串 | - -重要约束： +```text +~/.flocks/workspace/workflows/stream_alert_triage/.triage_cursor.json +``` -- 如果告警缺少 `dedup_key`，该告警会作为独立 work unit 研判，无法和其它告警复用。 -- `is_duplicate` 不决定是否研判。缓存命中策略只依赖 `dedup_key`。 -- `is_duplicate` 决定是否允许写入 SOC DB：只有明确为 `false` 的告警才有资格持久化；缺少该字段时按“未证明首见”处理，不写入 SOC DB。 -- SOC DB 持久化要求非空 `dedup_key`：批内只接受第一条，数据库再通过唯一索引保证跨执行全局唯一；无 `dedup_key` 的防御性研判结果不会写入 SOC DB。 -- `stream_alert_denoise` 只会把跨批次首见告警写入 JSONL；因此常规情况下本工作流读取到的是适合继续研判的首见告警。 -- 如果用户手工构造 JSONL，必须保证每行是独立 JSON 对象，不能是整文件 JSON 数组。 +日期切换会忽略旧日期游标，从新日期第一个文件开始。前一天未消费完的数据会丢弃，这是当前设计的明确取舍。 +游标只接受 `version=2`。除序号和偏移量外，还会校验 device ID、file ID、文件头 SHA-256 和游标前最多 4 KiB 内容的 SHA-256；身份或内容不匹配、旧版本以及字段损坏都会设置 `cursor_invalidated=true` 并从当前文件头安全重读。文件打开后会执行二次校验，避免校验与读取之间的同名替换竞态。 -## 5. 输出去向 +### 显式重放模式 -工作流返回： +传入 `input_path` 或 `input_paths`。显式模式： -| 输出字段 | 说明 | -| --- | --- | -| `enriched_alerts_with_triage` | 每条输入告警加上研判字段后的完整列表 | -| `triage_results` | 精简研判结果列表，不含 markdown 正文 | -| `triage_stats` | leader/follower、cache、并发、耗时、verdict 分布等统计 | -| `load_stats` | 输入文件加载统计 | -| `loaded_files` | 实际读取的上游文件 | -| `input_date` | 本次读取日期 | -| `triage_output_mode` | 本次生效的输出模式：`soc_db` / `jsonl` / `both` / `none` | -| `soc_db_result` / `soc_db_path` | 本次写入的 SOC DB 结果和路径；结果区分 `inserted_rows` 与 `updated_rows` | -| `output_paths` | 本次写入的研判 JSONL 文件列表；未启用 JSONL 时为空 | -| `output_dir` | 研判 JSONL 结果目录；未启用 JSONL 时为空 | -| `summary_report` | markdown 总览文本 | -| `summary_path` | 总览 markdown 落盘路径 | -| `top_attack_verdict` / `top_risk_level` / `top_report_title` / `top_triage_report` | top-risk 告警研判字段 | - -每条告警追加： - -- `has_dedup_key` -- `triage_source` -- `triage_status` -- `attack_verdict` -- `risk_level` -- `report_title` -- `triage_report` -- `attack_success` -- `triage_ms` -- `triage_error` - -默认 SOC DB 输出： - -```text -~/.flocks/data/soc.db -``` +- 不读取、不修改生产游标。 +- 仍限制每批最多 10 条和 32 MiB。 +- 按调用方路径顺序去重处理。 +- 路径不存在时记录 `missing_files`，不回退自动目录。 +- 将输出的 `next_cursor` 作为下一次输入的 `resume_cursor` 继续读取。 -默认写入表： +重放示例： -```text -alert_records +```json +{ + "input_path": "~/.flocks/workspace/workflows/stream_alert_denoise//dedup_result_001.jsonl", + "batch_max_records": 10, + "batch_max_bytes": 33554432, + "concurrency": 1, + "triage_output_mode": "soc_db" +} ``` -可选研判 JSONL 输出目录： +下一批： -```text -~/.flocks/workspace/workflows/stream_alert_triage//triage_result_NNN.jsonl +```json +{ + "input_path": "~/.flocks/workspace/workflows/stream_alert_denoise//dedup_result_001.jsonl", + "resume_cursor": { + "version": 2, + "date": "", + "file_seq": 1, + "file_name": "dedup_result_001.jsonl", + "byte_offset": 1837294, + "device_id": 16777234, + "file_id": 123456, + "head_hash": "<64 位 SHA-256>", + "boundary_start": 1833198, + "boundary_hash": "<64 位 SHA-256>", + "file_index": 0, + "path": "~/.flocks/workspace/workflows/stream_alert_denoise//dedup_result_001.jsonl" + } +} ``` -写入规则： - -- SOC DB 只接受明确 `is_duplicate=false`、包含非空 `dedup_key` 且该 key 在本批次首次出现的告警。 -- `is_duplicate=true`、缺少 `is_duplicate`、缺少 `dedup_key` 或批内重复 `dedup_key` 的告警均不会写入 SOC DB。 -- `alert_records.dedup_key` 使用部分唯一索引保证跨批次、跨执行只能存在一条告警记录。 -- 已存在的 `dedup_key` 再次回放时只更新研判字段和研判运行标记，不覆盖首次告警的时间、来源、行号、标识与原始事件字段。 -- SOC DB 建表、迁移或写入失败会使本次工作流失败，不会以“成功但写入 0 条”结束。 -- `soc_db_result` 记录本次持久化总数、新增数与更新数；`triage_stats.soc_db_filter_stats` 记录候选数及各类跳过数。 -- `triage_output_mode=soc_db`：默认写入 `soc.db`，不写 JSONL。 -- `triage_output_mode=jsonl`：只写 `triage_result_NNN.jsonl`，不写 `soc.db`。 -- `triage_output_mode=both`：同时写 `soc.db` 和 JSONL。 -- `triage_output_mode=none`：不写 `soc.db` 和 JSONL，但仍可能写缓存。 -- 旧参数 `persist_triage_output=true` 仍兼容：当 `triage_output_mode=soc_db` 时会额外写 JSONL，相当于 `both`。 -- 每个文件第一行是 file header，包含 `workflow`、`seq`、`run_id`、`batch_total`、`batch_triaged`、`batch_followers_reused`、`batch_cache_hit`、`batch_triage_failed`。 -- 每个文件最多 10000 条告警记录。 -- `.triage_counter.json` 记录当前文件序号和条数。 -- 未启用 JSONL 时不写 `triage_result_NNN.jsonl`，但仍可能触发 LLM 和缓存写入，除非全部 cache 命中或没有输入。 - -总览报告输出： +`resume_cursor` 应原样使用上一批返回的完整 `next_cursor`；以上数值仅展示字段结构。 -```text -~/.flocks/workspace/outputs//artifacts/stream_alert_triage_summary.md -``` - -注意： +## 4. 有界加载规则 -- 每条告警完整 markdown 在 `triage_report` 字段中。 -- 不存在 `report_path`。 -- 不会生成 `triage_report_*.md` 这类单告警 markdown 文件。 +| 参数 | 默认值 | 非法值处理 | +| --- | ---: | --- | +| `batch_max_records` | `10` | 回退 10 | +| `batch_max_bytes` | `33554432` | 回退 32 MiB | +| `concurrency` | `1` | 回退 1 | +| `max_triage_cache_size` | `100000` | 回退 100000 | -- `triage_report` 是带语义标签的 markdown 字符串，根标签为 ``。 -- 前端应按 ``、``、``、``、``、``、``、``、`` 切块后渲染标签内 markdown。 -- 工作流的报告生成 prompt 已包含攻击成功和攻击失败 few-shot；如果 LLM 未按标签输出，会自动使用确定性 fallback。 +- 多文件按文件名中的数字序号排序，不使用字符串排序。 +- 批次限制针对整个日期目录或整组显式路径，不是每个文件各算一批。 +- 有效 JSON 对象计入告警数。 +- header、空行、坏 JSON 和非对象不计入告警数，但完整行会推进待提交 offset。 +- 未达到单行字节上限且没有换行符的末尾半行不解析、不推进 offset。 +- 超大单行可跨多个批次分块跳过到下一换行符，每批实际读取量仍不超过 + `batch_max_bytes`；跨批状态由 `next_cursor` 自动携带，并记录 `oversized_lines`。 +- `loaded_files` 只包含本批实际触达的文件。 +## 5. 游标与提交条件 -## 6. 处理规则 +自动模式在读取游标前获取批次租约，覆盖完整的 `load -> triage -> persist -> commit`；重叠执行会以 `production_batch_lease_busy` 失败，提交或异常后释放。入口节点只计算 `pending_cursor` 和 `cursor_revision`。生产游标由 `commit_cursor` 在 `concurrent_triage` 成功后，在游标文件锁内通过 revision CAS 和单调性校验原子提交。 -默认参数： +提交生产游标的情况： -| 参数 | 推荐默认 | 代码行为和说明 | -| --- | --- | --- | -| `input_paths` | 无 | 显式路径列表，优先级最高 | -| `input_path` | 无 | 单个显式路径 | -| `input_date` | 今天 | 自动发现该日所有上游 `dedup_result_*.jsonl` | -| `concurrency` | `1` | 控制外层 work unit 和运行级 LLM 并发预算；`concurrent_triage` 会限制到 1 到 5 | -| `max_triage_cache_size` | `100000` | 小于 1 时回退 100000 | -| `triage_output_mode` | `soc_db` | 输出模式：`soc_db` / `jsonl` / `both` / `none` | -| `soc_db_path` | `~/.flocks/data/soc.db` | 默认 SOC DB 写入位置 | -| `persist_triage_output` | `false` | 旧兼容参数；设为 `true` 会在 `soc_db` 模式下额外写 JSONL | -| `jsonl_output_dir` | 空 | 可选 JSONL 输出目录；为空时使用工作流默认日期目录 | +- 所有告警研判完成，所有启用的持久化目标成功。 +- 部分单条研判失败，但失败状态已经形成。 +- 本批只消费了 header、空行或完整坏行。 -并发注意： +不提交生产游标的情况： -- 外层 `ThreadPoolExecutor(max_workers=concurrency)` 处理 unique work units。 -- 每个 leader 仍会执行 `survey`、`cve_related`、`cve_info`、`payload_analysis` 4 个 LLM 分支。 -- 所有 `llm.ask()` 共享运行级信号量，稳态 LLM 峰值不超过 `concurrency`，不会再与 4 个分支相乘。 -- 配置引导应默认显式给出 `concurrency=1`。如果用户要提高到 2 到 5，先说明 LLM 和上游工具压力，再确认。 -- `load_dedup_file` 在完全不传 `concurrency` 时同样输出 1，与文档和样例保持一致。 +- SOC DB 写入失败。 +- 启用 JSONL 时 JSONL 写入失败。 +- 研判缓存写入失败。 +- 当前游标 revision 与读取时不一致，或新位置违反单调性。 +- 节点超时、取消或进程退出。 +- 本批没有读取任何新字节。 +- 显式重放模式。 -leader/follower 规则： +游标写入使用包含 run ID 的同目录唯一临时文件、`flush()`、`os.fsync()` 和 `os.replace()`；CAS 失败返回 `stale_cursor_commit`。文件身份失效触发的明确重置仍要求 revision 未变化，但允许偏移回到文件头。单条研判失败不会阻塞流，失败告警保留 `triage_status=failed`、`triage_error`、`triage_attack_verdict=unknown` 和 `triage_attack_success=unknown`。模型判定节点直接生成这两个字段：`triage_attack_verdict` 表示是否攻击，枚举为 `attack | non_attack | unknown`；`triage_attack_success` 表示攻击结果，枚举为 `success | failed | unknown`。原始告警的 `attack_verdict`、`attack_success`、`threat_result` 等字段保持不变，不参与研判结果字段的生成。 -- 按 `dedup_key` 分组。 -- 每个分组首条为 leader。 -- leader 负责真实研判。 -- follower 复制 leader 的 `attack_verdict`、`risk_level`、`report_title`、`triage_report`、`attack_success`。 -- 无 `dedup_key` 告警各自独立研判，`triage_source` 会是 `no_dedup_key_triaged` 或 `no_dedup_key_failed`。 +## 6. 输出与副作用 -研判缓存： +默认 SOC DB： ```text -~/.flocks/workspace/workflows/stream_alert_triage/triage_cache.pkl -~/.flocks/workspace/workflows/stream_alert_triage/triage_cache.lock +~/.flocks/data/soc.db ``` -- key：`dedup_key`。 -- value：`attack_verdict`、`risk_level`、`report_title`、`triage_report`、`attack_success`。 -- cache 命中时只有在 `triage_report` 为 `soc.triage.markdown.v1` 标签化 markdown 时才直接复用，不调用 LLM。 -- 旧 `final_report` 缓存会按 miss 重新研判并写回新版字段。 -- cache 未命中时 leader 执行完整内联研判。 -- 新结果会合并写回 cache，文件锁 + 原子落盘。 -- 淘汰策略是 FIFO LRU，超过 `max_triage_cache_size` 时丢弃最旧条目。 - -研判逻辑： +可选 JSONL： -- 本工作流不调用 `tdp_alert_triage` 子工作流。 -- 研判逻辑是内联实现，语义同源于 `tdp_alert_triage` 文档版本。 -- 单条 leader 会执行情报准备、4 路 LLM 分析、攻击状态判断、verdict 归一化、标题生成和 markdown 聚合。 -- 如果后续要接入真实 TDP 平台检索或页面调查，必须按 `tdp-use` skill 处理，不得绕过对应 skill 直接调用 TDP 工具。 - -低层参数默认隐藏，不主动询问：LLM timeout、retry、verdict 映射、JSONL counter 文件名、file lock 细节。只有用户明确排障时再解释。 - -## 7. 样例验证 - -推荐样例 1：上游本次输出文件。 - -```json -{ - "input_paths": [ - "~/.flocks/workspace/workflows/stream_alert_denoise/2026-05-18/dedup_result_001.jsonl" - ], - "concurrency": 1, - "max_triage_cache_size": 100000, - "triage_output_mode": "soc_db" -} +```text +~/.flocks/workspace/workflows/stream_alert_triage//triage_result_NNN.jsonl ``` -推荐样例 2：按日期重放。 +总览报告： -```json -{ - "input_date": "2026-05-18", - "concurrency": 1, - "max_triage_cache_size": 100000, - "triage_output_mode": "soc_db" -} +```text +~/.flocks/workspace/outputs//artifacts/stream_alert_triage_summary.md ``` -轻量验证优先做只读检查： - -1. 路径是否存在。 -2. 首行是否为 `_type=file_header` 或第一条 JSON 告警。 -3. 后续每行是否能按 JSON 对象解析。 -4. 是否至少有 `dedup_key`、`sip`、`dip`、`req_http_url`、`threat_name` 中的关键字段。 -5. 按 `dedup_key` 估算 unique work units 和 follower 数。 -6. 确认运行级 LLM 峰值不超过 `concurrency`。 - -真实执行验证注意： +输出模式： -- 只要存在 cache miss，就会触发 LLM 和情报工具调用。 -- `triage_output_mode=none` 只是不写 SOC DB 和 JSONL，不代表不会调用 LLM，也不代表不会写 `triage_cache.pkl`。 -- `persist_triage_output=false` 只是旧 JSONL 开关为关；默认仍会按 `triage_output_mode=soc_db` 写入 SOC DB。 -- 如果要避免外部副作用，先只做文件解析和字段检查，不运行工作流。 -- 如果用户确认运行，建议用 1 到 3 条样例告警、`concurrency=1`、明确是否允许写缓存和输出文件。 - -最小期望输出： +| 模式 | 行为 | +| --- | --- | +| `soc_db` | 只写 SOC DB | +| `jsonl` | 只写 JSONL | +| `both` | 两者都写，任一失败都会使批次失败且不提交游标 | +| `none` | 不写业务持久化目标，但仍可能调用 LLM 和写研判缓存 | -- `load_stats.record_count > 0` -- `triage_stats.total == load_stats.record_count` -- `triage_stats.work_units <= triage_stats.total` -- `enriched_alerts_with_triage[*].triage_report` 存在于已研判或缓存命中的告警上 -- `soc_db_result.rows` 等于本次持久化的候选数（`inserted_rows + updated_rows`）；应与 `triage_stats.soc_db_first_seen_rows` 一致，除非 `triage_output_mode=jsonl/none` -- `triage_stats.soc_db_skipped_rows` 等于输入总数减去首见唯一告警数，并可通过 `triage_stats.soc_db_filter_stats` 查看具体跳过原因 -- `output_paths` 指向当日 `triage_result_NNN.jsonl`，仅在 `triage_output_mode=jsonl/both` 或旧参数 `persist_triage_output=true` 时存在 -- `summary_path` 指向 `outputs//artifacts/stream_alert_triage_summary.md` +`persist_triage_output=true` 是兼容参数，会在 `soc_db` 模式下额外启用 JSONL。 -## 8. 应用方式 +主要增量状态输出：`cursor_enabled`、`cursor_before`、`cursor_revision`、`cursor_invalidated`、`pending_cursor`、`next_cursor`、`cursor_committed`、`committed_cursor`、`has_more`、`batch_records`、`batch_bytes` 和 `load_stats`。批次租约句柄与 token 仅用于节点间内部传递，不属于业务输出。 -当前工作流目录包含 `config.json` 作为默认运行配置，默认输出到 `~/.flocks/data/soc.db`，不写 JSONL；`workflow.json` 中 `triggers` 为空。配置引导应把它当作“已有输出默认值、尚未声明发布/触发模板”的工作流。 +## 7. 调度与吞吐 -如果用户要配置运行入口： +工作流模板的 schedule trigger 为 5 分钟一次，保持 `noOverlap=true` 和单写者语义。每批最多 10 条时，理论上限为： -1. 优先引导为手动运行或 API run 输入参数模板。 -2. 如果用户要发布成 API 服务，应使用 `workflow_config_manage(action="get" 或 "sync" 或 "diff" 或 "put", workflow_id="stream_alert_triage")` 流程。 -3. 如果用户要开启定时触发，默认建议 3 分钟一次；应用前必须确认触发输入来源、输出模式、是否允许写入 `soc.db` 和是否允许触发 LLM，并使用 `workflow_config_manage(action="get" -> "diff" -> "put", workflow_id="stream_alert_triage", config_type="poller")` 读取和写入 poller 配置。 -4. 如果需要扩展工作流目录下的 `config.json`，必须使用 runtime 消费的结构：`kind: workflow.integration-config`，顶层包含 `publish` 和 `triggers`。 -5. 不要生成旧的 `publishTemplates` wrapper。 -6. 不要直接写 `config.json` 来表示发布、接入或触发配置已经生效。 -7. 启停、发布、取消发布等运行态动作必须调用运行时接口。 -8. 不要读取 `server_api_token`，不要用 curl 调 `/api/workflow/stream_alert_triage/poller-config` 读取或写入定时配置。 -9. 如果后端配置接口不可用，只能把目标配置保存为草稿到 outputs，并明确说明未应用、未发布、未启动。 +```text +10 × 12 × 24 = 2880 条/天 +``` -应用变更前必须展示： +LLM、情报查询、缓存 miss 和持久化耗时会降低实际吞吐。超过当天处理能力的积压会在日期切换时丢弃。启用定时触发前必须明确说明这一容量边界。 -- 计划。 -- 输入参数或 publish / triggers 模板 diff。 -- poller 配置变更时必须展示 `workflow_config_manage(config_type="poller")` 生成的 diff。 -- 是否会触发 LLM、情报工具、`triage_cache.pkl` 写入、`soc.db` 写入、`triage_result_NNN.jsonl` 写入。 -- question 工具确认：应用、保存草稿或暂不修改。 -- 用户确认应用后，使用 `workflow_config_manage(action="put", workflow_id="stream_alert_triage", config_type="", config={...})` 写入完整配置。 +## 8. 应用配置 -不要通过删除 `triage_cache.pkl` 来“重置配置”。缓存清理是运行数据操作，必须单独说明影响并取得确认。 +应用前展示： -## 9. 查配置 +- 自动增量或显式重放的输入模式。 +- 批次条数、字节上限、并发和缓存上限。 +- 输出模式及会发生的 SOC DB、JSONL、缓存和游标写入。 +- poller 的 5 分钟间隔与 `noOverlap=true`。 +- 完整 diff。 -只读检查顺序： +用户确认后才可调用 `workflow_config_manage(action="put", workflow_id="stream_alert_triage", config_type="poller", config={...})`。不要通过直接修改 `config.json` 冒充运行态配置已经生效，也不要通过删除 `triage_cache.pkl` 重置配置。 -1. 读取本文。 -2. 读取 `workflow.md` 和 `workflow.json`。 -3. 调用 `workflow_config_manage(action="get", workflow_id="stream_alert_triage")` 或 `workflow_config_manage(action="status", workflow_id="stream_alert_triage")`。 -4. 调用 `workflow_config_manage(action="get", workflow_id="stream_alert_triage", config_type="poller")` 或 `workflow_config_manage(action="status", workflow_id="stream_alert_triage", config_type="poller")`。 -5. 如果后端无配置，再检查工作流目录是否有 `config.json`。 -6. 汇总已配置项、缺失项和最推荐下一步。 +## 9. 只读查配置 -查配置时重点报告： +按以下顺序执行： -- 当前是否存在 `config.json` 或后端配置，默认输出方式是否为 `soc_db`。 -- `workflow.json.triggers` 是否为空。 -- 是否已经配置定时触发；如果没有，说明默认推荐是每 3 分钟一次，但需要用户确认后才应用。 -- 推荐输入方式：`input_paths`、`input_path`、`input_date` 或今日默认。 -- 推荐显式设置 `concurrency=1`。 -- 当前缓存路径和上限配置，但不要读取或修改大型 pickle 内容，除非用户明确要求排障。 -- 输出路径：SOC DB 默认在 `~/.flocks/data/soc.db`；triage JSONL 可选在 `~/.flocks/workspace/workflows/stream_alert_triage//`；总览报告在 `~/.flocks/workspace/outputs//artifacts/`。 +1. 读取本文、`workflow.md` 和 `workflow.json`。 +2. 调用 `workflow_config_manage(action="get", workflow_id="stream_alert_triage")`。 +3. 调用 `workflow_config_manage(action="get", workflow_id="stream_alert_triage", config_type="poller")`。 +4. 后端没有模板时，只读检查工作流目录的 `config.json`。 +5. 汇总输入模式、批次限制、输出方式、poller 状态、`noOverlap`、游标路径和剩余缺口。 -查配置不得修改文件、触发 LLM、调用情报工具、写缓存、启动监听、发布 API 或停止服务。 +查配置不得修改文件、触发研判、调用情报工具、写缓存、写 SOC DB、写 JSONL、修改生产游标或启动服务。 diff --git a/.flocks/flockshub/plugins/workflows/stream_alert_triage/manifest.json b/.flocks/flockshub/plugins/workflows/stream_alert_triage/manifest.json index 3dfcd21ca..896660293 100644 --- a/.flocks/flockshub/plugins/workflows/stream_alert_triage/manifest.json +++ b/.flocks/flockshub/plugins/workflows/stream_alert_triage/manifest.json @@ -6,7 +6,7 @@ "nameCn": "HTTP研判工作流", "description": "Downstream alert triage workflow that writes triage results to SOC DB by default with optional JSONL output.", "descriptionCn": "下游告警研判工作流，默认写入 SOC DB，并保留 JSONL 输出配置。", - "version": "1.0.0", + "version": "1.1.1", "author": "Flocks Team", "license": "MIT", "homepage": "", diff --git a/.flocks/flockshub/plugins/workflows/stream_alert_triage/workflow.json b/.flocks/flockshub/plugins/workflows/stream_alert_triage/workflow.json index b07e9474b..01a334536 100644 --- a/.flocks/flockshub/plugins/workflows/stream_alert_triage/workflow.json +++ b/.flocks/flockshub/plugins/workflows/stream_alert_triage/workflow.json @@ -1,53 +1,148 @@ { "name": "stream_alert_triage", "nameCn": "HTTP研判工作流", - "description": "你好", - "description_cn": "你好", + "description": "HTTP 告警增量研判工作流：按生产游标有界读取去重结果，完成研判与持久化后原子提交游标。", + "description_cn": "HTTP 告警增量研判工作流：按生产游标有界读取去重结果，完成研判与持久化后原子提交游标。", "start": "load_dedup_file", "nodes": [ { "id": "load_dedup_file", "type": "python", - "description": "一次性读取 stream_alert_denoise 写入的 JSONL 文件。输入优先级：input_paths > input_path > input_date（自动遍历该日所有 dedup_result_*.jsonl）> 当日默认。跳过 file_header 行，输出 enriched_alerts (list[dict])。", - "code": "\"\"\"\nload_dedup_file: 一次性读取 stream_alert_denoise 写入的 JSONL 文件。\n\n输入参数（按优先级）：\n - input_paths list[str] 显式文件路径列表（来自 stream_alert_denoise.outputs.output_paths）\n - input_path str 单个文件路径（来自 stream_alert_denoise.outputs.output_path）\n - input_date str 日期 YYYY-MM-DD；读取该日目录下全部 dedup_result_*.jsonl\n - 默认：取“今天”目录下全部 dedup_result_*.jsonl\n\n跳过首行 file_header（{_type: file_header}），其余每行为一条 enriched_alert。\n\n输出：\n - enriched_alerts list[dict] 含 dedup_key/is_duplicate 等字段\n - loaded_files list[str] 实际读取到的文件列表\n - load_stats dict 统计信息\n - concurrency int 外层并发数（默认 1）\n - max_triage_cache_size int 研判缓存 FIFO LRU 上限（默认 100000）\n\"\"\"\n\nimport datetime\nimport glob\nimport json\nimport os\nimport re\n\nfrom flocks.config import Config\n\nWORKFLOW_NAME = 'stream_alert_denoise'\n_JSONL_PREFIX = 'dedup_result'\n\n\ndef _dedup_root():\n flocks_root = Config().get_global().data_dir.parent\n return flocks_root / 'workspace' / 'workflows' / WORKFLOW_NAME\n\n\ndef _date_str(input_date):\n if input_date:\n s = str(input_date).strip()\n if re.match(r'^\\d{4}-\\d{2}-\\d{2}$', s):\n return s\n return datetime.datetime.now().strftime('%Y-%m-%d')\n\n\ndef _expand_paths(input_paths, input_path, input_date):\n paths = []\n if input_paths:\n if isinstance(input_paths, str):\n input_paths = [input_paths]\n for p in input_paths:\n if p:\n paths.append(os.path.expanduser(str(p)))\n if input_path:\n paths.append(os.path.expanduser(str(input_path)))\n\n if not paths:\n date_str = _date_str(input_date)\n day_dir = _dedup_root() / date_str\n pattern = str(day_dir / f'{_JSONL_PREFIX}_*.jsonl')\n paths = sorted(glob.glob(pattern))\n print(f'[load] auto-discovered date={date_str} dir={day_dir} files={len(paths)}')\n\n # Dedupe while preserving order; drop non-existent\n seen = set()\n final = []\n for p in paths:\n if p in seen:\n continue\n seen.add(p)\n if not os.path.exists(p):\n print(f'[load] WARNING: file not found, skipping: {p}')\n continue\n final.append(p)\n return final\n\n\ninput_paths = inputs.get('input_paths')\ninput_path = inputs.get('input_path')\ninput_date = inputs.get('input_date')\n\nfiles = _expand_paths(input_paths, input_path, input_date)\n\nenriched_alerts = []\nfile_stats = []\ntotal_skipped_headers = 0\ntotal_bad_lines = 0\n\nfor path in files:\n rec_count = 0\n skipped_headers = 0\n bad_lines = 0\n try:\n with open(path, 'r', encoding='utf-8') as f:\n for line in f:\n line = line.strip()\n if not line:\n continue\n try:\n obj = json.loads(line)\n except Exception:\n bad_lines += 1\n continue\n if isinstance(obj, dict) and obj.get('_type') == 'file_header':\n skipped_headers += 1\n continue\n enriched_alerts.append(obj)\n rec_count += 1\n except Exception as e:\n print(f'[load] WARNING: failed to read {path!r}: {e}')\n file_stats.append({'path': path, 'records': 0, 'headers': 0, 'bad_lines': 0, 'error': str(e)})\n continue\n total_skipped_headers += skipped_headers\n total_bad_lines += bad_lines\n file_stats.append({'path': path, 'records': rec_count, 'headers': skipped_headers, 'bad_lines': bad_lines})\n print(f'[load] {path}: {rec_count} records (+{skipped_headers} headers skipped)')\n\nload_stats = {\n 'file_count': len(files),\n 'record_count': len(enriched_alerts),\n 'header_skipped': total_skipped_headers,\n 'bad_lines': total_bad_lines,\n 'files': file_stats,\n}\nprint(f'[load] DONE files={len(files)} total_records={len(enriched_alerts)} '\n f'headers={total_skipped_headers} bad_lines={total_bad_lines}')\n\n# Down-stream runtime tunables\noutputs['enriched_alerts'] = enriched_alerts\noutputs['loaded_files'] = files\noutputs['load_stats'] = load_stats\noutputs['concurrency'] = max(1, int(inputs.get('concurrency', 1)))\noutputs['max_triage_cache_size'] = int(inputs.get('max_triage_cache_size', 100000))\noutputs['input_date'] = _date_str(input_date)\n" + "description": "按日期目录和生产游标增量读取 dedup_result_NNN.jsonl；单批最多 10 条且默认最多读取 32 MiB。显式路径进入有界重放模式，不修改生产游标。", + "code": "\"\"\"Bounded incremental loader for stream_alert_denoise JSONL output.\"\"\"\n\nimport datetime\nimport hashlib\nimport json\nimport os\nimport re\nimport sys\nimport time\n\nfrom flocks.config import Config\n\n\nIS_WINDOWS = sys.platform == 'win32'\nif IS_WINDOWS:\n import msvcrt # noqa: F401\nelse:\n import fcntl # noqa: F401\n\nSOURCE_WORKFLOW_NAME = 'stream_alert_denoise'\nTRIAGE_WORKFLOW_NAME = 'stream_alert_triage'\nDEFAULT_BATCH_MAX_RECORDS = 10\nDEFAULT_BATCH_MAX_BYTES = 32 * 1024 * 1024\n_READ_CHUNK_BYTES = 64 * 1024\n_CURSOR_ANCHOR_BYTES = 4 * 1024\n_CURSOR_VERSION = 2\n_FILENAME_RE = re.compile(r'^dedup_result_(\\d+)\\.jsonl$')\n\n\ndef _flocks_root():\n return Config().get_global().data_dir.parent\n\n\ndef _dedup_root():\n return _flocks_root() / 'workspace' / 'workflows' / SOURCE_WORKFLOW_NAME\n\n\ndef _cursor_path():\n return _flocks_root() / 'workspace' / 'workflows' / TRIAGE_WORKFLOW_NAME / '.triage_cursor.json'\n\n\n\ndef _batch_lease_path():\n state_dir = _flocks_root() / 'workspace' / 'workflows' / TRIAGE_WORKFLOW_NAME\n state_dir.mkdir(parents=True, exist_ok=True)\n return os.fspath(state_dir / '.triage_batch.lock')\n\n\ndef _acquire_batch_lease():\n path = _batch_lease_path()\n fd = os.open(path, os.O_RDWR | os.O_CREAT, 0o600)\n try:\n if IS_WINDOWS:\n if os.fstat(fd).st_size == 0:\n os.write(fd, b'0')\n os.fsync(fd)\n os.lseek(fd, 0, os.SEEK_SET)\n msvcrt.locking(fd, msvcrt.LK_NBLCK, 1)\n else:\n fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)\n except OSError as exc:\n os.close(fd)\n raise RuntimeError('production_batch_lease_busy') from exc\n\n try:\n token = f'{os.getpid()}-{time.time_ns()}'\n os.ftruncate(fd, 0)\n os.lseek(fd, 0, os.SEEK_SET)\n os.write(fd, token.encode('ascii'))\n os.fsync(fd)\n except BaseException:\n _release_batch_lease(fd)\n raise\n print(f'[load] acquired production batch lease token={token}')\n return fd, token\n\n\ndef _release_batch_lease(fd):\n if type(fd) is not int or fd < 0:\n return\n try:\n if IS_WINDOWS:\n try:\n os.lseek(fd, 0, os.SEEK_SET)\n msvcrt.locking(fd, msvcrt.LK_UNLCK, 1)\n except OSError:\n pass\n else:\n fcntl.flock(fd, fcntl.LOCK_UN)\n finally:\n os.close(fd)\n\ndef _date_str(input_date):\n if input_date:\n value = str(input_date).strip()\n if re.fullmatch(r'\\d{4}-\\d{2}-\\d{2}', value):\n return value\n return datetime.datetime.now().strftime('%Y-%m-%d')\n\n\ndef _positive_int(value, default):\n try:\n if isinstance(value, bool):\n raise ValueError\n if isinstance(value, float) and not value.is_integer():\n raise ValueError\n if isinstance(value, str) and not re.fullmatch(r'\\+?\\d+', value.strip()):\n raise ValueError\n parsed = int(value)\n if parsed <= 0:\n raise ValueError\n return parsed\n except (TypeError, ValueError, OverflowError):\n return default\n\n\ndef _file_seq(path, fallback):\n match = _FILENAME_RE.fullmatch(os.path.basename(path))\n return int(match.group(1)) if match else fallback\n\n\ndef _explicit_paths(input_paths, input_path):\n requested = []\n if input_paths:\n values = [input_paths] if isinstance(input_paths, str) else input_paths\n for value in values:\n if value:\n requested.append(os.path.expanduser(str(value)))\n if input_path:\n requested.append(os.path.expanduser(str(input_path)))\n\n seen = set()\n files = []\n missing = 0\n for path in requested:\n if path in seen:\n continue\n seen.add(path)\n if not os.path.isfile(path):\n missing += 1\n print(f'[load] WARNING: explicit file not found, skipping: {path}')\n continue\n files.append({'path': path, 'seq': _file_seq(path, len(files) + 1)})\n return files, len(seen), missing\n\n\ndef _auto_paths(date_str):\n day_dir = _dedup_root() / date_str\n files = []\n invalid_names = 0\n try:\n with os.scandir(day_dir) as entries:\n for entry in entries:\n if (\n not entry.is_file()\n or not entry.name.startswith('dedup_result_')\n or not entry.name.endswith('.jsonl')\n ):\n continue\n match = _FILENAME_RE.fullmatch(entry.name)\n if not match:\n invalid_names += 1\n continue\n files.append({'path': entry.path, 'seq': int(match.group(1))})\n except FileNotFoundError:\n pass\n except OSError as exc:\n print(f'[load] WARNING: failed to scan {day_dir}: {exc}')\n files.sort(key=lambda item: item['seq'])\n print(f'[load] auto-discovered date={date_str} dir={day_dir} files={len(files)} invalid={invalid_names}')\n return files, invalid_names\n\n\ndef _cursor_file_state(path, byte_offset, stream=None):\n owns_stream = stream is None\n if owns_stream:\n stream = open(path, 'rb')\n current_position = stream.tell()\n try:\n stat_result = os.fstat(stream.fileno())\n head_size = min(byte_offset, _CURSOR_ANCHOR_BYTES)\n stream.seek(0)\n head = stream.read(head_size)\n boundary_start = max(0, byte_offset - _CURSOR_ANCHOR_BYTES)\n stream.seek(boundary_start)\n boundary = stream.read(byte_offset - boundary_start)\n if len(head) != head_size or len(boundary) != byte_offset - boundary_start:\n raise OSError('cursor boundary is beyond the current file size')\n return {\n 'device_id': int(stat_result.st_dev),\n 'file_id': int(stat_result.st_ino),\n 'head_hash': hashlib.sha256(head).hexdigest(),\n 'boundary_start': boundary_start,\n 'boundary_hash': hashlib.sha256(boundary).hexdigest(),\n }\n finally:\n if owns_stream:\n stream.close()\n else:\n stream.seek(current_position)\n\n\ndef _cursor_matches_file(item, cursor, stream=None):\n try:\n state = _cursor_file_state(item['path'], cursor['byte_offset'], stream=stream)\n except OSError as exc:\n print(f'[load] WARNING: cursor file validation failed for {item[\"path\"]}: {exc}')\n return False\n for key in ('device_id', 'file_id', 'head_hash', 'boundary_start', 'boundary_hash'):\n if state[key] != cursor.get(key):\n print(f'[load] WARNING: cursor invalidated by file identity/content change: '\n f'{item[\"path\"]} field={key}')\n return False\n return True\n\n\ndef _cursor_revision_bytes(raw_bytes):\n return hashlib.sha256(raw_bytes).hexdigest()\n\n\n\ndef _normalise_cursor(raw_cursor, expected_date=None, explicit=False):\n if not isinstance(raw_cursor, dict):\n return None\n version = raw_cursor.get('version')\n cursor_date = raw_cursor.get('date')\n file_name = raw_cursor.get('file_name')\n file_seq = raw_cursor.get('file_seq')\n byte_offset = raw_cursor.get('byte_offset')\n device_id = raw_cursor.get('device_id')\n file_id = raw_cursor.get('file_id')\n head_hash = raw_cursor.get('head_hash')\n boundary_start = raw_cursor.get('boundary_start')\n boundary_hash = raw_cursor.get('boundary_hash')\n if type(version) is not int or version != _CURSOR_VERSION:\n return None\n if not isinstance(cursor_date, str) or not isinstance(file_name, str):\n return None\n integer_fields = (file_seq, byte_offset, device_id, file_id, boundary_start)\n if any(type(value) is not int for value in integer_fields):\n return None\n if expected_date and cursor_date != expected_date:\n return None\n if (\n not file_name\n or file_seq < 0\n or byte_offset < 0\n or device_id < 0\n or file_id < 0\n or boundary_start < 0\n or boundary_start > byte_offset\n or not isinstance(head_hash, str)\n or not re.fullmatch(r'[0-9a-f]{64}', head_hash)\n or not isinstance(boundary_hash, str)\n or not re.fullmatch(r'[0-9a-f]{64}', boundary_hash)\n ):\n return None\n\n cursor = {\n 'version': _CURSOR_VERSION,\n 'date': cursor_date,\n 'file_seq': file_seq,\n 'file_name': file_name,\n 'byte_offset': byte_offset,\n 'device_id': device_id,\n 'file_id': file_id,\n 'head_hash': head_hash,\n 'boundary_start': boundary_start,\n 'boundary_hash': boundary_hash,\n }\n updated_at = raw_cursor.get('updated_at')\n if isinstance(updated_at, str) and updated_at:\n cursor['updated_at'] = updated_at\n if raw_cursor.get('skipping_oversized_line') is True:\n cursor['skipping_oversized_line'] = True\n if explicit:\n cursor_path = raw_cursor.get('path')\n if isinstance(cursor_path, str) and cursor_path:\n cursor['path'] = os.path.expanduser(cursor_path)\n file_index = raw_cursor.get('file_index')\n if type(file_index) is int and file_index >= 0:\n cursor['file_index'] = file_index\n return cursor\n\n\n\ndef _load_production_cursor(path, date_str):\n try:\n with open(path, 'rb') as stream:\n raw_bytes = stream.read()\n except FileNotFoundError:\n return None, None\n except Exception as exc:\n print(f'[load] WARNING: failed to read cursor {path}: {exc}; restarting date from first file')\n return None, None\n\n revision = _cursor_revision_bytes(raw_bytes)\n try:\n raw_cursor = json.loads(raw_bytes.decode('utf-8'))\n except Exception as exc:\n print(f'[load] WARNING: invalid cursor {path}: {exc}; restarting date from first file')\n return None, revision\n\n cursor = _normalise_cursor(raw_cursor, expected_date=date_str)\n if cursor is None:\n print(f'[load] WARNING: stale, legacy, or malformed cursor {path}; '\n 'restarting date from first file')\n return cursor, revision\n\n\n\ndef _start_position(files, cursor, explicit):\n if not files or cursor is None:\n return 0, 0, False\n\n index = None\n if explicit:\n cursor_path = cursor.get('path')\n if cursor_path:\n index = next((i for i, item in enumerate(files) if item['path'] == cursor_path), None)\n if index is None:\n requested_index = cursor.get('file_index')\n if isinstance(requested_index, int) and 0 <= requested_index < len(files):\n if os.path.basename(files[requested_index]['path']) == cursor['file_name']:\n index = requested_index\n if index is None:\n index = next((\n i for i, item in enumerate(files)\n if os.path.basename(item['path']) == cursor['file_name']\n ), None)\n if index is None:\n print('[load] WARNING: replay cursor file is absent; restarting explicit paths from first file')\n return 0, 0, True\n else:\n index = next((\n i for i, item in enumerate(files)\n if os.path.basename(item['path']) == cursor['file_name']\n ), None)\n if index is None:\n index = next((i for i, item in enumerate(files) if item['seq'] > cursor['file_seq']), len(files))\n return index, 0, False\n\n if not _cursor_matches_file(files[index], cursor):\n return index, 0, True\n return index, cursor['byte_offset'], False\n\n\n\ndef _cursor_after(\n item,\n file_index,\n date_str,\n byte_offset,\n explicit,\n skipping_oversized_line=False,\n stream=None,\n):\n cursor = {\n 'version': _CURSOR_VERSION,\n 'date': date_str,\n 'file_seq': item['seq'],\n 'file_name': os.path.basename(item['path']),\n 'byte_offset': byte_offset,\n **_cursor_file_state(item['path'], byte_offset, stream=stream),\n }\n if skipping_oversized_line:\n cursor['skipping_oversized_line'] = True\n if explicit:\n cursor['file_index'] = file_index\n cursor['path'] = item['path']\n return cursor\n\n\ndef _read_bounded_line(stream, budget):\n chunks = []\n scanned = 0\n while scanned < budget:\n chunk = stream.readline(min(_READ_CHUNK_BYTES, budget - scanned))\n if not chunk:\n return ('eof' if scanned == 0 else 'partial'), None, scanned\n chunks.append(chunk)\n scanned += len(chunk)\n if chunk.endswith(b'\\n'):\n return 'line', b''.join(chunks), scanned\n return 'limit', None, scanned\n\n\ndef _discard_to_newline(stream, budget):\n discarded = 0\n while discarded < budget:\n chunk = stream.readline(min(_READ_CHUNK_BYTES, budget - discarded))\n if not chunk:\n return 'eof', discarded\n discarded += len(chunk)\n if chunk.endswith(b'\\n'):\n return 'line', discarded\n return 'limit', discarded\n\n\ndef _cursor_position(files, cursor):\n if cursor is None:\n return None\n cursor_path = cursor.get('path')\n if cursor_path:\n index = next((i for i, item in enumerate(files) if item['path'] == cursor_path), None)\n else:\n index = None\n if index is None:\n index = next((\n i for i, item in enumerate(files)\n if os.path.basename(item['path']) == cursor.get('file_name')\n ), None)\n if index is None:\n return None\n return index, int(cursor.get('byte_offset', 0))\n\n\ndef _has_more(files, start_index, start_offset, effective_cursor):\n position = _cursor_position(files, effective_cursor)\n if position is None:\n index, offset = start_index, start_offset\n else:\n index, offset = position\n for current_index in range(index, len(files)):\n current_offset = offset if current_index == index else 0\n try:\n if os.path.getsize(files[current_index]['path']) > current_offset:\n return True\n except OSError:\n return True\n return False\n\n\ninput_paths = inputs.get('input_paths')\ninput_path = inputs.get('input_path')\ninput_date = _date_str(inputs.get('input_date'))\nexplicit = bool(input_path) or bool(input_paths)\n_batch_lease_fd = None\n_batch_lease_token = ''\nif not explicit:\n _batch_lease_fd, _batch_lease_token = _acquire_batch_lease()\ntry:\n batch_max_records = _positive_int(inputs.get('batch_max_records'), DEFAULT_BATCH_MAX_RECORDS)\n batch_max_bytes = _positive_int(inputs.get('batch_max_bytes'), DEFAULT_BATCH_MAX_BYTES)\n\n invalid_file_names = 0\n missing_files = 0\n cursor_revision = None\n requested_file_count = 0\n if explicit:\n files, requested_file_count, missing_files = _explicit_paths(input_paths, input_path)\n cursor_before = _normalise_cursor(inputs.get('resume_cursor'), explicit=True)\n else:\n files, invalid_file_names = _auto_paths(input_date)\n requested_file_count = len(files)\n cursor_before, cursor_revision = _load_production_cursor(_cursor_path(), input_date)\n\n start_index, start_offset, cursor_reset = _start_position(files, cursor_before, explicit)\n if not explicit and cursor_before is None and cursor_revision is not None:\n cursor_reset = True\n enriched_alerts = []\n loaded_files = []\n file_stats = []\n pending_cursor = None\n if cursor_reset and start_index < len(files):\n pending_cursor = _cursor_after(files[start_index], start_index, input_date, start_offset, explicit)\n resume_oversized_position = None if cursor_reset else _cursor_position(files, cursor_before)\n batch_bytes = 0\n oversized_bytes_discarded = 0\n stop_batch = False\n stats = {\n 'candidate_file_count': len(files),\n 'requested_file_count': requested_file_count,\n 'touched_file_count': 0,\n 'record_count': 0,\n 'bytes_read': 0,\n 'header_skipped': 0,\n 'empty_lines': 0,\n 'bad_lines': 0,\n 'non_object_lines': 0,\n 'oversized_lines': 0,\n 'oversized_bytes_discarded': 0,\n 'partial_lines': 0,\n 'invalid_file_names': invalid_file_names,\n 'missing_files': missing_files,\n 'cursor_invalidated': cursor_reset,\n 'has_more': False,\n 'files': file_stats,\n }\n\n for file_index in range(start_index, len(files)):\n item = files[file_index]\n path = item['path']\n offset = start_offset if file_index == start_index else 0\n current = {\n 'path': path,\n 'records': 0,\n 'bytes_read': 0,\n 'headers': 0,\n 'empty_lines': 0,\n 'bad_lines': 0,\n 'non_object_lines': 0,\n 'oversized_lines': 0,\n 'partial_lines': 0,\n }\n file_stats.append(current)\n loaded_files.append(path)\n resume_oversized = (\n file_index == start_index\n and cursor_before is not None\n and cursor_before.get('skipping_oversized_line') is True\n and resume_oversized_position == (file_index, offset)\n )\n try:\n with open(path, 'rb') as stream:\n if (\n file_index == start_index\n and offset > 0\n and cursor_before is not None\n and not cursor_reset\n and not _cursor_matches_file(item, cursor_before, stream=stream)\n ):\n offset = 0\n cursor_reset = True\n stats['cursor_invalidated'] = True\n pending_cursor = _cursor_after(\n item, file_index, input_date, 0, explicit, stream=stream\n )\n resume_oversized = False\n stream.seek(offset)\n if resume_oversized:\n remaining = batch_max_bytes - batch_bytes\n discard_status, discarded = _discard_to_newline(stream, remaining)\n batch_bytes += discarded\n current['bytes_read'] += discarded\n oversized_bytes_discarded += discarded\n if discarded:\n pending_cursor = _cursor_after(\n item,\n file_index,\n input_date,\n stream.tell(),\n explicit,\n skipping_oversized_line=discard_status != 'line',\n stream=stream,\n )\n if discard_status != 'line' or batch_bytes >= batch_max_bytes:\n stop_batch = True\n\n while (\n not stop_batch\n and len(enriched_alerts) < batch_max_records\n and batch_bytes < batch_max_bytes\n ):\n line_start = stream.tell()\n remaining = batch_max_bytes - batch_bytes\n status, line, scanned = _read_bounded_line(stream, remaining)\n\n if status == 'eof':\n break\n if status == 'partial':\n batch_bytes += scanned\n current['bytes_read'] += scanned\n current['partial_lines'] += 1\n stats['partial_lines'] += 1\n stream.seek(line_start)\n stop_batch = True\n break\n if status == 'limit':\n batch_bytes += scanned\n current['bytes_read'] += scanned\n if batch_bytes > scanned:\n stream.seek(line_start)\n stop_batch = True\n break\n\n oversized_bytes_discarded += scanned\n current['oversized_lines'] += 1\n stats['oversized_lines'] += 1\n pending_cursor = _cursor_after(\n item,\n file_index,\n input_date,\n stream.tell(),\n explicit,\n skipping_oversized_line=True,\n stream=stream,\n )\n stop_batch = True\n break\n\n batch_bytes += scanned\n current['bytes_read'] += scanned\n pending_cursor = _cursor_after(\n item, file_index, input_date, stream.tell(), explicit, stream=stream\n )\n payload = line[:-1]\n if not payload.strip():\n current['empty_lines'] += 1\n stats['empty_lines'] += 1\n else:\n try:\n obj = json.loads(payload.decode('utf-8'))\n except Exception:\n current['bad_lines'] += 1\n stats['bad_lines'] += 1\n else:\n if not isinstance(obj, dict):\n current['non_object_lines'] += 1\n stats['non_object_lines'] += 1\n elif obj.get('_type') == 'file_header':\n current['headers'] += 1\n stats['header_skipped'] += 1\n else:\n enriched_alerts.append(obj)\n current['records'] += 1\n\n if len(enriched_alerts) >= batch_max_records or batch_bytes >= batch_max_bytes:\n stop_batch = True\n break\n except Exception as exc:\n current['error'] = str(exc)\n print(f'[load] WARNING: failed to read {path!r}: {exc}')\n stop_batch = True\n\n if stop_batch:\n break\n\n effective_cursor = pending_cursor or cursor_before\n has_more = _has_more(files, start_index, start_offset, effective_cursor)\n stats['touched_file_count'] = len(loaded_files)\n stats['file_count'] = len(loaded_files)\n stats['record_count'] = len(enriched_alerts)\n stats['bytes_read'] = batch_bytes\n stats['oversized_bytes_discarded'] = oversized_bytes_discarded\n stats['has_more'] = has_more\n\n print(f'[load] DONE candidates={len(files)} touched={len(loaded_files)} '\n f'records={len(enriched_alerts)} bytes={batch_bytes} has_more={has_more}')\n\n outputs['enriched_alerts'] = enriched_alerts\n outputs['loaded_files'] = loaded_files\n outputs['load_stats'] = stats\n outputs['concurrency'] = _positive_int(inputs.get('concurrency'), 1)\n outputs['max_triage_cache_size'] = _positive_int(inputs.get('max_triage_cache_size'), 100000)\n outputs['input_date'] = input_date\n outputs['cursor_enabled'] = not explicit\n outputs['cursor_before'] = cursor_before\n outputs['cursor_revision'] = cursor_revision\n outputs['cursor_invalidated'] = cursor_reset\n outputs['pending_cursor'] = pending_cursor\n outputs['next_cursor'] = pending_cursor or cursor_before\n outputs['has_more'] = has_more\n outputs['batch_records'] = len(enriched_alerts)\n outputs['batch_bytes'] = batch_bytes\n outputs['_triage_persistence_succeeded'] = False\n outputs['_run_id'] = inputs.get('_run_id')\n outputs['triage_output_mode'] = inputs.get('triage_output_mode')\n outputs['persist_triage_output'] = inputs.get('persist_triage_output')\n outputs['soc_db_path'] = inputs.get('soc_db_path')\n outputs['jsonl_output_dir'] = inputs.get('jsonl_output_dir')\n\n outputs['_batch_lease_fd'] = _batch_lease_fd\n outputs['batch_lease_token'] = _batch_lease_token\n outputs['_triage_state_dir'] = os.fspath(\n _flocks_root() / 'workspace' / 'workflows' / TRIAGE_WORKFLOW_NAME\n )\nexcept BaseException:\n _release_batch_lease(_batch_lease_fd)\n raise\n" }, { "id": "concurrent_triage", "type": "python", - "description": "Leader/follower 分组并发研判节点（自包含，内联 tdp_alert_triage 逻辑）。先按 dedup_key 把 alerts 分组：每组只对 leader 研判，follower 复用 leader 结果。外层 ThreadPoolExecutor(concurrency) 处理 unique work units（concurrency 取值 1–5，默认 1），单条告警仍执行 survey / cve_related / cve_info / payload_analysis 4 个分支，但所有 llm.ask() 共享运行级 concurrency 预算，总 LLM 峰值不超过 1–5。dedup_key 在 triage_cache.pkl 命中时直接复用历史 verdict/title/triage_report；未命中则 leader 执行完整研判（情报查询 + 4 个 LLM + attack_analysis + verdict + title + 聚合 markdown），完整研判 markdown 仅写入 alert 的 `triage_report` 字段，**不生成任何独立报告文件**。新结果合并写回 cache（FIFO LRU + 文件锁 + 原子落盘）。SOC DB 只接受明确 `is_duplicate=false`、包含 `dedup_key` 且批内首次出现的告警，并通过数据库唯一索引保证跨执行全局唯一；重复 key 只更新研判字段并保留首次事件元数据，持久化失败会使工作流失败。可通过工作流目录 `config.json` 或运行输入将 `triage_output_mode` 切换为 `jsonl` / `both` / `none`，保留 `triage_result_NNN.jsonl` 可选输出。", - "code": "\"\"\"\nconcurrent_triage: leader/follower 分组并发研判 + dedup_key 缓存复用（自包含）。\n\n去重模式：\n 1. 输入 alerts 先按 dedup_key 分组 → unique dedup_keys 列表\n 2. 每个 group 只对 leader（首条）做研判；followers 复用 leader 结果，不重复调 LLM\n 3. 无 dedup_key 的 alert 各自独立成 work unit（防御性研判，无法复用）\n\n并发结构：\n 外层 ThreadPoolExecutor(max_workers=concurrency)：处理 unique work unit\n （concurrency 取值 1–5，默认 1，由 inputs.concurrency 控制）\n 单条 alert 仍执行 4 个 LLM 分支\n （survey / cve_related / cve_info / payload_analysis — 保留 tdp_alert_triage\n 的 4 分支研判结构）\n 所有 llm.ask 调用共享运行级信号量，稳态 LLM 峰值不超过 concurrency\n\n研判产物只以字段形式附加到每条 alert（attack_verdict / risk_level /\nreport_title / triage_report ...），**不生成任何独立的 per-alert markdown 报告\n文件**，避免冗余落盘与跨日期路径失效。\n\ndedup_key 缓存（与 stream_alert_denoise 的 LSH 状态文件同根目录，逻辑独立）：\n ~/.flocks/workspace/workflows/stream_alert_triage/triage_cache.pkl\n - cache 命中：直接复用历史 verdict/title/triage_report，**不调用 LLM**\n - cache 未命中：leader 执行完整内联研判（情报 + 4 个 LLM 分支 + verdict + title + report）；\n follower 直接广播 leader 结果\n - 新结果合并写回 cache，FIFO LRU 淘汰，文件锁 + 原子落盘\n\"\"\"\n\nimport ipaddress\nimport json\nimport os\nimport pickle\nimport re\nimport sys\nimport threading\nimport time\nfrom concurrent.futures import ThreadPoolExecutor, as_completed\n\nIS_WINDOWS = sys.platform == 'win32'\nif IS_WINDOWS:\n import msvcrt # noqa: F401\nelse:\n import fcntl # noqa: F401\n\nWORKFLOW_NAME = 'stream_alert_triage'\n\n# Per-call LLM timeout and retry budget for every analysis branch in this\n# node. Workflow LLM calls share the dedicated ``flocks-workflow-llm-loop``;\n# a single hung call (e.g. provider 504, slow TLS handshake) without a\n# timeout would otherwise pin one of the (already concurrency-limited)\n# worker threads for up to httpx's DEFAULT read timeout (10 min), serially\n# blocking the rest of the alert pipeline. 120s + 1 retry covers normal\n# slow-but-alive responses while still recovering from transient hangs.\nLLM_CALL_TIMEOUT_S = 120.0\nLLM_CALL_MAX_RETRIES = 1\n\n# The user-facing concurrency setting is a run-wide LLM request budget. The\n# four logical branches still start together, but nested executors must not\n# multiply the actual provider load (5 outer workers used to become 20 calls).\n_llm_slots = None\n\nTRIAGE_FIELDS = (\n 'attack_verdict',\n 'risk_level',\n 'report_title',\n 'triage_report',\n 'attack_success',\n)\nVERDICT_LABELS = ('attack_success', 'attack_failed', 'attack', 'unknown', 'benign')\nVERDICT_RISK = {\n 'attack_success': 'High',\n 'attack_failed': 'Medium',\n 'attack': 'Medium',\n 'unknown': 'Medium',\n 'benign': 'Low',\n}\nVERDICT_CN = {\n 'attack_success': '攻击成功',\n 'attack_failed': '攻击失败',\n 'attack': '攻击',\n 'unknown': '未知',\n 'benign': '安全',\n}\nTRIAGE_REPORT_VERSION = 'soc.triage.markdown.v1'\nTRIAGE_REPORT_TAGS = (\n 'report_title',\n 'report_meta',\n 'analysis_steps',\n 'triage_conclusion',\n 'attack_payload',\n 'payload_explanation',\n 'response_evidence',\n 'key_evidence',\n 'disposal_recommendation',\n)\n\n\n# ── Cache persistence ─────────────────────────────────────────────────────────\n\ndef _cache_paths():\n from flocks.config import Config\n flocks_root = Config().get_global().data_dir.parent\n state_dir = flocks_root / 'workspace' / 'workflows' / WORKFLOW_NAME\n state_dir.mkdir(parents=True, exist_ok=True)\n return str(state_dir / 'triage_cache.pkl'), str(state_dir / 'triage_cache.lock')\n\n\ndef _acquire_lock(lock_path):\n fh = open(lock_path, 'w+')\n try:\n if IS_WINDOWS:\n fh.write('L'); fh.flush(); fh.seek(0)\n while True:\n try:\n msvcrt.locking(fh.fileno(), msvcrt.LK_LOCK, 1); break\n except OSError:\n continue\n else:\n fcntl.flock(fh.fileno(), fcntl.LOCK_EX)\n except BaseException:\n try:\n fh.close()\n except Exception:\n pass\n raise\n return fh\n\n\ndef _release_lock(fh):\n try:\n if IS_WINDOWS:\n try:\n fh.seek(0); msvcrt.locking(fh.fileno(), msvcrt.LK_UNLCK, 1)\n except OSError:\n pass\n else:\n fcntl.flock(fh.fileno(), fcntl.LOCK_UN)\n finally:\n fh.close()\n\n\ndef _load_cache(cache_path):\n if not os.path.exists(cache_path) or os.path.getsize(cache_path) == 0:\n return {}\n try:\n with open(cache_path, 'rb') as f:\n c = pickle.load(f)\n if not isinstance(c, dict):\n return {}\n print(f'[triage_cache] loaded {len(c)} entries from {cache_path}')\n return c\n except Exception as e:\n print(f'[triage_cache] WARNING: failed to load ({e}), starting fresh')\n return {}\n\n\ndef _save_cache_atomic(cache_path, cache):\n tmp = cache_path + '.tmp'\n try:\n with open(tmp, 'wb') as f:\n pickle.dump(cache, f); f.flush(); os.fsync(f.fileno())\n os.replace(tmp, cache_path)\n print(f'[triage_cache] saved {len(cache)} entries -> {cache_path}')\n except Exception as e:\n print(f'[triage_cache] WARNING: failed to save: {e}')\n if os.path.exists(tmp):\n try: os.remove(tmp)\n except Exception: pass\n\n\ndef _evict_lru(cache, max_keys):\n excess = len(cache) - max_keys\n if excess > 0:\n for k in list(cache.keys())[:excess]:\n del cache[k]\n return excess\n return 0\n\n\n# ── Runtime output config and persistence targets ──────────────────────────────\n#\n# Defaults are read from ~/.flocks/plugins/workflows/stream_alert_triage/config.json.\n# Runtime inputs override config values. The default mode is soc_db so SOC pages\n# read the same DB-backed dataset. JSONL remains available via config/input:\n# triage_output_mode = soc_db | jsonl | both | none\n# persist_triage_output = true (legacy alias that adds JSONL to soc_db)\n\nimport datetime as _datetime\n\nMAX_RECORDS_PER_FILE = 10000\n_TRIAGE_JSONL_PREFIX = 'triage_result'\n_TRIAGE_COUNTER_FILE = '.triage_counter.json'\n_WORKFLOW_CONFIG_PATH = os.path.expanduser('~/.flocks/plugins/workflows/stream_alert_triage/config.json')\n_DEFAULT_SOC_DB_PATH = os.path.expanduser('~/.flocks/data/soc.db')\n\n\ndef _load_workflow_config():\n try:\n with open(_WORKFLOW_CONFIG_PATH, 'r', encoding='utf-8') as f:\n cfg = json.load(f)\n if isinstance(cfg, dict):\n return cfg\n except FileNotFoundError:\n pass\n except Exception as e:\n print(f'[triage_config] WARNING: failed to read {_WORKFLOW_CONFIG_PATH}: {e}')\n return {}\n\n\ndef _configured_value(config, key, default=None):\n if key in inputs:\n value = inputs.get(key)\n if value is not None and not (isinstance(value, str) and not value.strip()):\n return value\n return config.get(key, default)\n\n\ndef _input_bool(value, default=False):\n if value is None:\n return default\n if isinstance(value, bool):\n return value\n if isinstance(value, (int, float)):\n return bool(value)\n text = str(value).strip().lower()\n if text in {'1', 'true', 'yes', 'y', 'on'}:\n return True\n if text in {'0', 'false', 'no', 'n', 'off'}:\n return False\n return default\n\n\ndef _select_first_seen_soc_alerts(alerts):\n selected = []\n seen_dedup_keys = set()\n stats = {\n 'input_rows': len(alerts),\n 'first_seen_rows': 0,\n 'skipped_not_first_seen_rows': 0,\n 'skipped_missing_dedup_key_rows': 0,\n 'skipped_repeated_dedup_key_rows': 0,\n }\n for alert in alerts:\n if not isinstance(alert, dict) or _input_bool(alert.get('is_duplicate'), True):\n stats['skipped_not_first_seen_rows'] += 1\n continue\n dedup_key = str(alert.get('dedup_key') or '').strip()\n if not dedup_key:\n stats['skipped_missing_dedup_key_rows'] += 1\n continue\n if dedup_key in seen_dedup_keys:\n stats['skipped_repeated_dedup_key_rows'] += 1\n continue\n seen_dedup_keys.add(dedup_key)\n selected.append(alert)\n stats['first_seen_rows'] = len(selected)\n return selected, stats\n\n\ndef _resolve_output_config():\n config = _load_workflow_config()\n raw_mode = str(_configured_value(config, 'triage_output_mode', 'soc_db') or 'soc_db').strip().lower()\n mode_alias = {\n 'db': 'soc_db',\n 'sqlite': 'soc_db',\n 'sqlite_db': 'soc_db',\n 'soc': 'soc_db',\n 'json': 'jsonl',\n 'file': 'jsonl',\n 'files': 'jsonl',\n 'off': 'none',\n 'disabled': 'none',\n }\n requested_mode = mode_alias.get(raw_mode, raw_mode)\n if requested_mode not in {'soc_db', 'jsonl', 'both', 'none'}:\n print(f'[triage_config] WARNING: invalid triage_output_mode={raw_mode!r}; using soc_db')\n requested_mode = 'soc_db'\n\n legacy_jsonl = _input_bool(_configured_value(config, 'persist_triage_output', False), False)\n write_soc_db = requested_mode in {'soc_db', 'both'}\n write_jsonl = requested_mode in {'jsonl', 'both'}\n effective_mode = requested_mode\n if requested_mode == 'soc_db' and legacy_jsonl:\n write_jsonl = True\n effective_mode = 'both'\n if requested_mode == 'none':\n write_soc_db = False\n write_jsonl = False\n effective_mode = 'none'\n\n soc_db_path = os.path.expanduser(str(\n _configured_value(config, 'soc_db_path', _DEFAULT_SOC_DB_PATH) or _DEFAULT_SOC_DB_PATH\n ))\n jsonl_output_dir = _configured_value(config, 'jsonl_output_dir', '') or ''\n jsonl_output_dir = os.path.expanduser(str(jsonl_output_dir)) if jsonl_output_dir else ''\n return {\n 'config_path': _WORKFLOW_CONFIG_PATH,\n 'requested_mode': requested_mode,\n 'mode': effective_mode,\n 'write_soc_db': write_soc_db,\n 'write_jsonl': write_jsonl,\n 'soc_db_path': soc_db_path,\n 'jsonl_output_dir': jsonl_output_dir,\n }\n\n\n# ── Persisted JSONL output (optional; mirrors stream_alert_denoise layout) ─────\n#\n# Directory : ~/.flocks/workspace/workflows/stream_alert_triage//\n# Filename : triage_result_NNN.jsonl (3-digit zero-padded seq)\n# Layout : line 1 = {\"_type\":\"file_header\", ...}, subsequent lines = one\n# enriched_with_triage alert per line.\n# Counter : .triage_counter.json sidecar tracks (seq, count) so we don't\n# rescan every existing file on each run; auto-rolls over to a\n# new file when reaching MAX_RECORDS_PER_FILE.\n\n\ndef _triage_output_dir(configured_dir=''):\n \"\"\"Return output directory for triage_result_*.jsonl.\"\"\"\n if configured_dir:\n out_dir = configured_dir\n os.makedirs(out_dir, exist_ok=True)\n return out_dir\n from flocks.config import Config\n flocks_root = Config().get_global().data_dir.parent\n date_str = _datetime.datetime.now().strftime('%Y-%m-%d')\n out_dir = flocks_root / 'workspace' / 'workflows' / WORKFLOW_NAME / date_str\n out_dir.mkdir(parents=True, exist_ok=True)\n return str(out_dir)\n\n\ndef _triage_get_counter(out_dir):\n path = os.path.join(out_dir, _TRIAGE_COUNTER_FILE)\n try:\n with open(path, 'r', encoding='utf-8') as f:\n d = json.load(f)\n return int(d.get('seq', 0)), int(d.get('count', 0))\n except Exception:\n return 0, 0\n\n\ndef _triage_set_counter(out_dir, seq, count):\n path = os.path.join(out_dir, _TRIAGE_COUNTER_FILE)\n tmp = path + '.tmp'\n try:\n with open(tmp, 'w', encoding='utf-8') as f:\n json.dump({'seq': seq, 'count': count}, f)\n os.replace(tmp, path)\n except Exception:\n pass\n\n\ndef _triage_find_active_file(out_dir):\n \"\"\"Locate the active (latest, not-yet-full) jsonl file; create if none.\"\"\"\n seq, count = _triage_get_counter(out_dir)\n if seq > 0:\n path = os.path.join(out_dir, f'{_TRIAGE_JSONL_PREFIX}_{seq:03d}.jsonl')\n if os.path.exists(path):\n return path, count, seq\n import glob as _glob\n existing = sorted(_glob.glob(os.path.join(out_dir, _TRIAGE_JSONL_PREFIX + '_*.jsonl')))\n if not existing:\n return None, 0, 0\n latest = existing[-1]\n try:\n seq = int(os.path.basename(latest).replace(_TRIAGE_JSONL_PREFIX + '_', '').replace('.jsonl', ''))\n except ValueError:\n seq = len(existing)\n count = 0\n try:\n with open(latest, 'r', encoding='utf-8') as f:\n for line in f:\n if line.strip() and '\"_type\"' not in line:\n count += 1\n except Exception:\n pass\n return latest, count, seq\n\n\ndef _triage_write_jsonl(out_dir, alerts, run_id, run_stats):\n \"\"\"Append all alerts to today's triage_result_NNN.jsonl, rolling over at\n MAX_RECORDS_PER_FILE. Returns the list of files that were written to.\"\"\"\n now = _datetime.datetime.now()\n written = []\n active_path, active_count, seq = _triage_find_active_file(out_dir)\n remaining = list(alerts)\n while remaining:\n available = MAX_RECORDS_PER_FILE - active_count\n if available <= 0 or active_path is None:\n seq += 1\n active_path = os.path.join(out_dir, f'{_TRIAGE_JSONL_PREFIX}_{seq:03d}.jsonl')\n active_count = 0\n available = MAX_RECORDS_PER_FILE\n header = {\n '_type': 'file_header',\n 'created_at': now.isoformat(),\n 'date': now.strftime('%Y-%m-%d'),\n 'workflow': WORKFLOW_NAME,\n 'seq': seq,\n 'run_id': run_id,\n 'batch_total': run_stats.get('total'),\n 'batch_triaged': run_stats.get('triaged'),\n 'batch_followers_reused':run_stats.get('followers_reused'),\n 'batch_cache_hit': run_stats.get('cache_hit'),\n 'batch_triage_failed': run_stats.get('triage_failed'),\n }\n with open(active_path, 'w', encoding='utf-8') as hf:\n hf.write(json.dumps(header, ensure_ascii=False) + '\\n')\n batch = remaining[:available]\n remaining = remaining[available:]\n with open(active_path, 'a', encoding='utf-8') as af:\n for alert in batch:\n af.write(json.dumps(alert, ensure_ascii=False) + '\\n')\n active_count += len(batch)\n if active_path not in written:\n written.append(active_path)\n if remaining:\n active_path = None\n active_count = 0\n if written:\n _triage_set_counter(out_dir, seq, active_count)\n return written\n\n\n# ── SOC DB output (default) ───────────────────────────────────────────────────\n\ndef _ensure_soc_db_schema(conn):\n conn.execute(\"\"\"\n CREATE TABLE IF NOT EXISTS alert_records (\n row_id TEXT PRIMARY KEY,\n record_id TEXT,\n asset_date TEXT NOT NULL,\n source_file TEXT NOT NULL,\n line_number INTEGER NOT NULL,\n event_time INTEGER,\n source_type TEXT,\n threat_name TEXT,\n dedup_key TEXT,\n is_duplicate INTEGER NOT NULL DEFAULT 0,\n record_json TEXT NOT NULL\n )\n \"\"\")\n columns = {row[1] for row in conn.execute('PRAGMA table_info(alert_records)')}\n dedup_key_added = 'dedup_key' not in columns\n if dedup_key_added:\n conn.execute('ALTER TABLE alert_records ADD COLUMN dedup_key TEXT')\n\n unique_index_name = 'idx_alert_records_first_seen_dedup_key'\n indexes = list(conn.execute('PRAGMA index_list(alert_records)'))\n unique_index_ready = any(row[1] == unique_index_name and bool(row[2]) for row in indexes)\n conn.execute('CREATE INDEX IF NOT EXISTS idx_alert_records_duplicate ON alert_records(is_duplicate)')\n has_persisted_duplicates = conn.execute(\"\"\"\n SELECT 1 FROM alert_records\n WHERE is_duplicate = 1\n AND dedup_key IS NOT NULL\n AND dedup_key <> ''\n LIMIT 1\n \"\"\").fetchone() is not None\n if not unique_index_ready or has_persisted_duplicates:\n if any(row[1] == unique_index_name for row in indexes):\n conn.execute(f'DROP INDEX {unique_index_name}')\n conn.execute(\"\"\"\n UPDATE alert_records\n SET dedup_key = CASE\n WHEN json_valid(record_json)\n THEN NULLIF(TRIM(CAST(json_extract(record_json, '$.dedup_key') AS TEXT)), '')\n ELSE NULL\n END\n WHERE dedup_key IS NULL OR TRIM(dedup_key) = ''\n \"\"\")\n conn.execute(\"\"\"\n UPDATE alert_records\n SET dedup_key = NULLIF(TRIM(dedup_key), '')\n WHERE dedup_key IS NOT NULL\n \"\"\")\n conn.execute(\"\"\"\n DELETE FROM alert_records\n WHERE dedup_key IS NOT NULL\n AND dedup_key <> ''\n AND rowid NOT IN (\n SELECT MIN(rowid)\n FROM alert_records\n WHERE dedup_key IS NOT NULL AND dedup_key <> ''\n AND is_duplicate = 0\n GROUP BY dedup_key\n )\n \"\"\")\n conn.execute(f\"\"\"\n CREATE UNIQUE INDEX {unique_index_name}\n ON alert_records(dedup_key)\n WHERE dedup_key IS NOT NULL AND dedup_key <> ''\n \"\"\")\n\n conn.execute('CREATE INDEX IF NOT EXISTS idx_alert_records_asset_date ON alert_records(asset_date)')\n conn.execute('CREATE INDEX IF NOT EXISTS idx_alert_records_event_time ON alert_records(event_time)')\n conn.execute('CREATE INDEX IF NOT EXISTS idx_alert_records_source_type ON alert_records(source_type)')\n conn.execute('CREATE INDEX IF NOT EXISTS idx_alert_records_threat_name ON alert_records(threat_name)')\n\n\ndef _event_time_value(alert):\n for key in ('time', 'event_time', 'timestamp', 'timestamp_real', 'occur_time', 'created_at'):\n value = alert.get(key)\n if value in (None, ''):\n continue\n if isinstance(value, (int, float)):\n ts = float(value)\n if ts > 100000000000:\n ts = ts / 1000.0\n return int(ts)\n text = str(value).strip()\n if not text:\n continue\n try:\n ts = float(text)\n if ts > 100000000000:\n ts = ts / 1000.0\n return int(ts)\n except Exception:\n pass\n normalized = text.replace('Z', '+00:00')\n try:\n return int(_datetime.datetime.fromisoformat(normalized).timestamp())\n except Exception:\n pass\n for fmt in ('%Y-%m-%d %H:%M:%S', '%Y/%m/%d %H:%M:%S', '%Y-%m-%d %H:%M', '%Y/%m/%d %H:%M'):\n try:\n return int(_datetime.datetime.strptime(text, fmt).timestamp())\n except Exception:\n continue\n return int(time.time())\n\n\ndef _asset_date_value(alert, event_time):\n value = alert.get('asset_date') or alert.get('_asset_date') or alert.get('date')\n if value:\n text = str(value).strip()\n if re.match(r'^\\d{4}-\\d{2}-\\d{2}$', text):\n return text\n try:\n return _datetime.datetime.fromtimestamp(int(event_time)).strftime('%Y-%m-%d')\n except Exception:\n return _datetime.datetime.now().strftime('%Y-%m-%d')\n\n\ndef _source_type_value(alert):\n for key in ('source_type', '_source_type', 'data_source', 'log_type', 'vendor', 'device_type'):\n value = alert.get(key)\n if value not in (None, ''):\n return str(value)\n return ''\n\n\ndef _record_id_value(alert):\n for key in ('record_id', 'id', 'uuid', 'event_id', 'dedup_key'):\n value = alert.get(key)\n if value not in (None, ''):\n return str(value)\n return ''\n\n\ndef _stable_row_id(alert, source_file, line_number, event_time):\n existing = alert.get('row_id') or alert.get('_row_id')\n if existing:\n return str(existing)\n import hashlib as _hashlib\n basis = {\n 'record_id': _record_id_value(alert),\n 'dedup_key': alert.get('dedup_key', ''),\n 'time': event_time,\n 'source_file': source_file,\n 'line_number': line_number,\n 'sip': alert.get('sip', ''),\n 'sport': alert.get('sport', ''),\n 'dip': alert.get('dip', ''),\n 'dport': alert.get('dport', ''),\n 'threat_rule_id': alert.get('threat_rule_id') or alert.get('rule_id') or '',\n }\n raw = json.dumps(basis, sort_keys=True, ensure_ascii=False)\n return _hashlib.sha256(raw.encode('utf-8')).hexdigest()\n\n\ndef _load_existing_soc_rows(conn, dedup_keys):\n existing = {}\n unique_keys = list(dict.fromkeys(str(key).strip() for key in dedup_keys if str(key).strip()))\n for start in range(0, len(unique_keys), 500):\n chunk = unique_keys[start:start + 500]\n placeholders = ','.join('?' for _ in chunk)\n rows = conn.execute(f\"\"\"\n SELECT row_id, record_id, asset_date, source_file, line_number,\n event_time, source_type, threat_name, dedup_key,\n is_duplicate, record_json\n FROM alert_records\n WHERE dedup_key IN ({placeholders})\n \"\"\", chunk)\n for row in rows:\n try:\n record = json.loads(row[10])\n except Exception:\n record = {}\n if not isinstance(record, dict):\n record = {}\n existing[row[8]] = {\n 'row_id': row[0],\n 'record_id': row[1],\n 'asset_date': row[2],\n 'source_file': row[3],\n 'line_number': row[4],\n 'event_time': row[5],\n 'source_type': row[6],\n 'threat_name': row[7],\n 'record': record,\n }\n return existing\n\n\ndef _merge_triage_record(existing_record, incoming_record):\n merged = dict(existing_record) if isinstance(existing_record, dict) else {}\n triage_fields = (\n 'has_dedup_key',\n 'triage_source',\n 'triage_status',\n 'attack_verdict',\n 'risk_level',\n 'report_title',\n 'triage_report',\n 'attack_success',\n 'triage_ms',\n 'triage_error',\n '_triage_run_id',\n '_triage_persisted_at',\n )\n for key in triage_fields:\n if key in incoming_record:\n merged[key] = incoming_record[key]\n elif key in {'triage_ms', 'triage_error'}:\n merged.pop(key, None)\n return merged\n\n\ndef _triage_write_soc_db(db_path, alerts, run_id):\n import sqlite3\n\n db_dir = os.path.dirname(db_path)\n if db_dir:\n os.makedirs(db_dir, exist_ok=True)\n default_source_file = ''\n loaded_files = inputs.get('loaded_files') or []\n if isinstance(loaded_files, list) and len(loaded_files) == 1:\n default_source_file = str(loaded_files[0])\n\n persisted_at = _datetime.datetime.now().isoformat()\n candidates = []\n seen_dedup_keys = set()\n for idx, alert in enumerate(alerts, 1):\n if not isinstance(alert, dict):\n continue\n dedup_key = str(alert.get('dedup_key') or '').strip()\n if not dedup_key or dedup_key in seen_dedup_keys:\n continue\n seen_dedup_keys.add(dedup_key)\n record = dict(alert)\n record['dedup_key'] = dedup_key\n record['is_duplicate'] = False\n record['_triage_run_id'] = run_id\n record['_triage_persisted_at'] = persisted_at\n source_file = str(\n record.get('source_file')\n or record.get('_source_file')\n or record.get('file_path')\n or default_source_file\n or 'stream_alert_triage'\n )\n try:\n line_number = int(record.get('line_number') or record.get('_line_number') or idx)\n except Exception:\n line_number = idx\n event_time = _event_time_value(record)\n candidates.append({\n 'row_id': _stable_row_id(record, source_file, line_number, event_time),\n 'record_id': _record_id_value(record),\n 'asset_date': _asset_date_value(record, event_time),\n 'source_file': source_file,\n 'line_number': line_number,\n 'event_time': event_time,\n 'source_type': _source_type_value(record),\n 'threat_name': str(record.get('threat_name') or record.get('rule_name') or ''),\n 'dedup_key': dedup_key,\n 'record': record,\n })\n\n insert_rows = []\n update_rows = []\n with sqlite3.connect(db_path, timeout=30) as conn:\n conn.execute('BEGIN IMMEDIATE')\n _ensure_soc_db_schema(conn)\n existing_by_key = _load_existing_soc_rows(\n conn, [candidate['dedup_key'] for candidate in candidates],\n )\n for candidate in candidates:\n existing = existing_by_key.get(candidate['dedup_key'])\n if existing:\n merged_record = _merge_triage_record(existing['record'], candidate['record'])\n merged_record['dedup_key'] = candidate['dedup_key']\n merged_record['is_duplicate'] = False\n update_rows.append((\n candidate['dedup_key'],\n json.dumps(merged_record, ensure_ascii=False),\n existing['row_id'],\n ))\n continue\n insert_rows.append((\n candidate['row_id'],\n candidate['record_id'],\n candidate['asset_date'],\n candidate['source_file'],\n candidate['line_number'],\n candidate['event_time'],\n candidate['source_type'],\n candidate['threat_name'],\n candidate['dedup_key'],\n 0,\n json.dumps(candidate['record'], ensure_ascii=False),\n ))\n\n if insert_rows:\n conn.executemany(\"\"\"\n INSERT INTO alert_records (\n row_id, record_id, asset_date, source_file, line_number,\n event_time, source_type, threat_name, dedup_key,\n is_duplicate, record_json\n ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)\n \"\"\", insert_rows)\n if update_rows:\n conn.executemany(\"\"\"\n UPDATE alert_records\n SET dedup_key = ?, is_duplicate = 0, record_json = ?\n WHERE row_id = ?\n \"\"\", update_rows)\n conn.commit()\n\n persisted_rows = len(insert_rows) + len(update_rows)\n return {\n 'path': db_path,\n 'table': 'alert_records',\n 'rows': persisted_rows,\n 'inserted_rows': len(insert_rows),\n 'updated_rows': len(update_rows),\n }\n\n\n# ── LLM provider warm-up (avoid cold-start race in _parallel_4_branches) ──────\n#\n# Background: when this node starts, the LLM provider (e.g. threatbook-cn-llm)\n# is lazy-initialized on the first call. Inside `_parallel_4_branches` we\n# submit 4 LLM calls to a ThreadPoolExecutor simultaneously; whichever one\n# wins the race may race against provider registration and fail with\n# \"provider 'xxx' not exists\" while subsequent calls succeed. A single\n# synchronous warm-up call before any concurrent fan-out forces the provider\n# to finish registering on the main thread, eliminating the race entirely.\n#\n# Failure of the warm-up is non-fatal — we just log and continue. The first\n# real LLM call will still see the same error and surface it normally.\n\ndef _warmup_llm():\n \"\"\"Force LLM provider lazy-init on the main thread before any fan-out.\"\"\"\n t0 = time.time()\n try:\n _ask_llm('ping')\n print(f'[triage] LLM provider warm-up OK in {round((time.time()-t0)*1000)}ms')\n return True\n except Exception as e:\n print(f'[triage] WARNING: LLM warm-up failed ({type(e).__name__}: '\n f'{str(e)[:200]}); proceeding anyway')\n return False\n\n\n# ── Inline triage helpers (mirroring tdp_alert_triage docs version) ───────────\n\ndef _strip_think(text):\n return re.sub(r'[\\s\\S]*?', '', str(text or ''), flags=re.IGNORECASE).strip()\n\n\ndef _is_public_ip(value):\n try:\n ip_obj = ipaddress.ip_address(value)\n except Exception:\n return False\n return not (ip_obj.is_private or ip_obj.is_loopback or ip_obj.is_reserved\n or ip_obj.is_link_local or ip_obj.is_multicast or ip_obj.is_unspecified)\n\n\ndef _pick(*values):\n for v in values:\n if v not in (None, '', [], {}):\n return v\n return ''\n\n\ndef _parse_alert(alert_input):\n \"\"\"Mirrors tdp_alert_triage.receive_alert. Supports three input shapes:\n nested TDP (net.http.url), flat TDP (net_http_url), normalized (req_http_url).\n \"\"\"\n if isinstance(alert_input, str):\n try:\n alert_input = json.loads(alert_input)\n except Exception:\n alert_input = {}\n if isinstance(alert_input, list):\n alert_data = alert_input[0] if alert_input else {}\n elif isinstance(alert_input, dict) and isinstance(alert_input.get('data'), list):\n alert_data = alert_input.get('data', [])[0] if alert_input.get('data') else {}\n else:\n alert_data = alert_input if isinstance(alert_input, dict) else {}\n\n net = alert_data.get('net', {}) or {}\n http = net.get('http', {}) or {}\n threat = alert_data.get('threat', {}) or {}\n assets = alert_data.get('assets', {}) or {}\n\n src_ip = _pick(\n alert_data.get('attacker'), alert_data.get('external_ip'),\n net.get('src_ip'), net.get('flow_src_ip'),\n alert_data.get('net_real_src_ip'),\n alert_data.get('sip'), alert_data.get('src_ip'), alert_data.get('src'),\n )\n dst_ip = _pick(\n alert_data.get('victim'), alert_data.get('machine'),\n alert_data.get('server_ip'), net.get('dest_ip'), net.get('flow_dest_ip'),\n alert_data.get('net_dest_ip'),\n alert_data.get('dip'), alert_data.get('dst_ip'), alert_data.get('dst'),\n )\n src_port = _pick(\n net.get('src_port'), net.get('flow_src_port'),\n alert_data.get('external_port'), alert_data.get('net_src_port'),\n alert_data.get('sport'), alert_data.get('src_port'), 0,\n )\n dst_port = _pick(\n net.get('dest_port'), net.get('flow_dest_port'),\n alert_data.get('server_port'), alert_data.get('machine_port'),\n alert_data.get('net_dest_port'),\n alert_data.get('dport'), alert_data.get('dst_port'), 0,\n )\n protocol = _pick(\n net.get('app_proto'), net.get('type'), net.get('proto'),\n alert_data.get('net_app_proto'), alert_data.get('protocol'),\n alert_data.get('event_type'), 'TCP',\n )\n alert_type = _pick(\n threat.get('name'), alert_data.get('threat_name'),\n alert_data.get('vuln_name'),\n alert_data.get('alert_type'), threat.get('topic'),\n alert_data.get('type'), 'unknown',\n )\n severity = _pick(\n threat.get('severity'), alert_data.get('threat_severity'),\n alert_data.get('severity'), threat.get('level'),\n alert_data.get('level'), 'medium',\n )\n\n req_line = _pick(http.get('reqs_line'), alert_data.get('req_line'), alert_data.get('net_http_reqs_line'))\n req_header = _pick(http.get('reqs_header'), alert_data.get('req_header'), alert_data.get('net_http_reqs_header'))\n req_body = _pick(http.get('req_body'), alert_data.get('req_body'), alert_data.get('net_http_reqs_body'))\n resp_line = _pick(http.get('resp_line'), alert_data.get('rsp_line'), alert_data.get('resp_line'),\n alert_data.get('net_http_resp_line'))\n resp_header = _pick(http.get('resp_header'), alert_data.get('rsp_header'), alert_data.get('resp_header'),\n alert_data.get('net_http_resp_header'))\n resp_body = _pick(http.get('resp_body'), alert_data.get('rsp_body'), alert_data.get('resp_body'),\n alert_data.get('net_http_resp_body'))\n status = _pick(http.get('status'), alert_data.get('http_status'),\n alert_data.get('net_http_status'), alert_data.get('rsp_status_code'), 0)\n\n host = _pick(http.get('reqs_host'), alert_data.get('url_host'), http.get('domain'),\n alert_data.get('req_host'), alert_data.get('net_http_reqs_host'), dst_ip)\n raw_url = _pick(http.get('raw_url'), http.get('url'),\n alert_data.get('url_path'),\n alert_data.get('net_http_url'), alert_data.get('req_http_url'),\n alert_data.get('uri'))\n url = ''\n if host and raw_url:\n scheme = 'https' if net.get('is_https') else 'http'\n url = raw_url if str(raw_url).startswith(('http://', 'https://')) else f'{scheme}://{host}{raw_url}'\n elif raw_url and str(raw_url).startswith(('http://', 'https://')):\n url = raw_url\n elif raw_url:\n url = raw_url\n\n payload = f'请求行: {req_line}\\n请求头: {req_header}\\n请求体: {req_body}'\n response = f'状态行: {resp_line}\\n响应头: {resp_header}\\n响应体: {resp_body}'\n threat_result = _pick(threat.get('result'), alert_data.get('threat_result'))\n threat_msg = _pick(threat.get('msg'), alert_data.get('threat_msg'))\n\n log_text = (\n f'[告警基本信息]\\n'\n f'告警类型: {alert_type}\\n严重级别: {severity}\\n'\n f'源地址: {src_ip}:{src_port}\\n目的地址: {dst_ip}:{dst_port}\\n'\n f'协议: {protocol}\\nURL: {url}\\nHTTP状态码: {status}\\n'\n f'TDP判定: {threat_result}\\nTDP消息: {threat_msg}\\n\\n'\n f'[HTTP请求内容]\\n{payload}\\n\\n'\n f'[HTTP响应内容]\\n{response}'\n )\n\n vuln_text = '\\n'.join(str(item) for item in [\n threat_msg, threat.get('topic', ''),\n alert_data.get('data', ''), url,\n json.dumps(threat.get('tag', []), ensure_ascii=False),\n ] if item)\n vuln_matches = sorted(set(re.findall(r'\\b(?:CVE|CNVD|CNNVD|XVE)-[A-Za-z0-9._-]+\\b', vuln_text, flags=re.I)))\n\n iocs = []\n for candidate in [src_ip, dst_ip]:\n if candidate:\n iocs.append({'type': 'ip', 'value': candidate})\n if url:\n iocs.append({'type': 'url', 'value': url})\n if host and not re.match(r'^\\d{1,3}(?:\\.\\d{1,3}){3}(?::\\d+)?$', str(host)):\n iocs.append({'type': 'domain', 'value': str(host).split(':')[0]})\n\n return {\n 'src_ip': src_ip, 'dst_ip': dst_ip, 'src_port': src_port, 'dst_port': dst_port,\n 'protocol': protocol, 'payload': payload, 'response': response,\n 'url': url, 'status': status,\n 'alert_type': alert_type, 'severity': severity,\n 'vuln_id': vuln_matches[0] if vuln_matches else '',\n 'vuln_candidates': vuln_matches,\n 'threat_result': threat_result, 'threat_msg': threat_msg,\n 'failed_by': threat.get('failed_by', []),\n 'asset_ip': assets.get('ip', ''), 'asset_name': assets.get('name', []),\n 'iocs': iocs, 'log_text': log_text,\n }\n\n\ndef _prepare_intel(parsed):\n \"\"\"Mirrors tdp_alert_triage.prepare_intel. Pre-fetches IP/domain/URL threat intel\n and CVE info so the parallel LLM tasks have concrete context to consume.\n \"\"\"\n iocs = parsed.get('iocs', [])\n intel_results = []\n seen = set()\n for ioc in iocs:\n ioc_type = ioc.get('type', '')\n ioc_value = str(ioc.get('value', '')).strip()\n key = (ioc_type, ioc_value)\n if not ioc_value or key in seen:\n continue\n seen.add(key)\n if ioc_type == 'ip':\n if not _is_public_ip(ioc_value):\n continue\n r = tool.run_safe('threatbook_ip_query', ip=ioc_value)\n if r['success']:\n intel_results.append({'source': 'threatbook', 'type': 'ip',\n 'value': ioc_value, 'result': r['text']})\n elif ioc_type == 'domain':\n r = tool.run_safe('threatbook_domain_query', domain=ioc_value)\n if r['success']:\n intel_results.append({'source': 'threatbook', 'type': 'domain',\n 'value': ioc_value, 'result': r['text']})\n elif ioc_type == 'url':\n r = tool.run_safe('threatbook_url_query', url=ioc_value)\n if r['success']:\n intel_results.append({'source': 'threatbook', 'type': 'url',\n 'value': ioc_value, 'result': r['text']})\n\n vuln_info = {}\n vuln_id = parsed.get('vuln_id', '')\n if vuln_id:\n r = tool.run_safe('__mcp_vuln_query', vuln_id=vuln_id)\n if r['success']:\n try:\n obj = r.get('obj')\n if isinstance(obj, str):\n obj = json.loads(obj)\n vuln_info = obj if isinstance(obj, dict) else {'raw_result': r.get('text', '')}\n except Exception:\n vuln_info = {'raw_result': r.get('text', '')}\n\n intel_content = '\\n'.join(\n f\"[{i['source']}/{i['type']}] {i['value']}\\n{i['result']}\" for i in intel_results\n ) or '（无可用情报数据）'\n vuln_content = (\n json.dumps(vuln_info, ensure_ascii=False, indent=2)\n if vuln_info else '（无可用漏洞情报数据）'\n )\n return intel_results, intel_content, vuln_info, vuln_content\n\n\n# ── 4 LLM analysis branches (sharing one run-wide request budget) ───────────\n\ndef _ask_llm(prompt):\n \"\"\"Wrap ``llm.ask`` with the workflow-wide timeout + retry budget.\n\n Centralizing this avoids a hung provider request (no TCP timeout from\n upstream) from blocking a worker thread indefinitely. Any call that does\n not need bespoke parameters should go through here.\n \"\"\"\n kwargs = {\n 'timeout_s': LLM_CALL_TIMEOUT_S,\n 'max_retries': LLM_CALL_MAX_RETRIES,\n }\n if _llm_slots is None:\n return llm.ask(prompt, **kwargs)\n with _llm_slots:\n return llm.ask(prompt, **kwargs)\n\n\ndef _llm_survey(log_text, intel_content):\n prompt = f'''你是一个专业的Web日志分析专家。请总结以下IP的情报数据中的空间测绘信息。\n1. 如果该IP没有测绘信息，则不列出。\n2. 如果IP有测绘信息，则以简短的语言对该IP的测绘信息进行总结，关键说明ip的标签和测绘信息显示有哪些服务或者应用资产。\n3. 多个IP的测绘信息以无序列表显示，每个ip数据描述占一行数据。\n4. 不需要生成其他额外的补充信息。\n\n## 情报参考信息\n{intel_content}\n\n## 用户的原始输入日志\n{log_text}\n'''\n return _strip_think(_ask_llm(prompt))\n\n\ndef _llm_cve_related(log_text):\n prompt = f'''请从以下的日志数据中提取漏洞编号。\n要求：\n1. 仅从日志文本中识别漏洞编号，不要做任何推测。\n2. 如果日志中存在漏洞编号，则用简短语言描述，如：\"日志中存在漏洞编号：CVE-****-****\"。\n3. 如果日志中不存在漏洞编号，则输出：\"日志中无关联漏洞情报\"。\n\n日志数据如下：\n{log_text}\n'''\n return _strip_think(_ask_llm(prompt))\n\n\ndef _llm_cve_info(log_text, vuln_content):\n prompt = f'''你是一个专业的Web日志分析专家。参考情报信息中的漏洞数据，简要说明关联的CVE漏洞信息。\n1. 不要输出任何解释说明，只输出漏洞基本信息。不需要生成漏洞的处置建议或修复措施等。\n\n## 情报参考信息\n{vuln_content}\n\n## 用户的原始输入日志\n{log_text}\n'''\n return _strip_think(_ask_llm(prompt))\n\n\ndef _llm_payload_analysis(log_text):\n prompt = f'''你是一个专业的Web日志分析专家。根据用户输入的日志进行攻击负载分析。\n1. 首先分析日志中是否包含攻击负载，并给出判定依据。\n2. 不要进行攻击意图分析、攻击影响分析。\n3. 用简短的语言在一段话中进行描述。\n\n## 用户的原始输入日志:\n{log_text}\n'''\n return _strip_think(_ask_llm(prompt))\n\n\ndef _parallel_4_branches(parsed, intel_content, vuln_content):\n \"\"\"Submit 4 logical branches; _ask_llm enforces the run-wide provider budget.\"\"\"\n log_text = parsed.get('log_text', '')\n with ThreadPoolExecutor(max_workers=4, thread_name_prefix='triage_branch') as pool:\n futs = {\n 'survey_result': pool.submit(_llm_survey, log_text, intel_content),\n 'cve_related_result': pool.submit(_llm_cve_related, log_text),\n 'cve_info_result': pool.submit(_llm_cve_info, log_text, vuln_content),\n 'payload_analysis_result': pool.submit(_llm_payload_analysis, log_text),\n }\n out = {}\n for name, fut in futs.items():\n try:\n out[name] = fut.result()\n except Exception as e:\n print(f'[triage] WARNING: branch {name} failed: {e}')\n out[name] = ''\n return out\n\n\n# ── Join-point LLM analyses (attack_analysis_result -> verdict -> title) ──────\n\ndef _llm_attack_analysis(log_text):\n prompt = f'''你是一名专业且经验丰富的网络安全分析师和Web日志分析专家，你对HTTP协议以及Web攻击有着深入的理解，并且你能够快速识别和应对各种网络威胁。你的任务是对提供的HTTP请求与响应内容进行详细的专业分析，并判断日志请求的攻击状态。\n\n请严格遵循以下指令进行思考和分析：\n1. 攻击状态只能从以下情况中选择一种：[\"攻击成功\", \"攻击失败\", \"攻击\", \"未知\", \"安全\"]。\n2. 从日志中提取出\"HTTP请求内容\"和\"HTTP响应内容\"。请注意，HTTP请求内容和HTTP响应内容是分开的，请不要混淆，有些日志中没有包含HTTP响应内容，请不要将HTTP请求内容和HTTP响应内容混淆。分析后请你记住哪些是HTTP请求内容，哪些是HTTP响应内容。\n3. 请检查HTTP响应状态码，2xx或者3xx状态码都代表本次HTTP请求成功，4xx或者5xx状态码大多数情况下都代表请求失败，只有在请求成功的情况下才能对攻击是否成功进行后续判断。\n\n各攻击状态的定义以及判定标准：\n1. 攻击成功：\n(1) 首先分析日志中是否含有清晰的\"HTTP响应内容\"，如果日志中没有\"HTTP响应内容\"，则肯定不属于攻击成功。\n(2) 如果日志中未提供\"HTTP响应内容\"，即使HTTP请求内容中包含攻击者预期的结果，也不能判定为攻击成功。\n(3) 从日志中提取出\"HTTP请求内容\"和\"HTTP响应内容\"。请深入分析\"HTTP响应内容\"，并判定其是否为\"HTTP请求内容\"攻击成功时的预期结果，这是判定攻击成功的强依据。请注意，HTTP响应码200仅表示网络连接成功，不代表攻击攻击成功。\n(4) 分析HTTP请求内容和HTTP响应内容，只有当HTTP响应内容中明确包含攻击载荷在目标机器上成功执行的证据，并且HTTP请求内容中包含攻击载荷的特征，则判定为\"攻击成功\"。\n(5) 请注意：攻击成功的判定必须包含HTTP响应内容。如果不包含HTTP响应内容，则肯定不属于攻击成功。\n(6) 请注意：如果不包含HTTP响应内容，即使HTTP请求内容是攻击，这也不属于攻击成功。\n2. 攻击失败：\n(1) 分析HTTP请求内容和HTTP响应内容，如果HTTP响应内容中明确包含攻击载荷在目标机器上执行失败或者被阻止的证据，并且HTTP请求内容中包含攻击载荷的特征，则判定为\"攻击失败\"。\n(2) 攻击失败的判定必须包含HTTP响应内容。如果不包含HTTP响应内容，则肯定不属于攻击失败。\n3. 攻击：\n(1) 在\"HTTP请求内容\"或\"HTTP响应内容\"中发现任何证明存在攻击意图的证据，即可判定为存在攻击行为。但如果不符合上述的攻击成功或者攻击失败的标准，则\"攻击状态\"为\"攻击\"。\n(2) 请注意：如果日志中只提供了\"HTTP请求内容\"，且没有提供\"HTTP响应内容\"，且HTTP的请求内容分析中是包含攻击行为的，则\"攻击状态\"为\"攻击\"。\n4. 未知：\n(1) 如果不能100%确定HTTP通信的攻击结果，那么请在\"攻击状态\"处给出\"未知\"。\n(2) 请注意：如果在你给的判定原因中存在\"可能\"等不确定词汇，都代表你不能对你的结论100%确定，那么请在\"攻击状态\"处给出\"未知\"。\n5. 安全：\n(1) 如果\"HTTP请求内容\"和\"HTTP响应内容\"中都没有任何攻击意图的证据，那么请在\"攻击状态\"处给出\"安全\"。\n\n## 日志内容\n{log_text}\n\n## 输出要求\n请按下列结构输出（中文）：\n1. 攻击状态: [攻击成功/攻击失败/攻击/未知/安全]\n2. 判定依据: 简要说明请求与响应的关键证据\n3. 详细分析: 不超过200字\n'''\n return _strip_think(_ask_llm(prompt))\n\n\ndef _llm_attack_verdict(attack_analysis_result):\n prompt = f'''你是一个专业的Web日志分析专家。请据参考信息，直接输出攻击判定类别：\nattack_success：表示攻击成功。\nattack_failed：表示攻击失败。\nattack：表示是日志内容是攻击。\nunknown：表示未知。\nbenign：是安全。\n不额外输出任何其他信息，包括解释、判定依据等。\n\n## 日志分析结果:\n{attack_analysis_result}\n'''\n raw = _strip_think(_ask_llm(prompt)).strip().lower()\n return next((v for v in VERDICT_LABELS if v in raw), 'unknown')\n\n\ndef _llm_report_title(alert_type, attack_verdict, attack_analysis_result):\n prompt = f'''你是一个专业的Web日志分析专家。请基于以下分析结果，生成一份不超过 30 字的中文报告标题。\n要求：\n1. 标题必须能体现\"攻击类型\"或\"攻击结果分析的结论\"。\n2. 不要带书名号、引号或其他标点。\n3. 只输出标题本身，不要任何解释或说明。\n\n## 攻击类型\n{alert_type}\n\n## 攻击判定\n{attack_verdict}\n\n## 攻击分析结果\n{attack_analysis_result}\n'''\n raw = _strip_think(_ask_llm(prompt)).strip()\n return raw.splitlines()[0].strip(' \"\\'《》[]【】') if raw else f'{alert_type} - {attack_verdict}'\n\n\ndef _clip_text(value, limit=3000):\n text = str(value or '').strip()\n if len(text) > limit:\n return text[:limit] + '\\n...（已截断）'\n return text or '未提供'\n\n\ndef _fence_text(value):\n text = _clip_text(value, 6000)\n return text.replace('```', '``\\\\u200b`')\n\n\ndef _extract_tagged_triage_report(text):\n text = _strip_think(text)\n m = re.search(r']*>[\\s\\S]*?', text, flags=re.I)\n return m.group(0).strip() if m else text.strip()\n\n\ndef _is_valid_triage_report(markdown):\n text = str(markdown or '')\n if not re.search(r']*version=[\"\\']soc\\.triage\\.markdown\\.v1[\"\\'][^>]*>', text, flags=re.I):\n return False\n if not re.search(r'', text, flags=re.I):\n return False\n for tag in TRIAGE_REPORT_TAGS:\n if not re.search(rf'<{tag}\\b[^>]*>', text, flags=re.I):\n return False\n if not re.search(rf'', text, flags=re.I):\n return False\n return True\n\n\ndef _is_current_triage_fields(fields):\n if not isinstance(fields, dict):\n return False\n return _is_valid_triage_report(fields.get('triage_report'))\n\n\ndef _format_intel_brief(intel_results):\n if not intel_results:\n return '未查询到外部威胁情报。'\n lines = []\n for intel in intel_results[:6]:\n lines.append(f\"- {intel.get('source', 'intel')} / {intel.get('type', 'ioc')}: {intel.get('value', '')} => {_clip_text(intel.get('result'), 500)}\")\n return '\\n'.join(lines)\n\n\ndef _build_default_tagged_triage_report(parsed, intel_results, vuln_info, branches,\n attack_analysis_result, attack_verdict,\n report_title, risk_level):\n verdict_cn = VERDICT_CN.get(attack_verdict, attack_verdict)\n title = report_title or f'{parsed.get(\"alert_type\", \"Web日志告警\")} - {verdict_cn}'\n payload = _fence_text(parsed.get('payload', ''))\n response = _fence_text(parsed.get('response', ''))\n url = parsed.get('url') or '未提供'\n threat_msg = parsed.get('threat_msg') or '未提供'\n status = parsed.get('status') or '未提供'\n survey = _clip_text(branches.get('survey_result'), 1500)\n cve_related = _clip_text(branches.get('cve_related_result'), 1500)\n cve_info = _clip_text(branches.get('cve_info_result'), 1500)\n payload_analysis = _clip_text(branches.get('payload_analysis_result'), 1500)\n attack_analysis = _clip_text(attack_analysis_result, 1500)\n intel_brief = _format_intel_brief(intel_results)\n vuln_brief = _clip_text(json.dumps(vuln_info, ensure_ascii=False, indent=2), 1800) if vuln_info else '未查询到漏洞详情。'\n\n if attack_verdict == 'attack_success':\n recommendation = '立即核查目标资产是否产生异常文件、进程、账号或敏感数据访问记录，并按成功入侵事件升级处置。'\n elif attack_verdict == 'attack_failed':\n recommendation = '保留拦截与响应证据，复核同源后续请求，并确认防护策略是否持续生效。'\n elif attack_verdict == 'benign':\n recommendation = '作为低风险事件留痕，结合资产白名单或业务访问记录确认是否可降噪。'\n else:\n recommendation = '补齐目标 Web 日志、响应体、主机侧进程和文件证据后再确认攻击成功性。'\n\n return f'''\n\n\n# {title}\n\n\n\n- 研判结论：{verdict_cn}\n- 风险等级：{risk_level}\n- 告警类型：{parsed.get('alert_type', 'unknown')}\n- 源 IP：{parsed.get('src_ip', 'N/A')}:{parsed.get('src_port', 'N/A')}\n- 目标资产：{parsed.get('dst_ip', 'N/A')}:{parsed.get('dst_port', 'N/A')}\n- URL：{url}\n- 响应码：{status}\n\n\n\n## 分析步骤\n\n### 1. 日志类型分析\n该告警按 Web 日志处理，已提取 HTTP 请求、响应、源地址、目标资产、URL、响应码和 TDP 判定字段。\n\n### 2. 情报信息\n{intel_brief}\n\n### 3. 测绘信息\n{survey}\n\n### 4. 告警关联漏洞情报\n{cve_related}\n\n### 5. 攻击负载分析\n{payload_analysis}\n\n### 6. 攻击分析结果\n{attack_analysis}\n\n\n\n## 研判结论\n当前研判结论为 **{verdict_cn}**，风险等级为 **{risk_level}**。TDP 消息为：{threat_msg}\n\n\n\n## 攻击payload\n\n```http\n{payload}\n```\n\n\n\n## 具体含义解释\n\n1. 请求命中的告警类型为 {parsed.get('alert_type', 'unknown')}。\n2. 请求 URL 为 {url}，需要结合参数、请求体和目标业务判断攻击意图。\n3. Payload 分析结果：{payload_analysis}\n\n\n\n## 响应证据\n\n```http\n{response}\n```\n\n响应码为 {status}。如果响应体未提供或没有执行成功证据，则不能仅凭请求侧 payload 判定攻击成功。\n\n\n\n## 重要证据\n\n1. 源地址：{parsed.get('src_ip', 'N/A')}:{parsed.get('src_port', 'N/A')}。\n2. 目标资产：{parsed.get('dst_ip', 'N/A')}:{parsed.get('dst_port', 'N/A')}。\n3. TDP 判定：{parsed.get('threat_result', '未提供')}；TDP 消息：{threat_msg}。\n4. 漏洞详情：{vuln_brief}\n\n\n\n## 处置建议\n\n1. {recommendation}\n2. 检索同源 IP、同一 dedup_key、同一 URL 或同一漏洞特征的横向告警。\n3. 结合目标资产 Web 访问日志、主机审计、EDR 与 WAF 日志补齐证据链。\n\n\n'''\n\n\ndef _llm_triage_report_markdown(parsed, intel_results, vuln_info, branches,\n attack_analysis_result, attack_verdict,\n report_title, risk_level):\n verdict_cn = VERDICT_CN.get(attack_verdict, attack_verdict)\n context = json.dumps({\n 'report_title': report_title,\n 'attack_verdict': attack_verdict,\n 'verdict_cn': verdict_cn,\n 'risk_level': risk_level,\n 'alert': {\n 'alert_type': parsed.get('alert_type'),\n 'severity': parsed.get('severity'),\n 'src_ip': parsed.get('src_ip'),\n 'src_port': parsed.get('src_port'),\n 'dst_ip': parsed.get('dst_ip'),\n 'dst_port': parsed.get('dst_port'),\n 'url': parsed.get('url'),\n 'status': parsed.get('status'),\n 'threat_result': parsed.get('threat_result'),\n 'threat_msg': parsed.get('threat_msg'),\n 'payload': parsed.get('payload'),\n 'response': parsed.get('response'),\n },\n 'survey_result': branches.get('survey_result'),\n 'cve_related_result': branches.get('cve_related_result'),\n 'cve_info_result': branches.get('cve_info_result'),\n 'payload_analysis_result': branches.get('payload_analysis_result'),\n 'attack_analysis_result': attack_analysis_result,\n 'intel_results': intel_results,\n 'vuln_info': vuln_info,\n }, ensure_ascii=False, indent=2)\n\n prompt = f'''你是一名资深 SOC 告警研判分析师。请根据输入上下文，生成一份供前端直接渲染的 SOC 告警研判报告 markdown。\n\n硬性要求：\n1. 只输出带语义标签的 markdown，不要输出 JSON，不要解释规则。\n2. 根标签必须是 。\n3. 必须按顺序输出并完整闭合这些标签：\n 、、、、、、、、。\n4. 标签外不得输出正文内容。标签内可以使用 markdown 标题、列表、引用、代码块。\n5. 段落标题必须贴近前端展示模板：分析步骤、研判结论、攻击payload、具体含义解释、响应证据、重要证据、处置建议。\n6. 如果没有 HTTP 响应体或没有明确响应证据，不得判定为攻击成功；需要写明“当前日志未提供有效响应证据”。\n7. 不要编造输入中不存在的 IP、域名、URL、CVE、账号、文件路径或响应内容。\n8. 攻击 payload 和响应证据必须分别放在对应标签中，不要混淆请求与响应。\n\nFew-shot 示例 1：攻击成功\n\n\n\n# 敏感文件泄露攻击成功分析报告\n\n\n\n- 研判结论：攻击成功\n- 风险等级：High\n- 告警类型：敏感文件访问\n- 源 IP：203.0.113.10:42131\n- 目标资产：198.51.100.20:80\n- URL：http://example.com/api/.env\n- 响应码：200\n\n\n\n## 分析步骤\n\n### 1. 日志类型分析\n该日志包含 HTTP 请求路径、响应码和响应体，可用于判断敏感文件是否被返回。\n\n### 2. 情报信息\n源 IP 命中扫描源标签，风险高。\n\n### 3. 测绘信息\n目标为公网 Web 服务，存在敏感路径暴露风险。\n\n### 4. 告警关联漏洞情报\n该行为与环境变量文件泄露场景一致。\n\n### 5. 攻击负载分析\n攻击者直接请求 /api/.env，目标是读取环境变量配置。\n\n### 6. 攻击分析结果\n响应码为 200，响应体中出现 DB_PASSWORD，支持攻击成功。\n\n\n\n## 研判结论\n攻击者成功读取敏感配置文件，响应体中包含数据库密码字段，结论为攻击成功。\n\n\n\n## 攻击payload\n\n```http\nGET /api/.env HTTP/1.1\nHost: example.com\n```\n\n\n\n## 具体含义解释\n\n1. /api/.env 是常见环境变量文件路径。\n2. 攻击者通过 GET 请求尝试直接读取配置文件。\n3. 该路径若返回真实内容，通常意味着敏感文件暴露。\n\n\n\n## 响应证据\n\n```http\nHTTP/1.1 200 OK\n\nDB_PASSWORD=example-secret\n```\n\n响应体出现 DB_PASSWORD，证明敏感配置内容已被返回。\n\n\n\n## 重要证据\n\n1. 请求路径为 /api/.env。\n2. 响应码为 200。\n3. 响应体包含 DB_PASSWORD。\n\n\n\n## 处置建议\n\n1. 立即下线或限制敏感文件访问。\n2. 轮换可能泄露的密钥和数据库密码。\n3. 检索同源 IP 和同路径访问记录。\n\n\n\n\nFew-shot 示例 2：攻击失败\n\n\n\n# SQL注入攻击失败分析报告\n\n\n\n- 研判结论：攻击失败\n- 风险等级：Medium\n- 告警类型：SQL注入\n- 源 IP：203.0.113.44:51002\n- 目标资产：198.51.100.30:443\n- URL：https://shop.example.com/item?id=1\n- 响应码：403\n\n\n\n## 分析步骤\n\n### 1. 日志类型分析\n该日志包含请求参数和响应码，能够确认请求侧存在 SQL 注入尝试。\n\n### 2. 情报信息\n源 IP 暂无高置信恶意标签。\n\n### 3. 测绘信息\n目标为公网电商 Web 服务。\n\n### 4. 告警关联漏洞情报\n当前日志未提供可确认具体 CVE 的证据。\n\n### 5. 攻击负载分析\n请求参数中包含 union select，存在明显 SQL 注入意图。\n\n### 6. 攻击分析结果\n响应码为 403，响应体显示请求被阻断，不支持攻击成功。\n\n\n\n## 研判结论\n该请求存在 SQL 注入攻击意图，但响应显示被拒绝，当前判断为攻击失败。\n\n\n\n## 攻击payload\n\n```http\nGET /item?id=1 union select user HTTP/1.1\nHost: shop.example.com\n```\n\n\n\n## 具体含义解释\n\n1. union select 是典型 SQL 注入关键字组合。\n2. 攻击者尝试拼接查询以读取数据库用户信息。\n3. 该 payload 证明攻击意图，但不等同于成功执行。\n\n\n\n## 响应证据\n\n```http\nHTTP/1.1 403 Forbidden\n\nblocked by waf\n```\n\n响应状态和内容说明请求被拦截，未见数据泄露或执行成功证据。\n\n\n\n## 重要证据\n\n1. 请求参数包含 union select。\n2. 响应码为 403。\n3. 响应体显示 blocked by waf。\n\n\n\n## 处置建议\n\n1. 保留 WAF 拦截证据。\n2. 检查同源 IP 是否持续尝试其他注入 payload。\n3. 确认目标接口参数化查询和安全策略仍然有效。\n\n\n\n\n## 待研判上下文\n```json\n{context}\n```\n\n请输出最终报告：\n'''\n return _extract_tagged_triage_report(_ask_llm(prompt))\n\n\ndef _generate_triage_report(parsed, intel_results, vuln_info, branches,\n attack_analysis_result, attack_verdict, report_title):\n # Aggregate everything into tagged markdown for frontend rendering.\n # The markdown is returned through `triage_report` and is not written as a\n # per-alert file; leader/follower/cache-hit paths reuse the same field.\n verdict_cn = VERDICT_CN.get(attack_verdict, attack_verdict)\n risk_level = VERDICT_RISK.get(attack_verdict, 'Medium')\n\n if not report_title:\n report_title = f'{parsed.get(\"alert_type\", \"Web日志告警\")} - {verdict_cn}'\n\n try:\n triage_report = _llm_triage_report_markdown(\n parsed, intel_results, vuln_info, branches,\n attack_analysis_result, attack_verdict, report_title, risk_level,\n )\n except Exception as e:\n print(f'[triage] WARNING: triage_report LLM generation failed: {e}')\n triage_report = ''\n\n if not _is_valid_triage_report(triage_report):\n print('[triage] WARNING: triage_report missing required semantic tags; using deterministic fallback')\n triage_report = _build_default_tagged_triage_report(\n parsed, intel_results, vuln_info, branches,\n attack_analysis_result, attack_verdict, report_title, risk_level,\n )\n\n return triage_report, report_title, risk_level\n\n\ndef _triage_single_alert(alert):\n \"\"\"End-to-end inline triage for a single alert. Returns triage_fields dict.\n\n No file I/O — the full markdown report lives in the returned `triage_report` field\n and is broadcast to followers / persisted via `triage_cache.pkl`.\n \"\"\"\n parsed = _parse_alert(alert)\n intel_results, intel_content, vuln_info, vuln_content = _prepare_intel(parsed)\n\n branches = _parallel_4_branches(parsed, intel_content, vuln_content)\n\n attack_analysis_result = _llm_attack_analysis(parsed['log_text'])\n attack_verdict = _llm_attack_verdict(attack_analysis_result)\n report_title = _llm_report_title(parsed.get('alert_type', 'unknown'),\n attack_verdict, attack_analysis_result)\n\n triage_report, report_title, risk_level = _generate_triage_report(\n parsed, intel_results, vuln_info, branches,\n attack_analysis_result, attack_verdict, report_title,\n )\n\n return {\n 'attack_verdict': attack_verdict,\n 'risk_level': risk_level,\n 'report_title': report_title,\n 'triage_report': triage_report,\n 'attack_success': attack_verdict == 'attack_success',\n }\n\n\n# ── Main: leader/follower batch deduplication ────────────────────────────────\n# When the input batch contains multiple alerts sharing the same dedup_key\n# (e.g. upstream emits is_duplicate=True alerts in the same batch, or LSH\n# clustering produces several alerts per cluster), we only triage the LEADER\n# (first occurrence of each dedup_key). All FOLLOWERS in the same group reuse\n# the leader's triage result without invoking the LLM again.\n#\n# Work unit types:\n# ('dk', dedup_key, leader_idx) — group of 1+ alerts sharing dedup_key\n# ('nokey', None, alert_idx) — single alert with no dedup_key\n# (cannot be deduplicated, always triaged)\n\nenriched_alerts = list(inputs.get('enriched_alerts', []) or [])\nconcurrency = min(5, max(1, int(inputs.get('concurrency', 1))))\n_llm_slots = threading.BoundedSemaphore(concurrency)\nmax_triage_cache_size = int(inputs.get('max_triage_cache_size', 100000))\nif max_triage_cache_size < 1:\n max_triage_cache_size = 100000\n\n# Group by dedup_key\ngroups = {} # dedup_key -> [alert_index, ...]\nno_key_indices = [] # alerts with no dedup_key\nfor i, a in enumerate(enriched_alerts):\n dk = a.get('dedup_key', '') if isinstance(a, dict) else ''\n if dk:\n groups.setdefault(dk, []).append(i)\n else:\n no_key_indices.append(i)\n\nwork_units = (\n [('dk', dk, group_indices[0]) for dk, group_indices in groups.items()]\n + [('nokey', None, idx) for idx in no_key_indices]\n)\nfollower_count = sum(len(v) - 1 for v in groups.values())\n\nprint(f'[triage] alerts={len(enriched_alerts)} '\n f'→ {len(groups)} unique dedup_keys ({follower_count} followers) + '\n f'{len(no_key_indices)} no-key alerts '\n f'= {len(work_units)} work units; outer_concurrency={concurrency} '\n f'llm_concurrency_limit={concurrency} '\n f'(4 branches share the run-wide LLM budget)')\n\ncache_path, lock_path = _cache_paths()\nlock_fh = _acquire_lock(lock_path)\ntry:\n triage_cache_snapshot = _load_cache(cache_path)\nfinally:\n _release_lock(lock_fh)\n\n# Only warm up the LLM when at least one work unit may actually need it.\n# A unit \"may need\" the LLM if it's a no-key alert OR a dedup_key unit whose\n# entry is not in the cache snapshot. Pure cache-hit batches skip warm-up.\n_needs_llm = any(\n unit_type == 'nokey' or not _is_current_triage_fields(triage_cache_snapshot.get(dk))\n for unit_type, dk, _ in work_units\n)\nif _needs_llm:\n _warmup_llm()\n\nresults_lock = threading.Lock()\nnew_results = {} # dedup_key -> triage_fields (only for freshly computed leaders)\ngroup_outcomes = {} # dedup_key -> (triage_fields, source) for broadcasting to followers\nnokey_outcomes = {} # alert_idx -> (triage_fields, source)\nstats = {\n 'total': len(enriched_alerts),\n 'unique_dedup_keys': len(groups),\n 'followers_reused': follower_count,\n 'no_dedup_key_alerts': len(no_key_indices),\n 'work_units': len(work_units),\n 'llm_concurrency_limit': concurrency,\n 'cache_hit': 0,\n 'triaged': 0,\n 'triage_failed': 0,\n 'verdict_counts': {},\n 'cache_size_before': len(triage_cache_snapshot),\n 'cache_size_after': 0,\n 'evicted': 0,\n}\n\n\ndef _bump_verdict(verdict):\n with results_lock:\n stats['verdict_counts'][verdict] = stats['verdict_counts'].get(verdict, 0) + 1\n\n\n_UNKNOWN_TRIAGE = {\n 'attack_verdict': 'unknown',\n 'risk_level': 'Medium',\n 'report_title': '',\n 'triage_report': '',\n 'attack_success': False,\n}\n\n\ndef _process_unit(unit_type, dedup_key, leader_idx):\n \"\"\"Triage one unique work unit. Returns (key, triage_fields, source, ms, error).\"\"\"\n leader_alert = enriched_alerts[leader_idx]\n t0 = time.time()\n\n # 1) cache lookup for dedup_key units\n if unit_type == 'dk':\n cached = triage_cache_snapshot.get(dedup_key)\n if cached:\n if _is_current_triage_fields(cached):\n with results_lock:\n stats['cache_hit'] += 1\n _bump_verdict(cached.get('attack_verdict', 'unknown'))\n return dedup_key, cached, 'cache', 0, None\n print(f'[triage_cache] stale entry for dedup_key={dedup_key}: '\n 'missing current triage_report markdown; treating as cache miss')\n\n # 2) cache miss or no-key → run full inline triage on the leader\n try:\n triage_fields = _triage_single_alert(leader_alert)\n ms = round((time.time() - t0) * 1000)\n with results_lock:\n stats['triaged'] += 1\n if dedup_key:\n new_results[dedup_key] = triage_fields\n _bump_verdict(triage_fields.get('attack_verdict', 'unknown'))\n source = 'triaged' if unit_type == 'dk' else 'no_dedup_key_triaged'\n return (dedup_key if unit_type == 'dk' else leader_idx), triage_fields, source, ms, None\n except Exception as e:\n import traceback\n ms = round((time.time() - t0) * 1000)\n with results_lock:\n stats['triage_failed'] += 1\n _bump_verdict('unknown')\n err = str(e)[:500]\n print(f'[triage] leader_idx={leader_idx} FAILED: {e}\\n{traceback.format_exc()}')\n source = 'failed' if unit_type == 'dk' else 'no_dedup_key_failed'\n return (dedup_key if unit_type == 'dk' else leader_idx), dict(_UNKNOWN_TRIAGE), source, ms, err\n\n\nt_start = time.time()\nunit_completions = {} # key -> (triage_fields, source, ms, error)\nif work_units:\n with ThreadPoolExecutor(max_workers=concurrency, thread_name_prefix='stream_triage') as pool:\n futures = [pool.submit(_process_unit, *u) for u in work_units]\n for done_count, fut in enumerate(as_completed(futures), 1):\n try:\n key, triage_fields, source, ms, err = fut.result()\n unit_completions[key] = (triage_fields, source, ms, err)\n if source == 'cache':\n group_outcomes[key] = (triage_fields, source)\n elif source.startswith('no_dedup_key'):\n nokey_outcomes[key] = (triage_fields, source)\n else:\n group_outcomes[key] = (triage_fields, source)\n except Exception as e:\n print(f'[triage] WARNING: unexpected worker exception: {e}')\n if done_count % 5 == 0 or done_count == len(futures):\n print(f'[triage] progress {done_count}/{len(futures)} '\n f'(cache_hit={stats[\"cache_hit\"]} triaged={stats[\"triaged\"]} '\n f'failed={stats[\"triage_failed\"]})')\n\n# ── Apply outcomes back to every alert (broadcast leader → followers) ─────────\nenriched_with_triage = [None] * len(enriched_alerts)\nfor i, alert in enumerate(enriched_alerts):\n out = dict(alert) if isinstance(alert, dict) else {'_raw': alert}\n dk = alert.get('dedup_key', '') if isinstance(alert, dict) else ''\n out['has_dedup_key'] = bool(dk)\n\n if dk:\n triage_fields, source = group_outcomes.get(dk, (dict(_UNKNOWN_TRIAGE), 'failed'))\n is_leader = (groups.get(dk, [i])[0] == i)\n # All triage fields are pure data (no file-path references); broadcast as-is.\n for k, v in triage_fields.items():\n out[k] = v\n if not is_leader and source != 'cache':\n out['triage_source'] = 'follower_reused'\n out['triage_status'] = 'reused_from_leader'\n else:\n out['triage_source'] = source\n out['triage_status'] = 'cached' if source == 'cache' else (\n 'ok' if source == 'triaged' else 'failed'\n )\n completion = unit_completions.get(dk)\n if completion and is_leader:\n _, _, ms, err = completion\n if ms:\n out['triage_ms'] = ms\n if err:\n out['triage_error'] = err\n else:\n # no-key alerts: each is its own unit, keyed by alert idx\n triage_fields, source = nokey_outcomes.get(i, (dict(_UNKNOWN_TRIAGE), 'no_dedup_key_failed'))\n for k, v in triage_fields.items():\n out[k] = v\n out['triage_source'] = source\n out['triage_status'] = 'ok' if source == 'no_dedup_key_triaged' else 'failed'\n completion = unit_completions.get(i)\n if completion:\n _, _, ms, err = completion\n if ms:\n out['triage_ms'] = ms\n if err:\n out['triage_error'] = err\n\n enriched_with_triage[i] = out\n\nelapsed_ms = round((time.time() - t_start) * 1000)\nprint(f'[triage] all done in {elapsed_ms}ms: cache_hit={stats[\"cache_hit\"]} '\n f'triaged={stats[\"triaged\"]} failed={stats[\"triage_failed\"]} '\n f'followers_reused={stats[\"followers_reused\"]} '\n f'no_dedup_key={stats[\"no_dedup_key_alerts\"]}')\n\n# Persist new triage results back to cache (merge with concurrent writers).\nif new_results:\n lock_fh = _acquire_lock(lock_path)\n try:\n cache = _load_cache(cache_path)\n for k, v in new_results.items():\n if k in cache:\n del cache[k] # LRU touch (move to end on rewrite)\n cache[k] = v\n evicted = _evict_lru(cache, max_triage_cache_size)\n if evicted:\n print(f'[triage_cache] LRU eviction: dropped {evicted} entries (max={max_triage_cache_size})')\n _save_cache_atomic(cache_path, cache)\n stats['cache_size_after'] = len(cache)\n stats['evicted'] = evicted\n finally:\n _release_lock(lock_fh)\nelse:\n stats['cache_size_after'] = stats['cache_size_before']\n\ntriage_results = []\nfor a in enriched_with_triage:\n triage_results.append({\n 'dedup_key': a.get('dedup_key', ''),\n 'has_dedup_key': a.get('has_dedup_key', False),\n 'threat_name': a.get('threat_name', ''),\n 'sip': a.get('sip', ''),\n 'dip': a.get('dip', ''),\n 'is_duplicate': a.get('is_duplicate'),\n 'triage_source': a.get('triage_source', ''),\n 'triage_status': a.get('triage_status', ''),\n 'attack_verdict': a.get('attack_verdict', ''),\n 'risk_level': a.get('risk_level', ''),\n 'report_title': a.get('report_title', ''),\n 'triage_ms': a.get('triage_ms'),\n 'triage_error': a.get('triage_error'),\n })\n\nstats['elapsed_ms'] = elapsed_ms\nstats['concurrency'] = concurrency\nstats['max_triage_cache_size'] = max_triage_cache_size\n\n# Persist enriched_with_triage according to workflow config. Default is SOC DB;\n# JSONL is still available by config/input for downstream pipelines or archival.\noutput_cfg = _resolve_output_config()\nrun_id = (inputs.get('_run_id')\n or os.environ.get('FLOCKS_RUN_ID')\n or str(int(time.time() * 1000)))\noutput_paths = []\noutput_dir = ''\nfirst_seen_soc_alerts, soc_db_filter_stats = _select_first_seen_soc_alerts(\n enriched_with_triage,\n)\nsoc_db_result = {\n 'path': output_cfg.get('soc_db_path', ''),\n 'table': 'alert_records',\n 'rows': 0,\n 'inserted_rows': 0,\n 'updated_rows': 0,\n}\n\nif output_cfg['write_soc_db'] and first_seen_soc_alerts:\n try:\n soc_db_result.update(_triage_write_soc_db(\n output_cfg['soc_db_path'], first_seen_soc_alerts, run_id,\n ))\n print(f'[triage] persisted {soc_db_result.get(\"rows\", 0)} globally unique alerts to '\n f'SOC DB {soc_db_result.get(\"path\")} '\n f'(inserted={soc_db_result.get(\"inserted_rows\", 0)}, '\n f'updated={soc_db_result.get(\"updated_rows\", 0)}); '\n f'filter={soc_db_filter_stats}')\n except Exception as e:\n import traceback\n print(f'[triage] ERROR: failed to persist triage results to SOC DB: {e}\\n{traceback.format_exc()}')\n raise\nelif output_cfg['write_soc_db']:\n print(f'[triage] no verified first-seen unique alerts to persist; filter={soc_db_filter_stats}')\nelse:\n print(f'[triage] SOC DB output disabled by triage_output_mode={output_cfg[\"requested_mode\"]!r}')\n\nif output_cfg['write_jsonl'] and enriched_with_triage:\n try:\n output_dir = _triage_output_dir(output_cfg.get('jsonl_output_dir', ''))\n output_paths = _triage_write_jsonl(\n output_dir, enriched_with_triage, run_id, stats,\n )\n print(f'[triage] wrote {len(enriched_with_triage)} enriched alerts to '\n f'{len(output_paths)} JSONL file(s) under {output_dir}')\n for p in output_paths:\n print(f' → {p}')\n except Exception as e:\n import traceback\n print(f'[triage] WARNING: failed to persist triage_result JSONL: {e}\\n{traceback.format_exc()}')\nelif not output_cfg['write_jsonl']:\n print(f'[triage] JSONL output disabled by triage_output_mode={output_cfg[\"requested_mode\"]!r}')\n\nstats['output_mode'] = output_cfg['mode']\nstats['requested_output_mode'] = output_cfg['requested_mode']\nstats['output_config_path'] = output_cfg['config_path']\nstats['soc_db_path'] = soc_db_result.get('path', '')\nstats['soc_db_rows'] = soc_db_result.get('rows', 0)\nstats['soc_db_inserted_rows'] = soc_db_result.get('inserted_rows', 0)\nstats['soc_db_updated_rows'] = soc_db_result.get('updated_rows', 0)\nstats['soc_db_first_seen_rows'] = soc_db_filter_stats['first_seen_rows']\nstats['soc_db_skipped_rows'] = (\n soc_db_filter_stats['input_rows'] - soc_db_filter_stats['first_seen_rows']\n)\nstats['soc_db_filter_stats'] = soc_db_filter_stats\nstats['output_paths'] = output_paths\nstats['output_dir'] = output_dir\n\nprint(f'[triage] stats={json.dumps(stats, ensure_ascii=False)}')\n\noutputs['enriched_alerts_with_triage'] = enriched_with_triage\noutputs['triage_results'] = triage_results\noutputs['triage_stats'] = stats\noutputs['load_stats'] = inputs.get('load_stats', {})\noutputs['loaded_files'] = inputs.get('loaded_files', [])\noutputs['input_date'] = inputs.get('input_date', '')\noutputs['triage_output_mode'] = output_cfg['mode']\noutputs['soc_db_result'] = soc_db_result\noutputs['soc_db_path'] = soc_db_result.get('path', '')\noutputs['output_config_path'] = output_cfg['config_path']\noutputs['output_paths'] = output_paths\noutputs['output_dir'] = output_dir\n" + "description": "Leader/follower 分组并发研判节点（自包含，内联 tdp_alert_triage 逻辑）。先按 dedup_key 把 alerts 分组：每组只对 leader 研判，follower 复用 leader 结果。外层 ThreadPoolExecutor(concurrency) 处理 unique work units（concurrency 取值 1–5，默认 1），单条告警仍执行 survey / cve_related / cve_info / payload_analysis 4 个分支，但所有 llm.ask() 共享运行级 concurrency 预算，总 LLM 峰值不超过 1–5。dedup_key 在 triage_cache.pkl 命中时直接复用历史 verdict/title/triage_report；未命中则 leader 执行完整研判（情报查询 + 4 个 LLM + attack_analysis + verdict + title + 聚合 markdown），完整研判 markdown 仅写入 alert 的 `triage_report` 字段，**不生成任何独立报告文件**。新结果合并写回 cache（FIFO LRU + 文件锁 + 原子落盘）。SOC DB 只接受明确 `is_duplicate=false`、包含 `dedup_key` 且批内首次出现的告警，并通过数据库唯一索引保证跨执行全局唯一；重复 key 只更新研判字段并保留首次事件元数据，持久化失败会使工作流失败。可通过工作流目录 `config.json` 或运行输入将 `triage_output_mode` 切换为 `jsonl` / `both` / `none`，保留 `triage_result_NNN.jsonl` 可选输出。 该节点在短生命周期宿主子进程中执行，结束后由操作系统回收解释器堆。", + "code": "\"\"\"\nconcurrent_triage: leader/follower 分组并发研判 + dedup_key 缓存复用（自包含）。\n\n去重模式：\n 1. 输入 alerts 先按 dedup_key 分组 → unique dedup_keys 列表\n 2. 每个 group 只对 leader（首条）做研判；followers 复用 leader 结果，不重复调 LLM\n 3. 无 dedup_key 的 alert 各自独立成 work unit（防御性研判，无法复用）\n\n并发结构：\n 外层按 3 个 work unit 分批创建 ThreadPoolExecutor，避免整批 future/上下文同时驻留\n （concurrency 取值 1–5，默认 1，由 inputs.concurrency 控制，每批实际 worker 不超过 3）\n 单条 alert 包含 4 个并行分支以及攻击分析、结论、标题、报告 4 个汇总调用\n （survey / cve_related / cve_info / payload_analysis — 保留 tdp_alert_triage\n 的 4 分支研判结构）\n 所有 llm.ask 调用共享运行级信号量，稳态 LLM 峰值不超过 concurrency\n\n模型攻击结论以两个正交维度附加到每条 alert：triage_attack_verdict 表示\n是否攻击（attack / non_attack / unknown），triage_attack_success 表示攻击\n结果（success / failed / unknown）。原始同名字段保持不变；其他研判产物包括 risk_level /\nreport_title / triage_report 等。**不生成任何独立的 per-alert markdown 报告\n文件**，避免冗余落盘与跨日期路径失效。\n\ndedup_key 缓存（与 stream_alert_denoise 的 LSH 状态文件同根目录，逻辑独立）：\n ~/.flocks/workspace/workflows/stream_alert_triage/triage_cache.pkl\n - cache 命中：直接复用历史 verdict/title/triage_report，**不调用 LLM**\n - cache 未命中：leader 执行完整内联研判（情报 + 4 个 LLM 分支 + verdict + title + report）；\n follower 直接广播 leader 结果\n - 新结果合并写回 cache，FIFO LRU 淘汰，文件锁 + 原子落盘\n\"\"\"\n\nimport ipaddress\nimport json\nimport os\nimport pickle\nimport re\nimport sys\nimport threading\nimport time\nfrom concurrent.futures import ThreadPoolExecutor, as_completed\n\nIS_WINDOWS = sys.platform == 'win32'\nif IS_WINDOWS:\n import msvcrt # noqa: F401\nelse:\n import fcntl # noqa: F401\n\nWORKFLOW_NAME = 'stream_alert_triage'\n\n# Per-call LLM timeout and retry budget for every analysis branch in this\n# node. Workflow LLM calls share the dedicated ``flocks-workflow-llm-loop``;\n# a single hung call (e.g. provider 504, slow TLS handshake) without a\n# timeout would otherwise pin one of the (already concurrency-limited)\n# worker threads for up to httpx's DEFAULT read timeout (10 min), serially\n# blocking the rest of the alert pipeline. 120s + 1 retry covers normal\n# slow-but-alive responses while still recovering from transient hangs.\nLLM_CALL_TIMEOUT_S = 120.0\nLLM_CALL_MAX_RETRIES = 1\n\n# Cache and work scheduling remain bounded without truncating alert or model content.\nMAX_TRIAGE_CACHE_BYTES = 128 * 1024 * 1024\nTRIAGE_SUB_BATCH_SIZE = 3\n\n# The user-facing concurrency setting is a run-wide LLM request budget. The\n# four logical branches still start together, but nested executors must not\n# multiply the actual provider load (5 outer workers used to become 20 calls).\n_llm_slots = None\n\nTRIAGE_FIELDS = (\n 'triage_attack_verdict',\n 'triage_attack_success',\n 'risk_level',\n 'report_title',\n 'triage_report',\n)\nATTACK_VERDICT_CN = {\n 'attack': '攻击',\n 'non_attack': '非攻击',\n 'unknown': '未知',\n}\nATTACK_SUCCESS_CN = {\n 'success': '攻击成功',\n 'failed': '攻击失败',\n 'unknown': '未知',\n}\nOUTCOME_RISK = {\n 'attack_success': 'High',\n 'attack_failed': 'Medium',\n 'attack': 'Medium',\n 'unknown': 'Medium',\n 'non_attack': 'Low',\n}\nOUTCOME_CN = {\n 'attack_success': '攻击成功',\n 'attack_failed': '攻击失败',\n 'attack': '攻击，结果未知',\n 'unknown': '未知',\n 'non_attack': '非攻击',\n}\nTRIAGE_REPORT_VERSION = 'soc.triage.markdown.v1'\nTRIAGE_REPORT_TAGS = (\n 'report_title',\n 'report_meta',\n 'analysis_steps',\n 'triage_conclusion',\n 'attack_payload',\n 'payload_explanation',\n 'response_evidence',\n 'key_evidence',\n 'disposal_recommendation',\n)\n\n\n# ── Cache persistence ─────────────────────────────────────────────────────────\n\ndef _cache_paths():\n configured_state_dir = inputs.get('_triage_state_dir')\n if configured_state_dir:\n from pathlib import Path\n state_dir = Path(os.path.expanduser(str(configured_state_dir)))\n else:\n from flocks.config import Config\n flocks_root = Config().get_global().data_dir.parent\n state_dir = flocks_root / 'workspace' / 'workflows' / WORKFLOW_NAME\n state_dir.mkdir(parents=True, exist_ok=True)\n return str(state_dir / 'triage_cache.pkl'), str(state_dir / 'triage_cache.lock')\n\n\ndef _acquire_lock(lock_path):\n fh = open(lock_path, 'w+')\n try:\n if IS_WINDOWS:\n fh.write('L'); fh.flush(); fh.seek(0)\n while True:\n try:\n msvcrt.locking(fh.fileno(), msvcrt.LK_LOCK, 1); break\n except OSError:\n continue\n else:\n fcntl.flock(fh.fileno(), fcntl.LOCK_EX)\n except BaseException:\n try:\n fh.close()\n except Exception:\n pass\n raise\n return fh\n\n\ndef _release_lock(fh):\n try:\n if IS_WINDOWS:\n try:\n fh.seek(0); msvcrt.locking(fh.fileno(), msvcrt.LK_UNLCK, 1)\n except OSError:\n pass\n else:\n fcntl.flock(fh.fileno(), fcntl.LOCK_UN)\n finally:\n fh.close()\n\n\ndef _batch_lease_path():\n cache_path, _ = _cache_paths()\n return os.path.join(os.path.dirname(cache_path), '.triage_batch.lock')\n\n\ndef _acquire_batch_lease():\n path = _batch_lease_path()\n fd = os.open(path, os.O_RDWR | os.O_CREAT, 0o600)\n try:\n if IS_WINDOWS:\n if os.fstat(fd).st_size == 0:\n os.write(fd, b'0')\n os.fsync(fd)\n os.lseek(fd, 0, os.SEEK_SET)\n msvcrt.locking(fd, msvcrt.LK_NBLCK, 1)\n else:\n fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)\n except OSError as exc:\n os.close(fd)\n raise RuntimeError('production_batch_lease_busy') from exc\n\n try:\n token = f'{os.getpid()}-{threading.get_ident()}-{time.time_ns()}'\n os.ftruncate(fd, 0)\n os.lseek(fd, 0, os.SEEK_SET)\n os.write(fd, token.encode('ascii'))\n os.fsync(fd)\n except BaseException:\n _release_batch_lease(fd)\n raise\n print(f'[triage] acquired production batch lease token={token}')\n return fd, token\n\n\ndef _release_batch_lease(fd):\n if type(fd) is not int or fd < 0:\n return\n try:\n if IS_WINDOWS:\n try:\n os.lseek(fd, 0, os.SEEK_SET)\n msvcrt.locking(fd, msvcrt.LK_UNLCK, 1)\n except OSError:\n pass\n else:\n fcntl.flock(fd, fcntl.LOCK_UN)\n finally:\n os.close(fd)\n\n\ndef _load_cache(cache_path):\n if not os.path.exists(cache_path):\n return {}\n cache_size = os.path.getsize(cache_path)\n if cache_size == 0:\n return {}\n if cache_size > MAX_TRIAGE_CACHE_BYTES:\n quarantine_path = f'{cache_path}.{time.time_ns()}.oversized'\n try:\n os.replace(cache_path, quarantine_path)\n print(f'[triage_cache] WARNING: skipped oversized cache '\n f'({cache_size} bytes > {MAX_TRIAGE_CACHE_BYTES}); '\n f'moved to {quarantine_path}')\n except OSError as e:\n print(f'[triage_cache] WARNING: skipped oversized cache '\n f'({cache_size} bytes > {MAX_TRIAGE_CACHE_BYTES}); '\n f'quarantine failed: {e}')\n return {}\n try:\n with open(cache_path, 'rb') as f:\n c = pickle.load(f)\n if not isinstance(c, dict):\n return {}\n print(f'[triage_cache] loaded {len(c)} entries from {cache_path}')\n return c\n except Exception as e:\n print(f'[triage_cache] WARNING: failed to load ({e}), starting fresh')\n return {}\n\n\n\ndef _save_cache_atomic(cache_path, cache):\n tmp = (\n f'{cache_path}.{os.getpid()}.{threading.get_ident()}.'\n f'{time.time_ns()}.tmp'\n )\n try:\n with open(tmp, 'wb') as f:\n pickle.dump(cache, f)\n f.flush()\n os.fsync(f.fileno())\n os.replace(tmp, cache_path)\n print(f'[triage_cache] saved {len(cache)} entries -> {cache_path}')\n except Exception as exc:\n try:\n if os.path.exists(tmp):\n os.remove(tmp)\n except OSError:\n pass\n raise RuntimeError(f'failed to save triage cache: {exc}') from exc\n\n\ndef _evict_lru(cache, max_keys, max_bytes=None):\n evicted = 0\n excess = len(cache) - max_keys\n if excess > 0:\n for k in list(cache.keys())[:excess]:\n del cache[k]\n evicted += 1\n\n if max_bytes is not None and max_bytes > 0 and cache:\n # Estimate one entry at a time so enforcing the cache budget never\n # creates a second serialized copy of the whole cache in memory.\n entry_sizes = {}\n estimated_bytes = 64\n for key, value in cache.items():\n size = len(pickle.dumps((key, value), protocol=pickle.HIGHEST_PROTOCOL)) + 16\n entry_sizes[key] = size\n estimated_bytes += size\n for key in list(cache.keys()):\n if estimated_bytes <= max_bytes:\n break\n estimated_bytes -= entry_sizes[key]\n del cache[key]\n evicted += 1\n return evicted\n\n\n# ── Runtime output config and persistence targets ──────────────────────────────\n#\n# Defaults are read from ~/.flocks/plugins/workflows/stream_alert_triage/config.json.\n# Runtime inputs override config values. The default mode is soc_db so SOC pages\n# read the same DB-backed dataset. JSONL remains available via config/input:\n# triage_output_mode = soc_db | jsonl | both | none\n# persist_triage_output = true (legacy alias that adds JSONL to soc_db)\n\nimport datetime as _datetime\n\nMAX_RECORDS_PER_FILE = 10000\n_TRIAGE_JSONL_PREFIX = 'triage_result'\n_TRIAGE_COUNTER_FILE = '.triage_counter.json'\n_WORKFLOW_CONFIG_PATH = os.path.expanduser('~/.flocks/plugins/workflows/stream_alert_triage/config.json')\n_DEFAULT_SOC_DB_PATH = os.path.expanduser('~/.flocks/data/soc.db')\n\n\ndef _load_workflow_config():\n try:\n with open(_WORKFLOW_CONFIG_PATH, 'r', encoding='utf-8') as f:\n cfg = json.load(f)\n if isinstance(cfg, dict):\n return cfg\n except FileNotFoundError:\n pass\n except Exception as e:\n print(f'[triage_config] WARNING: failed to read {_WORKFLOW_CONFIG_PATH}: {e}')\n return {}\n\n\ndef _configured_value(config, key, default=None):\n if key in inputs:\n value = inputs.get(key)\n if value is not None and not (isinstance(value, str) and not value.strip()):\n return value\n return config.get(key, default)\n\n\ndef _input_bool(value, default=False):\n if value is None:\n return default\n if isinstance(value, bool):\n return value\n if isinstance(value, (int, float)):\n return bool(value)\n text = str(value).strip().lower()\n if text in {'1', 'true', 'yes', 'y', 'on'}:\n return True\n if text in {'0', 'false', 'no', 'n', 'off'}:\n return False\n return default\n\n\ndef _select_first_seen_soc_alerts(alerts):\n selected = []\n seen_dedup_keys = set()\n stats = {\n 'input_rows': len(alerts),\n 'first_seen_rows': 0,\n 'skipped_not_first_seen_rows': 0,\n 'skipped_missing_dedup_key_rows': 0,\n 'skipped_repeated_dedup_key_rows': 0,\n }\n for alert in alerts:\n if not isinstance(alert, dict) or _input_bool(alert.get('is_duplicate'), True):\n stats['skipped_not_first_seen_rows'] += 1\n continue\n dedup_key = str(alert.get('dedup_key') or '').strip()\n if not dedup_key:\n stats['skipped_missing_dedup_key_rows'] += 1\n continue\n if dedup_key in seen_dedup_keys:\n stats['skipped_repeated_dedup_key_rows'] += 1\n continue\n seen_dedup_keys.add(dedup_key)\n selected.append(alert)\n stats['first_seen_rows'] = len(selected)\n return selected, stats\n\n\ndef _resolve_output_config():\n config = _load_workflow_config()\n raw_mode = str(_configured_value(config, 'triage_output_mode', 'soc_db') or 'soc_db').strip().lower()\n mode_alias = {\n 'db': 'soc_db',\n 'sqlite': 'soc_db',\n 'sqlite_db': 'soc_db',\n 'soc': 'soc_db',\n 'json': 'jsonl',\n 'file': 'jsonl',\n 'files': 'jsonl',\n 'off': 'none',\n 'disabled': 'none',\n }\n requested_mode = mode_alias.get(raw_mode, raw_mode)\n if requested_mode not in {'soc_db', 'jsonl', 'both', 'none'}:\n print(f'[triage_config] WARNING: invalid triage_output_mode={raw_mode!r}; using soc_db')\n requested_mode = 'soc_db'\n\n legacy_jsonl = _input_bool(_configured_value(config, 'persist_triage_output', False), False)\n write_soc_db = requested_mode in {'soc_db', 'both'}\n write_jsonl = requested_mode in {'jsonl', 'both'}\n effective_mode = requested_mode\n if requested_mode == 'soc_db' and legacy_jsonl:\n write_jsonl = True\n effective_mode = 'both'\n if requested_mode == 'none':\n write_soc_db = False\n write_jsonl = False\n effective_mode = 'none'\n\n soc_db_path = os.path.expanduser(str(\n _configured_value(config, 'soc_db_path', _DEFAULT_SOC_DB_PATH) or _DEFAULT_SOC_DB_PATH\n ))\n jsonl_output_dir = _configured_value(config, 'jsonl_output_dir', '') or ''\n jsonl_output_dir = os.path.expanduser(str(jsonl_output_dir)) if jsonl_output_dir else ''\n return {\n 'config_path': _WORKFLOW_CONFIG_PATH,\n 'requested_mode': requested_mode,\n 'mode': effective_mode,\n 'write_soc_db': write_soc_db,\n 'write_jsonl': write_jsonl,\n 'soc_db_path': soc_db_path,\n 'jsonl_output_dir': jsonl_output_dir,\n }\n\n\n# ── Persisted JSONL output (optional; mirrors stream_alert_denoise layout) ─────\n#\n# Directory : ~/.flocks/workspace/workflows/stream_alert_triage//\n# Filename : triage_result_NNN.jsonl (3-digit zero-padded seq)\n# Layout : line 1 = {\"_type\":\"file_header\", ...}, subsequent lines = one\n# enriched_with_triage alert per line.\n# Counter : .triage_counter.json sidecar tracks (seq, count) so we don't\n# rescan every existing file on each run; auto-rolls over to a\n# new file when reaching MAX_RECORDS_PER_FILE.\n\n\ndef _triage_output_dir(configured_dir=''):\n \"\"\"Return output directory for triage_result_*.jsonl.\"\"\"\n if configured_dir:\n out_dir = configured_dir\n os.makedirs(out_dir, exist_ok=True)\n return out_dir\n from flocks.config import Config\n flocks_root = Config().get_global().data_dir.parent\n date_str = _datetime.datetime.now().strftime('%Y-%m-%d')\n out_dir = flocks_root / 'workspace' / 'workflows' / WORKFLOW_NAME / date_str\n out_dir.mkdir(parents=True, exist_ok=True)\n return str(out_dir)\n\n\ndef _triage_get_counter(out_dir):\n path = os.path.join(out_dir, _TRIAGE_COUNTER_FILE)\n try:\n with open(path, 'r', encoding='utf-8') as f:\n d = json.load(f)\n return int(d.get('seq', 0)), int(d.get('count', 0))\n except Exception:\n return 0, 0\n\n\ndef _triage_set_counter(out_dir, seq, count):\n path = os.path.join(out_dir, _TRIAGE_COUNTER_FILE)\n tmp = path + '.tmp'\n with open(tmp, 'w', encoding='utf-8') as f:\n json.dump({'seq': seq, 'count': count}, f)\n os.replace(tmp, path)\n\n\ndef _triage_find_active_file(out_dir):\n \"\"\"Locate the active (latest, not-yet-full) jsonl file; create if none.\"\"\"\n seq, count = _triage_get_counter(out_dir)\n if seq > 0:\n path = os.path.join(out_dir, f'{_TRIAGE_JSONL_PREFIX}_{seq:03d}.jsonl')\n if os.path.exists(path):\n return path, count, seq\n import glob as _glob\n existing = sorted(_glob.glob(os.path.join(out_dir, _TRIAGE_JSONL_PREFIX + '_*.jsonl')))\n if not existing:\n return None, 0, 0\n latest = existing[-1]\n try:\n seq = int(os.path.basename(latest).replace(_TRIAGE_JSONL_PREFIX + '_', '').replace('.jsonl', ''))\n except ValueError:\n seq = len(existing)\n count = 0\n try:\n with open(latest, 'r', encoding='utf-8') as f:\n for line in f:\n if line.strip() and '\"_type\"' not in line:\n count += 1\n except Exception:\n pass\n return latest, count, seq\n\n\ndef _triage_write_jsonl(out_dir, alerts, run_id, run_stats):\n \"\"\"Append all alerts to today's triage_result_NNN.jsonl, rolling over at\n MAX_RECORDS_PER_FILE. Returns the list of files that were written to.\"\"\"\n now = _datetime.datetime.now()\n written = []\n active_path, active_count, seq = _triage_find_active_file(out_dir)\n remaining = list(alerts)\n while remaining:\n available = MAX_RECORDS_PER_FILE - active_count\n if available <= 0 or active_path is None:\n seq += 1\n active_path = os.path.join(out_dir, f'{_TRIAGE_JSONL_PREFIX}_{seq:03d}.jsonl')\n active_count = 0\n available = MAX_RECORDS_PER_FILE\n header = {\n '_type': 'file_header',\n 'created_at': now.isoformat(),\n 'date': now.strftime('%Y-%m-%d'),\n 'workflow': WORKFLOW_NAME,\n 'seq': seq,\n 'run_id': run_id,\n 'batch_total': run_stats.get('total'),\n 'batch_triaged': run_stats.get('triaged'),\n 'batch_followers_reused':run_stats.get('followers_reused'),\n 'batch_cache_hit': run_stats.get('cache_hit'),\n 'batch_triage_failed': run_stats.get('triage_failed'),\n }\n with open(active_path, 'w', encoding='utf-8') as hf:\n hf.write(json.dumps(header, ensure_ascii=False) + '\\n')\n batch = remaining[:available]\n remaining = remaining[available:]\n with open(active_path, 'a', encoding='utf-8') as af:\n for alert in batch:\n af.write(json.dumps(alert, ensure_ascii=False) + '\\n')\n active_count += len(batch)\n if active_path not in written:\n written.append(active_path)\n if remaining:\n active_path = None\n active_count = 0\n if written:\n _triage_set_counter(out_dir, seq, active_count)\n return written\n\n\n# ── SOC DB output (default) ───────────────────────────────────────────────────\n\ndef _ensure_soc_db_schema(conn):\n conn.execute(\"\"\"\n CREATE TABLE IF NOT EXISTS alert_records (\n row_id TEXT PRIMARY KEY,\n record_id TEXT,\n asset_date TEXT NOT NULL,\n source_file TEXT NOT NULL,\n line_number INTEGER NOT NULL,\n event_time INTEGER,\n source_type TEXT,\n threat_name TEXT,\n dedup_key TEXT,\n is_duplicate INTEGER NOT NULL DEFAULT 0,\n record_json TEXT NOT NULL\n )\n \"\"\")\n columns = {row[1] for row in conn.execute('PRAGMA table_info(alert_records)')}\n dedup_key_added = 'dedup_key' not in columns\n if dedup_key_added:\n conn.execute('ALTER TABLE alert_records ADD COLUMN dedup_key TEXT')\n\n unique_index_name = 'idx_alert_records_first_seen_dedup_key'\n indexes = list(conn.execute('PRAGMA index_list(alert_records)'))\n unique_index_ready = any(row[1] == unique_index_name and bool(row[2]) for row in indexes)\n conn.execute('CREATE INDEX IF NOT EXISTS idx_alert_records_duplicate ON alert_records(is_duplicate)')\n has_persisted_duplicates = conn.execute(\"\"\"\n SELECT 1 FROM alert_records\n WHERE is_duplicate = 1\n AND dedup_key IS NOT NULL\n AND dedup_key <> ''\n LIMIT 1\n \"\"\").fetchone() is not None\n if not unique_index_ready or has_persisted_duplicates:\n if any(row[1] == unique_index_name for row in indexes):\n conn.execute(f'DROP INDEX {unique_index_name}')\n conn.execute(\"\"\"\n UPDATE alert_records\n SET dedup_key = CASE\n WHEN json_valid(record_json)\n THEN NULLIF(TRIM(CAST(json_extract(record_json, '$.dedup_key') AS TEXT)), '')\n ELSE NULL\n END\n WHERE dedup_key IS NULL OR TRIM(dedup_key) = ''\n \"\"\")\n conn.execute(\"\"\"\n UPDATE alert_records\n SET dedup_key = NULLIF(TRIM(dedup_key), '')\n WHERE dedup_key IS NOT NULL\n \"\"\")\n conn.execute(\"\"\"\n DELETE FROM alert_records\n WHERE dedup_key IS NOT NULL\n AND dedup_key <> ''\n AND rowid NOT IN (\n SELECT MIN(rowid)\n FROM alert_records\n WHERE dedup_key IS NOT NULL AND dedup_key <> ''\n AND is_duplicate = 0\n GROUP BY dedup_key\n )\n \"\"\")\n conn.execute(f\"\"\"\n CREATE UNIQUE INDEX {unique_index_name}\n ON alert_records(dedup_key)\n WHERE dedup_key IS NOT NULL AND dedup_key <> ''\n \"\"\")\n\n conn.execute('CREATE INDEX IF NOT EXISTS idx_alert_records_asset_date ON alert_records(asset_date)')\n conn.execute('CREATE INDEX IF NOT EXISTS idx_alert_records_event_time ON alert_records(event_time)')\n conn.execute('CREATE INDEX IF NOT EXISTS idx_alert_records_source_type ON alert_records(source_type)')\n conn.execute('CREATE INDEX IF NOT EXISTS idx_alert_records_threat_name ON alert_records(threat_name)')\n\n\ndef _event_time_value(alert):\n for key in ('time', 'event_time', 'timestamp', 'timestamp_real', 'occur_time', 'created_at'):\n value = alert.get(key)\n if value in (None, ''):\n continue\n if isinstance(value, (int, float)):\n ts = float(value)\n if ts > 100000000000:\n ts = ts / 1000.0\n return int(ts)\n text = str(value).strip()\n if not text:\n continue\n try:\n ts = float(text)\n if ts > 100000000000:\n ts = ts / 1000.0\n return int(ts)\n except Exception:\n pass\n normalized = text.replace('Z', '+00:00')\n try:\n return int(_datetime.datetime.fromisoformat(normalized).timestamp())\n except Exception:\n pass\n for fmt in ('%Y-%m-%d %H:%M:%S', '%Y/%m/%d %H:%M:%S', '%Y-%m-%d %H:%M', '%Y/%m/%d %H:%M'):\n try:\n return int(_datetime.datetime.strptime(text, fmt).timestamp())\n except Exception:\n continue\n return int(time.time())\n\n\ndef _asset_date_value(alert, event_time):\n value = alert.get('asset_date') or alert.get('_asset_date') or alert.get('date')\n if value:\n text = str(value).strip()\n if re.match(r'^\\d{4}-\\d{2}-\\d{2}$', text):\n return text\n try:\n return _datetime.datetime.fromtimestamp(int(event_time)).strftime('%Y-%m-%d')\n except Exception:\n return _datetime.datetime.now().strftime('%Y-%m-%d')\n\n\ndef _source_type_value(alert):\n for key in ('source_type', '_source_type', 'data_source', 'log_type', 'vendor', 'device_type'):\n value = alert.get(key)\n if value not in (None, ''):\n return str(value)\n return ''\n\n\ndef _record_id_value(alert):\n for key in ('record_id', 'id', 'uuid', 'event_id', 'dedup_key'):\n value = alert.get(key)\n if value not in (None, ''):\n return str(value)\n return ''\n\n\ndef _stable_row_id(alert, source_file, line_number, event_time):\n existing = alert.get('row_id') or alert.get('_row_id')\n if existing:\n return str(existing)\n import hashlib as _hashlib\n basis = {\n 'record_id': _record_id_value(alert),\n 'dedup_key': alert.get('dedup_key', ''),\n 'time': event_time,\n 'source_file': source_file,\n 'line_number': line_number,\n 'sip': alert.get('sip', ''),\n 'sport': alert.get('sport', ''),\n 'dip': alert.get('dip', ''),\n 'dport': alert.get('dport', ''),\n 'threat_rule_id': alert.get('threat_rule_id') or alert.get('rule_id') or '',\n }\n raw = json.dumps(basis, sort_keys=True, ensure_ascii=False)\n return _hashlib.sha256(raw.encode('utf-8')).hexdigest()\n\n\ndef _load_existing_soc_rows(conn, dedup_keys):\n existing = {}\n unique_keys = list(dict.fromkeys(str(key).strip() for key in dedup_keys if str(key).strip()))\n for start in range(0, len(unique_keys), 500):\n chunk = unique_keys[start:start + 500]\n placeholders = ','.join('?' for _ in chunk)\n rows = conn.execute(f\"\"\"\n SELECT row_id, record_id, asset_date, source_file, line_number,\n event_time, source_type, threat_name, dedup_key,\n is_duplicate, record_json\n FROM alert_records\n WHERE dedup_key IN ({placeholders})\n \"\"\", chunk)\n for row in rows:\n try:\n record = json.loads(row[10])\n except Exception:\n record = {}\n if not isinstance(record, dict):\n record = {}\n existing[row[8]] = {\n 'row_id': row[0],\n 'record_id': row[1],\n 'asset_date': row[2],\n 'source_file': row[3],\n 'line_number': row[4],\n 'event_time': row[5],\n 'source_type': row[6],\n 'threat_name': row[7],\n 'record': record,\n }\n return existing\n\n\ndef _merge_triage_record(existing_record, incoming_record):\n merged = dict(existing_record) if isinstance(existing_record, dict) else {}\n triage_fields = (\n 'has_dedup_key',\n 'triage_source',\n 'triage_status',\n 'triage_attack_verdict',\n 'risk_level',\n 'report_title',\n 'triage_report',\n 'triage_attack_success',\n 'triage_ms',\n 'triage_error',\n '_triage_run_id',\n '_triage_persisted_at',\n )\n for key in triage_fields:\n if key in incoming_record:\n merged[key] = incoming_record[key]\n elif key in {'triage_ms', 'triage_error'}:\n merged.pop(key, None)\n return merged\n\n\ndef _triage_write_soc_db(db_path, alerts, run_id):\n import sqlite3\n\n db_dir = os.path.dirname(db_path)\n if db_dir:\n os.makedirs(db_dir, exist_ok=True)\n default_source_file = ''\n loaded_files = inputs.get('loaded_files') or []\n if isinstance(loaded_files, list) and len(loaded_files) == 1:\n default_source_file = str(loaded_files[0])\n\n persisted_at = _datetime.datetime.now().isoformat()\n candidates = []\n seen_dedup_keys = set()\n for idx, alert in enumerate(alerts, 1):\n if not isinstance(alert, dict):\n continue\n dedup_key = str(alert.get('dedup_key') or '').strip()\n if not dedup_key or dedup_key in seen_dedup_keys:\n continue\n seen_dedup_keys.add(dedup_key)\n record = dict(alert)\n record['dedup_key'] = dedup_key\n record['is_duplicate'] = False\n record['_triage_run_id'] = run_id\n record['_triage_persisted_at'] = persisted_at\n source_file = str(\n record.get('source_file')\n or record.get('_source_file')\n or record.get('file_path')\n or default_source_file\n or 'stream_alert_triage'\n )\n try:\n line_number = int(record.get('line_number') or record.get('_line_number') or idx)\n except Exception:\n line_number = idx\n event_time = _event_time_value(record)\n candidates.append({\n 'row_id': _stable_row_id(record, source_file, line_number, event_time),\n 'record_id': _record_id_value(record),\n 'asset_date': _asset_date_value(record, event_time),\n 'source_file': source_file,\n 'line_number': line_number,\n 'event_time': event_time,\n 'source_type': _source_type_value(record),\n 'threat_name': str(record.get('threat_name') or record.get('rule_name') or ''),\n 'dedup_key': dedup_key,\n 'record': record,\n })\n\n insert_rows = []\n update_rows = []\n with sqlite3.connect(db_path, timeout=30) as conn:\n conn.execute('BEGIN IMMEDIATE')\n _ensure_soc_db_schema(conn)\n existing_by_key = _load_existing_soc_rows(\n conn, [candidate['dedup_key'] for candidate in candidates],\n )\n for candidate in candidates:\n existing = existing_by_key.get(candidate['dedup_key'])\n if existing:\n merged_record = _merge_triage_record(existing['record'], candidate['record'])\n merged_record['dedup_key'] = candidate['dedup_key']\n merged_record['is_duplicate'] = False\n update_rows.append((\n candidate['dedup_key'],\n json.dumps(merged_record, ensure_ascii=False),\n existing['row_id'],\n ))\n continue\n insert_rows.append((\n candidate['row_id'],\n candidate['record_id'],\n candidate['asset_date'],\n candidate['source_file'],\n candidate['line_number'],\n candidate['event_time'],\n candidate['source_type'],\n candidate['threat_name'],\n candidate['dedup_key'],\n 0,\n json.dumps(candidate['record'], ensure_ascii=False),\n ))\n\n if insert_rows:\n conn.executemany(\"\"\"\n INSERT INTO alert_records (\n row_id, record_id, asset_date, source_file, line_number,\n event_time, source_type, threat_name, dedup_key,\n is_duplicate, record_json\n ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)\n \"\"\", insert_rows)\n if update_rows:\n conn.executemany(\"\"\"\n UPDATE alert_records\n SET dedup_key = ?, is_duplicate = 0, record_json = ?\n WHERE row_id = ?\n \"\"\", update_rows)\n conn.commit()\n\n persisted_rows = len(insert_rows) + len(update_rows)\n return {\n 'path': db_path,\n 'table': 'alert_records',\n 'rows': persisted_rows,\n 'inserted_rows': len(insert_rows),\n 'updated_rows': len(update_rows),\n }\n\n\n# ── LLM provider warm-up (avoid cold-start race in _parallel_4_branches) ──────\n#\n# Background: when this node starts, the LLM provider (e.g. threatbook-cn-llm)\n# is lazy-initialized on the first call. Inside `_parallel_4_branches` we\n# submit 4 LLM calls to a ThreadPoolExecutor simultaneously; whichever one\n# wins the race may race against provider registration and fail with\n# \"provider 'xxx' not exists\" while subsequent calls succeed. A single\n# synchronous warm-up call before any concurrent fan-out forces the provider\n# to finish registering on the main thread, eliminating the race entirely.\n#\n# Failure of the warm-up is non-fatal — we just log and continue. The first\n# real LLM call will still see the same error and surface it normally.\n\ndef _warmup_llm():\n \"\"\"Force LLM provider lazy-init on the main thread before any fan-out.\"\"\"\n t0 = time.time()\n try:\n _ask_llm('ping')\n print(f'[triage] LLM provider warm-up OK in {round((time.time()-t0)*1000)}ms')\n return True\n except Exception as e:\n print(f'[triage] WARNING: LLM warm-up failed ({type(e).__name__}: '\n f'{str(e)[:200]}); proceeding anyway')\n return False\n\n\n# ── Inline triage helpers (mirroring tdp_alert_triage docs version) ───────────\n\ndef _strip_think(text):\n return re.sub(r'[\\s\\S]*?', '', str(text or ''), flags=re.IGNORECASE).strip()\n\n\ndef _is_public_ip(value):\n try:\n ip_obj = ipaddress.ip_address(value)\n except Exception:\n return False\n return not (ip_obj.is_private or ip_obj.is_loopback or ip_obj.is_reserved\n or ip_obj.is_link_local or ip_obj.is_multicast or ip_obj.is_unspecified)\n\n\ndef _pick(*values):\n for v in values:\n if v not in (None, '', [], {}):\n return v\n return ''\n\n\ndef _parse_alert(alert_input):\n \"\"\"Mirrors tdp_alert_triage.receive_alert. Supports three input shapes:\n nested TDP (net.http.url), flat TDP (net_http_url), normalized (req_http_url).\n \"\"\"\n if isinstance(alert_input, str):\n try:\n alert_input = json.loads(alert_input)\n except Exception:\n alert_input = {}\n if isinstance(alert_input, list):\n alert_data = alert_input[0] if alert_input else {}\n elif isinstance(alert_input, dict) and isinstance(alert_input.get('data'), list):\n alert_data = alert_input.get('data', [])[0] if alert_input.get('data') else {}\n else:\n alert_data = alert_input if isinstance(alert_input, dict) else {}\n\n net = alert_data.get('net', {}) or {}\n http = net.get('http', {}) or {}\n threat = alert_data.get('threat', {}) or {}\n assets = alert_data.get('assets', {}) or {}\n\n src_ip = _pick(\n alert_data.get('attacker'), alert_data.get('external_ip'),\n net.get('src_ip'), net.get('flow_src_ip'),\n alert_data.get('net_real_src_ip'),\n alert_data.get('sip'), alert_data.get('src_ip'), alert_data.get('src'),\n )\n dst_ip = _pick(\n alert_data.get('victim'), alert_data.get('machine'),\n alert_data.get('server_ip'), net.get('dest_ip'), net.get('flow_dest_ip'),\n alert_data.get('net_dest_ip'),\n alert_data.get('dip'), alert_data.get('dst_ip'), alert_data.get('dst'),\n )\n src_port = _pick(\n net.get('src_port'), net.get('flow_src_port'),\n alert_data.get('external_port'), alert_data.get('net_src_port'),\n alert_data.get('sport'), alert_data.get('src_port'), 0,\n )\n dst_port = _pick(\n net.get('dest_port'), net.get('flow_dest_port'),\n alert_data.get('server_port'), alert_data.get('machine_port'),\n alert_data.get('net_dest_port'),\n alert_data.get('dport'), alert_data.get('dst_port'), 0,\n )\n protocol = _pick(\n net.get('app_proto'), net.get('type'), net.get('proto'),\n alert_data.get('net_app_proto'), alert_data.get('protocol'),\n alert_data.get('event_type'), 'TCP',\n )\n alert_type = _pick(\n threat.get('name'), alert_data.get('threat_name'),\n alert_data.get('vuln_name'),\n alert_data.get('alert_type'), threat.get('topic'),\n alert_data.get('type'), 'unknown',\n )\n severity = _pick(\n threat.get('severity'), alert_data.get('threat_severity'),\n alert_data.get('severity'), threat.get('level'),\n alert_data.get('level'), 'medium',\n )\n\n req_line = _pick(http.get('reqs_line'), alert_data.get('req_line'), alert_data.get('net_http_reqs_line'))\n req_header = _pick(http.get('reqs_header'), alert_data.get('req_header'), alert_data.get('net_http_reqs_header'))\n req_body = _pick(http.get('req_body'), alert_data.get('req_body'), alert_data.get('net_http_reqs_body'))\n resp_line = _pick(http.get('resp_line'), alert_data.get('rsp_line'), alert_data.get('resp_line'),\n alert_data.get('net_http_resp_line'))\n resp_header = _pick(http.get('resp_header'), alert_data.get('rsp_header'), alert_data.get('resp_header'),\n alert_data.get('net_http_resp_header'))\n resp_body = _pick(http.get('resp_body'), alert_data.get('rsp_body'), alert_data.get('resp_body'),\n alert_data.get('net_http_resp_body'))\n status = _pick(http.get('status'), alert_data.get('http_status'),\n alert_data.get('net_http_status'), alert_data.get('rsp_status_code'), 0)\n\n host = _pick(http.get('reqs_host'), alert_data.get('url_host'), http.get('domain'),\n alert_data.get('req_host'), alert_data.get('net_http_reqs_host'), dst_ip)\n raw_url = _pick(http.get('raw_url'), http.get('url'),\n alert_data.get('url_path'),\n alert_data.get('net_http_url'), alert_data.get('req_http_url'),\n alert_data.get('uri'))\n url = ''\n if host and raw_url:\n scheme = 'https' if net.get('is_https') else 'http'\n url = raw_url if str(raw_url).startswith(('http://', 'https://')) else f'{scheme}://{host}{raw_url}'\n elif raw_url and str(raw_url).startswith(('http://', 'https://')):\n url = raw_url\n elif raw_url:\n url = raw_url\n\n payload = f'请求行: {req_line}\\n请求头: {req_header}\\n请求体: {req_body}'\n response = f'状态行: {resp_line}\\n响应头: {resp_header}\\n响应体: {resp_body}'\n threat_result = _pick(threat.get('result'), alert_data.get('threat_result'))\n threat_msg = _pick(threat.get('msg'), alert_data.get('threat_msg'))\n\n log_text = (\n f'[告警基本信息]\\n'\n f'告警类型: {alert_type}\\n严重级别: {severity}\\n'\n f'源地址: {src_ip}:{src_port}\\n目的地址: {dst_ip}:{dst_port}\\n'\n f'协议: {protocol}\\nURL: {url}\\nHTTP状态码: {status}\\n'\n f'TDP判定: {threat_result}\\nTDP消息: {threat_msg}\\n\\n'\n f'[HTTP请求内容]\\n{payload}\\n\\n'\n f'[HTTP响应内容]\\n{response}'\n )\n\n vuln_text = '\\n'.join(str(item) for item in [\n threat_msg, threat.get('topic', ''),\n alert_data.get('data', ''), url,\n json.dumps(threat.get('tag', []), ensure_ascii=False),\n ] if item)\n vuln_matches = sorted(set(re.findall(r'\\b(?:CVE|CNVD|CNNVD|XVE)-[A-Za-z0-9._-]+\\b', vuln_text, flags=re.I)))\n\n iocs = []\n for candidate in [src_ip, dst_ip]:\n if candidate:\n iocs.append({'type': 'ip', 'value': candidate})\n if url:\n iocs.append({'type': 'url', 'value': url})\n if host and not re.match(r'^\\d{1,3}(?:\\.\\d{1,3}){3}(?::\\d+)?$', str(host)):\n iocs.append({'type': 'domain', 'value': str(host).split(':')[0]})\n\n return {\n 'src_ip': src_ip, 'dst_ip': dst_ip, 'src_port': src_port, 'dst_port': dst_port,\n 'protocol': protocol, 'payload': payload, 'response': response,\n 'url': url, 'status': status,\n 'alert_type': alert_type, 'severity': severity,\n 'vuln_id': vuln_matches[0] if vuln_matches else '',\n 'vuln_candidates': vuln_matches,\n 'threat_result': threat_result, 'threat_msg': threat_msg,\n 'failed_by': threat.get('failed_by', []),\n 'asset_ip': assets.get('ip', ''), 'asset_name': assets.get('name', []),\n 'iocs': iocs, 'log_text': log_text,\n }\n\n\ndef _prepare_intel(parsed):\n \"\"\"Mirrors tdp_alert_triage.prepare_intel. Pre-fetches IP/domain/URL threat intel\n and CVE info so the parallel LLM tasks have concrete context to consume.\n \"\"\"\n iocs = parsed.get('iocs', [])\n intel_results = []\n seen = set()\n for ioc in iocs:\n ioc_type = ioc.get('type', '')\n ioc_value = str(ioc.get('value', '')).strip()\n key = (ioc_type, ioc_value)\n if not ioc_value or key in seen:\n continue\n seen.add(key)\n if ioc_type == 'ip':\n if not _is_public_ip(ioc_value):\n continue\n r = tool.run_safe('threatbook_ip_query', ip=ioc_value)\n if r['success']:\n intel_results.append({'source': 'threatbook', 'type': 'ip',\n 'value': ioc_value, 'result': r['text']})\n elif ioc_type == 'domain':\n r = tool.run_safe('threatbook_domain_query', domain=ioc_value)\n if r['success']:\n intel_results.append({'source': 'threatbook', 'type': 'domain',\n 'value': ioc_value, 'result': r['text']})\n elif ioc_type == 'url':\n r = tool.run_safe('threatbook_url_query', url=ioc_value)\n if r['success']:\n intel_results.append({'source': 'threatbook', 'type': 'url',\n 'value': ioc_value, 'result': r['text']})\n\n vuln_info = {}\n vuln_id = parsed.get('vuln_id', '')\n if vuln_id:\n r = tool.run_safe('__mcp_vuln_query', vuln_id=vuln_id)\n if r['success']:\n try:\n obj = r.get('obj')\n if isinstance(obj, str):\n obj = json.loads(obj)\n vuln_info = obj if isinstance(obj, dict) else {'raw_result': r.get('text', '')}\n except Exception:\n vuln_info = {'raw_result': r.get('text', '')}\n\n intel_content = '\\n'.join(\n f\"[{i['source']}/{i['type']}] {i['value']}\\n{i['result']}\" for i in intel_results\n ) or '（无可用情报数据）'\n vuln_content = (\n json.dumps(vuln_info, ensure_ascii=False, indent=2)\n if vuln_info else '（无可用漏洞情报数据）'\n )\n return intel_results, intel_content, vuln_info, vuln_content\n\n\n# ── 4 LLM analysis branches (sharing one run-wide request budget) ───────────\n\ndef _ask_llm(prompt):\n \"\"\"Wrap ``llm.ask`` with the workflow-wide timeout + retry budget.\n\n Centralizing this avoids a hung provider request (no TCP timeout from\n upstream) from blocking a worker thread indefinitely. Any call that does\n not need bespoke parameters should go through here.\n \"\"\"\n kwargs = {\n 'timeout_s': LLM_CALL_TIMEOUT_S,\n 'max_retries': LLM_CALL_MAX_RETRIES,\n }\n if _llm_slots is None:\n return llm.ask(prompt, **kwargs)\n with _llm_slots:\n return llm.ask(prompt, **kwargs)\n\n\ndef _llm_survey(log_text, intel_content):\n prompt = f'''你是一个专业的Web日志分析专家。请总结以下IP的情报数据中的空间测绘信息。\n1. 如果该IP没有测绘信息，则不列出。\n2. 如果IP有测绘信息，则以简短的语言对该IP的测绘信息进行总结，关键说明ip的标签和测绘信息显示有哪些服务或者应用资产。\n3. 多个IP的测绘信息以无序列表显示，每个ip数据描述占一行数据。\n4. 不需要生成其他额外的补充信息。\n\n## 情报参考信息\n{intel_content}\n\n## 用户的原始输入日志\n{log_text}\n'''\n return _strip_think(_ask_llm(prompt))\n\n\ndef _llm_cve_related(log_text):\n prompt = f'''请从以下的日志数据中提取漏洞编号。\n要求：\n1. 仅从日志文本中识别漏洞编号，不要做任何推测。\n2. 如果日志中存在漏洞编号，则用简短语言描述，如：\"日志中存在漏洞编号：CVE-****-****\"。\n3. 如果日志中不存在漏洞编号，则输出：\"日志中无关联漏洞情报\"。\n\n日志数据如下：\n{log_text}\n'''\n return _strip_think(_ask_llm(prompt))\n\n\ndef _llm_cve_info(log_text, vuln_content):\n prompt = f'''你是一个专业的Web日志分析专家。参考情报信息中的漏洞数据，简要说明关联的CVE漏洞信息。\n1. 不要输出任何解释说明，只输出漏洞基本信息。不需要生成漏洞的处置建议或修复措施等。\n\n## 情报参考信息\n{vuln_content}\n\n## 用户的原始输入日志\n{log_text}\n'''\n return _strip_think(_ask_llm(prompt))\n\n\ndef _llm_payload_analysis(log_text):\n prompt = f'''你是一个专业的Web日志分析专家。根据用户输入的日志进行攻击负载分析。\n1. 首先分析日志中是否包含攻击负载，并给出判定依据。\n2. 不要进行攻击意图分析、攻击影响分析。\n3. 用简短的语言在一段话中进行描述。\n\n## 用户的原始输入日志:\n{log_text}\n'''\n return _strip_think(_ask_llm(prompt))\n\n\ndef _parallel_4_branches(parsed, intel_content, vuln_content):\n \"\"\"Submit 4 logical branches; _ask_llm enforces the run-wide provider budget.\"\"\"\n log_text = parsed.get('log_text', '')\n with ThreadPoolExecutor(max_workers=4, thread_name_prefix='triage_branch') as pool:\n futs = {\n 'survey_result': pool.submit(_llm_survey, log_text, intel_content),\n 'cve_related_result': pool.submit(_llm_cve_related, log_text),\n 'cve_info_result': pool.submit(_llm_cve_info, log_text, vuln_content),\n 'payload_analysis_result': pool.submit(_llm_payload_analysis, log_text),\n }\n out = {}\n for name, fut in futs.items():\n try:\n out[name] = fut.result()\n except Exception as e:\n print(f'[triage] WARNING: branch {name} failed: {e}')\n out[name] = ''\n return out\n\n\n# ── Join-point LLM analyses (attack_analysis_result -> verdict -> title) ──────\n\ndef _llm_attack_analysis(log_text):\n prompt = f'''你是一名专业且经验丰富的网络安全分析师和Web日志分析专家，你对HTTP协议以及Web攻击有着深入的理解，并且你能够快速识别和应对各种网络威胁。你的任务是对提供的HTTP请求与响应内容进行详细的专业分析，并判断日志请求的攻击状态。\n\n请严格遵循以下指令进行思考和分析：\n1. 分别判断是否为攻击（攻击/非攻击/未知）以及攻击结果（攻击成功/攻击失败/未知）。\n2. 从日志中提取出\"HTTP请求内容\"和\"HTTP响应内容\"。请注意，HTTP请求内容和HTTP响应内容是分开的，请不要混淆，有些日志中没有包含HTTP响应内容，请不要将HTTP请求内容和HTTP响应内容混淆。分析后请你记住哪些是HTTP请求内容，哪些是HTTP响应内容。\n3. 请检查HTTP响应状态码，2xx或者3xx状态码都代表本次HTTP请求成功，4xx或者5xx状态码大多数情况下都代表请求失败，只有在请求成功的情况下才能对攻击是否成功进行后续判断。\n\n各攻击状态的定义以及判定标准：\n1. 攻击成功：\n(1) 首先分析日志中是否含有清晰的\"HTTP响应内容\"，如果日志中没有\"HTTP响应内容\"，则肯定不属于攻击成功。\n(2) 如果日志中未提供\"HTTP响应内容\"，即使HTTP请求内容中包含攻击者预期的结果，也不能判定为攻击成功。\n(3) 从日志中提取出\"HTTP请求内容\"和\"HTTP响应内容\"。请深入分析\"HTTP响应内容\"，并判定其是否为\"HTTP请求内容\"攻击成功时的预期结果，这是判定攻击成功的强依据。请注意，HTTP响应码200仅表示网络连接成功，不代表攻击攻击成功。\n(4) 分析HTTP请求内容和HTTP响应内容，只有当HTTP响应内容中明确包含攻击载荷在目标机器上成功执行的证据，并且HTTP请求内容中包含攻击载荷的特征，则判定为\"攻击成功\"。\n(5) 请注意：攻击成功的判定必须包含HTTP响应内容。如果不包含HTTP响应内容，则肯定不属于攻击成功。\n(6) 请注意：如果不包含HTTP响应内容，即使HTTP请求内容是攻击，这也不属于攻击成功。\n2. 攻击失败：\n(1) 分析HTTP请求内容和HTTP响应内容，如果HTTP响应内容中明确包含攻击载荷在目标机器上执行失败或者被阻止的证据，并且HTTP请求内容中包含攻击载荷的特征，则判定为\"攻击失败\"。\n(2) 攻击失败的判定必须包含HTTP响应内容。如果不包含HTTP响应内容，则肯定不属于攻击失败。\n3. 攻击：\n(1) 在\"HTTP请求内容\"或\"HTTP响应内容\"中发现任何证明存在攻击意图的证据，即可判定为存在攻击行为。但如果不符合上述的攻击成功或者攻击失败的标准，则\"攻击状态\"为\"攻击\"。\n(2) 请注意：如果日志中只提供了\"HTTP请求内容\"，且没有提供\"HTTP响应内容\"，且HTTP的请求内容分析中是包含攻击行为的，则\"攻击状态\"为\"攻击\"。\n4. 未知：\n(1) 如果不能100%确定HTTP通信的攻击结果，那么请在\"攻击状态\"处给出\"未知\"。\n(2) 请注意：如果在你给的判定原因中存在\"可能\"等不确定词汇，都代表你不能对你的结论100%确定，那么请在\"攻击状态\"处给出\"未知\"。\n5. 安全：\n(1) 如果\"HTTP请求内容\"和\"HTTP响应内容\"中都没有任何攻击意图的证据，那么请在\"攻击状态\"处给出\"安全\"。\n\n## 日志内容\n{log_text}\n\n## 输出要求\n请按下列结构输出（中文）：\n1. 是否攻击: [攻击/非攻击/未知]\n2. 攻击结果: [攻击成功/攻击失败/未知]\n3. 判定依据: 简要说明请求与响应的关键证据\n4. 详细分析: 不超过200字\n'''\n return _strip_think(_ask_llm(prompt))\n\n\ndef _normalize_triage_outcome(attack_verdict, attack_success):\n verdict = str(attack_verdict or '').strip().lower()\n result = str(attack_success or '').strip().lower()\n if verdict not in {'attack', 'non_attack', 'unknown'}:\n verdict = 'unknown'\n if result not in {'success', 'failed', 'unknown'}:\n result = 'unknown'\n if verdict != 'attack':\n result = 'unknown'\n return verdict, result\n\n\ndef _triage_outcome_key(attack_verdict, attack_success):\n verdict, result = _normalize_triage_outcome(attack_verdict, attack_success)\n if verdict == 'attack':\n if result == 'success':\n return 'attack_success'\n if result == 'failed':\n return 'attack_failed'\n return 'attack'\n if verdict == 'non_attack':\n return 'non_attack'\n return 'unknown'\n\n\ndef _llm_attack_outcome(attack_analysis_result):\n prompt = f'''你是一个专业的 Web 日志分析专家。根据参考信息，输出且仅输出一个 JSON 对象：\n{{\n \"triage_attack_verdict\": \"attack | non_attack | unknown\",\n \"triage_attack_success\": \"success | failed | unknown\"\n}}\n\n字段规则：\n1. triage_attack_verdict 只表示是否为攻击：attack=攻击，non_attack=非攻击，unknown=无法判断。\n2. triage_attack_success 只表示攻击结果：success=攻击成功，failed=攻击失败，unknown=无法判断结果。\n3. 只有 triage_attack_verdict=attack 时，攻击结果才允许为 success 或 failed。\n4. 当 triage_attack_verdict 为 non_attack 或 unknown 时，triage_attack_success 必须为 unknown。\n5. 不要输出解释、Markdown 或其他字段。\n\n## 日志分析结果\n{attack_analysis_result}\n'''\n raw = _strip_think(_ask_llm(prompt)).strip()\n match = re.search(r'\\{.*?\\}', raw, flags=re.S)\n try:\n parsed = json.loads(match.group(0)) if match else {}\n except Exception:\n parsed = {}\n verdict, result = _normalize_triage_outcome(\n parsed.get('triage_attack_verdict'),\n parsed.get('triage_attack_success'),\n )\n return {\n 'triage_attack_verdict': verdict,\n 'triage_attack_success': result,\n }\n\n\ndef _llm_report_title(alert_type, attack_verdict, attack_success, attack_analysis_result):\n prompt = f'''你是一个专业的Web日志分析专家。请基于以下分析结果，生成一份不超过 30 字的中文报告标题。\n要求：\n1. 标题必须能体现\"攻击类型\"或\"攻击结果分析的结论\"。\n2. 不要带书名号、引号或其他标点。\n3. 只输出标题本身，不要任何解释或说明。\n\n## 攻击类型\n{alert_type}\n\n## 是否攻击\n{attack_verdict}\n\n## 攻击结果\n{attack_success}\n\n## 攻击分析结果\n{attack_analysis_result}\n'''\n raw = _strip_think(_ask_llm(prompt)).strip()\n outcome_cn = OUTCOME_CN.get(\n _triage_outcome_key(attack_verdict, attack_success),\n '未知',\n )\n return raw.splitlines()[0].strip(' \"\\'《》[]【】') if raw else f'{alert_type} - {outcome_cn}'\n\n\ndef _clip_text(value, limit=3000):\n text = str(value or '').strip()\n if len(text) > limit:\n return text[:limit] + '\\n...（已截断）'\n return text or '未提供'\n\n\ndef _fence_text(value):\n text = _clip_text(value, 6000)\n return text.replace('```', '``\\\\u200b`')\n\n\ndef _extract_tagged_triage_report(text):\n text = _strip_think(text)\n m = re.search(r']*>[\\s\\S]*?', text, flags=re.I)\n return m.group(0).strip() if m else text.strip()\n\n\ndef _is_valid_triage_report(markdown):\n text = str(markdown or '')\n if not re.search(r']*version=[\"\\']soc\\.triage\\.markdown\\.v1[\"\\'][^>]*>', text, flags=re.I):\n return False\n if not re.search(r'', text, flags=re.I):\n return False\n for tag in TRIAGE_REPORT_TAGS:\n if not re.search(rf'<{tag}\\b[^>]*>', text, flags=re.I):\n return False\n if not re.search(rf'', text, flags=re.I):\n return False\n return True\n\n\ndef _is_current_triage_fields(fields):\n if not isinstance(fields, dict):\n return False\n verdict = fields.get('triage_attack_verdict')\n result = fields.get('triage_attack_success')\n normalized = _normalize_triage_outcome(verdict, result)\n return (\n (verdict, result) == normalized\n and all(key in fields for key in TRIAGE_FIELDS)\n and _is_valid_triage_report(fields.get('triage_report'))\n )\n\n\ndef _format_intel_brief(intel_results):\n if not intel_results:\n return '未查询到外部威胁情报。'\n lines = []\n for intel in intel_results[:6]:\n lines.append(f\"- {intel.get('source', 'intel')} / {intel.get('type', 'ioc')}: {intel.get('value', '')} => {_clip_text(intel.get('result'), 500)}\")\n return '\\n'.join(lines)\n\n\ndef _build_default_tagged_triage_report(parsed, intel_results, vuln_info, branches,\n attack_analysis_result, attack_verdict,\n attack_success, report_title, risk_level):\n outcome = _triage_outcome_key(attack_verdict, attack_success)\n verdict_cn = OUTCOME_CN.get(outcome, '未知')\n attack_verdict_cn = ATTACK_VERDICT_CN.get(attack_verdict, '未知')\n attack_success_cn = ATTACK_SUCCESS_CN.get(attack_success, '未知')\n title = report_title or f'{parsed.get(\"alert_type\", \"Web日志告警\")} - {verdict_cn}'\n payload = _fence_text(parsed.get('payload', ''))\n response = _fence_text(parsed.get('response', ''))\n url = parsed.get('url') or '未提供'\n threat_msg = parsed.get('threat_msg') or '未提供'\n status = parsed.get('status') or '未提供'\n survey = _clip_text(branches.get('survey_result'), 1500)\n cve_related = _clip_text(branches.get('cve_related_result'), 1500)\n cve_info = _clip_text(branches.get('cve_info_result'), 1500)\n payload_analysis = _clip_text(branches.get('payload_analysis_result'), 1500)\n attack_analysis = _clip_text(attack_analysis_result, 1500)\n intel_brief = _format_intel_brief(intel_results)\n vuln_brief = _clip_text(json.dumps(vuln_info, ensure_ascii=False, indent=2), 1800) if vuln_info else '未查询到漏洞详情。'\n\n if outcome == 'attack_success':\n recommendation = '立即核查目标资产是否产生异常文件、进程、账号或敏感数据访问记录，并按成功入侵事件升级处置。'\n elif outcome == 'attack_failed':\n recommendation = '保留拦截与响应证据，复核同源后续请求，并确认防护策略是否持续生效。'\n elif outcome == 'non_attack':\n recommendation = '作为低风险事件留痕，结合资产白名单或业务访问记录确认是否可降噪。'\n else:\n recommendation = '补齐目标 Web 日志、响应体、主机侧进程和文件证据后再确认攻击成功性。'\n\n return f'''\n\n\n# {title}\n\n\n\n- 是否攻击：{attack_verdict_cn}\n- 攻击结果：{attack_success_cn}\n- 研判结论：{verdict_cn}\n- 风险等级：{risk_level}\n- 告警类型：{parsed.get('alert_type', 'unknown')}\n- 源 IP：{parsed.get('src_ip', 'N/A')}:{parsed.get('src_port', 'N/A')}\n- 目标资产：{parsed.get('dst_ip', 'N/A')}:{parsed.get('dst_port', 'N/A')}\n- URL：{url}\n- 响应码：{status}\n\n\n\n## 分析步骤\n\n### 1. 日志类型分析\n该告警按 Web 日志处理，已提取 HTTP 请求、响应、源地址、目标资产、URL、响应码和 TDP 判定字段。\n\n### 2. 情报信息\n{intel_brief}\n\n### 3. 测绘信息\n{survey}\n\n### 4. 告警关联漏洞情报\n{cve_related}\n\n### 5. 攻击负载分析\n{payload_analysis}\n\n### 6. 攻击分析结果\n{attack_analysis}\n\n\n\n## 研判结论\n当前研判结论为 **{verdict_cn}**，风险等级为 **{risk_level}**。TDP 消息为：{threat_msg}\n\n\n\n## 攻击payload\n\n```http\n{payload}\n```\n\n\n\n## 具体含义解释\n\n1. 请求命中的告警类型为 {parsed.get('alert_type', 'unknown')}。\n2. 请求 URL 为 {url}，需要结合参数、请求体和目标业务判断攻击意图。\n3. Payload 分析结果：{payload_analysis}\n\n\n\n## 响应证据\n\n```http\n{response}\n```\n\n响应码为 {status}。如果响应体未提供或没有执行成功证据，则不能仅凭请求侧 payload 判定攻击成功。\n\n\n\n## 重要证据\n\n1. 源地址：{parsed.get('src_ip', 'N/A')}:{parsed.get('src_port', 'N/A')}。\n2. 目标资产：{parsed.get('dst_ip', 'N/A')}:{parsed.get('dst_port', 'N/A')}。\n3. TDP 判定：{parsed.get('threat_result', '未提供')}；TDP 消息：{threat_msg}。\n4. 漏洞详情：{vuln_brief}\n\n\n\n## 处置建议\n\n1. {recommendation}\n2. 检索同源 IP、同一 dedup_key、同一 URL 或同一漏洞特征的横向告警。\n3. 结合目标资产 Web 访问日志、主机审计、EDR 与 WAF 日志补齐证据链。\n\n\n'''\n\n\ndef _llm_triage_report_markdown(parsed, intel_results, vuln_info, branches,\n attack_analysis_result, attack_verdict,\n attack_success, report_title, risk_level):\n verdict_cn = OUTCOME_CN.get(\n _triage_outcome_key(attack_verdict, attack_success),\n '未知',\n )\n context = json.dumps({\n 'report_title': report_title,\n 'triage_attack_verdict': attack_verdict,\n 'triage_attack_success': attack_success,\n 'verdict_cn': verdict_cn,\n 'risk_level': risk_level,\n 'alert': {\n 'alert_type': parsed.get('alert_type'),\n 'severity': parsed.get('severity'),\n 'src_ip': parsed.get('src_ip'),\n 'src_port': parsed.get('src_port'),\n 'dst_ip': parsed.get('dst_ip'),\n 'dst_port': parsed.get('dst_port'),\n 'url': parsed.get('url'),\n 'status': parsed.get('status'),\n 'threat_result': parsed.get('threat_result'),\n 'threat_msg': parsed.get('threat_msg'),\n 'payload': parsed.get('payload'),\n 'response': parsed.get('response'),\n },\n 'survey_result': branches.get('survey_result'),\n 'cve_related_result': branches.get('cve_related_result'),\n 'cve_info_result': branches.get('cve_info_result'),\n 'payload_analysis_result': branches.get('payload_analysis_result'),\n 'attack_analysis_result': attack_analysis_result,\n 'intel_results': intel_results,\n 'vuln_info': vuln_info,\n }, ensure_ascii=False, indent=2)\n\n prompt = f'''你是一名资深 SOC 告警研判分析师。请根据输入上下文，生成一份供前端直接渲染的 SOC 告警研判报告 markdown。\n\n硬性要求：\n1. 只输出带语义标签的 markdown，不要输出 JSON，不要解释规则。\n2. 根标签必须是 。\n3. 必须按顺序输出并完整闭合这些标签：\n 、、、、、、、、。\n4. 标签外不得输出正文内容。标签内可以使用 markdown 标题、列表、引用、代码块。\n5. 段落标题必须贴近前端展示模板：分析步骤、研判结论、攻击payload、具体含义解释、响应证据、重要证据、处置建议。\n6. 如果没有 HTTP 响应体或没有明确响应证据，不得判定为攻击成功；需要写明“当前日志未提供有效响应证据”。\n7. 不要编造输入中不存在的 IP、域名、URL、CVE、账号、文件路径或响应内容。\n8. 攻击 payload 和响应证据必须分别放在对应标签中，不要混淆请求与响应。\n9. 研判结论必须严格遵循上下文中的 triage_attack_verdict 和 triage_attack_success，不得使用原始告警同名字段替代。\n\nFew-shot 示例 1：攻击成功\n\n\n\n# 敏感文件泄露攻击成功分析报告\n\n\n\n- 研判结论：攻击成功\n- 风险等级：High\n- 告警类型：敏感文件访问\n- 源 IP：203.0.113.10:42131\n- 目标资产：198.51.100.20:80\n- URL：http://example.com/api/.env\n- 响应码：200\n\n\n\n## 分析步骤\n\n### 1. 日志类型分析\n该日志包含 HTTP 请求路径、响应码和响应体，可用于判断敏感文件是否被返回。\n\n### 2. 情报信息\n源 IP 命中扫描源标签，风险高。\n\n### 3. 测绘信息\n目标为公网 Web 服务，存在敏感路径暴露风险。\n\n### 4. 告警关联漏洞情报\n该行为与环境变量文件泄露场景一致。\n\n### 5. 攻击负载分析\n攻击者直接请求 /api/.env，目标是读取环境变量配置。\n\n### 6. 攻击分析结果\n响应码为 200，响应体中出现 DB_PASSWORD，支持攻击成功。\n\n\n\n## 研判结论\n攻击者成功读取敏感配置文件，响应体中包含数据库密码字段，结论为攻击成功。\n\n\n\n## 攻击payload\n\n```http\nGET /api/.env HTTP/1.1\nHost: example.com\n```\n\n\n\n## 具体含义解释\n\n1. /api/.env 是常见环境变量文件路径。\n2. 攻击者通过 GET 请求尝试直接读取配置文件。\n3. 该路径若返回真实内容，通常意味着敏感文件暴露。\n\n\n\n## 响应证据\n\n```http\nHTTP/1.1 200 OK\n\nDB_PASSWORD=example-secret\n```\n\n响应体出现 DB_PASSWORD，证明敏感配置内容已被返回。\n\n\n\n## 重要证据\n\n1. 请求路径为 /api/.env。\n2. 响应码为 200。\n3. 响应体包含 DB_PASSWORD。\n\n\n\n## 处置建议\n\n1. 立即下线或限制敏感文件访问。\n2. 轮换可能泄露的密钥和数据库密码。\n3. 检索同源 IP 和同路径访问记录。\n\n\n\n\nFew-shot 示例 2：攻击失败\n\n\n\n# SQL注入攻击失败分析报告\n\n\n\n- 研判结论：攻击失败\n- 风险等级：Medium\n- 告警类型：SQL注入\n- 源 IP：203.0.113.44:51002\n- 目标资产：198.51.100.30:443\n- URL：https://shop.example.com/item?id=1\n- 响应码：403\n\n\n\n## 分析步骤\n\n### 1. 日志类型分析\n该日志包含请求参数和响应码，能够确认请求侧存在 SQL 注入尝试。\n\n### 2. 情报信息\n源 IP 暂无高置信恶意标签。\n\n### 3. 测绘信息\n目标为公网电商 Web 服务。\n\n### 4. 告警关联漏洞情报\n当前日志未提供可确认具体 CVE 的证据。\n\n### 5. 攻击负载分析\n请求参数中包含 union select，存在明显 SQL 注入意图。\n\n### 6. 攻击分析结果\n响应码为 403，响应体显示请求被阻断，不支持攻击成功。\n\n\n\n## 研判结论\n该请求存在 SQL 注入攻击意图，但响应显示被拒绝，当前判断为攻击失败。\n\n\n\n## 攻击payload\n\n```http\nGET /item?id=1 union select user HTTP/1.1\nHost: shop.example.com\n```\n\n\n\n## 具体含义解释\n\n1. union select 是典型 SQL 注入关键字组合。\n2. 攻击者尝试拼接查询以读取数据库用户信息。\n3. 该 payload 证明攻击意图，但不等同于成功执行。\n\n\n\n## 响应证据\n\n```http\nHTTP/1.1 403 Forbidden\n\nblocked by waf\n```\n\n响应状态和内容说明请求被拦截，未见数据泄露或执行成功证据。\n\n\n\n## 重要证据\n\n1. 请求参数包含 union select。\n2. 响应码为 403。\n3. 响应体显示 blocked by waf。\n\n\n\n## 处置建议\n\n1. 保留 WAF 拦截证据。\n2. 检查同源 IP 是否持续尝试其他注入 payload。\n3. 确认目标接口参数化查询和安全策略仍然有效。\n\n\n\n\n## 待研判上下文\n```json\n{context}\n```\n\n请输出最终报告：\n'''\n return _extract_tagged_triage_report(_ask_llm(prompt))\n\n\ndef _generate_triage_report(parsed, intel_results, vuln_info, branches,\n attack_analysis_result, attack_verdict,\n attack_success, report_title):\n # Aggregate everything into tagged markdown for frontend rendering.\n # The markdown is returned through `triage_report` and is not written as a\n # per-alert file; leader/follower/cache-hit paths reuse the same field.\n outcome = _triage_outcome_key(attack_verdict, attack_success)\n verdict_cn = OUTCOME_CN.get(outcome, '未知')\n risk_level = OUTCOME_RISK.get(outcome, 'Medium')\n\n if not report_title:\n report_title = f'{parsed.get(\"alert_type\", \"Web日志告警\")} - {verdict_cn}'\n\n try:\n triage_report = _llm_triage_report_markdown(\n parsed, intel_results, vuln_info, branches,\n attack_analysis_result, attack_verdict, attack_success,\n report_title, risk_level,\n )\n except Exception as e:\n print(f'[triage] WARNING: triage_report LLM generation failed: {e}')\n triage_report = ''\n\n if not _is_valid_triage_report(triage_report):\n print('[triage] WARNING: triage_report missing required semantic tags; using deterministic fallback')\n triage_report = _build_default_tagged_triage_report(\n parsed, intel_results, vuln_info, branches,\n attack_analysis_result, attack_verdict, attack_success,\n report_title, risk_level,\n )\n\n return triage_report, report_title, risk_level\n\n\ndef _triage_single_alert(alert):\n \"\"\"End-to-end inline triage for a single alert. Returns triage_fields dict.\n\n No file I/O — the full markdown report lives in the returned `triage_report` field\n and is broadcast to followers / persisted via `triage_cache.pkl`.\n \"\"\"\n parsed = _parse_alert(alert)\n intel_results, intel_content, vuln_info, vuln_content = _prepare_intel(parsed)\n\n branches = _parallel_4_branches(parsed, intel_content, vuln_content)\n\n attack_analysis_result = _llm_attack_analysis(parsed['log_text'])\n outcome = _llm_attack_outcome(attack_analysis_result)\n attack_verdict = outcome['triage_attack_verdict']\n attack_success = outcome['triage_attack_success']\n report_title = _llm_report_title(\n parsed.get('alert_type', 'unknown'),\n attack_verdict,\n attack_success,\n attack_analysis_result,\n )\n\n triage_report, report_title, risk_level = _generate_triage_report(\n parsed, intel_results, vuln_info, branches,\n attack_analysis_result, attack_verdict, attack_success, report_title,\n )\n\n return {\n 'triage_attack_verdict': attack_verdict,\n 'triage_attack_success': attack_success,\n 'risk_level': risk_level,\n 'report_title': report_title,\n 'triage_report': triage_report,\n }\n\n\n# ── Main: leader/follower batch deduplication ────────────────────────────────\n# When the input batch contains multiple alerts sharing the same dedup_key\n# (e.g. upstream emits is_duplicate=True alerts in the same batch, or LSH\n# clustering produces several alerts per cluster), we only triage the LEADER\n# (first occurrence of each dedup_key). All FOLLOWERS in the same group reuse\n# the leader's triage result without invoking the LLM again.\n#\n# Work unit types:\n# ('dk', dedup_key, leader_idx) — group of 1+ alerts sharing dedup_key\n# ('nokey', None, alert_idx) — single alert with no dedup_key\n# (cannot be deduplicated, always triaged)\n\ndef _bump_verdict(verdict):\n with results_lock:\n stats['verdict_counts'][verdict] = stats['verdict_counts'].get(verdict, 0) + 1\n\n\n_UNKNOWN_TRIAGE = {\n 'triage_attack_verdict': 'unknown',\n 'triage_attack_success': 'unknown',\n 'risk_level': 'Medium',\n 'report_title': '',\n 'triage_report': '',\n}\n\n\ndef _apply_triage_fields(record, triage_fields):\n verdict = triage_fields.get('triage_attack_verdict', 'unknown')\n result = triage_fields.get('triage_attack_success', 'unknown')\n if verdict not in {'attack', 'non_attack', 'unknown'}:\n verdict = 'unknown'\n if result not in {'success', 'failed', 'unknown'} or verdict != 'attack':\n result = 'unknown'\n record['triage_attack_verdict'] = verdict\n record['triage_attack_success'] = result\n for key in ('risk_level', 'report_title', 'triage_report'):\n if key in triage_fields:\n record[key] = triage_fields[key]\n\n\ndef _process_unit(unit_type, dedup_key, leader_idx):\n \"\"\"Triage one unique work unit. Returns (key, triage_fields, source, ms, error).\"\"\"\n leader_alert = enriched_alerts[leader_idx]\n t0 = time.time()\n\n # 1) cache lookup for dedup_key units\n if unit_type == 'dk':\n cached = triage_cache_snapshot.get(dedup_key)\n if cached:\n if _is_current_triage_fields(cached):\n with results_lock:\n stats['cache_hit'] += 1\n _bump_verdict(_triage_outcome_key(\n cached.get('triage_attack_verdict'),\n cached.get('triage_attack_success'),\n ))\n return dedup_key, cached, 'cache', 0, None\n print(f'[triage_cache] stale entry for dedup_key={dedup_key}: '\n 'missing current two-field triage schema/report; treating as cache miss')\n\n # 2) cache miss or no-key → run full inline triage on the leader\n try:\n triage_fields = _triage_single_alert(leader_alert)\n ms = round((time.time() - t0) * 1000)\n with results_lock:\n stats['triaged'] += 1\n if dedup_key:\n new_results[dedup_key] = triage_fields\n _bump_verdict(_triage_outcome_key(\n triage_fields.get('triage_attack_verdict'),\n triage_fields.get('triage_attack_success'),\n ))\n source = 'triaged' if unit_type == 'dk' else 'no_dedup_key_triaged'\n return (dedup_key if unit_type == 'dk' else leader_idx), triage_fields, source, ms, None\n except Exception as e:\n import traceback\n ms = round((time.time() - t0) * 1000)\n with results_lock:\n stats['triage_failed'] += 1\n _bump_verdict('unknown')\n err = str(e)[:500]\n print(f'[triage] leader_idx={leader_idx} FAILED: {e}\\n{traceback.format_exc()}')\n source = 'failed' if unit_type == 'dk' else 'no_dedup_key_failed'\n return (dedup_key if unit_type == 'dk' else leader_idx), dict(_UNKNOWN_TRIAGE), source, ms, err\n\n\n\n_batch_lease_fd = inputs.get('_batch_lease_fd')\nif type(_batch_lease_fd) is not int or _batch_lease_fd < 0:\n _batch_lease_fd = None\n_batch_lease_token = str(inputs.get('batch_lease_token') or '')\nif bool(inputs.get('cursor_enabled')) and _batch_lease_fd is None:\n _batch_lease_fd, _batch_lease_token = _acquire_batch_lease()\ntry:\n enriched_alerts = list(inputs.get('enriched_alerts', []) or [])\n concurrency = min(5, max(1, int(inputs.get('concurrency', 1))))\n _llm_slots = threading.BoundedSemaphore(concurrency)\n max_triage_cache_size = int(inputs.get('max_triage_cache_size', 100000))\n if max_triage_cache_size < 1:\n max_triage_cache_size = 100000\n\n # Group by dedup_key\n groups = {} # dedup_key -> [alert_index, ...]\n no_key_indices = [] # alerts with no dedup_key\n for i, a in enumerate(enriched_alerts):\n dk = a.get('dedup_key', '') if isinstance(a, dict) else ''\n if dk:\n groups.setdefault(dk, []).append(i)\n else:\n no_key_indices.append(i)\n\n work_units = (\n [('dk', dk, group_indices[0]) for dk, group_indices in groups.items()]\n + [('nokey', None, idx) for idx in no_key_indices]\n )\n follower_count = sum(len(v) - 1 for v in groups.values())\n\n print(f'[triage] alerts={len(enriched_alerts)} '\n f'→ {len(groups)} unique dedup_keys ({follower_count} followers) + '\n f'{len(no_key_indices)} no-key alerts '\n f'= {len(work_units)} work units; outer_concurrency={concurrency} '\n f'llm_concurrency_limit={concurrency} '\n f'sub_batch_size={TRIAGE_SUB_BATCH_SIZE} '\n f'(all LLM calls share the run-wide budget)')\n\n cache_path, lock_path = _cache_paths()\n cache_batch_lock = _acquire_lock(lock_path)\n try:\n triage_cache_snapshot = _load_cache(cache_path)\n\n # Only warm up the LLM when at least one work unit may actually need it.\n # A unit \"may need\" the LLM if it's a no-key alert OR a dedup_key unit whose\n # entry is not in the cache snapshot. Pure cache-hit batches skip warm-up.\n _needs_llm = any(\n unit_type == 'nokey' or not _is_current_triage_fields(triage_cache_snapshot.get(dk))\n for unit_type, dk, _ in work_units\n )\n if _needs_llm:\n _warmup_llm()\n\n results_lock = threading.Lock()\n new_results = {} # dedup_key -> triage_fields (only for freshly computed leaders)\n group_outcomes = {} # dedup_key -> (triage_fields, source) for broadcasting to followers\n nokey_outcomes = {} # alert_idx -> (triage_fields, source)\n stats = {\n 'total': len(enriched_alerts),\n 'unique_dedup_keys': len(groups),\n 'followers_reused': follower_count,\n 'no_dedup_key_alerts': len(no_key_indices),\n 'work_units': len(work_units),\n 'llm_concurrency_limit': concurrency,\n 'cache_hit': 0,\n 'triaged': 0,\n 'triage_failed': 0,\n 'verdict_counts': {},\n 'cache_size_before': len(triage_cache_snapshot),\n 'cache_size_after': 0,\n 'evicted': 0,\n }\n\n\n t_start = time.time()\n unit_completions = {} # key -> (triage_fields, source, ms, error)\n if work_units:\n completed_count = 0\n for batch_start in range(0, len(work_units), TRIAGE_SUB_BATCH_SIZE):\n work_batch = work_units[batch_start:batch_start + TRIAGE_SUB_BATCH_SIZE]\n worker_count = min(concurrency, len(work_batch))\n with ThreadPoolExecutor(\n max_workers=worker_count,\n thread_name_prefix='stream_triage',\n ) as pool:\n futures = [pool.submit(_process_unit, *u) for u in work_batch]\n for fut in as_completed(futures):\n completed_count += 1\n try:\n key, triage_fields, source, ms, err = fut.result()\n unit_completions[key] = (triage_fields, source, ms, err)\n if source == 'cache':\n group_outcomes[key] = (triage_fields, source)\n elif source.startswith('no_dedup_key'):\n nokey_outcomes[key] = (triage_fields, source)\n else:\n group_outcomes[key] = (triage_fields, source)\n except Exception as e:\n print(f'[triage] WARNING: unexpected worker exception: {e}')\n if completed_count % 5 == 0 or completed_count == len(work_units):\n print(f'[triage] progress {completed_count}/{len(work_units)} '\n f'(cache_hit={stats[\"cache_hit\"]} triaged={stats[\"triaged\"]} '\n f'failed={stats[\"triage_failed\"]})')\n del futures\n\n # ── Apply outcomes back to every alert (broadcast leader → followers) ─────────\n enriched_with_triage = [None] * len(enriched_alerts)\n for i, alert in enumerate(enriched_alerts):\n out = dict(alert) if isinstance(alert, dict) else {'_raw': alert}\n dk = alert.get('dedup_key', '') if isinstance(alert, dict) else ''\n out['has_dedup_key'] = bool(dk)\n\n if dk:\n triage_fields, source = group_outcomes.get(dk, (dict(_UNKNOWN_TRIAGE), 'failed'))\n is_leader = (groups.get(dk, [i])[0] == i)\n # All triage fields are pure data (no file-path references); broadcast as-is.\n _apply_triage_fields(out, triage_fields)\n if not is_leader and source != 'cache':\n out['triage_source'] = 'follower_reused'\n out['triage_status'] = 'reused_from_leader'\n else:\n out['triage_source'] = source\n out['triage_status'] = 'cached' if source == 'cache' else (\n 'ok' if source == 'triaged' else 'failed'\n )\n completion = unit_completions.get(dk)\n if completion and is_leader:\n _, _, ms, err = completion\n if ms:\n out['triage_ms'] = ms\n if err:\n out['triage_error'] = err\n else:\n # no-key alerts: each is its own unit, keyed by alert idx\n triage_fields, source = nokey_outcomes.get(i, (dict(_UNKNOWN_TRIAGE), 'no_dedup_key_failed'))\n _apply_triage_fields(out, triage_fields)\n out['triage_source'] = source\n out['triage_status'] = 'ok' if source == 'no_dedup_key_triaged' else 'failed'\n completion = unit_completions.get(i)\n if completion:\n _, _, ms, err = completion\n if ms:\n out['triage_ms'] = ms\n if err:\n out['triage_error'] = err\n\n enriched_with_triage[i] = out\n\n elapsed_ms = round((time.time() - t_start) * 1000)\n print(f'[triage] all done in {elapsed_ms}ms: cache_hit={stats[\"cache_hit\"]} '\n f'triaged={stats[\"triaged\"]} failed={stats[\"triage_failed\"]} '\n f'followers_reused={stats[\"followers_reused\"]} '\n f'no_dedup_key={stats[\"no_dedup_key_alerts\"]}')\n\n # The batch lock has remained held since the initial load, so reuse\n # that snapshot instead of temporarily retaining a second full cache.\n cache = triage_cache_snapshot\n for k, v in new_results.items():\n if k in cache:\n del cache[k] # LRU touch (move to end on rewrite)\n cache[k] = v\n evicted = _evict_lru(\n cache,\n max_triage_cache_size,\n MAX_TRIAGE_CACHE_BYTES if new_results else None,\n )\n if evicted:\n print(f'[triage_cache] LRU eviction: dropped {evicted} entries '\n f'(max_keys={max_triage_cache_size}, max_bytes={MAX_TRIAGE_CACHE_BYTES})')\n if new_results or evicted:\n _save_cache_atomic(cache_path, cache)\n stats['cache_size_after'] = len(cache)\n stats['evicted'] = evicted\n finally:\n _release_lock(cache_batch_lock)\n\n triage_results = []\n top_triage_result = {}\n top_outcome_rank = -1\n outcome_order = {\n 'attack_success': 5,\n 'attack': 4,\n 'attack_failed': 3,\n 'unknown': 2,\n 'non_attack': 1,\n }\n for a in enriched_with_triage:\n item = {\n 'dedup_key': a.get('dedup_key', ''),\n 'has_dedup_key': a.get('has_dedup_key', False),\n 'threat_name': a.get('threat_name', ''),\n 'sip': a.get('sip', ''),\n 'dip': a.get('dip', ''),\n 'is_duplicate': a.get('is_duplicate'),\n 'triage_source': a.get('triage_source', ''),\n 'triage_status': a.get('triage_status', ''),\n 'triage_attack_verdict': a.get('triage_attack_verdict', ''),\n 'triage_attack_success': a.get('triage_attack_success', 'unknown'),\n 'risk_level': a.get('risk_level', ''),\n 'report_title': a.get('report_title', ''),\n 'triage_ms': a.get('triage_ms'),\n 'triage_error': a.get('triage_error'),\n }\n triage_results.append(item)\n outcome = _triage_outcome_key(\n item['triage_attack_verdict'],\n item['triage_attack_success'],\n )\n rank = outcome_order.get(outcome, 0)\n if rank > top_outcome_rank:\n top_outcome_rank = rank\n top_triage_result = dict(item)\n # Keep the selected report intact; only the number of reports\n # crossing the process boundary is reduced.\n top_triage_result['triage_report'] = a.get('triage_report', '')\n\n stats['elapsed_ms'] = elapsed_ms\n stats['concurrency'] = concurrency\n stats['triage_sub_batch_size'] = TRIAGE_SUB_BATCH_SIZE\n stats['max_triage_cache_size'] = max_triage_cache_size\n stats['max_triage_cache_bytes'] = MAX_TRIAGE_CACHE_BYTES\n\n # Persist enriched_with_triage according to workflow config. Default is SOC DB;\n # JSONL is still available by config/input for downstream pipelines or archival.\n output_cfg = _resolve_output_config()\n run_id = (inputs.get('_run_id')\n or os.environ.get('FLOCKS_RUN_ID')\n or str(int(time.time() * 1000)))\n output_paths = []\n output_dir = ''\n first_seen_soc_alerts, soc_db_filter_stats = _select_first_seen_soc_alerts(\n enriched_with_triage,\n )\n soc_db_result = {\n 'path': output_cfg.get('soc_db_path', ''),\n 'table': 'alert_records',\n 'rows': 0,\n 'inserted_rows': 0,\n 'updated_rows': 0,\n }\n\n if output_cfg['write_soc_db'] and first_seen_soc_alerts:\n try:\n soc_db_result.update(_triage_write_soc_db(\n output_cfg['soc_db_path'], first_seen_soc_alerts, run_id,\n ))\n print(f'[triage] persisted {soc_db_result.get(\"rows\", 0)} globally unique alerts to '\n f'SOC DB {soc_db_result.get(\"path\")} '\n f'(inserted={soc_db_result.get(\"inserted_rows\", 0)}, '\n f'updated={soc_db_result.get(\"updated_rows\", 0)}); '\n f'filter={soc_db_filter_stats}')\n except Exception as e:\n import traceback\n print(f'[triage] ERROR: failed to persist triage results to SOC DB: {e}\\n{traceback.format_exc()}')\n raise\n elif output_cfg['write_soc_db']:\n print(f'[triage] no verified first-seen unique alerts to persist; filter={soc_db_filter_stats}')\n else:\n print(f'[triage] SOC DB output disabled by triage_output_mode={output_cfg[\"requested_mode\"]!r}')\n\n if output_cfg['write_jsonl'] and enriched_with_triage:\n try:\n output_dir = _triage_output_dir(output_cfg.get('jsonl_output_dir', ''))\n output_paths = _triage_write_jsonl(\n output_dir, enriched_with_triage, run_id, stats,\n )\n print(f'[triage] wrote {len(enriched_with_triage)} enriched alerts to '\n f'{len(output_paths)} JSONL file(s) under {output_dir}')\n for p in output_paths:\n print(f' → {p}')\n except Exception as e:\n import traceback\n print(f'[triage] ERROR: failed to persist triage_result JSONL: {e}\\n{traceback.format_exc()}')\n raise\n elif not output_cfg['write_jsonl']:\n print(f'[triage] JSONL output disabled by triage_output_mode={output_cfg[\"requested_mode\"]!r}')\n\n stats['output_mode'] = output_cfg['mode']\n stats['requested_output_mode'] = output_cfg['requested_mode']\n stats['output_config_path'] = output_cfg['config_path']\n stats['soc_db_path'] = soc_db_result.get('path', '')\n stats['soc_db_rows'] = soc_db_result.get('rows', 0)\n stats['soc_db_inserted_rows'] = soc_db_result.get('inserted_rows', 0)\n stats['soc_db_updated_rows'] = soc_db_result.get('updated_rows', 0)\n stats['soc_db_first_seen_rows'] = soc_db_filter_stats['first_seen_rows']\n stats['soc_db_skipped_rows'] = (\n soc_db_filter_stats['input_rows'] - soc_db_filter_stats['first_seen_rows']\n )\n stats['soc_db_filter_stats'] = soc_db_filter_stats\n stats['output_paths'] = output_paths\n stats['output_dir'] = output_dir\n\n print(f'[triage] stats={json.dumps(stats, ensure_ascii=False)}')\n\n outputs['top_triage_result'] = top_triage_result\n outputs['triage_results'] = triage_results\n outputs['triage_stats'] = stats\n outputs['load_stats'] = inputs.get('load_stats', {})\n outputs['loaded_files'] = inputs.get('loaded_files', [])\n outputs['input_date'] = inputs.get('input_date', '')\n outputs['cursor_enabled'] = bool(inputs.get('cursor_enabled'))\n outputs['cursor_before'] = inputs.get('cursor_before')\n outputs['cursor_revision'] = inputs.get('cursor_revision')\n outputs['cursor_invalidated'] = bool(inputs.get('cursor_invalidated'))\n outputs['pending_cursor'] = inputs.get('pending_cursor')\n outputs['next_cursor'] = inputs.get('next_cursor')\n outputs['has_more'] = bool(inputs.get('has_more'))\n outputs['batch_records'] = int(inputs.get('batch_records', 0) or 0)\n outputs['batch_bytes'] = int(inputs.get('batch_bytes', 0) or 0)\n outputs['triage_output_mode'] = output_cfg['mode']\n outputs['soc_db_result'] = soc_db_result\n outputs['soc_db_path'] = soc_db_result.get('path', '')\n outputs['output_config_path'] = output_cfg['config_path']\n outputs['output_paths'] = output_paths\n outputs['output_dir'] = output_dir\n outputs['_batch_lease_fd'] = _batch_lease_fd\n outputs['batch_lease_token'] = _batch_lease_token\n outputs['_triage_persistence_succeeded'] = True\nexcept BaseException:\n _release_batch_lease(_batch_lease_fd)\n raise\n", + "processIsolated": true, + "processRetainFdKeys": [ + "_batch_lease_fd" + ], + "timeoutFatal": true + }, + { + "id": "commit_cursor", + "type": "python", + "description": "研判及所有启用的持久化目标成功后，原子提交自动目录模式的生产游标；显式重放或没有新字节时不写游标。", + "code": "\"\"\"Commit the production cursor after triage persistence succeeds.\"\"\"\n\nimport datetime\nimport hashlib\nimport json\nimport os\nimport sys\nimport uuid\n\nfrom flocks.config import Config\n\n\nIS_WINDOWS = sys.platform == 'win32'\nif IS_WINDOWS:\n import msvcrt # noqa: F401\nelse:\n import fcntl # noqa: F401\n\nWORKFLOW_NAME = 'stream_alert_triage'\n\n\ndef _cursor_path():\n flocks_root = Config().get_global().data_dir.parent\n return flocks_root / 'workspace' / 'workflows' / WORKFLOW_NAME / '.triage_cursor.json'\n\n\nclass StaleCursorCommitError(RuntimeError):\n def __init__(self, current_cursor):\n super().__init__('stale_cursor_commit')\n self.current_cursor = current_cursor\n\n\ndef _cursor_revision(path):\n try:\n with open(path, 'rb') as stream:\n return hashlib.sha256(stream.read()).hexdigest()\n except FileNotFoundError:\n return None\n\n\ndef _read_current_cursor(path):\n try:\n with open(path, 'r', encoding='utf-8') as stream:\n value = json.load(stream)\n except Exception:\n return None\n return value if isinstance(value, dict) else None\n\n\ndef _acquire_cursor_lock(path):\n lock_path = path + '.lock'\n fh = open(lock_path, 'a+b')\n try:\n if IS_WINDOWS:\n if os.fstat(fh.fileno()).st_size == 0:\n fh.write(b'0')\n fh.flush()\n fh.seek(0)\n msvcrt.locking(fh.fileno(), msvcrt.LK_LOCK, 1)\n else:\n fcntl.flock(fh.fileno(), fcntl.LOCK_EX)\n except BaseException:\n fh.close()\n raise\n return fh\n\n\ndef _release_cursor_lock(fh):\n try:\n if IS_WINDOWS:\n try:\n fh.seek(0)\n msvcrt.locking(fh.fileno(), msvcrt.LK_UNLCK, 1)\n except OSError:\n pass\n else:\n fcntl.flock(fh.fileno(), fcntl.LOCK_UN)\n finally:\n fh.close()\n\n\ndef _release_batch_lease(fd):\n if type(fd) is not int or fd < 0:\n return\n try:\n if IS_WINDOWS:\n try:\n os.lseek(fd, 0, os.SEEK_SET)\n msvcrt.locking(fd, msvcrt.LK_UNLCK, 1)\n except OSError:\n pass\n else:\n fcntl.flock(fd, fcntl.LOCK_UN)\n finally:\n os.close(fd)\n\n\ndef _cursor_is_monotonic(current, pending):\n if not isinstance(current, dict):\n return True\n current_date = str(current.get('date') or '')\n pending_date = str(pending.get('date') or '')\n if pending_date != current_date:\n return pending_date > current_date\n current_seq = int(current.get('file_seq', -1))\n pending_seq = int(pending.get('file_seq', -1))\n if pending_seq != current_seq:\n return pending_seq > current_seq\n if str(pending.get('file_name') or '') != str(current.get('file_name') or ''):\n return False\n return int(pending.get('byte_offset', -1)) >= int(current.get('byte_offset', -1))\n\n\n\ndef _cursor_reset_is_safe(current, pending):\n if not isinstance(current, dict):\n return True\n current_date = str(current.get('date') or '')\n pending_date = str(pending.get('date') or '')\n if pending_date != current_date:\n return pending_date > current_date\n try:\n current_seq = int(current.get('file_seq', -1))\n pending_seq = int(pending.get('file_seq', -1))\n except (TypeError, ValueError):\n return False\n if pending_seq != current_seq:\n return pending_seq > current_seq\n return str(pending.get('file_name') or '') == str(current.get('file_name') or '')\n\n\ndef _commit_cursor_atomic(\n path, pending_cursor, expected_revision=None, run_id=None, allow_reset=False\n):\n path = os.fspath(path)\n os.makedirs(os.path.dirname(path), exist_ok=True)\n lock_fh = _acquire_cursor_lock(path)\n try:\n current_revision = _cursor_revision(path)\n current_cursor = _read_current_cursor(path)\n position_is_valid = (\n _cursor_reset_is_safe(current_cursor, pending_cursor)\n if allow_reset\n else _cursor_is_monotonic(current_cursor, pending_cursor)\n )\n if current_revision != expected_revision or not position_is_valid:\n raise StaleCursorCommitError(current_cursor)\n\n committed = {\n 'version': 2,\n 'date': str(pending_cursor['date']),\n 'file_seq': int(pending_cursor['file_seq']),\n 'file_name': str(pending_cursor['file_name']),\n 'byte_offset': int(pending_cursor['byte_offset']),\n 'device_id': int(pending_cursor['device_id']),\n 'file_id': int(pending_cursor['file_id']),\n 'head_hash': str(pending_cursor['head_hash']),\n 'boundary_start': int(pending_cursor['boundary_start']),\n 'boundary_hash': str(pending_cursor['boundary_hash']),\n 'updated_at': datetime.datetime.now().astimezone().isoformat(timespec='seconds'),\n }\n if pending_cursor.get('skipping_oversized_line') is True:\n committed['skipping_oversized_line'] = True\n\n token = str(run_id or uuid.uuid4().hex)\n safe_token = ''.join(char if char.isalnum() or char in '-_' else '_' for char in token)\n tmp_path = f'{path}.{safe_token}.{uuid.uuid4().hex}.tmp'\n try:\n with open(tmp_path, 'w', encoding='utf-8') as stream:\n json.dump(committed, stream, ensure_ascii=False, indent=2)\n stream.write(chr(10))\n stream.flush()\n os.fsync(stream.fileno())\n os.replace(tmp_path, path)\n except BaseException:\n try:\n if os.path.exists(tmp_path):\n os.remove(tmp_path)\n except OSError:\n pass\n raise\n return committed\n finally:\n _release_cursor_lock(lock_fh)\n\n\ncursor_enabled = bool(inputs.get('cursor_enabled'))\ncursor_before = inputs.get('cursor_before')\ncursor_revision = inputs.get('cursor_revision')\ncursor_invalidated = bool(inputs.get('cursor_invalidated'))\npending_cursor = inputs.get('pending_cursor')\npersistence_succeeded = inputs.get('_triage_persistence_succeeded') is True\nbatch_lease_fd = inputs.get('_batch_lease_fd')\ncursor_committed = False\ncursor_commit_error = ''\ncommitted_cursor = cursor_before if cursor_enabled else None\n\ntry:\n if cursor_enabled and pending_cursor and persistence_succeeded:\n try:\n committed_cursor = _commit_cursor_atomic(\n _cursor_path(),\n pending_cursor,\n expected_revision=cursor_revision,\n run_id=inputs.get('_run_id') or inputs.get('batch_lease_token'),\n allow_reset=cursor_invalidated,\n )\n except StaleCursorCommitError as exc:\n committed_cursor = exc.current_cursor\n cursor_commit_error = 'stale_cursor_commit'\n print('[cursor] stale cursor commit rejected; production cursor unchanged')\n else:\n cursor_committed = True\n print(f'[cursor] committed {committed_cursor[\"file_name\"]} '\n f'offset={committed_cursor[\"byte_offset\"]}')\n elif cursor_enabled and pending_cursor:\n print('[cursor] triage persistence did not complete; production cursor unchanged')\n elif not cursor_enabled:\n print('[cursor] replay mode; production cursor unchanged')\n else:\n print('[cursor] no newly consumed bytes; production cursor unchanged')\nfinally:\n _release_batch_lease(batch_lease_fd)\n\noutputs['cursor_committed'] = cursor_committed\noutputs['cursor_revision'] = cursor_revision\noutputs['cursor_invalidated'] = cursor_invalidated\noutputs['cursor_commit_error'] = cursor_commit_error\noutputs['committed_cursor'] = committed_cursor\noutputs['next_cursor'] = inputs.get('next_cursor')\noutputs['has_more'] = bool(inputs.get('has_more'))\noutputs['batch_records'] = int(inputs.get('batch_records', 0) or 0)\noutputs['batch_bytes'] = int(inputs.get('batch_bytes', 0) or 0)\noutputs['cursor_enabled'] = cursor_enabled\noutputs['input_date'] = inputs.get('input_date', '')\noutputs['load_stats'] = inputs.get('load_stats', {})\noutputs['loaded_files'] = inputs.get('loaded_files', [])\noutputs['top_triage_result'] = inputs.get('top_triage_result', {})\noutputs['triage_results'] = inputs.get('triage_results', [])\noutputs['triage_stats'] = inputs.get('triage_stats', {})\noutputs['triage_output_mode'] = inputs.get('triage_output_mode')\noutputs['soc_db_result'] = inputs.get('soc_db_result', {})\noutputs['soc_db_path'] = inputs.get('soc_db_path', '')\noutputs['output_paths'] = inputs.get('output_paths', [])\noutputs['output_dir'] = inputs.get('output_dir', '')\n" }, { "id": "summarize", "type": "python", "description": "汇总输出：写 pipeline_summary.md 到 ~/.flocks/workspace/outputs//artifacts/，暴露 top-risk 告警的 verdict/title/triage_report 作为工作流的 final outputs。", - "code": "\"\"\"\nsummarize: 汇总 triage 结果，写 pipeline_summary.md，并暴露 top-risk 研判字段。\n\"\"\"\n\nimport datetime\nimport json\nimport os\n\nfrom flocks.workspace.manager import WorkspaceManager\n\nVERDICT_ORDER = {\n 'attack_success': 5,\n 'attack': 4,\n 'attack_failed': 3,\n 'unknown': 2,\n 'benign': 1,\n}\nVERDICT_CN = {\n 'attack_success': '攻击成功',\n 'attack_failed': '攻击失败',\n 'attack': '攻击',\n 'unknown': '未知',\n 'benign': '安全',\n}\nSOURCE_CN = {\n 'cache': '缓存复用',\n 'triaged': '新研判(leader)',\n 'follower_reused': '同批次复用(follower)',\n 'failed': '研判失败',\n 'no_dedup_key_triaged': '无dedup_key已研判',\n 'no_dedup_key_failed': '无dedup_key研判失败',\n}\n\nenriched = list(inputs.get('enriched_alerts_with_triage', []) or [])\ntriage_stats = dict(inputs.get('triage_stats', {}) or {})\nload_stats = dict(inputs.get('load_stats', {}) or {})\nloaded_files = list(inputs.get('loaded_files', []) or [])\ninput_date = inputs.get('input_date', '') or datetime.date.today().isoformat()\noutput_paths = list(inputs.get('output_paths', []) or [])\noutput_dir = inputs.get('output_dir', '') or ''\ntriage_output_mode = triage_stats.get('output_mode') or inputs.get('triage_output_mode') or 'soc_db'\nsoc_db_result = dict(inputs.get('soc_db_result', {}) or {})\nsoc_db_path = inputs.get('soc_db_path') or soc_db_result.get('path') or triage_stats.get('soc_db_path', '')\n\n# Pick the top-risk alert (first occurrence of the highest verdict tier).\ntop = None\nfor a in enriched:\n v = a.get('attack_verdict', 'unknown')\n if top is None:\n top = a\n continue\n cur = VERDICT_ORDER.get(top.get('attack_verdict', ''), 0)\n new = VERDICT_ORDER.get(v, 0)\n if new > cur:\n top = a\n\ntriage_report = (top.get('triage_report') or '') if top else ''\nreport_title = (top.get('report_title') or '') if top else ''\nattack_verdict = (top.get('attack_verdict') or '') if top else ''\nrisk_level = (top.get('risk_level') or '') if top else ''\n\nrows = []\nfor i, a in enumerate(enriched, 1):\n verdict = a.get('attack_verdict', 'unknown')\n source = a.get('triage_source', '')\n rows.append(\n f'| {i} '\n f'| {(a.get(\"dedup_key\") or \"\")[:8]} '\n f'| {(\"是\" if a.get(\"is_duplicate\") else \"否\")} '\n f'| {(a.get(\"threat_name\") or \"\")[:30]} '\n f'| {SOURCE_CN.get(source, source)} '\n f'| {VERDICT_CN.get(verdict, verdict)} '\n f'| {(a.get(\"report_title\") or \"\")[:30]} |'\n )\n\ntotal_alerts = triage_stats.get('total', 0) or 1\nreuse_rate = (\n triage_stats.get('cache_hit', 0)\n + triage_stats.get('followers_reused', 0)\n) / total_alerts\n\noutput_files_md = (\n '\\n'.join(f'- `{p}`' for p in output_paths) if output_paths else '_(not persisted)_'\n)\nsoc_db_md = (\n f'`{soc_db_path}` ({soc_db_result.get(\"rows\", triage_stats.get(\"soc_db_rows\", 0))} rows)'\n if soc_db_path else '_(not persisted)_'\n)\n\nsummary_md = (\n f'# Stream Alert Triage Summary\\n\\n'\n f'**Date**: {input_date}\\n'\n f'**Loaded files**: {load_stats.get(\"file_count\", 0)}\\n'\n f'**Loaded records**: {load_stats.get(\"record_count\", 0)}\\n\\n'\n f'## Persisted Results\\n'\n f'**Output mode**: `{triage_output_mode}`\\n'\n f'**SOC DB**: {soc_db_md}\\n'\n f'**JSONL output dir**: `{output_dir or \"(none)\"}`\\n'\n f'**JSONL files written**:\\n{output_files_md}\\n\\n'\n f'## Statistics\\n'\n f'- Total alerts: {triage_stats.get(\"total\", 0)}\\n'\n f'- Unique dedup_keys: {triage_stats.get(\"unique_dedup_keys\", 0)}\\n'\n f'- Work units (LLM-bound): {triage_stats.get(\"work_units\", 0)}\\n'\n f'- Cache hit (reused historic triage): {triage_stats.get(\"cache_hit\", 0)}\\n'\n f'- Followers reused (in-batch dedup): {triage_stats.get(\"followers_reused\", 0)}\\n'\n f'- New triages (leaders): {triage_stats.get(\"triaged\", 0)}\\n'\n f'- Triage failed: {triage_stats.get(\"triage_failed\", 0)}\\n'\n f'- No-dedup_key alerts: {triage_stats.get(\"no_dedup_key_alerts\", 0)}\\n'\n f'- Reuse rate (cache + followers): {reuse_rate:.1%}\\n'\n f'- Concurrency: {triage_stats.get(\"concurrency\", \"?\")}\\n'\n f'- Elapsed: {triage_stats.get(\"elapsed_ms\", 0)} ms\\n'\n f'- Cache size: {triage_stats.get(\"cache_size_before\", 0)} → {triage_stats.get(\"cache_size_after\", 0)} '\n f'(evicted {triage_stats.get(\"evicted\", 0)})\\n'\n f'- Verdict distribution: {json.dumps(triage_stats.get(\"verdict_counts\", {}), ensure_ascii=False)}\\n\\n'\n f'## Details\\n\\n'\n f'| # | dedup_key | dup | threat | source | verdict | title |\\n'\n f'|---|-----------|-----|--------|--------|---------|-------|\\n'\n + ('\\n'.join(rows) if rows else '| – | – | – | – | – | – | – |') + '\\n\\n'\n + (f'## Top-Risk Alert Report\\n\\n{triage_report}\\n' if triage_report else '')\n)\n\nws = WorkspaceManager.get_instance()\n_summary_root = str(ws.get_workspace_dir() / 'outputs' / datetime.date.today().isoformat())\nartifacts_dir = os.path.join(_summary_root, 'artifacts')\nos.makedirs(artifacts_dir, exist_ok=True)\nsummary_path = os.path.join(artifacts_dir, 'stream_alert_triage_summary.md')\ntry:\n with open(summary_path, 'w', encoding='utf-8') as f:\n f.write(summary_md)\n print(f'[summarize] wrote {summary_path}')\nexcept Exception as e:\n print(f'[summarize] WARNING: failed to write summary: {e}')\n summary_path = ''\n\nprint(f'[summarize] top_verdict={attack_verdict!r} title={report_title[:40]!r}')\n\noutputs['summary_report'] = summary_md\noutputs['summary_path'] = summary_path\noutputs['top_attack_verdict'] = attack_verdict\noutputs['top_risk_level'] = risk_level\noutputs['top_report_title'] = report_title\noutputs['top_triage_report'] = triage_report\noutputs['triage_results'] = inputs.get('triage_results', [])\noutputs['triage_stats'] = triage_stats\noutputs['load_stats'] = load_stats\noutputs['loaded_files'] = loaded_files\noutputs['enriched_alerts_with_triage'] = enriched\noutputs['triage_output_mode'] = triage_output_mode\noutputs['soc_db_result'] = soc_db_result\noutputs['soc_db_path'] = soc_db_path\noutputs['output_paths'] = output_paths\noutputs['output_dir'] = output_dir\n" + "code": "\"\"\"\nsummarize: 汇总 triage 结果，写 pipeline_summary.md，并暴露 top-risk 研判字段。\n\"\"\"\n\nimport datetime\nimport json\nimport os\n\nfrom flocks.workspace.manager import WorkspaceManager\n\nOUTCOME_ORDER = {\n 'attack_success': 5,\n 'attack': 4,\n 'attack_failed': 3,\n 'unknown': 2,\n 'non_attack': 1,\n}\nATTACK_VERDICT_CN = {\n 'attack': '攻击',\n 'non_attack': '非攻击',\n 'unknown': '未知',\n}\nATTACK_SUCCESS_CN = {\n 'success': '攻击成功',\n 'failed': '攻击失败',\n 'unknown': '未知',\n}\n\n\ndef _triage_outcome(alert):\n verdict = alert.get('triage_attack_verdict', 'unknown')\n result = alert.get('triage_attack_success', 'unknown')\n if verdict == 'attack':\n if result == 'success':\n return 'attack_success'\n if result == 'failed':\n return 'attack_failed'\n return 'attack'\n if verdict == 'non_attack':\n return 'non_attack'\n return 'unknown'\n\n\nSOURCE_CN = {\n 'cache': '缓存复用',\n 'triaged': '新研判(leader)',\n 'follower_reused': '同批次复用(follower)',\n 'failed': '研判失败',\n 'no_dedup_key_triaged': '无dedup_key已研判',\n 'no_dedup_key_failed': '无dedup_key研判失败',\n}\n\nenriched = list(inputs.get('triage_results', []) or [])\ntop = dict(inputs.get('top_triage_result', {}) or {})\ntriage_stats = dict(inputs.get('triage_stats', {}) or {})\nload_stats = dict(inputs.get('load_stats', {}) or {})\nloaded_files = list(inputs.get('loaded_files', []) or [])\ninput_date = inputs.get('input_date', '') or datetime.date.today().isoformat()\noutput_paths = list(inputs.get('output_paths', []) or [])\noutput_dir = inputs.get('output_dir', '') or ''\ntriage_output_mode = triage_stats.get('output_mode') or inputs.get('triage_output_mode') or 'soc_db'\nsoc_db_result = dict(inputs.get('soc_db_result', {}) or {})\nsoc_db_path = inputs.get('soc_db_path') or soc_db_result.get('path') or triage_stats.get('soc_db_path', '')\ncursor_enabled = bool(inputs.get('cursor_enabled'))\ncursor_committed = bool(inputs.get('cursor_committed'))\ncursor_commit_error = str(inputs.get('cursor_commit_error') or '')\ncursor_revision = inputs.get('cursor_revision')\ncursor_invalidated = bool(inputs.get('cursor_invalidated'))\ncommitted_cursor = inputs.get('committed_cursor')\nnext_cursor = inputs.get('next_cursor')\nhas_more = bool(inputs.get('has_more'))\nbatch_records = int(inputs.get('batch_records', 0) or 0)\nbatch_bytes = int(inputs.get('batch_bytes', 0) or 0)\n\ntriage_report = (top.get('triage_report') or '') if top else ''\nreport_title = (top.get('report_title') or '') if top else ''\nattack_verdict = (top.get('triage_attack_verdict') or '') if top else ''\nattack_success = (top.get('triage_attack_success') or '') if top else ''\nrisk_level = (top.get('risk_level') or '') if top else ''\n\nrows = []\nfor i, a in enumerate(enriched, 1):\n verdict = a.get('triage_attack_verdict', 'unknown')\n attack_success_value = a.get('triage_attack_success', 'unknown')\n source = a.get('triage_source', '')\n rows.append(\n f'| {i} '\n f'| {(a.get(\"dedup_key\") or \"\")[:8]} '\n f'| {(\"是\" if a.get(\"is_duplicate\") else \"否\")} '\n f'| {(a.get(\"threat_name\") or \"\")[:30]} '\n f'| {SOURCE_CN.get(source, source)} '\n f'| {ATTACK_VERDICT_CN.get(verdict, verdict)} '\n f'| {ATTACK_SUCCESS_CN.get(attack_success_value, attack_success_value)} '\n f'| {(a.get(\"report_title\") or \"\")[:30]} |'\n )\n\ntotal_alerts = triage_stats.get('total', 0) or 1\nreuse_rate = (\n triage_stats.get('cache_hit', 0)\n + triage_stats.get('followers_reused', 0)\n) / total_alerts\n\noutput_files_md = (\n '\\n'.join(f'- `{p}`' for p in output_paths) if output_paths else '_(not persisted)_'\n)\nsoc_db_md = (\n f'`{soc_db_path}` ({soc_db_result.get(\"rows\", triage_stats.get(\"soc_db_rows\", 0))} rows)'\n if soc_db_path else '_(not persisted)_'\n)\n\nsummary_md = (\n f'# Stream Alert Triage Summary\\n\\n'\n f'**Date**: {input_date}\\n'\n f'**Loaded files**: {load_stats.get(\"file_count\", 0)}\\n'\n f'**Loaded records**: {load_stats.get(\"record_count\", 0)}\\n'\n f'**Batch bytes**: {batch_bytes}\\n'\n f'**Has more input**: {has_more}\\n'\n f'**Production cursor committed**: {cursor_committed if cursor_enabled else \"replay mode\"}\\n\\n'\n f'## Persisted Results\\n'\n f'**Output mode**: `{triage_output_mode}`\\n'\n f'**SOC DB**: {soc_db_md}\\n'\n f'**JSONL output dir**: `{output_dir or \"(none)\"}`\\n'\n f'**JSONL files written**:\\n{output_files_md}\\n\\n'\n f'## Statistics\\n'\n f'- Total alerts: {triage_stats.get(\"total\", 0)}\\n'\n f'- Unique dedup_keys: {triage_stats.get(\"unique_dedup_keys\", 0)}\\n'\n f'- Work units (LLM-bound): {triage_stats.get(\"work_units\", 0)}\\n'\n f'- Cache hit (reused historic triage): {triage_stats.get(\"cache_hit\", 0)}\\n'\n f'- Followers reused (in-batch dedup): {triage_stats.get(\"followers_reused\", 0)}\\n'\n f'- New triages (leaders): {triage_stats.get(\"triaged\", 0)}\\n'\n f'- Triage failed: {triage_stats.get(\"triage_failed\", 0)}\\n'\n f'- No-dedup_key alerts: {triage_stats.get(\"no_dedup_key_alerts\", 0)}\\n'\n f'- Reuse rate (cache + followers): {reuse_rate:.1%}\\n'\n f'- Concurrency: {triage_stats.get(\"concurrency\", \"?\")}\\n'\n f'- Elapsed: {triage_stats.get(\"elapsed_ms\", 0)} ms\\n'\n f'- Cache size: {triage_stats.get(\"cache_size_before\", 0)} → {triage_stats.get(\"cache_size_after\", 0)} '\n f'(evicted {triage_stats.get(\"evicted\", 0)})\\n'\n f'- Verdict distribution: {json.dumps(triage_stats.get(\"verdict_counts\", {}), ensure_ascii=False)}\\n\\n'\n f'## Details\\n\\n'\n f'| # | dedup_key | dup | threat | source | attack verdict | attack result | title |\\n'\n f'|---|-----------|-----|--------|--------|----------------|---------------|-------|\\n'\n + ('\\n'.join(rows) if rows else '| – | – | – | – | – | – | – | – |') + '\\n\\n'\n + (f'## Top-Risk Alert Report\\n\\n{triage_report}\\n' if triage_report else '')\n)\n\nws = WorkspaceManager.get_instance()\n_summary_root = str(ws.get_workspace_dir() / 'outputs' / datetime.date.today().isoformat())\nartifacts_dir = os.path.join(_summary_root, 'artifacts')\nos.makedirs(artifacts_dir, exist_ok=True)\nsummary_path = os.path.join(artifacts_dir, 'stream_alert_triage_summary.md')\ntry:\n with open(summary_path, 'w', encoding='utf-8') as f:\n f.write(summary_md)\n print(f'[summarize] wrote {summary_path}')\nexcept Exception as e:\n print(f'[summarize] WARNING: failed to write summary: {e}')\n summary_path = ''\n\nprint(f'[summarize] top_verdict={attack_verdict!r} title={report_title[:40]!r}')\n\noutputs['summary_report'] = summary_md\noutputs['summary_path'] = summary_path\noutputs['top_attack_verdict'] = attack_verdict\noutputs['top_attack_success'] = attack_success\noutputs['top_risk_level'] = risk_level\noutputs['top_report_title'] = report_title\noutputs['top_triage_report'] = triage_report\noutputs['triage_results'] = inputs.get('triage_results', [])\noutputs['triage_stats'] = triage_stats\noutputs['load_stats'] = load_stats\noutputs['loaded_files'] = loaded_files\noutputs['top_triage_result'] = top\noutputs['triage_output_mode'] = triage_output_mode\noutputs['soc_db_result'] = soc_db_result\noutputs['soc_db_path'] = soc_db_path\noutputs['output_paths'] = output_paths\noutputs['output_dir'] = output_dir\noutputs['cursor_enabled'] = cursor_enabled\noutputs['cursor_committed'] = cursor_committed\noutputs['cursor_commit_error'] = cursor_commit_error\noutputs['cursor_revision'] = cursor_revision\noutputs['cursor_invalidated'] = cursor_invalidated\noutputs['committed_cursor'] = committed_cursor\noutputs['next_cursor'] = next_cursor\noutputs['has_more'] = has_more\noutputs['batch_records'] = batch_records\noutputs['batch_bytes'] = batch_bytes\n" } ], "edges": [ { "from": "load_dedup_file", "to": "concurrent_triage", - "order": 0 + "order": 0, + "mapping": { + "enriched_alerts": "enriched_alerts", + "loaded_files": "loaded_files", + "load_stats": "load_stats", + "concurrency": "concurrency", + "max_triage_cache_size": "max_triage_cache_size", + "input_date": "input_date", + "cursor_enabled": "cursor_enabled", + "cursor_before": "cursor_before", + "pending_cursor": "pending_cursor", + "next_cursor": "next_cursor", + "has_more": "has_more", + "batch_records": "batch_records", + "batch_bytes": "batch_bytes", + "_triage_persistence_succeeded": "_triage_persistence_succeeded", + "_run_id": "_run_id", + "triage_output_mode": "triage_output_mode", + "persist_triage_output": "persist_triage_output", + "soc_db_path": "soc_db_path", + "jsonl_output_dir": "jsonl_output_dir", + "cursor_revision": "cursor_revision", + "cursor_invalidated": "cursor_invalidated", + "_batch_lease_fd": "_batch_lease_fd", + "batch_lease_token": "batch_lease_token", + "_triage_state_dir": "_triage_state_dir" + } }, { "from": "concurrent_triage", + "to": "commit_cursor", + "order": 0, + "mapping": { + "cursor_enabled": "cursor_enabled", + "cursor_before": "cursor_before", + "pending_cursor": "pending_cursor", + "next_cursor": "next_cursor", + "has_more": "has_more", + "batch_records": "batch_records", + "batch_bytes": "batch_bytes", + "_triage_persistence_succeeded": "_triage_persistence_succeeded", + "input_date": "input_date", + "load_stats": "load_stats", + "loaded_files": "loaded_files", + "triage_results": "triage_results", + "triage_stats": "triage_stats", + "triage_output_mode": "triage_output_mode", + "soc_db_result": "soc_db_result", + "soc_db_path": "soc_db_path", + "output_paths": "output_paths", + "output_dir": "output_dir", + "cursor_revision": "cursor_revision", + "_batch_lease_fd": "_batch_lease_fd", + "batch_lease_token": "batch_lease_token", + "cursor_invalidated": "cursor_invalidated", + "top_triage_result": "top_triage_result" + } + }, + { + "from": "commit_cursor", "to": "summarize", - "order": 0 + "order": 0, + "mapping": { + "cursor_enabled": "cursor_enabled", + "cursor_committed": "cursor_committed", + "committed_cursor": "committed_cursor", + "next_cursor": "next_cursor", + "has_more": "has_more", + "batch_records": "batch_records", + "batch_bytes": "batch_bytes", + "input_date": "input_date", + "load_stats": "load_stats", + "loaded_files": "loaded_files", + "triage_results": "triage_results", + "triage_stats": "triage_stats", + "triage_output_mode": "triage_output_mode", + "soc_db_result": "soc_db_result", + "soc_db_path": "soc_db_path", + "output_paths": "output_paths", + "output_dir": "output_dir", + "cursor_commit_error": "cursor_commit_error", + "cursor_revision": "cursor_revision", + "cursor_invalidated": "cursor_invalidated", + "top_triage_result": "top_triage_result" + } } ], "metadata": { "node_timeout_s": 7200, "sampleInputs": { - "_comment_input": "三选一：input_paths（来自 stream_alert_denoise.outputs.output_paths）/ input_path（来自 stream_alert_denoise.outputs.output_path）/ input_date（YYYY-MM-DD，遍历该日所有 dedup_result_*.jsonl）；都不传时默认取“今天”目录下所有文件。", - "input_date": "2026-05-18", + "_comment_input": "自动模式不传 input_path/input_paths，每次动态读取今天目录并使用生产游标；显式路径为有界重放模式，不修改生产游标，可回传 next_cursor 作为 resume_cursor。", "concurrency": 1, "max_triage_cache_size": 100000, "persist_triage_output": false, "_comment_dedup": "同批次内多条 alert 共享 dedup_key 时只 LLM 研判 1 次（leader），其余 follower 直接复用结果；跨批次/跨进程的复用由 triage_cache.pkl 提供。", - "_comment_cache": "研判缓存位于 ~/.flocks/workspace/workflows/stream_alert_triage/triage_cache.pkl，FIFO LRU，文件锁 + 原子落盘，可跨进程/跨执行复用。dedup_key 即 stream_alert_denoise 生成的 MD5(strict_fields + lsh_cluster_id)。", + "_comment_cache": "研判缓存位于 ~/.flocks/workspace/workflows/stream_alert_triage/triage_cache.pkl，FIFO LRU，同时受 max_triage_cache_size 和 128 MiB 字节上限约束；文件锁 + 原子落盘，可跨进程/跨执行复用。dedup_key 即 stream_alert_denoise 生成的 MD5(strict_fields + lsh_cluster_id)。", "triage_output_mode": "soc_db", - "_comment_output": "默认只接受明确 is_duplicate=false、包含 dedup_key 且批内首次出现的告警，并由 soc.db 保证 dedup_key 跨执行全局唯一；如需 JSONL，设置 triage_output_mode=jsonl 或 both。" + "_comment_output": "默认只接受明确 is_duplicate=false、包含 dedup_key 且批内首次出现的告警，并由 soc.db 保证 dedup_key 跨执行全局唯一；如需 JSONL，设置 triage_output_mode=jsonl 或 both。", + "batch_max_records": 10, + "batch_max_bytes": 33554432 + }, + "runtime": { + "strict_edge_mapping": true, + "dataflow_mode": "vertex_cache" } }, "triggers": [ @@ -62,15 +157,16 @@ }, "mapping": {}, "inputs": { - "_comment_input": "三选一：input_paths（来自 stream_alert_denoise.outputs.output_paths）/ input_path（来自 stream_alert_denoise.outputs.output_path）/ input_date（YYYY-MM-DD，遍历该日所有 dedup_result_*.jsonl）；都不传时默认取“今天”目录下所有文件。", - "input_date": "2026-05-18", + "_comment_input": "自动模式不传 input_path/input_paths，每次动态读取今天目录并使用生产游标；显式路径为有界重放模式，不修改生产游标，可回传 next_cursor 作为 resume_cursor。", "concurrency": 1, "max_triage_cache_size": 100000, "persist_triage_output": false, "_comment_dedup": "同批次内多条 alert 共享 dedup_key 时只 LLM 研判 1 次（leader），其余 follower 直接复用结果；跨批次/跨进程的复用由 triage_cache.pkl 提供。", - "_comment_cache": "研判缓存位于 ~/.flocks/workspace/workflows/stream_alert_triage/triage_cache.pkl，FIFO LRU，文件锁 + 原子落盘，可跨进程/跨执行复用。dedup_key 即 stream_alert_denoise 生成的 MD5(strict_fields + lsh_cluster_id)。", + "_comment_cache": "研判缓存位于 ~/.flocks/workspace/workflows/stream_alert_triage/triage_cache.pkl，FIFO LRU，同时受 max_triage_cache_size 和 128 MiB 字节上限约束；文件锁 + 原子落盘，可跨进程/跨执行复用。dedup_key 即 stream_alert_denoise 生成的 MD5(strict_fields + lsh_cluster_id)。", "triage_output_mode": "soc_db", - "_comment_output": "默认只接受明确 is_duplicate=false、包含 dedup_key 且批内首次出现的告警，并由 soc.db 保证 dedup_key 跨执行全局唯一；如需 JSONL，设置 triage_output_mode=jsonl 或 both。" + "_comment_output": "默认只接受明确 is_duplicate=false、包含 dedup_key 且批内首次出现的告警，并由 soc.db 保证 dedup_key 跨执行全局唯一；如需 JSONL，设置 triage_output_mode=jsonl 或 both。", + "batch_max_records": 10, + "batch_max_bytes": 33554432 }, "concurrency": { "policy": "allow", @@ -85,15 +181,16 @@ { "name": "default", "payload": { - "_comment_input": "三选一：input_paths（来自 stream_alert_denoise.outputs.output_paths）/ input_path（来自 stream_alert_denoise.outputs.output_path）/ input_date（YYYY-MM-DD，遍历该日所有 dedup_result_*.jsonl）；都不传时默认取“今天”目录下所有文件。", - "input_date": "2026-05-18", + "_comment_input": "自动模式不传 input_path/input_paths，每次动态读取今天目录并使用生产游标；显式路径为有界重放模式，不修改生产游标，可回传 next_cursor 作为 resume_cursor。", "concurrency": 1, "max_triage_cache_size": 100000, "persist_triage_output": false, "_comment_dedup": "同批次内多条 alert 共享 dedup_key 时只 LLM 研判 1 次（leader），其余 follower 直接复用结果；跨批次/跨进程的复用由 triage_cache.pkl 提供。", - "_comment_cache": "研判缓存位于 ~/.flocks/workspace/workflows/stream_alert_triage/triage_cache.pkl，FIFO LRU，文件锁 + 原子落盘，可跨进程/跨执行复用。dedup_key 即 stream_alert_denoise 生成的 MD5(strict_fields + lsh_cluster_id)。", + "_comment_cache": "研判缓存位于 ~/.flocks/workspace/workflows/stream_alert_triage/triage_cache.pkl，FIFO LRU，同时受 max_triage_cache_size 和 128 MiB 字节上限约束；文件锁 + 原子落盘，可跨进程/跨执行复用。dedup_key 即 stream_alert_denoise 生成的 MD5(strict_fields + lsh_cluster_id)。", "triage_output_mode": "soc_db", - "_comment_output": "默认只接受明确 is_duplicate=false、包含 dedup_key 且批内首次出现的告警，并由 soc.db 保证 dedup_key 跨执行全局唯一；如需 JSONL，设置 triage_output_mode=jsonl 或 both。" + "_comment_output": "默认只接受明确 is_duplicate=false、包含 dedup_key 且批内首次出现的告警，并由 soc.db 保证 dedup_key 跨执行全局唯一；如需 JSONL，设置 triage_output_mode=jsonl 或 both。", + "batch_max_records": 10, + "batch_max_bytes": 33554432 }, "headers": {}, "query": {} diff --git a/.flocks/flockshub/plugins/workflows/stream_alert_triage/workflow.md b/.flocks/flockshub/plugins/workflows/stream_alert_triage/workflow.md index 6484a6d93..0a228636e 100644 --- a/.flocks/flockshub/plugins/workflows/stream_alert_triage/workflow.md +++ b/.flocks/flockshub/plugins/workflows/stream_alert_triage/workflow.md @@ -1,170 +1,143 @@ # stream_alert_triage -stream_alert_triage 是 NDR 告警流并发研判 Pipeline。 +`stream_alert_triage` 是 `stream_alert_denoise` 的下游 HTTP 告警研判工作流。它按游标增量读取去重结果，完成 leader/follower 研判和持久化后再提交生产游标，避免日期目录数据量增长时全量加载导致 OOM。 -核心能力： -- 读取上游 stream_alert_denoise 去重输出 -- 同批次 dedup_key 相同 → 只研判 leader，follower 复用结果 -- 跨批次 dedup_key 命中缓存 → 直接复用历史研判，不调 LLM -- 4 个 LLM 分支（survey / cve_related / cve_info / payload_analysis），共享运行级并发预算 -- 研判产物仅写入 triage_report 字段，不生成独立报告文件 - -**完全自包含**，研判逻辑直接内联，不依赖也不嵌入 `tdp_alert_triage`。 - -## 核心特性 - -* **跨批次复用**：dedup_key 命中持久化 cache 时直接复用历史 verdict/title/triage_report，不调 LLM -* **同批次去重**：批内多条 alert 共享 dedup_key 时，只对 **leader（首条）研判**，follower 广播复用 leader 结果 -* **保留 4 个 LLM 分支**（survey / cve_related / cve_info / payload_analysis）— 与 `tdp_alert_triage` 完全相同的研判语义；所有 LLM 调用共享运行级并发预算，避免与外层 work unit 并发相乘 -* **研判产物仅以字段形式附加**到每条 alert（`triage_report` 字段含带语义标签的完整 markdown），**不生成任何独立的 per-alert 报告文件** - -## 与上游的关系 - -## 1. 功能概览 - -基本信息: - -- 工作流 ID: `stream_alert_triage` -- 工作流目录: `~/.flocks/plugins/workflows/stream_alert_triage/` -- 分类: `default` -- 状态: `active` -- 入口节点: load_dedup_file (Python) -- 终点节点: summarize (Python) -- 生成时间: 2026/6/24 15:29:18 - -适合在这里写清楚: - -- 这个工作流解决什么问题。 -- 适合处理什么输入。 -- 不负责处理什么边界场景。 - -## 2. 原理和总体流程 - -核心原理是把输入按节点顺序逐步加工，每个节点只负责一个清晰职责。流程顺序如下: +## 流程 ```text -load_dedup_file -> concurrent_triage -> summarize +load_dedup_file -> concurrent_triage -> commit_cursor -> summarize ``` -流程表: - -| 顺序 | 节点 | 做什么 | 下一步 | -| --- | --- | --- | --- | -| 1 | load_dedup_file | 一次性读取 stream_alert_denoise 写入的 JSONL 文件。输入优先级：input_paths > input_path > input_date（自动遍历该日所有 dedup_result_*.jsonl）> 当日默认。跳过 file_header 行，输出 enriched_alerts (list[dict])。 | concurrent_triage | -| 2 | concurrent_triage | Leader/follower 分组并发研判节点（自包含，内联 tdp_alert_triage 逻辑）。先按 dedup_key 把 alerts 分组：每组只对 leader 研判，follower 复用 leader 结果。外层 ThreadPoolExecutor(concurrency) 处理 unique work units（concurrency 取值 1–5，默认 1）；单条告警仍执行 survey / cve_related / cve_info / payload_analysis 4 个分支，但所有 `llm.ask()` 共享运行级 concurrency 预算，因此总 LLM 峰值不超过 1–5，不再与分支数相乘。dedup_key 在 triage_cache.pkl 命中时直接复用历史 verdict/title/triage_report；未命中则 leader 执行完整研判（情报查询 + 4 个 LLM 分支 + attack_analysis + verdict + title + 聚合 markdown），完整研判 markdown 仅写入 alert 的 `triage_report` 字段，**不生成任何独立报告文件**。新结果合并写回 cache（FIFO LRU + 文件锁 + 原子落盘）。SOC DB 持久化只接受明确 `is_duplicate=false`、包含 `dedup_key` 且批内首次出现的告警，并通过数据库唯一索引保证跨执行全局唯一；重复 key 只更新研判字段并保留首次事件元数据，持久化失败会使工作流失败。可通过工作流目录 `config.json` 或运行输入将 `triage_output_mode` 切换为 `jsonl` / `both` / `none`，保留 `triage_result_NNN.jsonl` 可选输出。 | summarize | -| 3 | summarize | 汇总输出：写 pipeline_summary.md 到 ~/.flocks/workspace/outputs//artifacts/，暴露 top-risk 告警的 verdict/title/triage_report 作为工作流的 final outputs。 | 工作流最终输出 | - -编辑流程结构时，要同时确认节点顺序、边关系、字段映射和最终输出是否仍然一致。 - -## 3. 输入说明 - -本章用于说明工作流接受什么输入，以及入口节点如何理解这些输入。 - -当前工作流保存了这些样例输入，可以先照着这些字段测试: - -- _comment_input: 三选一：input_paths（来自 stream_alert_denoise.outputs.output_paths）/ input_path（来自 stream_alert_denoise.outputs.output_path）/ ... -- input_date: 2026-05-18 -- concurrency: 1 -- max_triage_cache_size: 100000 -- persist_triage_output: false -- triage_output_mode: soc_db -- _comment_output: 默认只接受明确 is_duplicate=false、包含 dedup_key 且批内首次出现的告警，并由 soc.db 保证 dedup_key 跨执行全局唯一；如需 JSONL，设置 triage_output_mode=jsonl 或 both。 -- _comment_dedup: 同批次内多条 alert 共享 dedup_key 时只 LLM 研判 1 次（leader），其余 follower 直接复用结果；跨批次/跨进程的复用由 triage_cache.pkl 提供。 -- _comment_cache: 研判缓存位于 ~/.flocks/workspace/workflows/stream_alert_triage/triage_cache.pkl，FIFO LRU，文件锁 + 原子落盘，可跨进程/跨执行复用。dedup_key 即 str... +节点边使用显式字段映射并启用严格映射模式，只保留下一节点必需的批次、游标、研判和汇总字段，不透传完整上游 payload。 -修改输入时，至少同步检查: - -- 入口节点是否能读取新字段。 -- 样例输入是否覆盖主要场景。 -- 下游节点是否还在引用旧字段名。 -- 发布方式中的参数说明是否需要更新。 - -## 4. 模块逻辑 - -本章按执行顺序解释每个节点。修改内部逻辑时，优先定位到对应节点，再检查它的上下游关系。 - -### 4.1 load_dedup_file - -职责: 一次性读取 stream_alert_denoise 写入的 JSONL 文件。输入优先级：input_paths > input_path > input_date（自动遍历该日所有 dedup_result_*.jsonl）> 当日默认。跳过 file_header 行，输出 enriched_alerts (list[dict])。 - -- 节点类型: Python -- 输入来源: 工作流输入 / 触发器输入 -- 输出去向: concurrent_triage -- 编辑重点: 修改去重阈值、状态保存、结果落盘路径或输出格式时，优先编辑这里。 -- 上游关系: 从工作流输入开始 -- 下游关系: load_dedup_file -> concurrent_triage - -### 4.2 concurrent_triage +| 节点 | 职责 | +| --- | --- | +| `load_dedup_file` | 以二进制方式有界读取 JSONL，只产生待提交游标，不写生产游标 | +| `concurrent_triage` | 按 `dedup_key` 分组研判、复用缓存，并写入所有启用的持久化目标 | +| `commit_cursor` | 上一步整体成功后原子提交生产游标；显式重放不写游标 | +| `summarize` | 生成总览并暴露最终结构化输出 | + +## 加载边界 + +默认单批限制： + +```json +{ + "batch_max_records": 10, + "batch_max_bytes": 33554432, + "concurrency": 1, + "max_triage_cache_size": 100000, + "triage_output_mode": "soc_db" +} +``` -职责: Leader/follower 分组并发研判节点（自包含，内联 tdp_alert_triage 逻辑）。先按 dedup_key 把 alerts 分组：每组只对 leader 研判，follower 复用 leader 结果。外层 ThreadPoolExecutor(concurrency) 处理 unique work units（concurrency 取值 1–5，默认 1）；单条告警仍执行 survey / cve_related / cve_info / payload_analysis 4 个分支，但所有 `llm.ask()` 共享运行级 concurrency 预算，因此总 LLM 峰值不超过 1–5，不再与分支数相乘。dedup_key 在 triage_cache.pkl 命中时直接复用历史 verdict/title/triage_report；未命中则 leader 执行完整研判（情报查询 + 4 个 LLM 分支 + attack_analysis + verdict + title + 聚合 markdown），完整研判 markdown 仅写入 alert 的 `triage_report` 字段，**不生成任何独立报告文件**。新结果合并写回 cache（FIFO LRU + 文件锁 + 原子落盘）。SOC DB 持久化只接受明确 `is_duplicate=false`、包含 `dedup_key` 且批内首次出现的告警，并通过数据库唯一索引保证跨执行全局唯一；重复 key 只更新研判字段并保留首次事件元数据，持久化失败会使工作流失败。可通过工作流目录 `config.json` 或运行输入将 `triage_output_mode` 切换为 `jsonl` / `both` / `none`，保留 `triage_result_NNN.jsonl` 可选输出。 +- `batch_max_records` 和 `batch_max_bytes` 必须是大于 0 的整数，非法值分别回退为 10 和 32 MiB。 +- 条数只统计有效 JSON 告警对象；header、空行、坏 JSON 和非对象会推进待提交 offset，但不占告警条数。 +- 未达到单行字节上限且没有换行符的末尾半行不解析、不推进 offset。 +- 超大单行使用固定大小分块跨批跳过，跨批状态保存在游标中；每批实际读取量仍受 + `batch_max_bytes` 限制，不会整体读入内存或永久阻塞消费。 +- 多文件按 `dedup_result_NNN.jsonl` 的数字序号排序，可正确跨越 999 → 1000。 +- `loaded_files` 只返回本批实际触达的文件。 -- 节点类型: Python -- 输入来源: load_dedup_file -- 输出去向: summarize -- 编辑重点: 修改去重阈值、状态保存、结果落盘路径或输出格式时，优先编辑这里。 -- 上游关系: load_dedup_file -> concurrent_triage -- 下游关系: concurrent_triage -> summarize +## 运行模式 -### 4.3 summarize +### 自动目录模式 -职责: 汇总输出：写 pipeline_summary.md 到 ~/.flocks/workspace/outputs//artifacts/，暴露 top-risk 告警的 verdict/title/triage_report 作为工作流的 final outputs。 +未传 `input_path` 和 `input_paths` 时启用。`input_date` 未设置则每次执行动态计算当天日期，并扫描： -- 节点类型: Python -- 输入来源: concurrent_triage -- 输出去向: 工作流最终输出 -- 编辑重点: 修改此步骤的输入、输出或执行逻辑时，先确认上下游字段是否同步变化。 -- 上游关系: concurrent_triage -> summarize -- 下游关系: 输出工作流结果 +```text +~/.flocks/workspace/workflows/stream_alert_denoise//dedup_result_NNN.jsonl +``` -## 5. 输出说明 +生产游标位于： -本章用于维护工作流最终返回什么，以及是否产生额外副作用。 +```text +~/.flocks/workspace/workflows/stream_alert_triage/.triage_cursor.json +``` -输出说明建议包含: +日期变化时直接从新日期的第一个文件开始；旧日期未消费完的数据按设计丢弃，不跨天补偿。 +游标只接受 `version=2`，除序号和偏移量外还保存 device ID、file ID、文件头 SHA-256 以及游标前最多 4 KiB 内容的 SHA-256。恢复时任一文件身份或内容锚点不匹配都会设置 `cursor_invalidated=true`，并从该文件头重新读取；旧版或损坏游标也按同样方式安全重置。打开文件后会再次校验，避免校验与读取之间发生同名替换而沿用旧偏移。 -- 返回给用户或调用方的核心字段。 -- 给下游系统继续消费的结构化字段。 -- 是否写文件、发通知、调用外部系统或更新状态。 -- 没有结果、部分失败、完全失败时分别返回什么。 +### 显式重放模式 -如果还不确定输出格式，先用一条样例跑通，再把真实返回字段补到这里。 +传入 `input_path` 或 `input_paths` 时进入重放模式： -## 6. 发布方式 +- 仍受 10 条和 32 MiB 限制。 +- 不读取、也不修改生产游标。 +- 将返回的 `next_cursor` 作为下一次的 `resume_cursor`，即可继续读取。 +- 显式路径不存在时只记录统计，不回退到自动日期目录。 -发布页会根据 `config.json` 模板和运行时状态决定展示哪些能力；`workflow.md` 只负责解释这些能力的用途。 +## 游标提交语义 -当前 `workflow.json` 里配置了这些触发器: +自动模式在读取游标前获取工作流级批次租约，租约覆盖 `load -> triage -> persist -> commit`，异常或提交完成后释放。`load_dedup_file` 只输出 `pending_cursor` 和读取时的 `cursor_revision`。只有 `concurrent_triage` 完成研判、缓存和所有启用的持久化目标成功后，`commit_cursor` 才在独立游标锁内执行 revision CAS 与单调性校验，并通过 run 级唯一临时文件、`flush`、`fsync` 和 `os.replace` 原子写入生产游标。文件身份失效触发的明确重置仍需通过 CAS，但允许偏移回到文件头。 -- syslog-default: syslog，启用 +| 情况 | 推进生产游标 | +| --- | --- | +| 所有告警研判与持久化成功 | 是 | +| 单条研判失败，但失败状态已形成 | 是 | +| 只消费 header、空行或完整坏行 | 是 | +| SOC DB 或启用的 JSONL 写入失败 | 否 | +| 研判缓存写入失败 | 否 | +| 游标 revision 已被其他执行修改 | 否，返回 `stale_cursor_commit` | +| 节点超时、取消或进程退出 | 否 | +| 没有读取任何新字节 | 否 | +| 显式重放模式 | 否 | + +单条研判失败会保存 `triage_status=failed`、`triage_error`、`triage_attack_verdict=unknown` 和 `triage_attack_success=unknown`，不会阻塞整个数据流。模型判定节点直接生成这两个字段：`triage_attack_verdict` 表示是否攻击，枚举为 `attack | non_attack | unknown`；`triage_attack_success` 表示攻击结果，枚举为 `success | failed | unknown`。原始告警的 `attack_verdict`、`attack_success`、`threat_result` 等字段保持不变，不参与研判结果字段的生成。整体持久化失败会抛出异常，使游标保持不变。 + +```json +{ + "triage_attack_verdict": "attack", + "triage_attack_success": "failed" +} +``` -发布相关编辑原则: +## 研判与持久化 -- 改展示模板: 修改 `config.json`。 -- 改运行启停状态: 通过发布页或后端运行时状态处理。 -- 改参数语义: 同步更新本章、输入说明和相关节点。 -- 不要把明文密钥、长期 token 或私人路径写进 `workflow.md` 或 `config.json`。 +- 同批相同 `dedup_key` 只研判 leader，followers 复用结果。 +- 跨批命中 `triage_cache.pkl` 时直接复用，不调用 LLM；缓存锁覆盖 cache miss、LLM 研判和保存阶段，避免并发执行重复研判相同批次。 +- 缓存使用 run 级唯一临时文件原子保存；保存失败会抛出异常并阻止生产游标提交。 +- 默认 `triage_output_mode=soc_db`；也支持 `jsonl`、`both` 和 `none`。 +- SOC DB 只接收明确 `is_duplicate=false`、有非空 `dedup_key` 且批内首次出现的告警。 +- 研判正文只保存在 `triage_report` 字段，不生成逐告警 markdown 文件。 -## 7. 编辑指南 +## 主要输出 -先判断你要改哪一类内容，再去找对应位置: +除原有研判输出外，增量加载还返回： -| 修改目标 | 优先查看 | +| 字段 | 说明 | | --- | --- | -| 输入格式、来源、样例 | 第 3 章和入口节点 | -| 字段映射、清洗、分类 | 第 4 章对应节点 | -| 分支、循环、节点增删 | `workflow.json` 和第 2 章流程表 | -| 输出字段、落盘、通知 | 第 5 章和终点节点 | -| API、Syslog、Kafka 等发布方式 | `config.json` 和第 6 章 | -| 字段重命名 | 所有上下游节点、样例输入和输出说明 | +| `cursor_enabled` | 是否为自动目录生产模式 | +| `cursor_before` | 本次读取前的生产或重放游标 | +| `cursor_revision` | 本次读取到的生产游标内容摘要，用于提交 CAS | +| `cursor_invalidated` | 文件身份、内容锚点或游标结构失效后是否从文件头重置 | +| `cursor_commit_error` | 游标提交错误；并发 revision 变化时为 `stale_cursor_commit` | +| `pending_cursor` | 本批成功消费字节之后的待提交位置 | +| `next_cursor` | 显式重放调用方可回传的续读位置 | +| `cursor_committed` | 本次是否实际写入生产游标 | +| `committed_cursor` | 提交后的生产游标 | +| `has_more` | 当前输入中是否还有未消费字节 | +| `batch_records` | 本批有效告警数，最大 10 | +| `batch_bytes` | 本批受字节预算约束的读取量 | + +## 调度与容量 + +工作流内置的 schedule trigger 间隔为 5 分钟，且保持 `noOverlap=true`。默认每批 10 条时理论最大吞吐为： -编辑后建议把改动说明写回相应章节，让下一个人可以直接看懂为什么这样改。 +```text +10 × 12 × 24 = 2880 条/天 +``` -## 8. 验证方式 +实际吞吐还受 LLM、情报查询、持久化耗时和 cache miss 数量影响。超过当天处理能力的积压会在日期切换时丢弃。 -最小验收清单: +## 验证清单 -- [ ] 用一条正常样例能跑通。 -- [ ] 输出字段符合你的预期。 -- [ ] 如果改了字段名，下游节点没有继续引用旧字段。 -- [ ] 如果改了发布方式，发布页只展示应该出现的能力。 -- [ ] 没有明文密钥、长期 token 或私人路径写进工作流目录。 +- 单文件 20 条分两批读取，提交后第二批不重复第一批。 +- 当前文件追加数据后从旧 EOF 继续。 +- 文件 001 剩余 6 条、002 有更多数据时，本批读取 6 + 4。 +- header、空行、坏 JSON、非对象、半行和超大行符合各自 offset 规则。 +- SOC DB 或 JSONL 写失败时游标不变。 +- 缓存写失败、重叠生产执行或 stale cursor commit 均不会推进游标。 +- 同名文件替换及同 inode 重写会使 v2 游标失效并从文件头重读。 +- 显式重放可通过 `next_cursor` 续读且不污染生产游标。 +- 工作流 JSON 可解析，所有节点 Python 代码可通过 AST 解析，相关测试通过。 diff --git a/.flocks/plugins/tools/device/qingteng_v3_4_1_66/_test.yaml b/.flocks/plugins/tools/device/qingteng_v3_4_1_66/_test.yaml index bb430e9c9..104af326d 100644 --- a/.flocks/plugins/tools/device/qingteng_v3_4_1_66/_test.yaml +++ b/.flocks/plugins/tools/device/qingteng_v3_4_1_66/_test.yaml @@ -12,14 +12,6 @@ connectivity: # `label` is the default (English) display string; `label_cn` is the # optional Chinese override picked when the WebUI runs in zh-CN. fixtures: - qingteng_login: - - label: "Login and retrieve session token" - label_cn: "登录并获取会话 Token" - tags: [smoke] - params: {} - assert: - success: true - qingteng_system_audit: - label: "List audit logs (page 1)" label_cn: "查询系统审计日志（第 1 页）" diff --git a/.flocks/plugins/tools/device/qingteng_v3_4_1_66/qingteng.handler.py b/.flocks/plugins/tools/device/qingteng_v3_4_1_66/qingteng.handler.py index e14a510cd..b4e95bc5e 100644 --- a/.flocks/plugins/tools/device/qingteng_v3_4_1_66/qingteng.handler.py +++ b/.flocks/plugins/tools/device/qingteng_v3_4_1_66/qingteng.handler.py @@ -1599,21 +1599,6 @@ async def _dispatch_group(ctx: ToolContext, group: str, action: str, **params: A return _request_signed_json(spec.method, path, query=query, body=body, action=f"{group}.{action}") -async def login(ctx: ToolContext, **kwargs: Any) -> ToolResult: - del ctx, kwargs - config = _load_runtime_config() - if not config: - return ToolResult( - success=False, - error="Missing configuration: qingteng base_url/qingteng_host, qingteng_username, qingteng_password", - ) - conn_cls, host, port, base_path, username, password = config - ok, result, payload = _login_request(conn_cls, host, port, base_path, username, password) - if not ok: - return ToolResult(success=False, error=str(result), output=payload) - return ToolResult(success=True, output=result, metadata={"source": "Qingteng", "api": "login", "path": "/v1/api/auth"}) - - async def system_audit( ctx: ToolContext, eventName: str | None = None, diff --git a/.flocks/plugins/tools/device/qingteng_v3_4_1_66/qingteng_login.handler.py b/.flocks/plugins/tools/device/qingteng_v3_4_1_66/qingteng_login.handler.py deleted file mode 100644 index 695df3dff..000000000 --- a/.flocks/plugins/tools/device/qingteng_v3_4_1_66/qingteng_login.handler.py +++ /dev/null @@ -1,20 +0,0 @@ -import importlib.util -from pathlib import Path - -from flocks.tool.registry import ToolContext, ToolResult - - -def _load_core_module(): - script_path = Path(__file__).with_name("qingteng.handler.py") - spec = importlib.util.spec_from_file_location("_flocks_qingteng_core", str(script_path)) - if spec is None or spec.loader is None: - raise ImportError(f"Cannot create import spec for {script_path}") - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - return module - - -async def login(ctx: ToolContext, **kwargs) -> ToolResult: - del kwargs - core = _load_core_module() - return await core.login(ctx) diff --git a/.flocks/plugins/tools/device/qingteng_v3_4_1_66/qingteng_login.yaml b/.flocks/plugins/tools/device/qingteng_v3_4_1_66/qingteng_login.yaml deleted file mode 100644 index 433b6ee7f..000000000 --- a/.flocks/plugins/tools/device/qingteng_v3_4_1_66/qingteng_login.yaml +++ /dev/null @@ -1,27 +0,0 @@ -name: qingteng_login -description: > - 青藤云安全平台登录认证。首次使用需调用此接口获取认证凭证， - 返回的 jwt、signKey、comId 用于后续 API 请求的签名验证。 - 当用户提到青藤云主机安全、漏洞检测、基线检查时需要先调用此工具。 -category: custom -enabled: true -requires_confirmation: false -provider: qingteng - -inputSchema: - type: object - properties: {} - required: [] - -handler: - type: script - script_file: qingteng_login.handler.py - function: login - -notes: | - 返回数据包含: - - signKey: 签名密钥，用于后续请求签名 - - jwt: JWT token - - comId: 公司ID - - 提示：用户名密码应配置在 SecretManager 中 diff --git a/flocks/agent/agent.py b/flocks/agent/agent.py index e224aa24b..12f547b01 100644 --- a/flocks/agent/agent.py +++ b/flocks/agent/agent.py @@ -4,7 +4,7 @@ Consolidates: - AgentInfo / AgentModel (previously flocks.agent.core.agent) - AgentPromptMetadata / DelegationTrigger / AvailableAgent / - AvailableTool / AvailableSkill / AvailableCategory + AvailableTool / AvailableSkill (previously flocks.agent.prompts.builder.dynamic) All other modules should import from here; the old locations are kept @@ -87,13 +87,6 @@ class AvailableSkill: location: str -@dataclass -class AvailableCategory: - """Task delegation category (model preset + domain label).""" - name: str - description: str - - @dataclass class AvailableWorkflow: """Workflow available for execution via run_workflow tool.""" @@ -154,7 +147,7 @@ class AgentInfo(BaseModel): delegatable: Optional[bool] = None # "module.path:function_name" called during phase-2 prompt injection. - # Signature: (agent_info, available_agents, tools, skills, categories) → None + # Signature: (agent_info, available_agents, tools, skills, workflows) → None # The function sets agent_info.prompt directly. prompt_builder: Optional[str] = Field(default=None) diff --git a/flocks/agent/agent_factory.py b/flocks/agent/agent_factory.py index 2c2ac9b37..1a70dd38a 100644 --- a/flocks/agent/agent_factory.py +++ b/flocks/agent/agent_factory.py @@ -23,6 +23,7 @@ from __future__ import annotations import importlib +import inspect import shutil from pathlib import Path from typing import Any, Dict, List, Optional, Tuple @@ -296,14 +297,13 @@ def inject_dynamic_prompts( available_agents: list, tools: list, skills: list, - categories: list, workflows: Optional[list] = None, ) -> None: """ Inject dynamic prompts for all agents that have a ``prompt_builder``. Dynamically imports each agent's prompt_builder module and calls its - ``inject(agent_info, available_agents, tools, skills, categories, workflows)`` + ``inject(agent_info, available_agents, tools, skills, workflows)`` function. The inject function is expected to set ``agent_info.prompt`` directly. @@ -316,7 +316,29 @@ def inject_dynamic_prompts( module_path, func_name = agent.prompt_builder.rsplit(":", 1) module = importlib.import_module(module_path) inject_fn = getattr(module, func_name) - inject_fn(agent, available_agents, tools, skills, categories, workflows or []) + signature = inspect.signature(inject_fn) + positional_parameters = [ + parameter + for parameter in signature.parameters.values() + if parameter.kind + in ( + inspect.Parameter.POSITIONAL_ONLY, + inspect.Parameter.POSITIONAL_OR_KEYWORD, + ) + ] + has_variadic_arguments = any( + parameter.kind == inspect.Parameter.VAR_POSITIONAL + for parameter in signature.parameters.values() + ) + uses_legacy_signature = ( + len(positional_parameters) >= 6 + or (has_variadic_arguments and len(positional_parameters) <= 4) + ) + + inject_args = [agent, available_agents, tools, skills] + if uses_legacy_signature: + inject_args.append([]) + inject_fn(*inject_args, workflows or []) log.debug("agent.factory.prompt_injected", {"name": name}) except Exception as e: log.error("agent.factory.prompt_inject_error", { diff --git a/flocks/agent/agents/hephaestus/prompt_builder.py b/flocks/agent/agents/hephaestus/prompt_builder.py index 31d91f859..9617affcb 100644 --- a/flocks/agent/agents/hephaestus/prompt_builder.py +++ b/flocks/agent/agents/hephaestus/prompt_builder.py @@ -14,7 +14,6 @@ AvailableAgent, AvailableTool, AvailableSkill, - AvailableCategory, ) @@ -23,7 +22,6 @@ def inject( available_agents: List["AvailableAgent"], tools: List["AvailableTool"], skills: List["AvailableSkill"], - categories: List["AvailableCategory"], workflows: Optional[list] = None, ) -> None: """Build and inject Hephaestus's dynamic system prompt.""" @@ -31,7 +29,6 @@ def inject( available_agents=available_agents, available_tools=tools, available_skills=skills, - available_categories=categories, use_task_system=False, ) @@ -40,7 +37,6 @@ def build_hephaestus_prompt( available_agents: List["AvailableAgent"], available_tools: List["AvailableTool"], available_skills: List["AvailableSkill"], - available_categories: List["AvailableCategory"], use_task_system: bool = False, ) -> str: from flocks.agent.prompt_utils import ( @@ -49,7 +45,7 @@ def build_hephaestus_prompt( build_tool_selection_table, build_explore_section, build_librarian_section, - build_category_skills_delegation_guide, + build_skills_delegation_guide, build_delegation_table, build_oracle_section, build_hard_blocks_section, @@ -61,7 +57,7 @@ def build_hephaestus_prompt( agent_selection = build_agent_selection_table(available_agents) explore_section = build_explore_section(available_agents) librarian_section = build_librarian_section(available_agents) - category_skills_guide = build_category_skills_delegation_guide(available_categories, available_skills) + skills_guide = build_skills_delegation_guide(available_skills) delegation_table = build_delegation_table(available_agents) oracle_section = build_oracle_section(available_agents) hard_blocks = build_hard_blocks_section() @@ -158,9 +154,7 @@ def build_hephaestus_prompt( **Delegation Check (MANDATORY before acting directly):** 1. Is there a specialized agent that perfectly matches this request? -2. If not, is there a `delegate_task` category that best describes this task? What skills are available to equip the agent with? - - If delegating by `category=...`, evaluate relevant skills and pass them via `load_skills=[...]`. - - If delegating by `subagent_type=...`, `load_skills` may be omitted unless a specific skill is clearly needed. +2. If so, delegate with `subagent_type=...` and evaluate which skills should be passed via `load_skills=[...]`. 3. Can I do it myself for the best result, FOR SURE? **Default Bias: DELEGATE for complex tasks. Work yourself ONLY when trivial.** @@ -197,7 +191,7 @@ def build_hephaestus_prompt( ## Phase 2 - Implementation -__CATEGORY_SKILLS_GUIDE__ +__SKILLS_GUIDE__ __DELEGATION_TABLE__ @@ -241,7 +235,7 @@ def build_hephaestus_prompt( prompt = prompt.replace("__AGENT_SELECTION__", agent_selection) prompt = prompt.replace("__EXPLORE_SECTION__", explore_section) prompt = prompt.replace("__LIBRARIAN_SECTION__", librarian_section) - prompt = prompt.replace("__CATEGORY_SKILLS_GUIDE__", category_skills_guide) + prompt = prompt.replace("__SKILLS_GUIDE__", skills_guide) prompt = prompt.replace("__DELEGATION_TABLE__", delegation_table) prompt = prompt.replace("__HARD_BLOCKS__", hard_blocks) prompt = prompt.replace("__ANTI_PATTERNS__", anti_patterns) diff --git a/flocks/agent/agents/librarian/prompt_builder.py b/flocks/agent/agents/librarian/prompt_builder.py index 0daad08d9..b29392dcf 100644 --- a/flocks/agent/agents/librarian/prompt_builder.py +++ b/flocks/agent/agents/librarian/prompt_builder.py @@ -9,7 +9,7 @@ from typing import TYPE_CHECKING, Optional if TYPE_CHECKING: - from flocks.agent.agent import AgentInfo, AvailableAgent, AvailableSkill, AvailableCategory, AvailableTool + from flocks.agent.agent import AgentInfo def inject( @@ -17,7 +17,6 @@ def inject( available_agents: list, tools: list, skills: list, - categories: list, workflows: Optional[list] = None, ) -> None: """Inject the year-aware prompt into agent_info.""" diff --git a/flocks/agent/agents/rex/prompt_builder.py b/flocks/agent/agents/rex/prompt_builder.py index a4bc0a01c..1e74df7ab 100644 --- a/flocks/agent/agents/rex/prompt_builder.py +++ b/flocks/agent/agents/rex/prompt_builder.py @@ -13,7 +13,6 @@ AvailableAgent, AvailableTool, AvailableSkill, - AvailableCategory, AvailableWorkflow, ) @@ -23,7 +22,6 @@ def inject( available_agents: List["AvailableAgent"], tools: List["AvailableTool"], skills: List["AvailableSkill"], - categories: List["AvailableCategory"], workflows: Optional[List["AvailableWorkflow"]] = None, ) -> None: """Build and inject Rex's dynamic system prompt.""" @@ -31,7 +29,6 @@ def inject( available_agents=available_agents, available_tools=tools, available_skills=skills, - available_categories=categories, available_workflows=workflows or [], use_task_system=False, ) @@ -41,7 +38,6 @@ def build_dynamic_rex_prompt( available_agents: List["AvailableAgent"], available_tools: List["AvailableTool"], available_skills: List["AvailableSkill"], - available_categories: List["AvailableCategory"], available_workflows: Optional[List["AvailableWorkflow"]] = None, use_task_system: bool = False, ) -> str: @@ -52,7 +48,7 @@ def build_dynamic_rex_prompt( build_anti_patterns_section, ) - _ = available_tools, available_categories + _ = available_tools key_triggers = build_key_triggers_section(available_agents, available_skills) agent_selection = build_agent_selection_table(available_agents) @@ -130,7 +126,7 @@ def build_dynamic_rex_prompt( 1. **Direct tools first**: if there is a short tool path, execute directly. 2. **Security exception**: for one IOC that only needs basic TI facts, prefer direct lookup. 3. **Delegate when needed**: use specialists for deep investigation, attribution, correlation, batching, external docs, or structured expert output. -4. **Do not guess**: if unsure whether something is a tool, skill, category, or subagent, use `tool_search` first. +4. **Do not guess**: if unsure whether something is a tool, skill, or subagent, use `tool_search` first. ## 3. Delegation diff --git a/flocks/agent/agents/rex_junior/prompt_builder.py b/flocks/agent/agents/rex_junior/prompt_builder.py index 56e847254..55bd6729e 100644 --- a/flocks/agent/agents/rex_junior/prompt_builder.py +++ b/flocks/agent/agents/rex_junior/prompt_builder.py @@ -17,7 +17,6 @@ def inject( available_agents: list, tools: list, skills: list, - categories: list, workflows: Optional[list] = None, ) -> None: """Inject the default rex-junior prompt into agent_info.""" @@ -27,14 +26,14 @@ def inject( def _build_prompt(prompt_append: Optional[str] = None) -> str: prompt = """ Rex-Junior - Focused executor. -Execute tasks directly. NEVER delegate or spawn other agents. +Execute implementation tasks directly. BLOCKED ACTIONS (will fail if attempted): - delegate_task for implementation work: BLOCKED -ALLOWED: delegate_task with `subagent_type="explore"` or `subagent_type="librarian"` for research only. +ALLOWED: delegate_task with `subagent_type="explore"` or `subagent_type="librarian"` for read-only research only. You work ALONE for implementation. No delegation of implementation tasks. diff --git a/flocks/agent/prompt_utils.py b/flocks/agent/prompt_utils.py index 9c7beb943..87d722efb 100644 --- a/flocks/agent/prompt_utils.py +++ b/flocks/agent/prompt_utils.py @@ -12,7 +12,7 @@ from typing import Dict, List, Optional -from flocks.agent.agent import AvailableAgent, AvailableCategory, AvailableSkill, AvailableTool, AvailableWorkflow +from flocks.agent.agent import AvailableAgent, AvailableSkill, AvailableTool, AvailableWorkflow # --------------------------------------------------------------------------- @@ -197,39 +197,24 @@ def build_delegation_table(agents: List[AvailableAgent]) -> str: return "\n".join(rows) -def build_category_skills_delegation_guide( - categories: List[AvailableCategory], - skills: List[AvailableSkill], -) -> str: - if not categories and not skills: +def build_skills_delegation_guide(skills: List[AvailableSkill]) -> str: + if not skills: return "" - category_rows = [f"| `{c.name}` | {c.description or c.name} |" for c in categories] skill_rows = [ f"| `{s.name}` | {s.description.split('.')[0] or s.description} |" for s in skills ] return ( - "### Category + Skills Delegation System\n\n" - "**delegate_task() combines categories and skills for optimal task execution.**\n\n" - "#### Available Categories (Domain-Optimized Models)\n\n" - "Each category is configured with a model optimized for that domain. Read the description to understand when to use it.\n\n" - "| Category | Domain / Best For |\n" - "|----------|-------------------|\n" - + "\n".join(category_rows) - + "\n\n#### Available Skills (Domain Expertise Injection)\n\n" + "### Skills Delegation System\n\n" + "#### Available Skills (Domain Expertise Injection)\n\n" "Skills inject specialized instructions into the subagent. Read the description to understand when each skill applies.\n\n" "| Skill | Expertise Domain |\n" "|-------|------------------|\n" + "\n".join(skill_rows) + "\n\n---\n\n" - "### MANDATORY: Category + Skill Selection Protocol\n\n" - "**STEP 1: Select Category**\n" - "- Read each category's description\n" - "- Match task requirements to category domain\n" - "- Select the category whose domain BEST fits the task\n\n" - "**STEP 2: Evaluate ALL Skills**\n" + "### MANDATORY: Skill Selection Protocol\n\n" "For EVERY skill listed above, ask yourself:\n" '> "Does this skill\'s expertise domain overlap with my task?"\n\n' "- If YES → INCLUDE in `load_skills=[...]`\n" @@ -252,14 +237,14 @@ def build_category_skills_delegation_guide( "### Delegation Pattern\n\n" "```typescript\n" "delegate_task(\n" - ' category="[selected-category]",\n' + ' subagent_type="[selected-agent]",\n' ' load_skills=["skill-1", "skill-2"], // Include ALL relevant skills\n' ' prompt="..."\n' ")\n" "```\n\n" "**ANTI-PATTERN (will produce poor results):**\n" "```typescript\n" - 'delegate_task(category="...", load_skills=[], prompt="...") // Empty load_skills without justification\n' + 'delegate_task(subagent_type="...", load_skills=[], prompt="...") // Empty load_skills without justification\n' "```" ) @@ -319,48 +304,10 @@ def build_anti_patterns_section() -> str: ) -def build_ultrawork_section( - agents: List[AvailableAgent], - categories: List[AvailableCategory], - skills: List[AvailableSkill], -) -> str: - lines: List[str] = [] - - if categories: - lines.append("**Categories** (for implementation tasks):") - for cat in categories: - short_desc = cat.description or cat.name - lines.append(f"- `{cat.name}`: {short_desc}") - lines.append("") - - if skills: - lines.append("**Skills** (combine with categories - EVALUATE ALL for relevance):") - for skill in skills: - short_desc = skill.description.split(".")[0] or skill.description - lines.append(f"- `{skill.name}`: {short_desc}") - lines.append("") - - if agents: - ultrawork_agent_priority = ["explore", "librarian", "plan", "oracle"] - sorted_agents = list(agents) - sorted_agents.sort( - key=lambda a: ultrawork_agent_priority.index(a.name) - if a.name in ultrawork_agent_priority - else 999 - ) - lines.append("**Agents** (for specialized consultation/exploration):") - for agent in sorted_agents: - short_desc = agent.description.split(".")[0] or agent.description - suffix = " (multiple)" if agent.name in ("explore", "librarian") else "" - lines.append(f"- `{agent.name}{suffix}`: {short_desc}") - - return "\n".join(lines) - - def build_workflows_section(workflows: List[AvailableWorkflow]) -> str: """Render the available workflows section for injection into system prompts. - Mirrors the pattern used by build_category_skills_delegation_guide() for + Mirrors the pattern used by build_skills_delegation_guide() for skills, so agents know which workflows exist before calling run_workflow. """ if not workflows: diff --git a/flocks/agent/registry.py b/flocks/agent/registry.py index b5b004b46..ee1b5ca32 100644 --- a/flocks/agent/registry.py +++ b/flocks/agent/registry.py @@ -4,7 +4,7 @@ Replaces the old flocks.agent.core.registry module (kept as a shim). Loading order: - ① Collect context: tools, skills, categories + ① Collect context: tools and skills ② scan_and_load() — built-in YAML agents + plugin YAML agents ③ Python plugin agents via PluginLoader + cfg.agent user overrides ④ inject_dynamic_prompts() — phase-2 dynamic prompt injection @@ -34,7 +34,6 @@ AgentModel, AgentPromptMetadata, AvailableAgent, - AvailableCategory, AvailableSkill, AvailableWorkflow, DelegationTrigger, @@ -76,9 +75,11 @@ def _set_agents_ref(agents: Dict[str, AgentInfo]) -> None: # --------------------------------------------------------------------------- def is_delegatable(agent_name: str) -> bool: - if _agents_ref and agent_name in _agents_ref: - return bool(_agents_ref[agent_name].delegatable) - return True # unknown → safe default + resolved = AGENT_ALIASES.get(agent_name, agent_name) + if not _agents_ref: + return False + agent = _agents_ref.get(resolved) + return bool(agent and agent.delegatable and not agent.hidden) def get_agent_mode(agent_name: str) -> Optional[str]: @@ -95,7 +96,7 @@ def is_hidden(agent_name: str) -> bool: def list_delegatable_agents() -> List[str]: if _agents_ref: - return [n for n, a in _agents_ref.items() if a.delegatable] + return [n for n, a in _agents_ref.items() if a.delegatable and not a.hidden] return [] @@ -242,13 +243,12 @@ async def _load_agents() -> Dict[str, AgentInfo]: """ 4-step agent loading: - ① Context — tools, skills, categories + ① Context — tools and skills ② YAML — scan_and_load() from built-in + plugin directories ③ Plugins — PluginLoader Python modules + cfg.agent overrides ④ Prompts — inject_dynamic_prompts() for phase-2 dynamic agents """ # Lazy imports to avoid circular dependencies - from flocks.tool.delegate_task_constants import CATEGORY_DESCRIPTIONS, DEFAULT_CATEGORIES from flocks.tool.registry import ToolRegistry cfg = await Config.get() @@ -264,19 +264,6 @@ async def _load_agents() -> Dict[str, AgentInfo]: AvailableSkill(name=s.name, description=s.description, location=s.source or "project") for s in skills ] - category_configs = {**DEFAULT_CATEGORIES, **(cfg.categories or {})} - available_categories = [ - AvailableCategory( - name=name, - description=( - cfg.categories.get(name).description - if cfg.categories and cfg.categories.get(name) - else CATEGORY_DESCRIPTIONS.get(name, name) - ), - ) - for name in category_configs.keys() - ] - # Discover available workflows (best-effort; failure must not block agent load) available_workflows: List[AvailableWorkflow] = [] try: @@ -430,7 +417,6 @@ def _permission_dict_to_tools(permission_cfg: Dict[str, Any]) -> List[str]: available_agents, categorized_tools, available_skills, - available_categories, available_workflows, ) diff --git a/flocks/channel/builtin/feishu/monitor.py b/flocks/channel/builtin/feishu/monitor.py index e2339a9d6..471576471 100644 --- a/flocks/channel/builtin/feishu/monitor.py +++ b/flocks/channel/builtin/feishu/monitor.py @@ -172,7 +172,7 @@ def _build_ws_client( native_client_cls = ws_module.Client class _Dispatcher: - def do_without_validation(self, payload: bytes) -> None: + def _do_without_validation(self, payload: bytes) -> None: try: data = json.loads(payload.decode("utf-8")) except Exception as e: @@ -181,6 +181,9 @@ def do_without_validation(self, payload: bytes) -> None: event_handler(data) return None + def do_without_validation(self, payload: bytes) -> None: + return self._do_without_validation(payload) + class _CompatWSClient: def __init__(self) -> None: self._client: Any | None = None diff --git a/flocks/channel/inbound/dispatcher.py b/flocks/channel/inbound/dispatcher.py index 1e516c91b..910020e0d 100644 --- a/flocks/channel/inbound/dispatcher.py +++ b/flocks/channel/inbound/dispatcher.py @@ -17,7 +17,10 @@ from typing import Any, Optional from flocks.channel.base import ChatType, InboundMessage, OutboundContext -from flocks.channel.inbound.session_binding import SessionBindingService +from flocks.channel.inbound.session_binding import ( + SessionBindingService, + is_channel_media_placeholder, +) from flocks.config.config import ChannelConfig from flocks.utils.log import Log @@ -872,6 +875,7 @@ async def _handle_session_command( from flocks.session.session import Session from flocks.channel.inbound.session_binding import ( _build_title, + _build_title_fallback, resolve_channel_session_owner_kwargs, ) @@ -881,10 +885,15 @@ async def _handle_session_command( return owner_kwargs = await resolve_channel_session_owner_kwargs(session) + title = ( + _build_title(msg, initial_text) + if initial_text and initial_text.strip() + else _build_title_fallback(msg) + ) new_session = await Session.create( project_id=session.project_id, directory=session.directory, - title=_build_title(msg), + title=title, agent=session.agent, **Session.inherited_model_kwargs(session), **owner_kwargs, @@ -895,26 +904,6 @@ async def _handle_session_command( agent_id=new_session.agent, scope_override=scope_override, ) - # Archive the previous session so it no longer shows up as an *active* - # IM session. The binding has already moved to ``new_session`` via - # ``rebind`` above, but the old session retains ``status="active"`` and - # the same ``[Feishu]/[Wecom]/[Dingtalk]`` title prefix. Without this, - # repeated ``/new`` leaves multiple active sessions for the same - # conversation, and unattended scheduled tasks (which resolve the IM - # target via ``session_list(status="active")`` + title prefix) can no - # longer tell which one is current — sending to the wrong session or - # failing outright. Best-effort: archiving failure must not abort /new. - try: - await Session.update( - session.project_id, - session.id, - status="archived", - ) - except Exception as exc: - log.warning("dispatcher.archive_previous_session_failed", { - "session_id": session.id, - "error": str(exc), - }) await self._trigger_command_hook( "new", session.id, @@ -1338,7 +1327,7 @@ async def _append_user_message_unchecked( parts = await Message.parts(message.id, session_id=session_id) for p in parts: if p.type == "text" and hasattr(p, "text") and p.text: - if _is_placeholder_text(p.text): + if is_channel_media_placeholder(p.text): updated = TextPart( id=p.id, sessionID=session_id, @@ -1589,23 +1578,6 @@ def register_inbound_media_downloader(channel_id: str, downloader: Any) -> None: _DOWNLOADERS[channel_id] = downloader -def _is_placeholder_text(text: str) -> bool: - """True if *text* is one of the channel-generated media placeholders.""" - if not text: - return False - placeholders = ( - "[图片消息]", - "[文件消息]", - "[Image]", - "[Attachment]", - "[图片]", - "[文件]", - ) - if text in placeholders: - return True - return text.startswith("[文件消息:") or text.startswith("[图片消息:") - - # Best-effort eager registration at import time. Channels that need # more control can call ``register_inbound_media_downloader`` themselves. try: diff --git a/flocks/channel/inbound/session_binding.py b/flocks/channel/inbound/session_binding.py index 3f023a358..fa8840ca2 100644 --- a/flocks/channel/inbound/session_binding.py +++ b/flocks/channel/inbound/session_binding.py @@ -17,6 +17,7 @@ import asyncio import os +import re import time from dataclasses import dataclass from typing import Literal, Optional @@ -29,6 +30,14 @@ log = Log.create(service="channel.binding") +_CHANNEL_TITLE_PLACEHOLDER_RE = re.compile( + r"^\[(?:图片消息|文件消息|图片|文件|音频|语音消息|视频|圆形视频|" + r"贴纸(?: [^\]]*)?|动图|位置|联系人|Image|Attachment)" + r"(?::[^\]]*)?\](?:(?:\s*:\s*|\s+)(.*))?$", + re.IGNORECASE, +) +_CHANNEL_TITLE_MAX_LENGTH = 50 + # Supported group session scope values (mirrors FeishuGroupConfig.group_session_scope) GroupSessionScope = Literal["group", "group_sender", "group_topic", "group_topic_sender"] @@ -455,13 +464,25 @@ async def get_bindings_by_session(self, session_id: str) -> list[SessionBinding] async def list_bindings( self, channel_id: Optional[str] = None, + session_ids: Optional[list[str]] = None, ) -> list[SessionBinding]: db = await _get_db() sql = "SELECT * FROM channel_bindings" - params: tuple = () + conditions: list[str] = [] + params: list[str] = [] if channel_id: - sql += " WHERE channel_id = ?" - params = (channel_id,) + conditions.append("channel_id = ?") + params.append(channel_id) + if session_ids is not None: + unique_session_ids = list(dict.fromkeys(session_ids)) + if not unique_session_ids: + return [] + conditions.append( + f"session_id IN ({','.join('?' for _ in unique_session_ids)})" + ) + params.extend(unique_session_ids) + if conditions: + sql += " WHERE " + " AND ".join(conditions) sql += " ORDER BY last_message_at DESC" cursor = await db.execute(sql, params) @@ -646,7 +667,39 @@ def _mark_cwd_fallback_warned() -> None: _CWD_FALLBACK_WARNED = True -def _build_title(msg: InboundMessage) -> str: +def is_channel_media_placeholder(text: str) -> bool: + """Return whether text is only a channel-generated media placeholder.""" + match = _CHANNEL_TITLE_PLACEHOLDER_RE.fullmatch(text.strip()) + return bool(match and not (match.group(1) or "").strip()) + + +def extract_channel_title_text(raw_text: str) -> str: + """Extract the first user-authored title candidate from channel text.""" + for line in raw_text.splitlines(): + candidate = line.strip() + if not candidate or candidate.startswith("__merge_forward_expand__"): + continue + if candidate == "[Merged forward message]": + continue + placeholder = _CHANNEL_TITLE_PLACEHOLDER_RE.fullmatch(candidate) + if placeholder: + caption = (placeholder.group(1) or "").strip() + if caption: + return caption + continue + return candidate + return "" + + +def format_channel_title(channel_id: str, title_text: str) -> str: + """Add the channel prefix and apply the session-title length limit.""" + title_text = title_text.strip() + if len(title_text) > _CHANNEL_TITLE_MAX_LENGTH: + title_text = title_text[:_CHANNEL_TITLE_MAX_LENGTH - 3] + "..." + return f"[{channel_id.capitalize()}] {title_text}" + + +def _build_title_fallback(msg: InboundMessage) -> str: prefix = msg.channel_id.capitalize() if msg.chat_type == ChatType.DIRECT: who = msg.sender_name or msg.sender_id @@ -654,6 +707,17 @@ def _build_title(msg: InboundMessage) -> str: return f"[{prefix}] {msg.chat_id}" +def _build_title(msg: InboundMessage, text_override: Optional[str] = None) -> str: + if text_override is None: + raw_text = msg.mention_text or msg.text or "" + else: + raw_text = text_override + title_text = extract_channel_title_text(raw_text) + if title_text: + return format_channel_title(msg.channel_id, title_text) + return _build_title_fallback(msg) + + def _resolve_session_key( msg: InboundMessage, scope_override: Optional[GroupSessionScope] = None, diff --git a/flocks/cli/main.py b/flocks/cli/main.py index 626b9f5a7..45981f2b5 100644 --- a/flocks/cli/main.py +++ b/flocks/cli/main.py @@ -284,7 +284,8 @@ def restart( webui_port: Optional[int] = typer.Option(None, "--webui-port", help="WebUI port"), ): """ - Restart Flocks service. + Restart Flocks service. Agents must use `flocks restart --server-only`; + bare restart stops the supervisor and terminates the running agent. """ try: if server_only: diff --git a/flocks/cli/session_runner.py b/flocks/cli/session_runner.py index 994042259..5d5aa6766 100644 --- a/flocks/cli/session_runner.py +++ b/flocks/cli/session_runner.py @@ -654,7 +654,7 @@ async def _on_tool_start(self, tool_name: str, arguments: Dict[str, Any]) -> Non self._flush_content() if tool_name in DELEGATE_TOOLS: - agent = arguments.get("subagent_type") or arguments.get("category") or "unknown" + agent = arguments.get("subagent_type") or "unknown" desc = arguments.get("description", "子任务") bg = " [dim](后台)[/dim]" if arguments.get("run_in_background") else "" self.console.print( diff --git a/flocks/config/config.py b/flocks/config/config.py index 99fa1f3af..d6342f0a5 100644 --- a/flocks/config/config.py +++ b/flocks/config/config.py @@ -149,19 +149,6 @@ def process_agent(self): return self -# ==================== Category Configuration ==================== - -class CategoryConfig(BaseModel): - """Delegate-task category configuration""" - model_config = {"extra": "allow", "populate_by_name": True} - - model: Optional[str] = None - variant: Optional[str] = None - prompt_append: Optional[str] = Field(None, alias="promptAppend") - description: Optional[str] = None - is_unstable_agent: Optional[bool] = Field(None, alias="isUnstableAgent") - - # ==================== Command Configuration ==================== class CommandConfig(BaseModel): @@ -665,7 +652,6 @@ class ConfigInfo(BaseModel): mode: Optional[Dict[str, AgentConfig]] = Field(None, description="@deprecated Use 'agent'") agent: Optional[Dict[str, AgentConfig]] = None provider: Optional[Dict[str, ProviderConfig]] = None - categories: Optional[Dict[str, CategoryConfig]] = None mcp: Optional[Dict[str, Union[McpConfig, Dict[str, Any]]]] = None formatter: Optional[Union[Literal[False], Dict[str, Any]]] = None lsp: Optional[Union[Literal[False], Dict[str, Any]]] = None @@ -742,6 +728,28 @@ class ConfigInfo(BaseModel): description="Console portal base URL used by OSS console account login redirect.", ) + @model_validator(mode="before") + @classmethod + def remove_legacy_delegate_categories(cls, data): + """Drop the removed categories config instead of preserving it as an extra.""" + if not isinstance(data, dict) or "categories" not in data: + return data + + cleaned = dict(data) + cleaned.pop("categories", None) + from flocks.utils.log import Log + + Log.create(service="config").warn( + "config.categories_removed", + { + "message": ( + "The categories configuration is no longer supported; " + "select a delegatable agent with subagent_type." + ) + }, + ) + return cleaned + @model_validator(mode='after') def post_process(self): """Post-processing like TypeScript""" diff --git a/flocks/contracts/access/driver.py b/flocks/contracts/access/driver.py index 25af7108a..3a6874fa8 100644 --- a/flocks/contracts/access/driver.py +++ b/flocks/contracts/access/driver.py @@ -5,7 +5,7 @@ import json import re import sqlite3 -from collections.abc import Iterable +from collections.abc import Callable, Iterable from datetime import datetime from pathlib import Path from typing import Any @@ -119,7 +119,11 @@ def execute(self, plan: QueryPlan) -> DriverResult: if record.get("is_duplicate") is True: duplicates += 1 continue - if not self._matches_predicates(record, plan.policy_plan.driver_predicates): + if not self._matches_predicates( + record, + plan.policy_plan.driver_predicates, + plan.binding.predicate_value_resolver, + ): continue record_id = _read_string(record.get("id"), "") @@ -152,8 +156,13 @@ def execute(self, plan: QueryPlan) -> DriverResult: def _assert_allowed(self, path: Path, allowlist_roots: tuple[Path, ...]) -> None: JsonlDriverExecutor()._assert_allowed(path, allowlist_roots) - def _matches_predicates(self, record: dict[str, Any], predicates: tuple[Predicate, ...]) -> bool: - return JsonlDriverExecutor()._matches_predicates(record, predicates) + def _matches_predicates( + self, + record: dict[str, Any], + predicates: tuple[Predicate, ...], + value_resolver: Callable[[dict[str, Any], str], Any] | None = None, + ) -> bool: + return JsonlDriverExecutor()._matches_predicates(record, predicates, value_resolver) def _matches_event_time_range(self, record: dict[str, Any], start_time: int | None, end_time: int | None) -> bool: return JsonlDriverExecutor()._matches_event_time_range(record, start_time, end_time) @@ -184,7 +193,11 @@ def execute(self, plan: QueryPlan) -> DriverResult: if record.get("is_duplicate") is True: duplicates += 1 continue - if not self._matches_predicates(record, plan.policy_plan.driver_predicates): + if not self._matches_predicates( + record, + plan.policy_plan.driver_predicates, + plan.binding.predicate_value_resolver, + ): continue record_id = _read_string(record.get("id"), "") @@ -303,9 +316,18 @@ def _iter_records(self, path: Path) -> Iterable[dict[str, Any] | None]: continue yield value if isinstance(value, dict) else None - def _matches_predicates(self, record: dict[str, Any], predicates: tuple[Predicate, ...]) -> bool: + def _matches_predicates( + self, + record: dict[str, Any], + predicates: tuple[Predicate, ...], + value_resolver: Callable[[dict[str, Any], str], Any] | None = None, + ) -> bool: for predicate in predicates: - value = record.get(predicate.field) + value = ( + value_resolver(record, predicate.field) + if value_resolver is not None + else record.get(predicate.field) + ) if predicate.operator == "in": allowed = {_normalize_compare(item) for item in predicate.values} if _normalize_compare(value) not in allowed: diff --git a/flocks/contracts/access/models.py b/flocks/contracts/access/models.py index fcf471ef7..1e7eb0b74 100644 --- a/flocks/contracts/access/models.py +++ b/flocks/contracts/access/models.py @@ -2,6 +2,7 @@ from __future__ import annotations +from collections.abc import Callable from dataclasses import dataclass, field from pathlib import Path from typing import Any, Literal @@ -76,6 +77,7 @@ class Binding: driver_available_fields: frozenset[str] driver_allowlist_roots: tuple[Path, ...] driver_options: dict[str, Any] = field(default_factory=dict) + predicate_value_resolver: Callable[[dict[str, Any], str], Any] | None = None capabilities: frozenset[str] = frozenset({"query"}) diff --git a/flocks/ingest/kafka/manager.py b/flocks/ingest/kafka/manager.py index 9b799a0d6..78c2c305e 100644 --- a/flocks/ingest/kafka/manager.py +++ b/flocks/ingest/kafka/manager.py @@ -22,6 +22,7 @@ import asyncio import hashlib import json +import threading import time import uuid from dataclasses import dataclass @@ -81,6 +82,14 @@ _MAX_PARTITION_FETCH_BYTES = 4 * 1024 * 1024 _MAX_POLL_RECORDS = 16 + +def _worker_count_for_trigger(trigger: TriggerDefinition) -> int: + return min(_MAX_CONCURRENT_EXECUTIONS, max(1, int(trigger.concurrency.maxParallel))) + + +def _queue_size_for_trigger(trigger: TriggerDefinition) -> int: + return min(_MAX_QUEUE_SIZE, max(1, int(trigger.concurrency.queueSize))) + _KAFKA_STORAGE_LIST_KEYS = DEFAULT_LARGE_LIST_KEYS | frozenset( { "duplicate_alerts", @@ -244,6 +253,10 @@ def __init__(self) -> None: self._queues: dict[str, asyncio.Queue] = {} # Per-workflow fixed worker pool draining the queue self._worker_pools: dict[str, List[asyncio.Task]] = {} + # One cancellation event per consumer generation. It exists before + # workers start, closing the stop-vs-run-registration race. + self._generation_cancel_events: dict[str, threading.Event] = {} + self._draining_workers: dict[str, set[asyncio.Task]] = {} # Per-workflow consumer runtime status for the kafka-status API. # State values: "connecting" | "running" | "failed" | "stopped". self._status: dict[str, Dict[str, Any]] = {} @@ -315,25 +328,54 @@ async def stop_all(self) -> None: await self.stop_workflow(workflow_id) async def _cleanup_runtime_resources(self, workflow_id: str) -> None: - # Cancel all worker pool tasks; pop first so callers observing a stopped - # consumer see an empty pool immediately. + abort = self._abort_events.get(workflow_id) + if abort is not None: + abort.set() + cancel_event = self._generation_cancel_events.pop(workflow_id, None) + if cancel_event is not None: + cancel_event.set() + + # Let workers drain the current ``to_thread`` workflow after the + # cooperative cancellation signal. Pending workers stay tracked. pool = self._worker_pools.pop(workflow_id, None) if pool: - for worker in pool: - if not worker.done(): - worker.cancel() - try: - await asyncio.wait_for( - asyncio.gather(*pool, return_exceptions=True), - timeout=5.0, - ) - except (asyncio.TimeoutError, asyncio.CancelledError): - pass + _, pending = await asyncio.wait(pool, timeout=5.0) + self._track_draining_workers(workflow_id, pending) self._queues.pop(workflow_id, None) self._abort_events.pop(workflow_id, None) self._ready.pop(workflow_id, None) + def _track_draining_workers( + self, + workflow_id: str, + workers: set[asyncio.Task], + ) -> None: + if not workers: + return + bucket = self._draining_workers.setdefault(workflow_id, set()) + bucket.update(workers) + + def _discard(done: asyncio.Task) -> None: + current = self._draining_workers.get(workflow_id) + if current is None: + return + current.discard(done) + if not current: + self._draining_workers.pop(workflow_id, None) + + for worker in workers: + worker.add_done_callback(_discard) + + def _active_draining_workers(self, workflow_id: str) -> set[asyncio.Task]: + workers = self._draining_workers.get(workflow_id, set()) + active = {worker for worker in workers if not worker.done()} + if active: + self._draining_workers[workflow_id] = active + else: + self._draining_workers.pop(workflow_id, None) + return active + def get_consumer_status(self, workflow_id: str) -> Dict[str, Any]: """Return a snapshot of the consumer runtime state for ``workflow_id``. @@ -384,6 +426,10 @@ async def restart_workflow( surface connection errors to the user. """ await self.stop_workflow(workflow_id) + if self._active_draining_workers(workflow_id): + error = "previous_workers_still_draining" + self._status[workflow_id] = {"state": "failed", "error": error} + return {"state": "failed", "error": error} try: data = await WorkflowStore.get_config(workflow_id, kind="workflow_kafka_config") except Exception as exc: @@ -434,11 +480,15 @@ async def restart_workflow( group_id = str(data.get("inputGroupId") or "").strip() or f"flocks-consumer-{workflow_id}" configured_inputs = _strip_execution_only_comments(trigger.inputs if isinstance(trigger.inputs, dict) else {}) - queue: asyncio.Queue = asyncio.Queue(maxsize=_MAX_QUEUE_SIZE) + queue_capacity = _queue_size_for_trigger(trigger) + worker_count = _worker_count_for_trigger(trigger) + queue: asyncio.Queue = asyncio.Queue(maxsize=queue_capacity) self._queues[workflow_id] = queue abort = asyncio.Event() self._abort_events[workflow_id] = abort + generation_cancel_event = threading.Event() + self._generation_cancel_events[workflow_id] = generation_cancel_event ready = asyncio.Event() self._ready[workflow_id] = ready @@ -451,10 +501,9 @@ async def restart_workflow( "groupId": group_id, } - # Fixed worker pool drains the queue (at most _MAX_CONCURRENT_EXECUTIONS - # concurrent runs). + # Trigger-configured worker pool, bounded by service safety caps. workers: List[asyncio.Task] = [] - for i in range(_MAX_CONCURRENT_EXECUTIONS): + for i in range(worker_count): workers.append( asyncio.create_task( self._worker_loop( @@ -465,6 +514,7 @@ async def restart_workflow( queue, abort, input_topic, + generation_cancel_event, ), name=f"kafka-worker-{workflow_id}-{i}", ) @@ -635,7 +685,13 @@ async def _worker_loop( queue: asyncio.Queue, abort: asyncio.Event, source: str, + generation_cancel_event: Optional[threading.Event] = None, ) -> None: + run_cancel_event = ( + generation_cancel_event + or self._generation_cancel_events.get(workflow_id) + or threading.Event() + ) while not abort.is_set(): try: msg = await asyncio.wait_for(queue.get(), timeout=0.5) @@ -654,6 +710,7 @@ async def _worker_loop( configured_inputs, trigger=trigger, source=source, + generation_cancel_event=run_cancel_event, ) except asyncio.CancelledError: return @@ -673,7 +730,9 @@ async def _trigger_workflow( *, trigger: Optional[TriggerDefinition] = None, source: Optional[str] = None, + generation_cancel_event: Optional[threading.Event] = None, ) -> None: + run_cancel_event = generation_cancel_event or threading.Event() trigger = trigger or TriggerDefinition.model_validate( { "id": "kafka-default", @@ -735,6 +794,7 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: run_id=exec_id, trace=False, execution_profile="high_frequency", + cancel=run_cancel_event.is_set, on_step_complete=step_recorder.on_step_complete, tool_context=tool_context, ) diff --git a/flocks/ingest/syslog/manager.py b/flocks/ingest/syslog/manager.py index b81f88e3e..184621f57 100644 --- a/flocks/ingest/syslog/manager.py +++ b/flocks/ingest/syslog/manager.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import threading import time import uuid from typing import Any, Dict, List @@ -53,6 +54,14 @@ _DROP_LOG_WINDOW_S = 1.0 +def _worker_count_for_trigger(trigger: TriggerDefinition) -> int: + return min(_MAX_CONCURRENT_EXECUTIONS, max(1, int(trigger.concurrency.maxParallel))) + + +def _queue_size_for_trigger(trigger: TriggerDefinition) -> int: + return min(_MAX_QUEUE_SIZE, max(1, int(trigger.concurrency.queueSize))) + + class _DropWarningThrottle: """Aggregate per-workflow ``QueueFull`` drops into windowed warnings. @@ -135,6 +144,11 @@ def __init__(self) -> None: self._queues: dict[str, asyncio.Queue] = {} # Per-workflow fixed worker pool draining the queue self._worker_pools: dict[str, List[asyncio.Task]] = {} + # One cancellation event per listener generation. It is installed + # before any worker starts, so stop cannot miss a run that is still + # creating its execution record or tool context. + self._generation_cancel_events: dict[str, threading.Event] = {} + self._draining_workers: dict[str, set[asyncio.Task]] = {} # Per-workflow listener runtime status for the syslog-status API. # Possible state values: "binding" | "listening" | "failed" | "stopped". self._listener_status: dict[str, Dict[str, Any]] = {} @@ -229,7 +243,40 @@ def get_listener_status(self, workflow_id: str) -> Dict[str, Any]: status["workerCount"] = sum(1 for t in pool if not t.done()) return status + def _track_draining_workers( + self, + workflow_id: str, + workers: set[asyncio.Task], + ) -> None: + if not workers: + return + bucket = self._draining_workers.setdefault(workflow_id, set()) + bucket.update(workers) + + def _discard(done: asyncio.Task) -> None: + current = self._draining_workers.get(workflow_id) + if current is None: + return + current.discard(done) + if not current: + self._draining_workers.pop(workflow_id, None) + + for worker in workers: + worker.add_done_callback(_discard) + + def _active_draining_workers(self, workflow_id: str) -> set[asyncio.Task]: + workers = self._draining_workers.get(workflow_id, set()) + active = {worker for worker in workers if not worker.done()} + if active: + self._draining_workers[workflow_id] = active + else: + self._draining_workers.pop(workflow_id, None) + return active + async def stop_workflow(self, workflow_id: str) -> None: + cancel_event = self._generation_cancel_events.pop(workflow_id, None) + if cancel_event is not None: + cancel_event.set() ev = self._abort_events.pop(workflow_id, None) if ev is not None: ev.set() @@ -240,20 +287,13 @@ async def stop_workflow(self, workflow_id: str) -> None: await task except asyncio.CancelledError: pass - # Cancel all worker pool tasks; pop first so callers observing a - # stopped listener see an empty pool immediately. + # Let workers drain their current ``to_thread`` workflow after the + # cooperative cancel signal. Pending workers stay tracked instead of + # becoming invisible orphan threads. pool = self._worker_pools.pop(workflow_id, None) if pool: - for w in pool: - if not w.done(): - w.cancel() - try: - await asyncio.wait_for( - asyncio.gather(*pool, return_exceptions=True), - timeout=5.0, - ) - except (asyncio.TimeoutError, asyncio.CancelledError): - pass + _, pending = await asyncio.wait(pool, timeout=5.0) + self._track_draining_workers(workflow_id, pending) self._queues.pop(workflow_id, None) self._listener_ready.pop(workflow_id, None) if workflow_id in self._listener_status: @@ -274,6 +314,10 @@ async def restart_workflow( user instead of silently leaving the listener in a failed state. """ await self.stop_workflow(workflow_id) + if self._active_draining_workers(workflow_id): + error = "previous_workers_still_draining" + self._listener_status[workflow_id] = {"state": "failed", "error": error} + return {"state": "failed", "error": error} try: data = await WorkflowStore.get_config(workflow_id, kind="workflow_syslog_config") except Exception as exc: @@ -316,11 +360,15 @@ async def restart_workflow( host = str(data.get("host") or "0.0.0.0") port = int(data.get("port") or 5140) protocol = str(data.get("protocol") or "udp").lower() - queue: asyncio.Queue = asyncio.Queue(maxsize=_MAX_QUEUE_SIZE) + queue_capacity = _queue_size_for_trigger(trigger) + worker_count = _worker_count_for_trigger(trigger) + queue: asyncio.Queue = asyncio.Queue(maxsize=queue_capacity) self._queues[workflow_id] = queue abort = asyncio.Event() self._abort_events[workflow_id] = abort + generation_cancel_event = threading.Event() + self._generation_cancel_events[workflow_id] = generation_cancel_event ready = asyncio.Event() self._listener_ready[workflow_id] = ready @@ -333,14 +381,19 @@ async def restart_workflow( "protocol": protocol, } - # Spin up a fixed worker pool: exactly _MAX_CONCURRENT_EXECUTIONS - # coroutines drain the queue. pending tasks cannot exceed this number, - # which is the actual backpressure invariant we want. + # Spin up the trigger-configured worker pool within service safety caps. workers: List[asyncio.Task] = [] - for i in range(_MAX_CONCURRENT_EXECUTIONS): + for i in range(worker_count): workers.append( asyncio.create_task( - self._worker_loop(workflow_id, workflow_plan, trigger, queue, abort), + self._worker_loop( + workflow_id, + workflow_plan, + trigger, + queue, + abort, + generation_cancel_event, + ), name=f"syslog-worker-{workflow_id}-{i}", ) ) @@ -489,6 +542,7 @@ async def _worker_loop( trigger: TriggerDefinition, queue: asyncio.Queue, abort: asyncio.Event, + generation_cancel_event: Optional[threading.Event] = None, ) -> None: """One worker drains the queue serially. @@ -496,6 +550,11 @@ async def _worker_loop( do not spawn additional asyncio.Tasks per message so the total number of in-flight workflow runs is exactly ``_MAX_CONCURRENT_EXECUTIONS``. """ + run_cancel_event = ( + generation_cancel_event + or self._generation_cancel_events.get(workflow_id) + or threading.Event() + ) while not abort.is_set(): try: msg = await asyncio.wait_for(queue.get(), timeout=0.5) @@ -511,6 +570,7 @@ async def _worker_loop( next(iter(trigger.mapping or {}), "syslog_message"), trigger=trigger, source=f"{(trigger.source or {}).get('protocol', 'udp')}://{(trigger.source or {}).get('host', '0.0.0.0')}:{(trigger.source or {}).get('port', 5140)}", + generation_cancel_event=run_cancel_event, ) except asyncio.CancelledError: return @@ -529,7 +589,9 @@ async def _trigger_workflow( *, trigger: Optional[TriggerDefinition] = None, source: Optional[str] = None, + generation_cancel_event: Optional[threading.Event] = None, ) -> None: + run_cancel_event = generation_cancel_event or threading.Event() trigger = trigger or TriggerDefinition.model_validate( { "id": "syslog-default", @@ -578,6 +640,7 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: run_id=exec_id, trace=False, execution_profile="high_frequency", + cancel=run_cancel_event.is_set, on_step_complete=step_recorder.on_step_complete, tool_context=tool_context, ) diff --git a/flocks/provider/catalog.json b/flocks/provider/catalog.json index 69f35518f..f4c9f60b7 100644 --- a/flocks/provider/catalog.json +++ b/flocks/provider/catalog.json @@ -52,6 +52,25 @@ "THREATBOOK_CN_LLM_API_KEY" ], "models": { + "deepseek-v4-flash-0731": { + "name": "deepseek-v4-flash-0731", + "family": "deepseek-v4", + "capabilities": { + "supports_tools": true, + "supports_streaming": true + }, + "limits": { + "context_window": 1000000, + "max_input_tokens": 1000000, + "max_output_tokens": 384000 + }, + "pricing": { + "input": 1.0, + "output": 2.0, + "cache_read": 0.2, + "currency": "CNY" + } + }, "kimi-k2.7-code": { "name": "kimi-k2.7-code", "family": "kimi-k2.7-code", @@ -289,6 +308,25 @@ "THREATBOOK_IO_LLM_API_KEY" ], "models": { + "deepseek-v4-flash-0731": { + "name": "deepseek-v4-flash-0731", + "family": "deepseek-v4", + "capabilities": { + "supports_tools": true, + "supports_streaming": true + }, + "limits": { + "context_window": 1000000, + "max_input_tokens": 1000000, + "max_output_tokens": 384000 + }, + "pricing": { + "input": 1.0, + "output": 2.0, + "cache_read": 0.2, + "currency": "CNY" + } + }, "kimi-k2.7-code": { "name": "kimi-k2.7-code", "family": "kimi-k2.7-code", diff --git a/flocks/provider/provider.py b/flocks/provider/provider.py index 65465ee5c..44084d270 100644 --- a/flocks/provider/provider.py +++ b/flocks/provider/provider.py @@ -49,6 +49,8 @@ def _model_info_signature(model: "ModelInfo") -> tuple: ( pricing.get("input") if isinstance(pricing, dict) else None, pricing.get("output") if isinstance(pricing, dict) else None, + pricing.get("cache_read") if isinstance(pricing, dict) else None, + pricing.get("cache_write") if isinstance(pricing, dict) else None, pricing.get("currency") if isinstance(pricing, dict) else None, ) if pricing is not None @@ -849,13 +851,20 @@ async def apply_config(cls, config: Optional[Any] = None, provider_id: Optional[ # Create ModelInfo from config _input_price = model_dict.get("input_price") _output_price = model_dict.get("output_price") + _cache_read_price = model_dict.get("cache_read_price") _pricing = None - if _input_price is not None or _output_price is not None: + if ( + _input_price is not None + or _output_price is not None + or _cache_read_price is not None + ): _pricing = { "input": float(_input_price or 0.0), "output": float(_output_price or 0.0), "currency": model_dict.get("currency", "USD"), } + if _cache_read_price is not None: + _pricing["cache_read"] = float(_cache_read_price) model_info = ModelInfo( id=model_id, name=model_dict.get("name", model_id), @@ -1233,6 +1242,8 @@ def _build_model_definition(self, model: "ModelInfo") -> "ModelDefinition": pricing = PriceConfig( input=model.pricing.get("input", 0.0), output=model.pricing.get("output", 0.0), + cache_read=model.pricing.get("cache_read"), + cache_write=model.pricing.get("cache_write"), currency=model.pricing.get("currency", "USD"), ) max_output = model.capabilities.max_tokens or 4096 @@ -1325,6 +1336,8 @@ def _apply_config_overrides(self, catalog_def: "ModelDefinition", model: "ModelI overridden.pricing = PriceConfig( input=model.pricing.get("input", 0.0), output=model.pricing.get("output", 0.0), + cache_read=model.pricing.get("cache_read"), + cache_write=model.pricing.get("cache_write"), currency=model.pricing.get("currency", "USD"), ) diff --git a/flocks/server/routes/custom_provider.py b/flocks/server/routes/custom_provider.py index 969d754ca..9e5bec81c 100644 --- a/flocks/server/routes/custom_provider.py +++ b/flocks/server/routes/custom_provider.py @@ -131,6 +131,7 @@ class CreateModelReq(BaseModel): supports_reasoning: bool = True input_price: float = Field(0.0, ge=0) output_price: float = Field(0.0, ge=0) + cache_read_price: Optional[float] = Field(None, ge=0) currency: str = "USD" @@ -144,6 +145,7 @@ class ModelResp(BaseModel): max_output_tokens: int input_price: float output_price: float + cache_read_price: Optional[float] = None currency: str created_at: str @@ -268,6 +270,7 @@ async def list_models(provider_id: str): max_output_tokens=mcfg.get("max_output_tokens", FALLBACK_MAX_OUTPUT_TOKENS), input_price=mcfg.get("input_price", 0.0), output_price=mcfg.get("output_price", 0.0), + cache_read_price=mcfg.get("cache_read_price"), currency=mcfg.get("currency", "USD"), created_at=mcfg.get("created_at", now_str), )) @@ -284,6 +287,9 @@ async def create_model(provider_id: str, body: CreateModelReq): models = raw.get("models", {}) existing_model = models.get(body.model_id) limits = await _resolve_model_limits(provider_id, body, raw) + cache_read_price = body.cache_read_price + if existing_model and "cache_read_price" not in body.model_fields_set: + cache_read_price = existing_model.get("cache_read_price") now = datetime.now(UTC).isoformat() model_config = { @@ -296,6 +302,7 @@ async def create_model(provider_id: str, body: CreateModelReq): "supports_reasoning": body.supports_reasoning, "input_price": body.input_price, "output_price": body.output_price, + "cache_read_price": cache_read_price, "currency": body.currency, "created_at": existing_model.get("created_at", now) if existing_model else now, } @@ -306,7 +313,7 @@ async def create_model(provider_id: str, body: CreateModelReq): # Add/update runtime _add_model_to_runtime( provider_id, - body, + body.model_copy(update={"cache_read_price": cache_read_price}), context_window=limits.context_window, max_output_tokens=limits.max_output_tokens, ) @@ -324,6 +331,7 @@ async def create_model(provider_id: str, body: CreateModelReq): context_window=limits.context_window, max_output_tokens=limits.max_output_tokens, input_price=body.input_price, output_price=body.output_price, + cache_read_price=cache_read_price, currency=body.currency, created_at=now, ) @@ -617,12 +625,18 @@ def _add_model_to_runtime( / OpenAICompatibleProvider (_config_models). """ _pricing = None - if body.input_price is not None or body.output_price is not None: + if ( + body.input_price is not None + or body.output_price is not None + or body.cache_read_price is not None + ): _pricing = { "input": float(body.input_price or 0.0), "output": float(body.output_price or 0.0), "currency": body.currency, } + if body.cache_read_price is not None: + _pricing["cache_read"] = float(body.cache_read_price) mi = ModelInfo( id=body.model_id, name=body.name, @@ -648,7 +662,8 @@ def _add_model_to_runtime( mi._explicit_keys = { "name", "context_window", "max_output_tokens", "supports_streaming", "supports_tools", "supports_vision", - "supports_reasoning", "input_price", "output_price", "currency", + "supports_reasoning", "input_price", "output_price", + "cache_read_price", "currency", } Provider._models[body.model_id] = mi p = Provider.get(provider_id) @@ -708,13 +723,20 @@ async def load_custom_providers_on_startup(): continue _input_price = mcfg.get("input_price") _output_price = mcfg.get("output_price") + _cache_read_price = mcfg.get("cache_read_price") _pricing = None - if _input_price is not None or _output_price is not None: + if ( + _input_price is not None + or _output_price is not None + or _cache_read_price is not None + ): _pricing = { "input": float(_input_price or 0.0), "output": float(_output_price or 0.0), "currency": mcfg.get("currency", "USD"), } + if _cache_read_price is not None: + _pricing["cache_read"] = float(_cache_read_price) mi = ModelInfo( id=model_id, name=mcfg.get("name", model_id), diff --git a/flocks/server/routes/onboarding.py b/flocks/server/routes/onboarding.py index 7ccb2a380..b90b18760 100644 --- a/flocks/server/routes/onboarding.py +++ b/flocks/server/routes/onboarding.py @@ -141,7 +141,7 @@ def _llm_provider_has_usable_credentials(provider_id: str) -> bool: "cn": { "activation_url": "https://x.threatbook.com/flocks/activate", "threatbook_llm_provider_id": "threatbook-cn-llm", - "threatbook_default_model_id": "kimi-k2.7-code", + "threatbook_default_model_id": "deepseek-v4-flash-0731", "threatbook_api_service_id": "threatbook-cn", "threatbook_mcp_name": "threatbook_mcp", "threatbook_mcp_url": "https://mcp.threatbook.cn/mcp?apikey={api_key}", @@ -151,7 +151,7 @@ def _llm_provider_has_usable_credentials(provider_id: str) -> bool: "global": { "activation_url": "https://threatbook.io/flocks/activate", "threatbook_llm_provider_id": "threatbook-io-llm", - "threatbook_default_model_id": "kimi-k2.7-code", + "threatbook_default_model_id": "deepseek-v4-flash-0731", "threatbook_api_service_id": "threatbook-io", "threatbook_mcp_name": None, "threatbook_mcp_url": None, diff --git a/flocks/server/routes/provider.py b/flocks/server/routes/provider.py index c35f13b49..4f28570cf 100644 --- a/flocks/server/routes/provider.py +++ b/flocks/server/routes/provider.py @@ -676,11 +676,14 @@ async def get_provider_catalog(): }, "limits": { "context_window": m.limits.context_window, + "max_input_tokens": m.limits.max_input_tokens, "max_output_tokens": m.limits.max_output_tokens, } if m.limits else None, "pricing": { "input": m.pricing.input, "output": m.pricing.output, + "cache_read": m.pricing.cache_read, + "cache_write": m.pricing.cache_write, "currency": m.pricing.currency, } if m.pricing else None, } @@ -2548,14 +2551,14 @@ async def test_provider_credentials( # Rank tools by required-parameter count (fewer = simpler); # prefer lightweight query/scan tools and avoid file/upload handlers. # - # NOTE: For services that expose a dedicated login/auth probe - # (e.g. qingteng_login, skyeye_login), we want it tried *first* — + # NOTE: For services that expose a dedicated login/auth probe, we + # want it tried *first* — # it is parameter-free and exercises the credential pipeline end # to end without invoking business APIs that may need extra # required fields beyond the JSON schema (e.g. assets.refresh # which the handler validates needs `resource`/`os_type`). # Match either the bare keyword (e.g. "login") or the conventional - # `_` suffix (e.g. "qingteng_login"). A loose + # `_` suffix (e.g. "service_login"). A loose # substring match would over-trigger on business tools whose names # merely contain the word — for example tdp_login_api_list and # tdp_login_weakpwd_list are query endpoints, not probes. diff --git a/flocks/server/routes/session.py b/flocks/server/routes/session.py index cbc6c3bdd..a7c920c04 100644 --- a/flocks/server/routes/session.py +++ b/flocks/server/routes/session.py @@ -10,6 +10,7 @@ import json import os import time +from dataclasses import dataclass from pathlib import Path from typing import List, Optional, Any, Dict, Literal, Union, Tuple from fastapi import APIRouter, HTTPException, status, Query, Request @@ -17,6 +18,7 @@ from pydantic import BaseModel, Field, ConfigDict from flocks.auth.context import AuthUser, get_current_auth_user, reset_current_auth_user, set_current_auth_user +from flocks.channel.base import ChatType from flocks.server.routes._timing import log_route_timing from flocks.audit import emit_audit_event from flocks.license import assert_license_active @@ -50,6 +52,7 @@ # Default agent name constant DEFAULT_AGENT = "rex" DEFAULT_MESSAGE_PAGE_LIMIT = 50 +_CHANNEL_BINDING_QUERY_BATCH_SIZE = 500 _CONTEXT_USAGE_CACHE_TTL_SECONDS = 5.0 _context_usage_cache: Dict[Tuple[str, int], Tuple[float, ContextUsageSnapshot]] = {} _context_usage_inflight: Dict[Tuple[str, int], asyncio.Task[ContextUsageSnapshot]] = {} @@ -280,6 +283,8 @@ class SessionListItem(BaseModel): title: str time: SessionTime category: str = "user" + channelID: Optional[str] = None + channelChatType: Optional[ChatType] = None status: str = "active" parentID: Optional[str] = None provider: Optional[str] = None @@ -293,6 +298,13 @@ class SessionListItem(BaseModel): isShared: bool = False +@dataclass(frozen=True) +class _ChannelSessionBinding: + channel_id: str + chat_id: str + chat_type: ChatType + + def _session_to_response( session: SessionModel, effective_project_id: Optional[str] = None, @@ -350,6 +362,8 @@ def _session_to_list_item( effective_project_id: Optional[str] = None, shared_project_ids: Optional[set[str]] = None, project_name: Optional[str] = None, + channel_binding: Optional[_ChannelSessionBinding] = None, + title_override: Optional[str] = None, ) -> SessionListItem: """Convert a session to the lightweight manager-list response shape.""" current_user = get_current_auth_user() @@ -366,7 +380,7 @@ def _session_to_list_item( projectName=project_name, effectiveProjectID=effective_project_id, directory=session.directory, - title=session.title, + title=title_override if title_override is not None else session.title, time=SessionTime( created=session.time.created, updated=session.time.updated, @@ -374,6 +388,8 @@ def _session_to_list_item( archived=session.time.archived, ), category=session.category, + channelID=channel_binding.channel_id if channel_binding else None, + channelChatType=channel_binding.chat_type if channel_binding else None, status=session.status, parentID=session.parent_id, provider=session.provider, @@ -388,6 +404,119 @@ def _session_to_list_item( ) +async def _latest_channel_bindings( + session_ids: List[str], +) -> Dict[str, _ChannelSessionBinding]: + """Return the most recently active channel metadata for each session.""" + if not session_ids: + return {} + try: + from flocks.channel.inbound.session_binding import SessionBindingService + + service = SessionBindingService() + unique_session_ids = list(dict.fromkeys(session_ids)) + channel_bindings: Dict[str, _ChannelSessionBinding] = {} + for start in range(0, len(unique_session_ids), _CHANNEL_BINDING_QUERY_BATCH_SIZE): + batch = unique_session_ids[start:start + _CHANNEL_BINDING_QUERY_BATCH_SIZE] + for binding in await service.list_bindings(session_ids=batch): + if binding.session_id in channel_bindings: + continue + try: + chat_type = ChatType(binding.chat_type) + except (TypeError, ValueError): + log.warn("session.list.unknown_channel_chat_type", { + "session_id": binding.session_id, + "chat_type": str(binding.chat_type), + }) + continue + channel_bindings[binding.session_id] = _ChannelSessionBinding( + channel_id=binding.channel_id, + chat_id=binding.chat_id, + chat_type=chat_type, + ) + return channel_bindings + except Exception as exc: + log.warn("session.list.channel_bindings_error", {"error": str(exc)}) + return {} + + +def _has_legacy_channel_title( + session: SessionModel, + binding: _ChannelSessionBinding, +) -> bool: + """Return whether a stored title matches a legacy channel fallback.""" + prefix = f"[{binding.channel_id.capitalize()}]" + if binding.chat_type == ChatType.DIRECT: + direct_prefix = f"{prefix} DM — " + title = session.title.strip() + return ( + title.casefold().startswith(direct_prefix.casefold()) + and bool(title[len(direct_prefix):].strip()) + ) + else: + expected = f"{prefix} {binding.chat_id}" + return session.title.strip().casefold() == expected.casefold() + + +async def _first_user_channel_title(session_id: str) -> Optional[str]: + """Read the first valid user-authored text stored for a channel session.""" + from flocks.channel.inbound.session_binding import extract_channel_title_text + from flocks.session.message import Message + + messages = await Message.list(session_id, include_archived=True) + for message in messages: + if message.role != "user": + continue + with_parts = await Message.get_with_parts_lazy(session_id, message.id) + if not with_parts: + continue + raw_text = "\n".join( + part.text + for part in with_parts.parts + if part.type == "text" + and not getattr(part, "synthetic", False) + and getattr(part, "text", "") + ) + title_text = extract_channel_title_text(raw_text) + if title_text: + return title_text + return None + + +async def _legacy_channel_title_overrides( + sessions: List[SessionModel], + channel_bindings: Dict[str, _ChannelSessionBinding], +) -> Dict[str, str]: + """Derive display titles for legacy fallback names without mutating storage.""" + from flocks.channel.inbound.session_binding import format_channel_title + + candidates = [] + for session in sessions: + binding = channel_bindings.get(session.id) + if binding and _has_legacy_channel_title(session, binding): + candidates.append(session) + + async def resolve(session: SessionModel) -> tuple[str, Optional[str]]: + try: + title_text = await _first_user_channel_title(session.id) + binding = channel_bindings[session.id] + title = ( + format_channel_title(binding.channel_id, title_text) + if title_text + else None + ) + return session.id, title + except Exception as exc: + log.warn("session.list.channel_title_override_error", { + "session_id": session.id, + "error": str(exc), + }) + return session.id, None + + resolved = await asyncio.gather(*(resolve(session) for session in candidates)) + return {session_id: title for session_id, title in resolved if title is not None} + + async def _session_to_response_with_goal( session: SessionModel, effective_project_id: Optional[str] = None, @@ -661,13 +790,12 @@ async def list_sessions( visible_project_ids = Project.visible_project_ids(current_user.id) shared_project_ids = Project.shared_project_ids() - filtered = [] - effective_project_ids: Dict[str, str] = {} - term = search.lower() if search else None + eligible = [] + eligible_project_ids: Dict[str, str] = {} + term = search.casefold() if search else None manager_categories = {"user", "workflow", "entity-config"} project_names = Project.registered_project_names() if view == "list" else {} - skip_remaining = offset or 0 - + for session in all_sessions: if session.status == "archived": if current_user.role != "admin" and not SessionPolicy.is_owner(session, current_user): @@ -700,8 +828,6 @@ async def list_sessions( continue if start is not None and session.time.updated < start: continue - if term is not None and term not in session.title.lower(): - continue if manager: if session.category not in manager_categories: continue @@ -712,23 +838,51 @@ async def list_sessions( # exclude test sessions from the default listing continue + eligible.append(session) + eligible_project_ids[session.id] = effective_project_id + + channel_bindings: Dict[str, _ChannelSessionBinding] = {} + channel_title_overrides: Dict[str, str] = {} + if view == "list" and term is not None: + # Search the same derived title that the lightweight list displays. + channel_bindings = await _latest_channel_bindings([session.id for session in eligible]) + channel_title_overrides = await _legacy_channel_title_overrides( + eligible, + channel_bindings, + ) + + filtered = [] + effective_project_ids: Dict[str, str] = {} + skip_remaining = offset or 0 + for session in eligible: + searchable_title = channel_title_overrides.get(session.id, session.title) + if term is not None and term not in searchable_title.casefold(): + continue if skip_remaining > 0: skip_remaining -= 1 continue - + filtered.append(session) - effective_project_ids[session.id] = effective_project_id - + effective_project_ids[session.id] = eligible_project_ids[session.id] + if limit is not None and len(filtered) >= limit: break if view == "list": + if term is None: + channel_bindings = await _latest_channel_bindings([session.id for session in filtered]) + channel_title_overrides = await _legacy_channel_title_overrides( + filtered, + channel_bindings, + ) response = [ _session_to_list_item( s, - effective_project_ids[s.id], - shared_project_ids, - project_names.get(s.project_id), + effective_project_id=effective_project_ids[s.id], + shared_project_ids=shared_project_ids, + project_name=project_names.get(s.project_id), + channel_binding=channel_bindings.get(s.id), + title_override=channel_title_overrides.get(s.id), ) for s in filtered ] @@ -2048,6 +2202,7 @@ class MessagePartInfo(BaseModel): sessionID: str type: str text: Optional[str] = None + time: Optional[Dict[str, Any]] = None synthetic: Optional[bool] = None tool: Optional[str] = None state: Optional[Dict[str, Any]] = None @@ -2093,6 +2248,14 @@ def _part_to_response_info( ) -> MessagePartInfo: text_value = getattr(part, "text", None) if part.type in ("text", "reasoning", "thinking") else None + raw_time = getattr(part, "time", None) + if hasattr(raw_time, "model_dump"): + time_value = raw_time.model_dump() + elif isinstance(raw_time, dict): + time_value = raw_time + else: + time_value = None + url_value = getattr(part, "url", None) if part.type == "file" else None state_value = None @@ -2110,6 +2273,7 @@ def _part_to_response_info( sessionID=session_id, type=part.type, text=text_value, + time=time_value, synthetic=getattr(part, "synthetic", None), tool=getattr(part, "tool", None) if part.type == "tool" else None, state=state_value, @@ -4181,6 +4345,7 @@ async def _create_user_message( model_info: Optional[Dict[str, str]] = None, *, agent_override: Optional[str] = None, + ignored: Optional[bool] = None, ) -> str: now_ms = int(_time.time() * 1000) user_msg_id = event.message_id or Identifier.create("message") @@ -4197,6 +4362,7 @@ async def _create_user_message( agent=message_agent, executionMode=event.execution_mode, **({"model": model_info} if model_info else {}), + ignored=ignored, part_id=user_part_id, ), expected_generation=lifecycle_generation, @@ -4219,6 +4385,7 @@ async def _create_user_message( "sessionID": sessionID, "type": "text", "text": user_text, + **({"ignored": ignored} if ignored is not None else {}), "time": {"start": now_ms}, } }) @@ -4226,7 +4393,7 @@ async def _create_user_message( async def _publish_direct_response(output_event, text: str) -> None: user_text = output_event.user_visible_text - parent_msg_id = await _create_user_message(user_text) + parent_msg_id = await _create_user_message(user_text, ignored=True) asst_now = int(_time.time() * 1000) asst_msg_id = Identifier.ascending("message") asst_part_id = Identifier.ascending("part") @@ -4243,6 +4410,7 @@ async def _publish_direct_response(output_event, text: str) -> None: providerID="builtin", agent=agent_name, finish="stop", + ignored=True, part_id=asst_part_id, ), expected_generation=lifecycle_generation, @@ -4274,6 +4442,7 @@ async def _publish_direct_response(output_event, text: str) -> None: "sessionID": sessionID, "type": "text", "text": text, + "ignored": True, "time": {"start": asst_now, "end": asst_now}, } }) @@ -4281,8 +4450,6 @@ async def _publish_direct_response(output_event, text: str) -> None: publish_event, sessionID, session=session, - provider_id="builtin", - model_id="command", ) async def _run_llm(output_event, prompt_text: str, display_text: Optional[str] = None) -> None: diff --git a/flocks/server/routes/workflow.py b/flocks/server/routes/workflow.py index 0ac9d0429..dce2cb020 100644 --- a/flocks/server/routes/workflow.py +++ b/flocks/server/routes/workflow.py @@ -1136,6 +1136,14 @@ async def _run_workflow_execution_task( "currentStepIndex": 0, "stepCount": 0, } + if req.session_id or req.message_id: + execution_summary.update( + { + "sessionId": req.session_id, + "messageId": req.message_id, + "agent": req.agent, + } + ) def _write_progress(update_fields: Dict[str, Any]) -> None: try: @@ -1737,6 +1745,16 @@ async def run_workflow_endpoint(workflow_id: str, req: WorkflowRunRequest): workflow_id, input_params=req.inputs or {}, ) + if req.session_id or req.message_id: + exec_data.update( + { + "sessionId": req.session_id, + "messageId": req.message_id, + "agent": req.agent, + "updatedAt": int(time.time() * 1000), + } + ) + await WorkflowStore.upsert_execution(compact_execution_summary(exec_data)) exec_id = str(exec_data["id"]) cancel_event = threading.Event() diff --git a/flocks/session/context_usage.py b/flocks/session/context_usage.py index 35084ee92..ca35e09d1 100644 --- a/flocks/session/context_usage.py +++ b/flocks/session/context_usage.py @@ -434,6 +434,8 @@ async def _estimate_message_breakdown(session_id: str, messages: List[Any]) -> t for part in parts: part_type = _field_value(part, "type", "") if part_type == "text": + if bool(_field_value(part, "ignored", False)): + continue tokens_by_key["conversation"] += SessionPrompt.count_tokens(_field_value(part, "text", "") or "") continue if part_type in {"reasoning", "thinking"}: @@ -540,6 +542,8 @@ def _resolve_message_model(messages: List[Any]) -> tuple[Optional[str], Optional if role == "assistant": provider_id = getattr(message, "providerID", None) model_id = getattr(message, "modelID", None) + if provider_id == "builtin" and model_id == "command": + continue if provider_id and model_id: return provider_id, model_id if role == "user": diff --git a/flocks/session/execution_mode.py b/flocks/session/execution_mode.py index fbad85d6f..13d12b154 100644 --- a/flocks/session/execution_mode.py +++ b/flocks/session/execution_mode.py @@ -111,7 +111,6 @@ def tool_call_denial_reason( subagent_type = str(arguments.get("subagent_type") or "").strip().lower() if ( subagent_type in PLAN_DELEGATABLE_AGENT_NAMES - and not arguments.get("category") and not arguments.get("session_id") ): return None diff --git a/flocks/session/message.py b/flocks/session/message.py index 0ca768a1a..2f3a8f4ed 100644 --- a/flocks/session/message.py +++ b/flocks/session/message.py @@ -1403,6 +1403,7 @@ async def create( # Pop TextPart-specific fields before message creation — they # belong on TextPart, not on UserMessageInfo/AssistantMessageInfo. _synthetic = kwargs.pop("synthetic", None) + _ignored = kwargs.pop("ignored", None) _part_metadata = kwargs.pop("part_metadata", None) # Create appropriate message type based on role @@ -1454,6 +1455,8 @@ async def create( _part_extras = {} if _synthetic is not None: _part_extras["synthetic"] = _synthetic + if _ignored is not None: + _part_extras["ignored"] = _ignored if _part_metadata is not None: _part_extras["metadata"] = _part_metadata part = TextPart( diff --git a/flocks/session/prompt.py b/flocks/session/prompt.py index 21fb0ebb3..1c0086e47 100644 --- a/flocks/session/prompt.py +++ b/flocks/session/prompt.py @@ -585,6 +585,8 @@ async def _tokens_for_message(cls, session_id: str, msg: Any) -> int: parts = await Message.parts(msg_id or "", session_id) for part in parts: if part.type == "text": + if getattr(part, "ignored", False): + continue total += cls.count_tokens(getattr(part, 'text', '')) elif part.type == "tool": state = getattr(part, 'state', None) diff --git a/flocks/session/runner.py b/flocks/session/runner.py index ab388b073..44dfc52d9 100644 --- a/flocks/session/runner.py +++ b/flocks/session/runner.py @@ -2852,6 +2852,8 @@ async def _to_chat_messages( # Text parts if part.type == "text" and hasattr(part, 'text'): + if getattr(part, "ignored", False): + continue assistant_content_parts.append(part.text) elif part.type == "reasoning" and hasattr(part, 'text'): assistant_reasoning_parts.append(part.text) diff --git a/flocks/tool/agent/delegate_task.py b/flocks/tool/agent/delegate_task.py index cf00616be..34280507a 100644 --- a/flocks/tool/agent/delegate_task.py +++ b/flocks/tool/agent/delegate_task.py @@ -1,6 +1,4 @@ -""" -delegate_task tool - category or subagent-based delegation (Oh-My-Flocks parity). -""" +"""delegate_task tool for direct subagent delegation.""" from __future__ import annotations @@ -16,18 +14,12 @@ ToolResult, ToolContext, ) -from flocks.tool.delegate_task_constants import ( - DEFAULT_CATEGORIES, - CATEGORY_PROMPT_APPENDS, - CATEGORY_DESCRIPTIONS, -) from flocks.session.session import Session from flocks.session.message import Message, MessageRole from flocks.session.session_loop import SessionLoop # 使用轻量级元数据查询，避免循环依赖 from flocks.agent.registry import is_delegatable from flocks.skill.skill import Skill -from flocks.config.config import Config from flocks.tool.subagent_result import ( _extract_message_error, format_sync_subagent_result, @@ -225,62 +217,6 @@ def _parse_model(model: Optional[str]) -> Optional[Dict[str, str]]: return {"modelID": model} -def _validate_category_model(category_model: Optional[Dict[str, str]], category: Optional[str]) -> Optional[Dict[str, str]]: - """Validate that the category model's provider is available and has the model registered. - - Returns the original model dict when valid, or None to signal the caller - should fall back to the parent session's model (via _resolve_model priority chain). - """ - if not category_model: - return None - - provider_id = category_model.get("providerID") - model_id = category_model.get("modelID") - if not provider_id or not model_id: - return category_model - - try: - from flocks.provider.provider import Provider - provider = Provider.get(provider_id) - if not provider: - log.warn("delegate_task.category_model_fallback", { - "category": category, - "provider": provider_id, - "model": model_id, - "reason": "provider not registered", - }) - return None - - if not provider.is_configured(): - log.warn("delegate_task.category_model_fallback", { - "category": category, - "provider": provider_id, - "model": model_id, - "reason": "provider not configured", - }) - return None - - registered_ids = {m.id for m in provider.get_models()} - if model_id not in registered_ids: - log.warn("delegate_task.category_model_fallback", { - "category": category, - "provider": provider_id, - "model": model_id, - "reason": "model not found in provider", - "available_models": list(registered_ids)[:10], - }) - return None - - except Exception as exc: - log.warn("delegate_task.category_model_validate_error", { - "category": category, - "error": str(exc), - }) - return None - - return category_model - - async def _find_completed_delegate( session_id: str, current_message_id: str, @@ -304,7 +240,7 @@ async def _find_completed_delegate( if getattr(state, "status", None) != "completed": continue inp = getattr(state, "input", {}) - prev_key = inp.get("subagent_type") or inp.get("category") + prev_key = inp.get("subagent_type") if prev_key == agent_key and inp.get("description") == description: output = getattr(state, "output", "") if isinstance(output, dict): @@ -354,7 +290,6 @@ def _derive_task_description( description: Optional[str], prompt: str, subagent_type: Optional[str] = None, - category: Optional[str] = None, session_id: Optional[str] = None, ) -> str: normalized = " ".join((description or "").split()) @@ -367,8 +302,6 @@ def _derive_task_description( if subagent_type: return f"delegate to {subagent_type}" - if category: - return f"delegate {category} task" if session_id: return f"continue task {session_id}" return "delegate task" @@ -378,12 +311,16 @@ def _derive_task_description( # Tool definition # ------------------------------------------------------------------ -DESCRIPTION = """Spawn agent task with category-based or direct agent selection. " +DESCRIPTION = """Spawn a task using a directly selected subagent. Use this tool when: -- The task requires multiple steps or research +- A specialized agent clearly matches the task - You need to explore code in parallel -- The task can be delegated to a specialized agent +- Independent work can run in parallel +- Isolating research or noisy intermediate work improves context quality + +Do not delegate trivial edits, direct one-tool operations, or tightly coupled +work that requires continuous coordination in the current context. Usage notes: - Provide a clear description (3-5 words) @@ -399,7 +336,7 @@ def _derive_task_description( REQUIRED: prompt. LOAD_SKILLS is optional and defaults to []. DESCRIPTION is optional and will be auto-derived when omitted. -USE EITHER subagent_type OR category — NEVER both simultaneously. +SUBAGENT_TYPE is required for new tasks. Omit it only when session_id continues an existing task. """ @ToolRegistry.register_function( @@ -426,16 +363,10 @@ def _derive_task_description( description="Full detailed prompt for the subagent.", required=True, ), - ToolParameter( - name="category", - type=ParameterType.STRING, - description="Category name. Mutually exclusive with subagent_type — use ONE or the other, never both.", - required=False, - ), ToolParameter( name="subagent_type", type=ParameterType.STRING, - description="Agent name. Mutually exclusive with category — use ONE or the other, never both. Must be a delegatable agent", + description="Delegatable agent name. Required for new tasks; omit when continuing with session_id.", required=False, ), ToolParameter( @@ -447,7 +378,7 @@ def _derive_task_description( ToolParameter( name="command", type=ParameterType.STRING, - description="Optional command name for tracking", + description="Deprecated command name retained for caller compatibility", required=False, ), ToolParameter( @@ -468,7 +399,6 @@ async def delegate_task_tool( # in-process call paths (e.g. `task.py` alias) may still pass it through. # This guard is the second line of defense. run_in_background: bool = False, - category: Optional[str] = None, subagent_type: Optional[str] = None, session_id: Optional[str] = None, command: Optional[str] = None, @@ -488,23 +418,21 @@ async def delegate_task_tool( return ToolResult(success=False, error="prompt is required") load_skills = [str(name).strip() for name in (load_skills or []) if str(name).strip()] - description = _derive_task_description(description, prompt, subagent_type, category, session_id) - if category and subagent_type: - return ToolResult(success=False, error="Provide EITHER category OR subagent_type, not both.") - if not category and not subagent_type and not session_id: - return ToolResult(success=False, error="Must provide either category or subagent_type.") + description = _derive_task_description(description, prompt, subagent_type, session_id) + if not subagent_type and not session_id: + return ToolResult(success=False, error="Must provide either subagent_type or session_id.") await ctx.ask( permission="delegate_task", - patterns=[category or subagent_type or "continue"], + patterns=[subagent_type or "continue"], always=["*"], - metadata={"description": description, "category": category, "subagent_type": subagent_type}, + metadata={"description": description, "subagent_type": subagent_type}, ) # Dedup: if an identical delegate_task already completed in this session, # return the previous result to prevent the LLM from re-delegating. if not session_id: - agent_key = subagent_type or category + agent_key = subagent_type prev = await _find_completed_delegate(ctx.session_id, ctx.message_id, agent_key, description) if prev is not None: log.info("delegate_task.dedup_hit", { @@ -518,10 +446,6 @@ async def delegate_task_tool( if skill_result["error"]: return ToolResult(success=False, error=skill_result["error"]) - cfg = await Config.get() - category_configs = {**DEFAULT_CATEGORIES, **(cfg.categories or {})} - category_prompt_append = None - category_model = None explicit_model = _parse_model(model) agent_to_use: Optional[str] = None @@ -558,47 +482,19 @@ async def delegate_task_tool( metadata={"sessionId": session.id}, ) - if category: - agent_to_use = "rex-junior" - config = category_configs.get(category) - if not config: - available = ", ".join(category_configs.keys()) - return ToolResult(success=False, error=f'Unknown category "{category}". Available: {available}') - raw_model = explicit_model or _parse_model(config.get("model") if isinstance(config, dict) else getattr(config, "model", None)) - category_model = _validate_category_model(raw_model, category) - if raw_model and not category_model: - log.info("delegate_task.using_parent_model", { - "category": category, - "original_model": raw_model, - "reason": "category model unavailable, inheriting parent session model", - }) - category_prompt_append = ( - (config.get("prompt_append") if isinstance(config, dict) else getattr(config, "prompt_append", None)) - or CATEGORY_PROMPT_APPENDS.get(category) - ) - elif subagent_type: + if subagent_type: # 使用轻量级元数据查询，避免循环依赖 # 不再调用 Agent.get()，而是使用 is_delegatable() if not is_delegatable(subagent_type): - # 针对特殊 Agent 提供更友好的错误提示 - if subagent_type.lower() in ["sisyphus-junior", "rex-junior"]: - return ToolResult( - success=False, - error=f'Cannot use subagent_type="{subagent_type}" directly. Use category parameter instead.', - ) - else: - return ToolResult( - success=False, - error=f'Agent "{subagent_type}" cannot be delegated to (it may be a primary agent or restricted).', - ) + return ToolResult( + success=False, + error=f'Agent "{subagent_type}" cannot be delegated to (it may be a primary agent or restricted).', + ) agent_to_use = subagent_type - category_model = explicit_model system_parts = [] if skill_result["content"]: system_parts.append(skill_result["content"]) - if category_prompt_append: - system_parts.append(category_prompt_append) system_content = "\n\n".join(system_parts) if system_parts else "" full_prompt = f"{system_content}\n\n{prompt}" if system_content else prompt @@ -620,11 +516,11 @@ async def delegate_task_tool( permission=await _subagent_session_permissions(agent_to_use), category="task", ) - if category_model and category_model.get("providerID") and category_model.get("modelID"): + if explicit_model and explicit_model.get("providerID") and explicit_model.get("modelID"): create_kwargs.update( - provider=category_model["providerID"], - model=category_model["modelID"], - model_pinned=bool(explicit_model), + provider=explicit_model["providerID"], + model=explicit_model["modelID"], + model_pinned=True, ) created = await Session.create(**create_kwargs) if ctx.extra.get("workflow_temp_parent") is True: @@ -651,8 +547,8 @@ async def delegate_task_tool( prompt=full_prompt, description=description, resumed=False, - provider_id=(category_model or {}).get("providerID"), - model_id=(category_model or {}).get("modelID"), + provider_id=(explicit_model or {}).get("providerID"), + model_id=(explicit_model or {}).get("modelID"), callbacks=forwarder.build_callbacks( event_publish_callback=ctx.event_publish_callback, ), diff --git a/flocks/tool/agent/task.py b/flocks/tool/agent/task.py index 0ef7ad922..891231729 100644 --- a/flocks/tool/agent/task.py +++ b/flocks/tool/agent/task.py @@ -51,13 +51,7 @@ ToolParameter( name="subagent_type", type=ParameterType.STRING, - description="Delegatable agent name. Mutually exclusive with category.", - required=False, - ), - ToolParameter( - name="category", - type=ParameterType.STRING, - description="Delegate category. Mutually exclusive with subagent_type.", + description="Delegatable agent name. Required for new tasks; omit when continuing with session_id.", required=False, ), ToolParameter( @@ -76,7 +70,7 @@ ToolParameter( name="command", type=ParameterType.STRING, - description="Optional command name for tracking", + description="Deprecated command name retained for caller compatibility", required=False, ), ToolParameter( @@ -92,7 +86,6 @@ async def task_tool( description: Optional[str] = None, prompt: Optional[str] = None, subagent_type: Optional[str] = None, - category: Optional[str] = None, load_skills: Optional[list] = None, run_in_background: bool = False, session_id: Optional[str] = None, @@ -106,7 +99,6 @@ async def task_tool( load_skills=load_skills, description=description, run_in_background=run_in_background, - category=category, subagent_type=subagent_type, session_id=session_id, command=command, diff --git a/flocks/tool/delegate_task_constants.py b/flocks/tool/delegate_task_constants.py deleted file mode 100644 index 3ef6c8596..000000000 --- a/flocks/tool/delegate_task_constants.py +++ /dev/null @@ -1,226 +0,0 @@ -""" -Constants for delegate_task tool (ported from oh-my-opencode). -""" - -VISUAL_CATEGORY_PROMPT_APPEND = """ -You are working on VISUAL/UI tasks. - -Design-first mindset: -- Bold aesthetic choices over safe defaults -- Unexpected layouts, asymmetry, grid-breaking elements -- Distinctive typography (avoid: Arial, Inter, Roboto, Space Grotesk) -- Cohesive color palettes with sharp accents -- High-impact animations with staggered reveals -- Atmosphere: gradient meshes, noise textures, layered transparencies - -AVOID: Generic fonts, purple gradients on white, predictable layouts, cookie-cutter patterns. -""" - -ULTRABRAIN_CATEGORY_PROMPT_APPEND = """ -You are working on DEEP LOGICAL REASONING / COMPLEX ARCHITECTURE tasks. - -**CRITICAL - CODE STYLE REQUIREMENTS (NON-NEGOTIABLE)**: -1. BEFORE writing ANY code, SEARCH the existing codebase to find similar patterns/styles -2. Your code MUST match the project's existing conventions - blend in seamlessly -3. Write READABLE code that humans can easily understand - no clever tricks -4. If unsure about style, explore more files until you find the pattern - -Strategic advisor mindset: -- Bias toward simplicity: least complex solution that fulfills requirements -- Leverage existing code/patterns over new components -- Prioritize developer experience and maintainability -- One clear recommendation with effort estimate (Quick/Short/Medium/Large) -- Signal when advanced approach warranted - -Response format: -- Bottom line (2-3 sentences) -- Action plan (numbered steps) -- Risks and mitigations (if relevant) -""" - -ARTISTRY_CATEGORY_PROMPT_APPEND = """ -You are working on HIGHLY CREATIVE / ARTISTIC tasks. - -Artistic genius mindset: -- Push far beyond conventional boundaries -- Explore radical, unconventional directions -- Surprise and delight: unexpected twists, novel combinations -- Rich detail and vivid expression -- Break patterns deliberately when it serves the creative vision - -Approach: -- Generate diverse, bold options first -- Embrace ambiguity and wild experimentation -- Balance novelty with coherence -- This is for tasks requiring exceptional creativity -""" - -QUICK_CATEGORY_PROMPT_APPEND = """ -You are working on SMALL / QUICK tasks. - -Efficient execution mindset: -- Fast, focused, minimal overhead -- Get to the point immediately -- No over-engineering -- Simple solutions for simple problems - -Approach: -- Minimal viable implementation -- Skip unnecessary abstractions -- Direct and concise - - - -THIS CATEGORY USES A LESS CAPABLE MODEL (claude-haiku-4-5). - -The model executing this task has LIMITED reasoning capacity. Your prompt MUST be: - -**EXHAUSTIVELY EXPLICIT** - Leave NOTHING to interpretation: -1. MUST DO: List every required action as atomic, numbered steps -2. MUST NOT DO: Explicitly forbid likely mistakes and deviations -3. EXPECTED OUTPUT: Describe exact success criteria with concrete examples - -**WHY THIS MATTERS:** -- Less capable models WILL deviate without explicit guardrails -- Vague instructions -> unpredictable results -- Implicit expectations -> missed requirements - -**PROMPT STRUCTURE (MANDATORY):** -``` -TASK: [One-sentence goal] - -MUST DO: -1. [Specific action with exact details] -2. [Another specific action] -... - -MUST NOT DO: -- [Forbidden action + why] -- [Another forbidden action] -... - -EXPECTED OUTPUT: -- [Exact deliverable description] -- [Success criteria / verification method] -``` - -If your prompt lacks this structure, REWRITE IT before delegating. -""" - -UNSPECIFIED_LOW_CATEGORY_PROMPT_APPEND = """ -You are working on tasks that don't fit specific categories but require moderate effort. - - -BEFORE selecting this category, VERIFY ALL conditions: -1. Task does NOT fit: quick (trivial), visual-engineering (UI), ultrabrain (deep logic), artistry (creative), writing (docs) -2. Task requires more than trivial effort but is NOT system-wide -3. Scope is contained within a few files/modules - -If task fits ANY other category, DO NOT select unspecified-low. -This is NOT a default choice - it's for genuinely unclassifiable moderate-effort work. - - - - -THIS CATEGORY USES A MID-TIER MODEL (claude-sonnet-4-5). - -**PROVIDE CLEAR STRUCTURE:** -1. MUST DO: Enumerate required actions explicitly -2. MUST NOT DO: State forbidden actions to prevent scope creep -3. EXPECTED OUTPUT: Define concrete success criteria -""" - -UNSPECIFIED_HIGH_CATEGORY_PROMPT_APPEND = """ -You are working on tasks that don't fit specific categories but require substantial effort. - - -BEFORE selecting this category, VERIFY ALL conditions: -1. Task does NOT fit: quick (trivial), visual-engineering (UI), ultrabrain (deep logic), artistry (creative), writing (docs) -2. Task requires substantial effort across multiple systems/modules -3. Changes have broad impact or require careful coordination -4. NOT just "complex" - must be genuinely unclassifiable AND high-effort - -If task fits ANY other category, DO NOT select unspecified-high. -If task is unclassifiable but moderate-effort, use unspecified-low instead. - -""" - -WRITING_CATEGORY_PROMPT_APPEND = """ -You are working on WRITING / PROSE tasks. - -Wordsmith mindset: -- Clear, flowing prose -- Appropriate tone and voice -- Engaging and readable -- Proper structure and organization - -Approach: -- Understand the audience -- Draft with care -- Polish for clarity and impact -- Documentation, READMEs, articles, technical writing -""" - -DEEP_CATEGORY_PROMPT_APPEND = """ -You are working on GOAL-ORIENTED AUTONOMOUS tasks. - -**CRITICAL - AUTONOMOUS EXECUTION MINDSET (NON-NEGOTIABLE)**: -You are NOT an interactive assistant. You are an autonomous problem-solver. - -**BEFORE making ANY changes**: -1. SILENTLY explore the codebase extensively (5-15 minutes of reading is normal) -2. Read related files, trace dependencies, understand the full context -3. Build a complete mental model of the problem space -4. DO NOT ask clarifying questions - the goal is already defined - -**Autonomous executor mindset**: -- You receive a GOAL, not step-by-step instructions -- Figure out HOW to achieve the goal yourself -- Thorough research before any action -- Fix hairy problems that require deep understanding -- Work independently without frequent check-ins - -**Approach**: -- Explore extensively, understand deeply, then act decisively -- Prefer comprehensive solutions over quick patches -- If the goal is unclear, make reasonable assumptions and proceed -- Document your reasoning in code comments only when non-obvious - -**Response format**: -- Minimal status updates (user trusts your autonomy) -- Focus on results, not play-by-play progress -- Report completion with summary of changes made -""" - -DEFAULT_CATEGORIES = { - "visual-engineering": {"model": "google/gemini-3-pro"}, - "ultrabrain": {"model": "openai/gpt-5.2-codex", "variant": "xhigh"}, - "deep": {"model": "openai/gpt-5.2-codex", "variant": "medium"}, - "artistry": {"model": "google/gemini-3-pro", "variant": "max"}, - "quick": {"model": "anthropic/claude-haiku-4-5"}, - "unspecified-low": {"model": "anthropic/claude-sonnet-4-6"}, - "unspecified-high": {"model": "anthropic/claude-opus-4-6", "variant": "max"}, - "writing": {"model": "google/gemini-3-flash"}, -} - -CATEGORY_PROMPT_APPENDS = { - "visual-engineering": VISUAL_CATEGORY_PROMPT_APPEND, - "ultrabrain": ULTRABRAIN_CATEGORY_PROMPT_APPEND, - "deep": DEEP_CATEGORY_PROMPT_APPEND, - "artistry": ARTISTRY_CATEGORY_PROMPT_APPEND, - "quick": QUICK_CATEGORY_PROMPT_APPEND, - "unspecified-low": UNSPECIFIED_LOW_CATEGORY_PROMPT_APPEND, - "unspecified-high": UNSPECIFIED_HIGH_CATEGORY_PROMPT_APPEND, - "writing": WRITING_CATEGORY_PROMPT_APPEND, -} - -CATEGORY_DESCRIPTIONS = { - "visual-engineering": "Frontend, UI/UX, design, styling, animation", - "ultrabrain": "Use ONLY for genuinely hard, logic-heavy tasks. Give clear goals only, not step-by-step instructions.", - "deep": "Goal-oriented autonomous problem-solving. Thorough research before action. For hairy problems requiring deep understanding.", - "artistry": "Complex problem-solving with unconventional, creative approaches - beyond standard patterns", - "quick": "Trivial tasks - single file changes, typo fixes, simple modifications", - "unspecified-low": "Tasks that don't fit other categories, low effort required", - "unspecified-high": "Tasks that don't fit other categories, high effort required", - "writing": "Documentation, prose, technical writing", -} diff --git a/flocks/updater/updater.py b/flocks/updater/updater.py index bfc418415..9dfc5a711 100644 --- a/flocks/updater/updater.py +++ b/flocks/updater/updater.py @@ -1450,6 +1450,42 @@ async def _download_with_fallback( raise RuntimeError(summary) +def _backup_path_is_excluded(parts: tuple[str, ...]) -> bool: + """Return whether a source-relative path should be excluded from backups.""" + if parts and parts[0] in _ROOT_RUNTIME_NAMES: + return True + return any(part in _PRESERVE_NAMES or part == "dist" for part in parts) + + +def _write_backup_archive(source_root: Path, archive_path: Path) -> None: + """Write a source-only tar archive from *source_root*.""" + + def _filter(info: tarfile.TarInfo) -> tarfile.TarInfo | None: + parts = info.name.split("/") + relative_parts = tuple(parts[1:]) if parts and parts[0] == "flocks" else tuple(parts) + if _backup_path_is_excluded(relative_parts): + return None + return info + + with tarfile.open(archive_path, "w:gz") as tar: + tar.add(str(source_root), arcname="flocks", filter=_filter) + + +def _copy_backup_snapshot(install_root: Path, snapshot_root: Path) -> None: + """Copy backup-eligible source files into a stable temporary snapshot.""" + + def _ignore(directory: str, names: list[str]) -> set[str]: + relative_dir = Path(directory).relative_to(install_root) + return {name for name in names if _backup_path_is_excluded((*relative_dir.parts, name))} + + shutil.copytree( + install_root, + snapshot_root, + ignore=_ignore, + ignore_dangling_symlinks=True, + ) + + def _backup_current_version( install_root: Path, current_version: str, @@ -1457,6 +1493,7 @@ def _backup_current_version( ) -> Path | None: """ Compress the current source tree into ~/.flocks/version/ . + Falls back to archiving a temporary source snapshot when direct archiving fails. Returns the backup path on success, None on failure. Preserved runtime/user directories are excluded. """ @@ -1464,25 +1501,40 @@ def _backup_current_version( ts = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S") backup_name = f"flocks-{current_version}-{ts}" backup_path = _BACKUP_DIR / f"{backup_name}.tar.gz" + partial_path = _BACKUP_DIR / f"{backup_name}.tar.gz.partial" + partial_path.unlink(missing_ok=True) - def _filter(info: tarfile.TarInfo) -> tarfile.TarInfo | None: - parts = info.name.split("/") - if len(parts) >= 2 and parts[0] == "flocks" and parts[1] in _ROOT_RUNTIME_NAMES: + try: + _write_backup_archive(install_root, partial_path) + except Exception as direct_exc: + partial_path.unlink(missing_ok=True) + log.warning("updater.backup.direct_failed", {"error": str(direct_exc)}) + snapshot_dir: Path | None = None + try: + snapshot_dir = Path(tempfile.mkdtemp(prefix="flocks-backup-")) + snapshot_root = snapshot_dir / "flocks" + _copy_backup_snapshot(install_root, snapshot_root) + _write_backup_archive(snapshot_root, partial_path) + except Exception as snapshot_exc: + partial_path.unlink(missing_ok=True) + log.warning( + "updater.backup.failed", + { + "direct_error": str(direct_exc), + "snapshot_error": str(snapshot_exc), + }, + ) return None - for part in parts: - if part in _PRESERVE_NAMES: - return None - if part == "dist": - return None - return info + finally: + if snapshot_dir is not None: + shutil.rmtree(snapshot_dir, ignore_errors=True) try: - with tarfile.open(backup_path, "w:gz") as tar: - tar.add(str(install_root), arcname="flocks", filter=_filter) + partial_path.replace(backup_path) except Exception as exc: + partial_path.unlink(missing_ok=True) log.warning("updater.backup.failed", {"error": str(exc)}) return None - _cleanup_old_backups(retain_count) return backup_path diff --git a/flocks/workflow/engine.py b/flocks/workflow/engine.py index a0dcffe27..dfcfba12c 100644 --- a/flocks/workflow/engine.py +++ b/flocks/workflow/engine.py @@ -6,6 +6,7 @@ import hashlib from itertools import islice import logging +import threading import time import traceback import uuid @@ -14,11 +15,11 @@ from .code_gen import CodeGen, SimpleCodeGen, LLMCodeGen from .edge_resolver import EdgeResolver -from .errors import MaxStepsExceededError, NodeExecutionError, RunCancelledError, RunTimeoutError +from .errors import MaxStepsExceededError, NodeExecutionError, NodeTimeoutError, RunCancelledError, RunTimeoutError from .execution_plan import WorkflowExecutionPlan from .execution_state import ExecutionResult, StepResult, WorkflowExecutionState from .models import Edge, Workflow, Node -from .repl_runtime import PythonExecRuntime, Runtime +from .repl_runtime import HostProcessPythonExecRuntime, PythonExecRuntime, Runtime _logger = logging.getLogger("flocks.workflow.engine") @@ -53,10 +54,7 @@ def _summarize_for_observability(value: Any, *, depth: int = 0) -> Any: return {"_type": "string", "chars": len(value), "preview": value[:200]} return value if isinstance(value, dict): - return { - key: _summarize_for_observability(item, depth=depth + 1) - for key, item in islice(value.items(), 50) - } + return {key: _summarize_for_observability(item, depth=depth + 1) for key, item in islice(value.items(), 50)} if isinstance(value, (list, tuple, set)): if isinstance(value, (list, tuple)): preview_items = value[:3] @@ -87,6 +85,23 @@ def _outputs_for_log(outputs: Dict[str, Any], *, max_chars: int = 4000) -> str: return text[:max_chars] + f"...[truncated:{len(text) - max_chars}]" +def _input_hash_for_dedup(node_id: str, inputs: Dict[str, Any]) -> str: + """Hash canonical inputs without retaining one full serialized copy.""" + try: + digest = hashlib.sha256() + encoder = json.JSONEncoder( + sort_keys=True, + ensure_ascii=False, + default=str, + separators=(",", ":"), + ) + for chunk in encoder.iterencode({"n": node_id, "i": inputs}): + digest.update(chunk.encode()) + return digest.hexdigest()[:16] + except Exception: + return "" + + def _default_workflow_loader(workflow_id: str) -> "Workflow": """Default loader: resolves workflow by ID from disk, then legacy KV.""" import asyncio @@ -179,6 +194,7 @@ def _get_isolated_runtime(self) -> "Runtime": tool_registry=self.runtime.tool_registry, cancel_checker=self.runtime.cancel_checker, cleanup_globals_after_execute=self.runtime.cleanup_globals_after_execute, + enable_cancel_trace=self.runtime.enable_cancel_trace, ) return self.runtime @@ -217,10 +233,24 @@ def run( timeout_executor: Optional[ThreadPoolExecutor] = None if step_timeout_s is not None: timeout_executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="wf-node") + timeout_cancel_event = threading.Event() previous_cancel_checker = None + previous_enable_cancel_trace = None if isinstance(self.runtime, PythonExecRuntime): previous_cancel_checker = self.runtime.cancel_checker - self.runtime.cancel_checker = cancel + previous_enable_cancel_trace = self.runtime.enable_cancel_trace + + def _runtime_cancel_requested() -> bool: + return timeout_cancel_event.is_set() or bool(cancel and cancel()) + + self.runtime.cancel_checker = ( + _runtime_cancel_requested if step_timeout_s is not None or cancel is not None else None + ) + # A timeout-only checker is consumed by isolated-process RPC calls. + # Enabling sys.settrace for every ordinary workflow line would add + # a large steady-state cost; host-thread timeouts keep their legacy + # nonfatal, best-effort semantics. + self.runtime.enable_cancel_trace = cancel is not None try: def _raise_cancelled() -> None: @@ -295,24 +325,11 @@ def merge_payload(src: str, payload: Dict[str, Any]) -> None: inputs = merged state.join_seen_sources.pop(node_id, None) - # Dedup: skip if same node already ran with identical inputs. - # Lightweight history mode is used by high-throughput ingest - # paths with large payloads; hashing full inputs there would - # serialize the same large alert lists we are trying not to - # retain. - if self.history_mode == "summary": - _input_hash = "" - else: - try: - _hash_raw = json.dumps( - {"n": node_id, "i": inputs}, - sort_keys=True, - ensure_ascii=False, - default=str, - ) - _input_hash = hashlib.sha256(_hash_raw.encode()).hexdigest()[:16] - except Exception: - _input_hash = "" + # History retention is observational and must not alter + # workflow graph semantics. Stream the hash so summary mode + # can still deduplicate large inputs without retaining a + # second complete JSON payload. + _input_hash = _input_hash_for_dedup(node_id, inputs) if _input_hash and node_id in _dedup_hashes and _dedup_hashes[node_id] == _input_hash: _logger.info( "wf.step.dedup_skip node=%s (identical input hash %s)", @@ -330,8 +347,24 @@ def merge_payload(src: str, payload: Dict[str, Any]) -> None: continue # ── Phase 2: execute ready items ────────────────────────── - use_parallel = self.max_parallel_workers > 1 and len(ready) > 1 + # Process-isolated nodes own a child process and a cancellation + # boundary. Keep their engine-level execution serial so a + # mixed batch cannot leave an isolated child running after a + # host-thread sibling times out. + use_parallel = ( + self.max_parallel_workers > 1 + and len(ready) > 1 + and not any(item[1].process_isolated for item in ready) + ) exec_results: List[_ExecOutcome] = [] + _outs: Optional[Dict[str, Any]] = None + _so = "" + _sfut = None + _exec_args = None + _fut_map = None + _fut = None + _f2 = None + _pool = None # Call on_step_start hooks (always from main thread) step_tokens: Dict[int, Any] = {} @@ -439,6 +472,9 @@ def _par_exec( _ExecOutcome(_idx, _outs, _so, None, None, (time.perf_counter() - _t0) * 1000.0) ) except FuturesTimeoutError as _fte: + _process_isolated = bool(_nd.process_isolated) + if _process_isolated: + timeout_cancel_event.set() _fte_msg = str(_fte).strip() _err = _fte_msg if _fte_msg else f"节点执行超时 ({self.node_timeout_s}s)" exec_results.append( @@ -448,10 +484,21 @@ def _par_exec( ) if timeout_executor is not None: try: - timeout_executor.shutdown(wait=False, cancel_futures=True) + timeout_executor.shutdown( + wait=_process_isolated, + cancel_futures=True, + ) except TypeError: - timeout_executor.shutdown(wait=False) - timeout_executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="wf-node") + timeout_executor.shutdown(wait=_process_isolated) + timeout_executor = None + if _process_isolated and not _nd.timeout_fatal: + timeout_cancel_event.clear() + if _nd.timeout_fatal: + break + timeout_executor = ThreadPoolExecutor( + max_workers=1, + thread_name_prefix="wf-node", + ) except RunCancelledError as _ce: _ce.execution_context = state.build_context() raise @@ -600,6 +647,19 @@ def _build_step_result( "wf.step_end.hook_error", extra={"run_id": rid, "step": _sn, "node_id": _nid}, ) + if _eo.is_timeout and _nd.timeout_fatal and _stop_exc is None: + _stop_exc = NodeTimeoutError( + node_id=_nid, + timeout_s=float(step_timeout_s or 0), + execution_context={ + "run_id": rid, + "steps": state.steps + len(exec_results), + "last_node_id": _nid, + "outputs": state.last_outputs, + "history": state.history, + }, + ) + continue if self.stop_on_error and _stop_exc is None and not _eo.is_timeout: _stop_exc = NodeExecutionError( node_id=_nid, @@ -679,7 +739,7 @@ def _build_step_result( # Enqueue downstream (skip for failed node when stop_on_error) if _stop_exc is None: - state.record_vertex_output(_nid, _eo.outputs) + state.record_vertex_output_keys(_nid, _eo.outputs) for edge, edge_inputs in edge_resolver.resolve( node=_nd, node_inputs=_inp, @@ -689,6 +749,26 @@ def _build_step_result( q.append((edge.to, edge_inputs, _nid)) state.steps += len(exec_results) + # Break references held by this iteration before the next + # downstream node starts. Explicit edge mappings only protect + # the queued payload; stale outcomes/futures would otherwise + # keep excluded values alive for another execution window. + ready.clear() + exec_results.clear() + step_tokens.clear() + if _exec_args is not None: + _exec_args.clear() + if _fut_map is not None: + _fut_map.clear() + _eo = None + _inp = {} + _outs = None + _so = "" + _sfut = None + _fut = None + _f2 = None + _pool = None + step_res = None if _stop_exc is not None: raise _stop_exc if cancel is not None and cancel(): @@ -708,13 +788,14 @@ def _build_step_result( raise NodeExecutionError(node_id=pending_joins[0][0], message=msg) return state.to_result() finally: - if isinstance(self.runtime, PythonExecRuntime): - self.runtime.cancel_checker = previous_cancel_checker if timeout_executor is not None: try: timeout_executor.shutdown(wait=False, cancel_futures=True) except TypeError: timeout_executor.shutdown(wait=False) + if isinstance(self.runtime, PythonExecRuntime): + self.runtime.cancel_checker = previous_cancel_checker + self.runtime.enable_cancel_trace = bool(previous_enable_cancel_trace) def run_node(self, node_id: str, inputs: Dict[str, Any]) -> "StepResult": """Public API: execute a single node by id and return a StepResult. @@ -767,6 +848,13 @@ def _execute_node( node_id = node.id if node.type == "python": assert node.code is not None + if node.process_isolated and isinstance(_rt, PythonExecRuntime): + _rt = HostProcessPythonExecRuntime( + tool_registry=_rt.tool_registry, + cancel_checker=_rt.cancel_checker, + inherited_fd_keys=tuple(node.process_inherit_fd_keys), + retained_fd_keys=tuple(node.process_retain_fd_keys), + ) return _rt.execute(node.code, inputs) if node.type == "logic": assert self.code_gen is not None @@ -781,6 +869,13 @@ def _execute_node( ) from gen_error if self.mutate_workflow: node.code = code + if node.process_isolated and isinstance(_rt, PythonExecRuntime): + _rt = HostProcessPythonExecRuntime( + tool_registry=_rt.tool_registry, + cancel_checker=_rt.cancel_checker, + inherited_fd_keys=tuple(node.process_inherit_fd_keys), + retained_fd_keys=tuple(node.process_retain_fd_keys), + ) return _rt.execute(code, inputs) if node.type in {"branch", "loop"}: return {}, "" @@ -993,7 +1088,9 @@ def _try_get_by_path_from_scopes( scopes: List[Dict[str, Any]], path: str, ) -> tuple[bool, Any]: - return EdgeResolver(dataflow_mode=self.dataflow_mode, trace=self.trace).try_get_by_path_from_scopes(scopes, path) + return EdgeResolver(dataflow_mode=self.dataflow_mode, trace=self.trace).try_get_by_path_from_scopes( + scopes, path + ) def _try_get_by_path(self, data: Any, path: str) -> tuple[bool, Any]: return EdgeResolver(dataflow_mode=self.dataflow_mode, trace=self.trace).try_get_by_path(data, path) diff --git a/flocks/workflow/errors.py b/flocks/workflow/errors.py index 0a611229e..67f2ec5bc 100644 --- a/flocks/workflow/errors.py +++ b/flocks/workflow/errors.py @@ -55,10 +55,19 @@ def __init__(self, run_id: str, timeout_s: float): self.timeout_s = timeout_s -class NodeTimeoutError(FlocksWorkflowError): - """Raised when a single node execution exceeds its time limit (skipped, workflow continues).""" +class NodeTimeoutError(NodeExecutionError): + """Raised when a node exceeds its time limit and aborts the workflow run.""" - def __init__(self, node_id: str, timeout_s: float): - super().__init__(f"节点执行超时 ({timeout_s}s): node_id={node_id}") - self.node_id = node_id + def __init__( + self, + node_id: str, + timeout_s: float, + *, + execution_context: Optional[dict] = None, + ): + super().__init__( + node_id=node_id, + message=f"节点执行超时 ({timeout_s}s): node_id={node_id}", + execution_context=execution_context, + ) self.timeout_s = timeout_s diff --git a/flocks/workflow/execution_state.py b/flocks/workflow/execution_state.py index e418217fa..45fb87c95 100644 --- a/flocks/workflow/execution_state.py +++ b/flocks/workflow/execution_state.py @@ -43,14 +43,14 @@ class WorkflowExecutionState: history: list[StepResult] = field(default_factory=list) join_inputs: Dict[str, Dict[str, Dict[str, Any]]] = field(default_factory=dict) join_seen_sources: Dict[str, Set[str]] = field(default_factory=dict) - vertex_outputs: Dict[str, Dict[str, Any]] = field(default_factory=dict) + vertex_output_keys: Dict[str, list[str]] = field(default_factory=dict) def retain_step(self, step: StepResult) -> None: if self.retain_history: self.history.append(step) - def record_vertex_output(self, node_id: str, outputs: Dict[str, Any]) -> None: - self.vertex_outputs[node_id] = outputs + def record_vertex_output_keys(self, node_id: str, outputs: Dict[str, Any]) -> None: + self.vertex_output_keys[node_id] = list(outputs)[:_VERTEX_OUTPUT_KEY_LIMIT] def build_context(self) -> Dict[str, Any]: return { @@ -59,10 +59,7 @@ def build_context(self) -> Dict[str, Any]: "last_node_id": self.last_node_id, "outputs": self.last_outputs, "history": self.history, - "vertex_output_keys": { - node_id: list(outputs.keys())[:_VERTEX_OUTPUT_KEY_LIMIT] - for node_id, outputs in self.vertex_outputs.items() - }, + "vertex_output_keys": dict(self.vertex_output_keys), } def to_result(self) -> ExecutionResult: diff --git a/flocks/workflow/execution_store.py b/flocks/workflow/execution_store.py index c9389c445..9d0ae1c79 100644 --- a/flocks/workflow/execution_store.py +++ b/flocks/workflow/execution_store.py @@ -4,6 +4,7 @@ import asyncio from itertools import islice +import sys import time import uuid from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple @@ -27,16 +28,17 @@ "raw_alerts", "normalized_alerts", "filtered_alerts", + "enriched_alerts_with_triage", } ) -# Lists smaller than this many items are passed through verbatim. The cap -# protects against accidentally stripping small metadata lists that happen -# to share a name with a known large-list key. +# Lists below the item threshold stay inspectable unless their estimated +# in-memory footprint exceeds the collection byte cap. DEFAULT_COMPACT_SIZE_THRESHOLD: int = 100 DEFAULT_GENERIC_SEQUENCE_THRESHOLD: int = 1_000 DEFAULT_MAX_INLINE_STRING_CHARS: int = 20_000 DEFAULT_MAX_INLINE_DICT_KEYS: int = 200 +DEFAULT_MAX_INLINE_COLLECTION_BYTES: int = 1 * 1024 * 1024 DEFAULT_PREVIEW_ITEMS: int = 3 DEFAULT_PREVIEW_CHARS: int = 500 @@ -76,6 +78,29 @@ def _summarize_large_value(value: Any, *, depth: int = 0) -> Dict[str, Any]: } +def _estimated_size_exceeds(value: Any, max_bytes: int) -> bool: + """Return whether a bounded recursive size estimate exceeds *max_bytes*.""" + remaining = max_bytes + stack = [value] + seen_containers: set[int] = set() + while stack: + item = stack.pop() + if isinstance(item, (dict, list, tuple, set)): + item_id = id(item) + if item_id in seen_containers: + continue + seen_containers.add(item_id) + remaining -= sys.getsizeof(item) + if remaining < 0: + return True + if isinstance(item, dict): + stack.extend(item.keys()) + stack.extend(item.values()) + elif isinstance(item, (list, tuple, set)): + stack.extend(item) + return False + + def _compact_value_for_storage( value: Any, *, @@ -85,12 +110,13 @@ def _compact_value_for_storage( generic_sequence_threshold: int, max_inline_string_chars: int, max_inline_dict_keys: int, + max_inline_collection_bytes: int, depth: int = 0, ) -> Any: if ( key in known_large_keys and isinstance(value, (list, tuple)) - and len(value) > size_threshold + and (len(value) > size_threshold or _estimated_size_exceeds(value, max_inline_collection_bytes)) ): return {f"_{key}_count": len(value)} @@ -100,12 +126,18 @@ def _compact_value_for_storage( return value if isinstance(value, (list, tuple, set)): - if len(value) > generic_sequence_threshold: + if len(value) > generic_sequence_threshold or _estimated_size_exceeds( + value, + max_inline_collection_bytes, + ): return _summarize_large_value(value) return value if isinstance(value, dict): - if len(value) > max_inline_dict_keys: + if len(value) > max_inline_dict_keys or _estimated_size_exceeds( + value, + max_inline_collection_bytes, + ): return _summarize_large_value(value) if depth >= 2: return value @@ -120,6 +152,7 @@ def _compact_value_for_storage( generic_sequence_threshold=generic_sequence_threshold, max_inline_string_chars=max_inline_string_chars, max_inline_dict_keys=max_inline_dict_keys, + max_inline_collection_bytes=max_inline_collection_bytes, depth=depth + 1, ) if isinstance(child_compacted, dict) and len(child_compacted) == 1: @@ -143,6 +176,7 @@ def compact_outputs_for_storage( generic_sequence_threshold: int = DEFAULT_GENERIC_SEQUENCE_THRESHOLD, max_inline_string_chars: int = DEFAULT_MAX_INLINE_STRING_CHARS, max_inline_dict_keys: int = DEFAULT_MAX_INLINE_DICT_KEYS, + max_inline_collection_bytes: int = DEFAULT_MAX_INLINE_COLLECTION_BYTES, ) -> Dict[str, Any]: """Return a bounded copy of *outputs* safe for execution records. @@ -164,6 +198,7 @@ def compact_outputs_for_storage( generic_sequence_threshold=generic_sequence_threshold, max_inline_string_chars=max_inline_string_chars, max_inline_dict_keys=max_inline_dict_keys, + max_inline_collection_bytes=max_inline_collection_bytes, ) if isinstance(value, dict) and len(value) == 1: marker_key = next(iter(value)) diff --git a/flocks/workflow/models.py b/flocks/workflow/models.py index 57f3761a4..79b449d2f 100644 --- a/flocks/workflow/models.py +++ b/flocks/workflow/models.py @@ -26,6 +26,21 @@ class Node(BaseModel): join_namespace_key: str = "__by_source__" input_schema: Optional[Dict[str, Any]] = Field(None, alias="inputSchema") output_schema: Optional[Dict[str, Any]] = Field(None, alias="outputSchema") + # Trusted host subprocess isolation for memory-heavy Python nodes. The + # child exits after the node, returning its allocator arenas to the OS. + process_isolated: bool = Field(False, alias="processIsolated") + process_inherit_fd_keys: List[str] = Field( + default_factory=list, + alias="processInheritFdKeys", + ) + # File descriptors that must remain owned by the parent while the child + # executes, then continue downstream unchanged. Unlike inherited FDs these + # are only opaque numeric values in the child, which is portable to Windows. + process_retain_fd_keys: List[str] = Field( + default_factory=list, + alias="processRetainFdKeys", + ) + timeout_fatal: bool = Field(False, alias="timeoutFatal") # tool 节点 tool_name: Optional[str] = None @@ -52,6 +67,12 @@ class Node(BaseModel): @model_validator(mode="after") def _validate_code(self) -> "Node": + overlapping_fd_keys = set(self.process_inherit_fd_keys).intersection(self.process_retain_fd_keys) + if overlapping_fd_keys: + raise ValueError( + "processInheritFdKeys and processRetainFdKeys must not overlap: " + + ", ".join(sorted(overlapping_fd_keys)) + ) if self.type == "python": if self.code is None or not str(self.code).strip(): raise ValueError("python node requires non-empty code") diff --git a/flocks/workflow/repl_runtime.py b/flocks/workflow/repl_runtime.py index 4b9618e96..86ede68d8 100644 --- a/flocks/workflow/repl_runtime.py +++ b/flocks/workflow/repl_runtime.py @@ -6,12 +6,16 @@ import io import json import os +import queue import shlex +import signal import subprocess import sys +import tempfile import threading import traceback import uuid +from concurrent.futures import ThreadPoolExecutor as _ThreadPoolExecutor from concurrent.futures import TimeoutError as _FuturesTimeoutError from dataclasses import dataclass, field from typing import Any, Callable, ClassVar, Dict, Optional, TextIO, Tuple @@ -56,6 +60,7 @@ class PythonExecRuntime(Runtime): tool_registry: Optional[Any] = None # FlocksToolAdapter or compatible cancel_checker: Optional[Callable[[], bool]] = None cleanup_globals_after_execute: bool = False + enable_cancel_trace: bool = True _RUNTIME_GLOBAL_KEYS: ClassVar[frozenset[str]] = frozenset( { @@ -146,7 +151,7 @@ def _cancel_trace(_frame: Any, event: str, _arg: Any) -> Any: try: try: - if self.cancel_checker is not None: + if self.cancel_checker is not None and self.enable_cancel_trace: previous_trace = sys.gettrace() sys.settrace(_cancel_trace) with contextlib.redirect_stdout(buf): @@ -213,7 +218,7 @@ def _cancel_trace(_frame: Any, event: str, _arg: Any) -> Any: raise NodeExecutionError(node_id="", message="`outputs` must be a dict") return dict(out_obj), buf.getvalue() finally: - if self.cancel_checker is not None: + if self.cancel_checker is not None and self.enable_cancel_trace: sys.settrace(previous_trace) if self.cleanup_globals_after_execute: for key in list(g.keys()): @@ -376,11 +381,13 @@ def reset(self) -> None: # Stateless runtime; each execute call runs isolated command. return - def _build_python_cmd( + def _build_python_source( self, *, code: str, bridge_token: str, + rpc_max_bytes: Optional[int] = _RPC_MAX_BYTES, + extra_site_packages: str = _WORKFLOW_SITE_PACKAGES, ) -> str: wrapped = f""" import contextlib @@ -388,55 +395,106 @@ def _build_python_cmd( import json import os import sys +import threading import traceback outputs = {{}} -_MAX = {_RPC_MAX_BYTES} +_MAX = {rpc_max_bytes!r} _TOKEN = {json.dumps(bridge_token, ensure_ascii=False)} def _read_json_line(): line = sys.stdin.readline() if not line: raise RuntimeError("Bridge channel closed") - if len(line) > _MAX: + if _MAX is not None and len(line) > _MAX: raise RuntimeError("Bridge payload too large") obj = json.loads(line) if not isinstance(obj, dict): raise RuntimeError("Bridge payload must be an object") return obj +def _fail_rpc_waiters(message): + with _rpc_waiters_lock: + waiters = list(_rpc_waiters.values()) + for waiter in waiters: + waiter["response"] = {{ + "type": "rpc_result", + "ok": False, + "error": message, + }} + waiter["event"].set() + +def _read_rpc_responses(): + try: + while True: + resp = _read_json_line() + req_id = str(resp.get("id") or "") + with _rpc_waiters_lock: + waiter = _rpc_waiters.get(req_id) + if waiter is not None: + waiter["response"] = resp + waiter["event"].set() + except Exception as exc: + _fail_rpc_waiters(str(exc)) + def _rpc_call(payload): global _rpc_seq - _rpc_seq += 1 + with _rpc_seq_lock: + _rpc_seq += 1 + req_id = str(_rpc_seq) + waiter = {{"event": threading.Event(), "response": None}} + with _rpc_waiters_lock: + _rpc_waiters[req_id] = waiter req = {{ "type": "rpc", "token": _TOKEN, - "id": str(_rpc_seq), + "id": req_id, "rpc": dict(payload), }} - # Always use original stdout for RPC control channel. - # User code stdout may be redirected for capture. - _out = sys.__stdout__ if getattr(sys, "__stdout__", None) is not None else sys.stdout - _out.write(json.dumps(req, ensure_ascii=False, default=str) + "\\n") - _out.flush() - resp = _read_json_line() + try: + # Serialize complete frames, while allowing multiple requests to be + # in flight and responses to arrive out of order. + _out = sys.__stdout__ if getattr(sys, "__stdout__", None) is not None else sys.stdout + with _rpc_write_lock: + _out.write(json.dumps(req, ensure_ascii=False, default=str) + "\\n") + _out.flush() + waiter["event"].wait() + resp = waiter["response"] or {{}} + finally: + with _rpc_waiters_lock: + _rpc_waiters.pop(req_id, None) if ( resp.get("type") != "rpc_result" or resp.get("token") != _TOKEN - or str(resp.get("id")) != str(_rpc_seq) + or str(resp.get("id")) != req_id ): raise RuntimeError("Invalid RPC response frame") if not resp.get("ok", False): raise RuntimeError(str(resp.get("error") or "Bridge request failed")) return resp.get("output") +def _cancel_requested(): + return bool(_rpc_call({{"kind": "cancelled"}})) + init = _read_json_line() if init.get("type") != "init" or init.get("token") != _TOKEN: raise RuntimeError("Invalid init frame") inputs = init.get("inputs", {{}}) if not isinstance(inputs, dict): raise RuntimeError("inputs must be an object") +_retained_fd_keys = inputs.pop("_process_retained_fd_keys", []) +if not isinstance(_retained_fd_keys, list): + raise RuntimeError("retained fd keys must be a list") +for _retained_fd_key in _retained_fd_keys: + if isinstance(_retained_fd_key, str) and _retained_fd_key in inputs: + inputs[_retained_fd_key] = os.open(os.devnull, os.O_RDWR) _rpc_seq = 0 +_rpc_seq_lock = threading.Lock() +_rpc_write_lock = threading.Lock() +_rpc_waiters_lock = threading.Lock() +_rpc_waiters = {{}} +_rpc_reader_thread = threading.Thread(target=_read_rpc_responses, daemon=True) +_rpc_reader_thread.start() class _ToolProxy: def run(self, name, **kwargs): @@ -502,20 +560,25 @@ def get_path(path, data=None): g = {{ "inputs": inputs, "outputs": outputs, + "cancelled": _cancel_requested, + "is_cancelled": _cancel_requested, "get_path": get_path, "tool": _ToolProxy(), "llm": _LLMProxy(), }} -_extra_site = {json.dumps(_WORKFLOW_SITE_PACKAGES, ensure_ascii=False)} +_extra_site = {json.dumps(extra_site_packages, ensure_ascii=False)} if _extra_site and os.path.isdir(_extra_site) and _extra_site not in sys.path: sys.path.insert(0, _extra_site) buf = io.StringIO() payload = {{"outputs": {{}}, "stdout": "", "error": None}} try: - with contextlib.redirect_stdout(buf): - exec({json.dumps(code, ensure_ascii=False)}, g, g) + try: + with contextlib.redirect_stdout(buf): + exec({json.dumps(code, ensure_ascii=False)}, g, g) + except SystemExit: + pass out = g.get("outputs", {{}}) if out is None: out = {{}} @@ -534,7 +597,24 @@ def get_path(path, data=None): sys.stdout.write(json.dumps({{"type": "final", "token": _TOKEN, "payload": payload}}, ensure_ascii=False) + "\\n") sys.stdout.flush() """ - return f"python3 -I -c {shlex.quote(wrapped)}" + return wrapped + + def _build_python_cmd( + self, + *, + code: str, + bridge_token: str, + rpc_max_bytes: Optional[int] = _RPC_MAX_BYTES, + extra_site_packages: str = _WORKFLOW_SITE_PACKAGES, + python_executable: str = "python3", + ) -> str: + wrapped = self._build_python_source( + code=code, + bridge_token=bridge_token, + rpc_max_bytes=rpc_max_bytes, + extra_site_packages=extra_site_packages, + ) + return f"{shlex.quote(python_executable)} -I -c {shlex.quote(wrapped)}" def _write_json_line(self, stream: TextIO, payload: Dict[str, Any]) -> None: stream.write(json.dumps(payload, ensure_ascii=False, default=str) + "\n") @@ -560,6 +640,10 @@ def _handle_rpc_request(self, *, msg: Dict[str, Any], token: str) -> Dict[str, A kind = str(rpc.get("kind") or "").strip().lower() try: + if kind == "cancelled": + output = bool(self.cancel_checker and self.cancel_checker()) + return {"type": "rpc_result", "token": token, "id": req_id, "ok": True, "output": output} + if kind in ("tool", "tool_safe"): name = str(rpc.get("name") or "").strip() if not name: @@ -570,6 +654,11 @@ def _handle_rpc_request(self, *, msg: Dict[str, Any], token: str) -> Dict[str, A if not isinstance(kwargs, dict): raise RuntimeError("Tool kwargs must be an object") registry = self.tool_registry or get_tool_registry() + if hasattr(registry, "cancel_checker"): + try: + registry.cancel_checker = self.cancel_checker + except Exception: + pass if kind == "tool_safe": output = registry.run_safe(name, **kwargs) else: @@ -624,6 +713,327 @@ def _handle_rpc_request(self, *, msg: Dict[str, Any], token: str) -> Dict[str, A return {"type": "rpc_result", "token": token, "id": req_id, "ok": False, "error": str(e)} +@dataclass +class HostProcessPythonExecRuntime(SandboxPythonExecRuntime): + """Run one trusted Python node in a short-lived host subprocess. + + Tool and LLM calls are bridged back to the configured parent registry, + while allocations made by node code live in the child and are returned to + the OS when it exits. Cancellation terminates the child instead of leaving + an unkillable workflow thread behind. + """ + + sandbox: Dict[str, Any] = field(default_factory=dict) + inherited_fd_keys: Tuple[str, ...] = () + retained_fd_keys: Tuple[str, ...] = () + + def execute(self, code: str, inputs: Dict[str, Any]) -> Tuple[Dict[str, Any], str]: + if not isinstance(code, str): + raise NodeExecutionError( + node_id="", + message=f"Code must be a string, got {type(code).__name__}", + ) + if not code.strip(): + raise NodeExecutionError(node_id="", message="Code cannot be empty or whitespace-only") + if not isinstance(inputs, dict): + raise NodeExecutionError( + node_id="", + message=f"Inputs must be a dict, got {type(inputs).__name__}", + ) + + inherited_fds = self._resolve_inherited_fds(inputs) + retained_fds = self._resolve_retained_fds(inputs) + managed_fds = tuple(dict.fromkeys((*inherited_fds, *retained_fds.values()))) + child_inputs = dict(inputs) + if retained_fds: + child_inputs["_process_retained_fd_keys"] = list(retained_fds) + if self._cancel_requested(): + self._close_parent_fds(managed_fds) + raise RunCancelledError("") + token = uuid.uuid4().hex + package_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + popen_kwargs: Dict[str, Any] = {} + if inherited_fds: + if os.name == "nt": + self._close_parent_fds(managed_fds) + raise NodeExecutionError( + node_id="", + message="Inherited workflow file descriptors are not supported on Windows", + ) + popen_kwargs["pass_fds"] = inherited_fds + + script_path: Optional[str] = None + if sys.platform == "win32": + python_source = self._build_python_source( + code=code, + bridge_token=token, + rpc_max_bytes=None, + extra_site_packages=package_root, + ) + with tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + suffix=".py", + delete=False, + ) as script: + script.write(python_source) + script_path = script.name + process_args = [sys.executable, "-I", script_path] + else: + python_cmd = self._build_python_cmd( + code=code, + bridge_token=token, + rpc_max_bytes=None, + extra_site_packages=package_root, + python_executable=sys.executable, + ) + process_args = ["sh", "-lc", python_cmd] + + try: + proc = subprocess.Popen( + process_args, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + bufsize=1, + env={**os.environ, "LC_ALL": "C.UTF-8", "LANG": "C.UTF-8"}, + start_new_session=(os.name != "nt"), + **popen_kwargs, + ) + except Exception as exc: + self._remove_temp_script(script_path) + self._close_parent_fds(managed_fds) + raise NodeExecutionError( + node_id="", + message=f"Isolated host execution failed to start: {exc}", + ) from exc + + if proc.stdin is None or proc.stdout is None or proc.stderr is None: + self._terminate_process(proc) + self._remove_temp_script(script_path) + self._close_parent_fds(managed_fds) + raise NodeExecutionError(node_id="", message="Isolated process stdio is unavailable") + + stderr_chunks: list[str] = [] + stderr_thread = threading.Thread( + target=_drain_text_stream, + args=(proc.stderr, stderr_chunks), + name="wf-process-stderr", + daemon=True, + ) + stderr_thread.start() + + stdout_lines: queue.Queue[Optional[str]] = queue.Queue() + + def _read_stdout() -> None: + try: + while True: + line = proc.stdout.readline() + if line == "": + break + stdout_lines.put(line) + finally: + stdout_lines.put(None) + + stdout_thread = threading.Thread( + target=_read_stdout, + name="wf-process-stdout", + daemon=True, + ) + stdout_thread.start() + + final_payload: Optional[Dict[str, Any]] = None + cancelled = False + stdin_lock = threading.Lock() + rpc_pool = _ThreadPoolExecutor(max_workers=32, thread_name_prefix="wf-process-rpc") + + def _handle_rpc(message: Dict[str, Any]) -> None: + response = self._handle_rpc_request(msg=message, token=token) + try: + with stdin_lock: + if proc.poll() is None: + self._write_json_line(proc.stdin, response) + except (BrokenPipeError, OSError, ValueError): + return + + try: + self._write_json_line(proc.stdin, {"type": "init", "token": token, "inputs": child_inputs}) + while True: + if self._cancel_requested(): + cancelled = True + self._terminate_process(proc) + break + try: + line = stdout_lines.get(timeout=0.05) + except queue.Empty: + continue + if line is None: + break + msg = self._parse_json_line(line) + if msg is None: + continue + msg_type = str(msg.get("type") or "").strip().lower() + if msg_type == "rpc": + rpc_pool.submit(_handle_rpc, msg) + elif msg_type == "final" and msg.get("token") == token: + payload = msg.get("payload") + final_payload = payload if isinstance(payload, dict) else {} + finally: + rpc_pool.shutdown( + wait=final_payload is not None and not cancelled, + cancel_futures=True, + ) + try: + with stdin_lock: + proc.stdin.close() + except Exception: + pass + if proc.poll() is None: + self._terminate_process(proc) + exit_code = proc.wait() + stdout_thread.join(timeout=1.0) + stderr_thread.join(timeout=1.0) + try: + proc.stdout.close() + except Exception: + pass + try: + proc.stderr.close() + except Exception: + pass + self._remove_temp_script(script_path) + + stderr_text = "".join(stderr_chunks) + if cancelled: + self._close_parent_fds(managed_fds) + raise RunCancelledError("") + if exit_code != 0: + self._close_parent_fds(managed_fds) + message = stderr_text.strip() or f"Isolated command exited with code {exit_code}" + raise NodeExecutionError( + node_id="", + message=f"Isolated host execution failed: {message}", + traceback=stderr_text, + ) + if final_payload is None: + self._close_parent_fds(managed_fds) + raise NodeExecutionError( + node_id="", + message="Isolated host execution did not produce final payload", + traceback=stderr_text, + ) + + stdout = str(final_payload.get("stdout") or "") + error = final_payload.get("error") + if error: + self._close_parent_fds(managed_fds) + raise NodeExecutionError( + node_id="", + message=f"Runtime error ({error.get('type', 'Exception')}): {error.get('message', '')}", + stdout=stdout, + traceback=str(error.get("traceback") or stderr_text or ""), + ) + outputs = final_payload.get("outputs", {}) + if outputs is None: + outputs = {} + if not isinstance(outputs, dict): + self._close_parent_fds(managed_fds) + raise NodeExecutionError(node_id="", message="`outputs` must be a dict") + for key, fd in retained_fds.items(): + outputs[key] = fd + return outputs, stdout + + def _cancel_requested(self) -> bool: + try: + return bool(self.cancel_checker and self.cancel_checker()) + except Exception: + return False + + def _resolve_inherited_fds(self, inputs: Dict[str, Any]) -> Tuple[int, ...]: + fds: list[int] = [] + for key in self.inherited_fd_keys: + value = inputs.get(key) + if value is None or (type(value) is int and value < 0): + continue + if type(value) is not int: + raise NodeExecutionError( + node_id="", + message=f"Process-isolated node requires an integer fd input: {key}", + ) + try: + os.fstat(value) + except OSError as exc: + raise NodeExecutionError( + node_id="", + message=f"Process-isolated node received a closed fd input: {key}", + ) from exc + if value not in fds: + fds.append(value) + return tuple(fds) + + def _resolve_retained_fds(self, inputs: Dict[str, Any]) -> Dict[str, int]: + fds: Dict[str, int] = {} + for key in self.retained_fd_keys: + value = inputs.get(key) + if value is None or (type(value) is int and value < 0): + continue + if type(value) is not int: + raise NodeExecutionError( + node_id="", + message=f"Process-isolated node requires an integer retained fd input: {key}", + ) + try: + os.fstat(value) + except OSError as exc: + raise NodeExecutionError( + node_id="", + message=f"Process-isolated node received a closed retained fd input: {key}", + ) from exc + fds[key] = value + return fds + + @staticmethod + def _close_parent_fds(fds: Tuple[int, ...]) -> None: + for fd in fds: + try: + os.close(fd) + except OSError: + pass + + @staticmethod + def _remove_temp_script(path: Optional[str]) -> None: + if not path: + return + try: + os.remove(path) + except OSError: + pass + + @staticmethod + def _terminate_process(proc: subprocess.Popen[str]) -> None: + if proc.poll() is not None: + return + if os.name != "nt": + try: + os.killpg(proc.pid, signal.SIGTERM) + except ProcessLookupError: + return + else: + proc.terminate() + try: + proc.wait(timeout=0.5) + except subprocess.TimeoutExpired: + if os.name != "nt": + try: + os.killpg(proc.pid, signal.SIGKILL) + except ProcessLookupError: + pass + else: + proc.kill() + proc.wait() + + @dataclass class PythonREPLRuntime(Runtime): _repl: Optional[object] = None diff --git a/pyproject.toml b/pyproject.toml index 384249fa7..4d6689711 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "flocks" -version = "v2026.7.29" +version = "v2026.8.4" description = "AI-Native SecOps platform with multi-agent collaboration" authors = [ {name = "Flocks Team", email = "team@example.com"} diff --git a/tests/agent/test_agent_factory.py b/tests/agent/test_agent_factory.py index 78f86fdac..25e67dcbf 100644 --- a/tests/agent/test_agent_factory.py +++ b/tests/agent/test_agent_factory.py @@ -421,7 +421,7 @@ def test_inject_sets_prompt_on_agent(self): (helper_path.parent / "__init__.py").touch() agents = {"dyn": agent} - inject_dynamic_prompts(agents, [], [], [], []) + inject_dynamic_prompts(agents, [], [], []) assert agent.prompt == "injected" def test_inject_handles_import_error_gracefully(self): @@ -429,13 +429,13 @@ def test_inject_handles_import_error_gracefully(self): agent = AgentInfo(name="broken", mode="subagent", native=False) agent.prompt_builder = "nonexistent.module.path:inject" agents = {"broken": agent} - inject_dynamic_prompts(agents, [], [], [], []) # Should not raise + inject_dynamic_prompts(agents, [], [], []) # Should not raise assert agent.prompt is None def test_inject_skips_agents_without_builder(self): agent = AgentInfo(name="static", mode="subagent", native=False, prompt="static prompt") agents = {"static": agent} - inject_dynamic_prompts(agents, [], [], [], []) + inject_dynamic_prompts(agents, [], [], []) assert agent.prompt == "static prompt" @pytest.mark.asyncio @@ -471,6 +471,9 @@ async def test_rex_junior_gets_injected_prompt(self): assert jr is not None assert jr.prompt is not None assert len(jr.prompt) > 50 + assert "Execute implementation tasks directly." in jr.prompt + assert "NEVER delegate or spawn other agents" not in jr.prompt + assert "for read-only research only" in jr.prompt @pytest.mark.asyncio async def test_static_prompt_agents(self): diff --git a/tests/agent/test_agent_metadata.py b/tests/agent/test_agent_metadata.py index 5222d3e39..2d607d0d1 100644 --- a/tests/agent/test_agent_metadata.py +++ b/tests/agent/test_agent_metadata.py @@ -49,9 +49,28 @@ def test_is_delegatable_special_agents(self): assert is_delegatable("rex-junior") is False def test_is_delegatable_unknown_agent(self): - """未知 agent 应该返回 True（保守策略，允许插件 agent 委派）""" - assert is_delegatable("unknown-agent") is True - assert is_delegatable("custom-agent-123") is True + """Unknown agents must be rejected instead of falling back at runtime.""" + assert is_delegatable("unknown-agent") is False + assert is_delegatable("custom-agent-123") is False + + def test_is_delegatable_hidden_agent(self, monkeypatch): + """Hidden agents must not be exposed as delegation targets.""" + from flocks.agent import registry + from flocks.agent.agent import AgentInfo + + hidden_agent = AgentInfo( + name="hidden-agent", + mode="subagent", + hidden=True, + delegatable=True, + ) + monkeypatch.setattr(registry, "_agents_ref", {"hidden-agent": hidden_agent}) + + assert is_delegatable("hidden-agent") is False + + def test_is_delegatable_resolves_legacy_aliases(self): + """Known aliases should use the target agent's delegation policy.""" + assert is_delegatable("sisyphus") is False def test_get_agent_mode(self): """测试获取 agent 模式""" diff --git a/tests/agent/test_agent_workflows.py b/tests/agent/test_agent_workflows.py index e251bb016..b7a626ee6 100644 --- a/tests/agent/test_agent_workflows.py +++ b/tests/agent/test_agent_workflows.py @@ -101,7 +101,7 @@ def test_workflows_passed_to_inject_function(self, tmp_path): """inject() receives the workflows list.""" received: list = [] builder_code = """ -def inject(agent_info, available_agents, tools, skills, categories, workflows=None): +def inject(agent_info, available_agents, tools, skills, workflows=None): agent_info.prompt = "injected" import builtins builtins._test_received_workflows = workflows @@ -119,7 +119,7 @@ def inject(agent_info, available_agents, tools, skills, categories, workflows=No prompt_builder=f"{tmp_path.name}.cap_builder:inject", ) workflows = _simple_workflows() - inject_dynamic_prompts({"cap_agent": agent}, [], [], [], [], workflows) + inject_dynamic_prompts({"cap_agent": agent}, [], [], [], workflows) import builtins result = getattr(builtins, "_test_received_workflows", None) assert result is not None @@ -134,7 +134,7 @@ def inject(agent_info, available_agents, tools, skills, categories, workflows=No def test_workflows_defaults_to_empty_list_when_none(self, tmp_path): """inject() receives [] when workflows=None.""" builder_code = """ -def inject(agent_info, available_agents, tools, skills, categories, workflows=None): +def inject(agent_info, available_agents, tools, skills, workflows=None): agent_info.prompt = repr(workflows) """ builder_path = tmp_path / "none_builder.py" @@ -149,15 +149,76 @@ def inject(agent_info, available_agents, tools, skills, categories, workflows=No native=False, prompt_builder=f"{tmp_path.name}.none_builder:inject", ) - inject_dynamic_prompts({"none_agent": agent}, [], [], [], [], None) + inject_dynamic_prompts({"none_agent": agent}, [], [], [], None) assert agent.prompt == "[]" finally: sys.path.pop(0) + def test_legacy_builder_receives_empty_categories_and_workflows(self, tmp_path): + """Legacy inject() signatures keep workflows in the correct argument.""" + builder_code = """ +def inject(agent_info, available_agents, tools, skills, category_context, workflows=None): + workflow_names = [workflow.name for workflow in (workflows or [])] + agent_info.prompt = repr((category_context, workflow_names)) +""" + builder_path = tmp_path / "legacy_builder.py" + builder_path.write_text(textwrap.dedent(builder_code), encoding="utf-8") + + import sys + sys.path.insert(0, str(tmp_path.parent)) + try: + agent = AgentInfo( + name="legacy_agent", + mode="subagent", + native=False, + prompt_builder=f"{tmp_path.name}.legacy_builder:inject", + ) + inject_dynamic_prompts( + {"legacy_agent": agent}, + [], + [], + [], + _simple_workflows(), + ) + assert agent.prompt == "([], ['ndr_triage', 'global_scan'])" + finally: + sys.path.pop(0) + + def test_legacy_variadic_builder_receives_both_compatibility_arguments(self, tmp_path): + """Variadic legacy builders receive categories before workflows.""" + builder_code = """ +def inject(agent_info, available_agents, tools, skills, *args): + categories, workflows = args + workflow_names = [workflow.name for workflow in workflows] + agent_info.prompt = repr((categories, workflow_names)) +""" + builder_path = tmp_path / "legacy_variadic_builder.py" + builder_path.write_text(textwrap.dedent(builder_code), encoding="utf-8") + + import sys + sys.path.insert(0, str(tmp_path.parent)) + try: + agent = AgentInfo( + name="legacy_variadic_agent", + mode="subagent", + native=False, + prompt_builder=f"{tmp_path.name}.legacy_variadic_builder:inject", + ) + inject_dynamic_prompts( + {"legacy_variadic_agent": agent}, + [], + [], + [], + _simple_workflows(), + ) + assert agent.prompt == "([], ['ndr_triage', 'global_scan'])" + finally: + sys.path.pop(0) + def test_inject_sets_prompt_on_agent(self, tmp_path): """inject() is expected to set agent_info.prompt.""" builder_code = """ -def inject(agent_info, available_agents, tools, skills, categories, workflows=None): +def inject(agent_info, available_agents, tools, skills, workflows=None): wf_names = ", ".join(w.name for w in (workflows or [])) agent_info.prompt = f"I know these workflows: {wf_names}" """ @@ -174,7 +235,7 @@ def inject(agent_info, available_agents, tools, skills, categories, workflows=No prompt_builder=f"{tmp_path.name}.prompt_builder:inject", ) workflows = _simple_workflows() - inject_dynamic_prompts({"wf_agent": agent}, [], [], [], [], workflows) + inject_dynamic_prompts({"wf_agent": agent}, [], [], [], workflows) assert "ndr_triage" in agent.prompt assert "global_scan" in agent.prompt finally: @@ -188,14 +249,14 @@ def test_agent_without_builder_unaffected(self): native=False, prompt="Static prompt", ) - inject_dynamic_prompts({"static_agent": agent}, [], [], [], [], _simple_workflows()) + inject_dynamic_prompts({"static_agent": agent}, [], [], [], _simple_workflows()) # Prompt should remain unchanged assert agent.prompt == "Static prompt" def test_builder_error_logged_not_raised(self, tmp_path): """A broken inject function logs the error but does not raise.""" builder_code = """ -def inject(agent_info, available_agents, tools, skills, categories, workflows=None): +def inject(agent_info, available_agents, tools, skills, workflows=None): raise RuntimeError("inject exploded") """ builder_path = tmp_path / "broken_builder.py" @@ -211,14 +272,14 @@ def inject(agent_info, available_agents, tools, skills, categories, workflows=No prompt_builder=f"{tmp_path.name}.broken_builder:inject", ) # Must not raise - inject_dynamic_prompts({"broken_agent": agent}, [], [], [], [], []) + inject_dynamic_prompts({"broken_agent": agent}, [], [], [], []) finally: sys.path.pop(0) def test_multiple_agents_all_injected(self, tmp_path): """All agents with a builder receive the workflows.""" builder_code = """ -def inject(agent_info, available_agents, tools, skills, categories, workflows=None): +def inject(agent_info, available_agents, tools, skills, workflows=None): agent_info.prompt = str(len(workflows or [])) """ for i in range(3): @@ -236,7 +297,7 @@ def inject(agent_info, available_agents, tools, skills, categories, workflows=No ) for i in range(3) } - inject_dynamic_prompts(agents, [], [], [], [], _simple_workflows()) + inject_dynamic_prompts(agents, [], [], [], _simple_workflows()) for agent in agents.values(): assert agent.prompt == "2" # 2 workflows in _simple_workflows() finally: diff --git a/tests/agent/test_prompt_utils.py b/tests/agent/test_prompt_utils.py index e90c8c685..6a1637e4b 100644 --- a/tests/agent/test_prompt_utils.py +++ b/tests/agent/test_prompt_utils.py @@ -16,16 +16,32 @@ import pytest -from flocks.agent.agent import AvailableAgent, AvailableCategory, AvailableSkill, AvailableTool, AvailableWorkflow +from flocks.agent.agent import AvailableAgent, AvailableSkill, AvailableTool, AvailableWorkflow from flocks.agent.prompt_utils import ( _format_tools_for_prompt, build_agent_selection_table, + build_skills_delegation_guide, build_tool_selection_table, build_workflows_section, categorize_tools, ) +def test_skills_delegation_guide_uses_subagent_routing(): + output = build_skills_delegation_guide([ + AvailableSkill( + name="security-review", + description="Review security-sensitive changes.", + location="/tmp/SKILL.md", + ) + ]) + + assert 'subagent_type="[selected-agent]"' in output + assert "security-review" in output + assert "category=" not in output + assert "Available Categories" not in output + + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- diff --git a/tests/channel/test_channel.py b/tests/channel/test_channel.py index 5b6c57430..ecf2e380c 100644 --- a/tests/channel/test_channel.py +++ b/tests/channel/test_channel.py @@ -650,11 +650,8 @@ async def fake_deliver(ctx, session_id=None): assert create_kwargs["title"] == "[Feishu] oc_group" assert "session_new" in delivered[0] assert "已开始全新对话。" in delivered[0] - # The previous session must be archived so it no longer appears in the - # active IM session list used for scheduled-task target resolution. - update_mock.assert_awaited_once() - assert update_mock.await_args.args == ("channel", "session_old") - assert update_mock.await_args.kwargs["status"] == "archived" + # Starting a new conversation must not archive the previous session. + update_mock.assert_not_awaited() @pytest.mark.asyncio async def test_new_command_inherits_auto_model_mode(self, monkeypatch): @@ -811,10 +808,10 @@ async def fake_deliver(ctx, session_id=None): ) ), ) - monkeypatch.setattr( - "flocks.session.session.Session.create", - AsyncMock(return_value=SimpleNamespace(id="session_new", agent="rex")), + create_mock = AsyncMock( + return_value=SimpleNamespace(id="session_new", agent="rex") ) + monkeypatch.setattr("flocks.session.session.Session.create", create_mock) monkeypatch.setattr( "flocks.session.session.Session.update", AsyncMock(return_value=None), @@ -834,6 +831,7 @@ async def fake_deliver(ctx, session_id=None): assert handled is True assert "已开始全新对话。" in delivered[0] + assert create_mock.await_args.kwargs["title"] == "[Slack] 叫我uuuu" append_mock.assert_awaited_once() assert append_mock.await_args.args[:3] == ("session_new", "叫我uuuu", msg) assert append_mock.await_args.kwargs["agent"] == "rex" @@ -1917,30 +1915,29 @@ async def test_unknown_channel_returns_none(self): # ------------------------------------------------------------------ -# _is_placeholder_text +# is_channel_media_placeholder # ------------------------------------------------------------------ -class TestIsPlaceholderText: +class TestIsChannelMediaPlaceholder: def test_recognises_channel_placeholders(self): - from flocks.channel.inbound.dispatcher import _is_placeholder_text - assert _is_placeholder_text("[图片消息]") is True - assert _is_placeholder_text("[图片消息: screenshot.png]") is True - assert _is_placeholder_text("[文件消息]") is True - assert _is_placeholder_text("[文件消息: report.pdf]") is True - assert _is_placeholder_text("[Image]") is True - assert _is_placeholder_text("[Attachment]") is True - assert _is_placeholder_text("[图片]") is True - assert _is_placeholder_text("[文件]") is True + from flocks.channel.inbound.session_binding import is_channel_media_placeholder + assert is_channel_media_placeholder("[图片消息]") is True + assert is_channel_media_placeholder("[图片消息: screenshot.png]") is True + assert is_channel_media_placeholder("[文件消息]") is True + assert is_channel_media_placeholder("[文件消息: report.pdf]") is True + assert is_channel_media_placeholder("[Image]") is True + assert is_channel_media_placeholder("[Attachment]") is True + assert is_channel_media_placeholder("[图片]") is True + assert is_channel_media_placeholder("[文件]") is True def test_does_not_match_normal_text(self): - from flocks.channel.inbound.dispatcher import _is_placeholder_text - assert _is_placeholder_text("hello") is False - assert _is_placeholder_text("看这个") is False - assert _is_placeholder_text("") is False - # Suffix beyond the placeholder is fine — the text starts with a - # placeholder token, so the dispatcher will still rewrite it. - assert _is_placeholder_text("[文件消息: x] extra text") is True - assert _is_placeholder_text("not a placeholder [文件消息: x]") is False + from flocks.channel.inbound.session_binding import is_channel_media_placeholder + assert is_channel_media_placeholder("hello") is False + assert is_channel_media_placeholder("看这个") is False + assert is_channel_media_placeholder("") is False + assert is_channel_media_placeholder("[文件消息: x] extra text") is False + assert is_channel_media_placeholder("[图片]\n请分析") is False + assert is_channel_media_placeholder("not a placeholder [文件消息: x]") is False # ------------------------------------------------------------------ @@ -1948,6 +1945,61 @@ def test_does_not_match_normal_text(self): # ------------------------------------------------------------------ class TestAppendUserMessagePerChannel: + @pytest.mark.asyncio + async def test_captioned_media_text_is_not_replaced(self, monkeypatch): + from flocks.channel.inbound.dispatcher import InboundDispatcher + from flocks.session.message import TextPart + + created_message = SimpleNamespace(id="m1") + caption = "[图片]: 请识别发票金额" + text_part = TextPart( + id="part_text", + sessionID="s1", + messageID="m1", + text=caption, + ) + store_part = AsyncMock() + monkeypatch.setattr( + "flocks.session.message.Message.create", + AsyncMock(return_value=created_message), + ) + monkeypatch.setattr( + "flocks.session.message.Message.store_part", + store_part, + ) + monkeypatch.setattr( + "flocks.session.message.Message.parts", + AsyncMock(return_value=[text_part]), + ) + + import flocks.channel.builtin.telegram.inbound_media as tg_inb + + async def fake_download(msg, config): + return SimpleNamespace( + filename="invoice.jpg", + mime="image/jpeg", + url="file:///tmp/invoice.jpg", + source={"channel": "telegram"}, + ) + + monkeypatch.setattr(tg_inb, "download_inbound_media", fake_download) + + await InboundDispatcher._append_user_message_unchecked( + "s1", + caption, + InboundMessage( + channel_id="telegram", + account_id="a", + message_id="m1", + sender_id="u", + media_url="telegram://photo/ABC", + ), + None, + ) + + assert store_part.await_count == 1 + assert store_part.await_args.args[2].type == "file" + @pytest.mark.asyncio async def test_wecom_pipeline_stores_file_part(self, monkeypatch): from flocks.channel.inbound.dispatcher import InboundDispatcher diff --git a/tests/channel/test_e2e_file_roundtrip.py b/tests/channel/test_e2e_file_roundtrip.py index ab4a71e6b..43d5873cd 100644 --- a/tests/channel/test_e2e_file_roundtrip.py +++ b/tests/channel/test_e2e_file_roundtrip.py @@ -47,7 +47,6 @@ from flocks.channel.inbound.dispatcher import ( InboundDispatcher, _download_channel_media, - _is_placeholder_text, ) diff --git a/tests/channel/test_feishu.py b/tests/channel/test_feishu.py index 1a44933dc..7d63b5b6e 100644 --- a/tests/channel/test_feishu.py +++ b/tests/channel/test_feishu.py @@ -549,7 +549,14 @@ def stop(self): assert ws_client.disconnected_error is disconnect_error -def test_build_ws_client_falls_back_to_modern_sdk(monkeypatch) -> None: +@pytest.mark.parametrize( + "dispatch_method", + ["do_without_validation", "_do_without_validation"], +) +def test_build_ws_client_falls_back_to_modern_sdk( + monkeypatch, + dispatch_method: str, +) -> None: dispatched: list[dict] = [] captured: dict[str, object] = {} @@ -570,7 +577,10 @@ def __init__(self, **kwargs): self._disconnect_called = False def start(self): - captured["event_handler"].do_without_validation(b'{"header":{"event_type":"ping"}}') + dispatcher = captured["event_handler"] + getattr(dispatcher, dispatch_method)( + b'{"header":{"event_type":"ping"}}', + ) async def _disconnect(self): self._disconnect_called = True diff --git a/tests/channel/test_session_binding.py b/tests/channel/test_session_binding.py index 0a09d8e43..acf1670d1 100644 --- a/tests/channel/test_session_binding.py +++ b/tests/channel/test_session_binding.py @@ -4,7 +4,137 @@ import pytest from flocks.channel.base import ChatType, InboundMessage -from flocks.channel.inbound.session_binding import SessionBinding, SessionBindingService +from flocks.channel.inbound.session_binding import ( + SessionBinding, + SessionBindingService, + _build_title, + _build_title_fallback, + extract_channel_title_text, + is_channel_media_placeholder, +) + + +def test_build_title_uses_first_user_input() -> None: + msg = InboundMessage( + channel_id="wecom", + account_id="default", + message_id="msg-1", + sender_id="user-1", + chat_id="room-1", + chat_type=ChatType.DIRECT, + text="你是谁", + ) + + assert _build_title(msg) == "[Wecom] 你是谁" + + +def test_build_title_prefers_mention_text_and_truncates_first_line() -> None: + msg = InboundMessage( + channel_id="feishu", + account_id="default", + message_id="msg-1", + sender_id="user-1", + chat_id="room-1", + chat_type=ChatType.GROUP, + text="@Rex ignored raw text", + mention_text=f"{'a' * 55}\nsecond line", + ) + + assert _build_title(msg) == f"[Feishu] {'a' * 47}..." + + +def test_build_title_falls_back_when_user_input_is_empty() -> None: + msg = InboundMessage( + channel_id="wecom", + account_id="default", + message_id="msg-1", + sender_id="user-1", + sender_name="Alice", + chat_id="room-1", + chat_type=ChatType.DIRECT, + text=" \n ", + ) + + assert _build_title(msg) == "[Wecom] DM — Alice" + + +@pytest.mark.parametrize("text", [ + "__merge_forward_expand__msg-1", + "[文件消息: report.pdf]", + "[Merged forward message]", +]) +def test_build_title_falls_back_for_channel_placeholders(text: str) -> None: + msg = InboundMessage( + channel_id="feishu", + account_id="default", + message_id="msg-1", + sender_id="user-1", + chat_id="room-1", + chat_type=ChatType.GROUP, + text=text, + ) + + assert _build_title(msg) == "[Feishu] room-1" + + +def test_build_title_uses_media_caption_or_following_text() -> None: + msg = InboundMessage( + channel_id="telegram", + account_id="default", + message_id="msg-1", + sender_id="user-1", + chat_id="room-1", + chat_type=ChatType.DIRECT, + text="[图片]\n请分析这张图", + ) + + assert _build_title(msg, "[文件: report.pdf]: 总结报告") == "[Telegram] 总结报告" + assert _build_title(msg) == "[Telegram] 请分析这张图" + + +def test_channel_title_helpers_share_placeholder_rules() -> None: + assert is_channel_media_placeholder("[图片消息]") is True + assert is_channel_media_placeholder("[文件: report.pdf]: 总结报告") is False + assert is_channel_media_placeholder("[图片]\n请分析这张图") is False + assert is_channel_media_placeholder("普通消息") is False + assert extract_channel_title_text("[图片消息]\n请分析这张图") == "请分析这张图" + assert extract_channel_title_text("[文件: report.pdf]: 总结报告") == "总结报告" + + +def test_build_title_fallback_does_not_reuse_command_text() -> None: + msg = InboundMessage( + channel_id="wecom", + account_id="default", + message_id="msg-1", + sender_id="user-1", + sender_name="Alice", + chat_id="room-1", + chat_type=ChatType.DIRECT, + text="/new", + ) + + assert _build_title_fallback(msg) == "[Wecom] DM — Alice" + + +@pytest.mark.asyncio +async def test_list_bindings_filters_by_session_ids() -> None: + cursor = SimpleNamespace(fetchall=AsyncMock(return_value=[])) + db = SimpleNamespace(execute=AsyncMock(return_value=cursor)) + + with patch( + "flocks.channel.inbound.session_binding._get_db", + AsyncMock(return_value=db), + ): + result = await SessionBindingService().list_bindings( + channel_id="wecom", + session_ids=["ses-1", "ses-2", "ses-1"], + ) + + assert result == [] + sql, params = db.execute.await_args.args + assert "channel_id = ?" in sql + assert "session_id IN (?,?)" in sql + assert params == ["wecom", "ses-1", "ses-2"] @pytest.mark.asyncio diff --git a/tests/cli/test_service_commands.py b/tests/cli/test_service_commands.py index d80928bc5..cbd3b4a54 100644 --- a/tests/cli/test_service_commands.py +++ b/tests/cli/test_service_commands.py @@ -31,6 +31,9 @@ def test_cli_help_lists_service_commands(monkeypatch, tmp_path) -> None: assert result.exit_code == 0 for command in ("start", "stop", "restart", "status", "logs", "session", "mcp", "task", "skills"): assert _help_contains_command(result.stdout, command) + assert "Agents must use `flocks restart" in result.stdout + assert "--server-only`;" in result.stdout + assert "bare restart stops the supervisor and terminates the running agent" in result.stdout for command in ("agent", "acp", "debug", "run", "serve", "service-watchdog", "service-daemon", "auth", "models"): assert not _help_contains_command(result.stdout, command) diff --git a/tests/config/test_config.py b/tests/config/test_config.py index eb99c29e1..d808ce1e0 100644 --- a/tests/config/test_config.py +++ b/tests/config/test_config.py @@ -5,6 +5,7 @@ import pytest import json from pathlib import Path +from unittest.mock import patch from flocks.config.config import Config, GlobalConfig, ConfigInfo, PermissionAction, PermissionConfig @@ -34,6 +35,26 @@ def test_global_config(): assert config.server_port == 8000 +def test_delegate_categories_are_not_part_of_config_schema(): + assert "categories" not in ConfigInfo.model_fields + + +def test_removed_delegate_categories_are_not_silently_preserved(): + with patch("flocks.utils.log.Log.create") as create_log: + config = ConfigInfo.model_validate( + { + "categories": { + "quick": { + "model": "anthropic/claude-haiku-4-5", + } + } + } + ) + + assert "categories" not in config.model_dump(exclude_none=True) + create_log.return_value.warn.assert_called_once() + + @pytest.mark.asyncio async def test_config_loading(): """Test configuration loading""" diff --git a/tests/hub/test_hub_catalog.py b/tests/hub/test_hub_catalog.py index 0eaab0640..4bd0bc11d 100644 --- a/tests/hub/test_hub_catalog.py +++ b/tests/hub/test_hub_catalog.py @@ -296,17 +296,17 @@ def test_catalog_uses_webui_workspace_version_for_inferred_installs( entry = {item.id: item for item in list_catalog(plugin_type="webui")}["soc_ui"] - assert entry.version == "1.1.4" + assert entry.version == "1.1.5" assert entry.state == "updateAvailable" assert entry.installedVersion == "1.0.0" - workspace["version"] = "1.1.4" + workspace["version"] = "1.1.5" workspace_path.write_text(json.dumps(workspace), encoding="utf-8") refreshed = {item.id: item for item in list_catalog(plugin_type="webui")}["soc_ui"] assert refreshed.state == "installed" - assert refreshed.installedVersion == "1.1.4" + assert refreshed.installedVersion == "1.1.5" def test_pentest_agents_are_listed_in_agent_catalog(): @@ -520,7 +520,11 @@ async def noop_refresh(_plugin_type, _changed_path=None): "threat_level": "high", "threat_phase": "exploit", "threat_type": "exploit", - "threat_result": "failed", + "threat_result": "success", + "attack_verdict": "attack_success", + "attack_success": True, + "triage_attack_verdict": "attack_failed", + "triage_attack_success": False, "_source_type": "tdp", "is_duplicate": False, } @@ -587,6 +591,11 @@ async def noop_refresh(_plugin_type, _changed_path=None): assert response.body["summary"]["totalRaw"] == 1 assert response.body["summary"]["attackFailed"] == 1 assert response.body["incidents"][0]["id"] == "alert-1" + assert response.body["incidents"][0]["triageAttackVerdict"] == "attack" + assert response.body["incidents"][0]["triageAttackSuccess"] == "failed" + assert response.body["incidents"][0]["conclusion"]["verdict"] == "attack" + assert response.body["incidents"][0]["tableCells"]["threat_result"]["value"] == "success" + assert response.body["incidents"][0]["tableCells"]["attack_success"]["value"] == "true" assert response.body["incidents"][0]["tableCells"]["_source_type"]["value"] == "tdp" assert response.body["incidents"][0]["tableCells"]["threat_severity"]["value"] == "critical" assert response.body["incidents"][0]["tableCells"]["threat_level"]["value"] == "high" @@ -611,6 +620,25 @@ async def noop_refresh(_plugin_type, _changed_path=None): assert filtered.body["summary"]["representativeCount"] == 1 assert [incident["id"] for incident in filtered.body["incidents"]] == ["alert-1"] + filtered_by_model_result = runtime.execute( + page_id="soc-alerts", + contract_id="soc.alerts.operations", + operation_name="list", + payload={ + "params": { + "filters": { + "triage_attack_verdict": ["attack"], + "triage_attack_success": ["failed"], + }, + "limit": 10, + } + }, + principal=AuthUser(id="u1", username="admin", role="admin"), + ) + + assert filtered_by_model_result.status_code == 200 + assert [incident["id"] for incident in filtered_by_model_result.body["incidents"]] == ["alert-1"] + for filters in ( {"threat_severity": ["low"]}, {"threat_level": ["low"]}, diff --git a/tests/hub/test_soc_dashboard_schema.py b/tests/hub/test_soc_dashboard_schema.py index 30e87ce49..957ed2cb1 100644 --- a/tests/hub/test_soc_dashboard_schema.py +++ b/tests/hub/test_soc_dashboard_schema.py @@ -2,7 +2,7 @@ import json import sqlite3 import sys -from datetime import datetime +from datetime import datetime, timedelta from pathlib import Path @@ -85,7 +85,10 @@ def test_soc_dashboard_migrates_legacy_alert_records_schema(tmp_path: Path): "_threat_type": "web_attack", "triage_status": "ok", "_triage_persisted_at": "2026-07-14T13:00:00", - "attack_verdict": "attack", + "attack_verdict": "benign", + "attack_success": False, + "triage_attack_verdict": "attack", + "triage_attack_success": "unknown", } with sqlite3.connect(db_path) as conn: conn.execute( @@ -139,6 +142,8 @@ def test_soc_dashboard_migrates_legacy_alert_records_schema(tmp_path: Path): updated_first_record = { **first_record, "_triage_persisted_at": "2026-07-14T13:02:00", + "triage_attack_verdict": "attack_failed", + "triage_attack_success": False, "triage_report": "# Updated report", } with sqlite3.connect(db_path) as conn: @@ -190,7 +195,8 @@ def test_soc_dashboard_migrates_legacy_alert_records_schema(tmp_path: Path): "SELECT meta_value FROM soc_dashboard_meta WHERE meta_key='schema_version'" ).fetchone()[0] updated_fact = conn.execute( - "SELECT triage_persisted_at, verdict FROM soc_dashboard_alert_facts " + "SELECT triage_persisted_at, triage_attack_verdict, triage_attack_success " + "FROM soc_dashboard_alert_facts " "WHERE row_key = '1'" ).fetchone() @@ -216,21 +222,53 @@ def test_soc_dashboard_migrates_legacy_alert_records_schema(tmp_path: Path): (2, "2", "hids", "Malware download", "malware", 1), ] assert "COALESCE(NULLIF(NEW.row_id, ''), CAST(NEW.rowid AS TEXT))" in trigger_sql - assert updated_fact == ("2026-07-14T13:02:00", "attack") - assert schema_version == "3" + assert updated_fact == ("2026-07-14T13:02:00", "attack", "failed") + assert schema_version == "4" def test_soc_dashboard_triage_outcomes_partition_records(tmp_path: Path): db_path = tmp_path / "soc.db" asset_date = "2026-07-14" records = [ - {"triage_status": "ok", "attack_verdict": "attack_success", "attack_success": True}, - {"triage_status": "ok", "attack_verdict": "attack"}, - {"triage_status": "ok", "attack_verdict": "attack_failed"}, - {"triage_status": "ok", "attack_verdict": "benign"}, - {"triage_status": "ok", "attack_verdict": "unknown"}, - {"triage_status": "failed", "attack_verdict": "unknown"}, - {"triage_status": "failed", "attack_verdict": "benign"}, + { + "triage_status": "ok", + "attack_verdict": "attack_failed", + "attack_success": False, + "triage_attack_verdict": "attack", + "triage_attack_success": "success", + }, + { + "triage_status": "ok", + "triage_attack_verdict": "attack", + "triage_attack_success": "unknown", + }, + { + "triage_status": "ok", + "attack_verdict": "attack_success", + "attack_success": True, + "triage_attack_verdict": "attack", + "triage_attack_success": "failed", + }, + { + "triage_status": "ok", + "triage_attack_verdict": "non_attack", + "triage_attack_success": "success", + }, + { + "triage_status": "ok", + "triage_attack_verdict": "unknown", + "triage_attack_success": "unknown", + }, + { + "triage_status": "failed", + "triage_attack_verdict": "unknown", + "triage_attack_success": "unknown", + }, + { + "triage_status": "failed", + "triage_attack_verdict": "non_attack", + "triage_attack_success": "unknown", + }, {"triage_status": "ok", "attack_verdict": "legacy", "attack_success": True}, ] severity_values = ["low", "critical", "high", "medium", "low", "critical", "high", "medium"] @@ -292,17 +330,22 @@ def test_soc_dashboard_triage_outcomes_partition_records(tmp_path: Path): "SELECT threat_name, threat_type, severity, risk_level " "FROM soc_dashboard_alert_facts ORDER BY alert_row_id LIMIT 1" ).fetchone() + normalized_non_attack = conn.execute( + "SELECT triage_attack_verdict, triage_attack_success " + "FROM soc_dashboard_alert_facts ORDER BY alert_row_id LIMIT 1 OFFSET 3" + ).fetchone() assert triage["totalRecords"] == 8 assert triage["newTriaged"] == 6 - assert triage["attackSuccess"] == 2 + assert triage["attackSuccess"] == 1 assert triage["attack"] == 1 assert triage["attackFailed"] == 1 - assert triage["attackTotal"] == 4 - assert triage["benign"] == 1 - assert triage["unknown"] == 1 + assert triage["attackTotal"] == 3 + assert triage["benign"] == 2 + assert triage["unknown"] == 3 assert triage["triageFailed"] == 2 assert first_fact == ("threat-name-0", "threat-type-0", "low", "High") + assert normalized_non_attack == ("non_attack", "unknown") assert dict(triage["threatTypeCounter"]) == { f"threat-type-{index}": 1 for index in range(len(records)) } @@ -313,10 +356,78 @@ def test_soc_dashboard_triage_outcomes_partition_records(tmp_path: Path): "medium": 2, } assert dict(triage["riskCounter"]) == {"high": 8} - assert closed_loop["pending"] == 3 - assert triage["attackTotal"] + triage["benign"] + closed_loop["pending"] == 8 + assert closed_loop["manualDecision"] == triage["unknown"] + assert closed_loop["pending"] == triage["unknown"] + assert triage["attackTotal"] + triage["benign"] + triage["unknown"] == 8 assert sum(timeline["attack"]) == triage["attackTotal"] + overview_handlers = _load_overview_handlers() + overview_handlers.DEFAULT_SQLITE_DB = db_path + overview_triage = overview_handlers._read_triage( + [ + overview_handlers._RecordSource( + path=db_path, + role="triage", + date=asset_date, + data_source="sqlite", + ) + ] + ) + + assert overview_triage["attackSuccess"] == 1 + assert overview_triage["attack"] == 1 + assert overview_triage["attackFailed"] == 1 + assert overview_triage["attackTotal"] == 3 + assert overview_triage["benign"] == 2 + assert overview_triage["unknown"] == 3 + + +def test_soc_dashboard_command_graph_uses_model_outcome_partition(): + page_path = ( + Path(__file__).resolve().parents[2] + / ".flocks" + / "flockshub" + / "plugins" + / "webuis" + / "soc_ui" + / "soc_dashboard" + / "src" + / "Page.tsx" + ) + source = page_path.read_text(encoding="utf-8") + start = source.index("function CommandGraph(") + end = source.index("\nfunction CommandMetric(", start) + command_graph = source[start:end] + + assert "value: stats.triage.attackTotal" in command_graph + assert "value: stats.triage.benign" in command_graph + assert "value: stats.triage.unknown" in command_graph + assert "value: stats.closedLoop.pending" not in command_graph + assert "'安全事件'" in command_graph + assert "'非安全事件'" in command_graph + assert "'待人工复核'" in command_graph + + +def test_soc_overview_uses_backend_five_class_verdicts(): + page_path = ( + Path(__file__).resolve().parents[2] + / ".flocks" + / "flockshub" + / "plugins" + / "webuis" + / "soc_ui" + / "soc_overview" + / "src" + / "index.tsx" + ) + source = page_path.read_text(encoding="utf-8") + + assert "verdicts?: CounterItem[]" in source + assert "verdicts: list(value.verdicts)" in source + assert "stats.triage.totalRecords - success - failed" not in source + for key in ("attack_success", "attack", "attack_failed", "non_attack", "unknown"): + assert key in source + def test_soc_overview_keeps_threat_names_and_types_separate(tmp_path: Path): db_path = tmp_path / "soc.db" @@ -375,6 +486,8 @@ def test_soc_dashboard_activity_does_not_mix_name_type_or_risk_fields(): { "_threat_type": "web_attack", "triage_status": "ok", + "triage_attack_verdict": "attack", + "triage_attack_success": "failed", "threat_level": "critical", } ), @@ -386,6 +499,9 @@ def test_soc_dashboard_activity_does_not_mix_name_type_or_risk_fields(): assert event["alert"]["threatType"] == "web_attack" assert event["result"]["threatSeverity"] == "" assert event["result"]["riskLevel"] == "" + assert event["result"]["verdict"] == "attack" + assert event["result"]["attackSuccess"] == "failed" + assert event["result"]["verdictLabel"] == "攻击失败" def test_soc_alert_verdict_does_not_fall_back_to_risk_or_threat_level(): @@ -394,7 +510,54 @@ def test_soc_alert_verdict_does_not_fall_back_to_risk_or_threat_level(): assert operations._verdict_bucket( {"risk_level": "attack_success", "threat_level": "attack_failed"} ) == "unknown" - assert operations._verdict_bucket({"attack_verdict": "attack_success"}) == "success" + assert operations._verdict_bucket( + {"triage_attack_verdict": "attack", "triage_attack_success": "success"} + ) == "success" + + +def test_soc_alert_verdict_uses_model_triage_instead_of_raw_attack_result(): + operations = _load_alert_operations() + + assert operations._verdict_bucket( + { + "attack_verdict": "attack_success", + "attack_success": True, + "threat_result": "success", + "triage_attack_verdict": "non_attack", + "triage_attack_success": "unknown", + } + ) == "benign" + assert operations._verdict_bucket( + { + "attack_verdict": "attack_success", + "attack_success": True, + "threat_result": "success", + "triage_attack_verdict": "attack", + "triage_attack_success": "failed", + } + ) == "failed" + assert operations._verdict_bucket( + { + "attack_verdict": "attack_failed", + "attack_success": False, + "threat_result": "failed", + "triage_attack_verdict": "attack", + "triage_attack_success": "success", + } + ) == "success" + assert operations._verdict_bucket( + {"triage_attack_verdict": "attack", "triage_attack_success": "unknown"} + ) == "attack" + assert operations._verdict_bucket( + {"triage_attack_verdict": "unknown", "triage_attack_success": "unknown"} + ) == "unknown" + assert operations._triage_attack_success( + {"triage_attack_verdict": "non_attack", "triage_attack_success": "success"} + ) == "unknown" + assert operations._verdict_bucket( + {"attack_verdict": "attack_success", "attack_success": True, "threat_result": "success"} + ) == "unknown" + assert operations._verdict_bucket({"threat_result": "success"}) == "unknown" def test_soc_dashboard_activity_exposes_live_denoise_workflow_progress(tmp_path: Path): @@ -438,6 +601,1003 @@ def test_soc_dashboard_activity_exposes_live_denoise_workflow_progress(tmp_path: } +def test_soc_dashboard_activity_exposes_triage_workflow_link_context(tmp_path: Path): + workflow_db = tmp_path / "workflow.db" + now_ms = int(datetime.now().timestamp() * 1000) + with sqlite3.connect(workflow_db) as conn: + conn.execute( + """ + CREATE TABLE workflow_executions ( + id TEXT PRIMARY KEY, + workflow_id TEXT NOT NULL, + status TEXT NOT NULL, + input_params TEXT NOT NULL DEFAULT '{}', + output_results TEXT NOT NULL DEFAULT '{}', + started_at INTEGER NOT NULL, + payload TEXT NOT NULL DEFAULT '{}' + ) + """ + ) + conn.execute( + """ + INSERT INTO workflow_executions + (id, workflow_id, status, input_params, output_results, started_at, payload) + VALUES (?, ?, ?, ?, ?, ?, ?) + """, + ( + "triage-running", + "stream_alert_triage", + "running", + json.dumps({"input_date": "2026-07-23"}), + json.dumps( + { + "enriched_alerts_with_triage": [ + {"threat_name": "远程命令执行攻击"} + ] + }, + ensure_ascii=False, + ), + now_ms, + json.dumps({"sessionId": "session-triage-1", "messageId": "msg-triage-1"}), + ), + ) + conn.commit() + + handlers = _load_dashboard_handlers() + handlers.WORKFLOW_DB = workflow_db + handlers.DEFAULT_SQLITE_DB = tmp_path / "missing-soc.db" + handlers._activity_pruned_at = float("inf") + + payload = handlers._get_activity({"bootstrap": "latest"}) + triage_events = [ + event for event in payload["workflowEvents"] + if event["workflowId"] == "stream_alert_triage" + ] + + assert len(triage_events) == 1 + triage = triage_events[0] + assert triage["stage"] == "triage" + assert triage["status"] == "running" + assert triage["alert"]["threatName"] == "远程命令执行攻击" + assert triage["sessionId"] == "session-triage-1" + assert triage["messageId"] == "msg-triage-1" + + +def test_soc_dashboard_activity_tolerates_empty_soc_db_with_workflow_events(tmp_path: Path): + soc_db = tmp_path / "soc.db" + soc_db.touch() + workflow_db = tmp_path / "workflow.db" + now_ms = int(datetime.now().timestamp() * 1000) + with sqlite3.connect(workflow_db) as conn: + conn.execute( + """ + CREATE TABLE workflow_executions ( + id TEXT PRIMARY KEY, + workflow_id TEXT NOT NULL, + status TEXT NOT NULL, + input_params TEXT NOT NULL DEFAULT '{}', + output_results TEXT NOT NULL DEFAULT '{}', + started_at INTEGER NOT NULL, + payload TEXT NOT NULL DEFAULT '{}' + ) + """ + ) + conn.execute( + "INSERT INTO workflow_executions VALUES (?, ?, ?, ?, ?, ?, ?)", + ( + "triage-running", + "stream_alert_triage", + "running", + "{}", + json.dumps({"triage_results": [{"alert_name": "Mock 研判"}]}, ensure_ascii=False), + now_ms, + json.dumps({"sessionId": "session-1", "messageId": "message-1"}), + ), + ) + conn.commit() + + handlers = _load_dashboard_handlers() + handlers.DEFAULT_SQLITE_DB = soc_db + handlers.WORKFLOW_DB = workflow_db + handlers._activity_pruned_at = float("inf") + + payload = handlers._get_activity({"bootstrap": "latest"}) + + assert "error" not in payload + assert payload["workflowEvents"][0]["alert"]["threatName"] == "Mock 研判" + assert payload["workflowEvents"][0]["sessionId"] == "session-1" + + +def test_soc_dashboard_task_center_summarizes_tasks_and_workflows(tmp_path: Path): + tasks_db = tmp_path / "tasks.db" + today_at_1100 = datetime.now().astimezone().replace( + hour=11, + minute=0, + second=0, + microsecond=0, + ) + today_at_1105 = today_at_1100.replace(minute=5) + today_at_1110 = today_at_1100.replace(minute=10) + tomorrow_at_1205 = (today_at_1100 + timedelta(days=1)).replace(hour=12, minute=5) + yesterday_at_1100 = today_at_1100 - timedelta(days=1) + yesterday_at_1105 = today_at_1105 - timedelta(days=1) + today_ms = int(today_at_1100.timestamp() * 1000) + yesterday_ms = int(yesterday_at_1100.timestamp() * 1000) + with sqlite3.connect(tasks_db) as conn: + conn.execute( + """ + CREATE TABLE task_schedulers ( + id TEXT PRIMARY KEY, + title TEXT NOT NULL, + mode TEXT NOT NULL, + status TEXT NOT NULL, + trigger TEXT NOT NULL, + execution_mode TEXT NOT NULL, + workflow_id TEXT, + updated_at TEXT NOT NULL + ) + """ + ) + conn.execute( + """ + CREATE TABLE task_executions ( + id TEXT PRIMARY KEY, + scheduler_id TEXT NOT NULL, + status TEXT NOT NULL, + queued_at TEXT, + started_at TEXT, + completed_at TEXT, + updated_at TEXT, + created_at TEXT, + session_id TEXT + ) + """ + ) + conn.execute( + "INSERT INTO task_schedulers VALUES (?, ?, ?, ?, ?, ?, ?, ?)", + ( + "sched-1", + "每日告警巡检", + "cron", + "active", + json.dumps({"cron": "*/5 * * * *", "nextRun": tomorrow_at_1205.isoformat(timespec="seconds")}), + "workflow", + "stream_alert_denoise", + today_at_1100.isoformat(timespec="seconds"), + ), + ) + conn.executemany( + "INSERT INTO task_executions VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)", + [ + ( + "exec-1", + "sched-1", + "completed", + yesterday_at_1100.isoformat(timespec="seconds"), + (yesterday_at_1100 + timedelta(seconds=1)).isoformat(timespec="seconds"), + (yesterday_at_1100 + timedelta(seconds=10)).isoformat(timespec="seconds"), + (yesterday_at_1100 + timedelta(seconds=10)).isoformat(timespec="seconds"), + yesterday_at_1100.isoformat(timespec="seconds"), + "session-a", + ), + ( + "exec-2", + "sched-1", + "failed", + yesterday_at_1105.isoformat(timespec="seconds"), + (yesterday_at_1105 + timedelta(seconds=1)).isoformat(timespec="seconds"), + (yesterday_at_1105 + timedelta(seconds=10)).isoformat(timespec="seconds"), + (yesterday_at_1105 + timedelta(seconds=10)).isoformat(timespec="seconds"), + yesterday_at_1105.isoformat(timespec="seconds"), + "session-a", + ), + ( + "exec-3", + "sched-1", + "running", + today_at_1110.isoformat(timespec="seconds"), + (today_at_1110 + timedelta(seconds=1)).isoformat(timespec="seconds"), + None, + (today_at_1110 + timedelta(seconds=1)).isoformat(timespec="seconds"), + today_at_1110.isoformat(timespec="seconds"), + "session-b", + ), + ], + ) + conn.commit() + + workflow_db = tmp_path / "workflow.db" + with sqlite3.connect(workflow_db) as conn: + conn.execute( + """ + CREATE TABLE workflow_stats ( + workflow_id TEXT PRIMARY KEY, + call_count INTEGER NOT NULL, + success_count INTEGER NOT NULL, + error_count INTEGER NOT NULL, + updated_at INTEGER NOT NULL + ) + """ + ) + conn.execute( + """ + CREATE TABLE workflow_executions ( + id TEXT PRIMARY KEY, + workflow_id TEXT NOT NULL, + status TEXT NOT NULL, + current_phase TEXT, + current_node_id TEXT, + current_node_type TEXT, + current_step_index INTEGER, + step_count INTEGER NOT NULL DEFAULT 0, + input_params TEXT NOT NULL DEFAULT '{}', + output_results TEXT NOT NULL DEFAULT '{}', + error_message TEXT, + started_at INTEGER NOT NULL, + finished_at INTEGER, + updated_at INTEGER, + payload TEXT NOT NULL DEFAULT '{}' + ) + """ + ) + conn.execute( + """ + CREATE TABLE workflow_configs ( + workflow_id TEXT NOT NULL, + kind TEXT NOT NULL, + version INTEGER, + config TEXT NOT NULL, + updated_at INTEGER, + PRIMARY KEY (workflow_id, kind) + ) + """ + ) + conn.executemany( + "INSERT INTO workflow_stats VALUES (?, ?, ?, ?, ?)", + [ + ("custom_workflow", 2, 1, 1, 1784775000000), + ("stream_alert_denoise", 10, 9, 1, today_ms), + ("stream_alert_triage", 1, 0, 0, yesterday_ms + 5000), + ("e6d5581a-b105-4c75-a102-1d8e6c97e1c1", 1, 0, 0, 1784775700000), + ], + ) + conn.executemany( + """ + INSERT INTO workflow_executions ( + id, workflow_id, status, current_phase, current_node_id, current_node_type, + current_step_index, step_count, input_params, output_results, error_message, + started_at, finished_at, updated_at, payload + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + [ + ( + "wf-1", + "custom_workflow", + "error", + "error", + "node-1", + "python", + 1, + 1, + "{}", + "{}", + "boom", + yesterday_ms, + yesterday_ms + 1000, + yesterday_ms + 1000, + "{}", + ), + ( + "wf-2", + "stream_alert_denoise", + "success", + "success", + "dedup", + "python", + 1, + 1, + json.dumps({"alerts": [{"threat_name": "SQL注入攻击"}]}), + json.dumps({"unique_alerts": [{"threat_name": "SQL注入攻击"}]}), + None, + today_ms, + today_ms + 1000, + today_ms + 1000, + "{}", + ), + ( + "wf-triage", + "stream_alert_triage", + "running", + "running", + "concurrent_triage", + "python", + 2, + 1, + json.dumps({"input_date": yesterday_at_1100.date().isoformat()}), + json.dumps( + { + "enriched_alerts_with_triage": [ + { + "threat_name": "远程命令执行攻击", + "report_title": "高危RCE研判", + } + ] + }, + ensure_ascii=False, + ), + None, + yesterday_ms + 5000, + None, + yesterday_ms + 6000, + json.dumps({"sessionId": "session-triage-1", "messageId": "msg-triage-1"}), + ), + ( + "wf-dynamic", + "e6d5581a-b105-4c75-a102-1d8e6c97e1c1", + "running", + "running", + "node-1", + "python", + 1, + 0, + "{}", + "{}", + None, + 1784775700000, + None, + 1784775700000, + "{}", + ), + ], + ) + conn.executemany( + "INSERT INTO workflow_configs VALUES (?, ?, ?, ?, ?)", + [ + ( + "custom_workflow", + "workflow.integration-config", + 1, + json.dumps({"workflow": {"id": "custom_workflow", "name": "自定义处置流"}}), + 1784775000000, + ), + ( + "stream_alert_denoise", + "workflow.integration-config", + 1, + json.dumps({"workflow": {"id": "stream_alert_denoise", "name": "流式告警降噪"}}), + 1784775600000, + ), + ( + "stream_alert_triage", + "workflow_poller_config", + 1, + json.dumps({"enabled": True, "timeoutSeconds": 604800}), + yesterday_ms + 6000, + ), + ], + ) + conn.commit() + + handlers = _load_dashboard_handlers() + handlers.TASK_DB = tasks_db + handlers.WORKFLOW_DB = workflow_db + + payload = handlers._get_task_center() + + assert payload["sessionCount"] == 2 + assert payload["scheduledTasks"] == [ + { + "id": "sched-1", + "name": "每日告警巡检", + "mode": "cron", + "status": "active", + "executionMode": "workflow", + "workflowId": "stream_alert_denoise", + "executionCount": 3, + "todayExecutionCount": 1, + "successCount": 1, + "successRate": 0.3333, + "activeCount": 1, + "lastStatus": "running", + "lastRunAt": (today_at_1110 + timedelta(seconds=1)).isoformat(timespec="seconds"), + "nextRunAt": tomorrow_at_1205.isoformat(timespec="seconds"), + "cron": "*/5 * * * *", + "cronDescription": "", + } + ] + assert payload["scheduledExecutionCount"] == 3 + assert payload["scheduledTodayExecutionCount"] == 1 + assert payload["workflowExecutionCount"] == 13 + assert payload["workflowTodayExecutionCount"] == 1 + assert [workflow["id"] for workflow in payload["workflows"]] == [ + "stream_alert_triage", + "stream_alert_denoise", + "custom_workflow", + ] + assert [workflow["name"] for workflow in payload["workflows"]] == [ + "告警研判工作流", + "告警降噪工作流", + "自定义处置流", + ] + triage = payload["workflows"][0] + assert triage["latestExecutionHash"] == "wf-triage" + assert triage["latestAlertName"] == "远程命令执行攻击" + assert triage["progressLabel"] == "第 2/4 步" + assert triage["progressPercent"] == 0.5 + assert triage["sessionId"] == "session-triage-1" + assert triage["messageId"] == "msg-triage-1" + denoise = payload["workflows"][1] + assert denoise["executionCount"] == 10 + assert denoise["todayExecutionCount"] == 1 + assert denoise["successCount"] == 9 + assert denoise["successRate"] == 0.9 + assert denoise["latestExecutionHash"] == "wf-2" + assert denoise["latestAlertName"] == "SQL注入攻击" + assert denoise["progressLabel"] == "已完成" + + +def test_soc_dashboard_task_center_hides_mock_pinned_workflows_by_default(tmp_path: Path): + workflow_db = tmp_path / "workflow.db" + with sqlite3.connect(workflow_db) as conn: + conn.execute( + """ + CREATE TABLE workflow_stats ( + workflow_id TEXT PRIMARY KEY, + call_count INTEGER NOT NULL, + success_count INTEGER NOT NULL, + error_count INTEGER NOT NULL, + updated_at INTEGER + ) + """ + ) + conn.execute( + """ + CREATE TABLE workflow_executions ( + id TEXT PRIMARY KEY, + workflow_id TEXT NOT NULL, + status TEXT NOT NULL, + started_at INTEGER NOT NULL + ) + """ + ) + conn.commit() + + handlers = _load_dashboard_handlers() + handlers.TASK_DB = tmp_path / "missing-tasks.db" + handlers.WORKFLOW_DB = workflow_db + + payload = handlers._get_task_center() + assert payload["workflows"] == [] + + mock_payload = handlers._get_task_center(include_mock=True) + assert [workflow["id"] for workflow in mock_payload["workflows"]] == [ + "stream_alert_denoise", + "stream_alert_triage", + ] + + +def test_soc_dashboard_token_usage_uses_grouped_cached_reads(tmp_path: Path, monkeypatch): + handlers = _load_dashboard_handlers() + usage_db = tmp_path / "flocks.db" + handlers.USAGE_DB = usage_db + handlers._token_usage_cache.update({"updatedAt": 0.0, "mtimeNs": 0, "value": None}) + + now = datetime.now().astimezone().replace(hour=10, minute=0, second=0, microsecond=0) + yesterday = now - timedelta(days=1) + with sqlite3.connect(usage_db) as conn: + conn.execute("CREATE TABLE usage_records (created_at TEXT NOT NULL, total_tokens INTEGER NOT NULL)") + conn.executemany( + "INSERT INTO usage_records VALUES (?, ?)", + [ + (handlers._usage_iso(now), 10), + (handlers._usage_iso(now + timedelta(minutes=5)), 15), + (handlers._usage_iso(yesterday), 7), + ], + ) + conn.commit() + + first = handlers._read_token_usage() + + assert first["totalTokens"] == 32 + assert first["todayTokens"] == 25 + assert first["todayRequests"] == 2 + assert first["dailySeries"][-1] == 25 + assert first["dailySeries"][-2] == 7 + + def fail_connect(*args, **kwargs): + raise AssertionError("token usage should be served from cache") + + monkeypatch.setattr(handlers.sqlite3, "connect", fail_connect) + assert handlers._read_token_usage() == first + + +def test_soc_dashboard_task_center_supports_legacy_workflow_execution_schema(tmp_path: Path): + workflow_db = tmp_path / "workflow.db" + today_at_1100 = datetime.now().astimezone().replace( + hour=11, + minute=0, + second=0, + microsecond=0, + ) + today_ms = int(today_at_1100.timestamp() * 1000) + with sqlite3.connect(workflow_db) as conn: + conn.execute( + """ + CREATE TABLE workflow_stats ( + workflow_id TEXT PRIMARY KEY, + call_count INTEGER NOT NULL, + success_count INTEGER NOT NULL, + error_count INTEGER NOT NULL + ) + """ + ) + conn.execute( + """ + CREATE TABLE workflow_executions ( + id TEXT PRIMARY KEY, + workflow_id TEXT NOT NULL, + status TEXT NOT NULL, + started_at INTEGER NOT NULL + ) + """ + ) + conn.execute( + "INSERT INTO workflow_stats VALUES (?, ?, ?, ?)", + ("stream_alert_denoise", 1, 0, 0), + ) + conn.execute( + "INSERT INTO workflow_executions VALUES (?, ?, ?, ?)", + ("wf-running", "stream_alert_denoise", "running", today_ms), + ) + conn.commit() + + handlers = _load_dashboard_handlers() + handlers.TASK_DB = tmp_path / "missing-tasks.db" + handlers.WORKFLOW_DB = workflow_db + + payload = handlers._get_task_center() + + denoise = next( + workflow for workflow in payload["workflows"] if workflow["id"] == "stream_alert_denoise" + ) + assert payload["workflowExecutionCount"] == 1 + assert payload["workflowTodayExecutionCount"] == 1 + assert denoise["executionCount"] == 1 + assert denoise["todayExecutionCount"] == 1 + assert denoise["activeCount"] == 1 + assert denoise["latestExecutionHash"] == "wf-running" + assert denoise["lastRunAt"] == today_ms + + +def test_soc_dashboard_task_center_ignores_disabled_trigger_workflow_runs(tmp_path: Path): + workflow_db = tmp_path / "workflow.db" + now_ms = int(datetime.now().astimezone().timestamp() * 1000) + with sqlite3.connect(workflow_db) as conn: + conn.execute( + """ + CREATE TABLE workflow_stats ( + workflow_id TEXT PRIMARY KEY, + call_count INTEGER NOT NULL, + success_count INTEGER NOT NULL, + error_count INTEGER NOT NULL, + updated_at INTEGER NOT NULL + ) + """ + ) + conn.execute( + """ + CREATE TABLE workflow_executions ( + id TEXT PRIMARY KEY, + workflow_id TEXT NOT NULL, + status TEXT NOT NULL, + current_step_index INTEGER, + step_count INTEGER, + started_at INTEGER NOT NULL, + updated_at INTEGER + ) + """ + ) + conn.execute( + """ + CREATE TABLE workflow_configs ( + workflow_id TEXT NOT NULL, + kind TEXT NOT NULL, + version INTEGER, + config TEXT NOT NULL, + updated_at INTEGER, + PRIMARY KEY (workflow_id, kind) + ) + """ + ) + conn.execute( + "INSERT INTO workflow_stats VALUES (?, ?, ?, ?, ?)", + ("stream_alert_triage", 1, 0, 0, now_ms), + ) + conn.execute( + "INSERT INTO workflow_executions VALUES (?, ?, ?, ?, ?, ?, ?)", + ( + "wf-disabled-running", + "stream_alert_triage", + "running", + 2, + 4, + now_ms - 60_000, + now_ms - 60_000, + ), + ) + conn.execute( + "INSERT INTO workflow_configs VALUES (?, ?, ?, ?, ?)", + ( + "stream_alert_triage", + "workflow_poller_config", + 1, + json.dumps({"enabled": False, "timeoutSeconds": 7200}), + now_ms, + ), + ) + conn.commit() + + handlers = _load_dashboard_handlers() + handlers.TASK_DB = tmp_path / "missing-tasks.db" + handlers.WORKFLOW_DB = workflow_db + + payload = handlers._get_task_center() + + triage = next( + workflow for workflow in payload["workflows"] if workflow["id"] == "stream_alert_triage" + ) + assert triage["executionCount"] == 1 + assert triage["activeCount"] == 0 + assert triage["lastStatus"] == "disabled" + assert triage["progressLabel"] == "已关闭" + assert triage["progressPercent"] == 0 + + +def test_soc_dashboard_task_center_uses_trigger_runtime_window_for_active_workflows(tmp_path: Path): + workflow_db = tmp_path / "workflow.db" + now_ms = int(datetime.now().astimezone().timestamp() * 1000) + old_ms = now_ms - 3 * 60 * 60 * 1000 + with sqlite3.connect(workflow_db) as conn: + conn.execute( + """ + CREATE TABLE workflow_stats ( + workflow_id TEXT PRIMARY KEY, + call_count INTEGER NOT NULL, + success_count INTEGER NOT NULL, + error_count INTEGER NOT NULL, + updated_at INTEGER NOT NULL + ) + """ + ) + conn.execute( + """ + CREATE TABLE workflow_executions ( + id TEXT PRIMARY KEY, + workflow_id TEXT NOT NULL, + status TEXT NOT NULL, + current_step_index INTEGER, + step_count INTEGER, + started_at INTEGER NOT NULL, + updated_at INTEGER + ) + """ + ) + conn.execute( + """ + CREATE TABLE workflow_configs ( + workflow_id TEXT NOT NULL, + kind TEXT NOT NULL, + version INTEGER, + config TEXT NOT NULL, + updated_at INTEGER, + PRIMARY KEY (workflow_id, kind) + ) + """ + ) + conn.executemany( + "INSERT INTO workflow_stats VALUES (?, ?, ?, ?, ?)", + [ + ("stream_alert_denoise", 1, 0, 0, now_ms), + ("stream_alert_triage", 1, 0, 0, old_ms), + ], + ) + conn.executemany( + "INSERT INTO workflow_executions VALUES (?, ?, ?, ?, ?, ?, ?)", + [ + ( + "wf-recent-running", + "stream_alert_denoise", + "running", + 1, + 4, + now_ms - 30_000, + now_ms - 30_000, + ), + ( + "wf-stale-running", + "stream_alert_triage", + "running", + 3, + 4, + old_ms, + old_ms, + ), + ], + ) + conn.executemany( + "INSERT INTO workflow_configs VALUES (?, ?, ?, ?, ?)", + [ + ( + "stream_alert_denoise", + "workflow_poller_config", + 1, + json.dumps({"enabled": True, "timeoutSeconds": 7200}), + now_ms, + ), + ( + "stream_alert_triage", + "workflow_poller_config", + 1, + json.dumps({"enabled": True, "timeoutSeconds": 7200}), + old_ms, + ), + ], + ) + conn.commit() + + handlers = _load_dashboard_handlers() + handlers.TASK_DB = tmp_path / "missing-tasks.db" + handlers.WORKFLOW_DB = workflow_db + + payload = handlers._get_task_center() + + denoise = next( + workflow for workflow in payload["workflows"] if workflow["id"] == "stream_alert_denoise" + ) + triage = next( + workflow for workflow in payload["workflows"] if workflow["id"] == "stream_alert_triage" + ) + assert denoise["activeCount"] == 1 + assert denoise["lastStatus"] == "running" + assert denoise["progressLabel"].startswith("第 ") + assert triage["activeCount"] == 0 + assert triage["lastStatus"] == "stale" + assert triage["progressLabel"] == "已停止" + + +def test_soc_dashboard_task_center_marks_stale_running_without_trigger_config(tmp_path: Path): + workflow_db = tmp_path / "workflow.db" + now_ms = int(datetime.now().astimezone().timestamp() * 1000) + old_ms = now_ms - 3 * 60 * 60 * 1000 + with sqlite3.connect(workflow_db) as conn: + conn.execute( + """ + CREATE TABLE workflow_stats ( + workflow_id TEXT PRIMARY KEY, + call_count INTEGER NOT NULL, + success_count INTEGER NOT NULL, + error_count INTEGER NOT NULL, + updated_at INTEGER NOT NULL + ) + """ + ) + conn.execute( + """ + CREATE TABLE workflow_executions ( + id TEXT PRIMARY KEY, + workflow_id TEXT NOT NULL, + status TEXT NOT NULL, + current_step_index INTEGER, + step_count INTEGER, + started_at INTEGER NOT NULL, + updated_at INTEGER + ) + """ + ) + conn.execute( + "INSERT INTO workflow_stats VALUES (?, ?, ?, ?, ?)", + ("custom_workflow", 1, 0, 0, old_ms), + ) + conn.execute( + "INSERT INTO workflow_executions VALUES (?, ?, ?, ?, ?, ?, ?)", + ("wf-stale-manual", "custom_workflow", "running", 1, 2, old_ms, old_ms), + ) + conn.commit() + + handlers = _load_dashboard_handlers() + handlers.TASK_DB = tmp_path / "missing-tasks.db" + handlers.WORKFLOW_DB = workflow_db + + payload = handlers._get_task_center() + + workflow = next(item for item in payload["workflows"] if item["id"] == "custom_workflow") + assert workflow["activeCount"] == 0 + assert workflow["lastStatus"] == "stale" + assert workflow["progressLabel"] == "已停止" + + +def test_soc_dashboard_task_center_orders_dynamic_rows(tmp_path: Path): + now = datetime.now().astimezone().replace(microsecond=0) + older = now - timedelta(hours=3) + recent = now - timedelta(minutes=5) + future = now + timedelta(hours=2) + older_ms = int(older.timestamp() * 1000) + recent_ms = int(recent.timestamp() * 1000) + + tasks_db = tmp_path / "tasks.db" + with sqlite3.connect(tasks_db) as conn: + conn.execute( + """ + CREATE TABLE task_schedulers ( + id TEXT PRIMARY KEY, + title TEXT NOT NULL, + mode TEXT NOT NULL, + status TEXT NOT NULL, + trigger TEXT NOT NULL, + execution_mode TEXT NOT NULL, + workflow_id TEXT, + updated_at TEXT NOT NULL + ) + """ + ) + conn.execute( + """ + CREATE TABLE task_executions ( + id TEXT PRIMARY KEY, + scheduler_id TEXT NOT NULL, + status TEXT NOT NULL, + queued_at TEXT, + started_at TEXT, + completed_at TEXT, + updated_at TEXT, + created_at TEXT, + session_id TEXT + ) + """ + ) + conn.executemany( + "INSERT INTO task_schedulers VALUES (?, ?, 'cron', 'active', ?, 'agent', '', ?)", + [ + ("sched-active", "正在执行任务", "{}", older.isoformat()), + ("sched-recent", "最近完成任务", "{}", recent.isoformat()), + ("sched-popular", "高频历史任务", "{}", older.isoformat()), + ( + "sched-future", + "未来待执行任务", + json.dumps({"nextRun": future.isoformat(timespec="seconds")}), + now.isoformat(), + ), + ], + ) + task_exec_rows = [ + ( + "exec-active", + "sched-active", + "running", + older.isoformat(timespec="seconds"), + older.isoformat(timespec="seconds"), + None, + older.isoformat(timespec="seconds"), + older.isoformat(timespec="seconds"), + "s-active", + ), + ( + "exec-recent", + "sched-recent", + "completed", + recent.isoformat(timespec="seconds"), + recent.isoformat(timespec="seconds"), + (recent + timedelta(seconds=3)).isoformat(timespec="seconds"), + (recent + timedelta(seconds=3)).isoformat(timespec="seconds"), + recent.isoformat(timespec="seconds"), + "s-recent", + ), + ] + for index in range(5): + when = older - timedelta(minutes=index) + task_exec_rows.append( + ( + f"exec-popular-{index}", + "sched-popular", + "completed", + when.isoformat(timespec="seconds"), + when.isoformat(timespec="seconds"), + (when + timedelta(seconds=2)).isoformat(timespec="seconds"), + (when + timedelta(seconds=2)).isoformat(timespec="seconds"), + when.isoformat(timespec="seconds"), + "s-popular", + ) + ) + conn.executemany( + "INSERT INTO task_executions VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)", + task_exec_rows, + ) + conn.commit() + + workflow_db = tmp_path / "workflow.db" + with sqlite3.connect(workflow_db) as conn: + conn.execute( + """ + CREATE TABLE workflow_stats ( + workflow_id TEXT PRIMARY KEY, + call_count INTEGER NOT NULL, + success_count INTEGER NOT NULL, + error_count INTEGER NOT NULL, + updated_at INTEGER NOT NULL + ) + """ + ) + conn.execute( + """ + CREATE TABLE workflow_executions ( + id TEXT PRIMARY KEY, + workflow_id TEXT NOT NULL, + status TEXT NOT NULL, + started_at INTEGER NOT NULL, + finished_at INTEGER, + updated_at INTEGER + ) + """ + ) + conn.execute( + """ + CREATE TABLE workflow_configs ( + workflow_id TEXT NOT NULL, + kind TEXT NOT NULL, + version INTEGER, + config TEXT NOT NULL, + updated_at INTEGER, + PRIMARY KEY (workflow_id, kind) + ) + """ + ) + conn.executemany( + "INSERT INTO workflow_stats VALUES (?, ?, ?, ?, ?)", + [ + ("wf-active", 1, 0, 0, older_ms), + ("wf-recent", 1, 1, 0, recent_ms), + ("wf-popular", 6, 6, 0, older_ms), + ], + ) + conn.executemany( + "INSERT INTO workflow_executions VALUES (?, ?, ?, ?, ?, ?)", + [ + ("run-active", "wf-active", "running", older_ms, None, older_ms), + ("run-recent", "wf-recent", "success", recent_ms, recent_ms + 1000, recent_ms + 1000), + ("run-popular", "wf-popular", "success", older_ms, older_ms + 1000, older_ms + 1000), + ], + ) + conn.execute( + "INSERT INTO workflow_configs VALUES (?, ?, ?, ?, ?)", + ( + "wf-active", + "workflow_poller_config", + 1, + json.dumps({"enabled": True, "timeoutSeconds": 14400}), + older_ms, + ), + ) + conn.commit() + + handlers = _load_dashboard_handlers() + handlers.TASK_DB = tasks_db + handlers.WORKFLOW_DB = workflow_db + + payload = handlers._get_task_center() + + assert [task["id"] for task in payload["scheduledTasks"][:4]] == [ + "sched-active", + "sched-recent", + "sched-popular", + "sched-future", + ] + assert [workflow["id"] for workflow in payload["workflows"][:3]] == [ + "wf-active", + "wf-recent", + "wf-popular", + ] + + def test_soc_dashboard_workflow_progress_marks_unavailable_database(tmp_path: Path): handlers = _load_dashboard_handlers() handlers.WORKFLOW_DB = tmp_path / "missing-workflow.db" diff --git a/tests/ingest/test_kafka_manager.py b/tests/ingest/test_kafka_manager.py index 111cdf34e..a3c2b0432 100644 --- a/tests/ingest/test_kafka_manager.py +++ b/tests/ingest/test_kafka_manager.py @@ -16,6 +16,8 @@ import asyncio import json import sys +import threading +import time from types import SimpleNamespace from unittest.mock import AsyncMock @@ -221,6 +223,31 @@ async def _fake_trigger(workflow_id, workflow_json, msg, input_key, producer=Non assert captured == [{"ok": True}] +def test_trigger_concurrency_config_is_honored_with_safety_caps() -> None: + trigger = TriggerDefinition.model_validate( + { + "id": "kafka-limited", + "type": "kafka", + "concurrency": {"maxParallel": 2, "queueSize": 40}, + } + ) + assert kafka_manager._worker_count_for_trigger(trigger) == 2 + assert kafka_manager._queue_size_for_trigger(trigger) == 40 + + oversized = TriggerDefinition.model_validate( + { + "id": "kafka-capped", + "type": "kafka", + "concurrency": {"maxParallel": 999, "queueSize": 999_999}, + } + ) + assert ( + kafka_manager._worker_count_for_trigger(oversized) + == kafka_manager._MAX_CONCURRENT_EXECUTIONS + ) + assert kafka_manager._queue_size_for_trigger(oversized) == kafka_manager._MAX_QUEUE_SIZE + + @pytest.mark.asyncio async def test_stop_workflow_cancels_worker_pool() -> None: """``stop_workflow`` must cancel and drain the worker pool cleanly.""" @@ -260,6 +287,153 @@ async def _noop_trigger(*args, **kwargs): # noqa: ANN001 assert manager._status[workflow_id]["state"] == "stopped" +@pytest.mark.asyncio +async def test_stop_workflow_signals_running_thread( + monkeypatch: pytest.MonkeyPatch, +) -> None: + manager = kafka_manager.KafkaManager() + workflow_id = "test-wf-running-stop" + queue: asyncio.Queue = asyncio.Queue(maxsize=2) + abort = asyncio.Event() + started = threading.Event() + stopped = threading.Event() + trigger = TriggerDefinition.model_validate( + {"id": "kafka-default", "type": "kafka", "mapping": {"message": "$.body"}} + ) + + async def _fake_create_execution_record(*args, **kwargs): # noqa: ANN001, ANN002, ANN003 + return {"id": "exec-running-stop"} + + async def _fake_record_execution_result(*args, **kwargs): # noqa: ANN001, ANN002, ANN003 + return None + + def _fake_run_workflow(**kwargs): # noqa: ANN003 + started.set() + while not kwargs["cancel"](): + time.sleep(0.01) + stopped.set() + return SimpleNamespace( + status="CANCELLED", + error=None, + outputs={}, + history=[], + last_node_id=None, + steps=0, + ) + + monkeypatch.setattr(kafka_manager, "create_execution_record", _fake_create_execution_record) + monkeypatch.setattr(kafka_manager, "record_execution_result", _fake_record_execution_result) + monkeypatch.setattr(kafka_manager, "run_workflow", _fake_run_workflow) + + manager._queues[workflow_id] = queue + manager._abort_events[workflow_id] = abort + generation_cancel_event = threading.Event() + manager._generation_cancel_events[workflow_id] = generation_cancel_event + queue.put_nowait({"message": "demo"}) + worker = asyncio.create_task( + manager._worker_loop( + workflow_id, + {}, + trigger, + {}, + queue, + abort, + "topic-a", + generation_cancel_event, + ), + name="running-stop-worker", + ) + manager._worker_pools[workflow_id] = [worker] + + assert await asyncio.wait_for(asyncio.to_thread(started.wait, 1.0), timeout=2.0) + await manager.stop_workflow(workflow_id) + + assert stopped.is_set() + assert worker.done() + assert workflow_id not in manager._generation_cancel_events + + +@pytest.mark.asyncio +async def test_stop_cancels_run_that_is_still_creating_execution_record( + monkeypatch: pytest.MonkeyPatch, +) -> None: + manager = kafka_manager.KafkaManager() + workflow_id = "test-kafka-late-run-registration" + queue: asyncio.Queue = asyncio.Queue(maxsize=1) + abort = asyncio.Event() + generation_cancel_event = threading.Event() + create_started = asyncio.Event() + allow_create = asyncio.Event() + cancel_state_seen: list[bool] = [] + trigger = TriggerDefinition.model_validate( + {"id": "kafka-default", "type": "kafka", "mapping": {"message": "$.body"}} + ) + + async def _slow_create(*args, **kwargs): # noqa: ANN001, ANN002, ANN003 + create_started.set() + await allow_create.wait() + return {"id": "exec-late-run"} + + async def _record(*args, **kwargs): # noqa: ANN001, ANN002, ANN003 + return None + + def _run_workflow(**kwargs): # noqa: ANN003 + cancel_state_seen.append(kwargs["cancel"]()) + return SimpleNamespace( + status="CANCELLED", + error=None, + outputs={}, + history=[], + last_node_id=None, + steps=0, + ) + + monkeypatch.setattr(kafka_manager, "create_execution_record", _slow_create) + monkeypatch.setattr(kafka_manager, "record_execution_result", _record) + monkeypatch.setattr(kafka_manager, "run_workflow", _run_workflow) + manager._queues[workflow_id] = queue + manager._abort_events[workflow_id] = abort + manager._generation_cancel_events[workflow_id] = generation_cancel_event + queue.put_nowait({"message": "demo"}) + worker = asyncio.create_task( + manager._worker_loop( + workflow_id, + {}, + trigger, + {}, + queue, + abort, + "topic-a", + generation_cancel_event, + ) + ) + manager._worker_pools[workflow_id] = [worker] + + await create_started.wait() + stop_task = asyncio.create_task(manager.stop_workflow(workflow_id)) + await asyncio.sleep(0) + allow_create.set() + await stop_task + + assert cancel_state_seen == [True] + assert worker.done() + + +@pytest.mark.asyncio +async def test_restart_refuses_to_overlap_draining_generation() -> None: + manager = kafka_manager.KafkaManager() + workflow_id = "test-kafka-draining-restart" + release = asyncio.Event() + draining = asyncio.create_task(release.wait()) + manager._draining_workers[workflow_id] = {draining} + try: + status = await manager.restart_workflow(workflow_id) + assert status == {"state": "failed", "error": "previous_workers_still_draining"} + finally: + release.set() + await draining + + @pytest.mark.asyncio async def test_restart_disabled_config_reports_stopped(monkeypatch: pytest.MonkeyPatch) -> None: """A disabled (or missing) config must leave the consumer ``stopped``.""" @@ -436,6 +610,8 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 assert captured_input_params["kafka_message"]["alarmData"]["chars"] == 50_000 assert captured_run_kwargs["run_id"] == "exec-compact" assert captured_run_kwargs["execution_profile"] == "high_frequency" + assert callable(captured_run_kwargs["cancel"]) + assert captured_run_kwargs["cancel"]() is False assert callable(captured_run_kwargs["on_step_complete"]) assert captured_exec_data["outputResults"] == { "_enriched_alerts_count": 1, diff --git a/tests/ingest/test_syslog_manager_backpressure.py b/tests/ingest/test_syslog_manager_backpressure.py index 5309e4be8..1982caf84 100644 --- a/tests/ingest/test_syslog_manager_backpressure.py +++ b/tests/ingest/test_syslog_manager_backpressure.py @@ -17,6 +17,8 @@ from __future__ import annotations import asyncio +import threading +import time from types import SimpleNamespace from unittest.mock import AsyncMock @@ -128,6 +130,31 @@ async def test_bounded_queue_drops_excess_on_full() -> None: assert queue.qsize() == 4 +def test_trigger_concurrency_config_is_honored_with_safety_caps() -> None: + trigger = TriggerDefinition.model_validate( + { + "id": "syslog-limited", + "type": "syslog", + "concurrency": {"maxParallel": 3, "queueSize": 25}, + } + ) + assert syslog_manager._worker_count_for_trigger(trigger) == 3 + assert syslog_manager._queue_size_for_trigger(trigger) == 25 + + oversized = TriggerDefinition.model_validate( + { + "id": "syslog-capped", + "type": "syslog", + "concurrency": {"maxParallel": 999, "queueSize": 999_999}, + } + ) + assert ( + syslog_manager._worker_count_for_trigger(oversized) + == syslog_manager._MAX_CONCURRENT_EXECUTIONS + ) + assert syslog_manager._queue_size_for_trigger(oversized) == syslog_manager._MAX_QUEUE_SIZE + + @pytest.mark.asyncio async def test_stop_workflow_cancels_worker_pool() -> None: """``stop_workflow`` must cancel and drain the worker pool cleanly. @@ -174,6 +201,149 @@ async def _noop_trigger(*args, **kwargs): # noqa: ANN001, D401 assert manager._listener_status[workflow_id]["state"] == "stopped" +@pytest.mark.asyncio +async def test_stop_workflow_signals_running_thread( + monkeypatch: pytest.MonkeyPatch, +) -> None: + manager = syslog_manager.SyslogManager() + workflow_id = "test-wf-running-stop" + queue: asyncio.Queue = asyncio.Queue(maxsize=2) + abort = asyncio.Event() + started = threading.Event() + stopped = threading.Event() + trigger = TriggerDefinition.model_validate( + {"id": "syslog-default", "type": "syslog", "mapping": {"message": "$.body"}} + ) + + async def _fake_create_execution_record(*args, **kwargs): # noqa: ANN001, ANN002, ANN003 + return {"id": "exec-running-stop"} + + async def _fake_record_execution_result(*args, **kwargs): # noqa: ANN001, ANN002, ANN003 + return None + + def _fake_run_workflow(**kwargs): # noqa: ANN003 + started.set() + while not kwargs["cancel"](): + time.sleep(0.01) + stopped.set() + return SimpleNamespace( + status="CANCELLED", + error=None, + outputs={}, + history=[], + last_node_id=None, + steps=0, + ) + + monkeypatch.setattr(syslog_manager, "create_execution_record", _fake_create_execution_record) + monkeypatch.setattr(syslog_manager, "record_execution_result", _fake_record_execution_result) + monkeypatch.setattr(syslog_manager, "run_workflow", _fake_run_workflow) + + manager._queues[workflow_id] = queue + manager._abort_events[workflow_id] = abort + generation_cancel_event = threading.Event() + manager._generation_cancel_events[workflow_id] = generation_cancel_event + queue.put_nowait({"message": "demo"}) + worker = asyncio.create_task( + manager._worker_loop( + workflow_id, + {}, + trigger, + queue, + abort, + generation_cancel_event, + ), + name="running-stop-worker", + ) + manager._worker_pools[workflow_id] = [worker] + + assert await asyncio.wait_for(asyncio.to_thread(started.wait, 1.0), timeout=2.0) + await manager.stop_workflow(workflow_id) + + assert stopped.is_set() + assert worker.done() + assert workflow_id not in manager._generation_cancel_events + + +@pytest.mark.asyncio +async def test_stop_cancels_run_that_is_still_creating_execution_record( + monkeypatch: pytest.MonkeyPatch, +) -> None: + manager = syslog_manager.SyslogManager() + workflow_id = "test-wf-late-run-registration" + queue: asyncio.Queue = asyncio.Queue(maxsize=1) + abort = asyncio.Event() + generation_cancel_event = threading.Event() + create_started = asyncio.Event() + allow_create = asyncio.Event() + cancel_state_seen: list[bool] = [] + trigger = TriggerDefinition.model_validate( + {"id": "syslog-default", "type": "syslog", "mapping": {"message": "$.body"}} + ) + + async def _slow_create(*args, **kwargs): # noqa: ANN001, ANN002, ANN003 + create_started.set() + await allow_create.wait() + return {"id": "exec-late-run"} + + async def _record(*args, **kwargs): # noqa: ANN001, ANN002, ANN003 + return None + + def _run_workflow(**kwargs): # noqa: ANN003 + cancel_state_seen.append(kwargs["cancel"]()) + return SimpleNamespace( + status="CANCELLED", + error=None, + outputs={}, + history=[], + last_node_id=None, + steps=0, + ) + + monkeypatch.setattr(syslog_manager, "create_execution_record", _slow_create) + monkeypatch.setattr(syslog_manager, "record_execution_result", _record) + monkeypatch.setattr(syslog_manager, "run_workflow", _run_workflow) + manager._queues[workflow_id] = queue + manager._abort_events[workflow_id] = abort + manager._generation_cancel_events[workflow_id] = generation_cancel_event + queue.put_nowait({"message": "demo"}) + worker = asyncio.create_task( + manager._worker_loop( + workflow_id, + {}, + trigger, + queue, + abort, + generation_cancel_event, + ) + ) + manager._worker_pools[workflow_id] = [worker] + + await create_started.wait() + stop_task = asyncio.create_task(manager.stop_workflow(workflow_id)) + await asyncio.sleep(0) + allow_create.set() + await stop_task + + assert cancel_state_seen == [True] + assert worker.done() + + +@pytest.mark.asyncio +async def test_restart_refuses_to_overlap_draining_generation() -> None: + manager = syslog_manager.SyslogManager() + workflow_id = "test-wf-draining-restart" + release = asyncio.Event() + draining = asyncio.create_task(release.wait()) + manager._draining_workers[workflow_id] = {draining} + try: + status = await manager.restart_workflow(workflow_id) + assert status == {"state": "failed", "error": "previous_workers_still_draining"} + finally: + release.set() + await draining + + @pytest.mark.asyncio async def test_trigger_workflow_applies_mapping_and_filter( monkeypatch: pytest.MonkeyPatch, @@ -252,6 +422,8 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 assert captured_run_kwargs["run_id"] == "exec-syslog" assert captured_run_kwargs["execution_profile"] == "high_frequency" assert captured_run_kwargs["tool_context"] is trigger_tool_context.context + assert callable(captured_run_kwargs["cancel"]) + assert captured_run_kwargs["cancel"]() is False assert callable(captured_run_kwargs["on_step_complete"]) trigger_tool_context.builder.assert_awaited_once_with( workflow_id="wf-syslog", diff --git a/tests/integration/test_capability_awareness.py b/tests/integration/test_capability_awareness.py index 0cf15ecc0..3a18d2a74 100644 --- a/tests/integration/test_capability_awareness.py +++ b/tests/integration/test_capability_awareness.py @@ -226,7 +226,7 @@ def test_workflow_section_injected_into_agent_prompt(self, tmp_path): builder_code = textwrap.dedent(""" from flocks.agent.prompt_utils import build_workflows_section - def inject(agent_info, available_agents, tools, skills, categories, workflows=None): + def inject(agent_info, available_agents, tools, skills, workflows=None): section = build_workflows_section(workflows or []) agent_info.prompt = f"## Capability Context\\n{section}" """) @@ -250,7 +250,7 @@ def inject(agent_info, available_agents, tools, skills, categories, workflows=No source="project", ) ] - inject_dynamic_prompts({"wf_test_agent": agent}, [], [], [], [], workflows) + inject_dynamic_prompts({"wf_test_agent": agent}, [], [], [], workflows) assert "my_wf" in agent.prompt assert "My integration workflow" in agent.prompt finally: diff --git a/tests/provider/test_chinese_providers.py b/tests/provider/test_chinese_providers.py index 1245b52fe..f5a625370 100644 --- a/tests/provider/test_chinese_providers.py +++ b/tests/provider/test_chinese_providers.py @@ -6,6 +6,7 @@ get_provider_default_url, get_provider_meta, get_provider_model_definitions, + get_raw_catalog, list_catalog_provider_ids, ) @@ -282,10 +283,11 @@ def test_threatbook_cn_llm_catalog(self): "qwen3-max", "kimi-k2.6", "deepseek-v4-flash", + "deepseek-v4-flash-0731", } + assert models[0].id == "deepseek-v4-flash-0731" kimi_code = next(m for m in models if m.id == "kimi-k2.7-code") - assert models[0].id == "kimi-k2.7-code" assert kimi_code.capabilities.supports_vision is True assert kimi_code.capabilities.supports_reasoning is True assert kimi_code.capabilities.interleaved["field"] == "reasoning_content" @@ -311,6 +313,19 @@ def test_threatbook_cn_llm_catalog(self): assert flash_cn.pricing.currency == "CNY" assert flash_cn.limits.context_window == 1000000 assert flash_cn.limits.max_output_tokens == 384000 + raw_models = get_raw_catalog()["threatbook-cn-llm"]["models"] + assert raw_models["deepseek-v4-flash-0731"] == { + **raw_models["deepseek-v4-flash"], + "name": "deepseek-v4-flash-0731", + "limits": { + **raw_models["deepseek-v4-flash"]["limits"], + "max_input_tokens": 1000000, + }, + "pricing": { + **raw_models["deepseek-v4-flash"]["pricing"], + "cache_read": 0.2, + }, + } kimi = next(m for m in models if m.id == "kimi-k2.6") assert kimi.capabilities.supports_vision is True @@ -339,10 +354,11 @@ def test_threatbook_io_llm_catalog(self): "qwen3.6-plus", "qwen3-max", "deepseek-v4-flash", + "deepseek-v4-flash-0731", } + assert models[0].id == "deepseek-v4-flash-0731" kimi_code = next(m for m in models if m.id == "kimi-k2.7-code") - assert models[0].id == "kimi-k2.7-code" assert kimi_code.capabilities.supports_vision is True assert kimi_code.capabilities.supports_reasoning is True assert kimi_code.capabilities.interleaved["field"] == "reasoning_content" @@ -363,6 +379,19 @@ def test_threatbook_io_llm_catalog(self): assert flash_io.pricing.currency == "CNY" assert flash_io.limits.context_window == 1000000 assert flash_io.limits.max_output_tokens == 384000 + raw_models = get_raw_catalog()["threatbook-io-llm"]["models"] + assert raw_models["deepseek-v4-flash-0731"] == { + **raw_models["deepseek-v4-flash"], + "name": "deepseek-v4-flash-0731", + "limits": { + **raw_models["deepseek-v4-flash"]["limits"], + "max_input_tokens": 1000000, + }, + "pricing": { + **raw_models["deepseek-v4-flash"]["pricing"], + "cache_read": 0.2, + }, + } m27 = next(m for m in models if m.id == "minimax-m2.7") assert m27.capabilities.interleaved["field"] == "reasoning_details" diff --git a/tests/provider/test_model_management.py b/tests/provider/test_model_management.py index 7f9b3ed66..3c132eb1c 100644 --- a/tests/provider/test_model_management.py +++ b/tests/provider/test_model_management.py @@ -279,6 +279,43 @@ def test_config_override_false_removes_reasoning_feature(self): assert overridden.capabilities.supports_reasoning is False assert ModelFeature.REASONING not in overridden.capabilities.features + def test_config_override_preserves_cache_read_pricing(self): + from flocks.provider.provider import BaseProvider, ModelCapabilities, ModelInfo + + catalog_def = ModelDefinition( + id="dummy-model", + name="Dummy Model", + provider_id="dummy", + pricing=PriceConfig(input=1.0, output=2.0), + ) + model = ModelInfo( + id="dummy-model", + name="Dummy Model", + provider_id="dummy", + capabilities=ModelCapabilities(), + pricing={ + "input": 1.0, + "output": 2.0, + "cache_read": 0.2, + "currency": "CNY", + }, + ) + model._explicit_keys = { + "input_price", + "output_price", + "cache_read_price", + "currency", + } + + overridden = BaseProvider("dummy", "Dummy")._apply_config_overrides( + catalog_def, + model, + ) + + assert overridden.pricing is not None + assert overridden.pricing.cache_read == 0.2 + assert overridden.pricing.currency == "CNY" + def test_configure_from_credential(self): from flocks.provider.provider import BaseProvider diff --git a/tests/provider/test_test_credentials.py b/tests/provider/test_test_credentials.py index 783692ded..e1687bf96 100644 --- a/tests/provider/test_test_credentials.py +++ b/tests/provider/test_test_credentials.py @@ -303,23 +303,23 @@ async def test_onesec_service_prefers_threat_probe_and_uses_enum_action(self): @pytest.mark.asyncio async def test_service_prefers_login_probe_over_action_dispatch_tool(self): - """When a parameter-free login tool exists (e.g. qingteng_login), it - should be tried before action-dispatch tools whose handler-side validation - requires extra fields beyond the JSON schema (e.g. qingteng_assets which - needs `resource` + `os_type` for `assets.refresh`). + """When a parameter-free login tool exists, it + should be tried before an action-dispatch tool whose handler-side validation + requires extra fields beyond the JSON schema (e.g. a grouped assets tool + that needs `resource` + `os_type` for `assets.refresh`). """ from flocks.server.routes.provider import test_provider_credentials login_tool = ToolInfo( - name="qingteng_login", - description="Qingteng login probe", + name="example_login", + description="Example login probe", category=ToolCategory.CUSTOM, parameters=[], requires_confirmation=False, ) assets_tool = ToolInfo( - name="qingteng_assets", - description="Qingteng assets dispatch", + name="example_assets", + description="Example assets dispatch", category=ToolCategory.CUSTOM, parameters=[ ToolParameter( @@ -340,27 +340,28 @@ async def test_service_prefers_login_probe_over_action_dispatch_tool(self): patch(_PATCH_SECRET_MGR, return_value=mock_secrets), patch(_PATCH_PROVIDER) as mock_provider_cls, patch(_PATCH_TOOL_REGISTRY) as mock_tr, - patch(_PATCH_TOOL_SOURCE, return_value=("api", "qingteng")), + patch(_PATCH_TOOL_SOURCE, return_value=("api", "example")), ): mock_provider_cls._ensure_initialized = MagicMock() mock_provider_cls.apply_config = AsyncMock() mock_provider_cls.get.return_value = None mock_tr.init = MagicMock() + mock_tr.init_async = AsyncMock() mock_tr.list_tools.return_value = [assets_tool, login_tool] mock_tr._dynamic_tools_by_module = { - "flocks.tool.generated.qingteng": ["qingteng_assets", "qingteng_login"], + "flocks.tool.generated.example": ["example_assets", "example_login"], } mock_tr.execute = AsyncMock(return_value=ToolResult( success=True, - output={"jwt": "fake", "signKey": "fake", "comId": "1"}, + output={"authenticated": True}, )) - result = await test_provider_credentials("qingteng") + result = await test_provider_credentials("example") assert result["success"] is True, result - assert result["tool_tested"] == "qingteng_login" - mock_tr.execute.assert_awaited_once_with(tool_name="qingteng_login") + assert result["tool_tested"] == "example_login" + mock_tr.execute.assert_awaited_once_with(tool_name="example_login") @pytest.mark.asyncio async def test_declared_manifest_probe_is_used_before_heuristic_tool_selection(self): diff --git a/tests/server/routes/test_custom_provider_routes.py b/tests/server/routes/test_custom_provider_routes.py index 05e96a00b..d2b8c60cd 100644 --- a/tests/server/routes/test_custom_provider_routes.py +++ b/tests/server/routes/test_custom_provider_routes.py @@ -66,6 +66,7 @@ async def test_create_custom_model_accepts_string_currency( "supports_reasoning": False, "input_price": 0.0, "output_price": 0.0, + "cache_read_price": 0.2, "currency": "USD", }, ) @@ -74,14 +75,66 @@ async def test_create_custom_model_accepts_string_currency( data = response.json() assert data["provider_id"] == "custom-tb-inner" assert data["model_id"] == "minimax:MiniMax-M2.7" + assert data["cache_read_price"] == 0.2 assert data["currency"] == "USD" raw = ConfigWriter.get_provider_raw("custom-tb-inner") assert raw is not None + assert raw["models"]["minimax:MiniMax-M2.7"]["cache_read_price"] == 0.2 assert raw["models"]["minimax:MiniMax-M2.7"]["currency"] == "USD" + assert Provider._models["minimax:MiniMax-M2.7"].pricing["cache_read"] == 0.2 assert Provider._models["minimax:MiniMax-M2.7"].pricing["currency"] == "USD" +@pytest.mark.asyncio +async def test_update_custom_model_preserves_omitted_cache_price( + client: AsyncClient, + temp_custom_provider_project, + monkeypatch: pytest.MonkeyPatch, +): + from flocks.config.config_writer import ConfigWriter + from flocks.provider.provider import Provider + + monkeypatch.setattr(Provider, "_models", Provider._models.copy()) + payload = { + "model_id": "cached-model", + "name": "Cached Model", + "input_price": 1.0, + "output_price": 2.0, + "cache_read_price": 0.2, + "currency": "CNY", + } + created = await client.post( + "/api/custom/models/custom-tb-inner", + json=payload, + ) + assert created.status_code == 201, created.text + + payload.pop("cache_read_price") + updated = await client.post( + "/api/custom/models/custom-tb-inner", + json=payload, + ) + + assert updated.status_code == 201, updated.text + assert updated.json()["cache_read_price"] == 0.2 + raw = ConfigWriter.get_provider_raw("custom-tb-inner") + assert raw is not None + assert raw["models"]["cached-model"]["cache_read_price"] == 0.2 + assert Provider._models["cached-model"].pricing["cache_read"] == 0.2 + + cleared = await client.post( + "/api/custom/models/custom-tb-inner", + json={**payload, "cache_read_price": None}, + ) + assert cleared.status_code == 201, cleared.text + assert cleared.json()["cache_read_price"] is None + raw = ConfigWriter.get_provider_raw("custom-tb-inner") + assert raw is not None + assert raw["models"]["cached-model"]["cache_read_price"] is None + assert "cache_read" not in Provider._models["cached-model"].pricing + + @pytest.mark.asyncio async def test_create_custom_model_defaults_reasoning_on( client: AsyncClient, diff --git a/tests/server/routes/test_custom_provider_runtime.py b/tests/server/routes/test_custom_provider_runtime.py index 0a66aeea2..5093265c4 100644 --- a/tests/server/routes/test_custom_provider_runtime.py +++ b/tests/server/routes/test_custom_provider_runtime.py @@ -42,6 +42,7 @@ class DummyProvider: supports_reasoning=True, input_price=0.0, output_price=0.0, + cache_read_price=0.2, currency="USD", ) @@ -54,7 +55,12 @@ class DummyProvider: saved = Provider._models[body.model_id] assert saved.capabilities.supports_reasoning is True - assert saved.pricing == {"input": 0.0, "output": 0.0, "currency": "USD"} + assert saved.pricing == { + "input": 0.0, + "output": 0.0, + "cache_read": 0.2, + "currency": "USD", + } assert provider._custom_models[0].pricing["currency"] == "USD" assert provider._config_models[0].capabilities.supports_reasoning is True finally: diff --git a/tests/server/routes/test_onboarding_routes.py b/tests/server/routes/test_onboarding_routes.py index 53675b541..763831b7f 100644 --- a/tests/server/routes/test_onboarding_routes.py +++ b/tests/server/routes/test_onboarding_routes.py @@ -198,12 +198,12 @@ async def fake_test_mcp(region: str, api_key: str): class TestOnboardingApplyRoutes: - def test_threatbook_region_presets_use_kimi_k27_code(self): + def test_threatbook_region_presets_use_deepseek_v4_flash_0731(self): assert onboarding_routes.ONBOARDING_REGION_PRESETS["cn"]["threatbook_default_model_id"] == ( - "kimi-k2.7-code" + "deepseek-v4-flash-0731" ) assert onboarding_routes.ONBOARDING_REGION_PRESETS["global"]["threatbook_default_model_id"] == ( - "kimi-k2.7-code" + "deepseek-v4-flash-0731" ) def test_ensure_threatbook_mcp_config_uses_explicit_secret_reference( @@ -310,14 +310,14 @@ async def fake_set_default_model(model_type, body): data = resp.json() assert data["success"] is True assert data["default_model"]["provider_id"] == "threatbook-cn-llm" - assert data["default_model"]["model_id"] == "kimi-k2.7-code" + assert data["default_model"]["model_id"] == "deepseek-v4-flash-0731" assert ("provider", "threatbook-cn-llm") in calls assert ("service", "threatbook-cn") in calls assert ("service_enabled", "threatbook-cn") in calls assert ("ensure_mcp", "cn") in calls assert ("mcp_credentials", "threatbook_mcp") in calls assert ("mcp_connect", "threatbook_mcp") in calls - assert ("default_model", "threatbook-cn-llm", "kimi-k2.7-code") in calls + assert ("default_model", "threatbook-cn-llm", "deepseek-v4-flash-0731") in calls @pytest.mark.asyncio async def test_apply_returns_400_when_validation_fails( diff --git a/tests/server/routes/test_provider_model_bootstrap.py b/tests/server/routes/test_provider_model_bootstrap.py index 5ab24fae7..1dbef2c83 100644 --- a/tests/server/routes/test_provider_model_bootstrap.py +++ b/tests/server/routes/test_provider_model_bootstrap.py @@ -11,6 +11,30 @@ class TestThreatBookProviderModelBootstrap: + @pytest.mark.asyncio + async def test_catalog_exposes_deepseek_v4_flash_0731_metadata(self): + result = await provider_routes.get_provider_catalog() + providers = {provider["id"]: provider for provider in result["providers"]} + + for provider_id in ("threatbook-cn-llm", "threatbook-io-llm"): + models = { + model["id"]: model + for model in providers[provider_id]["models"] + } + model = models["deepseek-v4-flash-0731"] + assert model["limits"] == { + "context_window": 1000000, + "max_input_tokens": 1000000, + "max_output_tokens": 384000, + } + assert model["pricing"] == { + "input": 1.0, + "output": 2.0, + "cache_read": 0.2, + "cache_write": None, + "currency": "CNY", + } + @pytest.mark.asyncio async def test_set_provider_credentials_bootstraps_kimi_k26_from_catalog( self, monkeypatch: pytest.MonkeyPatch diff --git a/tests/server/routes/test_session_routes.py b/tests/server/routes/test_session_routes.py index e20a70d07..4c9766552 100644 --- a/tests/server/routes/test_session_routes.py +++ b/tests/server/routes/test_session_routes.py @@ -24,6 +24,8 @@ from flocks.session.message import ( Message, MessageRole, + PartTime, + ReasoningPart, ToolPart, ToolStateError, ToolStateRunning, @@ -405,6 +407,8 @@ async def test_list_sessions_light_manager_filters_and_omits_heavy_fields(self, "title", "time", "category", + "channelID", + "channelChatType", "status", "parentID", "provider", @@ -423,6 +427,157 @@ async def test_list_sessions_light_manager_filters_and_omits_heavy_fields(self, assert "goal" not in row assert "summary" not in row + @pytest.mark.asyncio + async def test_manager_list_includes_channel_metadata_and_legacy_title( + self, + client: AsyncClient, + monkeypatch: pytest.MonkeyPatch, + ): + from flocks.channel.inbound.session_binding import SessionBindingService + + session_resp = await client.post("/api/session", json={"title": "[Wecom] room-1"}) + session_id = session_resp.json()["id"] + await Message.create( + session_id=session_id, + role=MessageRole.USER, + content="你是谁", + ) + list_bindings_mock = AsyncMock(return_value=[SimpleNamespace( + session_id=session_id, + channel_id="wecom", + chat_id="room-1", + chat_type="group", + )]) + monkeypatch.setattr( + SessionBindingService, + "list_bindings", + list_bindings_mock, + ) + + response = await client.get( + "/api/session", + params={"view": "list", "manager": "true", "roots": "true", "limit": "100"}, + ) + + assert response.status_code == status.HTTP_200_OK + row = next(item for item in response.json() if item["id"] == session_id) + assert row["channelID"] == "wecom" + assert row["channelChatType"] == "group" + assert row["title"] == "[Wecom] 你是谁" + assert session_id in list_bindings_mock.await_args.kwargs["session_ids"] + + @pytest.mark.asyncio + async def test_manager_list_recognizes_legacy_direct_title_with_sender_name( + self, + client: AsyncClient, + monkeypatch: pytest.MonkeyPatch, + ): + from flocks.channel.inbound.session_binding import SessionBindingService + + session_resp = await client.post( + "/api/session", + json={"title": "[Telegram] DM — Alice"}, + ) + session_id = session_resp.json()["id"] + await Message.create( + session_id=session_id, + role=MessageRole.USER, + content="你是谁", + ) + monkeypatch.setattr( + SessionBindingService, + "list_bindings", + AsyncMock(return_value=[SimpleNamespace( + session_id=session_id, + channel_id="telegram", + chat_id="12345", + chat_type="direct", + )]), + ) + + response = await client.get( + "/api/session", + params={"view": "list", "manager": "true", "roots": "true"}, + ) + + assert response.status_code == status.HTTP_200_OK + row = next(item for item in response.json() if item["id"] == session_id) + assert row["title"] == "[Telegram] 你是谁" + + @pytest.mark.asyncio + async def test_manager_search_matches_derived_channel_title_before_pagination( + self, + client: AsyncClient, + monkeypatch: pytest.MonkeyPatch, + ): + from flocks.channel.inbound.session_binding import SessionBindingService + + session_resp = await client.post("/api/session", json={"title": "[Wecom] room-1"}) + session_id = session_resp.json()["id"] + await Message.create( + session_id=session_id, + role=MessageRole.USER, + content="你是谁", + ) + decoy_resp = await client.post( + "/api/session", + json={"title": "Unrelated newer session"}, + ) + decoy_id = decoy_resp.json()["id"] + list_bindings_mock = AsyncMock(return_value=[SimpleNamespace( + session_id=session_id, + channel_id="wecom", + chat_id="room-1", + chat_type="group", + )]) + monkeypatch.setattr( + SessionBindingService, + "list_bindings", + list_bindings_mock, + ) + + response = await client.get( + "/api/session", + params={ + "view": "list", + "manager": "true", + "roots": "true", + "search": "你是谁", + "limit": "1", + "offset": "0", + }, + ) + + assert response.status_code == status.HTTP_200_OK + assert [item["id"] for item in response.json()] == [session_id] + assert response.json()[0]["title"] == "[Wecom] 你是谁" + queried_session_ids = list_bindings_mock.await_args.kwargs["session_ids"] + assert session_id in queried_session_ids + assert decoy_id in queried_session_ids + + @pytest.mark.asyncio + async def test_channel_binding_lookup_batches_large_session_lists( + self, + monkeypatch: pytest.MonkeyPatch, + ): + from flocks.channel.inbound.session_binding import SessionBindingService + from flocks.server.routes.session import _latest_channel_bindings + + list_bindings_mock = AsyncMock(return_value=[]) + monkeypatch.setattr( + SessionBindingService, + "list_bindings", + list_bindings_mock, + ) + + await _latest_channel_bindings([f"ses-{index}" for index in range(1001)]) + + assert list_bindings_mock.await_count == 3 + assert all( + len(call.kwargs["session_ids"]) <= 500 + for call in list_bindings_mock.await_args_list + ) + @pytest.mark.asyncio async def test_archive_hides_session_preserves_history_and_restores_tree(self, client: AsyncClient): parent_resp = await client.post("/api/session", json={"title": "Archive Parent"}) @@ -1226,6 +1381,32 @@ async def test_send_message_noReply(self, client: AsyncClient, session_id: str): for m in messages ) + @pytest.mark.asyncio + async def test_list_messages_preserves_reasoning_part_time( + self, + client: AsyncClient, + session_id: str, + ): + """Reloaded message history retains timing needed by the process summary.""" + message = await Message.create(session_id, MessageRole.ASSISTANT, "") + part = ReasoningPart( + id="part_timed_reasoning", + sessionID=session_id, + messageID=message.id, + text="Inspect the request.", + time=PartTime(start=1_000, end=4_500), + ) + await Message.store_part(session_id, message.id, part) + + response = await client.get(f"/api/session/{session_id}/message") + + assert response.status_code == status.HTTP_200_OK + messages = response.json() + reloaded_message = next(item for item in messages if item["info"]["id"] == message.id) + reloaded_part = next(item for item in reloaded_message["parts"] if item["id"] == part.id) + assert reloaded_part["time"]["start"] == 1_000 + assert reloaded_part["time"]["end"] == 4_500 + @pytest.mark.asyncio async def test_list_messages_keeps_running_tool_when_session_busy( self, diff --git a/tests/server/test_input_dispatcher.py b/tests/server/test_input_dispatcher.py index 3a2c767cc..084c1820d 100644 --- a/tests/server/test_input_dispatcher.py +++ b/tests/server/test_input_dispatcher.py @@ -68,6 +68,92 @@ async def test_direct_command_uses_direct_response(self): assert direct and "Available / commands:" in direct[0] assert not llm + @pytest.mark.asyncio + async def test_webui_direct_response_is_excluded_from_model_context( + self, + monkeypatch, + ): + from flocks.server.routes import session as session_routes + + created_messages = AsyncMock() + context_usage_update = AsyncMock() + published_events = [] + + async def fake_persist( + _session_id, + write, + *, + expected_generation=None, + ): + del expected_generation + return await write() + + async def fake_dispatch(event, sink): + await sink.publish_direct_response( + event, + "Available Tools\n" + ("x" * 16_000), + ) + + async def fake_publish(event_type, properties): + published_events.append((event_type, properties)) + + monkeypatch.setattr( + "flocks.input.dispatcher.dispatch_user_input", + fake_dispatch, + ) + monkeypatch.setattr( + "flocks.session.message.Message.create", + created_messages, + ) + monkeypatch.setattr( + "flocks.server.routes.event.publish_event", + fake_publish, + ) + monkeypatch.setattr( + session_routes, + "_persist_active_session_write", + fake_persist, + ) + monkeypatch.setattr( + session_routes, + "_publish_context_usage_update", + context_usage_update, + ) + monkeypatch.setattr( + session_routes.Session, + "lifecycle_generation", + lambda _session_id: 0, + ) + event = UserInputEvent( + source_type="webui", + sessionID="ses_direct_context", + text="/tools", + display_text="/tools", + ) + + await session_routes._dispatch_sse_input( + "ses_direct_context", + SimpleNamespace(id="ses_direct_context"), + event, + "/tmp/project", + ) + + assert created_messages.await_count == 2 + assert all( + call.kwargs["ignored"] is True + for call in created_messages.await_args_list + ) + text_part_events = [ + properties["part"] + for event_type, properties in published_events + if event_type == "message.part.updated" + ] + assert len(text_part_events) == 2 + assert all(part["ignored"] is True for part in text_part_events) + context_usage_update.assert_awaited_once() + assert "provider_id" not in context_usage_update.await_args.kwargs + assert "model_id" not in context_usage_update.await_args.kwargs + @pytest.mark.asyncio async def test_clear_uses_history_callback_without_direct_response(self): direct = [] diff --git a/tests/session/test_context_usage.py b/tests/session/test_context_usage.py index 47d0e7347..2bc9db71a 100644 --- a/tests/session/test_context_usage.py +++ b/tests/session/test_context_usage.py @@ -185,6 +185,47 @@ async def test_context_usage_does_not_scan_archived_history(context_usage_mocks) assert snapshot.excluded_segments == [] +@pytest.mark.asyncio +async def test_context_usage_ignores_ignored_text_parts(context_usage_mocks): + msg = _message("command-output", tokens=None) + context_usage_mocks["active"] = [msg] + context_usage_mocks["all"] = [msg] + context_usage_mocks["estimate"] = 0 + context_usage_mocks["parts"] = { + "command-output": [ + SimpleNamespace( + type="text", + text="Available Tools\n" + ("x" * 16_000), + ignored=True, + ), + ], + } + + snapshot = await context_usage.build_context_usage_snapshot("sess-1") + + assert snapshot.used_tokens == 0 + assert all(segment.key != "conversation" for segment in snapshot.segments) + + +def test_context_usage_uses_last_real_model_before_builtin_command(): + real_model = _message( + "assistant-1", + provider_id="openai", + model_id="gpt-5.6-luna", + ) + command_output = _message( + "command-output", + created=200, + provider_id="builtin", + model_id="command", + ) + + assert context_usage._resolve_message_model([real_model, command_output]) == ( + "openai", + "gpt-5.6-luna", + ) + + @pytest.mark.asyncio async def test_context_usage_splits_tool_parts_from_conversation(context_usage_mocks): msg = _message("assistant-1", tokens=None) diff --git a/tests/session/test_execution_mode.py b/tests/session/test_execution_mode.py index 005042d3c..0ef33c1de 100644 --- a/tests/session/test_execution_mode.py +++ b/tests/session/test_execution_mode.py @@ -116,7 +116,6 @@ def test_plan_delegation_only_allows_explore_and_librarian(tool_name) -> None: for arguments in ( {"subagent_type": "general"}, - {"category": "quick"}, {"session_id": "child-session"}, {}, ): diff --git a/tests/session/test_prompt_tokens.py b/tests/session/test_prompt_tokens.py index fddedf832..cf130cc66 100644 --- a/tests/session/test_prompt_tokens.py +++ b/tests/session/test_prompt_tokens.py @@ -460,3 +460,33 @@ async def test_message_content_is_counted(self, monkeypatch: pytest.MonkeyPatch) "ses_x", messages, ) assert result == 100 + + @pytest.mark.asyncio + async def test_ignored_text_part_is_not_counted( + self, + monkeypatch: pytest.MonkeyPatch, + ): + from flocks.session import message as message_mod + + async def _ignored_parts(message_id, session_id): # noqa: ARG001 + return [ + SimpleNamespace( + type="text", + text="Available Tools\n" + ("x" * 400), + ignored=True, + ), + ] + + monkeypatch.setattr( + message_mod.Message, + "parts", + staticmethod(_ignored_parts), + ) + messages = [{"id": "ignored-command-output", "content": ""}] + + result = await SessionPrompt.estimate_full_context_tokens( + "ses_x", + messages, + ) + + assert result == 0 diff --git a/tests/session/test_runner_step.py b/tests/session/test_runner_step.py index 6109a1aaa..82fb6aaec 100644 --- a/tests/session/test_runner_step.py +++ b/tests/session/test_runner_step.py @@ -1530,6 +1530,27 @@ async def test_to_chat_messages_invalidates_shared_cache_when_message_parts_chan assert second_messages[1].content == "Error: Tool execution was interrupted" +@pytest.mark.asyncio +async def test_to_chat_messages_excludes_ignored_assistant_text(): + session = await Session.create( + project_id="test_runner_ignored_command_output", + directory="/tmp/runner-ignored-command-output", + ) + assistant_message = await Message.create( + session_id=session.id, + role=MessageRole.ASSISTANT, + content="Available Tools\n" + ("x" * 400), + providerID="builtin", + modelID="command", + ignored=True, + ) + runner = SessionRunner(session=session, static_cache={}) + + chat_messages = await runner._to_chat_messages([assistant_message], []) + + assert chat_messages == [] + + @pytest.mark.asyncio async def test_to_chat_messages_preserves_assistant_reasoning_for_replay(): session = await Session.create( diff --git a/tests/tool/test_delegate_task_compat.py b/tests/tool/test_delegate_task_compat.py index 1229aa1f9..59be3c16f 100644 --- a/tests/tool/test_delegate_task_compat.py +++ b/tests/tool/test_delegate_task_compat.py @@ -11,14 +11,27 @@ def _make_ctx() -> ToolContext: class TestDelegateTaskTolerance: - def test_delegate_task_schema_allows_omitting_optional_fields(self): - schema = ToolRegistry.get_schema("delegate_task") + def test_delegate_description_requires_material_delegation_benefit(self): + from flocks.tool.agent.delegate_task import DESCRIPTION + + assert "A specialized agent clearly matches the task" in DESCRIPTION + assert "You need to explore code in parallel" in DESCRIPTION + assert "Isolating research or noisy intermediate work" in DESCRIPTION + assert "Do not delegate trivial edits" in DESCRIPTION + assert "The task requires multiple steps or research" not in DESCRIPTION + + @pytest.mark.parametrize("tool_name", ["delegate_task", "task"]) + def test_delegate_schema_exposes_only_subagent_routing(self, tool_name): + schema = ToolRegistry.get_schema(tool_name) assert schema is not None assert "prompt" in schema.required + assert "subagent_type" in schema.properties + assert "category" not in schema.properties assert "load_skills" not in schema.required assert "description" not in schema.required assert "run_in_background" not in schema.properties - # Legacy batch shape is gone: tasks=[...] is no longer a public option. + assert "command" in schema.properties + assert "command" not in schema.required assert "tasks" not in schema.properties @pytest.mark.asyncio @@ -35,7 +48,6 @@ async def test_delegate_task_derives_description_and_ignores_blank_skills(self): child_session = SimpleNamespace(id="ses-child") with ( patch("flocks.tool.agent.delegate_task._find_completed_delegate", AsyncMock(return_value=None)), - patch("flocks.tool.agent.delegate_task.Config.get", AsyncMock(return_value=SimpleNamespace(categories=None))), patch("flocks.tool.agent.delegate_task.is_delegatable", return_value=True), patch("flocks.tool.agent.delegate_task.Skill.get", AsyncMock()) as skill_get, patch("flocks.tool.agent.delegate_task.Session.get_by_id", AsyncMock(return_value=parent_session)), @@ -66,7 +78,7 @@ async def test_delegate_task_derives_description_and_ignores_blank_skills(self): assert "task" not in denied_permissions @pytest.mark.asyncio - async def test_delegate_task_category_model_uses_runtime_override_without_pinning(self): + async def test_delegate_task_explicit_model_override_is_pinned(self): parent_session = SimpleNamespace( id="test-session", project_id="proj", @@ -74,20 +86,10 @@ async def test_delegate_task_category_model_uses_runtime_override_without_pinnin provider=None, model=None, ) - child_session = SimpleNamespace(id="ses-quick") - cfg = SimpleNamespace(categories={ - "quick": { - "model": "anthropic/claude-haiku-4-5", - "prompt_append": None, - } - }) + child_session = SimpleNamespace(id="ses-child") with patch("flocks.tool.agent.delegate_task._find_completed_delegate", AsyncMock(return_value=None)), \ - patch("flocks.tool.agent.delegate_task.Config.get", AsyncMock(return_value=cfg)), \ - patch("flocks.tool.agent.delegate_task._validate_category_model", return_value={ - "providerID": "anthropic", - "modelID": "claude-haiku-4-5", - }), \ + patch("flocks.tool.agent.delegate_task.is_delegatable", return_value=True), \ patch("flocks.tool.agent.delegate_task.Session.get_by_id", AsyncMock(return_value=parent_session)), \ patch("flocks.tool.agent.delegate_task.Session.create", AsyncMock(return_value=child_session)) as create_session, \ patch("flocks.tool.agent.delegate_task.Message.create", AsyncMock()), \ @@ -99,18 +101,92 @@ async def test_delegate_task_category_model_uses_runtime_override_without_pinnin result = await ToolRegistry.execute( "delegate_task", ctx=_make_ctx(), - category="quick", + subagent_type="asset-survey", + model="anthropic/claude-haiku-4-5", prompt="Summarize the diff", - description="quick task", + description="summary task", ) assert result.success is True assert create_session.await_args.kwargs["provider"] == "anthropic" assert create_session.await_args.kwargs["model"] == "claude-haiku-4-5" - assert create_session.await_args.kwargs["model_pinned"] is False + assert create_session.await_args.kwargs["model_pinned"] is True assert loop_run.await_args.kwargs["provider_id"] == "anthropic" assert loop_run.await_args.kwargs["model_id"] == "claude-haiku-4-5" + @pytest.mark.asyncio + @pytest.mark.parametrize("tool_name", ["delegate_task", "task"]) + async def test_delegate_tools_reject_removed_category_parameter(self, tool_name): + result = await ToolRegistry.execute( + tool_name, + ctx=_make_ctx(), + category="quick", + prompt="Summarize the diff", + ) + + assert result.success is False + assert "unknown parameters: category" in (result.error or "") + + @pytest.mark.asyncio + @pytest.mark.parametrize("tool_name", ["delegate_task", "task"]) + async def test_delegate_tools_accept_deprecated_command_parameter(self, tool_name): + result = await ToolRegistry.execute( + tool_name, + ctx=_make_ctx(), + command="legacy-tracking-command", + prompt="Summarize the diff", + ) + + assert result.success is False + assert "unknown parameters: command" not in (result.error or "") + assert "subagent_type or session_id" in (result.error or "") + + @pytest.mark.asyncio + async def test_delegate_task_requires_subagent_for_new_task(self): + result = await ToolRegistry.execute( + "delegate_task", + ctx=_make_ctx(), + prompt="Summarize the diff", + ) + + assert result.success is False + assert "subagent_type or session_id" in (result.error or "") + + @pytest.mark.asyncio + async def test_delegate_task_rejects_restricted_subagent(self): + with patch("flocks.tool.agent.delegate_task.is_delegatable", return_value=False): + result = await ToolRegistry.execute( + "delegate_task", + ctx=_make_ctx(), + subagent_type="restricted-agent", + prompt="Summarize the diff", + ) + + assert result.success is False + assert 'Agent "restricted-agent" cannot be delegated to' in (result.error or "") + + @pytest.mark.asyncio + async def test_delegate_task_rejects_unknown_subagent(self, monkeypatch): + from flocks.agent import registry + from flocks.agent.agent import AgentInfo + + known_agent = AgentInfo( + name="known-agent", + mode="subagent", + delegatable=True, + ) + monkeypatch.setattr(registry, "_agents_ref", {"known-agent": known_agent}) + + result = await ToolRegistry.execute( + "delegate_task", + ctx=_make_ctx(), + subagent_type="explroe", + prompt="Summarize the diff", + ) + + assert result.success is False + assert 'Agent "explroe" cannot be delegated to' in (result.error or "") + @pytest.mark.asyncio async def test_delegate_task_rejects_background_execution(self): # run_in_background is not in the public schema, so the registry @@ -163,8 +239,7 @@ async def test_delegate_task_sync_continue_fails_when_last_message_missing(self) agent="asset-survey", ) - with patch("flocks.tool.agent.delegate_task.Config.get", AsyncMock(return_value=SimpleNamespace(categories=None))), \ - patch("flocks.tool.agent.delegate_task.Session.get_by_id", AsyncMock(return_value=session)), \ + with patch("flocks.tool.agent.delegate_task.Session.get_by_id", AsyncMock(return_value=session)), \ patch("flocks.tool.agent.delegate_task.Message.create", AsyncMock()), \ patch("flocks.tool.agent.delegate_task.SessionLoop.run", AsyncMock(return_value=SimpleNamespace( action="stop", diff --git a/tests/tool/test_probe_loader.py b/tests/tool/test_probe_loader.py index edb96950c..0c561f699 100644 --- a/tests/tool/test_probe_loader.py +++ b/tests/tool/test_probe_loader.py @@ -240,6 +240,13 @@ def test_onesec_manifest_uses_threat_connectivity_probe(self): assert spec.tool == "onesec_threat" assert spec.params == {"action": "threat_query_bd_version"} + def test_qingteng_manifest_uses_system_audit_connectivity_probe(self): + spec = get_connectivity_spec("qingteng_v3_4_1_66") + + assert spec is not None + assert spec.tool == "qingteng_system_audit" + assert spec.params == {} + class TestFixtureParsing: def setup_method(self): diff --git a/tests/tool/test_qingteng_api_tool.py b/tests/tool/test_qingteng_api_tool.py index 138617086..a32f5f441 100644 --- a/tests/tool/test_qingteng_api_tool.py +++ b/tests/tool/test_qingteng_api_tool.py @@ -19,7 +19,7 @@ def _load_tool(yaml_name: str): / ".flocks" / "plugins" / "tools" - / "api" + / "device" / _QINGTENG_PLUGIN_DIR / yaml_name ) @@ -33,7 +33,7 @@ def _load_handler_module(script_name: str, module_name: str): / ".flocks" / "plugins" / "tools" - / "api" + / "device" / _QINGTENG_PLUGIN_DIR / script_name ) @@ -743,6 +743,7 @@ async def test_qingteng_system_audit_uses_shared_handler_logic(): ] module.get_secret_manager = lambda: mock_secret_manager + module.ConfigWriter.get_api_service_raw = lambda service_id: {} module.httplib.HTTPConnection = _FakeHTTPConnection module.time.time = lambda: 1700000000 @@ -759,17 +760,7 @@ async def test_qingteng_system_audit_uses_shared_handler_logic(): assert result.success is True assert result.metadata["api"] == "system.audit" assert result.output["total"] == 1 - - -@pytest.mark.asyncio -async def test_qingteng_login_returns_clear_error_when_configuration_missing(): - module = _load_handler_module("qingteng.handler.py", "qingteng_login_handler_test") - mock_secret_manager = MagicMock() - mock_secret_manager.get.return_value = None - module.get_secret_manager = lambda: mock_secret_manager - module.ConfigWriter.get_api_service_raw = lambda service_id: {} - - result = await module.login(ToolContext(session_id="test", message_id="test")) - - assert result.success is False - assert result.error == "Missing configuration: qingteng base_url/qingteng_host, qingteng_username, qingteng_password" + assert len(_FakeHTTPConnection.created) == 2 + login_call = _FakeHTTPConnection.created[0].calls[0] + assert login_call["method"] == "POST" + assert login_call["url"] == "/v1/api/auth" diff --git a/tests/updater/test_updater.py b/tests/updater/test_updater.py index 58dec2b10..dfb742f18 100644 --- a/tests/updater/test_updater.py +++ b/tests/updater/test_updater.py @@ -1218,6 +1218,73 @@ def test_backup_current_version_excludes_all_dist_directories( assert "flocks/dist/ignored.txt" not in names +def test_backup_current_version_uses_filtered_snapshot_after_direct_failure( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + install_root = tmp_path / "install" + install_root.mkdir() + (install_root / "source.py").write_text("print('ok')", encoding="utf-8") + venv_dir = install_root / ".venv" + venv_dir.mkdir() + (venv_dir / "dependency.py").write_text("excluded", encoding="utf-8") + node_modules = install_root / "webui" / "node_modules" + node_modules.mkdir(parents=True) + (node_modules / "dependency.js").write_text("excluded", encoding="utf-8") + backup_dir = tmp_path / "backups" + archive_sources: list[Path] = [] + original_write_backup_archive = updater._write_backup_archive + + def fail_first_archive(source_root: Path, archive_path: Path) -> None: + archive_sources.append(source_root) + if len(archive_sources) == 1: + archive_path.write_text("partial", encoding="utf-8") + raise RuntimeError("unexpected end of data") + original_write_backup_archive(source_root, archive_path) + + monkeypatch.setattr(updater, "_BACKUP_DIR", backup_dir) + monkeypatch.setattr(updater, "_write_backup_archive", fail_first_archive) + + backup_path = updater._backup_current_version(install_root, "2026.7.29", retain_count=1) + + assert backup_path is not None + assert archive_sources[0] == install_root + assert archive_sources[1] != install_root + assert not archive_sources[1].exists() + assert not list(backup_dir.glob("*.partial")) + with tarfile.open(backup_path, "r:gz") as tar: + names = tar.getnames() + assert "flocks/source.py" in names + assert "flocks/.venv/dependency.py" not in names + assert "flocks/webui/node_modules/dependency.js" not in names + + +def test_backup_current_version_cleans_partial_when_snapshot_fallback_fails( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + install_root = tmp_path / "install" + install_root.mkdir() + (install_root / "source.py").write_text("print('ok')", encoding="utf-8") + backup_dir = tmp_path / "backups" + archive_attempts = 0 + + def fail_archive(_source_root: Path, archive_path: Path) -> None: + nonlocal archive_attempts + archive_attempts += 1 + archive_path.write_text("partial", encoding="utf-8") + raise RuntimeError("archive failed") + + monkeypatch.setattr(updater, "_BACKUP_DIR", backup_dir) + monkeypatch.setattr(updater, "_write_backup_archive", fail_archive) + + backup_path = updater._backup_current_version(install_root, "2026.7.29", retain_count=1) + + assert backup_path is None + assert archive_attempts == 2 + assert not list(backup_dir.iterdir()) + + @pytest.mark.asyncio async def test_build_updated_frontend_uses_current_install_root_and_region_mirror( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/workflow/test_execution_store_compact.py b/tests/workflow/test_execution_store_compact.py index 21a268e4a..f5d02a555 100644 --- a/tests/workflow/test_execution_store_compact.py +++ b/tests/workflow/test_execution_store_compact.py @@ -24,6 +24,7 @@ DEFAULT_COMPACT_SIZE_THRESHOLD, DEFAULT_GENERIC_SEQUENCE_THRESHOLD, DEFAULT_LARGE_LIST_KEYS, + DEFAULT_MAX_INLINE_COLLECTION_BYTES, _trim_execution_history, compact_history_for_storage, compact_execution_summary, @@ -76,6 +77,24 @@ def test_compact_outputs_keeps_small_lists_verbatim() -> None: assert "_enriched_alerts_count" not in compacted +def test_compact_outputs_strips_small_count_large_alert_lists() -> None: + large_record = {"body": "x" * DEFAULT_MAX_INLINE_COLLECTION_BYTES} + outputs = {"enriched_alerts_with_triage": [large_record]} + + compacted = compact_outputs_for_storage(outputs) + + assert compacted == {"_enriched_alerts_with_triage_count": 1} + + +def test_compact_outputs_summarizes_small_count_large_unknown_sequences() -> None: + outputs = {"unknown_payload": [{"body": "x" * DEFAULT_MAX_INLINE_COLLECTION_BYTES}]} + + compacted = compact_outputs_for_storage(outputs) + + assert compacted["unknown_payload"]["_type"] == "list" + assert compacted["unknown_payload"]["count"] == 1 + + def test_compact_outputs_summarizes_unknown_large_sequences() -> None: big_unknown = _make_alerts(DEFAULT_GENERIC_SEQUENCE_THRESHOLD + 1) outputs = {"some_other_alerts": big_unknown} @@ -351,6 +370,7 @@ def test_default_large_list_keys_cover_stream_alert_dedup_outputs() -> None: "raw_alerts", "normalized_alerts", "filtered_alerts", + "enriched_alerts_with_triage", } assert expected <= DEFAULT_LARGE_LIST_KEYS diff --git a/tests/workflow/test_stream_alert_triage_incremental_load.py b/tests/workflow/test_stream_alert_triage_incremental_load.py new file mode 100644 index 000000000..12cc67421 --- /dev/null +++ b/tests/workflow/test_stream_alert_triage_incremental_load.py @@ -0,0 +1,883 @@ +from __future__ import annotations + +import ast +import hashlib +import io +import json +import os +import time +from pathlib import Path +from types import SimpleNamespace + +import flocks.config +import pytest +from flocks.workflow import NodeTimeoutError, Workflow, WorkflowEngine +from flocks.workflow.repl_runtime import PythonExecRuntime + + +WORKFLOW_PATH = ( + Path(__file__).resolve().parents[2] + / ".flocks" + / "flockshub" + / "plugins" + / "workflows" + / "stream_alert_triage" + / "workflow.json" +) + + +def _workflow() -> dict[str, object]: + return json.loads(WORKFLOW_PATH.read_text(encoding="utf-8")) + + +def _node_code(node_id: str) -> str: + workflow = _workflow() + return next(node["code"] for node in workflow["nodes"] if node["id"] == node_id) + + +def _run_node(node_id: str, inputs: dict[str, object]) -> dict[str, object]: + namespace: dict[str, object] = {"inputs": inputs, "outputs": {}} + exec(compile(_node_code(node_id), str(WORKFLOW_PATH), "exec"), namespace) + return namespace["outputs"] + + +def _write_alerts(path: Path, ids: range | list[int]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + "".join(json.dumps({"id": value, "dedup_key": f"key-{value}"}) + "\n" for value in ids), + encoding="utf-8", + ) + + +def _write_named_alerts(path: Path, prefix: str, count: int = 20) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + "".join( + json.dumps({"id": f"{prefix}-{index:03d}", "dedup_key": f"key-{index:03d}"}) + "\n" + for index in range(count) + ), + encoding="utf-8", + ) + + +def _use_flocks_root(monkeypatch, root: Path) -> None: + class FakeConfig: + def get_global(self) -> SimpleNamespace: + return SimpleNamespace(data_dir=root / "data") + + monkeypatch.setattr(flocks.config, "Config", FakeConfig) + + +def _release_loader_lease(outputs: dict[str, object]) -> None: + lease_fd = outputs.get("_batch_lease_fd") + if not isinstance(lease_fd, int): + return + _run_node( + "commit_cursor", + { + "cursor_enabled": True, + "pending_cursor": None, + "_batch_lease_fd": lease_fd, + "batch_lease_token": outputs.get("batch_lease_token"), + }, + ) + + +def test_explicit_replay_is_bounded_resumable_and_reads_appends(tmp_path: Path) -> None: + input_path = tmp_path / "dedup_result_001.jsonl" + _write_alerts(input_path, range(20)) + + first = _run_node("load_dedup_file", {"input_path": str(input_path)}) + second = _run_node( + "load_dedup_file", + {"input_path": str(input_path), "resume_cursor": first["next_cursor"]}, + ) + + assert [item["id"] for item in first["enriched_alerts"]] == list(range(10)) + assert [item["id"] for item in second["enriched_alerts"]] == list(range(10, 20)) + assert first["cursor_enabled"] is False + assert first["has_more"] is True + assert second["has_more"] is False + + with input_path.open("a", encoding="utf-8") as stream: + for value in range(20, 30): + stream.write(json.dumps({"id": value, "dedup_key": f"key-{value}"}) + "\n") + + appended = _run_node( + "load_dedup_file", + {"input_path": str(input_path), "resume_cursor": second["next_cursor"]}, + ) + assert [item["id"] for item in appended["enriched_alerts"]] == list(range(20, 30)) + + +def test_replaced_input_file_invalidates_resume_cursor(tmp_path: Path) -> None: + input_path = tmp_path / "dedup_result_001.jsonl" + replacement = tmp_path / "replacement.jsonl" + _write_named_alerts(input_path, "aaa") + + first = _run_node("load_dedup_file", {"input_path": str(input_path)}) + _write_named_alerts(replacement, "bbb") + os.replace(replacement, input_path) + resumed = _run_node( + "load_dedup_file", + {"input_path": str(input_path), "resume_cursor": first["next_cursor"]}, + ) + + assert [item["id"] for item in resumed["enriched_alerts"]] == [ + f"bbb-{index:03d}" for index in range(10) + ] + assert resumed["load_stats"]["cursor_invalidated"] is True + + +def test_same_inode_rewrite_invalidates_resume_cursor(tmp_path: Path) -> None: + input_path = tmp_path / "dedup_result_001.jsonl" + _write_named_alerts(input_path, "aaa") + first = _run_node("load_dedup_file", {"input_path": str(input_path)}) + original_file_id = input_path.stat().st_ino + + payload = "".join( + json.dumps({"id": f"bbb-{index:03d}", "dedup_key": f"key-{index:03d}"}) + "\n" + for index in range(20) + ) + with input_path.open("r+", encoding="utf-8") as stream: + stream.seek(0) + stream.write(payload) + stream.truncate() + assert input_path.stat().st_ino == original_file_id + + resumed = _run_node( + "load_dedup_file", + {"input_path": str(input_path), "resume_cursor": first["next_cursor"]}, + ) + + assert [item["id"] for item in resumed["enriched_alerts"]] == [ + f"bbb-{index:03d}" for index in range(10) + ] + assert resumed["load_stats"]["cursor_invalidated"] is True + + +def test_same_inode_head_rewrite_invalidates_cursor_beyond_boundary_anchor(tmp_path: Path) -> None: + input_path = tmp_path / "dedup_result_001.jsonl" + input_path.write_text( + "".join( + json.dumps( + { + "id": f"aaa-{index:03d}", + "dedup_key": f"key-{index:03d}", + "padding": "x" * 1200, + } + ) + + "\n" + for index in range(20) + ), + encoding="utf-8", + ) + first = _run_node("load_dedup_file", {"input_path": str(input_path)}) + assert first["next_cursor"]["byte_offset"] > 8192 + original_file_id = input_path.stat().st_ino + + with input_path.open("r+", encoding="utf-8") as stream: + payload = stream.read() + stream.seek(0) + stream.write(payload.replace("aaa-000", "bbb-000", 1)) + stream.truncate() + assert input_path.stat().st_ino == original_file_id + + resumed = _run_node( + "load_dedup_file", + {"input_path": str(input_path), "resume_cursor": first["next_cursor"]}, + ) + + assert resumed["enriched_alerts"][0]["id"] == "bbb-000" + assert resumed["load_stats"]["cursor_invalidated"] is True + + +def test_auto_mode_sorts_numeric_sequences_and_spans_files(monkeypatch, tmp_path: Path) -> None: + root = tmp_path / "flocks-root" + _use_flocks_root(monkeypatch, root) + day_dir = root / "workspace" / "workflows" / "stream_alert_denoise" / "2026-07-30" + _write_alerts(day_dir / "dedup_result_1000.jsonl", range(6, 16)) + _write_alerts(day_dir / "dedup_result_999.jsonl", range(6)) + (day_dir / "dedup_result_bad.jsonl").write_text("{}\n", encoding="utf-8") + + outputs = _run_node( + "load_dedup_file", + {"input_date": "2026-07-30", "batch_max_records": 10}, + ) + + assert [item["id"] for item in outputs["enriched_alerts"]] == list(range(10)) + assert [Path(path).name for path in outputs["loaded_files"]] == [ + "dedup_result_999.jsonl", + "dedup_result_1000.jsonl", + ] + assert outputs["pending_cursor"]["file_name"] == "dedup_result_1000.jsonl" + assert outputs["load_stats"]["invalid_file_names"] == 1 + + +def test_complete_bad_lines_advance_but_partial_line_does_not(tmp_path: Path) -> None: + input_path = tmp_path / "dedup_result_001.jsonl" + consumed = ( + json.dumps({"_type": "file_header"}) + + "\n\nnot-json\n[]\n" + + json.dumps({"id": 1}) + + "\n" + ) + input_path.write_bytes(consumed.encode("utf-8") + b'{"id": 2') + + first = _run_node("load_dedup_file", {"input_path": str(input_path)}) + + assert first["enriched_alerts"] == [{"id": 1}] + assert first["pending_cursor"]["byte_offset"] == len(consumed.encode("utf-8")) + assert first["load_stats"]["header_skipped"] == 1 + assert first["load_stats"]["empty_lines"] == 1 + assert first["load_stats"]["bad_lines"] == 1 + assert first["load_stats"]["non_object_lines"] == 1 + assert first["load_stats"]["partial_lines"] == 1 + + with input_path.open("ab") as stream: + stream.write(b"}\n") + second = _run_node( + "load_dedup_file", + {"input_path": str(input_path), "resume_cursor": first["next_cursor"]}, + ) + assert second["enriched_alerts"] == [{"id": 2}] + + +def test_oversized_line_is_skipped_without_exceeding_batch_accounting(tmp_path: Path) -> None: + input_path = tmp_path / "dedup_result_001.jsonl" + oversized = b"x" * 160 + b"\n" + input_path.write_bytes(oversized + b'{"id": 1}\n') + + first = _run_node( + "load_dedup_file", + {"input_path": str(input_path), "batch_max_bytes": 64}, + ) + + assert first["enriched_alerts"] == [] + assert first["batch_bytes"] == 64 + assert first["load_stats"]["oversized_lines"] == 1 + assert first["load_stats"]["oversized_bytes_discarded"] == 64 + assert first["pending_cursor"]["byte_offset"] == 64 + assert first["pending_cursor"]["skipping_oversized_line"] is True + assert first["has_more"] is True + + second = _run_node( + "load_dedup_file", + { + "input_path": str(input_path), + "batch_max_bytes": 64, + "resume_cursor": first["next_cursor"], + }, + ) + assert second["enriched_alerts"] == [] + assert second["batch_bytes"] == 64 + assert second["load_stats"]["oversized_bytes_discarded"] == 64 + assert second["pending_cursor"]["byte_offset"] == 128 + assert second["pending_cursor"]["skipping_oversized_line"] is True + + third = _run_node( + "load_dedup_file", + { + "input_path": str(input_path), + "batch_max_bytes": 64, + "resume_cursor": second["next_cursor"], + }, + ) + assert third["enriched_alerts"] == [{"id": 1}] + assert third["batch_bytes"] <= 64 + assert "skipping_oversized_line" not in third["pending_cursor"] + + +def test_bounded_line_reader_does_not_reread_short_lines() -> None: + tree = ast.parse(_node_code("load_dedup_file")) + function = next( + node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == "_read_bounded_line" + ) + namespace = {"_READ_CHUNK_BYTES": 64 * 1024} + exec(compile(ast.Module(body=[function], type_ignores=[]), str(WORKFLOW_PATH), "exec"), namespace) + + class CountingBytesIO(io.BytesIO): + def __init__(self, initial_bytes: bytes) -> None: + super().__init__(initial_bytes) + self.returned_bytes = 0 + + def read(self, size: int = -1) -> bytes: + data = super().read(size) + self.returned_bytes += len(data) + return data + + def readline(self, size: int = -1) -> bytes: + data = super().readline(size) + self.returned_bytes += len(data) + return data + + payload = b"\n" * 1024 + stream = CountingBytesIO(payload) + consumed = 0 + while consumed < len(payload): + status, line, scanned = namespace["_read_bounded_line"](stream, len(payload) - consumed) + assert (status, line, scanned) == ("line", b"\n", 1) + consumed += scanned + + assert stream.returned_bytes == len(payload) + + +def test_line_that_exceeds_remaining_budget_is_retried_next_batch(tmp_path: Path) -> None: + input_path = tmp_path / "dedup_result_001.jsonl" + first_line = b'{"id":1}\n' + second_line = b'{"id":2,"value":"12345"}\n' + byte_limit = len(first_line) + len(second_line) - 1 + input_path.write_bytes(first_line + second_line) + + first = _run_node( + "load_dedup_file", + {"input_path": str(input_path), "batch_max_bytes": byte_limit}, + ) + second = _run_node( + "load_dedup_file", + { + "input_path": str(input_path), + "batch_max_bytes": byte_limit, + "resume_cursor": first["next_cursor"], + }, + ) + + assert first["enriched_alerts"] == [{"id": 1}] + assert first["pending_cursor"]["byte_offset"] == len(first_line) + assert first["batch_bytes"] == byte_limit + assert second["enriched_alerts"] == [{"id": 2, "value": "12345"}] + + +def test_explicit_missing_path_never_falls_back_to_auto_discovery(tmp_path: Path) -> None: + outputs = _run_node( + "load_dedup_file", + {"input_path": str(tmp_path / "missing.jsonl")}, + ) + + assert outputs["cursor_enabled"] is False + assert outputs["loaded_files"] == [] + assert outputs["enriched_alerts"] == [] + assert outputs["load_stats"]["missing_files"] == 1 + + +def test_production_cursor_is_only_advanced_by_commit_node(monkeypatch, tmp_path: Path) -> None: + root = tmp_path / "flocks-root" + _use_flocks_root(monkeypatch, root) + date = "2026-07-30" + input_path = ( + root + / "workspace" + / "workflows" + / "stream_alert_denoise" + / date + / "dedup_result_001.jsonl" + ) + _write_alerts(input_path, range(12)) + cursor_path = ( + root + / "workspace" + / "workflows" + / "stream_alert_triage" + / ".triage_cursor.json" + ) + + first = _run_node("load_dedup_file", {"input_date": date}) + with pytest.raises(RuntimeError, match="production_batch_lease_busy"): + _run_node("load_dedup_file", {"input_date": date}) + + assert not cursor_path.exists() + + commit_inputs = dict(first) + commit_inputs["triage_stats"] = {"triage_failed": 1} + commit_inputs["_triage_persistence_succeeded"] = True + committed = _run_node("commit_cursor", commit_inputs) + remaining = _run_node("load_dedup_file", {"input_date": date}) + + saved_cursor = json.loads(cursor_path.read_text(encoding="utf-8")) + assert committed["cursor_committed"] is True + assert saved_cursor["byte_offset"] == first["pending_cursor"]["byte_offset"] + assert saved_cursor["updated_at"] + assert [item["id"] for item in remaining["enriched_alerts"]] == [10, 11] + _release_loader_lease(remaining) + + +def test_triage_timeout_does_not_commit_production_cursor(monkeypatch, tmp_path: Path) -> None: + root = tmp_path / "flocks-root" + _use_flocks_root(monkeypatch, root) + cursor_path = ( + root + / "workspace" + / "workflows" + / "stream_alert_triage" + / ".triage_cursor.json" + ) + pending_cursor = { + "version": 1, + "date": "2026-07-30", + "file_seq": 1, + "file_name": "dedup_result_001.jsonl", + "byte_offset": 123, + } + workflow = Workflow.from_dict( + { + "name": "triage_timeout_cursor_guard", + "start": "load", + "nodes": [ + { + "id": "load", + "type": "python", + "description": "Provide a pending production cursor", + "code": ( + "outputs.update({" + f"'cursor_enabled': True, 'pending_cursor': {pending_cursor!r}, " + "'cursor_before': None, 'next_cursor': None, " + "'_triage_persistence_succeeded': False})" + ), + }, + { + "id": "triage", + "type": "python", + "description": "Exceed the node timeout before persistence succeeds", + "processIsolated": True, + "timeoutFatal": True, + "code": ( + "import time\n" + "triage_outputs = outputs\n" + "time.sleep(0.2)\n" + "triage_outputs['_triage_persistence_succeeded'] = True" + ), + }, + { + "id": "commit", + "type": "python", + "description": "Use the real cursor commit node", + "code": _node_code("commit_cursor"), + }, + ], + "edges": [ + {"from": "load", "to": "triage"}, + {"from": "triage", "to": "commit"}, + ], + } + ) + engine = WorkflowEngine( + workflow, + runtime=PythonExecRuntime(tool_registry=SimpleNamespace(cancel_checker=None)), + node_timeout_s=0.05, + history_mode="full", + max_parallel_workers=1, + ) + + with pytest.raises(NodeTimeoutError) as caught: + engine.run(initial_inputs={}, retain_history=True) + + history = caught.value.execution_context["history"] + triage_step = next(step for step in history if step.node_id == "triage") + assert "节点执行超时" in (triage_step.error or "") + assert all(step.node_id != "commit" for step in history) + assert not cursor_path.exists() + + +def test_commit_node_does_not_rewrite_cursor_without_new_bytes(monkeypatch, tmp_path: Path) -> None: + root = tmp_path / "flocks-root" + _use_flocks_root(monkeypatch, root) + cursor_path = ( + root + / "workspace" + / "workflows" + / "stream_alert_triage" + / ".triage_cursor.json" + ) + cursor_path.parent.mkdir(parents=True) + original = { + "version": 1, + "date": "2026-07-30", + "file_seq": 1, + "file_name": "dedup_result_001.jsonl", + "byte_offset": 123, + "updated_at": "2026-07-30T10:00:00+08:00", + } + cursor_path.write_text(json.dumps(original), encoding="utf-8") + + outputs = _run_node( + "commit_cursor", + {"cursor_enabled": True, "cursor_before": original, "pending_cursor": None}, + ) + + assert outputs["cursor_committed"] is False + assert outputs["committed_cursor"] == original + assert json.loads(cursor_path.read_text(encoding="utf-8")) == original + + +def test_cursor_date_change_and_truncation_restart_current_input(monkeypatch, tmp_path: Path) -> None: + root = tmp_path / "flocks-root" + _use_flocks_root(monkeypatch, root) + cursor_dir = root / "workspace" / "workflows" / "stream_alert_triage" + cursor_dir.mkdir(parents=True) + cursor_path = cursor_dir / ".triage_cursor.json" + cursor_path.write_text( + json.dumps( + { + "version": 1, + "date": "2026-07-29", + "file_seq": 1, + "file_name": "dedup_result_001.jsonl", + "byte_offset": 9999, + } + ), + encoding="utf-8", + ) + input_path = ( + root + / "workspace" + / "workflows" + / "stream_alert_denoise" + / "2026-07-30" + / "dedup_result_001.jsonl" + ) + _write_alerts(input_path, [1]) + + date_reset = _run_node("load_dedup_file", {"input_date": "2026-07-30"}) + assert date_reset["enriched_alerts"] == [{"id": 1, "dedup_key": "key-1"}] + _release_loader_lease(date_reset) + + cursor_path.write_text( + json.dumps( + { + "version": 1, + "date": "2026-07-30", + "file_seq": 1, + "file_name": "dedup_result_001.jsonl", + "byte_offset": 9999, + } + ), + encoding="utf-8", + ) + truncated = _run_node("load_dedup_file", {"input_date": "2026-07-30"}) + assert truncated["enriched_alerts"] == [{"id": 1, "dedup_key": "key-1"}] + _release_loader_lease(truncated) + + input_path.write_bytes(b'{"id": 2') + truncated_partial = _run_node("load_dedup_file", {"input_date": "2026-07-30"}) + + assert truncated_partial["enriched_alerts"] == [] + assert truncated_partial["load_stats"]["partial_lines"] == 1 + assert truncated_partial["pending_cursor"]["byte_offset"] == 0 + assert truncated_partial["next_cursor"]["byte_offset"] == 0 + assert truncated_partial["has_more"] is True + _release_loader_lease(truncated_partial) + + +def test_semantically_invalid_production_cursor_restarts_from_file_head(monkeypatch, tmp_path: Path) -> None: + root = tmp_path / "flocks-root" + _use_flocks_root(monkeypatch, root) + date = "2026-07-30" + input_path = ( + root + / "workspace" + / "workflows" + / "stream_alert_denoise" + / date + / "dedup_result_001.jsonl" + ) + _write_alerts(input_path, [1, 2]) + cursor_path = ( + root + / "workspace" + / "workflows" + / "stream_alert_triage" + / ".triage_cursor.json" + ) + cursor_path.parent.mkdir(parents=True) + base_cursor = { + "version": 1, + "date": date, + "file_seq": 1, + "file_name": input_path.name, + "byte_offset": 1, + } + malformed_cursors = [ + {**base_cursor, "byte_offset": 1.5}, + {**base_cursor, "byte_offset": True}, + {**base_cursor, "version": 2}, + ] + + for malformed in malformed_cursors: + cursor_path.write_text(json.dumps(malformed), encoding="utf-8") + outputs = _run_node("load_dedup_file", {"input_date": date}) + + assert outputs["cursor_before"] is None + assert [item["id"] for item in outputs["enriched_alerts"]] == [1, 2] + assert outputs["load_stats"]["bad_lines"] == 0 + _release_loader_lease(outputs) + + +def test_successful_triage_sets_persistence_gate(monkeypatch, tmp_path: Path) -> None: + root = tmp_path / "flocks-root" + _use_flocks_root(monkeypatch, root) + + outputs = _run_node( + "concurrent_triage", + {"enriched_alerts": [], "triage_output_mode": "none"}, + ) + + assert outputs["_triage_persistence_succeeded"] is True + + +def test_production_triage_batch_lease_blocks_overlapping_run(monkeypatch, tmp_path: Path) -> None: + root = tmp_path / "flocks-root" + _use_flocks_root(monkeypatch, root) + date = "2026-07-30" + input_path = ( + root + / "workspace" + / "workflows" + / "stream_alert_denoise" + / date + / "dedup_result_001.jsonl" + ) + _write_alerts(input_path, []) + + first = _run_node("load_dedup_file", {"input_date": date}) + assert isinstance(first["_batch_lease_fd"], int) + + with pytest.raises(RuntimeError, match="production_batch_lease_busy"): + _run_node("load_dedup_file", {"input_date": date}) + + triaged = _run_node( + "concurrent_triage", + {**first, "triage_output_mode": "none"}, + ) + _run_node( + "commit_cursor", + { + **triaged, + "_triage_persistence_succeeded": True, + }, + ) + + retry = _run_node("load_dedup_file", {"input_date": date}) + _run_node( + "commit_cursor", + { + **retry, + "pending_cursor": None, + }, + ) + + +def test_stale_cursor_commit_cannot_overwrite_newer_cursor(monkeypatch, tmp_path: Path) -> None: + root = tmp_path / "flocks-root" + _use_flocks_root(monkeypatch, root) + cursor_path = ( + root + / "workspace" + / "workflows" + / "stream_alert_triage" + / ".triage_cursor.json" + ) + identity = { + "version": 2, + "date": "2026-07-30", + "file_seq": 1, + "file_name": "dedup_result_001.jsonl", + "device_id": 1, + "file_id": 2, + "head_hash": hashlib.sha256(b"").hexdigest(), + "boundary_start": 0, + "boundary_hash": hashlib.sha256(b"").hexdigest(), + } + + newer = _run_node( + "commit_cursor", + { + "cursor_enabled": True, + "cursor_before": None, + "cursor_revision": None, + "pending_cursor": {**identity, "byte_offset": 200}, + "_triage_persistence_succeeded": True, + "_run_id": "newer", + }, + ) + stale = _run_node( + "commit_cursor", + { + "cursor_enabled": True, + "cursor_before": None, + "cursor_revision": None, + "pending_cursor": {**identity, "byte_offset": 100}, + "_triage_persistence_succeeded": True, + "_run_id": "stale", + }, + ) + + assert newer["cursor_committed"] is True + assert stale["cursor_committed"] is False + assert stale["cursor_commit_error"] == "stale_cursor_commit" + assert json.loads(cursor_path.read_text(encoding="utf-8"))["byte_offset"] == 200 + assert not list(cursor_path.parent.glob(".triage_cursor.json.*.tmp")) + + +def test_invalidated_cursor_reset_requires_cas_and_cannot_move_to_older_date( + monkeypatch, + tmp_path: Path, +) -> None: + root = tmp_path / "flocks-root" + _use_flocks_root(monkeypatch, root) + cursor_path = ( + root + / "workspace" + / "workflows" + / "stream_alert_triage" + / ".triage_cursor.json" + ) + identity = { + "version": 2, + "date": "2026-07-30", + "file_seq": 1, + "file_name": "dedup_result_001.jsonl", + "device_id": 1, + "file_id": 2, + "head_hash": hashlib.sha256(b"old-head").hexdigest(), + "boundary_start": 0, + "boundary_hash": hashlib.sha256(b"old-boundary").hexdigest(), + } + _run_node( + "commit_cursor", + { + "cursor_enabled": True, + "cursor_revision": None, + "pending_cursor": {**identity, "byte_offset": 200}, + "_triage_persistence_succeeded": True, + }, + ) + + current_revision = hashlib.sha256(cursor_path.read_bytes()).hexdigest() + reset = _run_node( + "commit_cursor", + { + "cursor_enabled": True, + "cursor_revision": current_revision, + "cursor_invalidated": True, + "pending_cursor": { + **identity, + "byte_offset": 100, + "device_id": 3, + "file_id": 4, + "head_hash": hashlib.sha256(b"new-head").hexdigest(), + "boundary_hash": hashlib.sha256(b"new-boundary").hexdigest(), + }, + "_triage_persistence_succeeded": True, + }, + ) + assert reset["cursor_committed"] is True + assert json.loads(cursor_path.read_text(encoding="utf-8"))["byte_offset"] == 100 + + reset_revision = hashlib.sha256(cursor_path.read_bytes()).hexdigest() + older = _run_node( + "commit_cursor", + { + "cursor_enabled": True, + "cursor_revision": reset_revision, + "cursor_invalidated": True, + "pending_cursor": {**identity, "date": "2026-07-29", "byte_offset": 0}, + "_triage_persistence_succeeded": True, + }, + ) + assert older["cursor_committed"] is False + assert older["cursor_commit_error"] == "stale_cursor_commit" + assert json.loads(cursor_path.read_text(encoding="utf-8"))["date"] == "2026-07-30" + + +def test_workflow_wires_commit_after_persistence_and_has_dynamic_samples() -> None: + workflow = _workflow() + + assert [node["id"] for node in workflow["nodes"]] == [ + "load_dedup_file", + "concurrent_triage", + "commit_cursor", + "summarize", + ] + assert [(edge["from"], edge["to"]) for edge in workflow["edges"]] == [ + ("load_dedup_file", "concurrent_triage"), + ("concurrent_triage", "commit_cursor"), + ("commit_cursor", "summarize"), + ] + assert "input_date" not in workflow["metadata"]["sampleInputs"] + assert workflow["metadata"]["sampleInputs"]["batch_max_records"] == 10 + assert workflow["metadata"]["sampleInputs"]["batch_max_bytes"] == 32 * 1024 * 1024 + for trigger in workflow["triggers"]: + assert trigger["runtime"]["noOverlap"] is True + assert "input_date" not in trigger["inputs"] + assert "input_date" not in trigger["testSamples"][0]["payload"] + + +def test_jsonl_persistence_failure_is_reraised() -> None: + tree = ast.parse(_node_code("concurrent_triage")) + persistence_try = next( + node + for node in ast.walk(tree) + if isinstance(node, ast.Try) + and any( + isinstance(child, ast.Call) + and isinstance(child.func, ast.Name) + and child.func.id == "_triage_write_jsonl" + for child in ast.walk(node) + ) + ) + + assert any( + isinstance(child, ast.Raise) + for handler in persistence_try.handlers + for child in ast.walk(handler) + ) + + +def test_jsonl_counter_write_failure_is_not_swallowed() -> None: + tree = ast.parse(_node_code("concurrent_triage")) + function = next( + node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == "_triage_set_counter" + ) + + assert not any(isinstance(node, ast.ExceptHandler) for node in ast.walk(function)) + + +def test_commit_writer_uses_atomic_replace_and_fsync(monkeypatch, tmp_path: Path) -> None: + root = tmp_path / "flocks-root" + _use_flocks_root(monkeypatch, root) + cursor_path = ( + root + / "workspace" + / "workflows" + / "stream_alert_triage" + / ".triage_cursor.json" + ) + + outputs = _run_node( + "commit_cursor", + { + "cursor_enabled": True, + "cursor_revision": None, + "_triage_persistence_succeeded": True, + "_run_id": "atomic-test", + "pending_cursor": { + "version": 2, + "date": "2026-07-30", + "file_seq": 3, + "file_name": "dedup_result_003.jsonl", + "byte_offset": 123, + "device_id": 1, + "file_id": 2, + "head_hash": hashlib.sha256(b"").hexdigest(), + "boundary_start": 0, + "boundary_hash": hashlib.sha256(b"").hexdigest(), + "skipping_oversized_line": True, + }, + }, + ) + committed = outputs["committed_cursor"] + + assert json.loads(cursor_path.read_text(encoding="utf-8")) == committed + assert committed["updated_at"] + assert committed["skipping_oversized_line"] is True + assert not list(cursor_path.parent.glob(".triage_cursor.json.*.tmp")) diff --git a/tests/workflow/test_stream_alert_triage_persistence.py b/tests/workflow/test_stream_alert_triage_persistence.py index 722f844e4..f475f16c7 100644 --- a/tests/workflow/test_stream_alert_triage_persistence.py +++ b/tests/workflow/test_stream_alert_triage_persistence.py @@ -4,6 +4,7 @@ import datetime as datetime_module import json import os +import pickle import re import sqlite3 import threading @@ -11,6 +12,8 @@ from concurrent.futures import ThreadPoolExecutor from pathlib import Path +import pytest + WORKFLOW_PATH = ( Path(__file__).resolve().parents[2] @@ -45,8 +48,10 @@ def _load_functions(*names: str) -> dict[str, object]: "inputs": {"loaded_files": []}, "json": json, "os": os, + "pickle": pickle, "re": re, "time": time, + "MAX_TRIAGE_CACHE_BYTES": 128 * 1024 * 1024, } exec(compile(ast.Module(body=body, type_ignores=[]), str(WORKFLOW_PATH), "exec"), namespace) return namespace @@ -108,6 +113,90 @@ def ask(self, _prompt: str, **_kwargs: object) -> str: ) +def test_triage_does_not_truncate_content_or_limit_output_tokens() -> None: + code = _concurrent_triage_code() + + for forbidden in ( + "MAX_HTTP_FIELD_CHARS", + "MAX_LOG_TEXT_CHARS", + "MAX_LLM_PROMPT_CHARS", + "MAX_LLM_RESPONSE_CHARS", + "MAX_TRIAGE_REPORT_CHARS", + "LLM_CALL_MAX_TOKENS", + "_clip_prompt_text", + "'max_tokens':", + ): + assert forbidden not in code + + +def test_cache_eviction_enforces_serialized_byte_budget() -> None: + evict = _load_functions("_evict_lru")["_evict_lru"] + cache = {f"key-{index}": {"report": chr(65 + index) * 2048} for index in range(3)} + + evicted = evict(cache, max_keys=100, max_bytes=2500) + + assert evicted == 2 + assert list(cache) == ["key-2"] + assert len(pickle.dumps(cache, protocol=pickle.HIGHEST_PROTOCOL)) <= 2500 + + +def test_oversized_cache_is_quarantined_before_unpickling(tmp_path: Path) -> None: + functions = _load_functions("_load_cache") + cache_path = tmp_path / "triage_cache.pkl" + cache_path.write_bytes(b"x" * 9) + + class FailIfLoaded: + @staticmethod + def load(_stream: object) -> object: + raise AssertionError("oversized cache must not be deserialized") + + functions["pickle"] = FailIfLoaded + functions["MAX_TRIAGE_CACHE_BYTES"] = 8 + + assert functions["_load_cache"](str(cache_path)) == {} + assert not cache_path.exists() + assert len(list(tmp_path.glob("triage_cache.pkl.*.oversized"))) == 1 + + +def test_cache_is_loaded_once_while_batch_lock_is_held() -> None: + code = _concurrent_triage_code() + tree = ast.parse(code) + cache_path_loads = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "_load_cache" + and len(node.args) == 1 + and isinstance(node.args[0], ast.Name) + and node.args[0].id == "cache_path" + ] + + assert len(cache_path_loads) == 1 + assert "if new_results or evicted:" in code + assert "MAX_TRIAGE_CACHE_BYTES if new_results else None" in code + + +def test_memory_heavy_triage_node_uses_fatal_process_isolation() -> None: + workflow = json.loads(WORKFLOW_PATH.read_text(encoding="utf-8")) + node = next(item for item in workflow["nodes"] if item["id"] == "concurrent_triage") + + assert node["processIsolated"] is True + assert node["processRetainFdKeys"] == ["_batch_lease_fd"] + assert "processInheritFdKeys" not in node + assert node["timeoutFatal"] is True + assert "outputs['enriched_alerts_with_triage']" not in node["code"] + assert "top_triage_result['triage_report'] = a.get('triage_report', '')" in node["code"] + + +def test_triage_work_units_are_submitted_in_small_batches() -> None: + code = _concurrent_triage_code() + + assert "TRIAGE_SUB_BATCH_SIZE = 3" in code + assert "range(0, len(work_units), TRIAGE_SUB_BATCH_SIZE)" in code + assert "pool.submit(_process_unit, *u) for u in work_batch" in code + + def test_soc_db_selects_only_verified_first_seen_unique_alerts() -> None: functions = _load_functions("_input_bool", "_select_first_seen_soc_alerts") select_alerts = functions["_select_first_seen_soc_alerts"] @@ -149,6 +238,40 @@ def test_soc_db_persistence_uses_filtered_first_seen_alerts() -> None: assert write_calls[0].args[1].id == "first_seen_soc_alerts" +def test_apply_triage_fields_preserves_original_alert_fields() -> None: + apply_triage_fields = _load_functions("_apply_triage_fields")["_apply_triage_fields"] + record = { + "attack_verdict": "raw-verdict", + "attack_success": True, + "threat_result": "raw-result", + } + + apply_triage_fields( + record, + { + "triage_attack_verdict": "attack", + "triage_attack_success": "failed", + "risk_level": "High", + "report_title": "Model report", + "triage_report": "# Model report", + "attack_verdict": "model-verdict", + "attack_success": False, + "threat_result": "model-result", + }, + ) + + assert record == { + "attack_verdict": "raw-verdict", + "attack_success": True, + "threat_result": "raw-result", + "triage_attack_verdict": "attack", + "triage_attack_success": "failed", + "risk_level": "High", + "report_title": "Model report", + "triage_report": "# Model report", + } + + def test_soc_db_persistence_failure_is_reraised() -> None: tree = ast.parse(_concurrent_triage_code()) persistence_try = next( @@ -170,6 +293,138 @@ def test_soc_db_persistence_failure_is_reraised() -> None: ) +def test_cache_save_failure_is_reraised_and_cleans_unique_temp_file(tmp_path: Path) -> None: + functions = _load_functions("_save_cache_atomic") + functions["pickle"] = pickle + functions["threading"] = threading + cache_path = tmp_path / "cache-target" + cache_path.mkdir() + + with pytest.raises(RuntimeError, match="failed to save triage cache"): + functions["_save_cache_atomic"](str(cache_path), {"key": {"value": 1}}) + + assert not list(tmp_path.glob("cache-target.*.tmp")) + + +def test_soc_db_merge_preserves_original_attack_fields_and_updates_triage_fields() -> None: + merge_record = _load_functions("_merge_triage_record")["_merge_triage_record"] + + merged = merge_record( + { + "attack_success": True, + "attack_verdict": "attack_success", + "threat_result": "success", + }, + {"triage_attack_success": "unknown", "triage_attack_verdict": "non_attack"}, + ) + + assert merged["attack_success"] is True + assert merged["attack_verdict"] == "attack_success" + assert merged["threat_result"] == "success" + assert merged["triage_attack_success"] == "unknown" + assert merged["triage_attack_verdict"] == "non_attack" + + +def test_apply_triage_fields_namespaces_model_attack_fields() -> None: + apply_fields = _load_functions("_apply_triage_fields")["_apply_triage_fields"] + record = {"attack_success": True, "attack_verdict": "attack_success"} + + apply_fields( + record, + { + "triage_attack_success": "unknown", + "triage_attack_verdict": "non_attack", + "risk_level": "Low", + }, + ) + + assert record["attack_success"] is True + assert record["attack_verdict"] == "attack_success" + assert record["triage_attack_success"] == "unknown" + assert record["triage_attack_verdict"] == "non_attack" + assert record["risk_level"] == "Low" + + +def test_apply_triage_fields_validates_direct_model_dimensions() -> None: + apply_fields = _load_functions("_apply_triage_fields")["_apply_triage_fields"] + cases = [ + (("attack", "success"), ("attack", "success")), + (("attack", "failed"), ("attack", "failed")), + (("attack", "unknown"), ("attack", "unknown")), + (("non_attack", "success"), ("non_attack", "unknown")), + (("unknown", "failed"), ("unknown", "unknown")), + (("invalid", "invalid"), ("unknown", "unknown")), + ] + + for (model_verdict, model_success), (attack_verdict, attack_success) in cases: + record: dict[str, object] = {} + apply_fields( + record, + { + "triage_attack_verdict": model_verdict, + "triage_attack_success": model_success, + }, + ) + assert record["triage_attack_verdict"] == attack_verdict + assert record["triage_attack_success"] == attack_success + + +def test_llm_attack_outcome_generates_the_two_persisted_fields_directly() -> None: + functions = _load_functions("_normalize_triage_outcome", "_llm_attack_outcome") + functions.update( + { + "_strip_think": lambda value: value, + "_ask_llm": lambda _prompt: json.dumps( + { + "triage_attack_verdict": "attack", + "triage_attack_success": "failed", + } + ), + } + ) + + assert functions["_llm_attack_outcome"]("analysis") == { + "triage_attack_verdict": "attack", + "triage_attack_success": "failed", + } + + +def test_current_triage_cache_requires_the_two_field_schema() -> None: + functions = _load_functions("_normalize_triage_outcome", "_is_current_triage_fields") + functions.update( + { + "TRIAGE_FIELDS": ( + "triage_attack_verdict", + "triage_attack_success", + "risk_level", + "report_title", + "triage_report", + ), + "_is_valid_triage_report": lambda value: value == "valid-report", + } + ) + is_current = functions["_is_current_triage_fields"] + + assert is_current( + { + "triage_attack_verdict": "attack", + "triage_attack_success": "success", + "risk_level": "High", + "report_title": "title", + "triage_report": "valid-report", + } + ) + assert not is_current( + { + "attack_verdict": "attack_success", + "attack_success": True, + "risk_level": "High", + "report_title": "title", + "triage_report": "valid-report", + } + ) + + def test_soc_db_writer_receives_only_selected_first_seen_alerts(tmp_path: Path) -> None: functions = _load_functions( "_input_bool", diff --git a/tests/workflow/test_stream_alert_workflow_edge_mappings.py b/tests/workflow/test_stream_alert_workflow_edge_mappings.py new file mode 100644 index 000000000..7094671f7 --- /dev/null +++ b/tests/workflow/test_stream_alert_workflow_edge_mappings.py @@ -0,0 +1,360 @@ +from __future__ import annotations + +import json +from pathlib import Path +from types import SimpleNamespace + +import flocks.config +import flocks.workspace.manager +import pytest + +from flocks.workflow import Workflow, WorkflowEngine +from flocks.workflow.edge_resolver import EdgeResolver +from flocks.workflow.execution_plan import resolve_workflow_dataflow_mode +from flocks.workflow.repl_runtime import PythonExecRuntime +from flocks.workflow.workflow_lint import lint_workflow + + +WORKFLOW_ROOT = ( + Path(__file__).resolve().parents[2] + / ".flocks" + / "flockshub" + / "plugins" + / "workflows" +) + +EXPECTED_MAPPING_KEYS = { + "stream_alert_denoise": { + ("receive_alert", "normalize"): { + "raw_alerts", + "input_mode", + "source_log_type", + "filter_enabled", + "dedup_enabled", + "dedup_threshold", + "strict_fields", + "lsh_fields", + "max_field_len", + "max_dedup_keys", + "stats", + }, + ("normalize", "filter_logs"): { + "normalized_alerts", + "input_mode", + "source_log_type", + "filter_enabled", + "dedup_enabled", + "dedup_threshold", + "strict_fields", + "lsh_fields", + "max_field_len", + "max_dedup_keys", + "stats", + }, + ("filter_logs", "dedup_and_write"): { + "filtered_alerts", + "input_mode", + "dedup_enabled", + "dedup_threshold", + "strict_fields", + "lsh_fields", + "max_field_len", + "max_dedup_keys", + "stats", + }, + }, + "stream_alert_triage": { + ("load_dedup_file", "concurrent_triage"): { + "enriched_alerts", + "loaded_files", + "load_stats", + "concurrency", + "max_triage_cache_size", + "input_date", + "cursor_enabled", + "cursor_before", + "pending_cursor", + "next_cursor", + "has_more", + "batch_records", + "batch_bytes", + "_triage_persistence_succeeded", + "_run_id", + "triage_output_mode", + "persist_triage_output", + "soc_db_path", + "jsonl_output_dir", + "cursor_revision", + "cursor_invalidated", + "_batch_lease_fd", + "batch_lease_token", + "_triage_state_dir", + }, + ("concurrent_triage", "commit_cursor"): { + "cursor_enabled", + "cursor_before", + "pending_cursor", + "next_cursor", + "has_more", + "batch_records", + "batch_bytes", + "_triage_persistence_succeeded", + "input_date", + "load_stats", + "loaded_files", + "top_triage_result", + "triage_results", + "triage_stats", + "triage_output_mode", + "soc_db_result", + "soc_db_path", + "output_paths", + "output_dir", + "cursor_revision", + "cursor_invalidated", + "_batch_lease_fd", + "batch_lease_token", + }, + ("commit_cursor", "summarize"): { + "cursor_enabled", + "cursor_committed", + "committed_cursor", + "next_cursor", + "has_more", + "batch_records", + "batch_bytes", + "input_date", + "load_stats", + "loaded_files", + "top_triage_result", + "triage_results", + "triage_stats", + "triage_output_mode", + "soc_db_result", + "soc_db_path", + "output_paths", + "output_dir", + "cursor_commit_error", + "cursor_revision", + "cursor_invalidated", + }, + }, +} + + +def _workflow_dict(workflow_id: str) -> dict[str, object]: + path = WORKFLOW_ROOT / workflow_id / "workflow.json" + return json.loads(path.read_text(encoding="utf-8")) + + +def _use_flocks_root(monkeypatch: pytest.MonkeyPatch, root: Path) -> None: + class FakeConfig: + def get_global(self) -> SimpleNamespace: + return SimpleNamespace(data_dir=root / "data") + + monkeypatch.setattr(flocks.config, "Config", FakeConfig) + + +@pytest.mark.parametrize("workflow_id", EXPECTED_MAPPING_KEYS) +def test_stream_workflow_edges_use_strict_explicit_mappings(workflow_id: str) -> None: + raw = _workflow_dict(workflow_id) + workflow = Workflow.from_dict(raw) + actual = { + (edge["from"], edge["to"]): edge.get("mapping", {}) + for edge in raw["edges"] + } + expected = { + edge: {field: field for field in fields} + for edge, fields in EXPECTED_MAPPING_KEYS[workflow_id].items() + } + + assert raw["metadata"]["runtime"] == { + "strict_edge_mapping": True, + "dataflow_mode": "vertex_cache", + } + assert resolve_workflow_dataflow_mode(workflow.metadata) == "vertex_cache" + assert actual == expected + assert lint_workflow(workflow) == [] + + +@pytest.mark.parametrize("workflow_id", EXPECTED_MAPPING_KEYS) +def test_stream_workflow_mappings_drop_unlisted_payload_fields(workflow_id: str) -> None: + raw = _workflow_dict(workflow_id) + workflow = Workflow.from_dict(raw) + nodes = workflow.nodes_by_id() + resolver = EdgeResolver(dataflow_mode="vertex_cache") + + for edge in workflow.edges: + source_values = {source: f"value-for-{source}" for source in edge.mapping.values()} + resolved = resolver.resolve( + node=nodes[edge.from_], + node_inputs={"unlisted_large_payload": [object()]}, + node_outputs=source_values, + edges=[edge], + ) + + assert len(resolved) == 1 + assert resolved[0][1] == { + destination: source_values[source] + for destination, source in edge.mapping.items() + } + assert "unlisted_large_payload" not in resolved[0][1] + + +def test_mappings_do_not_forward_obsolete_large_alert_lists() -> None: + denoise = _workflow_dict("stream_alert_denoise") + denoise_edges = { + (edge["from"], edge["to"]): set(edge["mapping"].values()) + for edge in denoise["edges"] + } + triage = _workflow_dict("stream_alert_triage") + triage_edges = { + (edge["from"], edge["to"]): set(edge["mapping"].values()) + for edge in triage["edges"] + } + + assert "raw_alerts" not in denoise_edges[("normalize", "filter_logs")] + assert not {"raw_alerts", "normalized_alerts"} & denoise_edges[("filter_logs", "dedup_and_write")] + assert "enriched_alerts" not in triage_edges[("concurrent_triage", "commit_cursor")] + assert "pending_cursor" not in triage_edges[("commit_cursor", "summarize")] + + +def test_denoise_execution_drops_consumed_payloads(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + _use_flocks_root(monkeypatch, tmp_path / "flocks-root") + raw = _workflow_dict("stream_alert_denoise") + workflow = Workflow.from_dict(raw) + engine = WorkflowEngine( + workflow, + runtime=PythonExecRuntime(tool_registry=SimpleNamespace(cancel_checker=None)), + node_timeout_s=None, + history_mode="full", + dataflow_mode=resolve_workflow_dataflow_mode(workflow.metadata), + max_parallel_workers=1, + ) + + result = engine.run( + initial_inputs={ + "alerts": [ + { + "net_type": "http", + "net_http_url": "/health", + "net_real_src_ip": "192.0.2.1", + "net_dest_ip": "198.51.100.1", + "threat_name": "mapping-test", + "threat_type": "web攻击", + } + ], + "filter_enabled": False, + "dedup_enabled": False, + "threshold": 0.42, + "unlisted_large_payload": ["sentinel"], + }, + retain_history=True, + ) + + steps = {step.node_id: step for step in result.history} + assert all(step.error is None for step in result.history) + assert "unlisted_large_payload" not in steps["normalize"].inputs + assert "raw_alerts" not in steps["filter_logs"].inputs + assert "normalized_alerts" not in steps["dedup_and_write"].inputs + assert steps["dedup_and_write"].inputs["dedup_threshold"] == 0.42 + + +def test_triage_execution_drops_loader_payload_before_commit_and_summary( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + _use_flocks_root(monkeypatch, tmp_path / "flocks-root") + + class FakeWorkspaceManager: + @classmethod + def get_instance(cls) -> FakeWorkspaceManager: + return cls() + + def get_workspace_dir(self) -> Path: + return tmp_path / "workspace" + + monkeypatch.setattr(flocks.workspace.manager, "WorkspaceManager", FakeWorkspaceManager) + input_path = tmp_path / "dedup_result_001.jsonl" + input_path.write_text("", encoding="utf-8") + raw = _workflow_dict("stream_alert_triage") + workflow = Workflow.from_dict(raw) + engine = WorkflowEngine( + workflow, + runtime=PythonExecRuntime(tool_registry=SimpleNamespace(cancel_checker=None)), + node_timeout_s=None, + history_mode="full", + dataflow_mode=resolve_workflow_dataflow_mode(workflow.metadata), + max_parallel_workers=1, + ) + + result = engine.run( + initial_inputs={ + "input_path": str(input_path), + "triage_output_mode": "none", + "unlisted_large_payload": ["sentinel"], + }, + retain_history=True, + ) + + steps = {step.node_id: step for step in result.history} + assert all(step.error is None for step in result.history) + assert "unlisted_large_payload" not in steps["concurrent_triage"].inputs + assert steps["concurrent_triage"].inputs["triage_output_mode"] == "none" + assert "enriched_alerts" not in steps["commit_cursor"].inputs + assert steps["commit_cursor"].inputs["_triage_persistence_succeeded"] is True + assert "pending_cursor" not in steps["summarize"].inputs + assert "top_triage_result" in steps["summarize"].inputs + assert "enriched_alerts_with_triage" not in steps["summarize"].inputs + + +def test_production_triage_lease_survives_node_boundaries_and_is_released( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + flocks_root = tmp_path / "flocks-root" + _use_flocks_root(monkeypatch, flocks_root) + + class FakeWorkspaceManager: + @classmethod + def get_instance(cls) -> FakeWorkspaceManager: + return cls() + + def get_workspace_dir(self) -> Path: + return tmp_path / "workspace" + + monkeypatch.setattr(flocks.workspace.manager, "WorkspaceManager", FakeWorkspaceManager) + date = "2026-07-30" + input_path = ( + flocks_root + / "workspace" + / "workflows" + / "stream_alert_denoise" + / date + / "dedup_result_001.jsonl" + ) + input_path.parent.mkdir(parents=True) + input_path.write_text("", encoding="utf-8") + + raw = _workflow_dict("stream_alert_triage") + workflow = Workflow.from_dict(raw) + engine = WorkflowEngine( + workflow, + runtime=PythonExecRuntime(tool_registry=SimpleNamespace(cancel_checker=None)), + node_timeout_s=None, + history_mode="full", + dataflow_mode=resolve_workflow_dataflow_mode(workflow.metadata), + max_parallel_workers=1, + ) + + for _ in range(2): + result = engine.run( + initial_inputs={"input_date": date, "triage_output_mode": "none"}, + retain_history=True, + ) + steps = {step.node_id: step for step in result.history} + assert all(step.error is None for step in result.history) + lease_fd = steps["load_dedup_file"].outputs["_batch_lease_fd"] + assert isinstance(lease_fd, int) + assert steps["concurrent_triage"].inputs["_batch_lease_fd"] == lease_fd + assert steps["commit_cursor"].inputs["_batch_lease_fd"] == lease_fd diff --git a/tests/workflow/test_workflow_history_mode.py b/tests/workflow/test_workflow_history_mode.py index af4412ea7..a7360e296 100644 --- a/tests/workflow/test_workflow_history_mode.py +++ b/tests/workflow/test_workflow_history_mode.py @@ -139,3 +139,127 @@ def test_python_runtime_can_cleanup_node_globals_after_execute() -> None: assert "temporary_payload" not in runtime.globals assert "inputs" not in runtime.globals assert "outputs" not in runtime.globals + + +def test_mapped_payload_is_not_retained_for_remaining_run() -> None: + workflow = { + "start": "produce", + "metadata": { + "runtime": { + "strict_edge_mapping": True, + "dataflow_mode": "vertex_cache", + } + }, + "nodes": [ + { + "id": "produce", + "type": "python", + "code": "\n".join( + [ + "import weakref", + "class Payload(list): pass", + "payload = Payload(range(100000))", + "outputs['large_payload'] = payload", + "outputs['payload_ref'] = weakref.ref(payload)", + ] + ), + }, + { + "id": "inspect", + "type": "python", + "code": "\n".join( + [ + "import gc", + "gc.collect()", + "outputs['large_payload_still_alive'] = inputs['payload_ref']() is not None", + ] + ), + }, + ], + "edges": [ + { + "from": "produce", + "to": "inspect", + "mapping": {"payload_ref": "payload_ref"}, + }, + ], + } + + result = run_workflow( + workflow=workflow, + history_mode="summary", + ensure_requirements=False, + ) + + assert result.status == "SUCCEEDED" + assert result.outputs == {"large_payload_still_alive": False} + + +def test_parallel_mapped_payloads_are_released_before_downstream_nodes() -> None: + producer_code = "\n".join( + [ + "import weakref", + "class Payload(list): pass", + "payload = Payload(range(100000))", + "outputs['large_payload'] = payload", + "outputs['payload_ref'] = weakref.ref(payload)", + ] + ) + inspect_code = "\n".join( + [ + "import gc", + "gc.collect()", + "outputs['large_payload_still_alive'] = inputs['payload_ref']() is not None", + ] + ) + workflow = { + "start": "seed", + "metadata": {"runtime": {"dataflow_mode": "vertex_cache"}}, + "nodes": [ + {"id": "seed", "type": "python", "code": "outputs['token'] = True"}, + {"id": "produce_a", "type": "python", "code": producer_code}, + {"id": "produce_b", "type": "python", "code": producer_code}, + {"id": "inspect_a", "type": "python", "code": inspect_code}, + {"id": "inspect_b", "type": "python", "code": inspect_code}, + { + "id": "collect", + "type": "python", + "join": True, + "code": "outputs.update(inputs)", + }, + ], + "edges": [ + {"from": "seed", "to": "produce_a", "mapping": {"token": "token"}}, + {"from": "seed", "to": "produce_b", "mapping": {"token": "token"}}, + { + "from": "produce_a", + "to": "inspect_a", + "mapping": {"payload_ref": "payload_ref"}, + }, + { + "from": "produce_b", + "to": "inspect_b", + "mapping": {"payload_ref": "payload_ref"}, + }, + { + "from": "inspect_a", + "to": "collect", + "mapping": {"alive_a": "large_payload_still_alive"}, + }, + { + "from": "inspect_b", + "to": "collect", + "mapping": {"alive_b": "large_payload_still_alive"}, + }, + ], + } + + result = run_workflow( + workflow=workflow, + history_mode="summary", + ensure_requirements=False, + max_parallel_workers=2, + ) + + assert result.status == "SUCCEEDED" + assert result.outputs == {"alive_a": False, "alive_b": False} diff --git a/tests/workflow/test_workflow_node_timeout.py b/tests/workflow/test_workflow_node_timeout.py index ffd598a9b..98d5e714b 100644 --- a/tests/workflow/test_workflow_node_timeout.py +++ b/tests/workflow/test_workflow_node_timeout.py @@ -1,13 +1,18 @@ -"""Test workflow node timeout: node times out is skipped and error is recorded.""" +"""Regression tests for isolated and compatibility node timeouts.""" + +import os +import threading +import time import pytest -from flocks.workflow import Workflow, WorkflowEngine, run_workflow -from flocks.workflow.repl_runtime import PythonExecRuntime +import flocks.workflow.repl_runtime as repl_runtime_module +from flocks.workflow import NodeExecutionError, NodeTimeoutError, Workflow, WorkflowEngine, run_workflow +from flocks.workflow.repl_runtime import HostProcessPythonExecRuntime, PythonExecRuntime -def test_node_timeout_skips_node_and_records_error(): - """When a node exceeds node_timeout_s, it is skipped and error is in history.""" +def test_node_timeout_aborts_run_without_orphan_thread_or_downstream(): + """A timed-out subprocess is killed before the run fails.""" workflow = Workflow.from_dict({ "name": "timeout_test", "start": "slow", @@ -16,7 +21,9 @@ def test_node_timeout_skips_node_and_records_error(): "id": "slow", "type": "python", "code": "import time; time.sleep(5); outputs['x'] = 1", - "description": "Sleep 5s", + "description": "Never finishes within the timeout", + "processIsolated": True, + "timeoutFatal": True, }, { "id": "fast", @@ -31,25 +38,272 @@ def test_node_timeout_skips_node_and_records_error(): engine = WorkflowEngine( workflow, runtime=rt, - node_timeout_s=1.0, + node_timeout_s=0.05, stop_on_error=False, ) - result = engine.run(initial_inputs={}, retain_history=True) + baseline_threads = {thread.ident for thread in threading.enumerate() if thread.name.startswith("wf-node")} + started = time.perf_counter() + with pytest.raises(NodeTimeoutError) as caught: + engine.run(initial_inputs={}, retain_history=True) + elapsed = time.perf_counter() - started - assert result.steps == 2 - assert len(result.history) == 2 + history = caught.value.execution_context["history"] + assert elapsed < 1.0 + assert caught.value.execution_context["steps"] == 1 + assert len(history) == 1 - step_slow = result.history[0] + step_slow = history[0] assert step_slow.node_id == "slow" assert step_slow.error is not None assert "节点执行超时" in step_slow.error - assert "1.0" in step_slow.error + assert "0.05" in step_slow.error assert step_slow.outputs == {} + assert {thread.ident for thread in threading.enumerate() if thread.name.startswith("wf-node")} == baseline_threads - step_fast = result.history[1] - assert step_fast.node_id == "fast" - assert step_fast.error is None - assert step_fast.outputs.get("y") == 10 # x missing, get('x', 0) = 0, 0+10=10 + +def test_process_isolated_timeout_can_remain_nonfatal_for_compatibility(): + """Fatal timeout semantics remain opt-in at the node level.""" + workflow = Workflow.from_dict({ + "name": "cooperative_timeout", + "start": "slow", + "nodes": [ + { + "id": "slow", + "type": "python", + "code": "import time; time.sleep(5)", + "processIsolated": True, + } + ], + "edges": [], + }) + engine = WorkflowEngine( + workflow, + runtime=PythonExecRuntime(), + node_timeout_s=0.05, + ) + + started = time.perf_counter() + result = engine.run(retain_history=True) + + assert time.perf_counter() - started < 0.5 + assert result.steps == 1 + assert "节点执行超时" in (result.history[0].error or "") + assert not any(thread.name.startswith("wf-node") for thread in threading.enumerate()) + + +def test_process_timeout_closes_inherited_parent_fd(tmp_path): + lease_fd = os.open(tmp_path / "lease.lock", os.O_RDWR | os.O_CREAT, 0o600) + workflow = Workflow.from_dict({ + "start": "slow", + "nodes": [ + { + "id": "slow", + "type": "python", + "code": "import os, time\nos.fstat(inputs['lease_fd'])\ntime.sleep(5)", + "processIsolated": True, + "processInheritFdKeys": ["lease_fd"], + "timeoutFatal": True, + } + ], + "edges": [], + }) + + with pytest.raises(NodeTimeoutError): + WorkflowEngine( + workflow, + runtime=PythonExecRuntime(), + node_timeout_s=0.05, + ).run({"lease_fd": lease_fd}) + + with pytest.raises(OSError): + os.fstat(lease_fd) + + +def test_process_rpc_bridge_matches_concurrent_responses_by_request_id(): + class Registry: + cancel_checker = None + + def __init__(self): + self.active = 0 + self.peak = 0 + self.lock = threading.Lock() + + def run(self, _name, *, value): + with self.lock: + self.active += 1 + self.peak = max(self.peak, self.active) + try: + time.sleep((8 - value) * 0.005) + return value + finally: + with self.lock: + self.active -= 1 + + registry = Registry() + workflow = Workflow.from_dict({ + "start": "parallel_rpc", + "nodes": [ + { + "id": "parallel_rpc", + "type": "python", + "processIsolated": True, + "code": ( + "from concurrent.futures import ThreadPoolExecutor\n" + "def call(value):\n" + " return tool.run('echo', value=value)\n" + "with ThreadPoolExecutor(max_workers=8) as pool:\n" + " outputs['values'] = list(pool.map(call, range(8)))" + ), + } + ], + "edges": [], + }) + + result = WorkflowEngine( + workflow, + runtime=PythonExecRuntime(tool_registry=registry), + node_timeout_s=3, + history_mode="full", + ).run() + + assert result.outputs["values"] == list(range(8)) + assert registry.peak > 1 + + +def test_process_isolated_runtime_exposes_cooperative_cancel_hooks(): + workflow = Workflow.from_dict({ + "start": "check_cancel", + "nodes": [ + { + "id": "check_cancel", + "type": "python", + "processIsolated": True, + "code": ( + "outputs['cancelled'] = cancelled()\n" + "outputs['is_cancelled'] = is_cancelled()" + ), + } + ], + "edges": [], + }) + + result = WorkflowEngine( + workflow, + runtime=PythonExecRuntime(), + node_timeout_s=3, + ).run() + + assert result.outputs == {"cancelled": False, "is_cancelled": False} + + +def test_process_isolated_system_exit_preserves_outputs(): + workflow = Workflow.from_dict({ + "start": "early_return", + "nodes": [ + { + "id": "early_return", + "type": "python", + "processIsolated": True, + "code": "outputs['x'] = 1\nraise SystemExit(0)", + } + ], + "edges": [], + }) + + result = WorkflowEngine( + workflow, + runtime=PythonExecRuntime(), + node_timeout_s=3, + ).run() + + assert result.outputs == {"x": 1} + + +def test_process_retained_fd_stays_in_parent_and_crosses_node_boundary(tmp_path): + lease_fd = os.open(tmp_path / "lease.lock", os.O_RDWR | os.O_CREAT, 0o600) + workflow = Workflow.from_dict({ + "start": "passthrough", + "nodes": [ + { + "id": "passthrough", + "type": "python", + "processIsolated": True, + "processRetainFdKeys": ["lease_fd"], + "code": ( + "outputs['child_fd'] = inputs.get('lease_fd')\n" + "outputs['lease_fd'] = inputs.get('lease_fd')" + ), + } + ], + "edges": [], + }) + + try: + result = WorkflowEngine( + workflow, + runtime=PythonExecRuntime(), + node_timeout_s=3, + ).run({"lease_fd": lease_fd}) + + # The child sees a harmless placeholder descriptor, while the parent + # restores and continues owning the actual lease descriptor. + assert isinstance(result.outputs["child_fd"], int) + assert result.outputs["lease_fd"] == lease_fd + os.fstat(lease_fd) + finally: + try: + os.close(lease_fd) + except OSError: + pass + + +def test_process_retained_fd_is_closed_when_child_fails(tmp_path): + lease_fd = os.open(tmp_path / "lease.lock", os.O_RDWR | os.O_CREAT, 0o600) + workflow = Workflow.from_dict({ + "start": "fail", + "nodes": [ + { + "id": "fail", + "type": "python", + "processIsolated": True, + "processRetainFdKeys": ["lease_fd"], + "code": "raise RuntimeError('boom')", + } + ], + "edges": [], + }) + + with pytest.raises(NodeExecutionError, match="boom"): + WorkflowEngine( + workflow, + runtime=PythonExecRuntime(), + node_timeout_s=3, + ).run({"lease_fd": lease_fd}) + + with pytest.raises(OSError): + os.fstat(lease_fd) + + +def test_host_process_windows_launch_does_not_require_posix_shell(monkeypatch): + real_popen = repl_runtime_module.subprocess.Popen + script_paths = [] + + def checking_popen(args, *popen_args, **popen_kwargs): + assert args[0] != "sh" + script_paths.append(args[-1]) + return real_popen(args, *popen_args, **popen_kwargs) + + monkeypatch.setattr(repl_runtime_module.sys, "platform", "win32") + monkeypatch.setattr(repl_runtime_module.subprocess, "Popen", checking_popen) + + outputs, _stdout = HostProcessPythonExecRuntime().execute( + "outputs['ok'] = True", + {}, + ) + + assert outputs == {"ok": True} + assert script_paths + assert all(not os.path.exists(path) for path in script_paths) def test_node_timeout_none_disabled(): @@ -81,8 +335,10 @@ def test_run_workflow_node_timeout_param(): { "id": "s", "type": "python", - "code": "import time; time.sleep(3); outputs['ok'] = 1", + "code": "import time; time.sleep(0.5); outputs['ok'] = 1", "description": "Slow", + "processIsolated": True, + "timeoutFatal": True, }, ], "edges": [], @@ -94,7 +350,8 @@ def test_run_workflow_node_timeout_param(): ensure_requirements=False, retain_history=True, ) - assert result.status == "SUCCEEDED" + assert result.status == "FAILED" + assert "NodeTimeoutError" in (result.error or "") assert len(result.history) == 1 assert result.history[0].get("error") is not None assert "节点执行超时" in result.history[0]["error"] @@ -111,6 +368,8 @@ def test_run_workflow_uses_metadata_node_timeout_default(): "type": "python", "code": "import time; time.sleep(0.2); outputs['ok'] = 1", "description": "Slow-ish", + "processIsolated": True, + "timeoutFatal": True, }, ], "edges": [], @@ -122,7 +381,8 @@ def test_run_workflow_uses_metadata_node_timeout_default(): ensure_requirements=False, retain_history=True, ) - assert result.status == "SUCCEEDED" + assert result.status == "FAILED" + assert "NodeTimeoutError" in (result.error or "") assert len(result.history) == 1 assert "节点执行超时" in (result.history[0].get("error") or "") @@ -138,6 +398,8 @@ def test_run_workflow_explicit_node_timeout_overrides_metadata(): "type": "python", "code": "import time; time.sleep(0.2); outputs['ok'] = 1", "description": "Slow-ish", + "processIsolated": True, + "timeoutFatal": True, }, ], "edges": [], diff --git a/tests/workflow/test_workflow_parallel.py b/tests/workflow/test_workflow_parallel.py index 1207d1a6d..26a58d5b7 100644 --- a/tests/workflow/test_workflow_parallel.py +++ b/tests/workflow/test_workflow_parallel.py @@ -8,6 +8,7 @@ from __future__ import annotations import asyncio +import threading import time from typing import Any, Dict from unittest.mock import patch @@ -15,6 +16,7 @@ import pytest from flocks.tool.registry import ParameterType, ToolCategory, ToolParameter, ToolRegistry, ToolResult +from flocks.workflow.errors import NodeTimeoutError from flocks.workflow.engine import ExecutionResult, StepResult, WorkflowEngine from flocks.workflow.models import Workflow from flocks.workflow.repl_runtime import PythonExecRuntime @@ -293,11 +295,11 @@ async def _workflow_parallel_shared_loop_tool(ctx, value: str) -> ToolResult: try: calls: list[str] = [] - def _spy_run_sync(coro): + def _spy_run_sync(coro, cancel_checker): calls.append(type(coro).__name__) - from flocks.workflow._async_runtime import run_sync + from flocks.workflow._async_runtime import run_sync_cancellable - return run_sync(coro) + return run_sync_cancellable(coro, cancel_checker) wf = Workflow.from_dict({ "name": "parallel_tool_test", @@ -339,7 +341,10 @@ def _spy_run_sync(coro): ], }) - with patch("flocks.workflow.tools_adapter._run_sync_on_shared_loop", side_effect=_spy_run_sync): + with patch( + "flocks.workflow.tools_adapter._run_sync_cancellable_on_shared_loop", + side_effect=_spy_run_sync, + ): result = WorkflowEngine( wf, runtime=PythonExecRuntime(), @@ -365,7 +370,8 @@ def _spy_run_sync(coro): class TestParallelDedup: """Dedup still works correctly with batch draining.""" - def test_dedup_with_parallel_batch(self): + @pytest.mark.parametrize("history_mode", ["full", "summary"]) + def test_dedup_with_parallel_batch(self, history_mode): """Identical inputs to the same node are deduped within a batch.""" wf = Workflow.from_dict({ "name": "dedup_par", @@ -383,7 +389,7 @@ def test_dedup_with_parallel_batch(self): wf, runtime=PythonExecRuntime(), max_parallel_workers=4, - history_mode="full", + history_mode=history_mode, ) result = engine.run(retain_history=True) b_steps = [s for s in result.history if s.node_id == "b"] @@ -393,18 +399,25 @@ def test_dedup_with_parallel_batch(self): class TestParallelTimeout: """Node timeout behaviour in parallel mode.""" - def test_parallel_timeout_marks_slow_node(self): - """A slow parallel node is marked as timed-out while fast ones succeed.""" + def test_parallel_timeout_aborts_after_draining_workers(self): + """A slow parallel node aborts the run after all workers exit.""" wf = Workflow.from_dict({ "name": "par_timeout", "start": "start", "nodes": [ {"id": "start", "type": "python", "code": "outputs['x'] = 1"}, - {"id": "fast", "type": "python", "code": "outputs['r'] = 'ok'"}, + { + "id": "fast", + "type": "python", + "code": "outputs['r'] = 'ok'", + "processIsolated": True, + }, { "id": "slow", "type": "python", "code": "import time; time.sleep(5); outputs['r'] = 'done'", + "processIsolated": True, + "timeoutFatal": True, }, ], "edges": [ @@ -420,22 +433,23 @@ def test_parallel_timeout_marks_slow_node(self): stop_on_error=False, ) t0 = time.perf_counter() - result = engine.run(retain_history=True) + with pytest.raises(NodeTimeoutError) as caught: + engine.run(retain_history=True) elapsed = time.perf_counter() - t0 - # Should complete near the timeout, not wait for the 5s sleep. assert elapsed < 2.0 - - fast_step = next(s for s in result.history if s.node_id == "fast") + history = caught.value.execution_context["history"] + fast_step = next(s for s in history if s.node_id == "fast") assert fast_step.error is None assert fast_step.outputs.get("r") == "ok" - slow_step = next(s for s in result.history if s.node_id == "slow") + slow_step = next(s for s in history if s.node_id == "slow") assert slow_step.error is not None assert "超时" in slow_step.error + assert not any(thread.name.startswith("wf-par") for thread in threading.enumerate()) - def test_parallel_timeout_is_non_fatal(self): - """Timeout in parallel does not trigger stop_on_error.""" + def test_parallel_isolated_timeout_can_be_nonfatal(self): + """timeoutFatal preserves the historical nonfatal default when false.""" wf = Workflow.from_dict({ "name": "par_timeout_nonfatal", "start": "start", @@ -445,8 +459,14 @@ def test_parallel_timeout_is_non_fatal(self): "id": "slow", "type": "python", "code": "import time; time.sleep(5); outputs['r'] = 'done'", + "processIsolated": True, + }, + { + "id": "fast", + "type": "python", + "code": "outputs['r'] = 'ok'", + "processIsolated": True, }, - {"id": "fast", "type": "python", "code": "outputs['r'] = 'ok'"}, ], "edges": [ {"from": "start", "to": "slow"}, @@ -458,9 +478,8 @@ def test_parallel_timeout_is_non_fatal(self): runtime=PythonExecRuntime(), max_parallel_workers=4, node_timeout_s=0.3, - stop_on_error=True, + stop_on_error=False, ) - # Should NOT raise even though stop_on_error=True, because timeout is non-fatal. result = engine.run(retain_history=True) errors = [s for s in result.history if s.error is not None] assert len(errors) == 1 diff --git a/uv.lock b/uv.lock index 5eb156e92..181d3829a 100644 --- a/uv.lock +++ b/uv.lock @@ -553,7 +553,7 @@ wheels = [ [[package]] name = "flocks" -version = "2026.7.29" +version = "2026.8.4" source = { editable = "." } dependencies = [ { name = "aiofiles" }, diff --git a/webui/package-lock.json b/webui/package-lock.json index 2f6078c9a..b635cdd32 100644 --- a/webui/package-lock.json +++ b/webui/package-lock.json @@ -15,6 +15,7 @@ "date-fns": "^3.3.1", "i18next": "^25.8.14", "i18next-browser-languagedetector": "^8.2.1", + "katex": "^0.18.1", "lucide-react": "^0.562.0", "pdfjs-dist": "^6.1.200", "qrcode.react": "^4.2.0", @@ -25,10 +26,12 @@ "react-router-dom": "^7.12.0", "recharts": "^2.15.0", "rehype-highlight": "^7.0.2", + "rehype-katex": "^7.0.1", "rehype-raw": "^7.0.0", "rehype-sanitize": "^6.0.0", "remark-breaks": "^4.0.0", "remark-gfm": "^4.0.1", + "remark-math": "^6.0.0", "tailwind-merge": "^2.2.1", "zustand": "^4.5.0" }, @@ -2221,6 +2224,12 @@ "dev": true, "license": "MIT" }, + "node_modules/@types/katex": { + "version": "0.16.8", + "resolved": "https://registry.npmjs.org/@types/katex/-/katex-0.16.8.tgz", + "integrity": "sha512-trgaNyfU+Xh2Tc+ABIb44a5AYUpicB3uwirOioeOkNPPbmgRNtcWyDeeFRzjPZENO9Vq8gvVqfhaaXWLlevVwg==", + "license": "MIT" + }, "node_modules/@types/mdast": { "version": "4.0.4", "resolved": "https://registry.npmjs.org/@types/mdast/-/mdast-4.0.4.tgz", @@ -4354,6 +4363,55 @@ "node": ">= 0.4" } }, + "node_modules/hast-util-from-dom": { + "version": "5.0.1", + "resolved": "https://registry.npmjs.org/hast-util-from-dom/-/hast-util-from-dom-5.0.1.tgz", + "integrity": "sha512-N+LqofjR2zuzTjCPzyDUdSshy4Ma6li7p/c3pA78uTwzFgENbgbUrm2ugwsOdcjI1muO+o6Dgzp9p8WHtn/39Q==", + "license": "ISC", + "dependencies": { + "@types/hast": "^3.0.0", + "hastscript": "^9.0.0", + "web-namespaces": "^2.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, + "node_modules/hast-util-from-html": { + "version": "2.0.3", + "resolved": "https://registry.npmjs.org/hast-util-from-html/-/hast-util-from-html-2.0.3.tgz", + "integrity": "sha512-CUSRHXyKjzHov8yKsQjGOElXy/3EKpyX56ELnkHH34vDVw1N1XSQ1ZcAvTyAPtGqLTuKP/uxM+aLkSPqF/EtMw==", + "license": "MIT", + "dependencies": { + "@types/hast": "^3.0.0", + "devlop": "^1.1.0", + "hast-util-from-parse5": "^8.0.0", + "parse5": "^7.0.0", + "vfile": "^6.0.0", + "vfile-message": "^4.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, + "node_modules/hast-util-from-html-isomorphic": { + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/hast-util-from-html-isomorphic/-/hast-util-from-html-isomorphic-2.0.0.tgz", + "integrity": "sha512-zJfpXq44yff2hmE0XmwEOzdWin5xwH+QIhMLOScpX91e/NSGPsAzNCvLQDIEPyO2TXi+lBmU6hjLIhV8MwP2kw==", + "license": "MIT", + "dependencies": { + "@types/hast": "^3.0.0", + "hast-util-from-dom": "^5.0.0", + "hast-util-from-html": "^2.0.0", + "unist-util-remove-position": "^5.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, "node_modules/hast-util-from-parse5": { "version": "8.0.3", "resolved": "https://registry.npmjs.org/hast-util-from-parse5/-/hast-util-from-parse5-8.0.3.tgz", @@ -4969,6 +5027,31 @@ "node": ">=6" } }, + "node_modules/katex": { + "version": "0.18.1", + "resolved": "https://registry.npmjs.org/katex/-/katex-0.18.1.tgz", + "integrity": "sha512-Td8GCYSxDAoMhHOlKmCFMJ/hz5qlAAb71n66Dryw9nfCVfumLo7nhuotbvKom/XPADmrYC3O5QR71EPq4DarJQ==", + "funding": [ + "https://opencollective.com/katex", + "https://github.com/sponsors/katex" + ], + "license": "MIT", + "dependencies": { + "commander": "^8.3.0" + }, + "bin": { + "katex": "cli.js" + } + }, + "node_modules/katex/node_modules/commander": { + "version": "8.3.0", + "resolved": "https://registry.npmjs.org/commander/-/commander-8.3.0.tgz", + "integrity": "sha512-OkTL9umf+He2DZkUq8f8J9of7yL6RJKI24dVITBmNfZBmri9zYZQrKkuXiKhyfPSu8tUhnVBB1iKXevvnlR4Ww==", + "license": "MIT", + "engines": { + "node": ">= 12" + } + }, "node_modules/keyv": { "version": "4.5.4", "resolved": "https://registry.npmjs.org/keyv/-/keyv-4.5.4.tgz", @@ -5289,6 +5372,25 @@ "url": "https://opencollective.com/unified" } }, + "node_modules/mdast-util-math": { + "version": "3.0.0", + "resolved": "https://registry.npmjs.org/mdast-util-math/-/mdast-util-math-3.0.0.tgz", + "integrity": "sha512-Tl9GBNeG/AhJnQM221bJR2HPvLOSnLE/T9cJI9tlc6zwQk2nPk/4f0cHkOdEixQPC/j8UtKDdITswvLAy1OZ1w==", + "license": "MIT", + "dependencies": { + "@types/hast": "^3.0.0", + "@types/mdast": "^4.0.0", + "devlop": "^1.0.0", + "longest-streak": "^3.0.0", + "mdast-util-from-markdown": "^2.0.0", + "mdast-util-to-markdown": "^2.1.0", + "unist-util-remove-position": "^5.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, "node_modules/mdast-util-mdx-expression": { "version": "2.0.1", "resolved": "https://registry.npmjs.org/mdast-util-mdx-expression/-/mdast-util-mdx-expression-2.0.1.tgz", @@ -5631,6 +5733,50 @@ "url": "https://opencollective.com/unified" } }, + "node_modules/micromark-extension-math": { + "version": "3.1.0", + "resolved": "https://registry.npmjs.org/micromark-extension-math/-/micromark-extension-math-3.1.0.tgz", + "integrity": "sha512-lvEqd+fHjATVs+2v/8kg9i5Q0AP2k85H0WUOwpIVvUML8BapsMvh1XAogmQjOCsLpoKRCVQqEkQBB3NhVBcsOg==", + "license": "MIT", + "dependencies": { + "@types/katex": "^0.16.0", + "devlop": "^1.0.0", + "katex": "^0.16.0", + "micromark-factory-space": "^2.0.0", + "micromark-util-character": "^2.0.0", + "micromark-util-symbol": "^2.0.0", + "micromark-util-types": "^2.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, + "node_modules/micromark-extension-math/node_modules/commander": { + "version": "8.3.0", + "resolved": "https://registry.npmjs.org/commander/-/commander-8.3.0.tgz", + "integrity": "sha512-OkTL9umf+He2DZkUq8f8J9of7yL6RJKI24dVITBmNfZBmri9zYZQrKkuXiKhyfPSu8tUhnVBB1iKXevvnlR4Ww==", + "license": "MIT", + "engines": { + "node": ">= 12" + } + }, + "node_modules/micromark-extension-math/node_modules/katex": { + "version": "0.16.47", + "resolved": "https://registry.npmjs.org/katex/-/katex-0.16.47.tgz", + "integrity": "sha512-Eeo8Ys1doU1z+x8AZsPpQu+p/QcZBI5PeOo7QGQdy2x2m0MU/hYagBbGOmXwr5KVbEfVuWv9LpnQWeehogurjg==", + "funding": [ + "https://opencollective.com/katex", + "https://github.com/sponsors/katex" + ], + "license": "MIT", + "dependencies": { + "commander": "^8.3.0" + }, + "bin": { + "katex": "cli.js" + } + }, "node_modules/micromark-factory-destination": { "version": "2.0.1", "resolved": "https://registry.npmjs.org/micromark-factory-destination/-/micromark-factory-destination-2.0.1.tgz", @@ -6898,6 +7044,50 @@ "url": "https://opencollective.com/unified" } }, + "node_modules/rehype-katex": { + "version": "7.0.1", + "resolved": "https://registry.npmjs.org/rehype-katex/-/rehype-katex-7.0.1.tgz", + "integrity": "sha512-OiM2wrZ/wuhKkigASodFoo8wimG3H12LWQaH8qSPVJn9apWKFSH3YOCtbKpBorTVw/eI7cuT21XBbvwEswbIOA==", + "license": "MIT", + "dependencies": { + "@types/hast": "^3.0.0", + "@types/katex": "^0.16.0", + "hast-util-from-html-isomorphic": "^2.0.0", + "hast-util-to-text": "^4.0.0", + "katex": "^0.16.0", + "unist-util-visit-parents": "^6.0.0", + "vfile": "^6.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, + "node_modules/rehype-katex/node_modules/commander": { + "version": "8.3.0", + "resolved": "https://registry.npmjs.org/commander/-/commander-8.3.0.tgz", + "integrity": "sha512-OkTL9umf+He2DZkUq8f8J9of7yL6RJKI24dVITBmNfZBmri9zYZQrKkuXiKhyfPSu8tUhnVBB1iKXevvnlR4Ww==", + "license": "MIT", + "engines": { + "node": ">= 12" + } + }, + "node_modules/rehype-katex/node_modules/katex": { + "version": "0.16.47", + "resolved": "https://registry.npmjs.org/katex/-/katex-0.16.47.tgz", + "integrity": "sha512-Eeo8Ys1doU1z+x8AZsPpQu+p/QcZBI5PeOo7QGQdy2x2m0MU/hYagBbGOmXwr5KVbEfVuWv9LpnQWeehogurjg==", + "funding": [ + "https://opencollective.com/katex", + "https://github.com/sponsors/katex" + ], + "license": "MIT", + "dependencies": { + "commander": "^8.3.0" + }, + "bin": { + "katex": "cli.js" + } + }, "node_modules/rehype-raw": { "version": "7.0.0", "resolved": "https://registry.npmjs.org/rehype-raw/-/rehype-raw-7.0.0.tgz", @@ -6960,6 +7150,22 @@ "url": "https://opencollective.com/unified" } }, + "node_modules/remark-math": { + "version": "6.0.0", + "resolved": "https://registry.npmjs.org/remark-math/-/remark-math-6.0.0.tgz", + "integrity": "sha512-MMqgnP74Igy+S3WwnhQ7kqGlEerTETXMvJhrUzDikVZ2/uogJCb+WHUg97hK9/jcfc0dkD73s3LN8zU49cTEtA==", + "license": "MIT", + "dependencies": { + "@types/mdast": "^4.0.0", + "mdast-util-math": "^3.0.0", + "micromark-extension-math": "^3.0.0", + "unified": "^11.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, "node_modules/remark-parse": { "version": "11.0.0", "resolved": "https://registry.npmjs.org/remark-parse/-/remark-parse-11.0.0.tgz", @@ -7662,6 +7868,20 @@ "url": "https://opencollective.com/unified" } }, + "node_modules/unist-util-remove-position": { + "version": "5.0.0", + "resolved": "https://registry.npmjs.org/unist-util-remove-position/-/unist-util-remove-position-5.0.0.tgz", + "integrity": "sha512-Hp5Kh3wLxv0PHj9m2yZhhLt58KzPtEYKQQ4yxfYFEO7EvHwzyDYnduhHnY1mDxoqr7VUwVuHXk9RXKIiYS1N8Q==", + "license": "MIT", + "dependencies": { + "@types/unist": "^3.0.0", + "unist-util-visit": "^5.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, "node_modules/unist-util-stringify-position": { "version": "4.0.0", "resolved": "https://registry.npmjs.org/unist-util-stringify-position/-/unist-util-stringify-position-4.0.0.tgz", diff --git a/webui/package.json b/webui/package.json index c84357bdf..2bbd113a5 100644 --- a/webui/package.json +++ b/webui/package.json @@ -21,6 +21,7 @@ "date-fns": "^3.3.1", "i18next": "^25.8.14", "i18next-browser-languagedetector": "^8.2.1", + "katex": "^0.18.1", "lucide-react": "^0.562.0", "pdfjs-dist": "^6.1.200", "qrcode.react": "^4.2.0", @@ -31,10 +32,12 @@ "react-router-dom": "^7.12.0", "recharts": "^2.15.0", "rehype-highlight": "^7.0.2", + "rehype-katex": "^7.0.1", "rehype-raw": "^7.0.0", "rehype-sanitize": "^6.0.0", "remark-breaks": "^4.0.0", "remark-gfm": "^4.0.1", + "remark-math": "^6.0.0", "tailwind-merge": "^2.2.1", "zustand": "^4.5.0" }, diff --git a/webui/src/api/session.ts b/webui/src/api/session.ts index fd7fd63cb..7cda77675 100644 --- a/webui/src/api/session.ts +++ b/webui/src/api/session.ts @@ -8,6 +8,11 @@ export interface SessionMessagePartPayload { sessionID: string; type: string; text?: string; + time?: { + start: number; + end?: number; + compacted?: number; + }; synthetic?: boolean; tool?: string; state?: Record; diff --git a/webui/src/components/common/ChannelIcon.tsx b/webui/src/components/common/ChannelIcon.tsx new file mode 100644 index 000000000..e8e66431d --- /dev/null +++ b/webui/src/components/common/ChannelIcon.tsx @@ -0,0 +1,74 @@ +import { MessageSquare } from 'lucide-react'; + +type ChannelIconSize = 'xs' | 'sm' | 'md'; + +const CHANNEL_ICON_SRC: Record = { + feishu: '/channel-feishu.png', + wecom: '/channel-wecom.png', + telegram: '/channel-telegram.png', + email: '/channel-email.png', + whatsapp: '/channel-whatsapp.png', + slack: '/channel-slack.png', +}; + +const CHANNEL_MASK_ICON: Record = { + dingtalk: { src: '/channel-dingtalk-transparent.png', color: '#1677ff' }, + weixin: { src: '/channel-weixin-transparent.png', color: '#07c160' }, +}; + +export default function ChannelIcon({ + channelId, + size = 'sm', +}: { + channelId: string; + size?: ChannelIconSize; +}) { + const id = channelId.trim().toLowerCase(); + const compact = size === 'xs'; + const containerSize = size === 'md' ? 'h-10 w-10' : compact ? 'h-3.5 w-3.5' : 'h-9 w-9'; + const iconSize = size === 'md' ? 'h-7 w-7' : compact ? 'h-3.5 w-3.5' : 'h-6 w-6'; + const src = CHANNEL_ICON_SRC[id]; + const maskIcon = CHANNEL_MASK_ICON[id]; + const containerClass = compact + ? `${containerSize} inline-flex shrink-0 items-center justify-center` + : `${containerSize} flex shrink-0 items-center justify-center rounded-xl border border-gray-100 shadow-sm`; + + if (!src && !maskIcon) { + const fallbackClass = compact + ? containerClass + : `${containerSize} flex shrink-0 items-center justify-center rounded-xl bg-gray-100`; + return ( + + + + ); + } + + return ( + + {maskIcon ? ( + + ) : ( + {id} + )} + + ); +} diff --git a/webui/src/components/common/ChatDialog.tsx b/webui/src/components/common/ChatDialog.tsx index 9c5ba6292..436064b94 100644 --- a/webui/src/components/common/ChatDialog.tsx +++ b/webui/src/components/common/ChatDialog.tsx @@ -37,7 +37,13 @@ export default function ChatDialog({ }: ChatDialogProps) { const { t } = useTranslation('common'); const supportsVision = useDefaultModelVision(); - const { sessionId, createAndSend, reset } = useSessionChat({ + const { + sessionId, + pendingOptimisticMessage, + createAndSend, + consumePendingOptimisticMessage, + reset, + } = useSessionChat({ title, }); @@ -84,6 +90,8 @@ export default function ChatDialog({ emptyText={t('chat.starting')} suggestions={suggestions} supportsVision={supportsVision} + initialOptimisticMessage={pendingOptimisticMessage} + onInitialOptimisticMessageConsumed={consumePendingOptimisticMessage} onCreateAndSend={!sessionId ? (text, imageParts) => createAndSend({ text, imageParts }) : undefined} welcomeContent={!sessionId ? (
diff --git a/webui/src/components/common/ChatPromptSelectors.test.tsx b/webui/src/components/common/ChatPromptSelectors.test.tsx index 8bbe976fb..225f6e09d 100644 --- a/webui/src/components/common/ChatPromptSelectors.test.tsx +++ b/webui/src/components/common/ChatPromptSelectors.test.tsx @@ -237,6 +237,32 @@ describe('ChatModelPicker', () => { }); describe('useChatModelOptions', () => { + it('includes cache-read pricing in model options', async () => { + listDefinitionsMock.mockResolvedValue({ + data: { + models: [{ + ...makeModelDefinition(), + pricing: { + input: 0, + output: 0, + cache_read: 0.2, + unit: 1000000, + currency: 'CNY', + }, + }], + }, + }); + getResolvedMock.mockResolvedValue({ + data: { provider_id: 'provider-1', model_id: 'model-1' }, + }); + + const { result } = renderHook(() => useChatModelOptions()); + + await waitFor(() => { + expect(result.current.options[0]?.pricingLabel).toBe('¥0/¥0/¥0.2/M'); + }); + }); + it('keeps Auto opt-in and clears it when a concrete model is selected', async () => { listDefinitionsMock.mockResolvedValue({ data: { models: [makeModelDefinition()] }, diff --git a/webui/src/components/common/ChatPromptSelectors.tsx b/webui/src/components/common/ChatPromptSelectors.tsx index c39b6bec2..2adcbc87f 100644 --- a/webui/src/components/common/ChatPromptSelectors.tsx +++ b/webui/src/components/common/ChatPromptSelectors.tsx @@ -11,6 +11,7 @@ import { import { useAgents } from '@/hooks/useAgents'; import { useProviders } from '@/hooks/useProviders'; import { getAgentDisplayDescription, getAgentDisplayName, isAgentUsableInChat } from '@/utils/agentDisplay'; +import { formatPricingPerMillion, isPricingFree } from '@/utils/modelPricing'; import type { ModelDefinitionV2 } from '@/types'; export type AgentSourceFilter = 'all' | 'builtin' | 'custom'; @@ -116,9 +117,8 @@ export function useChatModelOptions({ enableAuto = false }: { enableAuto?: boole const formatPricing = (pricing: ModelDefinitionV2['pricing']): string => { if (!pricing) return t('modelPicker.noCost'); - if (pricing.input === 0 && pricing.output === 0) return t('modelPicker.free'); - const currencySymbol = pricing.currency === 'CNY' ? '¥' : '$'; - return `${currencySymbol}${pricing.input}/${currencySymbol}${pricing.output}/M`; + if (isPricingFree(pricing)) return t('modelPicker.free'); + return formatPricingPerMillion(pricing); }; const formatContextWindow = (contextWindow?: number): string => { diff --git a/webui/src/components/common/DelegateTaskCard.test.tsx b/webui/src/components/common/DelegateTaskCard.test.tsx index 3db628e49..770777c53 100644 --- a/webui/src/components/common/DelegateTaskCard.test.tsx +++ b/webui/src/components/common/DelegateTaskCard.test.tsx @@ -1,7 +1,10 @@ import { fireEvent, render, screen } from '@testing-library/react'; import { describe, expect, it, vi } from 'vitest'; -import DelegateTaskCard, { shouldRenderDelegateTaskCard } from './DelegateTaskCard'; +import DelegateTaskCard, { + extractDelegateInfo, + shouldRenderDelegateTaskCard, +} from './DelegateTaskCard'; import type { MessagePart } from '../../types'; vi.mock('react-i18next', () => ({ @@ -82,7 +85,7 @@ describe('shouldRenderDelegateTaskCard', () => { expect(shouldRenderDelegateTaskCard(part)).toBe(true); }); - it('uses persisted child session metadata as a delegate fallback', () => { + it('keeps category fallback for persisted legacy delegate records', () => { const part = { id: 'part-legacy', type: 'tool', @@ -125,6 +128,23 @@ describe('shouldRenderDelegateTaskCard', () => { }); }); +describe('extractDelegateInfo', () => { + it('uses the legacy category as the historical agent label', () => { + const info = extractDelegateInfo( + { + status: 'completed', + input: { + category: 'quick', + description: 'Legacy task', + }, + }, + 'Sub-task', + ); + + expect(info.agentName).toBe('Quick'); + }); +}); + describe('DelegateTaskCard process step', () => { it('shows delegated agent context and opens the child execution sheet', () => { const part = { diff --git a/webui/src/components/common/EntitySheet.tsx b/webui/src/components/common/EntitySheet.tsx index 355f73ff0..e327be653 100644 --- a/webui/src/components/common/EntitySheet.tsx +++ b/webui/src/components/common/EntitySheet.tsx @@ -266,6 +266,8 @@ export default function EntitySheet({ error: sessionError, create: createRexSession, createAndSend: createAndSendRex, + pendingOptimisticMessage: pendingRexOptimisticMessage, + consumePendingOptimisticMessage: consumePendingRexOptimisticMessage, retry: retryRexSession, reset: resetRexSession, } = useSessionChat({ @@ -710,6 +712,8 @@ export default function EntitySheet({ centerToolbarSlot={rexCenterToolbarSlot} composerTextareaMinHeight={rexComposerTextareaMinHeight} composerTextareaMaxHeight={rexComposerTextareaMaxHeight} + initialOptimisticMessage={pendingRexOptimisticMessage} + onInitialOptimisticMessageConsumed={consumePendingRexOptimisticMessage} onCreateAndSend={!activeRexSessionId ? (text, imageParts, agentOverride, modelOverride, options) => createAndSendRex({ text, imageParts, diff --git a/webui/src/components/common/SessionChat.test.ts b/webui/src/components/common/SessionChat.test.ts index 20788ca0f..cd91351cf 100644 --- a/webui/src/components/common/SessionChat.test.ts +++ b/webui/src/components/common/SessionChat.test.ts @@ -20,7 +20,9 @@ import { getMessageBubbleClassName, getMessageErrorText, getMessageGroupClassName, + getProcessGroupDurationMs, getRenderableThinkingText, + getThinkingFirstSentence, getRenderableFileUrl, getRegenerateTruncateTarget, getStandaloneThinkingBubbleClassName, @@ -33,6 +35,7 @@ import { shouldForwardSSEEventToParent, shouldRefetchFinishedMessage, truncateToolDisplayText, + formatProcessDuration, } from './SessionChat'; import { areChatMessagePartsRenderEqual } from './sessionChatRenderEquality'; @@ -61,6 +64,7 @@ const tMock = (key: string, options?: Record) => { 'chat.thinking': '思考中...', 'chat.streaming': '继续输出中...', 'chat.process.title': '查看 {{count}} 个步骤', + 'chat.process.duration': '已处理 {{duration}}', 'chat.process.deepThinking': '深度思考', 'chat.process.reasoningCount': '{{count}} 段思考', 'chat.process.toolCount': '{{count}} 次工具调用', @@ -453,6 +457,40 @@ describe('SessionChat message loading state', () => { }); }); +describe('SessionChat focused message deep links', () => { + it('scrolls to and consumes a rendered focus target', async () => { + const onFocusMessageConsumed = vi.fn(); + useSessionMessagesMock.mockReturnValue({ + messages: [ + makeMessage({ + id: 'target-message', + parts: [{ id: 'part-1', type: 'text', text: 'target' }], + }), + ], + loading: false, + error: null, + refetch: vi.fn(), + addMessage: vi.fn(), + updateMessage: vi.fn(), + updateMessagePart: vi.fn(), + removeMessage: vi.fn(), + clearMessages: vi.fn(), + replaceMessageText: vi.fn(), + markMessageStopped: vi.fn(), + truncateAfterMessage: vi.fn(), + }); + + render(React.createElement(SessionChat, { + sessionId: 'sess-1', + focusMessageId: 'target-message', + onFocusMessageConsumed, + })); + + await waitFor(() => expect(window.HTMLElement.prototype.scrollIntoView).toHaveBeenCalled()); + expect(onFocusMessageConsumed).toHaveBeenCalledTimes(1); + }); +}); + describe('ChatToolPart file operation titles', () => { it('shows only the filename in the header while preserving the full path in details', () => { const fullPath = '/Users/example/.flocks/workspace/outputs/2026-07-24/gold_price_retrieval_plan.md'; @@ -1480,10 +1518,16 @@ describe('SessionChat composer controls', () => { ), })); - await user.click(screen.getByRole('button', { name: '添加' })); + const addButton = screen.getByRole('button', { name: '添加' }); + await user.click(addButton); expect(screen.getByRole('menu', { name: '添加' })).toBeInTheDocument(); - expect(screen.getByText('文件和图片')).toBeInTheDocument(); + expect(addButton).not.toHaveClass('border'); + expect(addButton.className).not.toContain('shadow-'); + const filesMenuItem = screen.getByRole('menuitem', { name: '文件和图片' }); + const filesIconContainer = filesMenuItem.querySelector('svg')?.parentElement; + expect(filesIconContainer).not.toHaveClass('rounded-lg', 'border', 'bg-white'); + expect(filesIconContainer?.className).not.toContain('shadow-'); await user.click(screen.getByRole('button', { name: '智能体' })); expect(within(screen.getByLabelText('已选择的资源')).getByText('explore')).toBeInTheDocument(); await user.click(screen.getByRole('button', { name: '技能' })); @@ -1615,6 +1659,73 @@ describe('getRenderableThinkingText', () => { }); }); +describe('getThinkingFirstSentence', () => { + it('extracts the first sentence from the first non-empty line', () => { + expect(getThinkingFirstSentence('先检查上下文。再读取文件。')).toBe('先检查上下文。'); + expect(getThinkingFirstSentence('Inspect the context.\nThen read the file.')).toBe('Inspect the context.'); + expect(getThinkingFirstSentence('\n用户问了两个问题：\n1. 第一个问题')).toBe('用户问了两个问题：'); + expect(getThinkingFirstSentence('1. Inspect the context before changing code.')).toBe( + '1. Inspect the context before changing code.', + ); + expect(getThinkingFirstSentence('用户问了两个问题： 1. 查询 IP 情报。 2. 查询金价。')).toBe( + '用户问了两个问题： 1. 查询 IP 情报。', + ); + expect(getThinkingFirstSentence('Dr. Smith checks the context. Then edits.')).toBe( + 'Dr. Smith checks the context.', + ); + expect(getThinkingFirstSentence('Check inputs, etc. Then continue.')).toBe( + 'Check inputs, etc.', + ); + }); +}); + +describe('process group duration', () => { + it('uses the full wall-clock range and the current time for an active step', () => { + const parts = [ + { id: 'reason', type: 'reasoning', time: { start: 1_000, end: 2_000 } }, + { id: 'tool', type: 'tool', state: { status: 'running', time: { start: 2_500 } } }, + ] as Message['parts']; + + expect(getProcessGroupDurationMs(parts, 5_000)).toBe(4_000); + expect(formatProcessDuration(500)).toBe('1s'); + expect(formatProcessDuration(7_600)).toBe('7s'); + expect(formatProcessDuration(260_900)).toBe('4m20s'); + expect(getProcessGroupDurationMs([{ id: 'legacy', type: 'reasoning' }] as Message['parts'])) + .toBeNull(); + }); + + it('updates the displayed duration while the last process step is active', () => { + vi.useFakeTimers(); + vi.setSystemTime(5_000); + try { + render(React.createElement(ChatMessageBubble, { + message: makeMessage({ + id: 'assistant-active-duration', + role: 'assistant', + parts: [{ + id: 'tool-active-duration', + type: 'tool', + tool: 'read', + state: { + status: 'running', + input: { filePath: 'workflow.md' }, + time: { start: 1_000 }, + }, + }] as Message['parts'], + }), + isActive: true, + collapseIntermediateSteps: true, + })); + + expect(screen.getByTestId('chat-process-duration')).toHaveTextContent('已处理 4s'); + act(() => vi.advanceTimersByTime(1_000)); + expect(screen.getByTestId('chat-process-duration')).toHaveTextContent('已处理 5s'); + } finally { + vi.useRealTimers(); + } + }); +}); + describe('ChatMessageBubble reasoning streaming', () => { it.each(['reasoning', 'thinking'] as const)( 'paces an active %s part after a tool and flushes the completed text', @@ -2011,7 +2122,8 @@ describe('SessionChat intermediate process collapse', () => { messageID: 'assistant-process', sessionID: 'sess-1', type: 'reasoning', - text: '需要先读取工作流文件', + text: '需要先读取工作流文件。然后检查配置。', + time: { start: 1_000, end: 2_000 }, } as any, { id: 'tool-1', @@ -2024,6 +2136,7 @@ describe('SessionChat intermediate process collapse', () => { status: 'completed', input: { filePath: 'workflow.md' }, output: 'workflow content', + time: { start: 2_000, end: 4_500 }, }, } as any, { @@ -2053,6 +2166,7 @@ describe('SessionChat intermediate process collapse', () => { const processGroup = screen.getByTestId('chat-process-group') as HTMLDetailsElement; expect(processGroup.open).toBe(false); expect(screen.getByText('查看 2 个步骤')).toBeInTheDocument(); + expect(screen.getByTestId('chat-process-duration')).toHaveTextContent('· 已处理 3s'); expect(processGroup.querySelector('summary')).toHaveClass('text-sm', 'font-medium'); expect(processGroup.querySelector('summary')).not.toHaveClass('font-semibold'); expect(processGroup.className).not.toContain('rounded-lg'); @@ -2065,9 +2179,35 @@ describe('SessionChat intermediate process collapse', () => { expect(screen.getByTestId('chat-process-timeline')).toBeInTheDocument(); expect(screen.getByTestId('chat-process-reasoning-step')).toHaveTextContent('深度思考'); expect(screen.getByTestId('chat-process-reasoning-step').querySelector('button')).toHaveClass('text-sm'); + expect(screen.getByTestId('chat-process-reasoning-preview')).toHaveTextContent('需要先读取工作流文件。'); + expect(screen.getByTestId('chat-process-reasoning-preview')).not.toHaveTextContent('然后检查配置'); + expect(screen.getByTestId('chat-process-reasoning-preview')).not.toHaveClass('flex-1'); + expect(screen.getByTestId('chat-process-reasoning-preview').nextElementSibling).toHaveClass('lucide-chevron-down'); expect(screen.getByTestId('chat-process-tool-step')).toHaveTextContent('读取文件'); }); + it('hides the reasoning preview while the reasoning body is expanded', () => { + render(React.createElement(ChatMessageBubble, { + message: makeMessage({ + id: 'assistant-expanded-reasoning', + role: 'assistant', + parts: [{ + id: 'reason-expanded', + messageID: 'assistant-expanded-reasoning', + sessionID: 'sess-1', + type: 'reasoning', + text: '用户问了两个问题：\n1. 第一个问题', + } as any], + }), + isActive: true, + collapseIntermediateSteps: true, + processGroupsDefaultOpen: true, + })); + + expect(screen.queryByTestId('chat-process-reasoning-preview')).not.toBeInTheDocument(); + expect(screen.getByTestId('chat-process-reasoning-step')).toHaveTextContent('1. 第一个问题'); + }); + it('opens process groups while an assistant message is active and collapses after completion', () => { const activeMessage = makeMessage({ id: 'assistant-active-process', @@ -2765,6 +2905,43 @@ describe('SessionChat intermediate process collapse', () => { }); describe('SessionChat optimistic message identity', () => { + it('seeds the optimistic first message before relying on history or SSE', async () => { + const addMessage = vi.fn(); + const onConsumed = vi.fn(); + const optimisticMessage = makeMessage({ + id: 'msg_000000000001abcdefghijklmn', + sessionID: 'sess-1', + role: 'user', + parts: [{ + id: 'temp-msg_000000000001abcdefghijklmn-text', + type: 'text', + text: 'hello', + } as Message['parts'][number]], + }); + useSessionMessagesMock.mockReturnValue({ + messages: [], + loading: false, + error: null, + refetch: vi.fn(), + addMessage, + updateMessage: vi.fn(), + updateMessagePart: vi.fn(), + removeMessage: vi.fn(), + clearMessages: vi.fn(), + replaceMessageText: vi.fn(), + truncateAfterMessage: vi.fn(), + }); + + render(React.createElement(SessionChat, { + sessionId: 'sess-1', + initialOptimisticMessage: optimisticMessage, + onInitialOptimisticMessageConsumed: onConsumed, + })); + + await waitFor(() => expect(addMessage).toHaveBeenCalledWith(optimisticMessage)); + expect(onConsumed).toHaveBeenCalledWith(optimisticMessage.id); + }); + it.each([ ['prompt', 'message that fails', '/api/session/sess-1/prompt_async'], ['slash command', '/tools', '/api/session/sess-1/command'], diff --git a/webui/src/components/common/SessionChat.tsx b/webui/src/components/common/SessionChat.tsx index 693b04cca..76d19611c 100644 --- a/webui/src/components/common/SessionChat.tsx +++ b/webui/src/components/common/SessionChat.tsx @@ -161,6 +161,14 @@ export interface SessionChatProps { initialDisplayText?: string | null; /** Called immediately after initialMessage has been consumed (sent) */ onInitialMessageConsumed?: () => void; + /** Optimistic first message created while the parent is creating a session. */ + initialOptimisticMessage?: Message | null; + /** Called after initialOptimisticMessage has been seeded into local history. */ + onInitialOptimisticMessageConsumed?: (messageId: string) => void; + /** Scroll to this existing message after messages load. */ + focusMessageId?: string | null; + /** Called after focusMessageId is consumed. */ + onFocusMessageConsumed?: () => void; /** Agent name to include in prompt_async requests */ agentName?: string; /** Model override to include in prompt_async requests */ @@ -268,6 +276,79 @@ export function getRenderableThinkingText(part: Pick line.trim()).find(Boolean) || ''; + for (const match of firstLine.matchAll(/[。！？!?]|\.(?=\s|$)/g)) { + const sentenceEnd = match.index; + if (match[0] === '.' && isNonTerminalPeriod(firstLine, sentenceEnd)) continue; + return firstLine.slice(0, sentenceEnd + 1); + } + return firstLine; +} + +function getProcessPartTime(part: MessagePart): { start: number; end?: number } | undefined { + return part.type === 'tool' ? part.state?.time : part.time; +} + +export function getProcessGroupDurationMs( + parts: readonly MessagePart[], + activeNowMs?: number, +): number | null { + let firstStart = Number.POSITIVE_INFINITY; + let lastEnd = Number.NEGATIVE_INFINITY; + + for (const part of parts) { + const time = getProcessPartTime(part); + if (!time || !Number.isFinite(time.start)) continue; + const end = Number.isFinite(time.end) ? time.end : activeNowMs; + if (end === undefined || !Number.isFinite(end)) continue; + firstStart = Math.min(firstStart, time.start); + lastEnd = Math.max(lastEnd, end); + } + + if (!Number.isFinite(firstStart) || !Number.isFinite(lastEnd)) return null; + return Math.max(0, lastEnd - firstStart); +} + +export function formatProcessDuration(durationMs: number): string { + const totalSeconds = Math.max(1, Math.floor(durationMs / 1_000)); + if (totalSeconds < 60) return `${totalSeconds}s`; + const minutes = Math.floor(totalSeconds / 60); + return `${minutes}m${totalSeconds % 60}s`; +} + +function useProcessElapsedClock(enabled: boolean): number { + const [now, setNow] = useState(() => Date.now()); + + useEffect(() => { + if (!enabled) return undefined; + setNow(Date.now()); + const timer = window.setInterval(() => setNow(Date.now()), 1_000); + return () => window.clearInterval(timer); + }, [enabled]); + + return now; +} + const StreamingReasoningText = memo(function StreamingReasoningText({ content, isStreaming, @@ -1515,6 +1596,7 @@ export default function SessionChat({ onStreamingDone, initialMessage, initialDisplayText, + initialOptimisticMessage, agentName, model, executionMode = 'build', @@ -1529,6 +1611,9 @@ export default function SessionChat({ onCreateAndSend, onCreateNewSession, onInitialMessageConsumed, + onInitialOptimisticMessageConsumed, + focusMessageId, + onFocusMessageConsumed, supportsVision, toolbarSlot, composerAddMenuSlot, @@ -1695,6 +1780,7 @@ export default function SessionChat({ const messagesEndRef = useRef(null); const messagesContentRef = useRef(null); const scrollContainerRef = useRef(null); + const focusedMessageRef = useRef(''); const isAtBottomRef = useRef(true); const scrollToBottomRafRef = useRef(null); const textareaRef = useRef(null); @@ -1823,6 +1909,46 @@ export default function SessionChat({ truncateAfterMessage, } = useSessionMessages(sessionId || undefined); + + const seededOptimisticMessageIdRef = useRef(''); + useEffect(() => { + if ( + !initialOptimisticMessage + || initialOptimisticMessage.sessionID !== sessionId + || seededOptimisticMessageIdRef.current === initialOptimisticMessage.id + ) return; + + seededOptimisticMessageIdRef.current = initialOptimisticMessage.id; + addMessage(initialOptimisticMessage); + onInitialOptimisticMessageConsumed?.(initialOptimisticMessage.id); + }, [ + addMessage, + initialOptimisticMessage, + onInitialOptimisticMessageConsumed, + sessionId, + ]); + + useEffect(() => { + const targetId = String(focusMessageId || '').trim(); + if (!targetId || focusedMessageRef.current === targetId) return; + if (loading) return; + const target = messagesContentRef.current?.querySelector( + `[data-message-id="${CSS.escape(targetId)}"]`, + ); + if (!target) { + focusedMessageRef.current = targetId; + onFocusMessageConsumed?.(); + return; + } + focusedMessageRef.current = targetId; + target.scrollIntoView({ block: 'center', behavior: 'smooth' }); + target.classList.add('ring-2', 'ring-sky-400', 'ring-offset-2', 'ring-offset-white', 'dark:ring-offset-zinc-950'); + window.setTimeout(() => { + target.classList.remove('ring-2', 'ring-sky-400', 'ring-offset-2', 'ring-offset-white', 'dark:ring-offset-zinc-950'); + }, 1800); + onFocusMessageConsumed?.(); + }, [focusMessageId, loading, messages.length, onFocusMessageConsumed]); + const contextUsageMessages = contextUsageRefreshing && !contextUsageSnapshot ? [] : messages; const contextUsageBreakdown = useMemo( () => buildContextUsageBreakdown(contextUsageMessages, input, contextUsageSnapshot), @@ -3958,10 +4084,10 @@ export default function SessionChat({ aria-label={t('chat.addMenu.title')} aria-haspopup="menu" aria-expanded={showComposerAddMenu} - className={`inline-flex h-8 w-8 items-center justify-center rounded-full border transition-all duration-150 disabled:cursor-not-allowed disabled:opacity-40 ${ + className={`inline-flex h-8 w-8 items-center justify-center rounded-full transition-all duration-150 disabled:cursor-not-allowed disabled:opacity-40 ${ showComposerAddMenu - ? 'border-zinc-300 bg-white text-zinc-900 shadow-[0_2px_8px_rgba(22,27,34,0.08)] dark:border-white/[0.14] dark:bg-white/[0.09] dark:text-white' - : 'border-transparent text-zinc-500 hover:border-zinc-200 hover:bg-white hover:text-zinc-900 dark:text-zinc-400 dark:hover:border-white/[0.10] dark:hover:bg-white/[0.07] dark:hover:text-white' + ? 'bg-zinc-100 text-zinc-900 dark:bg-white/[0.09] dark:text-white' + : 'text-zinc-500 hover:bg-white hover:text-zinc-900 dark:text-zinc-400 dark:hover:bg-white/[0.07] dark:hover:text-white' }`} > @@ -3985,7 +4111,7 @@ export default function SessionChat({ }} className="group flex h-10 w-full items-center gap-2.5 rounded-[9px] px-2 text-left text-[13px] font-medium text-zinc-700 transition-colors hover:bg-zinc-100/90 hover:text-zinc-950 dark:text-zinc-200 dark:hover:bg-white/[0.07] dark:hover:text-white" > - + {t('chat.addMenu.files')} @@ -4198,33 +4324,34 @@ function ChatMessageTimelineInner({ return ( <> {items.map(({ message, isActive }) => ( - +
+ +
))} ); @@ -4311,6 +4438,7 @@ function ChatMessageBubbleInner({ const { t, i18n } = useTranslation('session'); const isUser = message.role === 'user'; const parts: MessagePart[] = Array.isArray(message.parts) ? message.parts : []; + const processElapsedClock = useProcessElapsedClock(isActive && collapseIntermediateSteps); const { getPartExpanded, togglePart } = useReasoningToggle(parts, message.finish); // Lightbox state for inline image previews. Browsers block top-level // navigation to ``data:`` URLs (the format we send for chat images), so a @@ -4558,7 +4686,15 @@ function ChatMessageBubbleInner({ )} - {t('chat.process.deepThinking')} + {t('chat.process.deepThinking')} + {!isExpanded && ( + + {getThinkingFirstSentence(thinkingText)} + + )} {isExpanded && isVisible && ( @@ -4620,6 +4756,11 @@ function ChatMessageBubbleInner({ const renderProcessGroup = (group: Array<{ part: MessagePart; index: number }>, groupIndex: number) => { const processGroupOpen = processGroupsDefaultOpen || (processGroupsOpenWhileActive && isActive); const processGroupKey = `${message.id}:process:${groupIndex}`; + const processGroupActive = isActive && group.some(({ part }) => part === activeTailPart); + const processDurationMs = getProcessGroupDurationMs( + group.map(({ part }) => part), + processGroupActive ? processElapsedClock : undefined, + ); const hasStoredOpenState = !!processGroupOpenState && Object.prototype.hasOwnProperty.call(processGroupOpenState, processGroupKey); const effectiveProcessGroupOpen = hasStoredOpenState @@ -4639,6 +4780,14 @@ function ChatMessageBubbleInner({ {t('chat.process.title', { count: group.length })} + {processDurationMs !== null && ( + + · {t('chat.process.duration', { duration: formatProcessDuration(processDurationMs) })} + + )} )} > diff --git a/webui/src/components/common/StreamingMarkdown.test.tsx b/webui/src/components/common/StreamingMarkdown.test.tsx index d4bd29f2e..b19f5f5ee 100644 --- a/webui/src/components/common/StreamingMarkdown.test.tsx +++ b/webui/src/components/common/StreamingMarkdown.test.tsx @@ -4,6 +4,7 @@ import { renderHook, act, render } from '@testing-library/react'; import { StreamingMarkdown, fallbackSplitStreamingGraphemes, + normalizeLatexDelimiters, splitStreamingGraphemes, useStreamingContent, } from './StreamingMarkdown'; @@ -311,6 +312,36 @@ describe('useStreamingContent', () => { }); describe('StreamingMarkdown', () => { + it('normalizes LaTeX delimiters only outside Markdown code', () => { + const content = [ + String.raw`Inline \(x + y\).`, + '', + String.raw`\[`, + 'x^2', + String.raw`\]`, + '', + '`\\(inlineCode\\)`', + '', + '```text', + String.raw`\[fencedCode\]`, + '```', + ].join('\n'); + + expect(normalizeLatexDelimiters(content)).toBe([ + 'Inline $x + y$.', + '', + '$$', + 'x^2', + '$$', + '', + '`\\(inlineCode\\)`', + '', + '```text', + String.raw`\[fencedCode\]`, + '```', + ].join('\n')); + }); + it('constrains rendered Markdown to its message container', () => { const { container } = render( , @@ -344,4 +375,22 @@ describe('StreamingMarkdown', () => { ); expect(container.querySelector('code')).not.toHaveClass('font-semibold'); }); + + it('renders LaTeX delimiters used in assistant messages', () => { + const { container } = render( + , + ); + + expect(container.querySelector('.katex-display')).not.toBeNull(); + expect(container.querySelectorAll('.katex')).toHaveLength(2); + expect(container.textContent).not.toContain(String.raw`\[`); + expect(container.querySelector('.katex-html')?.textContent).toContain('⊤'); + }); }); diff --git a/webui/src/components/common/StreamingMarkdown.tsx b/webui/src/components/common/StreamingMarkdown.tsx index bb4e81cd0..bd7cab827 100644 --- a/webui/src/components/common/StreamingMarkdown.tsx +++ b/webui/src/components/common/StreamingMarkdown.tsx @@ -1,11 +1,14 @@ import { memo, useState, useEffect, useRef, useCallback } from 'react'; import ReactMarkdown from 'react-markdown'; import rehypeHighlight from 'rehype-highlight'; +import rehypeKatex from 'rehype-katex'; import rehypeRaw from 'rehype-raw'; import rehypeSanitize, { defaultSchema } from 'rehype-sanitize'; import remarkBreaks from 'remark-breaks'; import remarkGfm from 'remark-gfm'; +import remarkMath from 'remark-math'; import 'highlight.js/styles/github-dark.css'; +import 'katex/dist/katex.min.css'; const sanitizeSchema = { ...defaultSchema, @@ -18,6 +21,107 @@ const MAX_STREAMING_GRAPHEMES_PER_SECOND = 360; const MAX_STREAMING_GRAPHEMES_PER_FRAME = 8; const MAX_DRAIN_ELAPSED_MS = 50; +function replaceLatexDelimiters(value: string): string { + return value + .replace(/\\\[([\s\S]*?)\\\]/g, (_, equation: string) => `$$${equation}$$`) + .replace(/\\\(([\s\S]*?)\\\)/g, (_, equation: string) => `$${equation}$`); +} + +function countRun(value: string, start: number, character: string): number { + let end = start; + while (value[end] === character) end += 1; + return end - start; +} + +function findFenceEnd( + value: string, + contentStart: number, + marker: string, + minimumLength: number, +): number { + let lineStart = contentStart; + + while (lineStart < value.length) { + let markerStart = lineStart; + while (markerStart < lineStart + 3 && value[markerStart] === ' ') markerStart += 1; + + if (value[markerStart] === marker) { + const markerLength = countRun(value, markerStart, marker); + const lineEnd = value.indexOf('\n', markerStart + markerLength); + const suffixEnd = lineEnd === -1 ? value.length : lineEnd; + const suffix = value.slice(markerStart + markerLength, suffixEnd); + + if (markerLength >= minimumLength && /^[\t ]*\r?$/.test(suffix)) { + return lineEnd === -1 ? value.length : lineEnd + 1; + } + } + + const nextLine = value.indexOf('\n', lineStart); + if (nextLine === -1) return value.length; + lineStart = nextLine + 1; + } + + return value.length; +} + +/** Convert model-style LaTeX delimiters without touching Markdown code. */ +export function normalizeLatexDelimiters(value: string): string { + let output = ''; + let plainStart = 0; + let index = 0; + + while (index < value.length) { + const isLineStart = index === 0 || value[index - 1] === '\n'; + + if (isLineStart) { + let markerStart = index; + while (markerStart < index + 3 && value[markerStart] === ' ') markerStart += 1; + const marker = value[markerStart]; + + if (marker === '`' || marker === '~') { + const markerLength = countRun(value, markerStart, marker); + if (markerLength >= 3) { + const openingLineEnd = value.indexOf('\n', markerStart + markerLength); + const codeEnd = openingLineEnd === -1 + ? value.length + : findFenceEnd(value, openingLineEnd + 1, marker, markerLength); + output += replaceLatexDelimiters(value.slice(plainStart, index)); + output += value.slice(index, codeEnd); + index = codeEnd; + plainStart = codeEnd; + continue; + } + } + } + + if (value[index] === '`') { + const markerLength = countRun(value, index, '`'); + let closingStart = index + markerLength; + + while (closingStart < value.length) { + closingStart = value.indexOf('`', closingStart); + if (closingStart === -1) break; + const closingLength = countRun(value, closingStart, '`'); + if (closingLength === markerLength) break; + closingStart += closingLength; + } + + if (closingStart !== -1) { + const codeEnd = closingStart + markerLength; + output += replaceLatexDelimiters(value.slice(plainStart, index)); + output += value.slice(index, codeEnd); + index = codeEnd; + plainStart = codeEnd; + continue; + } + } + + index += 1; + } + + return output + replaceLatexDelimiters(value.slice(plainStart)); +} + interface SegmentData { segment: string; } @@ -247,11 +351,18 @@ export interface StreamingMarkdownProps { * limiting ReactMarkdown re-parses to ~60fps instead of every SSE chunk. */ const MarkdownContent = memo(function MarkdownContent({ content }: { content: string }) { + const normalizedContent = normalizeLatexDelimiters(content); + return (
- {content} + {normalizedContent}
); diff --git a/webui/src/components/layout/Layout.test.tsx b/webui/src/components/layout/Layout.test.tsx index 8428a991b..4b35498b7 100644 --- a/webui/src/components/layout/Layout.test.tsx +++ b/webui/src/components/layout/Layout.test.tsx @@ -1,6 +1,6 @@ import React from 'react'; import { beforeEach, describe, expect, it, vi } from 'vitest'; -import { act, render, screen, waitFor, within } from '@testing-library/react'; +import { act, fireEvent, render, screen, waitFor, within } from '@testing-library/react'; import userEvent from '@testing-library/user-event'; import { MemoryRouter, Route, Routes, useLocation } from 'react-router-dom'; import Layout from './Layout'; @@ -20,6 +20,7 @@ const { getNotificationAckStatus, flocksproUsersApi, consoleUpgradeApi, + productNameContextValue, updateModalMock, useAuth, useStats, @@ -55,6 +56,18 @@ const { consoleUpgradeApi: { getProPackageStatus: vi.fn(), }, + productNameContextValue: { + productName: 'Flocks', + proProductName: 'Flocks Pro', + configuredDisplayName: null as string | null, + faviconUrl: '/favicon.svg', + hasCustomFavicon: false, + loading: false, + refreshProductName: vi.fn(), + updateProductName: vi.fn(), + uploadProductFavicon: vi.fn(), + resetProductFavicon: vi.fn(), + }, updateModalMock: vi.fn(() => null), useAuth: vi.fn(), useStats: vi.fn(), @@ -118,6 +131,10 @@ vi.mock('@/contexts/AuthContext', () => ({ useAuth, })); +vi.mock('@/contexts/ProductNameContext', () => ({ + useProductName: () => productNameContextValue, +})); + vi.mock('@/hooks/useStats', () => ({ useStats, })); @@ -242,6 +259,9 @@ describe('Layout onboarding entry', () => { vi.clearAllMocks(); vi.useRealTimers(); localStorage.clear(); + productNameContextValue.productName = 'Flocks'; + productNameContextValue.proProductName = 'Flocks Pro'; + productNameContextValue.configuredDisplayName = null; checkUpdate.mockResolvedValue({ has_update: false, @@ -590,6 +610,61 @@ describe('Layout onboarding entry', () => { expect(screen.queryByText('newVersion')).not.toBeInTheDocument(); }); + it('adapts the sidebar for a custom long workbench name and hides the product update mark', async () => { + const customName = '超长的威胁研判自动化工作台名称用于验证左侧菜单栏不会溢出'; + localStorage.setItem('flocks_onboarding_dismissed', 'true'); + productNameContextValue.productName = customName; + productNameContextValue.proProductName = customName; + productNameContextValue.configuredDisplayName = customName; + checkUpdate.mockResolvedValue({ + has_update: true, + latest_version: '2026.04.29', + current_version: '2026.04.28', + release_notes: 'Release line', + release_url: 'https://example.com/release', + error: null, + }); + + const { container } = renderHomeWithLayout(); + + const aside = container.querySelector('aside') as HTMLElement | null; + expect(aside).not.toBeNull(); + await waitFor(() => { + expect(Number.parseInt(aside!.style.width, 10)).toBeGreaterThan(208); + }); + + const sidebarShell = container.querySelector('aside > div'); + const logoRow = sidebarShell?.firstElementChild as HTMLElement | null; + expect(logoRow).not.toBeNull(); + const productLabel = within(logoRow!).getByText(customName); + expect(productLabel.className).toContain('[overflow-wrap:anywhere]'); + expect(screen.queryByText('newVersion')).not.toBeInTheDocument(); + expect(screen.queryByRole('button', { name: 'hasNewVersion v2026.04.29' })).not.toBeInTheDocument(); + }); + + it('resizes and persists the expanded desktop sidebar width', async () => { + localStorage.setItem('flocks_onboarding_dismissed', 'true'); + + const { container } = renderHomeWithLayout(); + + const aside = container.querySelector('aside') as HTMLElement | null; + expect(aside).not.toBeNull(); + const resizeHandle = await screen.findByRole('separator', { name: 'resizeNav' }); + + fireEvent.pointerDown(resizeHandle, { pointerId: 1, clientX: 208 }); + await act(async () => { + const pointerMove = new Event('pointermove') as PointerEvent; + Object.defineProperty(pointerMove, 'clientX', { value: 320 }); + window.dispatchEvent(pointerMove); + }); + await act(async () => { + window.dispatchEvent(new Event('pointerup')); + }); + + expect(aside).toHaveStyle({ width: '320px' }); + expect(localStorage.getItem('flocks_layout_sidebar_width')).toBe('320'); + }); + it('opens the account menu with settings and logout actions', async () => { const user = userEvent.setup(); const logout = vi.fn().mockResolvedValue(undefined); @@ -1267,4 +1342,78 @@ describe('Layout WebUI contract pages navigation', () => { `/sessions?session=session-soc-custom-page&message=${encodeURIComponent('workspace.socCustomPageInitialMessage')}&display=${encodeURIComponent('workspace.socCustomPageDisplayLabel')}`, ); }); + + it('customizes the SOC dashboard title from the SOC workspace menu', async () => { + const user = userEvent.setup(); + localStorage.setItem('flocks_onboarding_dismissed', 'true'); + const titleChanged = vi.fn(); + window.addEventListener('soc-dashboard:title-changed', titleChanged); + + const socPages = [ + { + id: 'soc-dashboard', + title: '告警态势', + route: '/contracts/webui/soc-dashboard', + icon: 'Activity', + order: 10, + enabled: true, + placement: 'home.after', + buildHash: 'ready', + buildStatus: 'ready' as const, + workspaceId: 'soc_ui', + workspaceTitle: 'SOC 工作区', + workspaceRoute: '/contracts/webui/workspaces/soc_ui', + }, + ]; + useWebUIContractPages.mockReturnValue({ + pages: socPages, + workspaces: [ + { + id: 'soc_ui', + title: 'SOC 工作区', + route: '/contracts/webui/workspaces/soc_ui', + icon: 'ShieldCheck', + order: 10, + enabled: true, + placement: 'sceneWorkspace', + defaultPageId: 'soc-dashboard', + sections: [ + { + id: 'posture', + label: '态势', + pageIds: ['soc-dashboard'], + defaultPageId: 'soc-dashboard', + contentPadding: 'none', + themeOverride: 'dark', + }, + ], + pages: socPages, + }, + ], + loading: false, + error: null, + refetch: vi.fn(), + }); + + try { + renderHomeWithLayout(); + + await user.click(await screen.findByRole('link', { name: 'SOC 工作区' })); + const workspaceMenu = screen.getByRole('navigation', { name: 'workspace.sectionNavigation' }); + await user.click(within(workspaceMenu).getByRole('button', { name: 'workspace.customTitle' })); + + const input = screen.getByLabelText('workspace.customTitle'); + await user.clear(input); + await user.type(input, '自定义 SOC 态势中心'); + await user.click(screen.getByRole('button', { name: 'workspace.customTitleSave' })); + + expect(localStorage.getItem('soc-dashboard-custom-title-v1')).toBe('自定义 SOC 态势中心'); + expect(titleChanged).toHaveBeenCalledTimes(1); + expect(titleChanged.mock.calls[0][0]).toMatchObject({ + detail: { title: '自定义 SOC 态势中心' }, + }); + } finally { + window.removeEventListener('soc-dashboard:title-changed', titleChanged); + } + }); }); diff --git a/webui/src/components/layout/Layout.tsx b/webui/src/components/layout/Layout.tsx index f361c9a41..5df046b9e 100644 --- a/webui/src/components/layout/Layout.tsx +++ b/webui/src/components/layout/Layout.tsx @@ -26,7 +26,7 @@ import { type LucideIcon, } from 'lucide-react'; import { useState, useEffect, useLayoutEffect, useCallback, useMemo, useRef, lazy, Suspense } from 'react'; -import type { ComponentType } from 'react'; +import type { ComponentType, CSSProperties, KeyboardEvent as ReactKeyboardEvent, PointerEvent as ReactPointerEvent } from 'react'; import { useTranslation } from 'react-i18next'; // Modals are only rendered after the user clicks/triggers them; pulling them // into the eager Layout chunk costs ~1.7k LOC + i18n keys + lucide icons that @@ -35,6 +35,13 @@ import { useTranslation } from 'react-i18next'; // would force Rollup to bundle the whole module eagerly). const ONBOARDING_DISMISSED_KEY = 'flocks_onboarding_dismissed'; const COLLAPSED_NAV_SECTIONS_KEY = 'flocks_layout_collapsed_nav_sections'; +const SIDEBAR_WIDTH_KEY = 'flocks_layout_sidebar_width'; +const SIDEBAR_DEFAULT_WIDTH = 208; +const SIDEBAR_MIN_WIDTH = 176; +const SIDEBAR_MAX_WIDTH = 520; +const SOC_DASHBOARD_TITLE_KEY = 'soc-dashboard-custom-title-v1'; +const SOC_DASHBOARD_TITLE_CHANGED_EVENT = 'soc-dashboard:title-changed'; +const SOC_DASHBOARD_TITLE_MAX_LENGTH = 64; type LazyLayoutModule = { default: ComponentType }; @@ -78,6 +85,65 @@ function saveCollapsedNavSectionIds(sectionIds: Set): void { } } +function clampSidebarWidth(width: number): number { + if (!Number.isFinite(width)) return SIDEBAR_DEFAULT_WIDTH; + return Math.min(SIDEBAR_MAX_WIDTH, Math.max(SIDEBAR_MIN_WIDTH, Math.round(width))); +} + +function estimateSidebarWidthForName(name?: string | null): number { + const normalized = (name || '').trim(); + if (!normalized) return SIDEBAR_DEFAULT_WIDTH; + + const weightedLength = Array.from(normalized).reduce((total, char) => { + return total + (char.charCodeAt(0) > 255 ? 1 : 0.58); + }, 0); + + return clampSidebarWidth(weightedLength * 18 + 112); +} + +function readSidebarWidth(): number { + try { + const rawValue = localStorage.getItem(SIDEBAR_WIDTH_KEY); + if (!rawValue) return SIDEBAR_DEFAULT_WIDTH; + const parsedValue = Number.parseInt(rawValue, 10); + return Number.isFinite(parsedValue) ? clampSidebarWidth(parsedValue) : SIDEBAR_DEFAULT_WIDTH; + } catch { + return SIDEBAR_DEFAULT_WIDTH; + } +} + +function saveSidebarWidth(width: number): void { + try { + localStorage.setItem(SIDEBAR_WIDTH_KEY, String(clampSidebarWidth(width))); + } catch { + // Local storage can be unavailable in restricted browser contexts. + } +} + +function readSocDashboardTitle(): string { + try { + return localStorage.getItem(SOC_DASHBOARD_TITLE_KEY)?.trim() || ''; + } catch { + return ''; + } +} + +function saveSocDashboardTitle(title: string): void { + const normalizedTitle = title.trim(); + try { + if (normalizedTitle) { + localStorage.setItem(SOC_DASHBOARD_TITLE_KEY, normalizedTitle); + } else { + localStorage.removeItem(SOC_DASHBOARD_TITLE_KEY); + } + } catch { + // Local storage can be unavailable in restricted browser contexts. + } + window.dispatchEvent(new CustomEvent(SOC_DASHBOARD_TITLE_CHANGED_EVENT, { + detail: { title: normalizedTitle || null }, + })); +} + const OnboardingModal = lazyLayoutComponent(() => import('@/components/common/OnboardingModal')); const UpdateModal = lazyLayoutComponent(() => import('@/components/common/UpdateModal'), ['update']); const NotificationModal = lazyLayoutComponent(() => import('@/components/common/NotificationModal'), ['notification']); @@ -189,6 +255,8 @@ export default function Layout() { const { user, logout } = useAuth(); const [sidebarOpen, setSidebarOpen] = useState(false); const [collapsed, setCollapsed] = useState(false); + const [sidebarWidth, setSidebarWidth] = useState(readSidebarWidth); + const [resizingSidebar, setResizingSidebar] = useState(false); const [accountMenuOpen, setAccountMenuOpen] = useState(false); const accountMenuRef = useRef(null); const isHome = location.pathname === '/'; @@ -198,7 +266,7 @@ export default function Layout() { const { t: tWebUIContractPage } = useTranslation('webuiContractPage'); const { t: tAuth } = useTranslation('auth'); const toast = useToast(); - const { productName, proProductName } = useProductName(); + const { productName, proProductName, configuredDisplayName } = useProductName(); const [hasUpdate, setHasUpdate] = useState(false); const [latestVersion, setLatestVersion] = useState(null); const [currentVersion, setCurrentVersion] = useState(null); @@ -225,7 +293,14 @@ export default function Layout() { const [collapsedNavSectionIds, setCollapsedNavSectionIds] = useState>(readCollapsedNavSectionIds); const [collapsedWorkspaceSectionIds, setCollapsedWorkspaceSectionIds] = useState>(() => new Set()); const [creatingWorkspaceCustomPageSession, setCreatingWorkspaceCustomPageSession] = useState(false); + const [socTitleDialogOpen, setSocTitleDialogOpen] = useState(false); + const [socTitleDraft, setSocTitleDraft] = useState(readSocDashboardTitle); const workspaceMenuCloseTimerRef = useRef(null); + const hasCustomDisplayName = Boolean(configuredDisplayName?.trim()); + const expandedSidebarWidth = sidebarWidth; + const sidebarOffsetStyle = { + '--layout-sidebar-width': `${expandedSidebarWidth}px`, + } as CSSProperties; useEffect(() => { if (!accountMenuOpen) return undefined; @@ -239,6 +314,72 @@ export default function Layout() { return () => document.removeEventListener('pointerdown', handlePointerDown); }, [accountMenuOpen]); + useEffect(() => { + if (!hasCustomDisplayName) return; + const estimatedWidth = estimateSidebarWidthForName(configuredDisplayName); + setSidebarWidth((currentWidth) => { + const nextWidth = Math.max(currentWidth, estimatedWidth); + if (nextWidth !== currentWidth) { + saveSidebarWidth(nextWidth); + } + return nextWidth; + }); + }, [configuredDisplayName, hasCustomDisplayName]); + + const updateSidebarWidth = useCallback((width: number) => { + const nextWidth = clampSidebarWidth(width); + setSidebarWidth(nextWidth); + saveSidebarWidth(nextWidth); + }, []); + + const handleSidebarResizePointerDown = useCallback((event: ReactPointerEvent) => { + if (collapsed) return; + event.preventDefault(); + setResizingSidebar(true); + + const target = event.currentTarget; + const pointerId = event.pointerId; + try { + target.setPointerCapture?.(pointerId); + } catch { + // Pointer capture is best-effort; dragging still works through window listeners. + } + + const previousCursor = document.body.style.cursor; + const previousUserSelect = document.body.style.userSelect; + document.body.style.cursor = 'col-resize'; + document.body.style.userSelect = 'none'; + + const handlePointerMove = (moveEvent: PointerEvent) => { + updateSidebarWidth(moveEvent.clientX); + }; + + const stopResize = () => { + try { + target.releasePointerCapture?.(pointerId); + } catch { + // The pointer may already be released by the browser. + } + document.body.style.cursor = previousCursor; + document.body.style.userSelect = previousUserSelect; + setResizingSidebar(false); + window.removeEventListener('pointermove', handlePointerMove); + window.removeEventListener('pointerup', stopResize); + window.removeEventListener('pointercancel', stopResize); + }; + + window.addEventListener('pointermove', handlePointerMove); + window.addEventListener('pointerup', stopResize); + window.addEventListener('pointercancel', stopResize); + }, [collapsed, updateSidebarWidth]); + + const handleSidebarResizeKeyDown = useCallback((event: ReactKeyboardEvent) => { + if (collapsed) return; + if (event.key !== 'ArrowLeft' && event.key !== 'ArrowRight') return; + event.preventDefault(); + updateSidebarWidth(sidebarWidth + (event.key === 'ArrowRight' ? 16 : -16)); + }, [collapsed, sidebarWidth, updateSidebarWidth]); + // useLayoutEffect runs synchronously before paint, so there's no flash on initial load. // It also re-runs when the user navigates back to /, covering both cases in one place. useLayoutEffect(() => { @@ -584,8 +725,9 @@ export default function Layout() { const accountInitial = (user?.username || productName || 'F').trim().charAt(0).toUpperCase(); const accountRoleLabel = user?.role === 'admin' ? tAuth('admin.roleAdmin') : tAuth('admin.roleMember'); const hasVisibleUpdate = hasUpdate && canManageUpdates; + const hasVisibleProductUpdate = hasVisibleUpdate && !hasCustomDisplayName; const showFlocksproUpgradeEntry = canManageUpdates; - const productUpdateTitle = hasVisibleUpdate + const productUpdateTitle = hasVisibleProductUpdate ? t('hasNewVersion', { version: formatUpdateVersion(latestVersion) || '' }) : productName; const settingsReturnState = { @@ -610,6 +752,7 @@ export default function Layout() { ? getLocalizedWebUIContractTitle(activeWorkspaceMenu, i18n.language) : ''; const showWorkspaceCustomPageAction = canCreateWorkspaceCustomPage && isSocWorkspace(activeWorkspaceMenu); + const showSocDashboardTitleAction = isSocWorkspace(activeWorkspaceMenu); const cancelWorkspaceMenuClose = useCallback(() => { if (workspaceMenuCloseTimerRef.current === null) return; @@ -703,6 +846,33 @@ export default function Layout() { toast, ]); + const openSocTitleDialog = useCallback(() => { + setSocTitleDraft(readSocDashboardTitle()); + setSocTitleDialogOpen(true); + setOpenWorkspaceMenuId(null); + }, []); + + const closeSocTitleDialog = useCallback(() => { + setSocTitleDialogOpen(false); + }, []); + + const handleSaveSocTitle = useCallback(() => { + saveSocDashboardTitle(socTitleDraft); + setSocTitleDialogOpen(false); + if (location.pathname === '/contracts/webui/workspaces/soc_ui/soc-dashboard') { + window.setTimeout(() => window.location.reload(), 0); + } + }, [location.pathname, socTitleDraft]); + + const handleResetSocTitle = useCallback(() => { + setSocTitleDraft(''); + saveSocDashboardTitle(''); + setSocTitleDialogOpen(false); + if (location.pathname === '/contracts/webui/workspaces/soc_ui/soc-dashboard') { + window.setTimeout(() => window.location.reload(), 0); + } + }, [location.pathname]); + const openManualUpdateCheck = useCallback(() => { setAccountMenuOpen(false); setUpdateInfo(null); @@ -742,6 +912,64 @@ export default function Layout() { + {socTitleDialogOpen && ( +
{ + if (event.target === event.currentTarget) { + closeSocTitleDialog(); + } + }} + > +
{ + event.preventDefault(); + handleSaveSocTitle(); + }} + > +

+ {tWebUIContractPage('workspace.customTitleDialogTitle')} +

+ +
+ +
+ + +
+
+
+
+ )} + {sidebarOpen && (
{/* Logo */} -
+
{collapsed ? (
- {hasVisibleUpdate && ( + {hasVisibleProductUpdate && (
+ {!collapsed && ( +
+ )} + {/* Collapse tab (desktop) */}
) : null} + + {showSocDashboardTitleAction ? ( +
+ +
+ ) : null}
)} @@ -1154,7 +1417,8 @@ export default function Layout() { {/* Main content area */}
{isFullScreenPage ? ( diff --git a/webui/src/hooks/useSessionChat.test.ts b/webui/src/hooks/useSessionChat.test.ts index 5b961b8f3..589094aa9 100644 --- a/webui/src/hooks/useSessionChat.test.ts +++ b/webui/src/hooks/useSessionChat.test.ts @@ -21,7 +21,7 @@ vi.mock('@/api/client', () => ({ }, })); -import { renderHook, act } from '@testing-library/react'; +import { renderHook, act, waitFor } from '@testing-library/react'; import { useSessionChat } from './useSessionChat'; import type { ImagePartData } from '@/utils/imageUpload'; @@ -142,10 +142,96 @@ describe('useSessionChat.createAndSend — image forwarding', () => { '/api/session/existing-session/prompt_async', { executionMode: 'build', + messageID: expect.stringMatching(/^msg_/), parts: [{ type: 'text', text: 'continue' }], }, ); }); + + it('publishes the new session and matching optimistic message only after the prompt is accepted', async () => { + let acceptPrompt!: () => void; + const promptAccepted = new Promise((resolve) => { + acceptPrompt = resolve; + }); + mockPost.mockImplementation((url: string) => { + if (url === '/api/session') return Promise.resolve({ data: { id: SESSION_ID } }); + if (url === `/api/session/${SESSION_ID}/prompt_async`) { + return promptAccepted.then(() => ({ data: {} })); + } + return Promise.resolve({ data: {} }); + }); + + const { result } = renderHook(() => useSessionChat({ title: 'Test' })); + let sendPromise!: Promise; + act(() => { + sendPromise = result.current.createAndSend({ + text: 'internal prompt', + displayText: 'visible prompt', + agent: 'rex', + }); + }); + + await waitFor(() => { + expect(mockPost).toHaveBeenCalledWith( + `/api/session/${SESSION_ID}/prompt_async`, + expect.objectContaining({ messageID: expect.stringMatching(/^msg_/) }), + ); + }); + expect(result.current.sessionId).toBeNull(); + expect(result.current.pendingOptimisticMessage).toBeNull(); + + await act(async () => { + acceptPrompt(); + await sendPromise; + }); + + const promptCall = mockPost.mock.calls.find( + ([url]) => url === `/api/session/${SESSION_ID}/prompt_async`, + ); + const messageId = promptCall?.[1]?.messageID; + expect(result.current.sessionId).toBe(SESSION_ID); + expect(result.current.pendingOptimisticMessage).toMatchObject({ + id: messageId, + sessionID: SESSION_ID, + agent: 'rex', + parts: [expect.objectContaining({ type: 'text', text: 'visible prompt' })], + }); + + act(() => { + result.current.consumePendingOptimisticMessage(messageId); + }); + expect(result.current.pendingOptimisticMessage).toBeNull(); + }); + + it('keeps the draft session inactive and does not leave an optimistic ghost when sending fails', async () => { + const sendError = new Error('prompt rejected'); + mockPost.mockImplementation((url: string) => { + if (url === '/api/session') return Promise.resolve({ data: { id: SESSION_ID } }); + if (url === `/api/session/${SESSION_ID}/prompt_async`) return Promise.reject(sendError); + return Promise.resolve({ data: {} }); + }); + + const { result } = renderHook(() => useSessionChat({ title: 'Test' })); + + await act(async () => { + await expect(result.current.createAndSend({ text: 'retry me' })).rejects.toBe(sendError); + }); + + expect(result.current.sessionId).toBeNull(); + expect(result.current.pendingOptimisticMessage).toBeNull(); + expect(mockPost.mock.calls.filter(([url]) => url === '/api/session')).toHaveLength(1); + + mockPost.mockImplementation((url: string) => { + if (url === '/api/session') return Promise.resolve({ data: { id: SESSION_ID } }); + return Promise.resolve({ data: {} }); + }); + await act(async () => { + await result.current.createAndSend({ text: 'retry me' }); + }); + + expect(result.current.sessionId).toBe(SESSION_ID); + expect(mockPost.mock.calls.filter(([url]) => url === '/api/session')).toHaveLength(1); + }); }); describe('useSessionChat — Auto session creation', () => { @@ -201,6 +287,7 @@ describe('useSessionChat — Auto session creation', () => { `/api/session/${SESSION_ID}/prompt_async`, { executionMode: 'build', + messageID: expect.stringMatching(/^msg_/), parts: [{ type: 'text', text: 'hello' }], }, ); @@ -230,6 +317,7 @@ describe('useSessionChat — Auto session creation', () => { '/api/session/existing-session/prompt_async', { executionMode: 'build', + messageID: expect.stringMatching(/^msg_/), parts: [{ type: 'text', text: 'continue' }], }, ); @@ -251,6 +339,7 @@ describe('useSessionChat — Auto session creation', () => { '/api/session/existing-session/prompt_async', { executionMode: 'plan', + messageID: expect.stringMatching(/^msg_/), parts: [{ type: 'text', text: 'inspect the implementation' }], }, ); diff --git a/webui/src/hooks/useSessionChat.ts b/webui/src/hooks/useSessionChat.ts index 295dbcf88..a884aa366 100644 --- a/webui/src/hooks/useSessionChat.ts +++ b/webui/src/hooks/useSessionChat.ts @@ -1,6 +1,8 @@ import { useState, useCallback, useRef, useEffect } from 'react'; import client from '@/api/client'; import { buildPromptParts, type ImagePartData } from '@/utils/imageUpload'; +import { createMessageId } from '@/utils/messageId'; +import type { Message } from '@/types'; import type { SessionExecutionMode } from '@/utils/sessionExecutionMode'; export interface UseSessionChatOptions { @@ -42,13 +44,14 @@ export function useSessionChat({ const [sessionId, setSessionId] = useState(initialSessionId); const [loading, setLoading] = useState(false); const [error, setError] = useState(null); + const [pendingOptimisticMessage, setPendingOptimisticMessage] = useState(null); const sessionIdRef = useRef(initialSessionId); const createPromiseRef = useRef | null>(null); const optionsRef = useRef({ title, category, modelAuto, contextMessage, welcomeMessage }); optionsRef.current = { title, category, modelAuto, contextMessage, welcomeMessage }; - const create = useCallback( + const ensureSession = useCallback( async (overrides?: Partial): Promise => { if (sessionIdRef.current) return sessionIdRef.current; // Reuse in-flight creation promise to prevent duplicates (e.g. React StrictMode double-mount) @@ -83,10 +86,7 @@ export function useSessionChat({ createPromiseRef.current = promise; try { - const sid = await promise; - sessionIdRef.current = sid; - setSessionId(sid); - return sid; + return await promise; } catch (err: unknown) { createPromiseRef.current = null; setError( @@ -100,11 +100,22 @@ export function useSessionChat({ [], ); + const create = useCallback( + async (overrides?: Partial): Promise => { + const sid = await ensureSession(overrides); + sessionIdRef.current = sid; + setSessionId(sid); + return sid; + }, + [ensureSession], + ); + useEffect(() => { if (initialSessionId === sessionIdRef.current) return; sessionIdRef.current = initialSessionId; createPromiseRef.current = null; setSessionId(initialSessionId); + setPendingOptimisticMessage(null); setLoading(false); setError(null); }, [initialSessionId]); @@ -118,6 +129,7 @@ export function useSessionChat({ sessionIdRef.current = null; createPromiseRef.current = null; setSessionId(null); + setPendingOptimisticMessage(null); setLoading(false); setError(null); }, []); @@ -136,7 +148,7 @@ export function useSessionChat({ const effectiveModelAuto = typeof createModelAuto === 'boolean' ? createModelAuto : optionsRef.current.modelAuto; - const sid = await create( + const sid = await ensureSession( typeof createModelAuto === 'boolean' ? { modelAuto: createModelAuto } : undefined, ); if (resumedExistingSession && effectiveModelAuto) { @@ -148,19 +160,70 @@ export function useSessionChat({ const payload: Record = { parts: buildPromptParts(text, imageParts), }; + const messageId = createMessageId(); + payload.messageID = messageId; if (agent) payload.agent = agent; if (model) payload.model = model; if (displayText) payload.displayText = displayText; payload.executionMode = executionMode; - client.post(`/api/session/${sid}/prompt_async`, payload).catch(() => {}); + await client.post(`/api/session/${sid}/prompt_async`, payload); + + if (!resumedExistingSession) { + const visibleText = displayText || text; + const optimisticParts: Message['parts'] = []; + if (visibleText) { + optimisticParts.push({ + id: `temp-${messageId}-text`, + type: 'text', + text: visibleText, + }); + } + imageParts?.forEach((image, index) => { + optimisticParts.push({ + id: `temp-${messageId}-img-${index}`, + type: 'file', + url: image.url, + mime: image.mime, + filename: image.filename, + }); + }); + setPendingOptimisticMessage({ + id: messageId, + sessionID: sid, + role: 'user', + parts: optimisticParts.length > 0 + ? optimisticParts + : [{ id: `temp-${messageId}-part`, type: 'text', text: visibleText }], + timestamp: Date.now(), + agent, + }); + sessionIdRef.current = sid; + setSessionId(sid); + } return sid; }, - [create], + [ensureSession], ); + const consumePendingOptimisticMessage = useCallback((messageId: string) => { + setPendingOptimisticMessage((message) => ( + message?.id === messageId ? null : message + )); + }, []); + useEffect(() => { if (autoCreate) create().catch(() => {}); }, []); - return { sessionId, loading, error, create, createAndSend, retry, reset }; + return { + sessionId, + loading, + error, + pendingOptimisticMessage, + create, + createAndSend, + consumePendingOptimisticMessage, + retry, + reset, + }; } diff --git a/webui/src/locales/en-US/model.json b/webui/src/locales/en-US/model.json index 367976e2a..569a34d51 100644 --- a/webui/src/locales/en-US/model.json +++ b/webui/src/locales/en-US/model.json @@ -114,6 +114,7 @@ "pricing": "Pricing (per 1M tokens)", "input": "Input", "output": "Output", + "cacheRead": "Cache read", "currency": "Currency", "fillModelId": "Please fill in Model ID and Name", "testingConnection": "Testing connection…", diff --git a/webui/src/locales/en-US/nav.json b/webui/src/locales/en-US/nav.json index 8c42318b3..e5fe0542f 100644 --- a/webui/src/locales/en-US/nav.json +++ b/webui/src/locales/en-US/nav.json @@ -63,6 +63,7 @@ "logout": "Log out", "expandNav": "Expand navigation", "collapseNav": "Collapse navigation", + "resizeNav": "Resize navigation", "switchLanguage": "Switch Language", "switchToDarkTheme": "Switch to dark mode", "switchToLightTheme": "Switch to light mode", diff --git a/webui/src/locales/en-US/session.json b/webui/src/locales/en-US/session.json index 4b39163cf..cbadeacf4 100644 --- a/webui/src/locales/en-US/session.json +++ b/webui/src/locales/en-US/session.json @@ -83,6 +83,8 @@ }, "workflowSession": "Workflow Session", "configSession": "Config Session", + "channelDirectChat": "Direct chat", + "channelGroupChat": "Group chat", "showHistory": "Show conversation history", "hideHistory": "Hide conversation history", "realTimeOk": "Real-time connection active", @@ -148,7 +150,7 @@ "alertOperations": "Alert Operations", "threatHunting": "Threat Hunting", "incidentResponse": "Incident Response", - "alertOperationsSuggestion": "Help me configure the SOC workspace alert operations capability. Please proceed in order:\n\n1. First confirm whether the FlocksHub component `soc-workspace` is installed. If it is not installed, ask me whether to install the SOC workspace component. If I decline, stop. If I confirm, install the component.\n2. After the component is installed, read `~/.flocks/plugins/workflows/stream_alert_denoise/guide.md` and guide me through the denoise workflow configuration. Focus on Syslog settings, whether the Syslog listener should be enabled, and whether the corresponding security device is configured and sending Syslog data.\n3. Then read `~/.flocks/plugins/workflows/stream_alert_triage/guide.md` and guide me through the triage workflow configuration. Focus on whether output should be DB or JSONL, and ask whether to enable scheduled triggering. The default schedule is every 3 minutes.\n4. After both workflows are confirmed, tell me the SOC workspace capability is configured and can be modified later if needed.\n\nAsk only one key question at a time, and show a plan and impact summary before changing configuration.", + "alertOperationsSuggestion": "Help me configure the SOC workspace alert operations capability. Please proceed in order:\n\n1. First confirm whether the FlocksHub component `soc-workspace` is installed. If it is not installed, ask me whether to install the SOC workspace component. If I decline, stop. If I confirm, install the component.\n2. After the component is installed, read `~/.flocks/plugins/workflows/stream_alert_denoise/guide.md` and guide me through the denoise workflow configuration. Focus on Syslog settings, whether the Syslog listener should be enabled, and whether the corresponding security device is configured and sending Syslog data.\n3. Then read `~/.flocks/plugins/workflows/stream_alert_triage/guide.md` and guide me through the triage workflow configuration. Focus on whether output should be DB or JSONL, and ask whether to enable scheduled triggering. The default schedule is every 5 minutes.\n4. After both workflows are confirmed, tell me the SOC workspace capability is configured and can be modified later if needed.\n\nAsk only one key question at a time, and show a plan and impact summary before changing configuration.", "socComponentInstallConfirm": "The SOC workspace component is not installed. Install it now?", "socComponentMissing": "Please make sure component:soc-workspace exists in FlocksHub.", "socComponentMissingTitle": "SOC workspace component not found", @@ -217,6 +219,7 @@ "streaming": "Streaming...", "process": { "title": "View {{count}} steps", + "duration": "Processed in {{duration}}", "deepThinking": "Deep thinking", "reasoningCount": "{{count}} reasoning", "toolCount": "{{count}} tool calls", @@ -490,6 +493,11 @@ "copyPathFailed": "Failed to copy project path", "saveFailed": "Failed to save project" }, + "openTaskSearch": "Search tasks", + "closeTaskSearch": "Close task search", + "taskSearchDialog": "Task search", + "recentTasks": "Recent tasks", + "searchResults": "Search results", "filterConversations": "Search tasks", "noResults": "No matching tasks", "showMore": "Show {{count}} more", diff --git a/webui/src/locales/en-US/webuiContractPage.json b/webui/src/locales/en-US/webuiContractPage.json index 1a8bedffb..b927794b5 100644 --- a/webui/src/locales/en-US/webuiContractPage.json +++ b/webui/src/locales/en-US/webuiContractPage.json @@ -28,6 +28,12 @@ "collapseSidebar": "Collapse sidebar", "expandSidebar": "Expand sidebar", "customPage": "Custom Page", + "customTitle": "Custom Title", + "customTitleDialogTitle": "Customize posture page title", + "customTitlePlaceholder": "Flocks AI Alert Posture Center", + "customTitleReset": "Reset", + "customTitleCancel": "Cancel", + "customTitleSave": "Save", "customPageSessionTitle": "{{workspace}} Custom Page", "customPageCreateError": "Failed to create a custom page session. Please try again later.", "socCustomPageDisplayLabel": "Create SOC Custom Page", diff --git a/webui/src/locales/zh-CN/model.json b/webui/src/locales/zh-CN/model.json index aef45a275..1d8ddef75 100644 --- a/webui/src/locales/zh-CN/model.json +++ b/webui/src/locales/zh-CN/model.json @@ -114,6 +114,7 @@ "pricing": "价格（每百万 Token）", "input": "输入", "output": "输出", + "cacheRead": "缓存命中", "currency": "货币", "fillModelId": "请填写模型 ID 和名称", "testingConnection": "正在测试连接…", diff --git a/webui/src/locales/zh-CN/nav.json b/webui/src/locales/zh-CN/nav.json index 59590281e..dcf0c076d 100644 --- a/webui/src/locales/zh-CN/nav.json +++ b/webui/src/locales/zh-CN/nav.json @@ -63,6 +63,7 @@ "logout": "退出登录", "expandNav": "展开导航", "collapseNav": "收起导航", + "resizeNav": "调整导航宽度", "switchLanguage": "切换语言", "switchToDarkTheme": "切换到深色模式", "switchToLightTheme": "切换到浅色模式", diff --git a/webui/src/locales/zh-CN/session.json b/webui/src/locales/zh-CN/session.json index 329fd3640..849fe828b 100644 --- a/webui/src/locales/zh-CN/session.json +++ b/webui/src/locales/zh-CN/session.json @@ -83,6 +83,8 @@ }, "workflowSession": "工作流会话", "configSession": "配置会话", + "channelDirectChat": "私聊", + "channelGroupChat": "群聊", "showHistory": "显示对话历史", "hideHistory": "隐藏对话历史", "realTimeOk": "实时连接正常", @@ -148,7 +150,7 @@ "alertOperations": "告警运营", "threatHunting": "威胁狩猎", "incidentResponse": "应急响应", - "alertOperationsSuggestion": "请帮助我配置 SOC 工作区的告警运营能力。请按顺序执行：\n\n1. 先确认 FlocksHub 组件 `soc-workspace` 是否已经安装；如果未安装，请询问我是否安装 SOC 工作区组件。若我拒绝则直接结束；若我确认，则安装该组件。\n2. 组件安装完成后，读取 `~/.flocks/plugins/workflows/stream_alert_denoise/guide.md`，按照 guide 引导我配置降噪工作流，重点确认 Syslog 配置、是否开启 Syslog 监听，以及对应安全设备是否已配置并开启 Syslog 数据传输。\n3. 然后读取 `~/.flocks/plugins/workflows/stream_alert_triage/guide.md`，按照 guide 引导我配置研判工作流，重点确认输出方式是写入 DB 还是 JSONL，并询问是否开启定时触发。定时触发默认每 3 分钟执行一次。\n4. 两个工作流配置都确认完成后，告知我 SOC 工作区能力配置完成；如果后续需要调整，可以自行修改工作流配置。\n\n请一次只问一个最关键问题，涉及修改配置前先给出计划和影响说明。", + "alertOperationsSuggestion": "请帮助我配置 SOC 工作区的告警运营能力。请按顺序执行：\n\n1. 先确认 FlocksHub 组件 `soc-workspace` 是否已经安装；如果未安装，请询问我是否安装 SOC 工作区组件。若我拒绝则直接结束；若我确认，则安装该组件。\n2. 组件安装完成后，读取 `~/.flocks/plugins/workflows/stream_alert_denoise/guide.md`，按照 guide 引导我配置降噪工作流，重点确认 Syslog 配置、是否开启 Syslog 监听，以及对应安全设备是否已配置并开启 Syslog 数据传输。\n3. 然后读取 `~/.flocks/plugins/workflows/stream_alert_triage/guide.md`，按照 guide 引导我配置研判工作流，重点确认输出方式是写入 DB 还是 JSONL，并询问是否开启定时触发。定时触发默认每 5 分钟执行一次。\n4. 两个工作流配置都确认完成后，告知我 SOC 工作区能力配置完成；如果后续需要调整，可以自行修改工作流配置。\n\n请一次只问一个最关键问题，涉及修改配置前先给出计划和影响说明。", "socComponentInstallConfirm": "SOC 工作区组件尚未安装，是否现在安装？", "socComponentMissing": "请先确认 FlocksHub 中存在 component:soc-workspace。", "socComponentMissingTitle": "未找到 SOC 工作区组件", @@ -217,6 +219,7 @@ "streaming": "继续输出中...", "process": { "title": "查看 {{count}} 个步骤", + "duration": "已处理 {{duration}}", "deepThinking": "深度思考", "reasoningCount": "{{count}} 段思考", "toolCount": "{{count}} 次工具调用", @@ -490,6 +493,11 @@ "copyPathFailed": "复制项目路径失败", "saveFailed": "保存项目失败" }, + "openTaskSearch": "搜索任务", + "closeTaskSearch": "关闭任务搜索", + "taskSearchDialog": "任务搜索", + "recentTasks": "最近任务", + "searchResults": "搜索结果", "filterConversations": "搜索任务", "noResults": "没有匹配的任务", "showMore": "显示更多 {{count}} 条", diff --git a/webui/src/locales/zh-CN/webuiContractPage.json b/webui/src/locales/zh-CN/webuiContractPage.json index fdebbc304..6be57f46c 100644 --- a/webui/src/locales/zh-CN/webuiContractPage.json +++ b/webui/src/locales/zh-CN/webuiContractPage.json @@ -28,6 +28,12 @@ "collapseSidebar": "折叠侧边栏", "expandSidebar": "展开侧边栏", "customPage": "自定义页面", + "customTitle": "自定义标题", + "customTitleDialogTitle": "自定义态势页标题", + "customTitlePlaceholder": "Flocks AI 智能告警态势中心", + "customTitleReset": "恢复默认", + "customTitleCancel": "取消", + "customTitleSave": "保存", "customPageSessionTitle": "{{workspace}} 自定义页面", "customPageCreateError": "无法创建自定义页面会话，请稍后重试", "socCustomPageDisplayLabel": "创建 SOC 自定义页面", diff --git a/webui/src/pages/Agent/CreateAgentChat.tsx b/webui/src/pages/Agent/CreateAgentChat.tsx index 54cd573bc..9435bf4db 100644 --- a/webui/src/pages/Agent/CreateAgentChat.tsx +++ b/webui/src/pages/Agent/CreateAgentChat.tsx @@ -50,7 +50,13 @@ export default function CreateAgentChat({ open, onClose }: CreateAgentChatProps) const { t } = useTranslation(['agent', 'common']); const supportsVision = useDefaultModelVision(); - const { sessionId, createAndSend, reset } = useSessionChat({ + const { + sessionId, + pendingOptimisticMessage, + createAndSend, + consumePendingOptimisticMessage, + reset, + } = useSessionChat({ title: t('agent:chat.createTitle'), category: 'agent', contextMessage: buildContext(), @@ -97,6 +103,8 @@ export default function CreateAgentChat({ open, onClose }: CreateAgentChatProps) className="flex-1 min-h-0" suggestions={SUGGESTIONS} supportsVision={supportsVision} + initialOptimisticMessage={pendingOptimisticMessage} + onInitialOptimisticMessageConsumed={consumePendingOptimisticMessage} onCreateAndSend={!sessionId ? (text, imageParts) => createAndSend({ text, imageParts }) : undefined} welcomeContent={!sessionId ? (
diff --git a/webui/src/pages/Channel/index.tsx b/webui/src/pages/Channel/index.tsx index 4dfb62d2a..187531911 100644 --- a/webui/src/pages/Channel/index.tsx +++ b/webui/src/pages/Channel/index.tsx @@ -26,6 +26,7 @@ import { useTranslation } from 'react-i18next'; import PageHeader from '@/components/common/PageHeader'; import LoadingSpinner from '@/components/common/LoadingSpinner'; import EmptyState from '@/components/common/EmptyState'; +import ChannelIcon from '@/components/common/ChannelIcon'; import { useToast } from '@/components/common/Toast'; import client from '@/api/client'; @@ -765,20 +766,6 @@ function Section({ // Channel Card (left panel) // ============================================================================ -const CHANNEL_ICON_SRC: Record = { - feishu: '/channel-feishu.png', - wecom: '/channel-wecom.png', - telegram: '/channel-telegram.png', - email: '/channel-email.png', - whatsapp: '/channel-whatsapp.png', - slack: '/channel-slack.png', -}; - -const CHANNEL_MASK_ICON: Record = { - dingtalk: { src: '/channel-dingtalk-transparent.png', color: '#1677ff' }, - weixin: { src: '/channel-weixin-transparent.png', color: '#07c160' }, -}; - const FEISHU_GUIDE_PDF_URL = '/feishu-bot-guide.pdf'; const FEISHU_GUIDE_PDF_FILENAME = 'feishu-bot-guide.pdf'; const WECOM_GUIDE_PDF_URL = '/wecom-bot-guide.pdf'; @@ -787,41 +774,6 @@ const DINGTALK_GUIDE_PDF_URL = '/dingtalk-channel-guide.pdf'; const DINGTALK_GUIDE_PDF_FILENAME = 'dingtalk-channel-guide.pdf'; const SLACK_APPS_URL = 'https://api.slack.com/apps'; -function getChannelIcon(id: string, size: 'sm' | 'md' = 'sm') { - const dim = size === 'md' ? 'w-10 h-10' : 'w-9 h-9'; - const imgDim = size === 'md' ? 'w-7 h-7' : 'w-6 h-6'; - const src = CHANNEL_ICON_SRC[id]; - const maskIcon = CHANNEL_MASK_ICON[id]; - return src || maskIcon ? ( -
- {maskIcon ? ( - - ) : ( - {id} - )} -
- ) : ( -
- -
- ); -} - function GuideDownloadButton({ href, download, @@ -1072,7 +1024,7 @@ function ChannelCard({ meta, config, status, isSelected, onClick }: ChannelCardP : 'border-gray-200 bg-white hover:border-gray-300 hover:bg-gray-50' }`} > - {getChannelIcon(meta.id)} +
@@ -2905,7 +2857,7 @@ function DetailHeader({ return (
-
{getChannelIcon(meta.id, 'md')}
+

{t(`channelName.${meta.id}`, { defaultValue: meta.label })}

diff --git a/webui/src/pages/DeviceIntegration/index.tsx b/webui/src/pages/DeviceIntegration/index.tsx index 061ac9646..b161451ad 100644 --- a/webui/src/pages/DeviceIntegration/index.tsx +++ b/webui/src/pages/DeviceIntegration/index.tsx @@ -16,7 +16,7 @@ import { sessionApi } from '@/api/session'; import { providerAPI } from '@/api/provider'; import { deviceAPI, type DeviceIntegration, type DeviceGroup, type DeviceTemplate, type DeviceToolInfo } from '@/api/device'; import { hubAPI } from '@/api/hub'; -import type { APIServiceCredentialField, Tool } from '@/types'; +import type { APIServiceCredentialField, Message, Tool } from '@/types'; import { toolAPI } from '@/api/tool'; import ToolDetailModal from '../Tool/components/ToolDetailModal'; import { buildCustomDeviceModeRoutingPrompt } from './customDevice'; @@ -539,6 +539,8 @@ function buildDeviceConfigRexAssistPrompt(input: DeviceConfigRexAssistInput): Cr function DeviceAddRexPanel({ templates, sessionId, + initialOptimisticMessage, + onInitialOptimisticMessageConsumed, showBuiltInTemplates, setShowBuiltInTemplates, workbenchResetToken, @@ -550,6 +552,8 @@ function DeviceAddRexPanel({ }: { templates: DeviceTemplate[]; sessionId: string | null; + initialOptimisticMessage: Message | null; + onInitialOptimisticMessageConsumed: (messageId: string) => void; showBuiltInTemplates: boolean; setShowBuiltInTemplates: (show: boolean) => void; workbenchResetToken: number; @@ -710,6 +714,8 @@ function DeviceAddRexPanel({ ; sessionId: string | null; + initialOptimisticMessage: Message | null; + onInitialOptimisticMessageConsumed: (messageId: string) => void; createAndSend: (options: CreateAndSendOptions) => Promise; rexComposerControls: ReturnType; onApplyRexDraft: (draft: DeviceAddDraft) => void; @@ -1024,6 +1034,8 @@ function AddDeviceWizardPanel({ templates={templates} instanceCounts={instanceCounts} sessionId={sessionId} + initialOptimisticMessage={initialOptimisticMessage} + onInitialOptimisticMessageConsumed={onInitialOptimisticMessageConsumed} showBuiltInTemplates={showBuiltInTemplates} setShowBuiltInTemplates={setShowBuiltInTemplates} workbenchResetToken={workbenchResetToken} @@ -1983,7 +1995,9 @@ export default function DeviceIntegrationPage() { const rexContextMessage = useMemo(() => buildDeviceAddSessionContext(templates), [templates]); const { sessionId: rexSessionId, + pendingOptimisticMessage: pendingRexOptimisticMessage, createAndSend: createAndSendRex, + consumePendingOptimisticMessage: consumePendingRexOptimisticMessage, reset: resetRexSession, } = useSessionChat({ title: t('wizard.rex.title'), @@ -2606,6 +2620,8 @@ export default function DeviceIntegrationPage() { templates={templates} instanceCounts={instanceCounts} sessionId={rexSessionId} + initialOptimisticMessage={pendingRexOptimisticMessage} + onInitialOptimisticMessageConsumed={consumePendingRexOptimisticMessage} createAndSend={createAndSendRex} rexComposerControls={rexComposerControls} onApplyRexDraft={handleApplyRexDraft} diff --git a/webui/src/pages/Model/index.test.tsx b/webui/src/pages/Model/index.test.tsx index dbcac372b..a47fc94e3 100644 --- a/webui/src/pages/Model/index.test.tsx +++ b/webui/src/pages/Model/index.test.tsx @@ -18,6 +18,9 @@ const mocks = vi.hoisted(() => ({ getSummary: vi.fn(), getResolved: vi.fn(), listDefinitions: vi.fn(), + createDefinition: vi.fn(), + getModelSettings: vi.fn(), + updateModelSettings: vi.fn(), catalogList: vi.fn(), createProvider: vi.fn(), getCredentials: vi.fn(), @@ -123,7 +126,7 @@ vi.mock('@/api/provider', () => ({ }, modelV2API: { listDefinitions: mocks.listDefinitions, - createDefinition: vi.fn(), + createDefinition: mocks.createDefinition, deleteDefinition: vi.fn(), }, usageAPI: { @@ -133,8 +136,8 @@ vi.mock('@/api/provider', () => ({ createProvider: mocks.createProvider, }, modelSettingsAPI: { - get: vi.fn(), - update: vi.fn(), + get: mocks.getModelSettings, + update: mocks.updateModelSettings, }, catalogAPI: { list: mocks.catalogList, @@ -432,7 +435,7 @@ describe('ModelPage default model selector', () => { modalities: { input: ['text', 'image'], output: ['text'] }, }, limits: { context_window: 200000, max_output_tokens: 8192 }, - pricing: { input: 1.25, output: 5, unit: 1000000, currency: 'USD' }, + pricing: { input: 1.25, output: 5, cache_read: 0.25, unit: 1000000, currency: 'USD' }, }, ]; @@ -450,6 +453,11 @@ describe('ModelPage default model selector', () => { mocks.getSummary.mockResolvedValue({ data: null }); mocks.getResolved.mockResolvedValue({ data: { provider_id: 'openai', model_id: 'gpt-4o' } }); mocks.listDefinitions.mockResolvedValue({ data: { models, total: models.length } }); + mocks.createDefinition.mockResolvedValue({ data: {} }); + mocks.getModelSettings.mockResolvedValue({ + data: { enabled: true, default_parameters: {} }, + }); + mocks.updateModelSettings.mockResolvedValue({ data: {} }); mocks.getCredentials.mockResolvedValue({ data: null }); mocks.testCredentials.mockResolvedValue({ data: { success: true, latency_ms: 10 } }); }); @@ -482,6 +490,119 @@ describe('ModelPage default model selector', () => { expect(tooltip).toHaveTextContent(/200(?:K|,?000)/); expect(tooltip).toHaveTextContent(/1\.25/); expect(tooltip).toHaveTextContent(/\b5(?:\.0+)?\b/); + expect(tooltip).toHaveTextContent('$1.25/$5/$0.25/M'); expect(tooltip).toHaveTextContent(/USD|\$/); }); + + it('shows and saves cache-read pricing in model details', async () => { + const user = userEvent.setup(); + mocks.listDefinitions.mockResolvedValue({ + data: { + models: [ + models[0], + { + ...models[1], + pricing: { + input: 1, + output: 2, + cache_read: 0.2, + unit: 1000000, + currency: 'CNY', + }, + }, + ], + total: models.length, + }, + }); + renderWithRouter(); + + await user.click(await screen.findByText('MiniMax Vision M3')); + const cacheReadLabel = await screen.findByText('form.cacheRead'); + const cacheReadInput = cacheReadLabel.parentElement?.querySelector('input'); + const inputPrice = screen.getByText('form.input').parentElement?.querySelector('input'); + const outputPrice = screen.getByText('form.output').parentElement?.querySelector('input'); + const currencySelect = screen.getByText('form.currency').parentElement?.querySelector('select'); + expect(cacheReadInput).toHaveValue(0.2); + + await user.selectOptions(currencySelect as HTMLSelectElement, 'USD'); + + expect(inputPrice).toHaveValue(0.142857); + expect(outputPrice).toHaveValue(0.285714); + expect(cacheReadInput).toHaveValue(0.028571); + + await user.selectOptions(currencySelect as HTMLSelectElement, 'CNY'); + expect(inputPrice).toHaveValue(1); + expect(outputPrice).toHaveValue(2); + expect(cacheReadInput).toHaveValue(0.2); + + await user.selectOptions(currencySelect as HTMLSelectElement, 'USD'); + + await user.click(screen.getByRole('button', { name: 'Save' })); + + await waitFor(() => { + expect(mocks.createDefinition).toHaveBeenCalledWith( + 'minimax', + expect.objectContaining({ + input_price: 0.142857, + output_price: 0.285714, + cache_read_price: 0.028571, + currency: 'USD', + }), + ); + }); + }); + + it('keeps prices unchanged when the source currency is unsupported', async () => { + const user = userEvent.setup(); + mocks.listDefinitions.mockResolvedValue({ + data: { + models: [ + models[0], + { + ...models[1], + pricing: { + input: 1, + output: 2, + cache_read: 0.2, + unit: 1000000, + currency: 'EUR', + }, + }, + ], + total: models.length, + }, + }); + renderWithRouter(); + + await user.click(await screen.findByText('MiniMax Vision M3')); + const inputPrice = screen.getByText('form.input').parentElement?.querySelector('input'); + const outputPrice = screen.getByText('form.output').parentElement?.querySelector('input'); + const cacheReadPrice = screen.getByText('form.cacheRead').parentElement?.querySelector('input'); + const currencySelect = screen.getByText('form.currency').parentElement?.querySelector('select'); + + expect(currencySelect).toHaveValue('EUR'); + await user.selectOptions(currencySelect as HTMLSelectElement, 'CNY'); + + expect(currencySelect).toHaveValue('EUR'); + expect(inputPrice).toHaveValue(1); + expect(outputPrice).toHaveValue(2); + expect(cacheReadPrice).toHaveValue(0.2); + }); + + it('sends null when the cache-read price is cleared', async () => { + const user = userEvent.setup(); + renderWithRouter(); + + await user.click(await screen.findByText('MiniMax Vision M3')); + const cacheReadPrice = screen.getByText('form.cacheRead').parentElement?.querySelector('input'); + await user.clear(cacheReadPrice as HTMLInputElement); + await user.click(screen.getByRole('button', { name: 'Save' })); + + await waitFor(() => { + expect(mocks.createDefinition).toHaveBeenCalledWith( + 'minimax', + expect.objectContaining({ cache_read_price: null }), + ); + }); + }); }); diff --git a/webui/src/pages/Model/index.tsx b/webui/src/pages/Model/index.tsx index 5d506f3d3..b7dd8b3ee 100644 --- a/webui/src/pages/Model/index.tsx +++ b/webui/src/pages/Model/index.tsx @@ -24,7 +24,9 @@ import { customAPI, modelSettingsAPI, catalogAPI, defaultModelAPI, } from '@/api/provider'; import { hasPendingProviderCredentialChanges } from './providerCredentialUtils'; +import { formatPricingPerMillion, isPricingFree } from '@/utils/modelPricing'; import { + convertCurrencyAmount, formatTokenMillions, getConvertedTotalCost, getDefaultDashboardCurrency, @@ -69,6 +71,20 @@ function isAzureProviderId(providerId: string): boolean { return AZURE_PROVIDER_IDS.has(providerId); } +function convertEditablePrice( + value: string, + sourceCurrency: string, + targetCurrency: string, +): string { + if (value.trim() === '') return value; + const amount = Number(value); + if (!Number.isFinite(amount)) return value; + const converted = convertCurrencyAmount(amount, sourceCurrency, targetCurrency); + if (converted === null) return value; + const precision = targetCurrency === 'CNY' ? 4 : 6; + return String(Number(converted.toFixed(precision))); +} + // ==================== Connection Cache ==================== const CONNECTION_CACHE_KEY = 'flocks_provider_connection_cache'; @@ -957,7 +973,6 @@ function ModelCard({ model, enabled, testStatus, onOpenDetail, onTestModel, onTo : null; const pricing = model.pricing; - const currencySymbol = pricing?.currency === 'CNY' ? '¥' : '$'; return (
{contextK && {contextK}} - {pricing && pricing.input > 0 && ( - {currencySymbol}{pricing.input}/{pricing.output}/M + {pricing && !isPricingFree(pricing) && ( + {formatPricingPerMillion(pricing)} )} - {pricing && pricing.input === 0 && pricing.output === 0 && ( + {pricing && isPricingFree(pricing) && ( {t('status.free')} )} {enabled && ( @@ -1707,13 +1722,10 @@ function AddProviderDialog({ connectedIds, onClose, onAdded }: { : `${(model.limits.context_window / 1000).toFixed(0)}K`} ctx )} - {model.pricing && model.pricing.input > 0 && ( - - {model.pricing.currency === 'CNY' ? '¥' : '$'} - {model.pricing.input}/{model.pricing.currency === 'CNY' ? '¥' : '$'}{model.pricing.output}/M - + {model.pricing && !isPricingFree(model.pricing) && ( + {formatPricingPerMillion(model.pricing)} )} - {model.pricing && model.pricing.input === 0 && ( + {model.pricing && isPricingFree(model.pricing) && ( {t('status.free')} )}
@@ -1831,14 +1843,24 @@ function useModelForm() { const [supportsReasoning, setSupportsReasoning] = useState(true); const [inputPrice, setInputPrice] = useState('0'); const [outputPrice, setOutputPrice] = useState('0'); + const [cacheReadPrice, setCacheReadPrice] = useState(''); const [currency, setCurrency] = useState('USD'); + const changeCurrency = useCallback((nextCurrency: string) => { + if (nextCurrency === currency) return; + if (convertCurrencyAmount(1, currency, nextCurrency) === null) return; + setInputPrice(value => convertEditablePrice(value, currency, nextCurrency)); + setOutputPrice(value => convertEditablePrice(value, currency, nextCurrency)); + setCacheReadPrice(value => convertEditablePrice(value, currency, nextCurrency)); + setCurrency(nextCurrency); + }, [currency]); + const reset = useCallback(() => { setModelId(''); setName(''); setContextWindow(''); setMaxOutput(''); setSupportsVision(false); setSupportsTools(true); setSupportsStreaming(true); setSupportsReasoning(true); - setInputPrice('0'); setOutputPrice('0'); setCurrency('USD'); + setInputPrice('0'); setOutputPrice('0'); setCacheReadPrice(''); setCurrency('USD'); }, []); const toPayload = useCallback(() => { @@ -1853,6 +1875,10 @@ function useModelForm() { output_price: parseFloat(outputPrice) || 0, currency, }; + const parsedCacheReadPrice = parseFloat(cacheReadPrice); + if (Number.isFinite(parsedCacheReadPrice) && parsedCacheReadPrice >= 0) { + payload.cache_read_price = parsedCacheReadPrice; + } const parsedContextWindow = parseInt(contextWindow); if (Number.isFinite(parsedContextWindow) && parsedContextWindow > 0) { payload.context_window = parsedContextWindow; @@ -1862,7 +1888,7 @@ function useModelForm() { payload.max_output_tokens = parsedMaxOutput; } return payload; - }, [modelId, name, contextWindow, maxOutput, supportsVision, supportsTools, supportsStreaming, supportsReasoning, inputPrice, outputPrice, currency]); + }, [modelId, name, contextWindow, maxOutput, supportsVision, supportsTools, supportsStreaming, supportsReasoning, inputPrice, outputPrice, cacheReadPrice, currency]); const isValid = modelId.trim() !== '' && name.trim() !== ''; @@ -1872,7 +1898,8 @@ function useModelForm() { supportsVision, setSupportsVision, supportsTools, setSupportsTools, supportsStreaming, setSupportsStreaming, supportsReasoning, setSupportsReasoning, inputPrice, setInputPrice, outputPrice, setOutputPrice, - currency, setCurrency, + cacheReadPrice, setCacheReadPrice, + currency, setCurrency: changeCurrency, reset, toPayload, isValid, }; } @@ -1950,7 +1977,7 @@ function ModelFormFields({ form, testResult, testing, modelIdPlaceholder, modelI
-
+
+
+ + form.setCacheReadPrice(e.target.value)} + className="w-full px-3 py-2 border border-gray-300 rounded-lg focus:outline-none focus:ring-2 focus:ring-slate-400 text-sm" + placeholder="—" + /> +
setInputPrice(e.target.value)} className={inputCls} /> @@ -2843,9 +2896,23 @@ function ModelDetailSheet({ setOutputPrice(e.target.value)} className={inputCls} />
+
+ + setCacheReadPrice(e.target.value)} + className={inputCls} + placeholder="—" + /> +
- handleCurrencyChange(e.target.value)} className={inputCls}> + {currency !== 'USD' && currency !== 'CNY' && ( + + )} @@ -2925,9 +2992,8 @@ function formatModelPricing( ): string { const pricing = model.pricing; if (!pricing) return unavailableLabel; - if (pricing.input === 0 && pricing.output === 0) return freeLabel; - const symbol = pricing.currency === 'CNY' ? '¥' : pricing.currency === 'USD' ? '$' : `${pricing.currency} `; - return `${symbol}${pricing.input} / ${symbol}${pricing.output} / 1M`; + if (isPricingFree(pricing)) return freeLabel; + return formatPricingPerMillion(pricing); } function ModelSelectionInfo({ model }: { model: ModelDefinitionV2 }) { diff --git a/webui/src/pages/Model/usageDisplay.test.ts b/webui/src/pages/Model/usageDisplay.test.ts index 3c02a3379..3c1a993e0 100644 --- a/webui/src/pages/Model/usageDisplay.test.ts +++ b/webui/src/pages/Model/usageDisplay.test.ts @@ -1,6 +1,7 @@ import { describe, expect, it } from 'vitest'; import { + convertCurrencyAmount, formatTokenMillions, getConvertedTotalCost, getDefaultDashboardCurrency, @@ -41,6 +42,12 @@ describe('usageDisplay helpers', () => { expect(getConvertedTotalCost(usageStats, 'USD')).toBe('$3.2500'); }); + it('converts individual amounts between USD and CNY', () => { + expect(convertCurrencyAmount(1, 'USD', 'CNY')).toBe(7); + expect(convertCurrencyAmount(7, 'CNY', 'USD')).toBe(1); + expect(convertCurrencyAmount(1, 'EUR', 'CNY')).toBeNull(); + }); + it('toggles dashboard currency', () => { expect(toggleDashboardCurrency('USD')).toBe('CNY'); expect(toggleDashboardCurrency('CNY')).toBe('USD'); diff --git a/webui/src/pages/Model/usageDisplay.ts b/webui/src/pages/Model/usageDisplay.ts index e97418acb..c59cda7a2 100644 --- a/webui/src/pages/Model/usageDisplay.ts +++ b/webui/src/pages/Model/usageDisplay.ts @@ -22,6 +22,17 @@ export function getDefaultDashboardCurrency(language: string | undefined): Dashb return language?.toLowerCase().startsWith('zh') ? 'CNY' : 'USD'; } +export function convertCurrencyAmount( + amount: number, + sourceCurrency: string, + targetCurrency: string, +): number | null { + if (sourceCurrency === targetCurrency) return amount; + if (sourceCurrency === 'USD' && targetCurrency === 'CNY') return amount * USD_TO_CNY; + if (sourceCurrency === 'CNY' && targetCurrency === 'USD') return amount / USD_TO_CNY; + return null; +} + export function getConvertedTotalCost( usageStats: UsageStats | null, targetCurrency: DashboardCurrency, @@ -32,16 +43,11 @@ export function getConvertedTotalCost( } const total = buckets.reduce((sum, bucket) => { - if (bucket.currency === targetCurrency) { - return sum + bucket.total_cost; - } - if (bucket.currency === 'USD' && targetCurrency === 'CNY') { - return sum + (bucket.total_cost * USD_TO_CNY); - } - if (bucket.currency === 'CNY' && targetCurrency === 'USD') { - return sum + (bucket.total_cost / USD_TO_CNY); - } - return sum; + return sum + (convertCurrencyAmount( + bucket.total_cost, + bucket.currency, + targetCurrency, + ) ?? 0); }, 0); if (total <= 0) { diff --git a/webui/src/pages/Session/index.test.tsx b/webui/src/pages/Session/index.test.tsx index 5db9684c1..4a116b8f1 100644 --- a/webui/src/pages/Session/index.test.tsx +++ b/webui/src/pages/Session/index.test.tsx @@ -155,6 +155,8 @@ vi.mock('@/components/common/SessionChat', () => ({ welcomeContent, initialMessage, initialDisplayText, + initialOptimisticMessage, + focusMessageId, onCreateAndSend, onSSEEvent, agentName, @@ -180,6 +182,12 @@ vi.mock('@/components/common/SessionChat', () => ({ welcomeContent?: React.ReactNode | ((setInput: (text: string) => void) => React.ReactNode); initialMessage?: string | null; initialDisplayText?: string | null; + initialOptimisticMessage?: { + id: string; + sessionID: string; + parts: Array<{ type: string; text?: string }>; + } | null; + focusMessageId?: string | null; model?: { providerID: string; modelID: string } | null; executionMode?: 'build' | 'plan' | 'goal'; onExecutionModeAccepted?: (mode: 'build' | 'plan' | 'goal') => void; @@ -220,6 +228,9 @@ vi.mock('@/components/common/SessionChat', () => ({ data-hide-input={String(Boolean(hideInput))} data-initial-message={initialMessage ?? ''} data-initial-display={initialDisplayText ?? ''} + data-optimistic-id={initialOptimisticMessage?.id ?? ''} + data-optimistic-text={initialOptimisticMessage?.parts.find((part) => part.type === 'text')?.text ?? ''} + data-focus-message={focusMessageId ?? ''} > {sessionId ?? 'no-session'} @@ -463,9 +476,12 @@ describe('SessionPage session actions menu', () => { await screen.findByRole('button', { name: 'executionMode.title' }); const agentButton = screen.getByRole('button', { name: 'chat.addMenu.agent' }); + const agentIconContainer = agentButton.querySelector('svg')?.parentElement; expect(screen.getByTestId('session-chat')).toHaveAttribute('data-execution-mode', 'build'); expect(agentButton).toHaveAttribute('aria-haspopup', 'menu'); + expect(agentIconContainer).not.toHaveClass('rounded-lg', 'border', 'bg-white'); + expect(agentIconContainer?.className).not.toContain('shadow-'); }); it('persists Plan per session', async () => { @@ -543,6 +559,50 @@ describe('SessionPage session actions menu', () => { expect(localStorage.getItem('flocks:session-execution-mode:draft')).toBeNull(); }); + it('keeps the first new-session message optimistic with the persisted message id', async () => { + const user = userEvent.setup(); + renderSessionPage(); + + await user.click(screen.getByRole('button', { name: 'mock-create-and-send' })); + + await waitFor(() => { + expect(client.post).toHaveBeenCalledWith( + '/api/session/session-2/prompt_async', + expect.objectContaining({ messageID: expect.stringMatching(/^msg_/) }), + ); + }); + + const promptCall = client.post.mock.calls.find( + ([url]) => url === '/api/session/session-2/prompt_async', + ); + const messageId = promptCall?.[1]?.messageID; + const chat = screen.getByTestId('session-chat'); + expect(chat).toHaveAttribute('data-optimistic-id', messageId); + expect(chat).toHaveAttribute('data-optimistic-text', 'hello from empty session'); + }); + + it('does not switch sessions or leave an optimistic message when the first send fails', async () => { + const user = userEvent.setup(); + client.post.mockImplementation((url: string) => { + if (url === '/api/session') return Promise.resolve({ data: secondSession }); + if (url === '/api/session/session-2/prompt_async') { + return Promise.reject(new Error('prompt rejected')); + } + return Promise.resolve({ data: {} }); + }); + renderSessionPage(); + + await user.click(screen.getByRole('button', { name: 'mock-create-and-send' })); + + await waitFor(() => { + expect(toast.error).toHaveBeenCalledWith('chat.sendFailed', 'prompt rejected'); + }); + const chat = screen.getByTestId('session-chat'); + expect(chat).toHaveTextContent('no-session'); + expect(chat).toHaveAttribute('data-optimistic-id', ''); + expect(addSession).not.toHaveBeenCalled(); + }); + it('keeps the workbench visible and shows a page refresh state while sessions load', () => { useSessions.mockReturnValue({ sessions: [], @@ -578,10 +638,18 @@ describe('SessionPage session actions menu', () => { const tasksSection = tasksHeading.closest('section'); const projectsSection = projectsHeading.closest('section'); const newSessionButton = screen.getByRole('button', { name: 'newSession' }); - const searchInput = screen.getByPlaceholderText('filterConversations'); + const searchButton = screen.getByRole('button', { name: 'openTaskSearch' }); + expect(screen.queryByPlaceholderText('filterConversations')).not.toBeInTheDocument(); + await user.click(searchButton); + const searchDialog = screen.getByRole('dialog', { name: 'taskSearchDialog' }); + const searchInput = within(searchDialog).getByPlaceholderText('filterConversations'); + expect(searchDialog).toHaveClass('fixed', 'inset-0', 'justify-center'); + expect(searchDialog.firstElementChild).toHaveClass('max-w-[620px]', 'rounded-2xl'); expect(newSessionButton.previousElementSibling).toHaveClass('left-2', 'h-3.5', 'w-3.5'); - expect(searchInput.previousElementSibling).toHaveClass('left-2', 'h-3.5', 'w-3.5'); - expect(searchInput).toHaveClass('text-sm', 'font-medium'); + expect(searchButton.closest('div')).toContainElement(screen.getByText('managementTitle')); + expect(searchInput).toHaveFocus(); + expect(searchInput).toHaveClass('text-[15px]', 'font-medium'); + await user.click(within(searchDialog).getByRole('button', { name: 'closeTaskSearch' })); expect(tasksHeading.closest('div')).toHaveClass('px-2', 'text-xs', 'text-zinc-500'); expect(projectsHeading.closest('div')).toHaveClass('px-2', 'text-xs', 'text-zinc-500'); expect(tasksHeading.nextElementSibling).toHaveTextContent('(1)'); @@ -603,7 +671,12 @@ describe('SessionPage session actions menu', () => { expect(tasksToggle).toContainElement(tasksHeading); expect(tasksToggle.querySelector('svg')).toHaveClass('h-3.5', 'w-3.5'); expect(screen.queryByRole('button', { name: 'selectTasks' })).not.toBeInTheDocument(); - expect(screen.queryByRole('button', { name: 'createTaskSession' })).not.toBeInTheDocument(); + const createTaskButton = screen.getByRole('button', { name: 'createTaskSession' }); + expect(createTaskButton).toHaveClass( + 'opacity-0', + 'group-hover/tasks-section:opacity-100', + 'group-focus-within/tasks-section:opacity-100', + ); await user.click(tasksToggle); expect(screen.queryByText('Original Session')).not.toBeInTheDocument(); @@ -612,6 +685,56 @@ describe('SessionPage session actions menu', () => { expect(screen.getByText('Original Session')).toBeInTheDocument(); }); + it('shows the five most recently updated sessions when search opens', async () => { + const user = userEvent.setup(); + const recentSessions = Array.from({ length: 6 }, (_, index) => ({ + ...session, + id: `recent-${index + 1}`, + slug: `recent-${index + 1}`, + title: `Recent ${index + 1}`, + time: { + ...session.time, + updated: session.time.updated + index, + }, + })); + useSessions.mockReturnValue({ + sessions: recentSessions, + loading: false, + error: null, + refetch: refetchSessions, + updateSessionTitle, + removeSession, + removeSessions, + addSession, + }); + + renderSessionPage(); + await user.click(screen.getByRole('button', { name: 'openTaskSearch' })); + + const searchDialog = screen.getByRole('dialog', { name: 'taskSearchDialog' }); + expect(within(searchDialog).getAllByRole('button')).toHaveLength(6); + expect(within(searchDialog).queryByText('Recent 1')).not.toBeInTheDocument(); + expect(within(searchDialog).getByText('Recent 6')).toBeInTheDocument(); + expect(within(searchDialog).getByText('recentTasks')).toBeInTheDocument(); + + await user.click(within(searchDialog).getByText('Recent 6')); + expect(screen.queryByRole('dialog', { name: 'taskSearchDialog' })).not.toBeInTheDocument(); + }); + + it('creates a new session from the tasks row', async () => { + const user = userEvent.setup(); + renderSessionPage(); + + await screen.findByText('tasksSection'); + await user.click(screen.getByRole('button', { name: 'createTaskSession' })); + + await waitFor(() => { + expect(client.post).toHaveBeenCalledWith('/api/session', { + title: 'New Session', + }); + }); + }); + it('keeps the workbench canvas, sidebar, selected row, and dark palette classes stable', async () => { renderSessionPage('/sessions?session=session-1'); @@ -938,6 +1061,54 @@ describe('SessionPage session actions menu', () => { expect(sessionCard).toHaveClass('min-h-[34px]', 'rounded-lg', 'border-transparent'); }); + it('hides the channel title prefix until the session is renamed', async () => { + const user = userEvent.setup(); + useSessions.mockReturnValue({ + sessions: [{ + ...session, + title: '[Wecom] 你能干什么事情', + channelID: 'wecom', + channelChatType: 'direct', + }, { + ...session, + id: 'session-2', + title: '[Wecom] 群聊问题', + channelID: 'wecom', + channelChatType: 'group', + }], + loading: false, + error: null, + refetch: refetchSessions, + updateSessionTitle, + removeSession, + removeSessions, + addSession, + }); + + renderSessionPage(); + + const displayTitle = await screen.findByText('你能干什么事情'); + const sessionRow = displayTitle.closest('div.group'); + expect(sessionRow).not.toBeNull(); + expect(within(sessionRow as HTMLElement).getByRole('img', { name: 'wecom' })) + .toHaveAttribute('src', '/channel-wecom.png'); + expect(within(sessionRow as HTMLElement).getByRole('img', { name: 'channelDirectChat' })) + .toHaveAttribute('data-channel-chat-type', 'direct'); + expect(screen.queryByText('[Wecom] 你能干什么事情')).not.toBeInTheDocument(); + + const groupTitle = screen.getByText('群聊问题'); + const groupRow = groupTitle.closest('div.group'); + expect(groupRow).not.toBeNull(); + expect(within(groupRow as HTMLElement).getByRole('img', { name: 'channelGroupChat' })) + .toHaveAttribute('data-channel-chat-type', 'group'); + + await user.click(within(sessionRow as HTMLElement).getByRole('button', { name: 'moreActions' })); + await user.click(screen.getByRole('button', { name: 'rename' })); + + expect(screen.getByRole('textbox', { name: 'rename' })) + .toHaveValue('[Wecom] 你能干什么事情'); + }); + it('groups legacy sessions by the effective project returned by the backend', async () => { client.get.mockResolvedValue({ data: [{ @@ -1202,6 +1373,7 @@ describe('SessionPage session actions menu', () => { renderSessionPage(); + await user.click(screen.getByRole('button', { name: 'openTaskSearch' })); await user.type(screen.getByPlaceholderText('filterConversations'), 'nothing matches'); await user.click(await screen.findByRole('button', { name: 'projectDialog.createTitle' })); const nameInput = screen.getByLabelText('projectDialog.nameLabel'); @@ -1319,7 +1491,7 @@ describe('SessionPage session actions menu', () => { const projectRow = (await screen.findByText('Shared Labs')).closest('[class*="group/project"]'); expect(projectRow).not.toBeNull(); - expect(within(projectRow as HTMLElement).queryByRole('button', { name: 'createSessionInProject' })).not.toBeInTheDocument(); + expect(within(projectRow as HTMLElement).getByRole('button', { name: 'createSessionInProject' })).toBeDisabled(); await user.click(within(projectRow as HTMLElement).getByRole('button', { name: 'projectActions' })); expect(within(projectRow as HTMLElement).getByRole('menuitem', { name: 'projectDialog.copyPathAction' })).toBeInTheDocument(); expect(within(projectRow as HTMLElement).queryByRole('menuitem', { name: 'shareAction' })).not.toBeInTheDocument(); @@ -1341,6 +1513,7 @@ describe('SessionPage session actions menu', () => { renderSessionPage(); await screen.findByText('tasksSection'); + await userEvent.setup().click(screen.getByRole('button', { name: 'openTaskSearch' })); const searchInput = screen.getByPlaceholderText('filterConversations'); fireEvent.change(searchInput, { target: { value: 'a' } }); fireEvent.change(searchInput, { target: { value: 'ab' } }); @@ -1412,18 +1585,43 @@ describe('SessionPage session actions menu', () => { }); }); - it('does not show a create-session button on a project row', async () => { + it('creates a session from a specific project row', async () => { + const user = userEvent.setup(); const currentProject = { id: 'default', worktree: '/tmp/project', name: '默认', isDefault: true }; client.get.mockResolvedValue({ data: [currentProject, { id: 'prj_project2', worktree: '/tmp/labs', name: 'Labs' }], }); + client.post.mockResolvedValue({ + data: { + ...secondSession, + id: 'session-labs', + projectID: 'prj_project2', + title: 'New Session', + }, + }); renderSessionPage(); const projectLabel = await screen.findByText('Labs'); const projectRow = projectLabel.closest('[class*="group/project"]'); expect(projectRow).not.toBeNull(); - expect(within(projectRow as HTMLElement).queryByRole('button', { name: 'createSessionInProject' })).not.toBeInTheDocument(); + const projectActionsButton = within(projectRow as HTMLElement).getByRole('button', { name: 'projectActions' }); + const createProjectSessionButton = within(projectRow as HTMLElement).getByRole('button', { + name: 'createSessionInProject', + }); + expect(projectActionsButton.nextElementSibling).toBe(createProjectSessionButton); + await user.click(createProjectSessionButton); + + await waitFor(() => { + expect(client.post).toHaveBeenCalledWith('/api/session', { + title: 'New Session', + projectID: 'prj_project2', + }); + }); + expect(addSession).toHaveBeenCalledWith(expect.objectContaining({ + id: 'session-labs', + projectID: 'prj_project2', + })); }); it('opens the actions menu for a session item', async () => { @@ -1794,6 +1992,15 @@ describe('SessionPage session actions menu', () => { ); }); + it('passes URL focusMessage to chat without treating it as an initial message', async () => { + renderSessionPage('/sessions?session=session-1&focusMessage=message-42'); + + await waitFor(() => { + expect(screen.getByTestId('session-chat')).toHaveAttribute('data-focus-message', 'message-42'); + }); + expect(screen.getByTestId('session-chat')).toHaveAttribute('data-initial-message', ''); + }); + it('starts SOC alert operations setup when the component is already installed', async () => { const user = userEvent.setup(); @@ -2818,6 +3025,11 @@ describe('SessionPage session actions menu', () => { const workflowButton = screen.getByRole('button', { name: 'chat.addMenu.workflows' }); const menuButtons = screen.getAllByRole('button'); expect(menuButtons.indexOf(skillButton)).toBeLessThan(menuButtons.indexOf(workflowButton)); + for (const button of [skillButton, workflowButton]) { + const iconContainer = button.querySelector('svg')?.parentElement; + expect(iconContainer).not.toHaveClass('rounded-lg', 'border', 'bg-white'); + expect(iconContainer?.className).not.toContain('shadow-'); + } await user.click(workflowButton); expect(screen.queryByText('security')).not.toBeInTheDocument(); diff --git a/webui/src/pages/Session/index.tsx b/webui/src/pages/Session/index.tsx index 7999b42ce..d68e2e862 100644 --- a/webui/src/pages/Session/index.tsx +++ b/webui/src/pages/Session/index.tsx @@ -6,12 +6,13 @@ import { Workflow as WorkflowIcon, Settings2, CheckSquare, MoreHorizontal, PencilLine, Download, Share2, Cpu, Info, X, Check, FolderGit2, FolderPlus, FolderOpen, Copy, ArrowUp, HardDrive, BookOpen, - Hammer, ClipboardList, Target, + Hammer, ClipboardList, Target, UserRound, UsersRound, } from 'lucide-react'; import { useTranslation } from 'react-i18next'; import { useLocation, useNavigate, useSearchParams } from 'react-router-dom'; import { getAnchoredMenuLeftOffset } from '@/components/common/ChatPromptSelectors'; import LoadingSpinner from '@/components/common/LoadingSpinner'; +import ChannelIcon from '@/components/common/ChannelIcon'; import { useToast } from '@/components/common/Toast'; import SessionChat, { buildInstructionDisplayText, type PromptDisplayOptions, type SSEChatEvent, type SSEConnectionStatus } from '@/components/common/SessionChat'; import { useSSE } from '@/hooks/useSSE'; @@ -39,7 +40,9 @@ import { buildPromptParts, type ImagePartData } from '@/utils/imageUpload'; import { getAgentDisplayDescription, getAgentDisplayName, isAgentUsableInChat } from '@/utils/agentDisplay'; import { formatRelativeTime, formatSessionDate } from '@/utils/time'; import { getWorkflowDisplayName } from '@/utils/workflowDisplay'; -import type { ModelDefinitionV2, Session } from '@/types'; +import { formatPricingPerMillion, isPricingFree } from '@/utils/modelPricing'; +import type { Message, ModelDefinitionV2, Session } from '@/types'; +import { createMessageId } from '@/utils/messageId'; import { useAuth } from '@/contexts/AuthContext'; import { DEFAULT_SESSION_EXECUTION_MODE, @@ -60,11 +63,43 @@ function sanitizeSessionExportName(value: string) { .replace(/^-|-$/g, '') || 'session'; } +function getSessionDisplayTitle(session: Pick): string { + if (!session.channelID) return session.title; + const prefix = `[${session.channelID}]`; + if (session.title.slice(0, prefix.length).toLowerCase() !== prefix.toLowerCase()) { + return session.title; + } + return session.title.slice(prefix.length).trim() || session.title; +} + +function ChannelChatTypeBadge({ + chatType, + label, +}: { + chatType: NonNullable; + label: string; +}) { + const Icon = chatType === 'direct' ? UserRound : UsersRound; + return ( + + + ); +} + const LAST_SELECTED_SESSION_STORAGE_KEY = 'flocks:last-selected-session'; const SESSION_PAGE_VISITED_STORAGE_KEY = 'flocks:sessions:visited'; const SOC_WORKSPACE_COMPONENT_ID = 'soc-workspace'; const INSTALLED_HUB_STATES = new Set(['installed', 'localOnly', 'updateAvailable']); const SESSION_UPDATE_REFETCH_DEBOUNCE_MS = 500; +const RECENT_SEARCH_SESSION_LIMIT = 5; const AUTO_MODEL_KEY = '__flocks_auto__'; const TASK_SESSION_GROUP_ID = 'tasks'; const SESSION_EXECUTION_MODES: SessionExecutionMode[] = ['build', 'plan', 'goal']; @@ -161,7 +196,7 @@ function ComposerResourcePicker({ aria-haspopup="menu" aria-expanded={open} > - + {label} @@ -384,6 +419,12 @@ function SessionSidebarItemInner({ onCancelRename, onToggleMenu, }: SessionSidebarItemProps) { + const channelChatLabel = !session.channelChatType + ? session.channelID + : session.channelChatType === 'direct' + ? t('channelDirectChat') + : t('channelGroupChat'); + return (
onSelect(session.id)} @@ -407,6 +448,20 @@ function SessionSidebarItemInner({ className="flex-shrink-0 w-3.5 h-3.5 accent-blue-500 cursor-pointer rounded" /> )} + {session.channelID && ( + + + {session.channelChatType && ( + + )} + + )} {session.category === 'workflow' && ( @@ -436,7 +491,7 @@ function SessionSidebarItemInner({ /> ) : (

- {session.title} + {getSessionDisplayTitle(session)} {session.isShared && ( {t('sharedTag')} @@ -498,6 +553,8 @@ const SessionSidebarItem = memo(SessionSidebarItemInner, (prev, next) => ( prev.nested === next.nested && prev.session.title === next.session.title && prev.session.category === next.session.category && + prev.session.channelID === next.session.channelID && + prev.session.channelChatType === next.session.channelChatType && prev.session.isShared === next.session.isShared && prev.session.time?.updated === next.session.time?.updated && prev.selected === next.selected && @@ -597,6 +654,7 @@ export default function SessionPage() { const navigate = useNavigate(); const [searchParams, setSearchParams] = useSearchParams(); const [selectedSessionId, setSelectedSessionId] = useState(null); + const [pendingFocusMessageId, setPendingFocusMessageId] = useState(null); const [sidebarCollapsed, setSidebarCollapsed] = useState(false); const [selectedAgent, setSelectedAgent] = useState('rex'); const [showAgentOptions, setShowAgentOptions] = useState(false); @@ -626,6 +684,7 @@ export default function SessionPage() { const [suiteInstallProgress, setSuiteInstallProgress] = useState(null); const [pendingInitialMessage, setPendingInitialMessage] = useState(null); const [pendingInitialDisplayText, setPendingInitialDisplayText] = useState(null); + const [pendingOptimisticMessage, setPendingOptimisticMessage] = useState(null); const [selectMode, setSelectMode] = useState(false); const [checkedIds, setCheckedIds] = useState>(new Set()); const [projects, setProjects] = useState([]); @@ -669,6 +728,7 @@ export default function SessionPage() { const [renameSubmitting, setRenameSubmitting] = useState(false); const [downloadingSessionId, setDownloadingSessionId] = useState(null); const supportsVision = useDefaultModelVision(); + const [sessionSearchOpen, setSessionSearchOpen] = useState(false); const [searchQuery, setSearchQuery] = useState(''); const [agentSourceFilter, setAgentSourceFilter] = useState('all'); const [selectedSessionFallback, setSelectedSessionFallback] = useState(null); @@ -715,6 +775,14 @@ export default function SessionPage() { projectIds: sessionProjectIds, pageSize: sessionListPageSize, }); + const searchPanelSessions = useMemo(() => { + const sortedSessions = [...sessions].sort( + (left, right) => (right.time?.updated ?? 0) - (left.time?.updated ?? 0), + ); + return searchQuery.trim() + ? sortedSessions + : sortedSessions.slice(0, RECENT_SEARCH_SESSION_LIMIT); + }, [searchQuery, sessions]); const { agents, loading: loadingAgents } = useAgents(); const { providers, loading: loadingProviders } = useProviders(); const { @@ -778,9 +846,8 @@ export default function SessionPage() { const formatPricing = (pricing: ModelDefinitionV2['pricing']): string => { if (!pricing) return t('modelPicker.noCost'); - if (pricing.input === 0 && pricing.output === 0) return t('modelPicker.free'); - const currencySymbol = pricing.currency === 'CNY' ? '¥' : '$'; - return `${currencySymbol}${pricing.input}/${currencySymbol}${pricing.output}/M`; + if (isPricingFree(pricing)) return t('modelPicker.free'); + return formatPricingPerMillion(pricing); }; const formatContextWindow = (contextWindow?: number): string => { @@ -985,6 +1052,7 @@ export default function SessionPage() { [projects, sessions, t], ); const taskGroupCollapsed = collapsedProjectIds.has(TASK_SESSION_GROUP_ID); + const taskGroupSelected = selectedProjectId === TASK_SESSION_GROUP_ID; const taskSessionsCollapsedToFirstPage = collapsedLoadedSessionGroupIds.has(TASK_SESSION_GROUP_ID); const visibleTaskSessions = taskSessionsCollapsedToFirstPage ? taskSessionGroup.sessions.slice(0, sessionListPageSize) @@ -1193,6 +1261,7 @@ export default function SessionPage() { useEffect(() => { const sessionParam = searchParams.get('session'); const messageParam = searchParams.get('message'); + const focusMessageParam = searchParams.get('focusMessage'); const displayParam = searchParams.get('display'); if (!sessionParam) return; @@ -1207,6 +1276,7 @@ export default function SessionPage() { setPendingInitialMessage(null); setPendingInitialDisplayText(null); } + setPendingFocusMessageId(focusMessageParam || null); setSearchParams({}, { replace: true }); } }, [searchParams, selectedSessionId, setSearchParams]); @@ -1488,6 +1558,10 @@ export default function SessionPage() { } }, [creating, selectedProjectId, selectedSessionId, selectedModelAuto, addSession, fetchProjects, searchQuery, toast, t]); + const handleCreateSessionInProject = useCallback((projectId: string) => { + void handleCreateSession(projectId); + }, [handleCreateSession]); + const handleSelectModel = useCallback(async (option: ChatModelOption) => { const previousModelKey = selectedModelKey; setSelectedModelKey(option.key); @@ -1543,9 +1617,39 @@ export default function SessionPage() { ...(selectedModelAuto ? { model_auto: true } : {}), }); const newSessionId = response.data.id; + const messageId = createMessageId(); + const visibleText = options?.displayText || text; + const effectiveAgent = agentOverride || selectedAgent || 'rex'; + const optimisticParts: Message['parts'] = []; + if (visibleText) { + optimisticParts.push({ + id: `temp-${messageId}-text`, + type: 'text', + text: visibleText, + }); + } + imageParts?.forEach((image, index) => { + optimisticParts.push({ + id: `temp-${messageId}-img-${index}`, + type: 'file', + url: image.url, + mime: image.mime, + filename: image.filename, + }); + }); + + const payload: Record = { + parts: buildPromptParts(text, imageParts), + messageID: messageId, + }; + if (effectiveAgent) payload.agent = effectiveAgent; + if (!selectedModelAuto && modelOverride) payload.model = modelOverride; + if (options?.displayText) payload.displayText = options.displayText; + payload.executionMode = effectiveExecutionMode; + await client.post(`/api/session/${newSessionId}/prompt_async`, payload); addSession(response.data); - await fetchProjects(undefined, searchQuery); + void fetchProjects(undefined, searchQuery).catch(() => {}); setSelectedSessionFallback(response.data); executionModeHandoffRef.current = { sessionId: newSessionId, @@ -1564,17 +1668,17 @@ export default function SessionPage() { ); } setSelectedModelKey(selectedModelAuto ? AUTO_MODEL_KEY : null); + setPendingOptimisticMessage({ + id: messageId, + sessionID: newSessionId, + role: 'user', + parts: optimisticParts.length > 0 + ? optimisticParts + : [{ id: `temp-${messageId}-part`, type: 'text', text: visibleText }], + timestamp: Date.now(), + agent: effectiveAgent, + }); setSelectedSessionId(newSessionId); - - const payload: Record = { - parts: buildPromptParts(text, imageParts), - }; - const effectiveAgent = agentOverride || selectedAgent || 'rex'; - if (effectiveAgent) payload.agent = effectiveAgent; - if (!selectedModelAuto && modelOverride) payload.model = modelOverride; - if (options?.displayText) payload.displayText = options.displayText; - payload.executionMode = effectiveExecutionMode; - await client.post(`/api/session/${newSessionId}/prompt_async`, payload); if (effectiveExecutionMode === 'goal') { setSelectedExecutionMode(DEFAULT_SESSION_EXECUTION_MODE); writeSessionExecutionMode( @@ -2033,6 +2137,14 @@ export default function SessionPage() { setSelectedSessionId(sessionId); } }, [handleToggleCheck, selectMode]); + const handleCloseSessionSearch = useCallback(() => { + setSessionSearchOpen(false); + setSearchQuery(''); + }, []); + const handleSelectSearchResult = useCallback((sessionId: string) => { + handleSelectSessionRow(sessionId); + handleCloseSessionSearch(); + }, [handleCloseSessionSearch, handleSelectSessionRow]); const handleToggleSessionMenu = useCallback((sessionId: string, trigger: HTMLElement) => { setMovePickerSessionId(null); setOpenMenuSessionId((current) => { @@ -2197,14 +2309,14 @@ export default function SessionPage() {
{/* ── Sidebar ── */}
- {/* Header：始终显示标题、新建与搜索 */} + {/* Header */}
@@ -2215,9 +2327,19 @@ export default function SessionPage() { {t('sessionCount', { count: sessions.length })}
+
-
+
{creating ? @@ -2230,29 +2352,76 @@ export default function SessionPage() { {t('newSession')}
+
+
-
- - setSearchQuery(e.target.value)} - placeholder={t('filterConversations', 'Search tasks')} - className="h-full w-full rounded-lg border-0 bg-transparent pl-9 pr-8 text-sm font-medium text-[#474b51] outline-none placeholder:text-[#474b51] focus:bg-transparent dark:text-[#c3ccd6] dark:placeholder:text-[#c3ccd6]" - /> - {searchQuery && ( + {sessionSearchOpen && ( +
+
event.stopPropagation()} + className="w-full max-w-[620px] overflow-hidden rounded-2xl border border-black/[0.12] bg-[#fdfdfc] shadow-[0_24px_70px_rgba(22,27,34,0.24)] dark:border-white/[0.11] dark:bg-[#303030] dark:shadow-[0_28px_80px_rgba(0,0,0,0.55)]" + > +
+ + setSearchQuery(event.target.value)} + onKeyDown={(event) => { + if (event.key === 'Escape') { + event.preventDefault(); + handleCloseSessionSearch(); + } + }} + placeholder={t('filterConversations', 'Search tasks')} + className="h-full w-full border-0 bg-transparent pl-12 pr-12 text-[15px] font-medium text-[#3f444a] outline-none placeholder:text-[#858a91] dark:text-[#e1e5ea] dark:placeholder:text-[#9298a0]" + /> - )} +
+
+
+ {t(searchQuery.trim() ? 'searchResults' : 'recentTasks')} +
+ {searchPanelSessions.length > 0 ? ( +
+ {searchPanelSessions.map((session) => ( + + ))} +
+ ) : ( +
+ {t('noResults')} +
+ )} +
-
+ )} {/* Session list */}
)} +
{persistedProject && openProjectMenuId === group.id && (
: } +
{!taskGroupCollapsed && (
@@ -2671,10 +2872,18 @@ export default function SessionPage() { composerTextareaMinHeight={56} initialMessage={pendingInitialMessage} initialDisplayText={pendingInitialDisplayText} + initialOptimisticMessage={pendingOptimisticMessage} + focusMessageId={pendingFocusMessageId} + onFocusMessageConsumed={() => setPendingFocusMessageId(null)} onInitialMessageConsumed={() => { setPendingInitialMessage(null); setPendingInitialDisplayText(null); }} + onInitialOptimisticMessageConsumed={(messageId) => { + setPendingOptimisticMessage((message) => ( + message?.id === messageId ? null : message + )); + }} onSseStatusChange={activeChatSessionId ? setSseStatus : undefined} onSSEEvent={handleSSEEvent} onError={handleChatError} @@ -2721,7 +2930,7 @@ export default function SessionPage() { aria-haspopup="menu" aria-expanded={showAgentOptions} > - + {t('chat.addMenu.agent')} diff --git a/webui/src/pages/WorkflowCreate/CreateChatTab.tsx b/webui/src/pages/WorkflowCreate/CreateChatTab.tsx index 41b69ee08..f626bcf75 100644 --- a/webui/src/pages/WorkflowCreate/CreateChatTab.tsx +++ b/webui/src/pages/WorkflowCreate/CreateChatTab.tsx @@ -109,7 +109,14 @@ export default function CreateChatTab({ [...guideActions, ...exampleActions] ), [exampleActions, guideActions]); - const { sessionId, error, createAndSend, retry } = useSessionChat({ + const { + sessionId, + error, + pendingOptimisticMessage, + createAndSend, + consumePendingOptimisticMessage, + retry, + } = useSessionChat({ title: t('create.chat.sessionTitle'), category: 'workflow', modelAuto: selectedModelAuto, @@ -346,6 +353,8 @@ export default function CreateChatTab({ contextWindowTokens={effectiveModelOption?.contextWindowTokens ?? null} model={selectedPromptModel} modelAuto={selectedModelAuto} + initialOptimisticMessage={pendingOptimisticMessage} + onInitialOptimisticMessageConsumed={consumePendingOptimisticMessage} onStreamingDone={handleStreamingDone} onSSEEvent={handleSSEEvent} onCreateAndSend={!sessionId ? handleCreateAndSend : undefined} diff --git a/webui/src/pages/WorkflowDetail/tabs/ChatTab.tsx b/webui/src/pages/WorkflowDetail/tabs/ChatTab.tsx index 192c23c61..6f9735fef 100644 --- a/webui/src/pages/WorkflowDetail/tabs/ChatTab.tsx +++ b/webui/src/pages/WorkflowDetail/tabs/ChatTab.tsx @@ -202,6 +202,8 @@ export default function ChatTab({ error, create: createSession, createAndSend: createAndSendSession, + pendingOptimisticMessage, + consumePendingOptimisticMessage, reset: resetSession, } = useSessionChat({ title: t('detail.chat.sessionTitle', { name: workflowDisplayName }), @@ -595,6 +597,8 @@ export default function ChatTab({ onNodeRefDismiss={onNodeRefDismiss} onStreamingDone={handleStreamingDone} initialMessage={initialMessage} + initialOptimisticMessage={pendingOptimisticMessage} + onInitialOptimisticMessageConsumed={consumePendingOptimisticMessage} onSSEEvent={handleSSEEvent} supportsVision={effectiveSupportsVision} contextWindowTokens={effectiveModelOption?.contextWindowTokens ?? null} diff --git a/webui/src/types/index.ts b/webui/src/types/index.ts index 62a72882e..91e7ff1f5 100644 --- a/webui/src/types/index.ts +++ b/webui/src/types/index.ts @@ -19,6 +19,10 @@ export interface Session { revert?: SessionRevert; /** Session category: 'user' | 'workflow' | 'task' | 'entity-config' | ... */ category?: string; + /** Messaging channel bound to this session, when it originated from IM. */ + channelID?: string; + /** Conversation type reported by the messaging channel binding. */ + channelChatType?: 'direct' | 'group' | 'channel'; status?: 'active' | 'archived'; provider?: string; model?: string; @@ -136,6 +140,11 @@ export interface MessagePart { toolCall?: ToolCall; toolResult?: ToolResult; thinking?: string; + time?: { + start: number; + end?: number; + compacted?: number; + }; image?: { url: string; alt?: string; @@ -490,6 +499,7 @@ export interface ModelCapabilitiesV2 { export interface ModelLimitsV2 { context_window: number; + max_input_tokens?: number; max_output_tokens: number; } @@ -586,6 +596,7 @@ export interface CustomModelCreate { supports_reasoning?: boolean; input_price?: number; output_price?: number; + cache_read_price?: number | null; currency?: string; } @@ -602,6 +613,7 @@ export interface CustomModelInfo { supports_reasoning: boolean; input_price: number; output_price: number; + cache_read_price?: number | null; currency: string; created_at: string; } @@ -659,11 +671,14 @@ export interface CatalogModel { }; limits?: { context_window: number; + max_input_tokens?: number; max_output_tokens: number; }; pricing?: { input: number; output: number; + cache_read?: number; + cache_write?: number; currency: string; }; } diff --git a/webui/src/utils/modelPricing.test.ts b/webui/src/utils/modelPricing.test.ts new file mode 100644 index 000000000..44b0142b1 --- /dev/null +++ b/webui/src/utils/modelPricing.test.ts @@ -0,0 +1,31 @@ +import { describe, expect, it } from 'vitest'; + +import { formatPricingPerMillion, isPricingFree } from './modelPricing'; + +describe('formatPricingPerMillion', () => { + it('formats input and output prices', () => { + expect(formatPricingPerMillion({ + input: 1, + output: 2, + currency: 'CNY', + })).toBe('¥1/¥2/M'); + }); + + it('includes the cache-read price when configured', () => { + expect(formatPricingPerMillion({ + input: 1, + output: 2, + cache_read: 0.2, + currency: 'CNY', + })).toBe('¥1/¥2/¥0.2/M'); + }); + + it('does not mark cache-only pricing as free', () => { + expect(isPricingFree({ + input: 0, + output: 0, + cache_read: 0.2, + currency: 'CNY', + })).toBe(false); + }); +}); diff --git a/webui/src/utils/modelPricing.ts b/webui/src/utils/modelPricing.ts new file mode 100644 index 000000000..baf3faa28 --- /dev/null +++ b/webui/src/utils/modelPricing.ts @@ -0,0 +1,25 @@ +type PricingPerMillion = { + input: number; + output: number; + cache_read?: number | null; + cache_write?: number | null; + currency: string; +}; + +export function isPricingFree(pricing: PricingPerMillion): boolean { + return pricing.input === 0 + && pricing.output === 0 + && (pricing.cache_read ?? 0) === 0 + && (pricing.cache_write ?? 0) === 0; +} + +export function formatPricingPerMillion(pricing: PricingPerMillion): string { + const symbol = pricing.currency === 'CNY' + ? '¥' + : pricing.currency === 'USD' + ? '$' + : `${pricing.currency} `; + const prices = [pricing.input, pricing.output]; + if (pricing.cache_read != null) prices.push(pricing.cache_read); + return `${prices.map(price => `${symbol}${price}`).join('/')}/M`; +}