diff --git a/.flocks/flocks.json.example b/.flocks/flocks.json.example index c60d4d838..1f8b273c7 100644 --- a/.flocks/flocks.json.example +++ b/.flocks/flocks.json.example @@ -39,6 +39,9 @@ "readMaxBytes": 51200, "readMaxLineLength": 2000 }, + "toolFailure": { + "disableOnRepeatedFailure": true + }, "sandbox": { "mode": "off", "scope": "agent", diff --git a/.flocks/flockshub/index.json b/.flocks/flockshub/index.json index 226009aef..7ac3a1117 100644 --- a/.flocks/flockshub/index.json +++ b/.flocks/flockshub/index.json @@ -14573,7 +14573,7 @@ "type": "device", "name": "Chaitin SafeLine WAF", "description": "Chaitin SafeLine WAF OpenAPI integration.", - "version": "1.0.0", + "version": "1.0.1", "category": "integration", "tags": [ "waf", diff --git a/.flocks/flockshub/plugins/tools/device/chaitin_safeline_waf_v1_0_0/manifest.json b/.flocks/flockshub/plugins/tools/device/chaitin_safeline_waf_v1_0_0/manifest.json index 065286060..11ec79673 100644 --- a/.flocks/flockshub/plugins/tools/device/chaitin_safeline_waf_v1_0_0/manifest.json +++ b/.flocks/flockshub/plugins/tools/device/chaitin_safeline_waf_v1_0_0/manifest.json @@ -5,7 +5,7 @@ "name": "Chaitin SafeLine WAF", "description": "Chaitin SafeLine WAF OpenAPI integration.", "descriptionCn": "长亭雷池 WAF OpenAPI 接入。", - "version": "1.0.0", + "version": "1.0.1", "author": "Flocks Team", "license": "MIT", "category": "integration", diff --git a/.flocks/flockshub/plugins/tools/device/huaweicloud_waf_v39/hw_waf.handler.py b/.flocks/flockshub/plugins/tools/device/huaweicloud_waf_v39/hw_waf.handler.py index fd0e4409b..5f6739eb0 100644 --- a/.flocks/flockshub/plugins/tools/device/huaweicloud_waf_v39/hw_waf.handler.py +++ b/.flocks/flockshub/plugins/tools/device/huaweicloud_waf_v39/hw_waf.handler.py @@ -267,10 +267,10 @@ async def _request( if resp.status >= 400: return ToolResult( success=False, - data=resp_json, + output=resp_json, error=f"HTTP {resp.status}: {resp_text[:300]}", ) - return ToolResult(success=True, data=resp_json) + return ToolResult(success=True, output=resp_json) def _pick(params: dict[str, Any], *keys: str) -> dict[str, Any]: @@ -291,10 +291,9 @@ def _page_query(params: dict[str, Any]) -> dict[str, Any]: # --------------------------------------------------------------------------- -async def host(params: dict[str, Any], ctx: ToolContext) -> ToolResult: +async def host(ctx: ToolContext, action: str, **params: Any) -> ToolResult: cfg = _load_config(params.get("enterprise_project_id")) pid = cfg.project_id - action = params.get("action", "") if action == "host_list": q = _page_query(params) @@ -360,10 +359,9 @@ async def host(params: dict[str, Any], ctx: ToolContext) -> ToolResult: return ToolResult(success=False, error=f"Unknown action: {action}") -async def policy(params: dict[str, Any], ctx: ToolContext) -> ToolResult: +async def policy(ctx: ToolContext, action: str, **params: Any) -> ToolResult: cfg = _load_config(params.get("enterprise_project_id")) pid = cfg.project_id - action = params.get("action", "") if action == "policy_list": q = _page_query(params) @@ -400,7 +398,9 @@ async def policy(params: dict[str, Any], ctx: ToolContext) -> ToolResult: if action == "cc_rule_create": pol_id = params["policy_id"] - body = _pick(params, "url", "limit_num", "limit_period", "lock_time", "tag_type", "action") + body = _pick(params, "url", "limit_num", "limit_period", "lock_time", "tag_type") + if params.get("rule_action") is not None: + body["action"] = params["rule_action"] return await _request(cfg, "POST", f"/v1/{pid}/waf/policy/{pol_id}/cc", body=body) if action == "cc_rule_delete": @@ -415,7 +415,9 @@ async def policy(params: dict[str, Any], ctx: ToolContext) -> ToolResult: if action == "custom_rule_create": pol_id = params["policy_id"] - body = _pick(params, "name", "conditions", "action", "priority", "description") + body = _pick(params, "name", "conditions", "priority", "description") + if params.get("rule_action") is not None: + body["action"] = params["rule_action"] return await _request(cfg, "POST", f"/v1/{pid}/waf/policy/{pol_id}/custom", body=body) if action == "custom_rule_delete": @@ -446,16 +448,17 @@ async def policy(params: dict[str, Any], ctx: ToolContext) -> ToolResult: return ToolResult(success=False, error=f"Unknown action: {action}") -async def event(params: dict[str, Any], ctx: ToolContext) -> ToolResult: +async def event(ctx: ToolContext, action: str, **params: Any) -> ToolResult: cfg = _load_config(params.get("enterprise_project_id")) pid = cfg.project_id - action = params.get("action", "") if action == "event_list": q = _page_query(params) - for k in ("from", "to", "hosts", "attacks", "action"): + for k in ("from", "to", "hosts", "attacks"): if params.get(k) is not None: q[k] = params[k] + if params.get("event_action") is not None: + q["action"] = params["event_action"] return await _request(cfg, "GET", f"/v1/{pid}/waf/event/attack/logs", query=q) if action == "event_show": @@ -471,9 +474,11 @@ async def event(params: dict[str, Any], ctx: ToolContext) -> ToolResult: if action == "event_export_job": body = {} - for k in ("from", "to", "hosts", "attacks", "action"): + for k in ("from", "to", "hosts", "attacks"): if params.get(k) is not None: body[k] = params[k] + if params.get("event_action") is not None: + body["action"] = params["event_action"] return await _request(cfg, "POST", f"/v1/{pid}/waf/event/attack/log/job", body=body) if action == "threat_distribution": @@ -500,10 +505,9 @@ async def event(params: dict[str, Any], ctx: ToolContext) -> ToolResult: return ToolResult(success=False, error=f"Unknown action: {action}") -async def overview(params: dict[str, Any], ctx: ToolContext) -> ToolResult: +async def overview(ctx: ToolContext, action: str, **params: Any) -> ToolResult: cfg = _load_config(params.get("enterprise_project_id")) pid = cfg.project_id - action = params.get("action", "") def _time_query() -> dict[str, Any]: q: dict[str, Any] = {} diff --git a/.flocks/flockshub/plugins/tools/device/huaweicloud_waf_v39/hw_waf_event.yaml b/.flocks/flockshub/plugins/tools/device/huaweicloud_waf_v39/hw_waf_event.yaml index 0c6c99713..183d78778 100644 --- a/.flocks/flockshub/plugins/tools/device/huaweicloud_waf_v39/hw_waf_event.yaml +++ b/.flocks/flockshub/plugins/tools/device/huaweicloud_waf_v39/hw_waf_event.yaml @@ -20,7 +20,7 @@ inputSchema: - event_list 用途: 查询攻击事件列表（分页） 必填: 无 - 常用: `from`、`to`、`hosts`、`attacks`、`action`、`page`、`pagesize` + 常用: `from`、`to`、`hosts`、`attacks`、`event_action`、`page`、`pagesize` 风险提示: 只读查询接口；建议传 from/to 缩小范围 是否任务型: 否 - event_show @@ -38,7 +38,7 @@ inputSchema: - event_export_job 用途: 下发自定义导出攻击事件的异步任务 必填: `from`、`to` - 常用: `from`、`to`、`hosts`、`attacks`、`action` + 常用: `from`、`to`、`hosts`、`attacks`、`event_action` 风险提示: 写操作（下发异步任务）；任务完成后可通过 `event_log_download` 获取结果 是否任务型: 是 - threat_distribution @@ -89,9 +89,12 @@ inputSchema: items: type: string description: 攻击类型列表（如 `sqli`、`xss`、`cmdi`、`cc` 等） - action: + event_action: type: string description: 防护动作筛选，`block`（拦截）或 `log`（仅记录） + enum: + - block + - log top: type: integer description: 返回 Top N 条数，默认 5 diff --git a/.flocks/flockshub/plugins/tools/device/huaweicloud_waf_v39/hw_waf_policy.yaml b/.flocks/flockshub/plugins/tools/device/huaweicloud_waf_v39/hw_waf_policy.yaml index fb04016f7..e37f02ed5 100644 --- a/.flocks/flockshub/plugins/tools/device/huaweicloud_waf_v39/hw_waf_policy.yaml +++ b/.flocks/flockshub/plugins/tools/device/huaweicloud_waf_v39/hw_waf_policy.yaml @@ -61,8 +61,8 @@ inputSchema: 是否任务型: 否 - cc_rule_create 用途: 创建 CC 防护规则 - 必填: `policy_id`、`url`、`limit_num`、`limit_period`、`lock_time`、`tag_type`、`action` - 常用: `policy_id`、`url`、`limit_num`、`limit_period`、`lock_time`、`tag_type`、`action` + 必填: `policy_id`、`url`、`limit_num`、`limit_period`、`lock_time`、`tag_type`、`rule_action` + 常用: `policy_id`、`url`、`limit_num`、`limit_period`、`lock_time`、`tag_type`、`rule_action` 风险提示: 写操作；会新增 CC 规则 是否任务型: 否 - cc_rule_delete @@ -79,8 +79,8 @@ inputSchema: 是否任务型: 否 - custom_rule_create 用途: 创建精准防护规则 - 必填: `policy_id`、`name`、`conditions`、`action` - 常用: `policy_id`、`name`、`conditions`、`action`、`priority` + 必填: `policy_id`、`name`、`conditions`、`rule_action` + 常用: `policy_id`、`name`、`conditions`、`rule_action`、`priority` 风险提示: 写操作；会新增精准防护规则 是否任务型: 否 - custom_rule_delete @@ -168,7 +168,7 @@ inputSchema: tag_type: type: string description: CC 规则标签类型，如 `ip`、`cookie`、`header` 等 - action: + rule_action: type: object description: 规则匹配后的动作，含 `category`（`block`/`pass`/`log`）等字段 conditions: diff --git a/.flocks/flockshub/plugins/tools/device/huorong_edr_v1_0/huorong.handler.py b/.flocks/flockshub/plugins/tools/device/huorong_edr_v1_0/huorong.handler.py index 3168dfdbc..cdb1fdcd7 100644 --- a/.flocks/flockshub/plugins/tools/device/huorong_edr_v1_0/huorong.handler.py +++ b/.flocks/flockshub/plugins/tools/device/huorong_edr_v1_0/huorong.handler.py @@ -156,10 +156,10 @@ async def _post( if resp.status >= 400: return ToolResult( success=False, - data=resp_json, + output=resp_json, error=f"HTTP {resp.status}: {resp_text[:200]}", ) - return ToolResult(success=True, data=resp_json) + return ToolResult(success=True, output=resp_json) def _pick(params: dict[str, Any], *keys: str) -> dict[str, Any]: @@ -171,9 +171,8 @@ def _pick(params: dict[str, Any], *keys: str) -> dict[str, Any]: # --------------------------------------------------------------------------- -async def group(params: dict[str, Any], ctx: ToolContext) -> ToolResult: +async def group(ctx: ToolContext, action: str, **params: Any) -> ToolResult: base_url, timeout, secret_id, secret_key, verify_ssl = _resolve_runtime_config() - action = params.get("action", "") if action == "group_list": return await _post(base_url, "/api/group/_list", {}, secret_id, secret_key, timeout, verify_ssl) @@ -194,9 +193,8 @@ async def group(params: dict[str, Any], ctx: ToolContext) -> ToolResult: return ToolResult(success=False, error=f"Unknown action: {action}") -async def clnts(params: dict[str, Any], ctx: ToolContext) -> ToolResult: +async def clnts(ctx: ToolContext, action: str, **params: Any) -> ToolResult: base_url, timeout, secret_id, secret_key, verify_ssl = _resolve_runtime_config() - action = params.get("action", "") if action == "clnts_online": body = _pick(params, "offset") @@ -231,9 +229,8 @@ async def clnts(params: dict[str, Any], ctx: ToolContext) -> ToolResult: return ToolResult(success=False, error=f"Unknown action: {action}") -async def task(params: dict[str, Any], ctx: ToolContext) -> ToolResult: +async def task(ctx: ToolContext, action: str, **params: Any) -> ToolResult: base_url, timeout, secret_id, secret_key, verify_ssl = _resolve_runtime_config() - action = params.get("action", "") if action == "task_create": body = _pick(params, "offset") diff --git a/.flocks/flockshub/plugins/workflows/stream_alert_triage/guide.md b/.flocks/flockshub/plugins/workflows/stream_alert_triage/guide.md index 69ee83be6..108483378 100644 --- a/.flocks/flockshub/plugins/workflows/stream_alert_triage/guide.md +++ b/.flocks/flockshub/plugins/workflows/stream_alert_triage/guide.md @@ -218,7 +218,7 @@ alert_records | `input_paths` | 无 | 显式路径列表，优先级最高 | | `input_path` | 无 | 单个显式路径 | | `input_date` | 今天 | 自动发现该日所有上游 `dedup_result_*.jsonl` | -| `concurrency` | `1` | `workflow.md` 和 metadata 推荐 1；`concurrent_triage` 会限制到 1 到 5 | +| `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 写入位置 | @@ -228,10 +228,10 @@ alert_records 并发注意： - 外层 `ThreadPoolExecutor(max_workers=concurrency)` 处理 unique work units。 -- 内层每个 leader 会用 4 路并行 LLM 分支：`survey`、`cve_related`、`cve_info`、`payload_analysis`。 -- 稳态 LLM 峰值约为 `concurrency * 4`。 -- 配置引导应默认显式给出 `concurrency=1`。如果用户要提高到 2 到 5，先说明 LLM 并发和上游工具压力，再确认。 -- 当前 `load_dedup_file` 节点在完全不传 `concurrency` 时会输出 5；因此引导和样例中应显式传 `concurrency=1`，避免与文档推荐值不一致。 +- 每个 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，与文档和样例保持一致。 leader/follower 规则： @@ -298,7 +298,7 @@ leader/follower 规则： 3. 后续每行是否能按 JSON 对象解析。 4. 是否至少有 `dedup_key`、`sip`、`dip`、`req_http_url`、`threat_name` 中的关键字段。 5. 按 `dedup_key` 估算 unique work units 和 follower 数。 -6. 预估 `concurrency * 4` 的 LLM 峰值。 +6. 确认运行级 LLM 峰值不超过 `concurrency`。 真实执行验证注意： diff --git a/.flocks/flockshub/plugins/workflows/stream_alert_triage/meta.json b/.flocks/flockshub/plugins/workflows/stream_alert_triage/meta.json index 92f3b5ed8..08b32d2d2 100644 --- a/.flocks/flockshub/plugins/workflows/stream_alert_triage/meta.json +++ b/.flocks/flockshub/plugins/workflows/stream_alert_triage/meta.json @@ -1,7 +1,7 @@ { "name": "stream_alert_triage", "nameCn": "HTTP研判工作流", - "description": "Self-contained downstream pipeline for stream_alert_denoise. Loads enriched_alerts from JSONL files written by stream_alert_denoise, then runs leader/follower concurrent triage with the tdp_alert_triage logic INLINED (no sub-workflow invocation). Alerts sharing the same dedup_key in a batch form groups; only the LEADER (first occurrence) is triaged, FOLLOWERS reuse the leader result with no extra LLM calls. Each alert's 4 LLM analysis branches (survey / cve_related / cve_info / payload_analysis) still run in parallel via a nested 4-way ThreadPoolExecutor. The semantic-tagged markdown verdict report is attached to each alert via the `triage_report` field — NO per-alert markdown file is written to disk. Persistent cache (triage_cache.pkl, FIFO LRU, file-locked, atomic write): historic dedup_key hits reuse the cached verdict/title/triage_report instantly; misses run the full inline triage and persist new results. Default persistence writes enriched triage alerts into ~/.flocks/data/soc.db alert_records; optional JSONL output is controlled by config.json or triage_output_mode=jsonl/both.", + "description": "Self-contained downstream pipeline for stream_alert_denoise. Loads enriched_alerts from JSONL files written by stream_alert_denoise, then runs leader/follower concurrent triage with the tdp_alert_triage logic INLINED (no sub-workflow invocation). Alerts sharing the same dedup_key in a batch form groups; only the LEADER (first occurrence) is triaged, FOLLOWERS reuse the leader result with no extra LLM calls. Each alert keeps the 4 LLM analysis branches (survey / cve_related / cve_info / payload_analysis), while every llm.ask call shares the run-wide concurrency budget so nested executors cannot multiply provider load. The semantic-tagged markdown verdict report is attached to each alert via the `triage_report` field — NO per-alert markdown file is written to disk. Persistent cache (triage_cache.pkl, FIFO LRU, file-locked, atomic write): historic dedup_key hits reuse the cached verdict/title/triage_report instantly; misses run the full inline triage and persist new results. Default persistence writes enriched triage alerts into ~/.flocks/data/soc.db alert_records; optional JSONL output is controlled by config.json or triage_output_mode=jsonl/both.", "category": "default", "status": "active", "createdBy": null, diff --git a/.flocks/flockshub/plugins/workflows/stream_alert_triage/workflow.json b/.flocks/flockshub/plugins/workflows/stream_alert_triage/workflow.json index 5ff3479c4..b07e9474b 100644 --- a/.flocks/flockshub/plugins/workflows/stream_alert_triage/workflow.json +++ b/.flocks/flockshub/plugins/workflows/stream_alert_triage/workflow.json @@ -9,13 +9,13 @@ "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 外层并发数（默认 5）\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', 5)))\noutputs['max_triage_cache_size'] = int(inputs.get('max_triage_cache_size', 100000))\noutputs['input_date'] = _date_str(input_date)\n" + "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" }, { "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），内层 ThreadPoolExecutor(4) 并行 survey / cve_related / cve_info / payload_analysis。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 内层 ThreadPoolExecutor(max_workers=4)：单条 alert 内 4 个 LLM 并行\n （survey / cve_related / cve_info / payload_analysis — 保留 tdp_alert_triage\n 的 4 并行分支结构）\n 稳态 LLM 并发峰值 ≈ concurrency × 4（仅 4 并行分支阶段，且分支全部启用时）\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\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 llm.ask('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 analyses (the parallel branches from tdp_alert_triage) ─────────────\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 return llm.ask(\n prompt,\n timeout_s=LLM_CALL_TIMEOUT_S,\n max_retries=LLM_CALL_MAX_RETRIES,\n )\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 \"\"\"4 LLM analyses run in parallel — keeps tdp_alert_triage's fan-out structure.\"\"\"\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))))\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'(per-alert 4 LLM branches run in parallel)')\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 '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 外层 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" }, { "id": "summarize", diff --git a/.flocks/flockshub/plugins/workflows/stream_alert_triage/workflow.md b/.flocks/flockshub/plugins/workflows/stream_alert_triage/workflow.md index 7f9dd0bd9..6484a6d93 100644 --- a/.flocks/flockshub/plugins/workflows/stream_alert_triage/workflow.md +++ b/.flocks/flockshub/plugins/workflows/stream_alert_triage/workflow.md @@ -6,7 +6,7 @@ stream_alert_triage 是 NDR 告警流并发研判 Pipeline。 - 读取上游 stream_alert_denoise 去重输出 - 同批次 dedup_key 相同 → 只研判 leader，follower 复用结果 - 跨批次 dedup_key 命中缓存 → 直接复用历史研判，不调 LLM -- 4 并行 LLM 分支（survey / cve_related / cve_info / payload_analysis） +- 4 个 LLM 分支（survey / cve_related / cve_info / payload_analysis），共享运行级并发预算 - 研判产物仅写入 triage_report 字段，不生成独立报告文件 **完全自包含**，研判逻辑直接内联，不依赖也不嵌入 `tdp_alert_triage`。 @@ -15,7 +15,7 @@ stream_alert_triage 是 NDR 告警流并发研判 Pipeline。 * **跨批次复用**：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` 完全相同的研判语义 +* **保留 4 个 LLM 分支**（survey / cve_related / cve_info / payload_analysis）— 与 `tdp_alert_triage` 完全相同的研判语义；所有 LLM 调用共享运行级并发预算，避免与外层 work unit 并发相乘 * **研判产物仅以字段形式附加**到每条 alert（`triage_report` 字段含带语义标签的完整 markdown），**不生成任何独立的 per-alert 报告文件** ## 与上游的关系 @@ -51,7 +51,7 @@ load_dedup_file -> concurrent_triage -> 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），内层 ThreadPoolExecutor(4) 并行 survey / cve_related / cve_info / payload_analysis。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 | +| 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。 | 工作流最终输出 | 编辑流程结构时，要同时确认节点顺序、边关系、字段映射和最终输出是否仍然一致。 @@ -96,7 +96,7 @@ load_dedup_file -> concurrent_triage -> summarize ### 4.2 concurrent_triage -职责: Leader/follower 分组并发研判节点（自包含，内联 tdp_alert_triage 逻辑）。先按 dedup_key 把 alerts 分组：每组只对 leader 研判，follower 复用 leader 结果。外层 ThreadPoolExecutor(concurrency) 处理 unique work units（concurrency 取值 1–5，默认 1），内层 ThreadPoolExecutor(4) 并行 survey / cve_related / cve_info / payload_analysis。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` 可选输出。 +职责: 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` 可选输出。 - 节点类型: Python - 输入来源: load_dedup_file diff --git a/.flocks/plugins/skills/browser-use/SKILL.md b/.flocks/plugins/skills/browser-use/SKILL.md index b993f4e6e..d326e126e 100644 --- a/.flocks/plugins/skills/browser-use/SKILL.md +++ b/.flocks/plugins/skills/browser-use/SKILL.md @@ -61,9 +61,9 @@ flocks browser --doctor | 结果 | 触发条件 | 一线修复 | 仍失败兜底 | |---|---|---|---| | **A** | `next action` 以 `ready` 开头 | 立即确定 `CDP 直连`,阅读 `references/cdp-direct.md`,之后不再切到 `agent-browser` | — | -| **B** | `next action` 以 `attach` 开头 | 不要先反复 `--setup`；按输出执行 `flocks browser -c 'print(page_info())'` 或 `flocks browser -c 'print(list_tabs(include_chrome=False))'` 触发一次实际连接/观察 | 如果 `-c` 失败或仍无连接，执行 `flocks browser --reload` 清理旧 daemon，再执行 `flocks browser --setup`；setup 可能需要多次，直到用户完成浏览器 Allow/inspect 授权或错误信息稳定 | -| **C** | `next action` 以 `setup` 开头 | 先执行 `flocks browser --setup`（不包短超时），再运行 `--doctor` 确认 | 如提示 remote debugging 未启用、`DevToolsActivePort` 缺失、403 handshake 或 not live yet，再提示用户打开对应 inspect 页面并 Allow；用户确认后可多次 `--setup` | -| **D** | `next action` 提示启动浏览器或提供 endpoint | 明确告知需要先启动 Chrome/Chromium/Edge 或提供 CDP endpoint | **不**擅自降级到 curl/webfetch；坚持告知 skill 边界 | +| **B** | `next action` 以 `attach` 开头 | 不要先反复 `--setup`；按输出执行 `flocks browser -c 'print(page_info())'` 或 `flocks browser -c 'print(list_tabs(include_chrome=False))'` 触发一次实际连接/观察 | 如果 `-c` 失败或仍无连接，执行 `flocks browser --reload` 清理旧 daemon，再执行 `flocks browser --setup`；若 setup 输出本地 remote debugging 不可达，按 `references/cdp-setup.md` 的固定端口流程处理 | +| **C** | `next action` 以 `setup` 开头 | 先执行 `flocks browser --setup`（不包短超时），再运行 `--doctor` 确认 | 如提示 remote debugging 不可达、`DevToolsActivePort` 缺失、403 handshake 或 not live yet，按 `references/cdp-setup.md` 指引使用非默认 `--user-data-dir` 启动 Chromium 系浏览器，再访问 `/json/version` 验证 | +| **D** | `next action` 提示启动浏览器或提供 endpoint | 明确告知需要先用 remote-debugging 参数启动 Chrome/Chromium/Edge/Brave，或提供 CDP endpoint | **不**擅自降级到 curl/webfetch；坚持告知 skill 边界 | ### Step 4: 跨模式通用失败 diff --git a/.flocks/plugins/skills/browser-use/references/cdp-direct.md b/.flocks/plugins/skills/browser-use/references/cdp-direct.md index 0ab0de7ec..98b6bf685 100644 --- a/.flocks/plugins/skills/browser-use/references/cdp-direct.md +++ b/.flocks/plugins/skills/browser-use/references/cdp-direct.md @@ -350,11 +350,11 @@ print({"cookies": [c["name"] for c in cookies], "localStorage": js("Object.keys( 1. 先运行 `flocks browser --doctor` 看版本、安装模式、daemon 和浏览器状态；不要只看退出码，优先读 `next action`，再看 `browser running`、`daemon alive`、`active browser connections`。 2. `next action` 为 `attach`，或 `daemon alive` ok 但 `active browser connections` 为 0 时，不要先反复 `--setup`。先用一次实际命令触发连接/观察：`flocks browser -c 'print(page_info())'` 或 `flocks browser -c 'print(list_tabs(include_chrome=False))'`。 -3. 如果上一步失败或仍无连接，再执行 `flocks browser --reload` 清旧 daemon，然后执行 `flocks browser --setup`；setup 可能需要多次，因为用户可能需要完成浏览器 inspect/Allow 授权，或浏览器需要时间写入 remote debugging 状态。 +3. 如果上一步失败或仍无连接，再执行 `flocks browser --reload` 清旧 daemon，然后执行 `flocks browser --setup`。若 setup 输出本地 remote debugging 不可达，按 `references/cdp-setup.md` 的固定端口流程处理。 4. 首次安装、冷启动、daemon 不存在/不通，且浏览器已经运行或配置了 `BU_CDP_URL` / `BU_CDP_WS` 时，优先运行 `flocks browser --setup`。 -5. Chrome / Chromium / Edge 未运行且没有显式 CDP endpoint 时，只提示用户启动浏览器或提供 endpoint；不要直接让用户改设置。 -6. 只有在明确提示 remote debugging 未启用、`DevToolsActivePort` 缺失、403 handshake、remote-debugging page 或 not live yet 时，才让用户打开对应浏览器的 inspect 页面（例如 `chrome://inspect/#remote-debugging` 或 `edge://inspect/#remote-debugging`）并勾选 Allow remote debugging。 -7. 用户刚开启 remote debugging 时，不要立刻再次运行 `flocks browser --doctor`；先执行一次 `flocks browser --setup`，或直接执行 `flocks browser -c 'print(page_info())'` 触发 daemon attach，再用 `--doctor` 做只读确认。 +5. Chrome / Chromium / Edge / Brave 未运行且没有显式 CDP endpoint 时，按 `references/cdp-setup.md` 提供对应平台的 remote-debugging 启动命令。 +6. 只有在明确提示 remote debugging 不可达、`DevToolsActivePort` 缺失、403 handshake、remote-debugging page 或 not live yet 时，才让用户使用非默认 `--user-data-dir` 和 `--remote-debugging-port=9222` 启动 Chromium 系浏览器，并访问 `http://127.0.0.1:9222/json/version` 验证。 +7. 用户刚按固定端口流程启动浏览器时，不要立刻再次运行 `flocks browser --doctor`；先执行一次 `flocks browser --setup`，或直接执行 `flocks browser -c 'print(page_info())'` 触发 daemon attach，再用 `--doctor` 做只读确认。 8. `connection refused`、`DevTools not live yet`、`/json/version` 404 通常是浏览器正在启动，轮询等待，不要重启。 9. stale websocket / stale socket 时执行一次： diff --git a/.flocks/plugins/skills/browser-use/references/cdp-setup.md b/.flocks/plugins/skills/browser-use/references/cdp-setup.md index 3f72f5825..0626cdb0d 100644 --- a/.flocks/plugins/skills/browser-use/references/cdp-setup.md +++ b/.flocks/plugins/skills/browser-use/references/cdp-setup.md @@ -1,6 +1,6 @@ # Flocks browser setup -浏览器已运行，但 daemon 不存在/不通，或 active browser connection 不可用 +本地固定端口 CDP 设置：daemon 不存在/不通，active browser connection 不可用，或浏览器尚未以 remote debugging 参数启动。 先区分两种情况： @@ -11,15 +11,44 @@ 2. daemon 不存在/不通，且浏览器已运行或配置了 `BU_CDP_URL` / `BU_CDP_WS`： - 执行 `flocks browser --setup` 触发 attach，不要用短超时包装该命令。 -只有在错误明确指向 remote debugging 未启用、`DevToolsActivePort` 缺失、403 handshake 或 not live yet 时，才提示用户： +只有在错误明确指向 remote debugging 不可达、`DevToolsActivePort` 缺失、403 handshake 或 not live yet 时，才提示用户走本地固定端口流程： ```text -browser: not connected — 请确保 Chrome / Chromium / Edge 已打开，然后访问对应浏览器的 inspect 页面（例如 chrome://inspect/#remote-debugging 或 edge://inspect/#remote-debugging）并勾选 Allow remote debugging +不要从 chrome://inspect 查找 webSocketDebuggerUrl。关闭对应 Chromium 系浏览器后，使用非默认 --user-data-dir 和 --remote-debugging-port=9222 启动浏览器，再访问 http://127.0.0.1:9222/json/version 验证。 ``` -然后等待用户进一步指示，不要直接操作。 +候选命令按平台选择一个即可；如果浏览器安装路径不同，替换可执行文件路径： -当用户确认已开启remote-debugging后: -1. 执行 `flocks browser --setup` 触发交互式 attach，不要用短超时包装该命令 +Windows PowerShell： + +```powershell +& "$env:ProgramFiles\Google\Chrome\Application\chrome.exe" --remote-debugging-port=9222 --user-data-dir="$env:USERPROFILE\.flocks\chrome-debug-profile" +& "${env:ProgramFiles(x86)}\Microsoft\Edge\Application\msedge.exe" --remote-debugging-port=9222 --user-data-dir="$env:USERPROFILE\.flocks\edge-debug-profile" +chromium.exe --remote-debugging-port=9222 --user-data-dir="$env:USERPROFILE\.flocks\chromium-debug-profile" +& "$env:ProgramFiles\BraveSoftware\Brave-Browser\Application\brave.exe" --remote-debugging-port=9222 --user-data-dir="$env:USERPROFILE\.flocks\brave-debug-profile" +``` + +macOS： + +```bash +/Applications/Google\ Chrome.app/Contents/MacOS/Google\ Chrome --remote-debugging-port=9222 --user-data-dir="$HOME/.flocks/chrome-debug-profile" +/Applications/Microsoft\ Edge.app/Contents/MacOS/Microsoft\ Edge --remote-debugging-port=9222 --user-data-dir="$HOME/.flocks/edge-debug-profile" +/Applications/Chromium.app/Contents/MacOS/Chromium --remote-debugging-port=9222 --user-data-dir="$HOME/.flocks/chromium-debug-profile" +/Applications/Brave\ Browser.app/Contents/MacOS/Brave\ Browser --remote-debugging-port=9222 --user-data-dir="$HOME/.flocks/brave-debug-profile" +``` + +Linux： + +```bash +google-chrome --remote-debugging-port=9222 --user-data-dir="$HOME/.flocks/chrome-debug-profile" +microsoft-edge --remote-debugging-port=9222 --user-data-dir="$HOME/.flocks/edge-debug-profile" +chromium --remote-debugging-port=9222 --user-data-dir="$HOME/.flocks/chromium-debug-profile" +brave-browser --remote-debugging-port=9222 --user-data-dir="$HOME/.flocks/brave-debug-profile" +``` + +输出命令后等待用户进一步指示，不要占用当前终端盲目重试。 + +当用户确认 `http://127.0.0.1:9222/json/version` 已可访问后: +1. 执行 `flocks browser --setup` 触发 attach，不要用短超时包装该命令 2. 再运行 `flocks browser --doctor` 做只读确认。 -3. 如果还失败，先执行 `flocks browser --reload` 清理旧 daemon，再重新执行 `flocks browser --setup`，避免因为残留 daemon 造成干扰。setup 可能需要多次，直到用户完成浏览器 Allow/inspect 授权或错误信息稳定。 +3. 如果还失败，先执行 `flocks browser --reload` 清理旧 daemon，再重新执行 `flocks browser --setup`，避免因为残留 daemon 造成干扰。 diff --git a/.flocks/plugins/skills/device-integration-guide/SKILL.md b/.flocks/plugins/skills/device-integration-guide/SKILL.md index 010594ca3..f02795a9f 100644 --- a/.flocks/plugins/skills/device-integration-guide/SKILL.md +++ b/.flocks/plugins/skills/device-integration-guide/SKILL.md @@ -18,67 +18,33 @@ description: 指导 Flocks 新建、添加和接入安全设备。Use when the u ## 核心原则 -- 先确认用户是在**新建设备实例**、**整理配置草稿**，还是**测试连通性**。 +- 先确认用户是在**新建设备实例**、**页面配置**，还是**测试连通性**。 - 不要要求用户在聊天里粘贴密码、Token、Cookie、API Key 等敏感凭证。 -- 不要在 skill 中优先引导通过工具写入设备配置；设备接入页面表单和 JSON 草稿是配置写入的主路径。 +- 独立会话中，已安装模板的新设备优先通过 `device_manage(action="create")` 创建；只有敏感凭证需要回到设备接入页面填写。 - 每次修改后，优先用标准连通性测试验证结果。 - 保持回答简短，给出当前动作、结果和下一步。 ## 决策流程 -1. 用户已经在设备列表里有目标设备，并提供了 `device_id`： +1. 已有 `device_id`：非敏感配置用 `device_manage(action="update")`，测试或排障用 `connectivity_test`；敏感字段回到设备接入页面填写。 +2. 没有 `device_id`：先用 `device_manage(action="list")` 排除已有实例，再调用 `device_manage(action="list_templates")` 查询模板。设备实例为空不代表模板不存在。 +3. 按名称、厂商、`plugin_id`、`service_id`、`storage_key` 和描述匹配模板： + - `installed=true`：按 `credential_schema` 整理非敏感字段，然后进入创建流程。 + - `installed=false`：引导用户在 FlockHub 安装返回的 `plugin_id`，安装后重新查询；不要创建自定义模板。 + - 没有匹配模板：进入自定义 API、浏览器或 Workflow 接入。 -- 如果是配置变更，回到设备接入页面表单或输出页面可回填 JSON 草稿。 -- 如果是测试或排障，先走 `device_manage(action="connectivity_test")`。 +## 使用模板直接创建设备 -2. 用户说“添加设备”“接入设备”“新建设备”，但还没有设备实例： +只有 `list_templates` 已返回匹配模板且 `installed=true` 时，才能调用 `device_manage(action="create")`。`create` 会再次校验模板状态，不能用 `update` 代替创建。 -- 如果有已安装模板，引导用户在设备接入页面填写表单。 -- 如果没有合适模板，进入自定义设备接入路径。 -- 如果涉及密钥、密码、Token、Cookie 或浏览器登录态，只说明应该填到页面表单，不要在聊天中收集真实值。 +从模板收集设备名称、机房、SSL 偏好和已声明的非敏感字段；不询问或传递 `storage=secret`、`sensitive=true`、`input_type=password` 的字段。 -3. 用户只描述产品、厂商、控制台地址或 API 文档： +创建成功后： -- 先判断已有模板是否可用。 -- 未安装模板需要先去 FlockHub 安装。 -- 没有合适模板时，按“自定义接入路由”选择 API、浏览器或 Workflow。 - -## 设备列表与目标确认 - -如果用户没有给出 `device_id`，先调用： - -```python -device_manage(action="list") -``` - -从返回结果里确认目标设备、机房、工具集和 `device_id`。 - -如果设备不存在，提醒用户前往「设备接入」页面添加设备；不要伪造 `device_id` 或直接调用业务工具。 - -## 新建设备与页面回填 - -用户在设备接入页面创建或配置设备时，目标是帮助页面得到清晰的表单信息。 - -需要收集的信息： - -- 设备名称。 -- 已安装模板的 `storage_key`。 -- Base URL、Host、端口、协议、租户或区域等非敏感字段。 -- SSL 证书验证偏好：`verify_ssl=true/false`。 -- 需要填写哪些敏感字段，但不收集真实值。 - -当信息足够，并且当前任务是在设备接入页面生成配置草稿时，在回复末尾输出 JSON 代码块供页面一键回填： - -```json -{"storage_key":"","device_name":"<设备名称>","fields":{"base_url":"https://example.local"},"verify_ssl":false} -``` - -JSON 草稿规则： - -- 只包含非敏感字段。 -- 不写真实 API Key、Secret、Token、Cookie、密码。 -- 敏感字段留空或省略，并提示用户稍后在页面表单中填写。 -- 如果没有合适模板，不要输出设备配置 JSON，先进入自定义接入路径。 +- 记录并报告返回的 `device_id`，后续操作始终使用它。 +- 如果返回 `sensitive_fields_to_complete`，告诉用户前往该设备的配置表单填写这些字段，不要在聊天中索要真实值。 +- 用户确认敏感字段已填写后，调用 `device_manage(action="connectivity_test", device_id="")`。 +- 同一轮会话已经拿到成功返回的 `device_id` 后，不要因重试或继续对话再次创建。 ## 自定义接入路由 @@ -90,9 +56,9 @@ JSON 草稿规则： 如果用户已经明确选择 API、浏览器或 Workflow，不要重复询问接入方式。只有无法判断时，才用一句话澄清。 -## 配置草稿与表单更新 +## 页面配置与更新 -如果用户要写入或更新设备配置，应根据当前页面上下文整理表单草稿，让设备接入页面负责落库。 +如果用户正在设备接入页面配置设备，帮助确认需要填写的表单字段，让页面负责保存。独立会话中的已有设备非敏感配置更新使用 `device_manage(action="update")`。 可以整理的非敏感字段包括： @@ -104,22 +70,21 @@ JSON 草稿规则： - `tenant` - `region` -不要在 JSON 草稿或聊天中写入敏感字段： +不要在聊天中索要或回显敏感字段： - `api_key` - `secret` - `password` - `token` - `cookie` -- `auth_state` 如果用户的目标是补填密钥、修改密码、刷新 Token 或重新登录，只说明应该在设备接入页面对应字段中处理。 -当需要给页面回填时，使用“新建设备与页面回填”中的 JSON 草稿格式。对于已有设备编辑，也只输出非敏感字段和 `verify_ssl`，并说明敏感字段在页面表单内填写。 +对于已有设备编辑，只处理非敏感字段和 `verify_ssl`，并说明敏感字段应在页面表单内填写。 ## 连通性与冒烟验证 -配置在设备接入页面保存后，除非用户明确不需要，继续调用： +设备通过 `create` 返回 `device_id`，或在设备接入页面保存后，除非用户明确不需要，继续调用： ```python device_manage(action="connectivity_test", device_id="") @@ -149,6 +114,5 @@ device_manage(action="connectivity_test", device_id="") - 不要在聊天中索要、保存或复述真实密钥。 - 不要把自定义设备误做成普通 API 服务。 -- 不要对未安装模板输出可回填 JSON。 - 不要跳过 `device_manage(action="connectivity_test")` 就声称设备已可用。 - 不要把卡片状态建立在普通业务工具结果上；卡片状态以标准连通性测试写入结果为准。 diff --git a/.flocks/plugins/skills/sangfor-edr-use/SKILL.md b/.flocks/plugins/skills/sangfor-edr-use/SKILL.md index a51058991..479c8a18a 100644 --- a/.flocks/plugins/skills/sangfor-edr-use/SKILL.md +++ b/.flocks/plugins/skills/sangfor-edr-use/SKILL.md @@ -1,99 +1,37 @@ ---- -name: sangfor-edr-use -description: 用于处理深信服 EDR（终端检测与响应）相关任务，通过浏览器（CDP 直连）进行以下任务：终端状态查询、终端概况统计、失陷设备排查、设备运行状态查看等。只要用户提到 深信服 EDR、EDR、sangfor EDR 等需求时，必须先加载本 skill。本 skill 是 EDR 平台操作的唯一决策入口：在未阅读本 skill 前，不要直接使用 browser-use skill。 ---- - -# 深信服 EDR Use - -## First +--- +name: sangfor-edr-use +description: 用于处理深信服 EDR（终端检测与响应）相关任务，通过 Flocks browser/CDP 完成登录态复用、终端状态查询、概况统计、失陷设备排查和设备运行状态查看。只要用户提到深信服 EDR、EDR 或 sangfor EDR，必须先加载本 skill；本 skill 是 EDR 任务的唯一入口，未阅读前不要直接使用 browser-use。 +--- -### 登录态处理规则 +# 深信服 EDR Use -当用户需要打开深信服 EDR 页面，或需要通过 Web2CLI 抓取 EDR 页面请求时，必须按下面顺序处理登录态，不要一开始就要求用户提供账密： +本 skill 只负责任务入口和登录分流；CDP 操作、浏览器启动、验证码、selector、tab/iframe 处理和页面数据提取见 [references/cdp-workflow.md](references/cdp-workflow.md)。 -1. 先调用 `sangfor_edr_auth` 的 `action=status_auth_state` 或 `action=validate_auth_state`，检查 `~/.flocks/browser/sangfor-edr/auth-state.json` 是否存在且可用。 -2. 如果登录态可用，直接复用该 state，继续打开 EDR 页面或执行 Web2CLI 抓取流程。 -3. 如果 state 不存在或已失效，但本地已经保存了 EDR 地址、用户名和密码，调用 `sangfor_edr_auth` 的 `action=ensure_auth_state` 自动刷新登录态。 -4. 如果没有可用 state，也没有保存账密配置，再引导用户选择： - - 提供 EDR 访问地址、用户名和密码，自动登录并保存账密配置，后续 state 失效时可直接自动刷新； - - 不提供账密，走浏览器手动登录流程。打开 EDR 页面后由用户在可视化浏览器中完成登录，登录成功后保存完整浏览器登录态（包括 cookies、localStorage 等）。 +## 登录分流 -无论采用哪种方式，只要获得可用登录态，就继续原有浏览器 / Web2CLI 流程：加载登录态、打开目标 EDR 页面、按需注入 Web2CLI hook、执行页面操作并导出捕获到的请求。后续再次打开页面时，仍必须先校验 `auth-state.json`；若登录态失效且已保存账密，则自动重新走 CDP 登录并刷新 state。 +需要打开 EDR 页面、抓取页面请求或调用 EDR API 采集工具时，按以下顺序执行： -若自动登录过程中出现验证码识别失败、MFA 校验、页面选择器变化、未检测到登录成功或有效 `sessionid` 等情况，立即回退到原有浏览器手动登录流程。 - -> ⚠️ **EDR 没有开放 API**，所有操作必须通过浏览器（CDP 直连）完成。 - -## 浏览器模式使用指南 - -请阅读以下文档获取完整流程： -- [references/cdp-workflow.md](references/cdp-workflow.md) - -### CDP 模式适用场景 - -- **首页仪表盘**（`/ui/#/index`）：设备 CPU/内存/硬盘使用率、终端概况（在线/离线/服务器/PC）、失陷设备统计 -- **威胁资产分析**：已失陷终端列表（需点击"已失陷终端"标签页，不是默认的"全部"） -- 页面详情、交互式筛选 - -### 可用工具脚本 - -| 脚本路径 | 功能 | 必需参数 | -|---------|------|---------| -| `references/fetch_edr_system_state.py` | 设备状态抓取 | `--url {EDR_URL}` | - -脚本位于 skill 目录的 `references/` 下，无硬编码 URL 或敏感信息。 - -### 执行示例 - -脚本位于 `/skills/sangfor-edr-use/references/fetch_edr_system_state.py`，请按当前平台选择对应命令。 - -**Windows（PowerShell）** - -```powershell -powershell -Command "& '\Scripts\python.exe' '\skills\sangfor-edr-use\references\fetch_edr_system_state.py' --url '{EDR_URL}'" -``` - -**macOS / Linux（bash / zsh）** - -```bash -"/bin/python" "/skills/sangfor-edr-use/references/fetch_edr_system_state.py" --url "{EDR_URL}" -``` - -**占位符说明** - -| 占位符 | Windows 典型值 | macOS/Linux 典型值 | -|--------|---------------|-------------------| -| `` | `D:\Flocks Project\flocks\.venv` | `~/Flocks/flocks/.venv`（取决于实际安装位置） | -| `` | `%USERPROFILE%\.flocks\plugins` | `~/.flocks/plugins` | - -> 必须使用 Flocks 虚拟环境（`.venv`）执行；系统 Python 可能缺少依赖。如不确定 venv 位置，先执行 `flocks --version` 或检查 Flocks 安装目录。 - -## 关键坑点（必须避免） - -| 坑 | 原因 | 解法 | -|---|---|---| -| `flocks browser -c js(...)` 返回空文本 | daemon session 指向错误的 tab | 用 Python socket 直连 daemon，通过 `Runtime.evaluate` 在正确 context 执行 | -| `flocks browser -c new_tab()` 后后续命令无响应 | tab 切换导致 session 错位 | 用 `switch_tab(targetId)` 明确切到 EDR tab | -| 多行代码转义失败 | PowerShell 引号嵌套 | 使用 `fetch_edr_system_state.py` 脚本，无需手动转义 | -| EDR 页面数据为空 | EDR 内容在跨域 iframe 中 | 用 CDP direct 方式 attach 到 EDR tab，在正确 frame context 执行 JS | -| 失陷设备数量不匹配 | 读取的是"全部"筛选而非"已失陷"筛选 | 需点击"已失陷终端"标签页获取准确数量 | - -## 失陷设备查询 SOP - -**问题**：首页仪表盘显示"已失陷 N 台"，但威胁资产分析页面默认只列出部分。 - -**成功路径**： -1. 在 EDR 首页仪表盘确认"已失陷 N 台"的数量 -2. 如需具体清单，进入 `威胁资产分析` → 点击 `已失陷终端` 标签页（不是默认的"全部"） -3. 只有点击"已失陷终端"标签页，列表数量才会与仪表盘一致 - -**⚠️ 必须避免**：直接读取威胁资产分析的默认"全部"筛选结果作为失陷设备清单，这是错误的。 - -## 执行规范 - -**必须使用 Flocks 虚拟环境（`.venv`）执行 Python 脚本，禁止使用系统 Python。** - -- ✅ 正确：`/bin/python`（Unix）或 `\Scripts\python.exe`（Windows） -- ❌ 禁止：`python script.py` / `python3 script.py`（直接调用 PATH 中的 Python） - -**原因**：Flocks 虚拟环境包含了所有项目依赖，系统 Python 可能缺少必要的包。完整跨平台示例见上一节"执行示例"。 +1. 调用 `sangfor_edr_auth(action=status_auth_state)`，确认 `validation.valid`、`can_auto_refresh` 和 `has_saved_token`。不要先向用户索要账密。 +2. `validation.valid=true`： + - 仅浏览器/CDP任务：直接继续； + - 需要 API token：`has_saved_token=true` 才继续，否则调用 `refresh_auth_state` 补齐 token。 +3. `validation.valid=false` 且 `can_auto_refresh=true`：调用 `ensure_auth_state` 自动登录。 +4. 没有可用 state 或账密时：调用 `ensure_auth_state` 打开登录页；返回 `manual_login_required` 后，让用户在该工具打开的同一个 EDR tab 中完成登录、MFA 或 UKey。 +5. 用户完成手动登录后，调用 `complete_manual_login`；只有返回 `manual_login_captured_auth_state` 且 `token_saved=true`，才认为浏览器登录和 API 登录准备均完成。 +6. 返回 `browser_daemon_not_ready` 或 `auth_state_load_failed_browser_daemon_not_ready`：依次执行 `flocks browser --setup`、`flocks browser --doctor`，再重试原 action。 + +认证工具会在登录提交前监听 `launch_login.php` 的 fetch/XHR 响应，将 `data.token` 保存到 Secret Manager；token 不得回显、写入日志或写入 `auth-state.json`。后续 API 工具应从 Secret Manager 读取对应 token，并同时使用同一 EDR 的 cookies。 + +`bu.port` 是 Flocks browser daemon 的 IPC 端口文件，不是 Chrome 的 remote-debugging 端口；禁止手工创建或修改它。 + +## 任务边界 + +- 当前 `sangfor_edr_v1_0_0` 的页面业务仍通过浏览器/CDP完成；需要 API 采集时，必须先完成 token readiness 检查。 +- 首页仪表盘可用于设备 CPU/内存/硬盘、终端概况和失陷统计。 +- 查询失陷终端清单时，必须进入“威胁资产分析”并选择“已失陷终端”，不能使用默认“全部”列表。 +- 设备状态抓取脚本位于 `references/fetch_edr_system_state.py`；执行前必须完成上述登录分流。 + +## 执行约束 + +- 运行 skill 内 Python 脚本时必须使用 Flocks 虚拟环境；不要使用系统 Python。 +- 需要具体 CDP 命令、平台启动方式、验证码/selector 配置、tab/iframe 处理或页面关键词时，读取 `references/cdp-workflow.md`，不要在本文件重复展开。 diff --git a/.flocks/plugins/skills/sangfor-edr-use/references/cdp-workflow.md b/.flocks/plugins/skills/sangfor-edr-use/references/cdp-workflow.md index d8b224596..aa07e31e1 100644 --- a/.flocks/plugins/skills/sangfor-edr-use/references/cdp-workflow.md +++ b/.flocks/plugins/skills/sangfor-edr-use/references/cdp-workflow.md @@ -10,12 +10,17 @@ ## 零、前置条件 -### 1. 确保浏览器 daemon 可用 -```bash -flocks browser --doctor -``` +### 1. 确保浏览器 daemon 可用 + +`sangfor_edr_auth` 在需要浏览器时会自动确保 daemon 运行，并尝试替换陈旧实例。仅在工具返回 `browser_daemon_not_ready` 时执行下列诊断： + +```bash +flocks browser --doctor +``` -如果 `active browser connections` 为 0，需用户开启浏览器 remote debugging。 +如果 `active browser connections` 为 0，需用户开启浏览器 remote debugging。 + +> 禁止手工创建或改写 `~/.flocks/browser/bu.port`。它指向 Flocks browser daemon 的自定义 IPC 端口，不是 Chrome 的 `--remote-debugging-port`。 ### 2. 开启 Chrome Remote Debugging @@ -47,9 +52,9 @@ chromium --remote-debugging-port=9222 1. 调用 `sangfor_edr_auth`，`action=status_auth_state`，确认固定 state 是否存在、是否可用，以及是否已有可自动刷新的账密配置。 2. 如果返回的 `validation.valid` 为 `true`，直接复用已保存 state。 -3. 如果 state 不存在或失效，但 `can_auto_refresh` 为 `true`，调用 `sangfor_edr_auth`，`action=ensure_auth_state`，通过 browser daemon / CDP 驱动真实 EDR 登录页自动登录。验证码图片在浏览器会话中获取，OCR 识别后填入页面；登录成功后保存浏览器 state。 +3. 如果 state 不存在或失效，但 `can_auto_refresh` 为 `true`，调用 `sangfor_edr_auth`，`action=ensure_auth_state`，通过 browser daemon / CDP 驱动真实 EDR 登录页自动登录。验证码图片在浏览器会话中获取，OCR 识别后填入页面；登录成功后保存 browser state，并从 `launch_login.php` 响应捕获 `data.token` 到 Secret Manager（`sangfor_edr_token`）。 4. 如果没有可用 state，也没有保存账密配置，再询问用户：提供 EDR 地址、用户名、密码后自动登录并保存，或不提供账密改走手动登录。 -5. 用户不提供账密、自动登录失败、MFA/验证码/OCR/DOM 选择器异常时，由用户在 Chrome 中手动登录 EDR（如需 MFA，完成认证），登录成功后保存浏览器 state。 +5. 用户不提供账密、自动登录失败、MFA/验证码/OCR/DOM 选择器异常时，由用户在已打开的 Chrome 中完成登录，再调用 `action=complete_manual_login` 自动校验并保存 browser state。 --- @@ -60,10 +65,13 @@ chromium --remote-debugging-port=9222 - `https://edr.example.com/` - `https://edr.company.com/` -### Step 2：检查 daemon -```bash -flocks browser --doctor -``` +### Step 2：检查 daemon + +先直接调用 `sangfor_edr_auth`。该工具会自动启动或恢复 daemon。仅当返回 `browser_daemon_not_ready` 或 `auth_state_load_failed_browser_daemon_not_ready` 时，执行： + +```bash +flocks browser --doctor +``` 如果 daemon 未运行，执行： ```bash @@ -78,11 +86,11 @@ flocks browser --setup 推荐先调用工具做状态检查： -- `action=status_auth_state`：返回 `auth_state_exists`、`validation.valid`、`can_auto_refresh` 等非敏感状态。 +- `action=status_auth_state`：返回 `auth_state_exists`、`validation.valid`、`can_auto_refresh`、`has_saved_token` 等非敏感状态。 - `validation.valid=true`：直接继续 Step 4。 - `validation.valid=false` 且 `can_auto_refresh=true`：调用 `action=ensure_auth_state` 自动打开真实登录页、识别验证码、填入账密并保存新的 state。 - `validation.valid=false` 且 `can_auto_refresh=false`：再询问用户是提供账密自动登录并保存，还是手动登录。 -- 用户拒绝提供账密或自动登录失败：提示用户在 Chrome 中手动登录，登录成功后执行 `flocks browser state save ~/.flocks/browser/sangfor-edr/auth-state.json`。 +- 用户拒绝提供账密或自动登录失败：提示用户在已打开的 Chrome 中完成登录，然后调用 `action=complete_manual_login`；成功后继续原任务。 ### Step 4：打开目标页面 用户或工具在 Chrome 中打开： @@ -196,7 +204,8 @@ print(text_result["result"]["result"]["value"]) | EDR tab 未找到 | 页面未打开或 URL 不匹配 | 确保 Chrome 中打开了 EDR 首页 | | 页面数据为空 | EDR 内容在跨域 iframe 中 | 用 CDP direct 方式 attach 到 EDR tab，在正确 frame context 执行 JS | | 页面显示登录框 | 会话已失效 | 若已保存账密，调用 `sangfor_edr_auth` 自动刷新 state；否则提示用户手动登录 EDR | -| 自动登录失败 | 验证码识别失败、MFA、页面选择器变化或登录成功检测失败 | 回退到浏览器手动登录，登录后保存 state | +| 自动登录失败 | 验证码识别失败、MFA、页面选择器变化或登录成功检测失败 | 在已打开的浏览器中完成登录，再调用 `action=complete_manual_login` | +| `browser_daemon_not_ready` | daemon 无法启动或连接 Chrome | 执行 `flocks browser --setup` 和 `flocks browser --doctor`，然后重试；不要修改 `bu.port` | --- diff --git a/.flocks/plugins/skills/tdp-use/SKILL.md b/.flocks/plugins/skills/tdp-use/SKILL.md index bcf4bab74..15f556540 100644 --- a/.flocks/plugins/skills/tdp-use/SKILL.md +++ b/.flocks/plugins/skills/tdp-use/SKILL.md @@ -38,9 +38,10 @@ API 参数和适用场景见 [references/api-reference.md](references/api-refere - API 调用必须以当前 tool schema 为准，优先使用 schema 暴露的顶层语义化参数；列表类工具常见 `keyword`、`severity`、`cur_page`、`page_size`、`sort_by`，但 `tdp_log_search.sql` 是过滤表达式不是完整 SQL，禁止 `SELECT/FROM`，控制返回数量用 `size`，`terms` 可不传 `sql`，外部攻击结果筛选用 `result_list`。 - 用户说“告警”“告警记录”“告警日志”“明细记录”“查某 IP 的告警”时，默认走 `tdp_log_search` - 用户说“看板”“概览”“趋势”“统计”时，先用 `tdp_dashboard_status` -- 用户说“威胁事件”“外部攻击”“攻击事件”“事件总览”“事件趋势”时，优先用 `tdp_incident_list` 或 `tdp_threat_inbound_attack` -- 用户说“实时监控”“威胁监控列表”时，优先用正式 API 工具 `tdp_threat_monitor_list`；不要用 Cookie 认证的旧 Web CLI 代替 -- 用户说“告警主机”“受害主机”“主机下的事件”时，优先用 `tdp_host_threat_list` +- 用户明确说“智能聚合”“聚合攻击事件”“智能聚合事件”“攻击过程”“事件攻击时间线”时，优先用 `tdp_threat_intelligent_aggregation` +- 用户说“外部攻击”“外部攻击严重性分布”时，优先用 `tdp_threat_external_attack` +- 用户说“威胁事件列表”“实时监控”“威胁实时监控列表”“威胁监控列表”时，优先用正式 API 工具 `tdp_threat_monitor_list`；不要用 Cookie 认证的旧 Web CLI 代替 +- 用户说“告警主机”“受害主机”“主机下的事件”时，优先用 `tdp_threat_alert_host` - 用户说“系统状态”“核心服务状态”“数据库状态”时，优先用 `tdp_system_status` - 用户说“MDR”“研判结果”“研判统计”时，优先用 `tdp_mdr_alert_list` - 用户说“脆弱性”“弱口令”“登录入口”“上传接口”“API 风险”“隐私数据”时，走对应资产或风险类工具 diff --git a/.flocks/plugins/skills/tdp-use/references/api-reference.md b/.flocks/plugins/skills/tdp-use/references/api-reference.md index 3bd3bef99..9aa51339b 100644 --- a/.flocks/plugins/skills/tdp-use/references/api-reference.md +++ b/.flocks/plugins/skills/tdp-use/references/api-reference.md @@ -9,11 +9,11 @@ | 看安全态势、趋势、TOP 统计 | `tdp_dashboard_status` | `status` / `security` / `threat_event` / `alert_level_trend` | 通常可空参；列表类 action 可补 `machine_type`、`severity`、分页参数 | | 查原始告警日志 | `tdp_log_search` | `search` | `time_from`、`time_to`、`sql` | | 查字段聚合统计 | `tdp_log_search` | `terms` | `time_from`、`time_to`、`term` | -| 查威胁事件列表 | `tdp_incident_list` | `search` | `time_from`、`time_to`；可补 `severity`、`phase`、`result`、`keyword`、分页参数 | -| 查威胁实时监控列表 | `tdp_threat_monitor_list` | 默认 | 时间范围不超过 24 小时；可补 `sql`、`net_data_type`、`assets_group`、分页参数 | -| 看事件时间线 / 结果分布 / 攻击者明细 | `tdp_incident_list` | `timeline` / `result_distribution` / `attacker_ip_detail` | 通常先要 `incident_id` | -| 查外部攻击严重性分布 | `tdp_threat_inbound_attack` | 默认 | `time_from`、`time_to`；可补 `severity`、`result_list`、`keyword` | -| 查告警主机汇总 / 主机下事件 | `tdp_host_threat_list` | `summary` / `events` | `summary` 可补 `severity`、`direction`、`threat_type`、`keyword`；`events` 至少要 `asset_machine` | +| 查智能聚合事件 / 聚合攻击事件 | `tdp_threat_intelligent_aggregation` | `incident_search` | `time_from`、`time_to`；可补 `severity`、`phase`、`result`、`keyword`、分页参数 | +| 查威胁事件列表 / 威胁实时监控列表 | `tdp_threat_monitor_list` | 默认 | 时间范围不超过 24 小时；可补 `sql`、`net_data_type`、`assets_group`、分页参数 | +| 看事件时间线 / 结果分布 / 攻击者明细 | `tdp_threat_intelligent_aggregation` | `attack_timeline` / `attack_result_distribution` / `attacker_ip_detail` | 通常先要 `incident_id` | +| 查外部攻击严重性分布 | `tdp_threat_external_attack` | 默认 | `time_from`、`time_to`；可补 `severity`、`result_list`、`keyword` | +| 查告警主机列表 / 指定主机威胁列表 | `tdp_threat_alert_host` | `alert_host_list` / `host_threat_list` | `alert_host_list` 可补 `severity`、`direction`、`threat_type`、`keyword`；`host_threat_list` 至少要 `asset_machine` | | 查脆弱性 | `tdp_vulnerability_list` | 默认 | 常见补 `time_from`、`time_to`、`severity`、`status`、`keyword`、分页参数 | | 查弱口令 | `tdp_login_weakpwd_list` | 默认 | 常见补 `time_from`、`time_to`、`data`、`result`、`app_class`、`keyword` | | 查服务 / 主机 / 框架资产 | `tdp_machine_asset_list` | `service_list` / `host_asset_list` / `web_app_framework_list` | 可空参；常见补 `service`、`service_class`、`is_public`、`keyword` | @@ -90,8 +90,8 @@ tdp_log_search(time_from=today_start, time_to=today_end, sql="...") | `tdp_log_search` | `size` | `page_size`、`cur_page` | `/api/v1/log/searchBySql` 只支持返回数量控制，不支持分页参数 | | `tdp_log_search.sql` | 过滤表达式 | `SELECT * FROM alert` | TDP 只接受类似 WHERE 条件的表达式，不接受完整 SQL 查询 | | `tdp_log_search(action="terms")` | 可只传 `term` | 强制补 `sql` | `/api/v1/log/terms` 的 `sql` 是可选过滤条件 | -| `tdp_threat_inbound_attack` | `result_list` | `result` | API 文档字段为 `condition.result_list` | -| `tdp_incident_list` | `result` | `result_list` | 事件搜索 API 文档字段为 `condition.result` | +| `tdp_threat_external_attack` | `result_list` | `result` | API 文档字段为 `condition.result_list` | +| `tdp_threat_intelligent_aggregation` | `result` | `result_list` | 事件搜索 API 文档字段为 `condition.result` | 示例： @@ -109,19 +109,21 @@ tdp_log_search(time_from=today_start, time_to=today_end, sql="...") | 用户实际要查什么 | 推荐 tool | 说明 | |---|---|---| -| 威胁事件 / 攻击事件 / 事件总览 | `tdp_incident_list` / `tdp_threat_inbound_attack` | 事件维度，平台已聚合 | -| 实时监控 / 威胁监控列表 | `tdp_threat_monitor_list` | 实时威胁监控列表，单次时间范围不得超过 24 小时 | +| 智能聚合事件 / 聚合攻击事件 / 攻击过程 | `tdp_threat_intelligent_aggregation` | 智能聚合维度，一条事件可关联多条告警 | +| 外部攻击严重性分布 | `tdp_threat_external_attack` | 外部攻击统计维度 | +| 威胁事件列表 / 实时监控 / 威胁监控列表 | `tdp_threat_monitor_list` | 威胁实时监控列表，单次时间范围不得超过 24 小时 | | 告警 / 告警日志 / 原始检测记录 | `tdp_log_search` | 告警维度，一条就是一条原始记录 | -| 告警主机 / 受害主机 / 主机下事件 | `tdp_host_threat_list` | 主机维度，按主机聚合 | +| 告警主机 / 受害主机 / 主机下事件 | `tdp_threat_alert_host` | 主机维度，按主机聚合 | | 漏洞、弱口令、登录入口、API 风险等 | 对应资产或风险类 tool | 资产/风险维度，不要混进日志查询 | 建议按下面的用词来路由： - 提到“告警”“最近一小时告警”“查某 IP 的告警”时，默认优先 `tdp_log_search` -- 提到“威胁事件”“攻击事件”“看下最近有什么事件”时，优先 `tdp_incident_list` -- 提到“实时监控”“威胁监控列表”时，优先 `tdp_threat_monitor_list` -- 提到“哪些主机被打了”“告警主机”“受害主机”时，优先 `tdp_host_threat_list` -- 用户没说清时，默认把“明细”理解为告警日志，把“总览/聚合”理解为事件 +- 明确提到“智能聚合”“聚合攻击事件”“智能聚合事件”“攻击过程”时，优先 `tdp_threat_intelligent_aggregation` +- 提到“外部攻击”“外部攻击严重性分布”时，优先 `tdp_threat_external_attack` +- 提到“威胁事件列表”“实时监控”“威胁实时监控列表”“威胁监控列表”时，优先 `tdp_threat_monitor_list` +- 提到“哪些主机被打了”“告警主机”“受害主机”时，优先 `tdp_threat_alert_host` +- 用户没说清时，默认把“事件列表/实时监控”理解为威胁监控列表，把“总览/聚合”理解为智能聚合事件 ## 高频场景 @@ -279,18 +281,18 @@ machine = '192.168.1.100' - 统计某时间段内最多的威胁名称 - 聚合源 IP、目的 IP、URL、威胁类型 -### 4. 威胁事件列表 +### 4. 智能聚合事件查询 推荐： -- `tdp_incident_list` -- `action=search` +- `tdp_threat_intelligent_aggregation` +- `action=incident_search` 最小可运行参数集： ```json { - "action": "search", + "action": "incident_search", "time_from": 1741536000, "time_to": 1741622400, "severity": [3, 4], @@ -308,12 +310,12 @@ machine = '192.168.1.100' 高频 action： -- `search`: 事件列表 +- `incident_search`: 查询智能聚合后的攻击事件列表；不是“威胁-实时监控”中的威胁事件列表 - `top_attacked_entity`: 受攻击实体 -- `result`: 事件研判结果 -- `timeline`: 时间线 -- `alert_search`: 事件下告警列表 -- `result_distribution`: 结果分布 +- `attack_success`: 攻击成功统计 +- `attack_timeline`: 攻击过程时间线 +- `attack_alert_list`: 攻击过程告警详情 +- `attack_result_distribution`: 成功、失败、尝试的数量分布 - `attacker_ip_list`: 攻击者 IP 列表 - `attacker_ip_detail`: 攻击者 IP 详情 @@ -321,14 +323,14 @@ machine = '192.168.1.100' ```json { - "action": "timeline", + "action": "attack_timeline", "incident_id": "b62899499fec914d6246137eed3b6ec4-1777334454", "time_from": 1777305600, "time_to": 1777391999 } ``` -`timeline` 会默认传 `show_attack=true`。如果 TDP 仍返回 `show_attack is false`，说明该事件当前不支持攻击过程展开，改用 `result`、`alert_search` 或回退浏览器查看事件详情。 +`attack_timeline` 会默认传 `show_attack=true`。如果 TDP 仍返回 `show_attack is false`，说明该事件当前不支持攻击过程展开，改用 `attack_success`、`attack_alert_list` 或回退浏览器查看事件详情。 返回结果重点关注： @@ -357,13 +359,13 @@ machine = '192.168.1.100' } ``` -注意：这里的攻击结果字段是 `result_list`，来自 API 文档的 `condition.result_list`；不要使用事件列表里的 `result` 参数名。 +注意：这里的攻击结果字段是 `result_list`，来自 API 文档的 `condition.result_list`；不要使用智能聚合事件查询里的 `result` 参数名。 告警主机汇总： ```json { - "action": "summary", + "action": "alert_host_list", "time_from": 1741536000, "time_to": 1741622400, "severity": [3, 4], @@ -374,13 +376,13 @@ machine = '192.168.1.100' } ``` -`summary` 默认覆盖全部常见威胁类型，且 `threat_characters` 默认不限制。用户明确要“失陷主机”时再传 `threat_characters=["is_compromised"]`。 +`alert_host_list` 默认覆盖全部常见威胁类型，且 `threat_characters` 默认不限制。用户明确要“失陷主机”时再传 `threat_characters=["is_compromised"]`。 某主机下事件： ```json { - "action": "events", + "action": "host_threat_list", "asset_machine": "asset-123", "time_from": 1741536000, "time_to": 1741622400 diff --git a/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr.handler.py b/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr.handler.py index d456738fa..2680c140e 100644 --- a/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr.handler.py +++ b/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr.handler.py @@ -24,6 +24,7 @@ from urllib.parse import urljoin, urlparse from flocks.browser import helpers +from flocks.browser.admin import ensure_daemon from flocks.config.config_writer import ConfigWriter from flocks.tool.registry import ToolContext, ToolResult @@ -31,6 +32,7 @@ LEGACY_SERVICE_ID = "sangfor_edr" USERNAME_SECRET_ID = "sangfor_edr_username" PASSWORD_SECRET_ID = "sangfor_edr_password" +TOKEN_SECRET_ID = "sangfor_edr_token" DEFAULT_AUTH_STATE_PATH = "~/.flocks/browser/sangfor-edr/auth-state.json" DEFAULT_LOGIN_PATH = "/ui/login.php" DEFAULT_INDEX_PATH = "/ui/#/index" @@ -245,6 +247,7 @@ def _saved_auto_login_status(params: dict[str, Any]) -> dict[str, Any]: "has_base_url": has_base_url, "has_saved_username": has_username, "has_saved_password": has_password, + "has_saved_token": bool(secrets.get(TOKEN_SECRET_ID)), "can_auto_refresh": has_base_url and has_username and has_password, } @@ -326,6 +329,106 @@ def _open_page(url: str) -> None: helpers.wait_for_load(timeout=15) +def _ensure_browser_daemon() -> None: + """Start the browser daemon or replace a stale/incompatible instance.""" + ensure_daemon(wait=15.0, _open_inspect=False) + + +def _browser_daemon_result(exc: Exception, auth_state_path: Path) -> dict[str, Any]: + return { + "success": False, + "valid": False, + "status": "browser_daemon_not_ready", + "reason": "browser_daemon_not_ready", + "error": str(exc), + "next_action": "run `flocks browser --setup`, then `flocks browser --doctor`, then retry", + "auth_state_path": str(auth_state_path), + } + + +def _browser_page_open_result(exc: Exception, cfg: RuntimeConfig) -> dict[str, Any]: + return { + "success": False, + "valid": False, + "status": "browser_login_page_open_failed", + "reason": "browser_login_page_open_failed", + "error": str(exc), + "next_action": "verify base_url and EDR connectivity, then retry", + "login_url": _login_url(cfg), + "auth_state_path": str(cfg.auth_state_path), + } + + +def _manual_login_result(cfg: RuntimeConfig, reason: str, error: str = "") -> dict[str, Any]: + return { + "success": False, + "valid": False, + "status": "manual_login_required", + "reason": reason, + "error": error, + "login_url": _login_url(cfg), + "auth_state_path": str(cfg.auth_state_path), + "browser_left_open": True, + "next_action": "complete login in the opened browser, then call `complete_manual_login`", + } + + +def _install_login_token_capture() -> bool: + script = r"""(() => { + const key = "__flocks_sangfor_edr_token"; + sessionStorage.removeItem(key); + if (window.__flocksSangforEdrTokenHooked) return true; + window.__flocksSangforEdrTokenHooked = true; + const capture = (text) => { + try { + const token = JSON.parse(text)?.data?.token; + if (token !== undefined && token !== null && String(token)) sessionStorage.setItem(key, String(token)); + } catch (_) {} + }; + const originalFetch = window.fetch; + window.fetch = async function(...args) { + const response = await originalFetch.apply(this, args); + const url = String(args[0]?.url || args[0] || ""); + if (url.includes("/launch_login.php")) { + try { capture(await response.clone().text()); } catch (_) {} + } + return response; + }; + const originalOpen = XMLHttpRequest.prototype.open; + XMLHttpRequest.prototype.open = function(method, url, ...rest) { + this.__flocksEdrLogin = String(url || "").includes("/launch_login.php"); + return originalOpen.call(this, method, url, ...rest); + }; + const originalSend = XMLHttpRequest.prototype.send; + XMLHttpRequest.prototype.send = function(...args) { + if (this.__flocksEdrLogin) this.addEventListener("load", () => { + try { capture(this.responseType && this.responseType !== "text" ? JSON.stringify(this.response) : this.responseText); } + catch (_) {} + }, {once: true}); + return originalSend.apply(this, args); + }; + return true; +})()""" + try: + return bool(helpers.js(script)) + except Exception: + return False + + +def _save_captured_login_token(timeout: float = 2.0) -> bool: + deadline = time.time() + timeout + while time.time() < deadline: + try: + token = str(helpers.js("sessionStorage.getItem('__flocks_sangfor_edr_token') || ''") or "").strip() + if token: + _get_secret_manager().set(TOKEN_SECRET_ID, token) + return True + except Exception: + return False + time.sleep(0.1) + return False + + def _page_text() -> str: try: return str(helpers.js("document.body ? document.body.innerText : ''") or "") @@ -393,6 +496,13 @@ def _validate_auth_state(cfg: RuntimeConfig) -> dict[str, Any]: "reason": "auth_state_not_found", "auth_state_path": str(cfg.auth_state_path), } + try: + _ensure_browser_daemon() + except Exception as exc: + result = _browser_daemon_result(exc, cfg.auth_state_path) + result["reason"] = "auth_state_load_failed_browser_daemon_not_ready" + return result + try: loaded = helpers.load_state(cfg.auth_state_path, url=_index_url(cfg)) helpers.wait_for_load(timeout=15) @@ -770,6 +880,15 @@ def _wait_for_login_success(cfg: RuntimeConfig) -> bool: return False +def _manual_login_reason() -> str: + text = _page_text().lower() + if any(marker in text for marker in ("动态口令", "短信验证码", "二次认证", "mfa", "otp")): + return "mfa_required" + if any(marker in text for marker in ("ukey", "usb key", "证书登录")): + return "certificate_login_required" + return "login_form_not_ready" + + def _missing_login_inputs(cfg: RuntimeConfig) -> list[str]: missing = [] if not cfg.username: @@ -781,33 +900,36 @@ def _missing_login_inputs(cfg: RuntimeConfig) -> list[str]: def _refresh_auth_state_with_cdp_login(cfg: RuntimeConfig, captcha_code: str = "") -> dict[str, Any]: missing = _missing_login_inputs(cfg) + try: + _ensure_browser_daemon() + except Exception as exc: + return _browser_daemon_result(exc, cfg.auth_state_path) + try: + _open_page(_login_url(cfg)) + except Exception as exc: + return _browser_page_open_result(exc, cfg) + _install_login_token_capture() if missing: - return { - "success": False, - "status": "manual_login_required", - "reason": "missing_cdp_login_credentials", - "missing": missing, - "auth_state_path": str(cfg.auth_state_path), - } - - _open_page(_login_url(cfg)) - form_state = _wait_for_login_form_ready(cfg) + result = _manual_login_result(cfg, "missing_cdp_login_credentials") + result["missing"] = missing + return result + try: + form_state = _wait_for_login_form_ready(cfg) + except Exception as exc: + return _manual_login_result(cfg, _manual_login_reason(), str(exc)) last_error = "" for attempt in range(1, cfg.max_captcha_retry + 1): try: + _install_login_token_capture() code = captcha_code.strip() if not code: if not cfg.auto_ocr_code: - return { - "success": False, - "status": "manual_login_required", - "reason": "captcha_code_required", - "auth_state_path": str(cfg.auth_state_path), - } + return _manual_login_result(cfg, "captcha_code_required") code = _ocr_verify_code(_captcha_image_from_browser(cfg)) filled = _set_login_form_values(cfg, code) if _wait_for_login_success(cfg): + token_saved = _save_captured_login_token() saved = _save_auth_state(cfg) return { "success": True, @@ -817,6 +939,7 @@ def _refresh_auth_state_with_cdp_login(cfg: RuntimeConfig, captcha_code: str = " "form": form_state, "filled": {key: bool(value) for key, value in filled.items()}, "saved": saved, + "token_saved": token_saved, } last_error = "login_success_check_timeout" except Exception as exc: @@ -830,13 +953,39 @@ def _refresh_auth_state_with_cdp_login(cfg: RuntimeConfig, captcha_code: str = " except Exception: pass - return { - "success": False, - "status": "manual_login_required", - "reason": "browser_cdp_login_failed", - "last_error": last_error, - "auth_state_path": str(cfg.auth_state_path), - } + return _manual_login_result(cfg, "browser_cdp_login_failed", last_error) + + +def _complete_manual_login(cfg: RuntimeConfig) -> dict[str, Any]: + try: + _ensure_browser_daemon() + except Exception as exc: + return _browser_daemon_result(exc, cfg.auth_state_path) + try: + info = helpers.page_info() + if urlparse(str(info.get("url") or "")).netloc != urlparse(cfg.base_url).netloc: + _open_page(_index_url(cfg)) + if not _is_logged_in(cfg): + return _manual_login_result(cfg, "manual_login_not_completed") + token_saved = _save_captured_login_token() + return { + "success": True, + "valid": True, + "status": "manual_login_captured_auth_state", + "auth_state_path": str(cfg.auth_state_path), + "saved": _save_auth_state(cfg), + "token_saved": token_saved, + } + except Exception as exc: + return { + "success": False, + "valid": False, + "status": "manual_login_capture_failed", + "reason": "manual_login_capture_failed", + "error": str(exc), + "auth_state_path": str(cfg.auth_state_path), + "next_action": "keep the browser open and retry `complete_manual_login`", + } # ── Tool actions ───────────────────────────────────────────────────────────── @@ -881,12 +1030,24 @@ async def handle(ctx: ToolContext) -> ToolResult: if action == "validate_auth_state": validation = _validate_auth_state(cfg) - return ToolResult(success=bool(validation.get("valid")), output=validation) + return ToolResult( + success=bool(validation.get("valid")), + output=validation, + error=None if validation.get("valid") else str(validation.get("reason") or "auth_state_invalid"), + ) + + if action == "complete_manual_login": + result = _complete_manual_login(cfg) + return ToolResult( + success=bool(result.get("success")), + output=result, + error=None if result.get("success") else result.get("reason"), + ) if action not in {"ensure_auth_state", "refresh_auth_state"}: return ToolResult( success=False, - error="Unsupported Sangfor EDR auth action. Use status_auth_state, ensure_auth_state, validate_auth_state, or refresh_auth_state.", + error="Unsupported Sangfor EDR auth action. Use status_auth_state, ensure_auth_state, validate_auth_state, refresh_auth_state, or complete_manual_login.", ) force_refresh = action == "refresh_auth_state" or _coerce_bool(params.get("force_refresh"), default=False) diff --git a/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr_auth.yaml b/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr_auth.yaml index 730338b81..1b3d060ba 100644 --- a/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr_auth.yaml +++ b/.flocks/plugins/tools/device/sangfor_edr_webcli/sangfor_edr_auth.yaml @@ -3,11 +3,13 @@ description: > Manage Sangfor EDR browser authentication state. Use status_auth_state or validate_auth_state before opening pages; use ensure_auth_state to reuse a valid saved auth-state or automatically refresh full browser state through - browser daemon / CDP when credentials are configured. + browser daemon / CDP when credentials are configured. Browser actions start + or recover the daemon automatically and return structured recovery guidance. description_cn: > 管理深信服 EDR 浏览器登录态。先用 status_auth_state 或 validate_auth_state 检查已保存 state；当 state 缺失或失效且已配置账密时，ensure_auth_state 通过 browser daemon / CDP 驱动真实登录页自动登录并重新保存完整浏览器 state。 + 需要浏览器的 action 会自动启动或恢复 daemon，失败时返回结构化恢复指引。 category: custom enabled: true requires_confirmation: false @@ -17,13 +19,14 @@ inputSchema: properties: action: type: string - enum: [status_auth_state, ensure_auth_state, validate_auth_state, refresh_auth_state] + enum: [status_auth_state, ensure_auth_state, validate_auth_state, refresh_auth_state, complete_manual_login] default: ensure_auth_state description: > Authentication action. status_auth_state returns non-sensitive saved state/credential availability; ensure_auth_state reuses existing state when it appears valid and refreshes when needed; validate_auth_state - checks only; refresh_auth_state always performs CDP-assisted browser login. + checks only; refresh_auth_state always performs CDP-assisted browser login; + complete_manual_login validates a user-completed login and saves its browser state. captcha_code: type: string description: Optional manual captcha code. When omitted, OCR is used if enabled. diff --git a/.flocks/plugins/tools/device/skyeye_v4_0_14_0_SP2/skyeye.handler.py b/.flocks/plugins/tools/device/skyeye_v4_0_14_0_SP2/skyeye.handler.py index e6f8f94d4..f00129ce6 100644 --- a/.flocks/plugins/tools/device/skyeye_v4_0_14_0_SP2/skyeye.handler.py +++ b/.flocks/plugins/tools/device/skyeye_v4_0_14_0_SP2/skyeye.handler.py @@ -1,4 +1,5 @@ import base64 +import gzip import hashlib import json import random @@ -189,6 +190,12 @@ def _clean_params(params: dict[str, Any]) -> dict[str, Any]: return cleaned +def _encode_ip_param(value: str | None) -> str | None: + if value is None or not value.strip() or value.startswith("H4sI"): + return value + return base64.b64encode(gzip.compress(value.encode("utf-8"))).decode("ascii") + + def _payload_error(payload: Any) -> str | None: if not isinstance(payload, dict): return None @@ -498,8 +505,8 @@ async def alarm_list( "status": status, "serial_num": serial_num, "data_source": data_source, - "alarm_sip": alarm_sip, - "attack_sip": attack_sip, + "alarm_sip": _encode_ip_param(alarm_sip), + "attack_sip": _encode_ip_param(attack_sip), "ioc": ioc, "threat_name": threat_name, "host": host, diff --git a/.flocks/plugins/tools/device/tdp_v3_3_10/_test.yaml b/.flocks/plugins/tools/device/tdp_v3_3_10/_test.yaml index 1eae486e1..97439e990 100644 --- a/.flocks/plugins/tools/device/tdp_v3_3_10/_test.yaml +++ b/.flocks/plugins/tools/device/tdp_v3_3_10/_test.yaml @@ -49,12 +49,12 @@ fixtures: assert: success: true - tdp_incident_list: - - label: "Search incidents (page 1)" - label_cn: "查询威胁事件列表（第 1 页）" + tdp_threat_intelligent_aggregation: + - label: "Search intelligent aggregation incidents (page 1)" + label_cn: "查询智能聚合威胁事件（第 1 页）" tags: [smoke, incidents] params: - action: search + action: incident_search cur_page: 1 page_size: 20 assert: @@ -66,12 +66,12 @@ fixtures: params: action: top_attacked_entity - tdp_host_threat_list: - - label: "Get threat host summary" - label_cn: "查询威胁主机摘要" + tdp_threat_alert_host: + - label: "List alert hosts" + label_cn: "查询告警主机列表" tags: [smoke, threats] params: - action: summary + action: alert_host_list cur_page: 1 page_size: 20 assert: @@ -147,15 +147,15 @@ fixtures: action: search size: 20 - tdp_threat_inbound_attack: - - label: "Get inbound attack overview (last 24 hours)" - label_cn: "查询外部入侵攻击概览（最近 24 小时）" + tdp_threat_external_attack: + - label: "Get external attack severity distribution (last 24 hours)" + label_cn: "查询外部攻击严重性分布（最近 24 小时）" tags: [smoke, threats] params: {} tdp_threat_monitor_list: - - label: "List real-time threats (last 24 hours)" - label_cn: "查询威胁实时监控列表（最近 24 小时）" + - label: "List monitored threats (last 24 hours)" + label_cn: "查询威胁监控列表（最近 24 小时）" tags: [smoke, threats] params: cur_page: 1 diff --git a/.flocks/plugins/tools/device/tdp_v3_3_10/tdp.handler.py b/.flocks/plugins/tools/device/tdp_v3_3_10/tdp.handler.py index 1e39d02b2..649355021 100644 --- a/.flocks/plugins/tools/device/tdp_v3_3_10/tdp.handler.py +++ b/.flocks/plugins/tools/device/tdp_v3_3_10/tdp.handler.py @@ -479,7 +479,7 @@ def _condition_time_body(body: dict[str, Any]) -> dict[str, Any]: def _incident_search_body(body: dict[str, Any]) -> dict[str, Any]: defaults = _body_with_condition_time( { - "condition": {"duration": {"begin_duration": 0, "end_duration": 24}}, + "condition": {"duration": {"begin_duration": 0}}, "page": {"cur_page": 1, "page_size": 20, "sort": [{"sort_by": "last_time", "sort_order": "desc"}]}, } ) @@ -845,17 +845,17 @@ def _passthrough_body(body: Any) -> Any: } THREAT_HOST_ACTIONS: JsonActionMap = { - "summary": ("/api/v1/host/getFallHostSumList", _threat_host_summary_body, "host_get_fall_host_sum_list"), - "events": ("/api/v1/host/threat/list", _threat_host_event_body, "host_threat_list"), + "alert_host_list": ("/api/v1/host/getFallHostSumList", _threat_host_summary_body, "host_get_fall_host_sum_list"), + "host_threat_list": ("/api/v1/host/threat/list", _threat_host_event_body, "host_threat_list"), } INCIDENT_ACTIONS: JsonActionMap = { - "search": ("/api/v1/incident/search", _incident_search_body, "incident_search"), + "incident_search": ("/api/v1/incident/search", _incident_search_body, "incident_search"), "top_attacked_entity": ("/api/v1/incident/topAttackedEntity", _condition_time_body, "incident_top_attacked_entity"), - "result": ("/api/v1/incident/result", _dashboard_body, "incident_result"), - "timeline": ("/api/v1/incident/timeline", _dashboard_body, "incident_timeline"), - "alert_search": ("/api/v1/alert/search", _condition_time_body, "incident_alert_search"), - "result_distribution": ( + "attack_success": ("/api/v1/incident/result", _dashboard_body, "incident_result"), + "attack_timeline": ("/api/v1/incident/timeline", _dashboard_body, "incident_timeline"), + "attack_alert_list": ("/api/v1/alert/search", _condition_time_body, "incident_alert_search"), + "attack_result_distribution": ( "/api/v1/incident/result/distribution", _condition_time_body, "incident_result_distribution", @@ -1267,7 +1267,7 @@ async def dashboard_status( async def threat_host_list( context: ToolContext, - action: str = "summary", + action: str = "alert_host_list", condition: dict[str, Any] | None = None, page: dict[str, Any] | None = None, time_from: int | None = None, @@ -1289,7 +1289,7 @@ async def threat_host_list( sort_order: str | None = None, ) -> ToolResult: del context - selected_action = _normalize_action(action, "summary") + selected_action = _normalize_action(action, "alert_host_list") body = _compose_payload(condition=condition, page=page) condition_dict = body.setdefault("condition", {}) _set_if_present(condition_dict, "time_from", time_from) @@ -1306,16 +1306,21 @@ async def threat_host_list( _set_list_if_present(condition_dict, "host_type", host_type) _set_keyword_fuzzy(condition_dict, keyword, HOST_THREAT_FUZZY_FIELDS) _set_page_overrides(body, cur_page=cur_page, page_size=page_size, sort_by=sort_by, sort_order=sort_order) - if selected_action == "events": + if selected_action == "host_threat_list": validation_error = _validate_required_body_fields(selected_action, body, "condition.asset_machine") if validation_error: return validation_error - return await _run_action_json_tool(THREAT_HOST_ACTIONS, default_action="summary", action=selected_action, body=body) + return await _run_action_json_tool( + THREAT_HOST_ACTIONS, + default_action="alert_host_list", + action=selected_action, + body=body, + ) async def incident_list( context: ToolContext, - action: str = "search", + action: str = "incident_search", condition: dict[str, Any] | None = None, page: dict[str, Any] | None = None, incident_id: str | None = None, @@ -1339,9 +1344,9 @@ async def incident_list( sort_order: str | None = None, ) -> ToolResult: del context - selected_action = _normalize_action(action, "search") + selected_action = _normalize_action(action, "incident_search") body = _compose_payload(condition=condition, page=page) - if selected_action == "search": + if selected_action == "incident_search": condition_dict = body.setdefault("condition", {}) _set_if_present(condition_dict, "time_from", time_from) _set_if_present(condition_dict, "time_to", time_to) @@ -1354,7 +1359,7 @@ async def incident_list( _set_if_present(duration, "begin_duration", begin_duration) _set_if_present(duration, "end_duration", end_duration) _set_keyword_fuzzy(condition_dict, keyword, INCIDENT_SEARCH_FUZZY_FIELDS) - elif selected_action == "alert_search": + elif selected_action == "attack_alert_list": condition_dict = body.setdefault("condition", {}) _set_if_present(condition_dict, "time_from", time_from) _set_if_present(condition_dict, "time_to", time_to) @@ -1372,17 +1377,17 @@ async def incident_list( _set_if_present(body, "time_from", time_from) _set_if_present(body, "time_to", time_to) _set_if_present(body, "include_risk", include_risk) - if selected_action == "timeline": + if selected_action == "attack_timeline": body["show_attack"] = True if show_attack is None else show_attack if attacker is not None: body["attacker"] = list(attacker) _set_page_overrides(body, cur_page=cur_page, page_size=page_size, sort_by=sort_by, sort_order=sort_order) validation_map: dict[str, tuple[str, ...]] = { "top_attacked_entity": ("incident_id",), - "result": ("incident_id",), - "timeline": ("incident_id",), - "alert_search": ("condition.id", "page"), - "result_distribution": ("incident_id",), + "attack_success": ("incident_id",), + "attack_timeline": ("incident_id",), + "attack_alert_list": ("condition.id", "page"), + "attack_result_distribution": ("incident_id",), "attacker_ip_list": ("condition.incident_id",), "attacker_ip_detail": ("incident_id", "attacker"), } @@ -1391,7 +1396,12 @@ async def incident_list( validation_error = _validate_required_body_fields(selected_action, body, *required_fields) if validation_error: return validation_error - return await _run_action_json_tool(INCIDENT_ACTIONS, default_action="search", action=selected_action, body=body) + return await _run_action_json_tool( + INCIDENT_ACTIONS, + default_action="incident_search", + action=selected_action, + body=body, + ) async def service_asset_list( diff --git a/.flocks/plugins/tools/device/tdp_v3_3_10/tdp_incident_list.yaml b/.flocks/plugins/tools/device/tdp_v3_3_10/tdp_incident_list.yaml deleted file mode 100644 index b724a7106..000000000 --- a/.flocks/plugins/tools/device/tdp_v3_3_10/tdp_incident_list.yaml +++ /dev/null @@ -1,109 +0,0 @@ -name: tdp_incident_list -description: > - 查询 TDP 智能聚合相关接口。默认返回事件列表（注意不是告警），也可通过 action 切换到时间线、 - 结果分布、攻击实体、攻击者 IP 明细等子接口。优先使用顶层语义化参数，不要手写 - `condition.id`、`condition.duration`、`condition.fuzzy` 这类嵌套字段；只有在显式参数未覆盖时才使用 `condition` / `page`。 -category: custom -enabled: true -requires_confirmation: false -provider: tdp_api - -inputSchema: - type: object - properties: - action: - type: string - description: > - 智能聚合子接口动作。`search` 会自动补齐 condition.time_from、 - condition.time_to、duration 和分页默认值；`top_attacked_entity`、 - `result`、`timeline`、`result_distribution` 需要传 - `incident_id`；`alert_search` 需要传 `alert_ids` 或 `condition.id`， - 且应同时提供 `page`；`attacker_ip_list` 需要传 `condition.incident_id`； - `attacker_ip_detail` 需要传 `incident_id` 和 `attacker`。 - enum: [search, top_attacked_entity, result, timeline, alert_search, result_distribution, attacker_ip_list, attacker_ip_detail] - default: search - incident_id: - type: string - description: 事件 ID，适用于 `top_attacked_entity`、`result`、`timeline`、`result_distribution`、`attacker_ip_detail`；应先从 `search` 返回结果中获取。 - time_from: - type: integer - description: 开始时间戳，Unix 秒级时间戳，按动作写入根级或 `condition.time_from`。可以使用python datetime动态计算。 - time_to: - type: integer - description: 结束时间戳，Unix 秒级时间戳，按动作写入根级或 `condition.time_to`。可以使用python datetime动态计算。 - include_risk: - type: boolean - description: 是否包含风险告警；`timeline` 等动作写入根级，`alert_search` 写入 `condition.include_risk`。 - show_attack: - type: boolean - description: > - 是否展示攻击过程，适用于 `timeline`。handler 默认传 `true`； - 若显式传 `false`，TDP 可能返回 `show_attack is false`。 - include_action: - type: boolean - description: 是否包含事件告警，适用于 `alert_search`，会写入 `condition.include_action`。 - alert_ids: - type: array - items: - type: string - description: 告警 ID 列表，适用于 `alert_search`；会转换写入 `condition.id`。 - attacker: - type: array - items: - type: string - description: 攻击者 IP 列表，适用于 `attacker_ip_detail`；通常先从 `attacker_ip_list` 返回结果中获取。 - severity: - type: array - items: - type: integer - enum: [0, 1, 2, 3, 4] - description: 严重级别列表，适用于 `search`。 - phase: - type: array - items: - type: string - description: 攻击阶段列表，适用于 `search`，例如 `recon`、`exploit`、`control`、`attack_out`、`post_exploit`。 - result: - type: array - items: - type: string - enum: [success, failed, unknown] - description: 攻击结果列表，适用于 `search`。 - is_target_attack: - type: boolean - description: 是否只看针对性攻击事件，适用于 `search`。 - begin_duration: - type: integer - description: 持续时间起始小时数，适用于 `search`，会写入 `condition.duration.begin_duration`。 - end_duration: - type: integer - description: 持续时间结束小时数，适用于 `search`，会写入 `condition.duration.end_duration`。 - keyword: - type: string - description: > - 关键词搜索，适用于 `search`。handler 会自动组装 - `condition.fuzzy.fieldlist=["attacker_ip","host_ip","attack_name","attack_tool","incident_id"]`。 - cur_page: - type: integer - description: 页码，适用于 `search`、`alert_search`、`attacker_ip_list`。 - page_size: - type: integer - description: 每页条数，适用于 `search`、`alert_search`、`attacker_ip_list`。 - sort_by: - type: string - description: 排序字段，适用于列表类 action。 - sort_order: - type: string - enum: [asc, desc] - description: 排序方向，适用于列表类 action。 - condition: - type: object - description: 高级兼容参数。仅在顶层参数未覆盖目标字段时使用。 - page: - type: object - description: 高级兼容分页参数。`alert_search` 如需复杂排序可直接透传。 - required: [] -handler: - type: script - script_file: tdp.handler.py - function: incident_list diff --git a/.flocks/plugins/tools/device/tdp_v3_3_10/tdp_host_threat_list.yaml b/.flocks/plugins/tools/device/tdp_v3_3_10/tdp_threat_alert_host.yaml similarity index 67% rename from .flocks/plugins/tools/device/tdp_v3_3_10/tdp_host_threat_list.yaml rename to .flocks/plugins/tools/device/tdp_v3_3_10/tdp_threat_alert_host.yaml index 84829590f..ca11f5091 100644 --- a/.flocks/plugins/tools/device/tdp_v3_3_10/tdp_host_threat_list.yaml +++ b/.flocks/plugins/tools/device/tdp_v3_3_10/tdp_threat_alert_host.yaml @@ -1,7 +1,7 @@ -name: tdp_host_threat_list +name: tdp_threat_alert_host description: > - 查询 TDP 告警主机相关数据。`summary` 返回告警主机列表，适合按严重级别、威胁性质、威胁类型、业务组和关键词筛选；`events` - 返回某台告警主机下的威胁事件列表，通常需要先从 `summary` 结果里拿到 `assets_machine` 再继续查询。优先使用顶层语义化参数，不要手写 + 查询 TDP 告警主机相关数据。`alert_host_list` 返回告警主机列表；`host_threat_list` + 返回指定告警主机的威胁列表，通常需要先从 `alert_host_list` 结果里取得 `assets_machine`。优先使用顶层语义化参数，不要手写 `condition` / `page`，只有在需要透传尚未覆盖的底层字段时才使用这两个高级兼容参数。 category: custom enabled: true @@ -14,10 +14,10 @@ inputSchema: action: type: string description: > - 告警主机子接口动作。`summary` 会自动补齐默认时间范围、全部告警主机的威胁类型、空的威胁特征和默认分页； - `events` 用于查看单台主机下的威胁事件，要求提供 `asset_machine`，级联场景可额外传 `device_id`。 - enum: [summary, events] - default: summary + `alert_host_list` 查询告警主机列表，可按严重级别、威胁性质、威胁类型、业务组和关键词筛选； + `host_threat_list` 查询指定 `asset_machine` 的威胁列表，级联场景可额外传 `device_id`。 + enum: [alert_host_list, host_threat_list] + default: alert_host_list time_from: type: integer description: 开始时间戳，Unix 秒级时间戳，会写入 `condition.time_from`。可以使用python datetime动态计算。 @@ -26,7 +26,7 @@ inputSchema: description: 结束时间戳，Unix 秒级时间戳，会写入 `condition.time_to`。可以使用python datetime动态计算。 asset_machine: type: string - description: 资产机器标识，`events` 动作必填；应从 `summary` 结果中的主机标识字段获取。 + description: 资产机器标识，`host_threat_list` 必填；应从 `alert_host_list` 结果中的主机标识字段获取。 device_id: type: string description: 设备 ID，级联平台场景使用；非级联平台通常可省略。 @@ -52,7 +52,7 @@ inputSchema: type: array items: type: string - description: 威胁类型列表。`summary` 默认已覆盖全部常见威胁类型；仅在需要缩小范围时传入。 + description: 威胁类型列表。`alert_host_list` 默认已覆盖全部常见威胁类型；仅在需要缩小范围时传入。 disposal_status: type: array items: @@ -64,21 +64,21 @@ inputSchema: type: array items: type: string - description: 主机类型列表，适用于 `summary`，常见值对应服务器或终端视角。 + description: 主机类型列表，适用于 `alert_host_list`，常见值对应服务器或终端视角。 assets_group: type: array items: type: integer - description: 业务组 ID 列表，适用于 `summary`。 + description: 业务组 ID 列表，适用于 `alert_host_list`。 host_type: type: array items: type: string - description: 资产类型列表，适用于 `summary`。 + description: 资产类型列表，适用于 `alert_host_list`。 keyword: type: string description: > - 关键词搜索，适用于 `summary` 和 `events`。handler 会自动组装固定的 + 关键词搜索，适用于 `alert_host_list` 和 `host_threat_list`。handler 会自动组装固定的 `condition.fuzzy.fieldlist=["threat.name","external_ip","machine","assets.name","data","machine_name"]`。 cur_page: type: integer @@ -88,7 +88,7 @@ inputSchema: description: 每页条数，默认 20。 sort_by: type: string - description: 排序字段。`summary` 默认按 `severity`，`events` 默认按 `max_severity`。 + description: 排序字段。`alert_host_list` 默认按 `severity`，`host_threat_list` 默认按 `max_severity`。 sort_order: type: string enum: [asc, desc] diff --git a/.flocks/plugins/tools/device/tdp_v3_3_10/tdp_threat_inbound_attack.yaml b/.flocks/plugins/tools/device/tdp_v3_3_10/tdp_threat_external_attack.yaml similarity index 98% rename from .flocks/plugins/tools/device/tdp_v3_3_10/tdp_threat_inbound_attack.yaml rename to .flocks/plugins/tools/device/tdp_v3_3_10/tdp_threat_external_attack.yaml index 3f1e11611..b982f2bb5 100644 --- a/.flocks/plugins/tools/device/tdp_v3_3_10/tdp_threat_inbound_attack.yaml +++ b/.flocks/plugins/tools/device/tdp_v3_3_10/tdp_threat_external_attack.yaml @@ -1,4 +1,4 @@ -name: tdp_threat_inbound_attack +name: tdp_threat_external_attack description: > 查询 TDP 外部攻击相关数据，当前提供外部攻击严重性分布统计。优先使用顶层语义化参数， 例如严重级别、攻击结果、级联资产组和关键词；`condition` 仅作为高级兼容入口。 diff --git a/.flocks/plugins/tools/device/tdp_v3_3_10/tdp_threat_intelligent_aggregation.yaml b/.flocks/plugins/tools/device/tdp_v3_3_10/tdp_threat_intelligent_aggregation.yaml new file mode 100644 index 000000000..889c6f9f6 --- /dev/null +++ b/.flocks/plugins/tools/device/tdp_v3_3_10/tdp_threat_intelligent_aggregation.yaml @@ -0,0 +1,109 @@ +name: tdp_threat_intelligent_aggregation +description: > + 查询 TDP“威胁-智能聚合”中的聚合攻击事件及其详情。默认查询智能聚合后的攻击事件，也可查询指定聚合事件的 + 攻击成功统计、攻击时间线、攻击过程告警、结果分布、受攻击实体和攻击者 IP 信息。本 tool 不用于查询 + “威胁-实时监控”中的威胁事件列表；该意图应使用 `tdp_threat_monitor_list`。优先使用顶层语义化参数，不要手写 + `condition.id`、`condition.duration`、`condition.fuzzy` 这类嵌套字段；只有在显式参数未覆盖时才使用 `condition` / `page`。 +category: custom +enabled: true +requires_confirmation: false +provider: tdp_api + +inputSchema: + type: object + properties: + action: + type: string + description: > + `incident_search` 查询智能聚合后的攻击事件列表；`top_attacked_entity` 查询指定聚合事件 TOP 受攻击域名或主机； + `attack_success` 查询指定聚合事件的攻击成功统计；`attack_timeline` 查询指定聚合事件的攻击过程时间线； + `attack_alert_list` 根据攻击过程返回的告警 ID 查询对应告警详情；`attack_result_distribution` 查询指定聚合事件中 + 成功、失败、尝试的数量分布；`attacker_ip_list` 查询指定聚合事件的攻击来源 IP； + `attacker_ip_detail` 查询指定聚合事件中指定攻击 IP 的情报和活跃时间段。 + enum: [incident_search, top_attacked_entity, attack_success, attack_timeline, attack_alert_list, attack_result_distribution, attacker_ip_list, attacker_ip_detail] + default: incident_search + incident_id: + type: string + description: 智能聚合事件 ID，适用于除 `incident_search` 和 `attack_alert_list` 外的详情动作；应先从 `incident_search` 返回结果中获取。 + time_from: + type: integer + description: 开始时间戳，Unix 秒级时间戳，按动作写入根级或 `condition.time_from`。可以使用python datetime动态计算。 + time_to: + type: integer + description: 结束时间戳，Unix 秒级时间戳，按动作写入根级或 `condition.time_to`。可以使用python datetime动态计算。 + include_risk: + type: boolean + description: 是否包含风险告警；`attack_timeline` 写入根级，`attack_alert_list` 写入 `condition.include_risk`。 + show_attack: + type: boolean + description: > + 是否展示攻击过程，适用于 `attack_timeline`。handler 默认传 `true`； + 若显式传 `false`，TDP 可能返回 `show_attack is false`。 + include_action: + type: boolean + description: 是否包含敏感行为，适用于 `attack_alert_list`，会写入 `condition.include_action`。 + alert_ids: + type: array + items: + type: string + description: 攻击过程返回的告警 ID 列表，适用于 `attack_alert_list`；会转换写入 `condition.id`。 + attacker: + type: array + items: + type: string + description: 攻击者 IP 列表，适用于 `attacker_ip_detail`；通常先从 `attacker_ip_list` 返回结果中获取。 + severity: + type: array + items: + type: integer + enum: [0, 1, 2, 3, 4] + description: 严重级别列表，适用于 `incident_search`。 + phase: + type: array + items: + type: string + description: 攻击阶段列表，适用于 `incident_search`，例如 `recon`、`exploit`、`control`、`attack_out`、`post_exploit`。 + result: + type: array + items: + type: string + enum: [success, failed, unknown] + description: 攻击结果列表，适用于 `incident_search`。 + is_target_attack: + type: boolean + description: 是否只看针对性攻击事件，适用于 `incident_search`。 + begin_duration: + type: integer + description: 持续时间起始小时数，适用于 `incident_search`，会写入 `condition.duration.begin_duration`。 + end_duration: + type: integer + description: 持续时间结束小时数，适用于 `incident_search`，会写入 `condition.duration.end_duration`。 + keyword: + type: string + description: > + 关键词搜索，适用于 `incident_search`。handler 会自动组装 + `condition.fuzzy.fieldlist=["attacker_ip","host_ip","attack_name","attack_tool","incident_id"]`。 + cur_page: + type: integer + description: 页码，适用于 `incident_search`、`attack_alert_list`、`attacker_ip_list`。 + page_size: + type: integer + description: 每页条数，适用于 `incident_search`、`attack_alert_list`、`attacker_ip_list`。 + sort_by: + type: string + description: 排序字段，适用于列表类 action。 + sort_order: + type: string + enum: [asc, desc] + description: 排序方向，适用于列表类 action。 + condition: + type: object + description: 高级兼容参数。仅在顶层参数未覆盖目标字段时使用。 + page: + type: object + description: 高级兼容分页参数。`attack_alert_list` 如需复杂排序可直接透传。 + required: [] +handler: + type: script + script_file: tdp.handler.py + function: incident_list diff --git a/.flocks/plugins/tools/device/tdp_v3_3_10/tdp_threat_monitor_list.yaml b/.flocks/plugins/tools/device/tdp_v3_3_10/tdp_threat_monitor_list.yaml index 2830e3a1f..0409fa813 100644 --- a/.flocks/plugins/tools/device/tdp_v3_3_10/tdp_threat_monitor_list.yaml +++ b/.flocks/plugins/tools/device/tdp_v3_3_10/tdp_threat_monitor_list.yaml @@ -1,6 +1,6 @@ name: tdp_threat_monitor_list description: > - 查询 TDP 威胁-实时监控列表，对应正式 API `/api/v1/monitor/threat/list`。 + 查询 TDP 威胁事件列表，即“威胁-实时监控”中的威胁实时监控列表，对应正式 API `/api/v1/monitor/threat/list`。 查询时间范围不得超过 24 小时；未传时间时 handler 动态使用最近 24 小时。 `sql` 是 TDP 过滤表达式，不是包含 SELECT/FROM 的完整 SQL。 category: custom diff --git a/README.md b/README.md index 32e0acd68..5f5b0d6c7 100644 --- a/README.md +++ b/README.md @@ -149,7 +149,6 @@ macOS / Linux ```bash docker run -d \ --name flocks \ - -p 8000:8000 \ -p 5173:5173 \ --shm-size 4gb \ -v "${HOME}/.flocks:/home/flocks/.flocks" \ @@ -160,14 +159,13 @@ Windows PowerShell ```powershell docker run -d ` --name flocks ` - -p 8000:8000 ` -p 5173:5173 ` --shm-size 4gb ` -v "${env:USERPROFILE}\.flocks:/home/flocks/.flocks" ` ghcr.io/agentflocks/flocks:latest ``` -`EXPOSE` in the image only documents container ports. You still need `-p 8000:8000 -p 5173:5173` to access the service from the host browser. +`EXPOSE` in the image only documents container ports. You still need `-p 5173:5173` to access the service from the host browser. ## 4. FAQ diff --git a/README_zh.md b/README_zh.md index 84ef87929..7ebba572f 100644 --- a/README_zh.md +++ b/README_zh.md @@ -149,7 +149,6 @@ macOS / Linux docker run -d \ --name flocks \ -e TZ=Asia/Shanghai \ - -p 8000:8000 \ -p 5173:5173 \ --shm-size 4gb \ -v "${HOME}/.flocks:/home/flocks/.flocks" \ @@ -161,14 +160,13 @@ Windows PowerShell docker run -d ` --name flocks ` -e TZ=Asia/Shanghai ` - -p 8000:8000 ` -p 5173:5173 ` --shm-size 4gb ` -v "${env:USERPROFILE}\.flocks:/home/flocks/.flocks" ` ghcr.io/agentflocks/flocks:latest ``` -镜像中的 `EXPOSE` 仅用于声明容器端口；要从宿主机浏览器访问服务，仍需使用 `-p 8000:8000 -p 5173:5173` 映射端口。 +镜像中的 `EXPOSE` 仅用于声明容器端口；要从宿主机浏览器访问服务，仍需使用 `-p 5173:5173` 映射端口。 ## 4. 常见问题 diff --git a/docker/Dockerfile b/docker/Dockerfile index 1b0384f12..a5c39f923 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -10,7 +10,7 @@ ARG VITE_APP_VERSION=docker ENV DEBIAN_FRONTEND=noninteractive \ FLOCKS_DEPLOY_MODE=docker \ FLOCKS_BACKEND_HOST=0.0.0.0 \ - FLOCKS_FRONTEND_HOST=0.0.0.0 \ + FLOCKS_BACKEND_PORT=5173 \ FLOCKS_FRONTEND_DIST_DIR=/opt/flocks/webui/dist \ FLOCKS_ROOT=/home/flocks/.flocks \ FLOCKS_CONFIG_DIR=/home/flocks/.flocks/config \ @@ -88,9 +88,9 @@ RUN sed -i 's/^# *zh_CN.UTF-8 UTF-8/zh_CN.UTF-8 UTF-8/' /etc/locale.gen \ RUN mkdir -p "$HOME/.flocks" "$HOME/.flocks/config" "$HOME/.agent-browser" \ && chown -R flocks:flocks /opt/flocks "$HOME" -EXPOSE 8000 5173 +EXPOSE 5173 USER flocks -ENTRYPOINT ["/usr/bin/tini", "--"] +ENTRYPOINT ["/usr/bin/tini", "-g", "--"] CMD ["bash", "./scripts/container-start.sh"] diff --git a/flocks/agent/agent_factory.py b/flocks/agent/agent_factory.py index d8c08f265..2c2ac9b37 100644 --- a/flocks/agent/agent_factory.py +++ b/flocks/agent/agent_factory.py @@ -223,8 +223,9 @@ def scan_and_load(dirs: Optional[List[Path]] = None) -> Dict[str, AgentInfo]: """ Scan agent directories and load all valid agents. - Scans built-in agents first, then user plugin agents, then project plugin - agents. Name conflicts are skipped with a warning (first wins). + Scans built-in agents first, then user plugin agents, then project-bundled + agents. The first match wins, so user customizations take precedence over + project bundles while core built-ins remain protected. ``native`` is determined by the source directory, not by agent.yaml: - ``_BUILTIN_AGENTS_DIR`` → native=True @@ -238,20 +239,21 @@ def scan_and_load(dirs: Optional[List[Path]] = None) -> Dict[str, AgentInfo]: Returns: Dict mapping agent name → AgentInfo. """ - # (directory, is_native) pairs — order determines first-wins conflict resolution - search_dirs: List[tuple[Path, bool]] = [(_BUILTIN_AGENTS_DIR, True)] + # (directory, is_native, source) tuples — order determines first-wins priority. + search_dirs: List[tuple[Path, bool, str]] = [(_BUILTIN_AGENTS_DIR, True, "builtin")] if _PLUGIN_AGENTS_DIR.exists(): - search_dirs.append((_PLUGIN_AGENTS_DIR, False)) + search_dirs.append((_PLUGIN_AGENTS_DIR, False, "user")) # Project-level plugin agents: resolved at call time because cwd is dynamic project_agents_dir = Path.cwd() / ".flocks" / "plugins" / "agents" if project_agents_dir.exists() and project_agents_dir != _PLUGIN_AGENTS_DIR: - search_dirs.append((project_agents_dir, True)) + search_dirs.append((project_agents_dir, True, "project")) if dirs: - search_dirs.extend((d, False) for d in dirs) + search_dirs.extend((d, False, "extra") for d in dirs) result: Dict[str, AgentInfo] = {} + selected_sources: Dict[str, tuple[str, Path]] = {} - for scan_dir, is_native in search_dirs: + for scan_dir, is_native, source in search_dirs: if not scan_dir.is_dir(): continue for agent_dir in _iter_agent_dirs(scan_dir): @@ -259,13 +261,21 @@ def scan_and_load(dirs: Optional[List[Path]] = None) -> Dict[str, AgentInfo]: if agent is None: continue if agent.name in result: - log.warn("agent.factory.name_conflict", { + selected_source, selected_dir = selected_sources[agent.name] + details = { "name": agent.name, - "existing_source": "previous scan", - "skipped_source": str(agent_dir), - }) + "selected_source": selected_source, + "selected_path": str(selected_dir), + "skipped_source": source, + "skipped_path": str(agent_dir), + } + if source != selected_source and source != "extra": + log.info("agent.factory.lower_priority_skipped", details) + else: + log.warn("agent.factory.name_conflict", details) continue result[agent.name] = agent + selected_sources[agent.name] = (source, agent_dir) log.debug("agent.factory.loaded", { "name": agent.name, "dir": str(agent_dir), diff --git a/flocks/browser/admin.py b/flocks/browser/admin.py index 0ed418e08..c8a4c0065 100644 --- a/flocks/browser/admin.py +++ b/flocks/browser/admin.py @@ -20,7 +20,7 @@ VERSION_CACHE = Path(tempfile.gettempdir()) / "flocks-browser-version-cache.json" VERSION_CACHE_TTL = 24 * 3600 DOCTOR_TEXT_LIMIT = 140 -# run_setup: at most two daemon/CDP attach attempts to avoid repeated Allow prompts. +# run_setup retries transient attach failures once, but exits after manual setup guidance. _SETUP_ATTACH_WAIT = 20.0 _SETUP_RETRY_WAIT = 30.0 @@ -49,10 +49,12 @@ def _log_tail(name: str | None): def _needs_chrome_remote_debugging_prompt(msg: str | None) -> bool: - """Return True when a local browser needs the inspect-page permission flow.""" + """Return True when a local browser needs remote-debugging setup guidance.""" lower = (msg or "").lower() return ( "devtoolsactiveport not found" in lower + or "chrome remote debugging is not reachable" in lower + or "chromium-based browser remote debugging is not reachable" in lower or "remote-debugging page" in lower or "inspect/#remote-debugging" in lower or "not live yet" in lower @@ -60,6 +62,63 @@ def _needs_chrome_remote_debugging_prompt(msg: str | None) -> bool: ) +def _local_debugging_setup_lines(system: str | None = None) -> list[str]: + """Return concise manual setup guidance for local Chromium-based debugging.""" + import platform + + system = system or platform.system() + lines = [ + "Do not look for webSocketDebuggerUrl in chrome://inspect; that page does not reliably show it.", + "Close the matching Chromium-based browser if it is already open, then run one command below.", + ] + if system == "Windows": + lines.extend( + [ + "On Windows PowerShell:", + '& "$env:ProgramFiles\\Google\\Chrome\\Application\\chrome.exe" --remote-debugging-port=9222 --user-data-dir="$env:USERPROFILE\\.flocks\\chrome-debug-profile"', + '& "${env:ProgramFiles(x86)}\\Microsoft\\Edge\\Application\\msedge.exe" --remote-debugging-port=9222 --user-data-dir="$env:USERPROFILE\\.flocks\\edge-debug-profile"', + 'chromium.exe --remote-debugging-port=9222 --user-data-dir="$env:USERPROFILE\\.flocks\\chromium-debug-profile"', + '& "$env:ProgramFiles\\BraveSoftware\\Brave-Browser\\Application\\brave.exe" --remote-debugging-port=9222 --user-data-dir="$env:USERPROFILE\\.flocks\\brave-debug-profile"', + "If your browser is installed elsewhere, replace the executable path.", + ] + ) + elif system == "Darwin": + lines.extend( + [ + "On macOS:", + "/Applications/Google\\ Chrome.app/Contents/MacOS/Google\\ Chrome --remote-debugging-port=9222 --user-data-dir=\"$HOME/.flocks/chrome-debug-profile\"", + "/Applications/Microsoft\\ Edge.app/Contents/MacOS/Microsoft\\ Edge --remote-debugging-port=9222 --user-data-dir=\"$HOME/.flocks/edge-debug-profile\"", + "/Applications/Chromium.app/Contents/MacOS/Chromium --remote-debugging-port=9222 --user-data-dir=\"$HOME/.flocks/chromium-debug-profile\"", + "/Applications/Brave\\ Browser.app/Contents/MacOS/Brave\\ Browser --remote-debugging-port=9222 --user-data-dir=\"$HOME/.flocks/brave-debug-profile\"", + ] + ) + else: + lines.extend( + [ + "On Linux:", + "google-chrome --remote-debugging-port=9222 --user-data-dir=\"$HOME/.flocks/chrome-debug-profile\"", + "microsoft-edge --remote-debugging-port=9222 --user-data-dir=\"$HOME/.flocks/edge-debug-profile\"", + "chromium --remote-debugging-port=9222 --user-data-dir=\"$HOME/.flocks/chromium-debug-profile\"", + "brave-browser --remote-debugging-port=9222 --user-data-dir=\"$HOME/.flocks/brave-debug-profile\"", + ] + ) + lines.extend( + [ + "Then verify http://127.0.0.1:9222/json/version; Flocks reads the WebSocket URL from that endpoint.", + "After the endpoint responds, rerun `flocks browser --setup`.", + ] + ) + return lines + + +def _print_local_debugging_setup(out=None) -> None: + import sys + + out = out or sys.stderr + for line in _local_debugging_setup_lines(): + print(line, file=out) + + def _configured_cdp_endpoint(env: dict | None = None) -> tuple[str | None, str | None]: """Return the explicit CDP endpoint env var name and value, if configured.""" merged_env = {**os.environ, **(env or {})} @@ -216,7 +275,7 @@ def _doctor_short_text(value, limit: int | None = None) -> str: def ensure_daemon( - wait: float = 60.0, name: str | None = None, env: dict | None = None, _open_inspect: bool = True + wait: float = 60.0, name: str | None = None, env: dict | None = None, _show_debugging_guidance: bool = True ) -> None: """Ensure a healthy daemon is running, restarting stale sessions when needed.""" effective_name = ipc.runtime_paths(name or NAME).name @@ -269,17 +328,10 @@ def spawn_daemon(): msg = _log_tail(effective_name) or "" if local and attempt == 0 and _needs_chrome_remote_debugging_prompt(msg): restart_daemon(effective_name) - if not _open_inspect: - raise RuntimeError( - msg or f"daemon {effective_name} didn't come up -- check {ipc.log_path(effective_name)}" - ) - _open_browser_inspect() - print( - f"{BROWSER_LABEL}: click Allow on your browser's inspect page " - "(for example chrome://inspect or edge://inspect), and tick the checkbox if shown", - file=sys.stderr, - ) - continue + if _show_debugging_guidance: + print(f"{BROWSER_LABEL}: local Chromium-based browser remote debugging is not reachable.", file=sys.stderr) + _print_local_debugging_setup(sys.stderr) + raise RuntimeError(msg or f"daemon {effective_name} didn't come up -- check {ipc.log_path(effective_name)}") raise RuntimeError(msg or f"daemon {effective_name} didn't come up -- check {ipc.log_path(effective_name)}") @@ -376,49 +428,15 @@ def _chrome_running() -> bool: try: if system == "Windows": output = subprocess.check_output(["tasklist"], timeout=5) - names = ("chrome.exe", "chromium.exe", "msedge.exe") + names = ("chrome.exe", "chromium.exe", "msedge.exe", "brave.exe") else: output = subprocess.check_output(["ps", "-A", "-o", "comm="], text=True, timeout=5) - names = ("Google Chrome", "chrome", "chromium", "Microsoft Edge", "msedge") + names = ("Google Chrome", "chrome", "chromium", "Microsoft Edge", "msedge", "Brave Browser", "brave") return _output_contains_process_names(output, names) except Exception: return False -def _open_browser_inspect() -> None: - import platform - import subprocess - import webbrowser - - inspect_targets = [ - ("Google Chrome", "chrome://inspect/#remote-debugging"), - ("Microsoft Edge", "edge://inspect/#remote-debugging"), - ] - if platform.system() == "Darwin": - for app_name, url in inspect_targets: - try: - subprocess.run( - [ - "osascript", - "-e", - f'tell application "{app_name}" to activate', - "-e", - f'tell application "{app_name}" to open location "{url}"', - ], - timeout=5, - check=False, - ) - return - except Exception: - continue - for _app_name, url in inspect_targets: - try: - if webbrowser.open(url, new=2): - return - except Exception: - continue - - def run_setup() -> int: """Interactively attach to the running browser.""" import sys @@ -439,28 +457,28 @@ def run_setup() -> int: print("daemon already running but browser connection is stale; restarting.") restart_daemon() if not endpoint_name and not _chrome_running(): - print("no Chrome/Chromium/Edge process detected. please start your browser and rerun `flocks browser --setup`.") + print("no Chrome/Chromium/Edge/Brave process detected.") + print("start a Chromium-based browser with remote debugging, then rerun `flocks browser --setup`.") + _print_local_debugging_setup(sys.stdout) return 1 try: - ensure_daemon(wait=_SETUP_ATTACH_WAIT, _open_inspect=False) + ensure_daemon(wait=_SETUP_ATTACH_WAIT, _show_debugging_guidance=False) print("daemon is up.") return 0 except RuntimeError as error: first_err = str(error) - needs_inspect = _is_local_chrome_mode() and _needs_chrome_remote_debugging_prompt(first_err) - if needs_inspect: - print("browser remote debugging is not enabled on the current profile.") - print("opening your browser's inspect page -- in the tab that opens:") - print(" 1. if the browser shows the profile picker, pick your normal profile;") - print(" 2. tick 'Discover network targets' and click Allow if prompted.") - _open_browser_inspect() + needs_manual_debugging_setup = _is_local_chrome_mode() and _needs_chrome_remote_debugging_prompt(first_err) + if needs_manual_debugging_setup: + print("Chromium-based browser remote debugging is not reachable for the current profile.") + _print_local_debugging_setup(sys.stdout) + return 1 else: print(f"attach failed: {first_err}") print("retrying once (the browser may still be starting up)...") try: - ensure_daemon(wait=_SETUP_RETRY_WAIT, _open_inspect=False) + ensure_daemon(wait=_SETUP_RETRY_WAIT, _show_debugging_guidance=False) print("daemon is up.") return 0 except RuntimeError as error: @@ -503,7 +521,10 @@ def row(label: str, ok: bool, detail: str = "") -> None: elif target_available: next_action = "setup; run `flocks browser --setup`" else: - next_action = "start Chrome/Chromium/Edge or provide BU_CDP_URL/BU_CDP_WS, then run `flocks browser --setup`" + next_action = ( + "start Chrome/Chromium/Edge/Brave or provide BU_CDP_URL/BU_CDP_WS, " + "then run `flocks browser --setup`" + ) print(f"{BROWSER_LABEL} doctor") print(f" platform {platform.system()} {platform.release()}") @@ -522,14 +543,14 @@ def row(label: str, ok: bool, detail: str = "") -> None: row( "browser running", browser_running, - "" if browser_running else "start Chrome, Chromium, or Edge and rerun `flocks browser --setup`", + "" if browser_running else "start Chrome, Chromium, Edge, or Brave and rerun `flocks browser --setup`", ) row( "daemon alive", daemon, "" if daemon - else "not running; run `flocks browser --setup` to attach; if setup reports remote debugging is disabled, follow its inspect-page prompt", + else "not running; run `flocks browser --setup` to attach; if setup reports remote debugging is not reachable, follow its debug-browser instructions", ) row("active browser connections", bool(connections), str(len(connections))) for conn in connections: diff --git a/flocks/browser/daemon.py b/flocks/browser/daemon.py index c735a6d50..8a60a5469 100644 --- a/flocks/browser/daemon.py +++ b/flocks/browser/daemon.py @@ -23,6 +23,12 @@ from .utils import load_env_file +FLOCKS_DEBUG_PROFILE_NAMES = ( + "chrome-debug-profile", + "edge-debug-profile", + "chromium-debug-profile", + "brave-debug-profile", +) AGENT_WORKSPACE = Path(os.environ.get("BH_AGENT_WORKSPACE", DEFAULT_AGENT_WORKSPACE)).expanduser() BUF = 500 INTERNAL = INTERNAL_URL_PREFIXES @@ -53,9 +59,11 @@ def profile_dirs( system = system or platform.system() home = home or Path.home() environ = os.environ if environ is None else environ + flocks_debug_profiles = [home / ".flocks" / name for name in FLOCKS_DEBUG_PROFILE_NAMES] if system == "Darwin": support = home / "Library/Application Support" return [ + *flocks_debug_profiles, support / "Google/Chrome", support / "Google/Chrome Canary", support / "Comet", @@ -70,6 +78,7 @@ def profile_dirs( if system == "Windows": local = Path(environ.get("LOCALAPPDATA", str(home / "AppData/Local"))).expanduser() return [ + *flocks_debug_profiles, local / "Google/Chrome/User Data", local / "Google/Chrome Beta/User Data", local / "Google/Chrome Dev/User Data", @@ -84,6 +93,7 @@ def profile_dirs( local / "BraveSoftware/Brave-Browser-Nightly/User Data", ] return [ + *flocks_debug_profiles, home / ".config/google-chrome", home / ".config/google-chrome-beta", home / ".config/google-chrome-unstable", @@ -125,7 +135,9 @@ def _local_cdp_ws_url(port: int, devtools_path: str | None = None) -> str: if not isinstance(payload, dict): raise ValueError("invalid CDP version response") websocket_url = payload.get("webSocketDebuggerUrl") - if not isinstance(websocket_url, str): + if not isinstance(websocket_url, str) and devtools_path and devtools_path.startswith("/"): + websocket_url = f"ws://127.0.0.1:{port}{devtools_path}" + elif not isinstance(websocket_url, str): raise ValueError("invalid webSocketDebuggerUrl") except urllib.error.HTTPError: if not devtools_path or not devtools_path.startswith("/"): @@ -198,15 +210,14 @@ def get_ws_url() -> str: if profile_errors: details = "; ".join(dict.fromkeys(profile_errors)) raise RuntimeError( - "The browser's remote-debugging page is open, but DevTools is not live yet for any detected profile: " - f"{details} — if the browser opened a profile picker, choose your normal profile first, then tick the " - "checkbox and click Allow if shown" + "Chromium-based browser remote debugging is not reachable for any detected profile: " + f"{details} — for manual setup, start the browser with --remote-debugging-port and a non-default " + "--user-data-dir, then verify http://127.0.0.1:9222/json/version" ) raise RuntimeError( "DevToolsActivePort not found in " - f"{[str(path) for path in profiles]} — enable your browser's remote-debugging page " - "(for example chrome://inspect/#remote-debugging or edge://inspect/#remote-debugging), " - "or set BU_CDP_WS for a remote browser" + f"{[str(path) for path in profiles]} — start a Chromium-based browser with --remote-debugging-port and a non-default " + "--user-data-dir, or set BU_CDP_WS for a remote browser" ) @@ -280,7 +291,9 @@ async def start(self) -> None: f"{hint}" ) from error raise RuntimeError( - f"CDP WS handshake failed: {error} -- click Allow in your browser if prompted, then retry" + f"CDP WS handshake failed: {error} -- restart your Chromium-based browser with " + "--remote-debugging-port and a non-default --user-data-dir, then verify " + "http://127.0.0.1:9222/json/version" ) await self.attach_first_page() orig = self.cdp._event_registry.handle_event diff --git a/flocks/channel/base.py b/flocks/channel/base.py index 67e96148c..ffd79dc4e 100644 --- a/flocks/channel/base.py +++ b/flocks/channel/base.py @@ -39,6 +39,7 @@ class ChannelCapabilities: reactions: bool = False edit: bool = False rich_text: bool = False + self_managed_connection: bool = False @dataclass @@ -86,6 +87,10 @@ class DeliveryResult: retryable: bool = False +class NonRetryableChannelError(RuntimeError): + """Raised when reconnecting cannot fix a channel startup failure.""" + + @dataclass class ChannelStatus: """Runtime health status of a channel connection.""" diff --git a/flocks/channel/builtin/slack/__init__.py b/flocks/channel/builtin/slack/__init__.py new file mode 100644 index 000000000..bc4986df2 --- /dev/null +++ b/flocks/channel/builtin/slack/__init__.py @@ -0,0 +1,5 @@ +"""Slack channel integration.""" + +from flocks.channel.builtin.slack.channel import SlackChannel + +__all__ = ["SlackChannel"] diff --git a/flocks/channel/builtin/slack/channel.py b/flocks/channel/builtin/slack/channel.py new file mode 100644 index 000000000..512ff7fb1 --- /dev/null +++ b/flocks/channel/builtin/slack/channel.py @@ -0,0 +1,439 @@ +"""Slack ChannelPlugin using Socket Mode.""" + +from __future__ import annotations + +import asyncio +from collections import OrderedDict +from typing import Any, Awaitable, Callable, Optional + +from flocks.channel.base import ( + ChannelCapabilities, + ChannelMeta, + ChannelPlugin, + ChatType, + DeliveryResult, + InboundMessage, + NonRetryableChannelError, + OutboundContext, +) +from flocks.storage.storage import Storage +from flocks.utils.log import Log + +from .config import ( + normalize_target, + resolve_app_token, + resolve_bot_token, + resolve_home_channel, + should_reply_broadcast, + should_reply_in_thread, +) +from .format import markdown_to_slack_mrkdwn +from .inbound import build_inbound_message, slack_thread_cache_key + +try: + from slack_bolt.async_app import AsyncApp + from slack_bolt.adapter.socket_mode.async_handler import AsyncSocketModeHandler + from slack_sdk.web.async_client import AsyncWebClient + + SLACK_AVAILABLE = True +except ImportError: # pragma: no cover - exercised by validate_config tests + AsyncApp = None # type: ignore[assignment] + AsyncSocketModeHandler = None # type: ignore[assignment] + AsyncWebClient = None # type: ignore[assignment] + SLACK_AVAILABLE = False + + +log = Log.create(service="channel.slack") + +_KNOWN_THREADS_STORAGE_KEY = "channel:slack:known_threads" + + +_PERMANENT_SLACK_ERROR_CODES = { + "invalid_auth", + "not_authed", + "no_auth", + "token_revoked", + "account_inactive", + "not_allowed_token_type", + "missing_scope", +} + + +def _is_token_reference(value: str) -> bool: + return value.startswith("{secret:") or value.startswith("{env:") + + +def _slack_response_get(response: Any, key: str, default: str = "") -> str: + getter = getattr(response, "get", None) + if callable(getter): + return str(getter(key, default) or default) + data = getattr(response, "data", None) + if isinstance(data, dict): + return str(data.get(key) or default) + return default + + +def _slack_error_code(exc: Exception) -> str: + response = getattr(exc, "response", None) + data = getattr(response, "data", None) + if isinstance(data, dict): + return str(data.get("error") or "").strip() + if isinstance(response, dict): + return str(response.get("error") or "").strip() + + message = str(exc).lower().replace("-", "_") + for code in _PERMANENT_SLACK_ERROR_CODES: + if code in message: + return code + if "invalid app token" in message: + return "invalid_auth" + return "" + + +def _is_permanent_slack_error(exc: Exception) -> bool: + return _slack_error_code(exc) in _PERMANENT_SLACK_ERROR_CODES + + +def _format_slack_start_error(exc: Exception, *, phase: str) -> str: + code = _slack_error_code(exc) + detail = f"Slack error={code}" if code else str(exc) + + if phase == "bot_auth": + if code == "missing_scope": + return ( + "Slack Bot Token 权限不足：请在 Slack 开发者管理后台的 OAuth & Permissions " + "确认 Bot Token 具备 Manifest 中的 bot scopes，重新安装 App 后复制 " + "Bot User OAuth Token（xoxb- 开头）。" + ) + return ( + "Slack Bot Token 验证失败：请在 Slack 开发者管理后台的 OAuth & Permissions " + "复制 Bot User OAuth Token（xoxb- 开头），确认 App 已安装到 workspace 后重新保存。" + f" 原始错误：{detail}" + ) + + if code == "missing_scope": + return ( + "Slack App Token 权限不足：请在 Basic Information > App-Level Tokens " + "创建或更新具备 connections:write scope 的 App-Level Token（xapp- 开头）。" + ) + if code in {"invalid_auth", "not_authed", "no_auth", "token_revoked", "account_inactive", "not_allowed_token_type"}: + return ( + "Slack App Token 验证失败：请在 Basic Information > App-Level Tokens " + "复制 App-Level Token（xapp- 开头），并确认它包含 connections:write scope。" + f" 原始错误：{detail}" + ) + return f"Slack Socket Mode 连接失败：{detail}" + + +class SlackChannel(ChannelPlugin): + """Slack bot channel — Socket Mode inbound and Web API outbound.""" + + MAX_MESSAGE_LENGTH = 39000 + + def __init__(self) -> None: + super().__init__() + self._app: Any = None + self._handler: Any = None + self._socket_task: Optional[asyncio.Task] = None + self._bot_user_id: Optional[str] = None + self._known_thread_ids: OrderedDict[str, None] = OrderedDict() + self._known_thread_ids_max = 5000 + + def meta(self) -> ChannelMeta: + return ChannelMeta(id="slack", label="Slack", aliases=["sl"], order=50) + + def capabilities(self) -> ChannelCapabilities: + return ChannelCapabilities( + chat_types=[ChatType.DIRECT, ChatType.GROUP, ChatType.CHANNEL], + media=True, + threads=True, + reactions=False, + edit=False, + rich_text=True, + self_managed_connection=True, + ) + + def validate_config(self, config: dict) -> Optional[str]: + if not SLACK_AVAILABLE: + return "Missing Python dependency: slack-bolt" + bot_token = resolve_bot_token(config) + app_token = resolve_app_token(config) + if not bot_token: + return "Missing required config: botToken" + if not app_token: + return "Missing required config: appToken" + if not _is_token_reference(bot_token) and not bot_token.startswith("xoxb-"): + return ( + "Slack Bot Token 必须是 Bot User OAuth Token（xoxb- 开头），" + "不能使用 User Token（xoxp-）或其他 token。" + ) + if not _is_token_reference(app_token) and not app_token.startswith("xapp-"): + return "Slack App Token 必须是 App-Level Token（xapp- 开头）。" + return None + + @property + def text_chunk_limit(self) -> int: + return self.MAX_MESSAGE_LENGTH + + @property + def rate_limit(self) -> tuple[float, int]: + return (1.0, 5) + + def normalize_target(self, raw: str) -> Optional[str]: + target = normalize_target(raw, fallback=resolve_home_channel(self._config)) + return target or None + + def target_hint(self) -> str: + return " (C..., G..., or D...)" + + def format_message(self, text: str, format_hint: str = "markdown") -> str: + if format_hint == "plain": + return text + return markdown_to_slack_mrkdwn(text) + + async def start( + self, + config: dict, + on_message: Callable[[InboundMessage], Awaitable[None]], + abort_event: Optional[asyncio.Event] = None, + ) -> None: + self._config = config + self._on_message = on_message + + error = self.validate_config(config) + if error: + self.mark_disconnected(error) + log.error("slack.start.invalid_config", {"error": error}) + raise NonRetryableChannelError(error) + + bot_token = resolve_bot_token(config) + app_token = resolve_app_token(config) + assert AsyncApp is not None + assert AsyncSocketModeHandler is not None + + self._app = AsyncApp(token=bot_token) + abort = abort_event or asyncio.Event() + try: + try: + auth = await self._app.client.auth_test() + except Exception as exc: + message = _format_slack_start_error(exc, phase="bot_auth") + self.mark_disconnected(message) + if _is_permanent_slack_error(exc): + raise NonRetryableChannelError(message) from exc + raise RuntimeError(message) from exc + + self._bot_user_id = _slack_response_get(auth, "user_id") + if not _slack_response_get(auth, "bot_id"): + message = ( + "Slack Bot Token 验证失败：当前 token 认证为用户而不是 Bot。" + "请在 Slack 开发者管理后台 OAuth & Permissions 中复制 " + "Bot User OAuth Token（xoxb- 开头），不要使用 xoxp- User Token。" + ) + self.mark_disconnected(message) + raise NonRetryableChannelError(message) + team_id = _slack_response_get(auth, "team_id", "default") + await self._load_known_threads() + + self._register_handlers() + await self._connect_socket_mode(app_token) + log.info("slack.connected", {"team_id": team_id, "bot_user_id": self._bot_user_id}) + await abort.wait() + except asyncio.CancelledError: + raise + except NonRetryableChannelError: + raise + except Exception as exc: + message = str(exc) + if self.status.connected or self.status.last_error != message: + self.mark_disconnected(message) + log.warning("slack.socket.stopped", {"error": message}) + raise + finally: + await self.stop() + + async def _connect_socket_mode(self, app_token: str) -> None: + """Connect Socket Mode once and mark connected only after success.""" + assert AsyncSocketModeHandler is not None + await self._verify_app_token(app_token) + self._handler = AsyncSocketModeHandler(self._app, app_token) + timeout_seconds = float(self._config.get("socketConnectTimeoutSeconds") or 15.0) + try: + await asyncio.wait_for( + self._handler.connect_async(), + timeout=max(timeout_seconds, 1.0), + ) + except asyncio.TimeoutError as exc: + message = f"Slack Socket Mode connection timed out after {timeout_seconds:.1f}s" + self.mark_disconnected(message) + raise TimeoutError(message) from exc + except Exception as exc: + message = _format_slack_start_error(exc, phase="socket_mode") + self.mark_disconnected(message) + if _is_permanent_slack_error(exc): + raise NonRetryableChannelError(message) from exc + raise RuntimeError(message) from exc + self.mark_connected() + + async def _verify_app_token(self, app_token: str) -> None: + """Validate the app-level token once before Socket Mode's retry loop.""" + assert AsyncWebClient is not None + client = AsyncWebClient(token=app_token) + try: + await client.apps_connections_open(app_token=app_token) + except Exception as exc: + message = _format_slack_start_error(exc, phase="socket_mode") + self.mark_disconnected(message) + if _is_permanent_slack_error(exc): + raise NonRetryableChannelError(message) from exc + raise RuntimeError(message) from exc + + async def stop(self) -> None: + handler = self._handler + self._handler = None + if handler is not None: + try: + await handler.close_async() + except Exception as exc: + log.warning("slack.close.failed", {"error": str(exc)}) + + task = self._socket_task + self._socket_task = None + if task is not None and not task.done(): + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + except Exception: + pass + + self.mark_disconnected() + + async def send_text(self, ctx: OutboundContext) -> DeliveryResult: + if not self._app: + return DeliveryResult( + channel_id="slack", + message_id="", + success=False, + error="Slack channel is not connected", + retryable=True, + ) + + target = normalize_target(ctx.to, fallback=resolve_home_channel(self._config)) + if not target: + return DeliveryResult( + channel_id="slack", + message_id="", + success=False, + error="Slack send requires 'to' (channel ID C..., private channel G..., or DM D...)", + ) + + try: + payload: dict[str, Any] = { + "channel": target, + "text": ctx.text, + "mrkdwn": ctx.format_hint != "plain", + } + thread_ts = self._resolve_thread_ts(ctx) + if thread_ts: + payload["thread_ts"] = thread_ts + if should_reply_broadcast(self._config): + payload["reply_broadcast"] = True + + result = await self._app.client.chat_postMessage(**payload) + message_id = str(result.get("ts") or "") + if thread_ts: + self._remember_thread(ctx.account_id or "default", target, thread_ts) + elif message_id: + self._remember_thread(ctx.account_id or "default", target, message_id) + await self._persist_known_threads() + self.record_message() + return DeliveryResult( + channel_id="slack", + message_id=message_id, + chat_id=target, + success=True, + ) + except Exception as exc: + message = str(exc) + retryable = any(term in message.lower() for term in ("timeout", "rate", "temporarily", "connection")) + return DeliveryResult( + channel_id="slack", + message_id="", + success=False, + error=f"Slack send failed: {message}", + retryable=retryable, + ) + + def _register_handlers(self) -> None: + if not self._app: + return + + @self._app.event("message") + async def handle_message_event(event, body, say): + await self._handle_slack_event(event, body=body) + + @self._app.event("app_mention") + async def handle_app_mention_event(event, body, say): + await self._handle_slack_event(event, body=body) + + async def _handle_slack_event( + self, + event: dict[str, Any], + *, + body: Optional[dict[str, Any]] = None, + ) -> None: + if not self._on_message: + return + inbound = build_inbound_message( + event, + bot_user_id=self._bot_user_id, + config=self._config, + known_thread_ids=set(self._known_thread_ids.keys()), + body=body, + ) + if inbound is None: + return + self.record_message() + await self._on_message(inbound) + + def _resolve_thread_ts(self, ctx: OutboundContext) -> Optional[str]: + if ctx.thread_id: + return ctx.thread_id + if should_reply_in_thread(self._config): + return ctx.reply_to_id + return None + + def _remember_thread(self, account_id: str, channel_id: str, thread_ts: str) -> None: + if not thread_ts: + return + thread_key = slack_thread_cache_key(account_id, channel_id, thread_ts) + self._known_thread_ids[thread_key] = None + self._known_thread_ids.move_to_end(thread_key) + if len(self._known_thread_ids) > self._known_thread_ids_max: + while len(self._known_thread_ids) > self._known_thread_ids_max // 2: + self._known_thread_ids.popitem(last=False) + + async def _load_known_threads(self) -> None: + try: + stored = await Storage.get(_KNOWN_THREADS_STORAGE_KEY) + except Exception as exc: + log.warning("slack.known_threads.load_failed", {"error": str(exc)}) + return + if not isinstance(stored, list): + return + for thread_ts in stored[-self._known_thread_ids_max:]: + if isinstance(thread_ts, str) and thread_ts: + self._known_thread_ids[thread_ts] = None + while len(self._known_thread_ids) > self._known_thread_ids_max: + self._known_thread_ids.popitem(last=False) + + async def _persist_known_threads(self) -> None: + try: + await Storage.set( + _KNOWN_THREADS_STORAGE_KEY, + list(self._known_thread_ids.keys()), + ) + except Exception as exc: + log.warning("slack.known_threads.persist_failed", {"error": str(exc)}) diff --git a/flocks/channel/builtin/slack/config.py b/flocks/channel/builtin/slack/config.py new file mode 100644 index 000000000..7029d752d --- /dev/null +++ b/flocks/channel/builtin/slack/config.py @@ -0,0 +1,75 @@ +"""Configuration helpers for the Slack channel.""" + +from __future__ import annotations + +import re +from typing import Any, Optional + + +_TARGET_PREFIX_RE = re.compile(r"^(?:slack|channel|dm|user):", re.IGNORECASE) + + +def coerce_str(value: Any) -> str: + if value is None: + return "" + return str(value).strip() + + +def coerce_bool(value: Any, default: bool = False) -> bool: + if value is None: + return default + if isinstance(value, bool): + return value + if isinstance(value, (int, float)): + return bool(value) + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"1", "true", "yes", "y", "on"}: + return True + if normalized in {"0", "false", "no", "n", "off"}: + return False + return default + + +def strip_target_prefix(raw: str) -> str: + value = raw.strip() + while value: + next_value = _TARGET_PREFIX_RE.sub("", value, count=1).strip() + if next_value == value: + return value + value = next_value + return value + + +def resolve_bot_token(config: dict[str, Any]) -> str: + return coerce_str(config.get("botToken") or config.get("slackBotToken")) + + +def resolve_app_token(config: dict[str, Any]) -> str: + return coerce_str(config.get("appToken") or config.get("slackAppToken")) + + +def resolve_home_channel(config: dict[str, Any]) -> str: + return coerce_str(config.get("homeChannel") or config.get("home_channel")) + + +def should_reply_in_thread(config: dict[str, Any]) -> bool: + return coerce_bool(config.get("replyInThread"), True) + + +def should_reply_broadcast(config: dict[str, Any]) -> bool: + return coerce_bool(config.get("replyBroadcast"), False) + + +def allow_bot_messages(config: dict[str, Any]) -> str: + value = coerce_str(config.get("allowBots")).lower() + if value in {"none", "mentions", "all"}: + return value + return "none" + + +def normalize_target(raw: str, *, fallback: Optional[str] = None) -> str: + target = strip_target_prefix(raw or "") + if target: + return target + return fallback or "" diff --git a/flocks/channel/builtin/slack/format.py b/flocks/channel/builtin/slack/format.py new file mode 100644 index 000000000..f92d97615 --- /dev/null +++ b/flocks/channel/builtin/slack/format.py @@ -0,0 +1,74 @@ +"""Markdown to Slack mrkdwn helpers.""" + +from __future__ import annotations + +import re +from html import escape +from urllib.parse import quote + + +_LINK_RE = re.compile(r"\[([^\]]+)\]\((https?://[^)\s]+)\)") +_SAFE_ENTITY_RE = re.compile(r"&(amp|lt|gt);") + + +def _escape_slack_text(value: str) -> str: + placeholders: list[str] = [] + + def protect_entity(match: re.Match[str]) -> str: + placeholders.append(match.group(0)) + return f"\x02SLACK_ENTITY_{len(placeholders) - 1}\x02" + + escaped = _SAFE_ENTITY_RE.sub(protect_entity, value) + escaped = escape(escaped, quote=False) + for index, entity in enumerate(placeholders): + escaped = escaped.replace(f"\x02SLACK_ENTITY_{index}\x02", entity) + return escaped + + +def _escape_slack_link_url(value: str) -> str: + return quote(value, safe=":/?#[]@!$&'()*+,;=%") + + +def markdown_to_slack_mrkdwn(text: str) -> str: + """Convert common Markdown to Slack mrkdwn without touching code fences.""" + if not text: + return "" + + code_placeholders: list[str] = [] + + def protect(match: re.Match[str]) -> str: + code_placeholders.append(match.group(0)) + return f"\x01SLACK_CODE_{len(code_placeholders) - 1}\x01" + + link_placeholders: list[str] = [] + + def convert_link(match: re.Match[str]) -> str: + label = _escape_slack_text(match.group(1)) + url = _escape_slack_link_url(match.group(2)) + link_placeholders.append(f"<{url}|{label}>") + return f"\x01SLACK_LINK_{len(link_placeholders) - 1}\x01" + + result = re.sub(r"```.*?```", protect, text, flags=re.DOTALL) + result = re.sub(r"`[^`\n]+`", protect, result) + + result = _LINK_RE.sub(convert_link, result) + result = re.sub(r"\*\*(.+?)\*\*", r"*\1*", result, flags=re.DOTALL) + result = re.sub(r"__([^_\n]+)__", r"*\1*", result) + result = re.sub(r"~~(.+?)~~", r"~\1~", result, flags=re.DOTALL) + result = re.sub(r"^#{1,6}\s+(.+)$", r"*\1*", result, flags=re.MULTILINE) + result = _escape_slack_text(result) + + for index, fragment in enumerate(link_placeholders): + result = result.replace(f"\x01SLACK_LINK_{index}\x01", fragment) + for index, fragment in enumerate(code_placeholders): + result = result.replace(f"\x01SLACK_CODE_{index}\x01", fragment) + + return result + + +def strip_bot_mention(text: str, bot_user_id: str) -> str: + if not bot_user_id: + return text.strip() + pattern = re.compile(rf"<@{re.escape(bot_user_id)}(?:\|[^>]+)?>") + cleaned = pattern.sub(" ", text) + return re.sub(r"\s+", " ", cleaned).strip() diff --git a/flocks/channel/builtin/slack/inbound.py b/flocks/channel/builtin/slack/inbound.py new file mode 100644 index 000000000..2587e2012 --- /dev/null +++ b/flocks/channel/builtin/slack/inbound.py @@ -0,0 +1,358 @@ +"""Inbound Slack event parsing.""" + +from __future__ import annotations + +import json +from typing import Any, Optional + +from flocks.channel.base import ChatType, InboundMessage + +from .config import allow_bot_messages +from .format import strip_bot_mention + + +def slack_thread_cache_key(account_id: str, channel_id: str, thread_ts: str) -> str: + return f"{account_id or 'default'}:{channel_id}:{thread_ts}" + + +def _join_parts(parts: list[str]) -> str: + return "\n\n".join(part for part in (p.strip() for p in parts) if part) + + +def _dedupe_part(part: str, existing_text: str) -> str: + candidate = part.strip() + existing = existing_text.strip() + if not candidate: + return "" + if not existing: + return candidate + if candidate == existing or candidate in existing: + return "" + if candidate.startswith(existing): + return candidate[len(existing):].strip() + return candidate + + +def _slack_text_object(value: Any) -> str: + if isinstance(value, str): + return value.strip() + if isinstance(value, dict): + text = value.get("text") + if isinstance(text, str): + return text.strip() + return "" + + +def _rich_text_element_text(element: Any) -> str: + if not isinstance(element, dict): + return "" + + element_type = str(element.get("type") or "") + if element_type == "text": + return str(element.get("text") or "") + if element_type == "user": + return f"<@{element.get('user_id')}>" + if element_type == "channel": + return f"<#{element.get('channel_id')}>" + if element_type == "usergroup": + return f"" + if element_type == "broadcast": + return f"" + if element_type == "emoji": + return f":{element.get('name')}:" + if element_type == "link": + text = str(element.get("text") or "").strip() + url = str(element.get("url") or "").strip() + return f"{text} ({url})" if text and url else text or url + if element_type == "date": + return str(element.get("fallback") or element.get("timestamp") or "") + + children = element.get("elements") + if isinstance(children, list): + child_text = "".join(_rich_text_element_text(child) for child in children).strip() + if element_type == "rich_text_preformatted" and child_text: + return f"```\n{child_text}\n```" + if element_type == "rich_text_quote" and child_text: + return "\n".join(f"> {line}" for line in child_text.splitlines()) + return child_text + return "" + + +def _extract_blocks_text(blocks: Any) -> str: + if not isinstance(blocks, list): + return "" + + parts: list[str] = [] + for block in blocks: + if not isinstance(block, dict): + continue + + block_type = str(block.get("type") or "") + if block_type == "rich_text": + elements = block.get("elements") + if isinstance(elements, list): + parts.append(_join_parts([ + _rich_text_element_text(element) + for element in elements + ])) + continue + + for key in ("text", "title"): + value = _slack_text_object(block.get(key)) + if value: + parts.append(value) + + fields = block.get("fields") + if isinstance(fields, list): + parts.extend(_slack_text_object(field) for field in fields) + + elements = block.get("elements") + if isinstance(elements, list): + parts.extend(_slack_text_object(element) for element in elements) + + return _join_parts(parts) + + +def _extract_attachments_text(attachments: Any) -> str: + if not isinstance(attachments, list): + return "" + + parts: list[str] = [] + for attachment in attachments: + if not isinstance(attachment, dict): + continue + + for key in ("pretext", "title", "text", "fallback"): + value = _slack_text_object(attachment.get(key)) + if value: + parts.append(value) + + fields = attachment.get("fields") + if isinstance(fields, list): + for field in fields: + if not isinstance(field, dict): + continue + title = _slack_text_object(field.get("title")) + value = _slack_text_object(field.get("value")) + if title and value: + parts.append(f"{title}: {value}") + elif value: + parts.append(value) + + blocks = _extract_blocks_text(attachment.get("blocks")) + if blocks: + parts.append(blocks) + + return _join_parts(parts) + + +def _file_display_name(item: dict[str, Any]) -> str: + return str(item.get("name") or item.get("title") or item.get("id") or "file") + + +def _extract_files_text(files: Any) -> str: + if not isinstance(files, list) or not files: + return "" + + names = [ + _file_display_name(item) + for item in files + if isinstance(item, dict) + ] + if not names: + return "" + + first = next((item for item in files if isinstance(item, dict)), {}) + mimetype = str(first.get("mimetype") or "").lower() + label = "图片消息" if mimetype.startswith("image/") else "文件消息" + return f"[{label}: " + ", ".join(names) + "]" + + +def _first_file_media(files: Any) -> tuple[Optional[str], Optional[str]]: + if not isinstance(files, list): + return None, None + for item in files: + if not isinstance(item, dict): + continue + url = str(item.get("url_private_download") or item.get("url_private") or "").strip() + file_id = str(item.get("id") or "").strip() + if not url and file_id: + url = f"slack://file/{file_id}" + if not url: + continue + mimetype = str(item.get("mimetype") or item.get("filetype") or "").strip() or None + return url, mimetype + return None, None + + +def extract_text(event: dict[str, Any]) -> str: + parts: list[str] = [] + combined = "" + text = event.get("text") + if isinstance(text, str) and text.strip(): + combined = text.strip() + parts.append(combined) + + blocks = _dedupe_part(_extract_blocks_text(event.get("blocks")), combined) + if blocks: + parts.append(blocks) + combined = _join_parts(parts) + + attachments = _dedupe_part(_extract_attachments_text(event.get("attachments")), combined) + if attachments: + parts.append(attachments) + combined = _join_parts(parts) + + files = _dedupe_part(_extract_files_text(event.get("files")), combined) + if files: + parts.append(files) + + rendered = _join_parts(parts) + if rendered: + return rendered + + if event.get("blocks") or event.get("attachments"): + try: + return "[Slack structured message]\n" + json.dumps( + { + "blocks": event.get("blocks"), + "attachments": event.get("attachments"), + }, + ensure_ascii=False, + separators=(",", ":"), + ) + except (TypeError, ValueError): + return "[Slack structured message]" + + return "" + + +def _extract_body_team_id(body: Optional[dict[str, Any]]) -> str: + if not isinstance(body, dict): + return "" + for key in ("team_id", "team"): + value = body.get(key) + if isinstance(value, str) and value.strip(): + return value.strip() + if isinstance(value, dict): + nested = str(value.get("id") or "").strip() + if nested: + return nested + authorizations = body.get("authorizations") + if isinstance(authorizations, list): + for authorization in authorizations: + if not isinstance(authorization, dict): + continue + team_id = str(authorization.get("team_id") or "").strip() + if team_id: + return team_id + return "" + + +def is_own_or_ignored_bot_message( + event: dict[str, Any], + *, + bot_user_id: Optional[str], + config: dict[str, Any], +) -> bool: + user_id = str(event.get("user") or event.get("bot_id") or "") + if bot_user_id and user_id == bot_user_id: + return True + + if not (event.get("bot_id") or event.get("subtype") == "bot_message"): + return False + + policy = allow_bot_messages(config) + if policy == "all": + return False + if policy == "mentions": + return not bool(bot_user_id and f"<@{bot_user_id}>" in str(event.get("text") or "")) + return True + + +def build_inbound_message( + event: dict[str, Any], + *, + bot_user_id: Optional[str], + config: dict[str, Any], + known_thread_ids: set[str], + body: Optional[dict[str, Any]] = None, +) -> Optional[InboundMessage]: + """Convert a Slack message/app_mention event into a Flocks InboundMessage.""" + subtype = event.get("subtype") + if subtype in {"message_changed", "message_deleted"}: + return None + + if is_own_or_ignored_bot_message(event, bot_user_id=bot_user_id, config=config): + return None + + channel_id = str(event.get("channel") or "") + user_id = str(event.get("user") or event.get("bot_id") or "") + ts = str(event.get("ts") or event.get("event_ts") or "") + if not channel_id or not user_id or not ts: + return None + + text = extract_text(event) + if not text: + return None + + channel_type = str(event.get("channel_type") or "").lower() + if not channel_type and channel_id.startswith("D"): + channel_type = "im" + + if channel_type == "im": + chat_type = ChatType.DIRECT + elif channel_type == "mpim": + chat_type = ChatType.GROUP + else: + # Slack public/private channels are not one-to-one chats. Treat them + # as channel surfaces while preserving group-trigger behavior in the + # dispatcher (all non-DIRECT chats use the same trigger policy). + chat_type = ChatType.CHANNEL + account_id = str( + event.get("team") + or event.get("team_id") + or _extract_body_team_id(body) + or "default" + ) + thread_ts = str(event.get("thread_ts") or "") + is_thread_reply = bool(thread_ts and thread_ts != ts) + mentioned = bool(bot_user_id and f"<@{bot_user_id}>" in text) + thread_key = slack_thread_cache_key(account_id, channel_id, thread_ts) if thread_ts else "" + continues_bot_thread = bool( + is_thread_reply + and ( + thread_key in known_thread_ids + or (bot_user_id and str(event.get("parent_user_id") or "") == bot_user_id) + ) + ) + mention_text = strip_bot_mention(text, bot_user_id or "") if mentioned else "" + + if chat_type == ChatType.DIRECT: + mentioned = False + mention_text = "" + elif continues_bot_thread: + mentioned = True + if not mention_text: + mention_text = text + + sender_name = str(event.get("username") or event.get("user_name") or "") or None + media_url, media_mime = _first_file_media(event.get("files")) + + return InboundMessage( + channel_id="slack", + account_id=account_id, + message_id=ts, + sender_id=user_id, + sender_name=sender_name, + chat_id=channel_id, + chat_type=chat_type, + text=text, + media_url=media_url, + media_mime=media_mime, + reply_to_id=str(event.get("parent_user_id") or "") or None, + thread_id=thread_ts or None, + mentioned=mentioned, + mention_text=mention_text, + raw=event, + ) diff --git a/flocks/channel/builtin/slack/inbound_media.py b/flocks/channel/builtin/slack/inbound_media.py new file mode 100644 index 000000000..03af97f07 --- /dev/null +++ b/flocks/channel/builtin/slack/inbound_media.py @@ -0,0 +1,256 @@ +"""Slack inbound media download helpers.""" + +from __future__ import annotations + +import datetime +import mimetypes +import os +import re +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Optional +from urllib.parse import unquote, urlparse + +import httpx + +from flocks.channel.base import InboundMessage +from flocks.channel.media_filename import sanitize_filename +from flocks.utils.log import Log + +from .config import resolve_bot_token + +log = Log.create(service="channel.slack.media") + +_DEFAULT_MAX_INBOUND_MEDIA_BYTES = 30 * 1024 * 1024 + + +class SlackInboundMediaTooLarge(ValueError): + """Slack 入站媒体超过允许大小。""" + + +@dataclass +class DownloadedInboundMedia: + filename: str + mime: str + url: str + source: dict + + +def _media_storage_dir(account_id: str) -> Path: + return ( + Path.home() + / ".flocks" + / "data" + / "channel_media" + / "slack" + / account_id + / datetime.date.today().isoformat() + ) + + +def _sanitize_filename(name: str) -> str: + return sanitize_filename(name) + + +def _guess_mime_from_ext(filename: str) -> Optional[str]: + _, ext = os.path.splitext(filename) + if ext: + return mimetypes.guess_type(filename)[0] + return None + + +def _filename_from_content_disposition(value: str) -> Optional[str]: + match = re.search(r'filename\*?=(?:UTF-8\'\')?"?([^";]+)"?', value, re.I) + if not match: + return None + return unquote(match.group(1).strip()) + + +def _max_size_error(max_bytes: int) -> SlackInboundMediaTooLarge: + return SlackInboundMediaTooLarge( + f"Slack inbound media too large: >{max_bytes // (1024 * 1024)}MB" + ) + + +def _iter_raw_files(msg: InboundMessage) -> list[dict[str, Any]]: + raw = msg.raw if isinstance(msg.raw, dict) else {} + files = raw.get("files") + if not isinstance(files, list): + return [] + return [item for item in files if isinstance(item, dict)] + + +def _find_file_object(msg: InboundMessage, media_url: str) -> dict[str, Any]: + parsed = urlparse(media_url) + ref_file_id = parsed.path.lstrip("/") if parsed.scheme == "slack" else "" + for item in _iter_raw_files(msg): + if ref_file_id and str(item.get("id") or "") == ref_file_id: + return item + for key in ("url_private_download", "url_private"): + if media_url and str(item.get(key) or "") == media_url: + return item + return {} + + +def _guess_filename( + msg: InboundMessage, + media_url: str, + file_obj: dict[str, Any], + cd_filename: Optional[str] = None, +) -> str: + for key in ("name", "title"): + candidate = str(file_obj.get(key) or "").strip() + if candidate: + return _sanitize_filename(candidate) + if cd_filename: + return _sanitize_filename(cd_filename) + basename = os.path.basename(urlparse(media_url).path) + if basename and "." in basename: + return _sanitize_filename(unquote(basename)) + file_id = str(file_obj.get("id") or "").strip() + suffix = file_id or msg.message_id or "unknown" + return _sanitize_filename(f"slack_file_{suffix[:16]}") + + +async def _resolve_slack_file_url( + *, + media_url: str, + bot_token: str, +) -> tuple[str, dict[str, Any]]: + parsed = urlparse(media_url) + if parsed.scheme != "slack": + return media_url, {} + file_id = parsed.path.lstrip("/") + if not file_id: + return "", {} + try: + from slack_sdk.web.async_client import AsyncWebClient + except Exception as exc: + raise RuntimeError("slack-sdk is required to resolve Slack file metadata") from exc + + response = await AsyncWebClient(token=bot_token).files_info(file=file_id) + file_obj = response.get("file") if hasattr(response, "get") else None + if not isinstance(file_obj, dict): + raise RuntimeError("Slack files.info returned no file object") + url = str(file_obj.get("url_private_download") or file_obj.get("url_private") or "").strip() + if not url: + raise RuntimeError("Slack files.info returned no private download URL") + return url, file_obj + + +async def _download_file_limited( + *, + media_url: str, + bot_token: str, + max_bytes: int, +) -> tuple[bytes, Optional[str], Optional[str]]: + chunks: list[bytes] = [] + total = 0 + async with httpx.AsyncClient(timeout=60.0, follow_redirects=True) as client: + async with client.stream( + "GET", + media_url, + headers={"Authorization": f"Bearer {bot_token}"}, + ) as resp: + resp.raise_for_status() + headers = resp.headers + content_type = headers.get("content-type") or headers.get("Content-Type") + if content_type and "text/html" in content_type.lower(): + raise RuntimeError( + "Slack returned HTML instead of media; check bot token scopes and file permissions" + ) + content_length = headers.get("content-length") or headers.get("Content-Length") + if content_length and content_length.isdigit() and int(content_length) > max_bytes: + raise _max_size_error(max_bytes) + content_disposition = ( + headers.get("content-disposition") + or headers.get("Content-Disposition") + or "" + ) + cd_filename = _filename_from_content_disposition(content_disposition) + async for chunk in resp.aiter_bytes(8192): + total += len(chunk) + if total > max_bytes: + raise _max_size_error(max_bytes) + chunks.append(chunk) + return b"".join(chunks), cd_filename, content_type + + +async def download_inbound_media( + msg: InboundMessage, + config: dict, + *, + max_bytes: int = _DEFAULT_MAX_INBOUND_MEDIA_BYTES, +) -> Optional[DownloadedInboundMedia]: + media_url = msg.media_url or "" + if not media_url: + return None + + bot_token = resolve_bot_token(config) + if not bot_token: + log.warning("slack.media.no_token", { + "account_id": msg.account_id, + "message_id": msg.message_id, + }) + return None + + file_obj = _find_file_object(msg, media_url) + try: + resolved_url, resolved_file = await _resolve_slack_file_url( + media_url=media_url, + bot_token=bot_token, + ) + if resolved_file: + file_obj = {**file_obj, **resolved_file} + if not resolved_url: + return None + buffer, cd_filename, response_mime = await _download_file_limited( + media_url=resolved_url, + bot_token=bot_token, + max_bytes=max_bytes, + ) + except SlackInboundMediaTooLarge as exc: + log.warning("slack.media.file_too_large", { + "message_id": msg.message_id, + "error": str(exc), + }) + return None + except Exception as exc: + log.warning("slack.media.download_failed", { + "message_id": msg.message_id, + "media_url": media_url[:200], + "error": str(exc), + }) + return None + + filename = _guess_filename(msg, media_url, file_obj, cd_filename) + mime = ( + msg.media_mime + or str(file_obj.get("mimetype") or "").strip() + or response_mime + or _guess_mime_from_ext(filename) + or "application/octet-stream" + ) + if "." not in filename: + ext = mimetypes.guess_extension(mime.split(";", 1)[0].strip()) or "" + if ext: + filename = f"{filename}{ext}" + + storage_dir = _media_storage_dir(msg.account_id or "default") + storage_dir.mkdir(parents=True, exist_ok=True) + msg_id = msg.message_id or "unknown" + file_path = storage_dir / _sanitize_filename(f"{msg_id}_{filename}") + file_path.write_bytes(buffer) + + return DownloadedInboundMedia( + filename=filename, + mime=mime, + url=file_path.resolve().as_uri(), + source={ + "channel": "slack", + "account_id": msg.account_id, + "message_id": msg.message_id, + "media_url": msg.media_url, + "file_id": str(file_obj.get("id") or "") or None, + }, + ) diff --git a/flocks/channel/builtin/slack/manifest.py b/flocks/channel/builtin/slack/manifest.py new file mode 100644 index 000000000..e39405494 --- /dev/null +++ b/flocks/channel/builtin/slack/manifest.py @@ -0,0 +1,77 @@ +"""Slack app manifest helpers. + +The shape mirrors the Hermes Slack gateway setup flow: users create a Slack +app from a manifest, install it, then paste the Bot User OAuth token and +App-level token into the channel config. +""" + +from __future__ import annotations + + +def build_slack_app_manifest( + *, + app_name: str = "Flocks", + description: str = "Connect Flocks agents to Slack via Socket Mode.", +) -> dict: + """Return a Slack app manifest suitable for the Flocks Slack channel.""" + bot_scopes = [ + "app_mentions:read", + "channels:history", + "channels:read", + "chat:write", + "files:read", + "groups:history", + "groups:read", + "im:history", + "im:read", + "im:write", + "mpim:history", + "mpim:read", + "users:read", + ] + bot_events = [ + "app_mention", + "message.channels", + "message.groups", + "message.im", + "message.mpim", + ] + + return { + "_metadata": { + "major_version": 1, + "minor_version": 1, + }, + "display_information": { + "name": app_name[:35], + "description": description[:140], + "background_color": "#0f172a", + }, + "features": { + "app_home": { + "home_tab_enabled": False, + "messages_tab_enabled": True, + "messages_tab_read_only_enabled": False, + }, + "bot_user": { + "display_name": app_name[:80], + "always_online": True, + }, + }, + "oauth_config": { + "scopes": { + "bot": bot_scopes, + }, + }, + "settings": { + "event_subscriptions": { + "bot_events": bot_events, + }, + "interactivity": { + "is_enabled": False, + }, + "org_deploy_enabled": False, + "socket_mode_enabled": True, + "token_rotation_enabled": False, + }, + } diff --git a/flocks/channel/builtin/weixin/channel.py b/flocks/channel/builtin/weixin/channel.py index cb1e30d5e..fb1c51c7b 100644 --- a/flocks/channel/builtin/weixin/channel.py +++ b/flocks/channel/builtin/weixin/channel.py @@ -237,7 +237,7 @@ async def start( "base_url": self._base_url, }) if self._group_policy != "disabled": - log.warning("weixin.group_policy.note", { + log.info("weixin.group_policy.note", { "group_policy": self._group_policy, "note": ( "QR-login connects an iLink bot identity (e.g. ...@im.bot), not a " diff --git a/flocks/channel/gateway/manager.py b/flocks/channel/gateway/manager.py index 2d268560a..fa1eed803 100644 --- a/flocks/channel/gateway/manager.py +++ b/flocks/channel/gateway/manager.py @@ -12,7 +12,7 @@ import time from typing import Callable, Optional -from flocks.channel.base import ChannelPlugin, ChannelStatus +from flocks.channel.base import ChannelPlugin, ChannelStatus, NonRetryableChannelError from flocks.channel.inbound.dispatcher import InboundDispatcher from flocks.channel.registry import ChannelRegistry, default_registry from flocks.utils.log import Log @@ -281,8 +281,11 @@ async def _run_with_reconnect( await abort_event.wait() break - # Long-running start() (WebSocket / polling mode) - await self._mark_connected(plugin, channel_id) + # Long-running start() (WebSocket / polling mode). Some active + # channels only know they are connected after a platform + # handshake; let those plugins publish their own status. + if not plugin.capabilities().self_managed_connection: + await self._mark_connected(plugin, channel_id) await start_task if abort_event.is_set(): @@ -299,6 +302,14 @@ async def _run_with_reconnect( except asyncio.CancelledError: break + except NonRetryableChannelError as e: + self._record_error(plugin, channel_id, e) + log.error("gateway.non_retryable", { + "channel": channel_id, + "error": str(e), + "attempt": attempt + 1, + }) + break except Exception as e: self._record_error(plugin, channel_id, e) @@ -361,14 +372,16 @@ def _record_error( channel_id: str, error: Exception, ) -> None: - plugin.mark_disconnected(error=str(error)) + message = str(error) + if plugin.status.connected or plugin.status.last_error != message: + plugin.mark_disconnected(error=message) try: from flocks.channel.events import ChannelDisconnected from flocks.bus.bus import Bus asyncio.ensure_future( Bus.publish(ChannelDisconnected, { "channel_id": channel_id, - "reason": str(error), + "reason": message, }) ) except Exception: diff --git a/flocks/channel/inbound/dispatcher.py b/flocks/channel/inbound/dispatcher.py index 16ff95321..44a88f566 100644 --- a/flocks/channel/inbound/dispatcher.py +++ b/flocks/channel/inbound/dispatcher.py @@ -163,7 +163,10 @@ def _check_allowlist(msg: InboundMessage, config: ChannelConfig) -> bool: non-empty, but are otherwise unrestricted. """ allow_from = config.allow_from - dm_policy = config.dm_policy or "open" + dm_policy = config.dm_policy + if msg.channel_id == "slack" and dm_policy is None and allow_from: + dm_policy = "allowlist" + dm_policy = dm_policy or "open" if msg.chat_type == ChatType.DIRECT: if dm_policy == "open": @@ -720,6 +723,8 @@ async def _run_session_control(_event, parsed) -> bool: msg=msg, callbacks=callbacks, scope_override=scope_override, + channel_config=channel_config, + initial_text=command_args or None, ) return True if parsed.canonical_name == "compact": @@ -861,6 +866,8 @@ async def _handle_session_command( msg: InboundMessage, callbacks: ChannelDeliveryCallbacks, scope_override: Optional[str], + channel_config: ChannelConfig, + initial_text: Optional[str] = None, ) -> None: from flocks.session.session import Session from flocks.channel.inbound.session_binding import ( @@ -928,6 +935,26 @@ async def _handle_session_command( ] ) ) + if initial_text and initial_text.strip(): + text = initial_text.strip() + resolved_model = await _resolve_session_model(new_binding.session_id) + await self._append_user_message( + new_binding.session_id, + text, + msg, + channel_config, + model=resolved_model, + agent=new_binding.agent_id, + ) + if msg.channel_id == "feishu": + await self._run_agent_with_typing( + new_binding, + new_callbacks, + msg, + channel_config, + ) + else: + await self._run_agent(new_binding, new_callbacks) @staticmethod async def _handle_compact_command( @@ -1531,7 +1558,7 @@ def _is_placeholder_text(text: str) -> bool: ) if text in placeholders: return True - return text.startswith("[文件消息:") + return text.startswith("[文件消息:") or text.startswith("[图片消息:") # Best-effort eager registration at import time. Channels that need @@ -1560,6 +1587,12 @@ def _is_placeholder_text(text: str) -> bool: except Exception: # pragma: no cover pass +try: + from flocks.channel.builtin.slack import inbound_media as _slack_inbound_media + register_inbound_media_downloader("slack", _slack_inbound_media.download_inbound_media) +except Exception: # pragma: no cover + pass + async def _download_channel_media(msg: "InboundMessage", config: dict) -> Any: """Dispatch inbound media download to the appropriate channel handler. @@ -1582,4 +1615,7 @@ async def _download_channel_media(msg: "InboundMessage", config: dict) -> Any: if channel_id == "telegram": from flocks.channel.builtin.telegram import inbound_media as _telegram_inbound_media return await _telegram_inbound_media.download_inbound_media(msg, config) + if channel_id == "slack": + from flocks.channel.builtin.slack import inbound_media as _slack_inbound_media + return await _slack_inbound_media.download_inbound_media(msg, config) return None diff --git a/flocks/channel/registry.py b/flocks/channel/registry.py index e3a138438..2a8116ff9 100644 --- a/flocks/channel/registry.py +++ b/flocks/channel/registry.py @@ -93,6 +93,7 @@ def _register_builtin_channels(self) -> None: from flocks.channel.builtin.dingtalk.channel import DingTalkChannel from flocks.channel.builtin.email.channel import EmailChannel from flocks.channel.builtin.feishu.channel import FeishuChannel + from flocks.channel.builtin.slack.channel import SlackChannel from flocks.channel.builtin.telegram.channel import TelegramChannel from flocks.channel.builtin.wecom.channel import WeComChannel from flocks.channel.builtin.whatsapp.channel import WhatsAppChannel @@ -103,6 +104,7 @@ def _register_builtin_channels(self) -> None: self.register(WhatsAppChannel()) self.register(DingTalkChannel()) self.register(WeixinChannel()) + self.register(SlackChannel()) self.register(EmailChannel()) def _register_plugin_extension_point(self) -> None: diff --git a/flocks/cli/commands/doctor.py b/flocks/cli/commands/doctor.py index e6e9a9dec..940561006 100644 --- a/flocks/cli/commands/doctor.py +++ b/flocks/cli/commands/doctor.py @@ -2,8 +2,8 @@ from __future__ import annotations +import asyncio import os -import shlex import shutil import subprocess import sys @@ -18,14 +18,12 @@ def doctor_command() -> None: - """Run the source installer from the Flocks source directory.""" + """Repair the active source installation using the shared core installer.""" source_root = _find_source_root() - script = _select_source_install_script(source_root) - command = _build_source_install_command(script) env = _build_source_install_env() console.print(f"[cyan]Flocks source directory:[/cyan] {source_root}") - console.print(f"[cyan]Source install command:[/cyan] {_format_command(command)}") + console.print("[cyan]Repairing Flocks core installation...[/cyan]") if _needs_windows_handoff(): _start_windows_handoff(source_root, env=env) @@ -35,12 +33,41 @@ def doctor_command() -> None: ) return - _run_source_install(command, source_root=source_root, env=env) + _run_core_install(source_root, env=env) console.print("[green]安装正常[/green]") _print_service_diagnosis() +def _run_core_install(source_root: Path, *, env: dict[str, str] | None = None) -> None: + """Run the shared updater/doctor core installation entry point.""" + from flocks.updater import updater + + uv_path = shutil.which("uv") + if not uv_path: + script_name = "install.ps1" if _is_windows() else "install.sh" + console.print( + f"[red]uv was not found. Run scripts/{script_name} to install system prerequisites.[/red]" + ) + raise typer.Exit(1) + + install_env = env or os.environ + try: + asyncio.run( + updater.install_or_repair_source( + install_root=source_root, + uv_path=uv_path, + version=updater.get_current_version(), + uv_default_index=install_env.get("FLOCKS_UV_DEFAULT_INDEX"), + npm_registry=install_env.get("FLOCKS_NPM_REGISTRY"), + sync_timeout=300, + ) + ) + except RuntimeError as error: + console.print(f"[red]Core installation repair failed: {error}[/red]") + raise typer.Exit(1) from error + + def _find_source_root(start: Path | None = None) -> Path: """Find the repository root that owns the source install scripts.""" current = (start or Path(__file__)).resolve() @@ -55,44 +82,6 @@ def _find_source_root(start: Path | None = None) -> Path: raise typer.BadParameter("Could not locate the Flocks source directory.") -def _select_source_install_script(source_root: Path) -> Path: - """Select the platform-specific source install script.""" - suffix = ".ps1" if _is_windows() else ".sh" - script = source_root / "scripts" / f"install{suffix}" - - if not script.is_file(): - raise typer.BadParameter(f"Source installer not found: {script}") - - return script - - -def _build_source_install_command(script: Path) -> list[str]: - """Build the subprocess command for the selected installer.""" - if script.suffix == ".ps1": - powershell = _find_powershell() - return [ - powershell, - "-NoProfile", - "-ExecutionPolicy", - "Bypass", - "-File", - str(script), - ] - - return ["bash", str(script)] - - -def _run_source_install(command: list[str], *, source_root: Path, env: dict[str, str] | None) -> None: - """Run the selected source installer synchronously.""" - try: - subprocess.run(command, cwd=source_root, check=True, env=env) - except FileNotFoundError as error: - console.print(f"[red]Failed to start installer: {error}[/red]") - raise typer.Exit(1) from error - except subprocess.CalledProcessError as error: - raise typer.Exit(error.returncode or 1) from error - - def _build_source_install_env() -> dict[str, str] | None: """Build installer environment from the persisted install language.""" if not is_cn_install_language(): @@ -195,11 +184,6 @@ def _service_status_is_healthy(status_lines: list[str]) -> bool: return backend_running and webui_running -def _format_command(command: list[str]) -> str: - """Return a shell-readable representation of the command.""" - return " ".join(shlex.quote(part) for part in command) - - def _is_windows() -> bool: """Return whether the current platform should use PowerShell installers.""" return sys.platform.startswith("win") diff --git a/flocks/cli/commands/session.py b/flocks/cli/commands/session.py index 625cde4e9..fb13111f3 100644 --- a/flocks/cli/commands/session.py +++ b/flocks/cli/commands/session.py @@ -8,6 +8,7 @@ import asyncio import json import os +from functools import cache from typing import Optional import typer @@ -76,9 +77,13 @@ async def _list_sessions( all_sessions = await Session.list_all() if project_filter == "": - result = await Project.from_directory(os.getcwd()) - current_project_id = result["project"].id - all_sessions = [s for s in all_sessions if s.project_id == current_project_id] + resolve_worktree = cache(Project.worktree_for_directory) + current_worktree = resolve_worktree(os.getcwd()) + all_sessions = [ + session + for session in all_sessions + if session.directory and resolve_worktree(session.directory) == current_worktree + ] elif project_filter: all_sessions = [s for s in all_sessions if s.project_id == project_filter] diff --git a/flocks/cli/commands/stats.py b/flocks/cli/commands/stats.py index ec3093198..74eca6244 100644 --- a/flocks/cli/commands/stats.py +++ b/flocks/cli/commands/stats.py @@ -6,6 +6,7 @@ import os from dataclasses import dataclass, field from datetime import UTC, datetime, timedelta +from functools import cache from typing import Dict, List, Optional, Sequence import typer @@ -118,8 +119,13 @@ async def _resolve_project_sessions(project_filter: Optional[str]) -> List[Sessi if project_filter is None: return sessions if project_filter == "": - result = await Project.from_directory(os.getcwd()) - project_filter = result["project"].id + resolve_worktree = cache(Project.worktree_for_directory) + current_worktree = resolve_worktree(os.getcwd()) + return [ + session + for session in sessions + if session.directory and resolve_worktree(session.directory) == current_worktree + ] return [session for session in sessions if session.project_id == project_filter] diff --git a/flocks/cli/commands/update.py b/flocks/cli/commands/update.py index c288b9755..fd9978bb7 100644 --- a/flocks/cli/commands/update.py +++ b/flocks/cli/commands/update.py @@ -35,7 +35,7 @@ def update_command( async def _update(check: bool, yes: bool, force: bool = False, region: str | None = None) -> None: - from flocks.updater import build_updated_frontend, check_update, perform_update, detect_deploy_mode + from flocks.updater import check_update, detect_deploy_mode, perform_update from flocks.cli.install_profile import is_cn_install_language if region is None and is_cn_install_language(): @@ -97,26 +97,15 @@ async def _load_update_info(selected_region: str | None): console.print("[yellow]已取消[/yellow]") return - from flocks.cli.service_manager import ServiceError, stop_all - - try: - stop_all(console) - except ServiceError as error: - console.print(f"[red]执行 flocks stop 失败：{error}[/red]") - raise typer.Exit(1) - - append_upgrade_text_log(f"OK cli_update_stopped_services_before_upgrade version={version_to_apply}") - console.print("[green]✓ 已执行 flocks stop[/green]") console.print() stage_labels = { "fetching": "下载最新源码包", "backing_up": "备份当前版本", - "applying": f"应用 v{info.latest_version}", - "syncing": "同步依赖", + "applying": "应用新版本", "restarting": "重启服务", "done": "完成", } - total_steps = 6 + total_steps = 3 seen_stages: set[str] = set() step = 0 active_stage: str | None = None @@ -137,8 +126,9 @@ def _finish_active(success: bool = True) -> None: tarball_url=info.tarball_url, bundle_sha256=info.bundle_sha256, bundle_format=info.bundle_format, - restart=False, + restart=True, region=region, + wait_for_handoff=True, ): if progress.stage == "error": _finish_active(success=False) @@ -149,6 +139,9 @@ def _finish_active(success: bool = True) -> None: _finish_active(success=True) continue + if progress.stage == "restarting": + continue + if progress.stage not in seen_stages: _finish_active(success=True) seen_stages.add(progress.stage) @@ -157,23 +150,8 @@ def _finish_active(success: bool = True) -> None: console.print(f"[cyan][{step}/{total_steps}] {label}...[/cyan] ", end="") active_stage = progress.stage - step += 1 - console.print(f"[cyan][{step}/{total_steps}] 构建前端...[/cyan] ", end="") - try: - await build_updated_frontend(region=region) - except Exception as error: - console.print("[red]✗[/red]") - console.print(f"\n[red]✗ 前端构建失败：{error}[/red]") - raise typer.Exit(1) - console.print("[green]✓[/green]") - - step += 1 - console.print(f"[cyan][{step}/{total_steps}] 完成[/cyan] ", end="") - console.print("[green]✓[/green]") - append_upgrade_text_log(f"OK cli_update_completed version={version_to_apply}") console.print(f"\n[green]✓ 升级完成 → v{version_to_apply}[/green]") - console.print("[dim]如需恢复服务，请执行 [bold]flocks start[/bold][/dim]") def _print_version_table(info) -> None: diff --git a/flocks/cli/service_control.py b/flocks/cli/service_control.py index 34320fa9f..50bd7ae74 100644 --- a/flocks/cli/service_control.py +++ b/flocks/cli/service_control.py @@ -183,23 +183,6 @@ def request_restart_webui( return parse_supervisor_status(data) -def request_prepare_upgrade(*, paths=None, timeout: float | None = 30.0) -> SupervisorStatus: - """Ask the supervisor daemon to pause managed services for upgrade handoff.""" - payload = _post_control_json("/upgrade/prepare", paths=paths, timeout=timeout) - return parse_supervisor_status(payload) - - -def request_resume_upgrade( - config: ServiceConfig, - *, - paths=None, - timeout: float | None = 180.0, -) -> SupervisorStatus: - """Ask the supervisor daemon to resume managed services after upgrade handoff.""" - payload = _post_control_json("/upgrade/resume", payload=service_config_payload(config), paths=paths, timeout=timeout) - return parse_supervisor_status(payload) - - def read_logs( *, service: str, diff --git a/flocks/cli/service_manager.py b/flocks/cli/service_manager.py index 46134d52f..cd52842ef 100644 --- a/flocks/cli/service_manager.py +++ b/flocks/cli/service_manager.py @@ -50,7 +50,7 @@ LOG_TRIM_CHUNK_BYTES = 1024 * 1024 WEBUI_DIRECT_BACKEND_URLS_ENV = "FLOCKS_WEBUI_DIRECT_BACKEND_URLS" DEFAULT_FLOCKS_CONSOLE_BASE_URL = "https://portalflocks.threatbook.cn" -DEFAULT_VITE_ADDITIONAL_SERVER_ALLOWED_HOSTS = "portalflocks.threatbook.cn" +DEFAULT_VITE_ADDITIONAL_SERVER_ALLOWED_HOSTS = "127.0.0.1,localhost,portalflocks.threatbook.cn" MISSING_PORT_OWNER_TOOLS_WARNING = ( "未检测到 lsof 或 fuser，无法解析端口占用 PID；将退回到 bind 检查。" "可尝试安装：apt/yum install lsof -y" @@ -89,21 +89,6 @@ class RuntimeRecord: started_at: float | None = None -@dataclass(frozen=True) -class UpgradeRuntimeInfo: - payload_present: bool = False - pid_file_present: bool = False - upgrade_pid: int | None = None - frontend_host: str | None = None - frontend_port: int | None = None - listener_pids: tuple[int, ...] = () - page_active: bool = False - - @property - def has_artifacts(self) -> bool: - return self.payload_present or self.pid_file_present - - def repo_root() -> Path: """Return the installed repository root.""" override = os.getenv("FLOCKS_REPO_ROOT") @@ -723,51 +708,6 @@ def _console_print(console, message: str) -> None: console.print(message) -def _read_upgrade_runtime_info(frontend_port: int | None = None) -> UpgradeRuntimeInfo: - try: - from flocks.updater import updater as updater_module - - payload = updater_module.read_upgrade_runtime_state(frontend_port=frontend_port) - except Exception: - return UpgradeRuntimeInfo(frontend_port=frontend_port) - - listener_pids = tuple(int(pid) for pid in payload.get("listener_pids", []) if isinstance(pid, int)) - return UpgradeRuntimeInfo( - payload_present=bool(payload.get("payload_present")), - pid_file_present=bool(payload.get("pid_file_present")), - upgrade_pid=payload.get("upgrade_pid") if isinstance(payload.get("upgrade_pid"), int) else None, - frontend_host=payload.get("frontend_host") if isinstance(payload.get("frontend_host"), str) else None, - frontend_port=payload.get("frontend_port") if isinstance(payload.get("frontend_port"), int) else frontend_port, - listener_pids=listener_pids, - page_active=bool(payload.get("page_active")), - ) - - -def _resolve_upgrade_runtime(console, *, frontend_port: int, attempt_recover: bool) -> dict[str, object]: - upgrade_info = _read_upgrade_runtime_info(frontend_port) - if not upgrade_info.has_artifacts: - return {"action": "noop", "error": None} - - from flocks.updater import updater as updater_module - - _console_print(console, "[flocks] 检测到升级临时页残留，正在尝试恢复或清理...") - result = updater_module.resolve_upgrade_runtime_state( - attempt_recover=attempt_recover, - frontend_port=upgrade_info.frontend_port or frontend_port, - ) - - action = str(result.get("action") or "noop") - error = result.get("error") - if action == "recovered": - _console_print(console, "[flocks] 已恢复未完成升级，正式 WebUI 将继续接管端口。") - elif action != "noop": - _console_print(console, "[flocks] 已清理升级临时页残留。") - - if isinstance(error, str) and error: - _console_print(console, f"[flocks] 未完成升级的自动恢复失败，已清理临时升级页: {error}") - return result - - def cleanup_stale_pid_file(pid_file: Path) -> None: """Remove pid files that no longer point to running processes.""" if not pid_file.exists(): @@ -1569,7 +1509,6 @@ def _start_all_without_stop(config: ServiceConfig, console) -> None: def _start_all_unlocked(config: ServiceConfig, console, *, paths: RuntimePaths) -> None: """Ensure the supervisor daemon is running; caller must hold lifecycle lock.""" - _resolve_upgrade_runtime(console, frontend_port=config.frontend_port, attempt_recover=False) if supervisor_is_running(paths): status = None try: diff --git a/flocks/cli/service_supervisor.py b/flocks/cli/service_supervisor.py index e94397300..4a615b59a 100644 --- a/flocks/cli/service_supervisor.py +++ b/flocks/cli/service_supervisor.py @@ -119,8 +119,6 @@ def __init__( self._shutdown_requested = threading.Event() self._server: ThreadingHTTPServer | None = None self._server_thread: threading.Thread | None = None - self._backend_paused = False - self._webui_paused = False self.backend = ManagedService( name="backend", label="后端", @@ -277,15 +275,6 @@ def do_POST(self) -> None: if parsed.path == "/stop/webui": self._send_json({"error": "static WebUI is served by Flocks service and cannot be stopped separately"}, status=409) return - if parsed.path == "/upgrade/prepare": - daemon.prepare_upgrade(reason="control upgrade prepare") - self._send_json(daemon.status_payload()) - return - if parsed.path == "/upgrade/resume": - daemon.update_config(payload) - daemon.resume_upgrade(reason="control upgrade resume") - self._send_json(daemon.status_payload()) - return self._send_json({"error": "not found"}, status=404) except _CLIENT_DISCONNECT_ERRORS: return @@ -320,8 +309,8 @@ def status_payload(self) -> dict[str, object]: "state": "stopping" if self._shutdown_requested.is_set() else "running", "log_path": str(supervisor_log_path(self.paths)), }, - "backend": _service_payload(self.backend, paused=self._backend_paused), - "webui": _service_payload(self.webui, paused=self._webui_paused), + "backend": _service_payload(self.backend), + "webui": _service_payload(self.webui), "config": service_config_payload(self.config), } @@ -411,22 +400,18 @@ def _log_paths_for_service(self, service_name: str) -> list[tuple[str, Path]]: def restart_all(self, *, reason: str) -> None: with self._lock: - self._backend_paused = False - self._webui_paused = False self._restart_service(self.backend, reason=reason, immediate=True) self._start_backend_locked(immediate=True) self._sync_static_webui_state() def restart_backend(self, *, reason: str) -> None: with self._lock: - self._backend_paused = False self._restart_service(self.backend, reason=reason, immediate=True) self._start_backend_locked(immediate=True) self._sync_static_webui_state() def restart_webui(self, *, reason: str, force_frontend_build: bool = False) -> None: with self._lock: - self._webui_paused = False if force_frontend_build: from flocks.cli.service_config import with_frontend_build @@ -435,27 +420,6 @@ def restart_webui(self, *, reason: str, force_frontend_build: bool = False) -> N self._start_backend_locked(immediate=True) self._sync_static_webui_state() - def prepare_upgrade(self, *, reason: str) -> None: - with self._lock: - self._backend_paused = True - self._webui_paused = True - _daemon_log("service_pause", {"service": "backend", "reason": reason}) - _daemon_log("service_pause", {"service": "webui", "reason": reason}) - self.backend.last_error = reason - self.webui.last_error = reason - self._stop_service(self.backend) - self.webui.state = "paused" - - def resume_upgrade(self, *, reason: str) -> None: - with self._lock: - self._backend_paused = False - self._webui_paused = False - _daemon_log("service_resume", {"service": "backend", "reason": reason}) - _daemon_log("service_resume", {"service": "webui", "reason": reason}) - self._probe_backend_locked() - self._start_backend_locked(immediate=True) - self._sync_static_webui_state() - def shutdown_children(self) -> None: with self._lock: self._stop_service(self.backend) @@ -463,10 +427,8 @@ def shutdown_children(self) -> None: def tick(self) -> None: with self._lock: - if not self._backend_paused: - self._probe_backend_locked() - if not self._backend_paused: - self._start_backend_locked(immediate=False) + self._probe_backend_locked() + self._start_backend_locked(immediate=False) self._sync_static_webui_state() def _restart_service(self, service: ManagedService, *, reason: str, immediate: bool) -> None: @@ -547,9 +509,6 @@ def _sync_static_webui_state(self) -> None: self.webui.log_path = self.paths.backend_log self.webui.process = None self.webui.command = () - if self._webui_paused: - self.webui.state = "paused" - return if self.backend.state == "healthy": self.webui.state = "static" self.webui.last_error = None diff --git a/flocks/cli/session_runner.py b/flocks/cli/session_runner.py index ccbf254cf..ff57c4b1d 100644 --- a/flocks/cli/session_runner.py +++ b/flocks/cli/session_runner.py @@ -12,6 +12,7 @@ import asyncio import json import os +from functools import cache from pathlib import Path from typing import Optional, Dict, Any @@ -253,9 +254,18 @@ async def _get_or_create_session( continue_session: bool = False, ) -> SessionInfo: """Get or create session.""" + resolve_worktree = cache(Project.worktree_for_directory) + current_worktree = resolve_worktree(str(self.directory)) + + def belongs_to_current_worktree(session: SessionInfo) -> bool: + return bool( + session.directory + and resolve_worktree(session.directory) == current_worktree + ) + if session_id: session = await Session.get(project_id, session_id) - if session: + if session and belongs_to_current_worktree(session): return session self.console.print(f"[yellow]Session {session_id} not found, creating new[/yellow]") @@ -263,7 +273,7 @@ async def _get_or_create_session( sessions = await Session.list(project_id) if sessions: for s in sessions: - if not s.parent_id: + if not s.parent_id and belongs_to_current_worktree(s): return s return await Session.create( diff --git a/flocks/config/config.py b/flocks/config/config.py index 9e5fc41fe..a5f234f9c 100644 --- a/flocks/config/config.py +++ b/flocks/config/config.py @@ -349,6 +349,20 @@ class ToolOutputConfig(BaseModel): ) +class ToolFailureConfig(BaseModel): + """Repeated tool-failure handling.""" + + model_config = {"populate_by_name": True} + + disable_on_repeated_failure: bool = Field( + True, + alias="disableOnRepeatedFailure", + description=( + "Disable a standalone custom tool after repeated identical failures." + ), + ) + + class EnterpriseConfig(BaseModel): """Enterprise configuration""" @@ -675,6 +689,11 @@ class ConfigInfo(BaseModel): alias="toolOutput", description="Tool output size limits (read, truncation caps).", ) + tool_failure: Optional[ToolFailureConfig] = Field( + None, + alias="toolFailure", + description="Repeated tool-failure handling.", + ) experimental: Optional[ExperimentalConfig] = None # Memory system configuration (added for memory system integration) @@ -1372,7 +1391,13 @@ async def resolve_default_llm(cls) -> Optional[Dict[str, str]]: return None @classmethod - async def update(cls, config: ConfigInfo, project_dir: Optional[Path] = None) -> None: + async def update( + cls, + config: ConfigInfo, + project_dir: Optional[Path] = None, + *, + channel_allow_from_deletions: Optional[set[str]] = None, + ) -> None: """ Update configuration @@ -1380,6 +1405,9 @@ async def update(cls, config: ConfigInfo, project_dir: Optional[Path] = None) -> config: New configuration project_dir: Deprecated and ignored. Config is always written to the unified user config directory. + channel_allow_from_deletions: Channel IDs whose persisted + allowFrom field should be removed after a successful full + config validation and merge. """ _ = project_dir @@ -1396,6 +1424,13 @@ async def update(cls, config: ConfigInfo, project_dir: Optional[Path] = None) -> # Write config_data = merged.model_dump(by_alias=True, exclude_none=True, mode="json") + if channel_allow_from_deletions: + channels = config_data.get("channels") + if isinstance(channels, dict): + for channel_id in channel_allow_from_deletions: + channel_cfg = channels.get(channel_id) + if isinstance(channel_cfg, dict): + channel_cfg.pop("allowFrom", None) config_file.write_text(json.dumps(config_data, indent=2), encoding="utf-8") # Clear cache diff --git a/flocks/hub/catalog.py b/flocks/hub/catalog.py index d82e001f5..87a45ce51 100644 --- a/flocks/hub/catalog.py +++ b/flocks/hub/catalog.py @@ -740,6 +740,23 @@ def _version_tuple(value: str) -> tuple[int, ...]: return tuple(int(part) for part in parts) if parts else (0,) +def _catalog_install_state( + install_path: Optional[Path], + record: Optional[local.InstalledPluginRecord], + available_version: str, +) -> tuple[str, Optional[str]]: + if install_path is None: + return "available", None + if record is None: + # An inferred install has no trustworthy package version because Hub + # manifests are not copied into the install directory. Reconcile it + # through the update path instead of assuming it is already current. + return "updateAvailable", None + if _version_tuple(record.version) < _version_tuple(available_version): + return "updateAvailable", record.version + return "installed", record.version + + def _os_compatible(manifest: HubPluginManifest) -> bool: allowed = {item.lower() for item in manifest.compatibility.os} if not allowed: @@ -751,15 +768,18 @@ def _os_compatible(manifest: HubPluginManifest) -> bool: def _entry_from_manifest(manifest: HubPluginManifest) -> HubCatalogEntry: record = local.get_record(manifest.type, manifest.id) - install_path = _resolve_install_path(manifest.type, manifest.id, record, local.infer_local_installs()) - - state = "available" - installed_version: Optional[str] = None - if install_path: - installed_version = record.version if record else manifest.version - state = "installed" - if _version_tuple(installed_version) < _version_tuple(manifest.version): - state = "updateAvailable" + install_path, record = _resolve_install_path( + manifest.type, + manifest.id, + record, + local.infer_local_installs(), + ) + + state, installed_version = _catalog_install_state( + install_path, + record, + manifest.version, + ) if not _os_compatible(manifest): state = "incompatible" @@ -792,15 +812,16 @@ def _resolve_install_path( plugin_id: str, record: Optional[local.InstalledPluginRecord], inferred_installs: Optional[dict[tuple[PluginType, str], Path]] = None, -) -> Optional[Path]: +) -> tuple[Optional[Path], Optional[local.InstalledPluginRecord]]: if record and record.installPath: path = Path(record.installPath) if local.has_install_payload(plugin_type, path): - return path + return path, record local.remove_installed_record(plugin_type, plugin_id) + record = None if inferred_installs is not None: - return inferred_installs.get((plugin_type, plugin_id)) - return local.infer_local_install(plugin_type, plugin_id) + return inferred_installs.get((plugin_type, plugin_id)), record + return local.infer_local_install(plugin_type, plugin_id), record def _entry_from_index( @@ -809,15 +830,18 @@ def _entry_from_index( inferred_installs: dict[tuple[PluginType, str], Path], ) -> HubCatalogEntry: record = records.get(f"{item.type}:{item.id}") - install_path = _resolve_install_path(item.type, item.id, record, inferred_installs) + install_path, record = _resolve_install_path( + item.type, + item.id, + record, + inferred_installs, + ) - state = "available" - installed_version: Optional[str] = None - if install_path: - installed_version = record.version if record else item.version - state = "installed" - if _version_tuple(installed_version) < _version_tuple(item.version): - state = "updateAvailable" + state, installed_version = _catalog_install_state( + install_path, + record, + item.version, + ) return HubCatalogEntry( id=item.id, @@ -875,11 +899,10 @@ def _entry_from_bundled_tool( """Catalog entry for a tool bundled inside flockshub. State is computed exactly like :func:`_entry_from_index`: - * ``installed`` — the user installed it under - ``~/.flocks/plugins/tools///`` - (record or on-disk payload found). - * ``updateAvailable`` — installed copy is older than the bundled - version. + * ``installed`` — an installed record confirms the local copy is + current. + * ``updateAvailable`` — the installed copy is older, or an unrecorded + on-disk payload needs reconciliation. * ``available`` — bundled in flockshub, no local install. The runtime ``ToolRegistry`` only picks up installed tools, so an "available" entry @@ -889,15 +912,18 @@ def _entry_from_bundled_tool( * ``incompatible`` — manifest declares an incompatible OS. """ record = records.get(f"{manifest.type}:{manifest.id}") - install_path = _resolve_install_path(manifest.type, manifest.id, record, inferred_installs) - - state = "available" - installed_version: Optional[str] = None - if install_path: - installed_version = record.version if record else manifest.version - state = "installed" - if _version_tuple(installed_version) < _version_tuple(manifest.version): - state = "updateAvailable" + install_path, record = _resolve_install_path( + manifest.type, + manifest.id, + record, + inferred_installs, + ) + + state, installed_version = _catalog_install_state( + install_path, + record, + manifest.version, + ) if not _os_compatible(manifest): state = "incompatible" diff --git a/flocks/hub/installer.py b/flocks/hub/installer.py index bcbbf2407..07b9a30b0 100644 --- a/flocks/hub/installer.py +++ b/flocks/hub/installer.py @@ -90,21 +90,54 @@ def _resolve_install_destination( return local.install_dir(plugin_type, plugin_id, scope) -def _copy_package(src: Path, dst: Path) -> None: +def _remove_path(path: Path) -> None: + if path.is_dir() and not path.is_symlink(): + shutil.rmtree(path) + elif path.exists() or path.is_symlink(): + path.unlink() + + +def _replace_prepared_path(prepared: Path, dst: Path) -> Path | None: + backup: Path | None = None + if dst.exists() or dst.is_symlink(): + backup = dst.parent / f".{dst.name}.bak" + _remove_path(backup) + dst.replace(backup) + try: + prepared.replace(dst) + except Exception: + if backup is not None and (backup.exists() or backup.is_symlink()): + backup.replace(dst) + raise + return backup + + +def _commit_replacement(backup: Path | None) -> None: + if backup is None: + return + try: + _remove_path(backup) + except OSError: + pass + + +def _rollback_replacement(dst: Path, backup: Path | None) -> None: + _remove_path(dst) + if backup is not None and (backup.exists() or backup.is_symlink()): + backup.replace(dst) + + +def _copy_package(src: Path, dst: Path, *, retain_backup: bool = False) -> Path | None: parent = dst.parent parent.mkdir(parents=True, exist_ok=True) tmp = Path(tempfile.mkdtemp(prefix=f".{dst.name}.", dir=str(parent))) try: _copy_package_contents(src, tmp) - backup = None - if dst.exists(): - backup = parent / f".{dst.name}.bak" - if backup.exists(): - shutil.rmtree(backup) - dst.replace(backup) - tmp.replace(dst) - if backup and backup.exists(): - shutil.rmtree(backup) + backup = _replace_prepared_path(tmp, dst) + if not retain_backup: + _commit_replacement(backup) + return None + return backup except Exception: if tmp.exists(): shutil.rmtree(tmp, ignore_errors=True) @@ -122,32 +155,26 @@ def _copy_package_contents(src: Path, dst: Path) -> None: shutil.copy2(item, target) -def _replace_dir(src: Path, dst: Path) -> None: - parent = dst.parent - backup = None - if dst.exists(): - backup = parent / f".{dst.name}.bak" - if backup.exists(): - shutil.rmtree(backup) - dst.replace(backup) - src.replace(dst) - if backup and backup.exists(): - shutil.rmtree(backup) - - def _contracts_access_dir(plugin_id: str, scope: str) -> Path: return local.install_root("webui", scope).parent / "access" / plugin_id -def _copy_attached_access_contracts(plugin_type: PluginType, plugin_id: str, src: Path, scope: str) -> Path | None: +def _copy_attached_access_contracts( + plugin_type: PluginType, + plugin_id: str, + src: Path, + scope: str, + *, + retain_backup: bool = False, +) -> tuple[Path, Path | None] | None: if plugin_type != "webui": return None access_src = src / "access" if not access_src.is_dir(): return None access_dst = _contracts_access_dir(plugin_id, scope) - _copy_package(access_src, access_dst) - return access_dst + backup = _copy_package(access_src, access_dst, retain_backup=retain_backup) + return access_dst, backup def _remove_attached_access_contracts(plugin_type: PluginType, plugin_id: str, scope: str) -> None: @@ -182,21 +209,34 @@ def _build_webui_pages(plugin_id: str, install_dir: Path) -> None: raise RuntimeError(f"Failed to build WebUI page bundle for {plugin_id}/{page.id}{detail}") -def _copy_webui_package_with_build(plugin_id: str, src: Path, dst: Path) -> None: +def _copy_webui_package_with_build( + plugin_id: str, + src: Path, + dst: Path, + *, + retain_backup: bool = False, +) -> Path | None: parent = dst.parent parent.mkdir(parents=True, exist_ok=True) tmp = Path(tempfile.mkdtemp(prefix=f".{dst.name}.", dir=str(parent))) try: _copy_package_contents(src, tmp) _build_webui_pages(plugin_id, tmp) - _replace_dir(tmp, dst) + backup = _replace_prepared_path(tmp, dst) + if not retain_backup: + _commit_replacement(backup) + return None + return backup except Exception: if tmp.exists(): shutil.rmtree(tmp, ignore_errors=True) raise -async def _refresh_runtime(plugin_type: PluginType) -> None: +async def _refresh_runtime( + plugin_type: PluginType, + changed_path: Path | None = None, +) -> None: if plugin_type == "skill": from flocks.skill.skill import Skill @@ -223,7 +263,13 @@ async def _refresh_runtime(plugin_type: PluginType) -> None: from flocks.tool.registry import ToolRegistry await ToolRegistry.init_async() - await asyncio.to_thread(ToolRegistry.refresh_plugin_tools) + if changed_path is None: + await asyncio.to_thread(ToolRegistry.refresh_plugin_tools) + else: + await asyncio.to_thread( + ToolRegistry.refresh_plugin_tools, + changed_path=changed_path, + ) clear_device_template_cache() # Drop the descriptor cache so freshly installed/uninstalled # API plugins surface in ``_load_provider_yaml_metadata`` (and @@ -331,16 +377,6 @@ async def _rollback_component_ref_installs( continue -def _rollback_install_path(plugin_type: PluginType, plugin_id: str, install_path: Path, scope: str) -> None: - if ".flocks/plugins" in install_path.as_posix(): - if install_path.is_dir(): - shutil.rmtree(install_path, ignore_errors=True) - elif install_path.exists(): - install_path.unlink() - _remove_attached_access_contracts(plugin_type, plugin_id, scope) - local.remove_installed_record(plugin_type, plugin_id) - - def _is_project_install_path(plugin_type: PluginType, install_path: Path) -> bool: try: project_root = local.install_root(plugin_type, "project").resolve() @@ -496,16 +532,31 @@ async def install_plugin( validate_package(src, manifest) dst = _resolve_install_destination(plugin_type, plugin_id, src, scope) component_key = f"component:{plugin_id}" - component_had_install = plugin_type == "component" and local.infer_local_install(plugin_type, plugin_id) is not None component_ref_installs: list[tuple[PluginType, str]] = [] + previous_record = local.get_record(plugin_type, plugin_id) + package_backup: Path | None = None + package_replaced = False + access_replacement: tuple[Path, Path | None] | None = None try: if plugin_type == "component": component_ref_installs = await _install_component_refs(manifest, scope=scope, progress=progress) if plugin_type == "webui": - _copy_webui_package_with_build(plugin_id, src, dst) + package_backup = _copy_webui_package_with_build( + plugin_id, + src, + dst, + retain_backup=True, + ) else: - _copy_package(src, dst) - _copy_attached_access_contracts(plugin_type, plugin_id, src, scope) + package_backup = _copy_package(src, dst, retain_backup=True) + package_replaced = True + access_replacement = _copy_attached_access_contracts( + plugin_type, + plugin_id, + src, + scope, + retain_backup=True, + ) record = local.make_record( plugin_type=plugin_type, plugin_id=plugin_id, @@ -518,15 +569,30 @@ async def install_plugin( ) local.save_installed_record(record) clear_catalog_caches() - await _refresh_runtime(plugin_type) + await _refresh_runtime(plugin_type, dst) if plugin_type == "component": await _emit_component_progress(progress, manifest, "complete", record=record, message="Installed") + if access_replacement is not None: + _commit_replacement(access_replacement[1]) + _commit_replacement(package_backup) return record except Exception: + if access_replacement is not None: + _rollback_replacement(*access_replacement) + if package_replaced: + _rollback_replacement(dst, package_backup) + if previous_record is None: + local.remove_installed_record(plugin_type, plugin_id) + else: + local.save_installed_record(previous_record) + clear_catalog_caches() if plugin_type == "component": - if not component_had_install: - _rollback_install_path(plugin_type, plugin_id, dst, scope) await _rollback_component_ref_installs(component_ref_installs, component_key) + if package_replaced: + try: + await _refresh_runtime(plugin_type, dst) + except Exception: + pass raise @@ -584,13 +650,24 @@ async def uninstall_plugin(plugin_type: PluginType, plugin_id: str) -> bool: record = local.get_record(plugin_type, plugin_id) install_path = Path(record.installPath) if record and record.installPath else local.infer_local_install(plugin_type, plugin_id) if install_path is None or not install_path.exists(): + changed_path = install_path + if changed_path is None and record is not None: + try: + changed_path = _resolve_install_destination( + plugin_type, + plugin_id, + plugin_root(plugin_type, plugin_id), + record.scope, + ) + except Exception: + changed_path = local.install_dir(plugin_type, plugin_id, record.scope) children_removed = await _uninstall_component_refs(manifest) if manifest is not None else False had_record = record is not None local.remove_installed_record(plugin_type, plugin_id) clear_catalog_caches() _clear_device_template_cache_if_needed(plugin_type) if children_removed or had_record: - await _refresh_runtime(plugin_type) + await _refresh_runtime(plugin_type, changed_path) return children_removed or had_record project_root = local.install_root(plugin_type, "project").resolve() resolved_install_path = install_path.resolve() @@ -618,5 +695,5 @@ async def uninstall_plugin(plugin_type: PluginType, plugin_id: str) -> bool: local.remove_installed_record(plugin_type, plugin_id) clear_catalog_caches() _cleanup_orphan_api_services(orphan_keys) - await _refresh_runtime(plugin_type) + await _refresh_runtime(plugin_type, install_path) return True diff --git a/flocks/hub/security.py b/flocks/hub/security.py index c630fdf41..3707add24 100644 --- a/flocks/hub/security.py +++ b/flocks/hub/security.py @@ -23,6 +23,30 @@ def _sha256(path: Path) -> str: return f"sha256:{digest.hexdigest()}" +def _validate_python_syntax(package_dir: Path, manifest: HubPluginManifest) -> None: + python_files = { + package_dir / entrypoint + for entrypoint in manifest.entrypoints + if Path(entrypoint).suffix == ".py" + } + if manifest.type in {"tool", "device"}: + python_files.update( + path + for path in package_dir.rglob("*.py") + if not any(part in SKIP_NAMES for part in path.relative_to(package_dir).parts) + ) + + for path in sorted(python_files): + if not path.is_file(): + continue + relative_path = path.relative_to(package_dir).as_posix() + try: + source = path.read_text(encoding="utf-8") + compile(source, str(path), "exec") + except (SyntaxError, UnicodeDecodeError) as exc: + raise ValueError(f"Invalid Python source in package: {relative_path}: {exc}") from exc + + def validate_package(package_dir: Path, manifest: HubPluginManifest) -> None: base = package_dir.resolve() if not base.is_dir(): @@ -41,6 +65,8 @@ def validate_package(package_dir: Path, manifest: HubPluginManifest) -> None: if not entry.exists(): raise ValueError(f"Missing entrypoint: {entrypoint}") + _validate_python_syntax(base, manifest) + for rel_path, expected in manifest.checksums.items(): if not expected: continue diff --git a/flocks/ingest/kafka/manager.py b/flocks/ingest/kafka/manager.py index 16acca4b6..9b799a0d6 100644 --- a/flocks/ingest/kafka/manager.py +++ b/flocks/ingest/kafka/manager.py @@ -42,6 +42,10 @@ from flocks.workflow.models import Workflow from flocks.workflow.runner import run_workflow from flocks.workflow.store import WorkflowStore +from flocks.workflow.tool_context import ( + build_workflow_tool_context, + cleanup_workflow_tool_context, +) from flocks.ingest.kafka.constants import WORKFLOW_KAFKA_CONFIG_PREFIX from flocks.workflow.triggers.compat import legacy_kafka_trigger_from_config @@ -297,7 +301,14 @@ async def start_all(self) -> None: if not workflow_id: continue if isinstance(data, dict) and data.get("enabled"): - await self.restart_workflow(workflow_id) + try: + await self.restart_workflow(workflow_id, startup=True) + except Exception as exc: + self._status[workflow_id] = {"state": "failed", "error": str(exc)} + log.warning( + "kafka.start_failed", + {"workflow_id": workflow_id, "error": str(exc)}, + ) async def stop_all(self) -> None: for workflow_id in list(self._tasks.keys()): @@ -360,7 +371,12 @@ async def stop_workflow(self, workflow_id: str) -> None: if workflow_id in self._status: self._status[workflow_id] = {"state": "stopped", "error": None} - async def restart_workflow(self, workflow_id: str) -> Dict[str, Any]: + async def restart_workflow( + self, + workflow_id: str, + *, + startup: bool = False, + ) -> Dict[str, Any]: """Restart the consumer and return its post-connect runtime status. Blocks until the consumer connects, the connection fails, or @@ -377,6 +393,21 @@ async def restart_workflow(self, workflow_id: str) -> Dict[str, Any]: self._status[workflow_id] = {"state": "stopped", "error": None} return {"state": "stopped", "error": None} + # Load and cache the workflow JSON once; avoids a disk read per message. + wf_data = read_workflow_from_fs(workflow_id) + if not wf_data: + err = "workflow_not_found" + if startup: + self._status[workflow_id] = {"state": "stopped", "error": err} + log.info("kafka.workflow_not_found_on_start", { + "workflow_id": workflow_id, + "action": "stale_config_skipped", + }) + return {"state": "stopped", "error": err} + self._status[workflow_id] = {"state": "failed", "error": err} + log.warning("kafka.workflow_not_found", {"workflow_id": workflow_id}) + return {"state": "failed", "error": err} + input_broker = str(data.get("inputBroker") or "").strip() input_topic = str(data.get("inputTopic") or "").strip() if not input_broker or not input_topic: @@ -385,13 +416,6 @@ async def restart_workflow(self, workflow_id: str) -> Dict[str, Any]: log.warning("kafka.config_incomplete", {"workflow_id": workflow_id}) return {"state": "failed", "error": err} - # Load and cache the workflow JSON once; avoids a disk read per message. - wf_data = read_workflow_from_fs(workflow_id) - if not wf_data: - err = "workflow_not_found" - self._status[workflow_id] = {"state": "failed", "error": err} - log.warning("kafka.workflow_not_found_on_start", {"workflow_id": workflow_id}) - return {"state": "failed", "error": err} workflow_json = wf_data.get("workflowJson") if not workflow_json: err = "workflow_json_missing" @@ -698,7 +722,12 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: input_keys=trigger_input_keys, ), ) + tool_context = None try: + tool_context = await build_workflow_tool_context( + workflow_id=workflow_id, + action_name=f"trigger:{trigger.type}", + ) result = await asyncio.to_thread( run_workflow, workflow=workflow_plan, @@ -707,6 +736,7 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: trace=False, execution_profile="high_frequency", on_step_complete=step_recorder.on_step_complete, + tool_context=tool_context, ) status, error_msg = resolve_execution_outcome(result) duration = time.time() - start_time @@ -754,6 +784,7 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: } ) finally: + await cleanup_workflow_tool_context(tool_context) try: await record_execution_result(workflow_id, exec_id, exec_data) except Exception as exc: diff --git a/flocks/ingest/syslog/manager.py b/flocks/ingest/syslog/manager.py index 3fccaed88..b81f88e3e 100644 --- a/flocks/ingest/syslog/manager.py +++ b/flocks/ingest/syslog/manager.py @@ -20,6 +20,10 @@ from flocks.workflow.models import Workflow from flocks.workflow.runner import run_workflow from flocks.workflow.store import WorkflowStore +from flocks.workflow.tool_context import ( + build_workflow_tool_context, + cleanup_workflow_tool_context, +) from flocks.ingest.syslog.constants import WORKFLOW_SYSLOG_CONFIG_PREFIX from flocks.ingest.syslog.listener import run_tcp_syslog_server, run_udp_syslog_server @@ -186,7 +190,17 @@ async def start_all(self) -> None: if not workflow_id: continue if isinstance(data, dict) and data.get("enabled"): - await self.restart_workflow(workflow_id) + try: + await self.restart_workflow(workflow_id, startup=True) + except Exception as exc: + self._listener_status[workflow_id] = { + "state": "failed", + "error": str(exc), + } + log.warning( + "syslog.start_failed", + {"workflow_id": workflow_id, "error": str(exc)}, + ) async def stop_all(self) -> None: for workflow_id in list(self._tasks.keys()): @@ -245,7 +259,12 @@ async def stop_workflow(self, workflow_id: str) -> None: if workflow_id in self._listener_status: self._listener_status[workflow_id] = {"state": "stopped", "error": None} - async def restart_workflow(self, workflow_id: str) -> Dict[str, Any]: + async def restart_workflow( + self, + workflow_id: str, + *, + startup: bool = False, + ) -> Dict[str, Any]: """Restart the listener and return its post-bind runtime status. This call blocks until the underlying socket either binds successfully, @@ -268,8 +287,15 @@ async def restart_workflow(self, workflow_id: str) -> Dict[str, Any]: wf_data = read_workflow_from_fs(workflow_id) if not wf_data: err = "workflow_not_found" + if startup: + self._listener_status[workflow_id] = {"state": "stopped", "error": err} + log.info("syslog.workflow_not_found_on_start", { + "workflow_id": workflow_id, + "action": "stale_config_skipped", + }) + return {"state": "stopped", "error": err} self._listener_status[workflow_id] = {"state": "failed", "error": err} - log.warning("syslog.workflow_not_found_on_start", {"workflow_id": workflow_id}) + log.warning("syslog.workflow_not_found", {"workflow_id": workflow_id}) return {"state": "failed", "error": err} workflow_json = wf_data.get("workflowJson") if not workflow_json: @@ -286,6 +312,10 @@ async def restart_workflow(self, workflow_id: str) -> Dict[str, Any]: self._listener_status[workflow_id] = {"state": "failed", "error": err} log.warning("syslog.workflow_plan_failed", {"workflow_id": workflow_id, "error": str(exc)}) return self.get_listener_status(workflow_id) + + 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) self._queues[workflow_id] = queue @@ -295,9 +325,6 @@ async def restart_workflow(self, workflow_id: str) -> Dict[str, Any]: ready = asyncio.Event() self._listener_ready[workflow_id] = ready - 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() self._listener_status[workflow_id] = { "state": "binding", "error": None, @@ -538,7 +565,12 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: ) start_time = time.time() trigger_meta = mapped_inputs.get("_flocks", {}).get("trigger", {}) + tool_context = None try: + tool_context = await build_workflow_tool_context( + workflow_id=workflow_id, + action_name=f"trigger:{trigger.type}", + ) result = await asyncio.to_thread( run_workflow, workflow=workflow_plan, @@ -547,6 +579,7 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: trace=False, execution_profile="high_frequency", on_step_complete=step_recorder.on_step_complete, + tool_context=tool_context, ) status, error_msg = resolve_execution_outcome(result) duration = time.time() - start_time @@ -594,6 +627,7 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: } ) finally: + await cleanup_workflow_tool_context(tool_context) try: await record_execution_result(workflow_id, exec_id, exec_data) except Exception as exc: diff --git a/flocks/mcp/client.py b/flocks/mcp/client.py index c2bf675ae..4a4bf5443 100644 --- a/flocks/mcp/client.py +++ b/flocks/mcp/client.py @@ -703,15 +703,7 @@ async def list_resources(self) -> List[McpResource]: Raises: RuntimeError: If not connected """ - try: - result = await self._submit_command("list_resources") - return result - except Exception as exc: - log.error("mcp.client.list_resources_error", { - "server": self.name, - "error": str(exc), - }) - raise + return await self._submit_command("list_resources") async def read_resource(self, uri: str) -> Any: """ @@ -789,11 +781,30 @@ async def _submit_command_on_owner_loop( command = _ClientCommand(action=action, payload=payload, response=response) await self._command_queue.put(command) - if owner_task.done() and not response.done(): - owner_error = self._owner_error or RuntimeError(f"Client not connected: {self.name}") - response.set_exception(owner_error) + try: + done, _ = await asyncio.wait( + {response, owner_task}, + return_when=asyncio.FIRST_COMPLETED, + ) + except asyncio.CancelledError: + response.cancel() + raise + + if response in done: + return await response + + response.cancel() + try: + owner_error = owner_task.exception() + except asyncio.CancelledError: + owner_error = None + + if owner_error is not None: + raise RuntimeError( + f"Connection lost: {self.name}: {_extract_root_cause(owner_error)}" + ) from owner_error - return await response + raise RuntimeError(f"Connection lost: {self.name}") async def _await_task(self, task: asyncio.Task[Any]) -> None: """Await a task safely from the owner loop or another loop.""" diff --git a/flocks/mcp/errors.py b/flocks/mcp/errors.py new file mode 100644 index 000000000..7849bf81a --- /dev/null +++ b/flocks/mcp/errors.py @@ -0,0 +1,35 @@ +"""Helpers for classifying MCP protocol errors.""" + +from __future__ import annotations + +from mcp.types import METHOD_NOT_FOUND + + +def is_method_not_found_error(exc: BaseException) -> bool: + """Return whether an exception chain contains JSON-RPC method-not-found.""" + return _contains_method_not_found(exc, seen=set()) + + +def _contains_method_not_found(exc: BaseException, *, seen: set[int]) -> bool: + exc_id = id(exc) + if exc_id in seen: + return False + seen.add(exc_id) + + error = getattr(exc, "error", None) + if getattr(error, "code", None) == METHOD_NOT_FOUND: + return True + + if isinstance(exc, BaseExceptionGroup): + if any(_contains_method_not_found(child, seen=seen) for child in exc.exceptions): + return True + + cause = exc.__cause__ + if cause is not None and _contains_method_not_found(cause, seen=seen): + return True + context = exc.__context__ + if context is not None and _contains_method_not_found(context, seen=seen): + return True + + message = str(exc).strip().lower() + return message == "method not found" or message.endswith(": method not found") diff --git a/flocks/mcp/server.py b/flocks/mcp/server.py index b7a20cd10..ae6dd0240 100644 --- a/flocks/mcp/server.py +++ b/flocks/mcp/server.py @@ -8,6 +8,7 @@ import time from typing import Dict, Optional, Any, List from flocks.mcp.client import McpClient +from flocks.mcp.errors import is_method_not_found_error from flocks.mcp.types import ( McpStatus, McpStatusInfo, @@ -226,10 +227,13 @@ async def _connect_and_register( resources = await client.list_resources() self._resources_cache[name] = resources except Exception as e: - log.warn("mcp.resources_unavailable", { - "server": name, - "error": str(e) - }) + if is_method_not_found_error(e): + log.info("mcp.resources_unsupported", {"server": name}) + else: + log.warn("mcp.resources_unavailable", { + "server": name, + "error": str(e) + }) self._resources_cache[name] = [] # 5. Register tools diff --git a/flocks/notifications/service.py b/flocks/notifications/service.py index 931fc2fc0..3cf3a02cd 100644 --- a/flocks/notifications/service.py +++ b/flocks/notifications/service.py @@ -206,7 +206,6 @@ async def list_active( *, user_id: str, locale: str | None = None, - current_version: str | None = None, ) -> list[NotificationResponse]: target_locale = _normalize_locale(locale) config_notifications = await cls._load_config_notifications() diff --git a/flocks/project/project.py b/flocks/project/project.py index 4be59d220..edfa1aa19 100644 --- a/flocks/project/project.py +++ b/flocks/project/project.py @@ -1,330 +1,668 @@ -""" -Project management module +"""User-scoped project registry backed by JSON files. -Handles project discovery, metadata, and lifecycle +Projects are UI/session groupings bound to existing directories. Legacy project +records in Storage are intentionally ignored; sessions that do not reference a +registered project remain ordinary sessions rather than becoming a project. """ +from __future__ import annotations + +import asyncio +import hashlib +import json import os import subprocess -from typing import Optional, Dict, Any, List +import sys +import tempfile +import time +import uuid +from contextlib import contextmanager from datetime import datetime -from pydantic import BaseModel, Field +from pathlib import Path +from typing import Any, ClassVar, Dict, Iterator, List, Literal, Optional, Tuple + +from pydantic import BaseModel, ConfigDict, Field -from flocks.storage.storage import Storage from flocks.utils.log import Log -from flocks.utils.id import Identifier log = Log.create(service="project") +# Legacy stored project ID used by ordinary sessions. It is not a ProjectInfo +# entry and never resolves to a worktree. +DEFAULT_PROJECT_ID = "default" +# Session-manager wire ID for the Tasks section; also not a project. +TASK_SESSION_GROUP_ID = "tasks" +DEFAULT_PROJECT_NAME = "默认" +_RESERVED_PROJECT_NAMES = {"default", DEFAULT_PROJECT_NAME.casefold()} + + +def _platform_file_lock(fd: int) -> None: + if sys.platform == "win32": # pragma: no cover - Windows only + import msvcrt + + msvcrt.locking(fd, msvcrt.LK_LOCK, 1) + else: + import fcntl + + fcntl.flock(fd, fcntl.LOCK_EX) + + +def _platform_file_unlock(fd: int) -> None: + if sys.platform == "win32": # pragma: no cover - Windows only + import msvcrt + + msvcrt.locking(fd, msvcrt.LK_UNLCK, 1) + else: + import fcntl + + fcntl.flock(fd, fcntl.LOCK_UN) + + +@contextmanager +def _registry_cross_process_lock(registry_path: Path) -> Iterator[None]: + """Serialize registry read-modify-write operations across processes.""" + + registry_path.parent.mkdir(parents=True, exist_ok=True) + lock_path = registry_path.with_suffix(".json.lock") + fd = os.open(str(lock_path), os.O_RDWR | os.O_CREAT, 0o600) + locked = False + try: + _platform_file_lock(fd) + locked = True + yield + finally: + try: + if locked: + _platform_file_unlock(fd) + finally: + os.close(fd) + + +class ProjectNameConflictError(ValueError): + """Raised when a project name is already registered for the user.""" + + +class ProjectPathConflictError(ValueError): + """Raised when a project directory is already registered.""" + + def __init__(self, project: "ProjectInfo") -> None: + super().__init__(f"Project directory is already registered as '{project.name}'") + self.project = project + + +class ProjectDeletionError(ValueError): + """Raised when a protected project cannot be deleted.""" + class ProjectIcon(BaseModel): - """Project icon configuration""" + """Project icon configuration.""" + url: Optional[str] = None override: Optional[str] = None color: Optional[str] = None class ProjectTime(BaseModel): - """Project time metadata""" + """Project time metadata.""" + created: int updated: int initialized: Optional[int] = None class ProjectInfo(BaseModel): - """Project information""" + """Project information returned by project and session APIs.""" + + model_config = ConfigDict(populate_by_name=True) + id: str worktree: str - vcs: Optional[str] = None # "git" or None + vcs: Optional[str] = None name: Optional[str] = None icon: Optional[ProjectIcon] = None time: ProjectTime sandboxes: List[str] = Field(default_factory=list) + is_default: bool = Field(False, alias="isDefault") + path_status: Literal["available", "missing", "unreadable"] = Field( + "available", + alias="pathStatus", + ) + session_count: int = Field(0, alias="sessionCount") + matched_session_count: int = Field(0, alias="matchedSessionCount") + last_activity_at: Optional[int] = Field(None, alias="lastActivityAt") + owner_user_id: Optional[str] = Field(None, alias="ownerUserID") + can_write: bool = Field(True, alias="canWrite") + can_delete: bool = Field(True, alias="canDelete") + is_shared: bool = Field(False, alias="isShared") + + +class ProjectRegistryEntry(BaseModel): + """Persisted project entry.""" + + model_config = ConfigDict(populate_by_name=True) + + id: str + name: str + worktree: str + created_at: int = Field(alias="createdAt") + updated_at: int = Field(alias="updatedAt") + owner_user_id: Optional[str] = Field(None, alias="ownerUserID") + shared_local: bool = Field(False, alias="sharedLocal") + + +class ProjectRegistry(BaseModel): + """Per-user project registry file.""" + + model_config = ConfigDict(populate_by_name=True) + + default_worktree: Optional[str] = Field(None, alias="defaultWorktree") + projects: List[ProjectRegistryEntry] = Field(default_factory=list) class Project: - """Project management namespace""" - - _current: Optional[ProjectInfo] = None - + """User-scoped registry of explicitly created projects.""" + + _lock = asyncio.Lock() + _session_stats_cache: ClassVar[ + Dict[Tuple[str, str], Tuple[float, Dict[str, Tuple[int, int, Optional[int]]]]] + ] = {} + _session_stats_cache_ttl_seconds = 30.0 + + @staticmethod + def _now_ms() -> int: + return int(datetime.now().timestamp() * 1000) + + @staticmethod + def _flocks_root() -> Path: + return Path(os.getenv("FLOCKS_ROOT", str(Path.home() / ".flocks"))).expanduser() + @classmethod - async def from_directory(cls, directory: str) -> Dict[str, Any]: - """ - Create or load project from directory - - Args: - directory: Directory path - - Returns: - Dict with 'project' and 'sandbox' keys - """ - log.info("from_directory", {"directory": directory}) - - # Find git repository - git_info = await cls._find_git_repo(directory) - - project_id = git_info["id"] - worktree = git_info["worktree"] - sandbox = git_info["sandbox"] - vcs = git_info["vcs"] - - # Try to load existing project - existing = await Storage.read(["project", project_id], ProjectInfo) - - if not existing: - # Create new project - existing = ProjectInfo( - id=project_id, - worktree=worktree, - vcs=vcs, - sandboxes=[], - time=ProjectTime( - created=int(datetime.now().timestamp() * 1000), - updated=int(datetime.now().timestamp() * 1000), - ) - ) - - # Update project info - result = ProjectInfo( - id=existing.id, - worktree=worktree, - vcs=vcs, - name=existing.name, - icon=existing.icon, - time=ProjectTime( - created=existing.time.created, - updated=int(datetime.now().timestamp() * 1000), - initialized=existing.time.initialized, - ), - sandboxes=existing.sandboxes.copy() if existing.sandboxes else [], - ) - - # Add sandbox if not in list - if sandbox != result.worktree and sandbox not in result.sandboxes: - result.sandboxes.append(sandbox) - - # Filter out non-existent sandboxes - result.sandboxes = [s for s in result.sandboxes if os.path.exists(s)] - - # Save project - await Storage.write(["project", project_id], result) - - cls._current = result - - return { - "project": result, - "sandbox": sandbox, - } - + def _owner_hash(cls, owner_id: str) -> str: + return hashlib.sha256(owner_id.encode("utf-8")).hexdigest()[:24] + @classmethod - async def _find_git_repo(cls, directory: str) -> Dict[str, Any]: - """ - Find git repository information - - Args: - directory: Starting directory - - Returns: - Dict with id, worktree, sandbox, vcs - """ - # Look for .git directory - current = os.path.abspath(directory) - git_dir = None - - while current != "/": - git_path = os.path.join(current, ".git") - if os.path.exists(git_path): - git_dir = git_path - break - parent = os.path.dirname(current) - if parent == current: - break - current = parent - - if not git_dir: - # No git repository found - return { - "id": "global", - "worktree": directory, - "sandbox": directory, - "vcs": None, - } - - sandbox = os.path.dirname(git_dir) - - # Try to get git root commit for ID + def registry_path(cls, owner_id: str) -> Path: + """Return the project registry path for a user.""" + + return cls._flocks_root() / "projects" / f"{cls._owner_hash(owner_id)}.json" + + @staticmethod + def _normalized_worktree(worktree: str) -> str: + return os.path.normcase(str(Path(worktree).expanduser().resolve(strict=True))) + + @staticmethod + def _directory_context(directory: str) -> Tuple[Path, Path, Optional[str]]: + path = Path(directory).expanduser().resolve() + worktree = path + vcs: Optional[str] = None try: result = subprocess.run( - ["git", "rev-list", "--max-parents=0", "--all"], - cwd=sandbox, + ["git", "rev-parse", "--show-toplevel"], + cwd=path, capture_output=True, text=True, timeout=5, + check=False, ) - - if result.returncode == 0: - roots = [line.strip() for line in result.stdout.split("\n") if line.strip()] - roots.sort() - - if roots: - project_id = roots[0] - - # Cache the ID - id_file = os.path.join(git_dir, "flocks") - try: - with open(id_file, "w") as f: - f.write(project_id) - except Exception: - pass - - # Get worktree - try: - result = subprocess.run( - ["git", "rev-parse", "--show-toplevel"], - cwd=sandbox, - capture_output=True, - text=True, - timeout=5, - ) - if result.returncode == 0: - worktree = result.stdout.strip() - else: - worktree = sandbox - except Exception: - worktree = sandbox - - return { - "id": project_id, - "worktree": worktree, - "sandbox": sandbox, - "vcs": "git", - } - except Exception as e: - log.warn("git.error", {"error": str(e)}) - - # Check for cached ID - id_file = os.path.join(git_dir, "flocks") - if os.path.exists(id_file): + if result.returncode == 0 and result.stdout.strip(): + worktree = Path(result.stdout.strip()).resolve() + vcs = "git" + except (OSError, subprocess.SubprocessError): + pass + return path, worktree, vcs + + @classmethod + def worktree_for_directory(cls, directory: str) -> str: + """Return a stable runtime scope for CLI sessions in a directory.""" + + _, worktree, _ = cls._directory_context(directory) + return os.path.normcase(str(worktree)) + + @staticmethod + def _path_status(worktree: str) -> Literal["available", "missing", "unreadable"]: + path = Path(worktree) + if not path.is_dir(): + return "missing" + if not os.access(path, os.R_OK | os.X_OK): + return "unreadable" + return "available" + + @classmethod + def allowed_roots(cls, default_worktree: Optional[str] = None) -> List[Path]: + """Return canonical roots available to the server-side folder picker.""" + + configured = os.getenv("FLOCKS_PROJECT_ROOTS", "").strip() + candidates = [Path(item).expanduser() for item in configured.split(os.pathsep) if item] + if not candidates: + candidates = [Path.home()] + if default_worktree: + candidates.append(Path(default_worktree).expanduser()) + + roots: List[Path] = [] + for candidate in candidates: try: - with open(id_file, "r") as f: - project_id = f.read().strip() - if project_id: - return { - "id": project_id, - "worktree": sandbox, - "sandbox": sandbox, - "vcs": "git", - } - except Exception: - pass - - # Fallback to global - return { - "id": "global", - "worktree": sandbox, - "sandbox": sandbox, - "vcs": "git", - } - + resolved = candidate.resolve(strict=True) + except (OSError, RuntimeError): + continue + if not resolved.is_dir(): + continue + if any(resolved == root or resolved.is_relative_to(root) for root in roots): + continue + roots = [root for root in roots if not root.is_relative_to(resolved)] + roots.append(resolved) + return sorted(roots, key=lambda item: str(item).casefold()) + @classmethod - async def list(cls) -> List[ProjectInfo]: - """ - List all projects - - Returns: - List of project info - """ + def validate_worktree( + cls, + worktree: str, + *, + default_worktree: Optional[str] = None, + create_if_missing: bool = False, + ) -> str: + """Validate and canonicalize a project directory.""" + + if not worktree or not worktree.strip(): + raise ValueError("Project directory cannot be empty") + requested_path = Path(worktree).expanduser() + if not requested_path.is_absolute(): + raise ValueError("Project directory must be an absolute path") + requested_path = Path(os.path.abspath(requested_path)) + roots = cls.allowed_roots(default_worktree) + + if create_if_missing and not requested_path.exists(): + ancestor = requested_path + missing_parts: List[str] = [] + while not ancestor.exists() and ancestor != ancestor.parent: + missing_parts.append(ancestor.name) + ancestor = ancestor.parent + try: + resolved_ancestor = ancestor.resolve(strict=True) + except (OSError, RuntimeError) as exc: + raise ValueError("Project directory cannot be resolved") from exc + target = resolved_ancestor.joinpath(*reversed(missing_parts)) + if roots and not any(target == root or target.is_relative_to(root) for root in roots): + raise ValueError("Project directory is outside the allowed roots") + try: + target.mkdir(parents=True, exist_ok=True) + except OSError as exc: + raise ValueError("Project directory could not be created") from exc + try: - projects = await Storage.list_by_prefix(["project"]) - result = [] - - for key, value in projects.items(): - if isinstance(value, dict): - try: - project = ProjectInfo(**value) - result.append(project) - except Exception as e: - log.warn("project.parse.error", {"key": key, "error": str(e)}) - - # Sort by updated time (newest first) - result.sort(key=lambda p: p.time.updated, reverse=True) - - return result - except Exception as e: - log.error("project.list.error", {"error": str(e)}) - return [] - + path = requested_path.resolve(strict=True) + except (OSError, RuntimeError) as exc: + raise ValueError("Project directory does not exist") from exc + if not path.is_dir(): + raise ValueError("Project path must be a directory") + if not os.access(path, os.R_OK | os.X_OK): + raise ValueError("Project directory is not readable") + + if roots and not any(path == root or path.is_relative_to(root) for root in roots): + raise ValueError("Project directory is outside the allowed roots") + return os.path.normcase(str(path)) + @classmethod - async def get(cls, project_id: str) -> Optional[ProjectInfo]: - """ - Get project by ID - - Args: - project_id: Project ID - - Returns: - Project info or None - """ + def _read_registry_file(cls, path: Path) -> Optional[ProjectRegistry]: + if not path.exists(): + return None try: - return await Storage.read(["project", project_id], ProjectInfo) + return ProjectRegistry.model_validate_json(path.read_text(encoding="utf-8")) + except (OSError, ValueError) as exc: + log.error("project.registry.read_failed", {"path": str(path), "error": str(exc)}) + return None + + @classmethod + def _read_registry( + cls, + owner_id: str, + ) -> ProjectRegistry: + path = cls.registry_path(owner_id) + registry = cls._read_registry_file(path) + if registry is not None: + # The pre-task/project split stored a virtual default worktree. + # Keep reading that field for compatibility but never use it. + registry.default_worktree = None + return registry + return ProjectRegistry(projects=[]) + + @classmethod + def _write_registry(cls, owner_id: str, registry: ProjectRegistry) -> None: + path = cls.registry_path(owner_id) + path.parent.mkdir(parents=True, exist_ok=True) + + fd, temporary_path = tempfile.mkstemp( + dir=str(path.parent), + prefix=f".{path.stem}.", + suffix=".tmp", + ) + try: + with os.fdopen(fd, "w", encoding="utf-8") as handle: + json.dump( + registry.model_dump(by_alias=True, exclude_none=True), + handle, + ensure_ascii=False, + indent=2, + ) + handle.write("\n") + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary_path, path) + try: + directory_fd = os.open(path.parent, os.O_RDONLY) + try: + os.fsync(directory_fd) + finally: + os.close(directory_fd) + except OSError: + pass except Exception: + try: + os.unlink(temporary_path) + except OSError: + pass + raise + + @classmethod + def _entry_to_info( + cls, + entry: ProjectRegistryEntry, + *, + can_write: bool = True, + ) -> ProjectInfo: + return ProjectInfo( + id=entry.id, + worktree=entry.worktree, + name=entry.name, + vcs="git" if (Path(entry.worktree) / ".git").exists() else None, + time=ProjectTime(created=entry.created_at, updated=entry.updated_at), + pathStatus=cls._path_status(entry.worktree), + ownerUserID=entry.owner_user_id, + canWrite=can_write, + canDelete=can_write, + isShared=entry.shared_local, + ) + + @classmethod + async def ensure_registry( + cls, + owner_id: str, + ) -> ProjectRegistry: + """Load a registry and create an empty one when needed.""" + + async with cls._lock: + path = cls.registry_path(owner_id) + with _registry_cross_process_lock(path): + registry = cls._read_registry(owner_id) + if not path.exists(): + cls._write_registry(owner_id, registry) + return registry + + @classmethod + async def create( + cls, + *, + owner_id: str, + name: Optional[str], + worktree: str, + ) -> ProjectInfo: + """Register an existing directory for a user.""" + + async with cls._lock: + with _registry_cross_process_lock(cls.registry_path(owner_id)): + registry = cls._read_registry(owner_id) + normalized_worktree = cls.validate_worktree( + worktree, + create_if_missing=True, + ) + + for entry in registry.projects: + if cls._normalized_worktree(entry.worktree) == normalized_worktree: + raise ProjectPathConflictError(cls._entry_to_info(entry)) + + normalized_name = (name or Path(normalized_worktree).name).strip() + if not normalized_name: + raise ValueError("Project name cannot be empty") + if normalized_name.casefold() in _RESERVED_PROJECT_NAMES: + raise ProjectNameConflictError("Project name is reserved") + if any(entry.name.strip().casefold() == normalized_name.casefold() for entry in registry.projects): + raise ProjectNameConflictError(f"Project name '{normalized_name}' already exists") + + now = cls._now_ms() + project_id = f"prj_{uuid.uuid4()}" + entry = ProjectRegistryEntry( + id=project_id, + name=normalized_name, + worktree=normalized_worktree, + createdAt=now, + updatedAt=now, + ownerUserID=owner_id, + ) + registry.projects.append(entry) + cls._write_registry(owner_id, registry) + cls.invalidate_session_stats(owner_id) + log.info("project.created", {"id": project_id, "owner": cls._owner_hash(owner_id)}) + return cls._entry_to_info(entry) + + @classmethod + async def list( + cls, + *, + owner_id: str, + ) -> List[ProjectInfo]: + registry = await cls.ensure_registry(owner_id) + projects = [cls._entry_to_info(entry) for entry in registry.projects] + projects.sort(key=lambda item: item.time.updated, reverse=True) + return projects + + @classmethod + def _all_registry_entries(cls) -> List[ProjectRegistryEntry]: + """Return valid entries across local user registries.""" + + registry_dir = cls._flocks_root() / "projects" + if not registry_dir.is_dir(): + return [] + entries: List[ProjectRegistryEntry] = [] + for path in registry_dir.glob("*.json"): + registry = cls._read_registry_file(path) + if registry is not None: + entries.extend(registry.projects) + return entries + + @classmethod + def shared_project_ids(cls) -> set[str]: + """Return IDs of all projects shared to local accounts.""" + + return {entry.id for entry in cls._all_registry_entries() if entry.shared_local} + + @classmethod + async def list_visible( + cls, + *, + owner_id: str, + ) -> List[ProjectInfo]: + """List owned projects plus projects shared by other local users.""" + + owned = await cls.list(owner_id=owner_id) + owned_ids = {project.id for project in owned} + shared = [ + cls._entry_to_info(entry, can_write=False) + for entry in cls._all_registry_entries() + if entry.shared_local and entry.owner_user_id != owner_id and entry.id not in owned_ids + ] + shared.sort(key=lambda item: item.time.updated, reverse=True) + return [*owned, *shared] + + @classmethod + def visible_project_ids(cls, owner_id: str) -> set[str]: + """Return project IDs visible to a local user.""" + + return cls.registered_project_ids(owner_id) | cls.shared_project_ids() + + @classmethod + def is_local_shared(cls, owner_id: Optional[str], project_id: Optional[str]) -> bool: + """Return whether an owner's registered project is locally shared.""" + + if not owner_id or not project_id or project_id == DEFAULT_PROJECT_ID: + return False + registry = cls._read_registry(owner_id) + return any(entry.id == project_id and entry.shared_local for entry in registry.projects) + + @classmethod + async def set_local_shared( + cls, + project_id: str, + *, + owner_id: str, + shared: bool, + ) -> ProjectInfo: + """Share or unshare an owned project with all local accounts.""" + + if project_id == DEFAULT_PROJECT_ID: + raise ProjectDeletionError("The default project cannot be shared") + async with cls._lock: + with _registry_cross_process_lock(cls.registry_path(owner_id)): + registry = cls._read_registry(owner_id) + entry = next((item for item in registry.projects if item.id == project_id), None) + if entry is None: + raise ValueError(f"Project {project_id} not found") + entry.owner_user_id = owner_id + entry.shared_local = shared + entry.updated_at = cls._now_ms() + cls._write_registry(owner_id, registry) + cls.invalidate_session_stats() + log.info( + "project.sharing.updated", + {"id": project_id, "owner": cls._owner_hash(owner_id), "shared": shared}, + ) + return cls._entry_to_info(entry) + + @classmethod + async def get( + cls, + project_id: str, + *, + owner_id: str, + ) -> Optional[ProjectInfo]: + registry = await cls.ensure_registry(owner_id) + entry = next((item for item in registry.projects if item.id == project_id), None) + return cls._entry_to_info(entry) if entry else None + + @classmethod + async def update(cls, project_id: str, *, owner_id: str, name: str) -> ProjectInfo: + """Rename a registered project.""" + + if project_id == DEFAULT_PROJECT_ID: + raise ProjectDeletionError("The default project cannot be renamed") + normalized_name = name.strip() + if not normalized_name: + raise ValueError("Project name cannot be empty") + if normalized_name.casefold() in _RESERVED_PROJECT_NAMES: + raise ProjectNameConflictError("Project name is reserved") + + async with cls._lock: + with _registry_cross_process_lock(cls.registry_path(owner_id)): + registry = cls._read_registry(owner_id) + entry = next((item for item in registry.projects if item.id == project_id), None) + if entry is None: + raise ValueError(f"Project {project_id} not found") + if any( + item.id != project_id and item.name.strip().casefold() == normalized_name.casefold() + for item in registry.projects + ): + raise ProjectNameConflictError(f"Project name '{normalized_name}' already exists") + + entry.name = normalized_name + entry.updated_at = cls._now_ms() + cls._write_registry(owner_id, registry) + cls.invalidate_session_stats(owner_id) + log.info("project.updated", {"id": project_id}) + return cls._entry_to_info(entry) + + @classmethod + async def delete(cls, project_id: str, *, owner_id: str) -> bool: + """Remove a project registration without changing sessions or files.""" + + if project_id == DEFAULT_PROJECT_ID: + raise ProjectDeletionError("The default project cannot be deleted") + async with cls._lock: + with _registry_cross_process_lock(cls.registry_path(owner_id)): + registry = cls._read_registry(owner_id) + remaining = [entry for entry in registry.projects if entry.id != project_id] + if len(remaining) == len(registry.projects): + raise ValueError(f"Project {project_id} not found") + registry.projects = remaining + cls._write_registry(owner_id, registry) + cls.invalidate_session_stats(owner_id) + log.info("project.deleted", {"id": project_id}) + return True + + @classmethod + def registered_project_ids(cls, owner_id: str) -> set[str]: + registry = cls._read_registry(owner_id) + return {entry.id for entry in registry.projects} + + @classmethod + def effective_project_id(cls, owner_id: str, stored_project_id: Optional[str]) -> str: + if stored_project_id and stored_project_id in cls.visible_project_ids(owner_id): + return stored_project_id + return TASK_SESSION_GROUP_ID + + @classmethod + def get_session_stats_cache( + cls, + owner_id: str, + search: str, + ) -> Optional[Dict[str, Tuple[int, int, Optional[int]]]]: + """Return a short-lived copy of cached project session statistics.""" + + key = (owner_id, search.casefold()) + cached = cls._session_stats_cache.get(key) + if cached is None: + return None + cached_at, stats = cached + if time.monotonic() - cached_at >= cls._session_stats_cache_ttl_seconds: + cls._session_stats_cache.pop(key, None) return None - + return dict(stats) + @classmethod - async def update(cls, project_id: str, **kwargs) -> ProjectInfo: - """ - Update project - - Args: - project_id: Project ID - **kwargs: Fields to update - - Returns: - Updated project info - """ - project = await cls.get(project_id) - if not project: - raise ValueError(f"Project {project_id} not found") - - # Update fields - update_data = project.model_dump() - - for key, value in kwargs.items(): - if key in update_data and value is not None: - if key == "icon" and isinstance(value, dict): - update_data["icon"] = ProjectIcon(**value) - else: - update_data[key] = value - - # Update timestamp - update_data["time"]["updated"] = int(datetime.now().timestamp() * 1000) - - updated_project = ProjectInfo(**update_data) - await Storage.write(["project", project_id], updated_project) - - if cls._current and cls._current.id == project_id: - cls._current = updated_project - - log.info("project.updated", {"id": project_id}) - - return updated_project - + def set_session_stats_cache( + cls, + owner_id: str, + search: str, + stats: Dict[str, Tuple[int, int, Optional[int]]], + ) -> None: + """Cache project session statistics for an owner and search term.""" + + cls._session_stats_cache[(owner_id, search.casefold())] = ( + time.monotonic(), + dict(stats), + ) + @classmethod - def current(cls) -> Optional[ProjectInfo]: - """ - Get current project - - Returns: - Current project info or None - """ - return cls._current - + def invalidate_session_stats(cls, owner_id: Optional[str] = None) -> None: + """Invalidate cached project counts after session or registry changes.""" + + if owner_id is None: + cls._session_stats_cache.clear() + return + for key in [key for key in cls._session_stats_cache if key[0] == owner_id]: + cls._session_stats_cache.pop(key, None) + @classmethod - def set_current(cls, project: ProjectInfo) -> None: - """ - Set current project - - Args: - project: Project info - """ - cls._current = project + async def from_directory(cls, directory: str) -> Dict[str, Any]: + """Build an ephemeral runtime context without persisting a project.""" + + path, worktree, vcs = cls._directory_context(directory) + + now = cls._now_ms() + project = ProjectInfo( + id=DEFAULT_PROJECT_ID, + worktree=str(worktree), + vcs=vcs, + name=DEFAULT_PROJECT_NAME, + time=ProjectTime(created=now, updated=now), + isDefault=True, + pathStatus=cls._path_status(str(path)), + ) + return {"project": project, "sandbox": str(worktree)} diff --git a/flocks/provider/catalog.json b/flocks/provider/catalog.json index 363d1408d..d6af63c94 100644 --- a/flocks/provider/catalog.json +++ b/flocks/provider/catalog.json @@ -52,6 +52,33 @@ "THREATBOOK_CN_LLM_API_KEY" ], "models": { + "kimi-k2.7-code": { + "name": "kimi-k2.7-code", + "family": "kimi-k2.7-code", + "capabilities": { + "supports_tools": true, + "supports_vision": true, + "supports_reasoning": true, + "interleaved": { + "field": "reasoning_content", + "echo": "tool_calls", + "placeholder": " ", + "cross_provider_policy": "placeholder" + }, + "supports_streaming": true + }, + "limits": { + "context_window": 256000, + "max_input_tokens": 224000, + "max_output_tokens": 16000 + }, + "pricing": { + "input": 6.5, + "output": 27.0, + "cache_read": 1.3, + "currency": "CNY" + } + }, "minimax-m3": { "name": "minimax-m3", "family": "minimax", @@ -262,6 +289,33 @@ "THREATBOOK_IO_LLM_API_KEY" ], "models": { + "kimi-k2.7-code": { + "name": "kimi-k2.7-code", + "family": "kimi-k2.7-code", + "capabilities": { + "supports_tools": true, + "supports_vision": true, + "supports_reasoning": true, + "interleaved": { + "field": "reasoning_content", + "echo": "tool_calls", + "placeholder": " ", + "cross_provider_policy": "placeholder" + }, + "supports_streaming": true + }, + "limits": { + "context_window": 256000, + "max_input_tokens": 224000, + "max_output_tokens": 16000 + }, + "pricing": { + "input": 6.5, + "output": 27.0, + "cache_read": 1.3, + "currency": "CNY" + } + }, "minimax-m3": { "name": "minimax-m3", "family": "minimax", @@ -1016,47 +1070,6 @@ "DEEPSEEK_API_KEY" ], "models": { - "deepseek-chat": { - "name": "DeepSeek Chat", - "family": "deepseek-v3", - "capabilities": { - "supports_tools": true, - "supports_streaming": true - }, - "limits": { - "context_window": 128000, - "max_output_tokens": 128000 - }, - "pricing": { - "input": 2.0, - "output": 3.0, - "currency": "CNY" - } - }, - "deepseek-reasoner": { - "name": "DeepSeek Reasoner", - "family": "deepseek-r1", - "capabilities": { - "supports_tools": true, - "supports_reasoning": true, - "interleaved": { - "field": "reasoning_content", - "echo": "tool_calls", - "placeholder": " ", - "cross_provider_policy": "placeholder" - }, - "supports_streaming": true - }, - "limits": { - "context_window": 128000, - "max_output_tokens": 128000 - }, - "pricing": { - "input": 4.0, - "output": 16.0, - "currency": "CNY" - } - }, "deepseek-v4-flash": { "name": "DeepSeek V4 Flash", "family": "deepseek-v4", diff --git a/flocks/provider/interleaved.py b/flocks/provider/interleaved.py index ad1b869b5..727628a4e 100644 --- a/flocks/provider/interleaved.py +++ b/flocks/provider/interleaved.py @@ -38,7 +38,6 @@ "deepseek-v4", "deepseek-v4-pro", "deepseek-v4-flash", - "deepseek-reasoner", "deepseek-r1", "reasoner", "r1-0528", diff --git a/flocks/security/channel_secrets.py b/flocks/security/channel_secrets.py index 48e0f48cf..99feb4107 100644 --- a/flocks/security/channel_secrets.py +++ b/flocks/security/channel_secrets.py @@ -18,6 +18,9 @@ SENSITIVE_FIELD_NAMES: Set[str] = { "appSecret", "botToken", + "appToken", + "slackAppToken", + "slackBotToken", "secret", "clientSecret", "apiKey", diff --git a/flocks/server/app.py b/flocks/server/app.py index 78271d130..d60990514 100644 --- a/flocks/server/app.py +++ b/flocks/server/app.py @@ -338,6 +338,7 @@ def _sync_catalog_models_phase() -> None: try: from flocks.server.routes.workflow import ( reconcile_published_workflow_api_services, + start_workflow_api_health_monitor, sync_workflows_from_filesystem, ) @@ -346,6 +347,7 @@ async def _sync_workflows_phase() -> None: log.info("workflow.sync.done", {"imported": imported}) api_services = await reconcile_published_workflow_api_services() log.info("workflow.api_services.reconciled", api_services) + app.state.startup_background_tasks.append(start_workflow_api_health_monitor()) _schedule_startup_phase(app, log, "workflow.sync_filesystem", _sync_workflows_phase) except Exception as e: @@ -472,18 +474,6 @@ async def _delayed_trigger_runtime_start() -> None: except Exception as e: log.warning("workflow.trigger_runtime.start_failed", {"error": str(e)}) - try: - from flocks.updater.updater import recover_upgrade_state - - await _run_startup_phase( - log, - "updater.recover_upgrade_state", - lambda: asyncio.to_thread(recover_upgrade_state), - ) - log.info("updater.recovery.checked") - except Exception as e: - log.warning("updater.recovery.failed", {"error": str(e)}) - blocking_startup_ms = int((time.perf_counter() - startup_started_at) * 1000) log.info("server.startup.ready", { "blocking_duration_ms": blocking_startup_ms, @@ -1111,7 +1101,8 @@ async def validation_exception_handler(request: Request, exc: RequestValidationE @app.exception_handler(StarletteHTTPException) async def http_exception_handler(request: Request, exc: StarletteHTTPException): """Handle HTTP exceptions""" - log.error("http.error", { + http_log = log.warn if 400 <= exc.status_code < 500 else log.error + http_log("http.error", { "path": request.url.path, "status": exc.status_code, "detail": exc.detail, diff --git a/flocks/server/routes/channel.py b/flocks/server/routes/channel.py index 02349aaff..dfa0facc1 100644 --- a/flocks/server/routes/channel.py +++ b/flocks/server/routes/channel.py @@ -4,10 +4,11 @@ from __future__ import annotations +import json from typing import Optional from fastapi import APIRouter, HTTPException, Request -from fastapi.responses import JSONResponse +from fastapi.responses import JSONResponse, Response from pydantic import BaseModel from flocks.channel.gateway.manager import default_manager @@ -212,6 +213,27 @@ async def list_channels(): ] +@router.get("/slack/manifest") +async def slack_manifest(): + """Return a Slack app manifest for creating the Flocks Slack app.""" + from flocks.channel.builtin.slack.manifest import build_slack_app_manifest + + return build_slack_app_manifest() + + +@router.get("/slack/manifest/download") +async def slack_manifest_download(): + """Download a Slack app manifest for creating the Flocks Slack app.""" + from flocks.channel.builtin.slack.manifest import build_slack_app_manifest + + payload = json.dumps(build_slack_app_manifest(), indent=2, ensure_ascii=False) + return Response( + content=payload + "\n", + media_type="application/json", + headers={"Content-Disposition": 'attachment; filename="flocks-slack-manifest.json"'}, + ) + + @router.post("/{channel_id}/record-inbound") async def record_inbound(channel_id: str): """Notify the gateway that a message was received on this channel. diff --git a/flocks/server/routes/config.py b/flocks/server/routes/config.py index 7f403f92c..58aee09ea 100644 --- a/flocks/server/routes/config.py +++ b/flocks/server/routes/config.py @@ -35,6 +35,36 @@ log = Log.create(service="routes.config") +def _channel_allow_from_deletion_ids(config_data: Dict[str, Any]) -> set[str]: + """Return channel IDs whose PATCH explicitly removes allowFrom.""" + channels = config_data.get("channels") + if not isinstance(channels, dict): + return set() + + return { + channel_id + for channel_id, channel_cfg in channels.items() + if isinstance(channel_cfg, dict) + and "allowFrom" in channel_cfg + and channel_cfg.get("allowFrom") is None + } + + +def _normalize_slack_dm_policy(config_data: Dict[str, Any]) -> None: + """Keep Slack allowFrom and dmPolicy aligned for DM access control.""" + channels = config_data.get("channels") + if not isinstance(channels, dict): + return + slack = channels.get("slack") + if not isinstance(slack, dict) or "allowFrom" not in slack: + return + allow_from = slack.get("allowFrom") + if isinstance(allow_from, list) and len(allow_from) > 0: + slack["dmPolicy"] = "allowlist" + else: + slack["dmPolicy"] = "open" + + def _build_model_from_config( provider_id: str, model_id: str, @@ -141,6 +171,17 @@ class UIConfigUpdateRequest(BaseModel): display_name: Optional[str] = Field(None, alias="displayName") +class ToolFailurePreference(BaseModel): + """Repeated tool-failure preference exposed to the WebUI.""" + + model_config = {"populate_by_name": True} + + disable_on_repeated_failure: bool = Field( + ..., + alias="disableOnRepeatedFailure", + ) + + DEFAULT_UI_DISPLAY_NAME = "Flocks" DEFAULT_UI_PRO_DISPLAY_NAME = "Flocks Pro" FAVICON_MAX_BYTES = 512 * 1024 @@ -400,6 +441,12 @@ def _persist_ui_section(data: Dict[str, Any], ui_section: Dict[str, Any]) -> Non ConfigWriter._write_raw(data) +def _effective_tool_failure_preference(config: ConfigInfoModel) -> bool: + if config.tool_failure is None: + return True + return config.tool_failure.disable_on_repeated_failure + + @router.get("/ui-display", response_model=UIDisplayResponse, summary="Get public UI display name") async def get_ui_display() -> UIDisplayResponse: """Return only the effective WebUI display name for public screens.""" @@ -521,6 +568,47 @@ async def reset_ui_favicon() -> UIDisplayResponse: return await get_ui_display() +@router.get( + "/tool-failure", + response_model=ToolFailurePreference, + summary="Get repeated tool-failure preference", +) +async def get_tool_failure_preference() -> ToolFailurePreference: + """Return whether repeated identical failures automatically disable tools.""" + try: + config = await Config.get() + return ToolFailurePreference( + disableOnRepeatedFailure=_effective_tool_failure_preference(config) + ) + except Exception as e: + log.error("config.tool_failure.get.error", {"error": str(e)}) + raise HTTPException(status_code=500, detail=str(e)) + + +@router.patch( + "/tool-failure", + response_model=ToolFailurePreference, + summary="Update repeated tool-failure preference", +) +async def update_tool_failure_preference( + request: ToolFailurePreference, +) -> ToolFailurePreference: + """Update only the repeated-failure switch in flocks.json.""" + try: + data = ConfigWriter._read_raw() + existing = data.get("toolFailure", data.get("tool_failure", {})) + section = dict(existing) if isinstance(existing, dict) else {} + section.pop("disable_on_repeated_failure", None) + section["disableOnRepeatedFailure"] = request.disable_on_repeated_failure + data.pop("tool_failure", None) + data["toolFailure"] = section + ConfigWriter._write_raw(data) + return await get_tool_failure_preference() + except Exception as e: + log.error("config.tool_failure.update.error", {"error": str(e)}) + raise HTTPException(status_code=400, detail=str(e)) + + @router.get("", summary="Get configuration") async def get_config() -> Dict[str, Any]: """ @@ -554,6 +642,9 @@ async def update_config(config_data: Dict[str, Any]) -> Dict[str, Any]: flocks.json, so that plaintext secrets never land in that file. """ try: + channel_allow_from_deletions = _channel_allow_from_deletion_ids(config_data) + _normalize_slack_dm_policy(config_data) + # Extract channel sensitive fields into .secret.json before persisting if "channels" in config_data and isinstance(config_data.get("channels"), dict): from flocks.security.channel_secrets import extract_channel_secrets @@ -563,7 +654,10 @@ async def update_config(config_data: Dict[str, Any]) -> Dict[str, Any]: config = ConfigInfoModel.model_validate(config_data) # Update project config - await Config.update(config) + await Config.update( + config, + channel_allow_from_deletions=channel_allow_from_deletions, + ) # Clear cache to reload Config.clear_cache() diff --git a/flocks/server/routes/console_upgrade.py b/flocks/server/routes/console_upgrade.py index 107140f35..93b74dfb2 100644 --- a/flocks/server/routes/console_upgrade.py +++ b/flocks/server/routes/console_upgrade.py @@ -525,11 +525,18 @@ async def _run_auto_upgrade_install(record: dict[str, Any]) -> dict[str, Any]: final_stage = "" final_message = "" - async for progress in perform_pro_bundle_install(restart=False): + async for progress in perform_pro_bundle_install(restart=True): final_stage = progress.stage final_message = progress.message if progress.stage == "error": raise ValueError(progress.message) + if progress.stage == "restarting": + details["auto_install_result"] = "restarting" + details["auto_install_message"] = progress.message + record["updated_at"] = datetime.now(UTC).isoformat() + request_id = str(record.get("request_id") or "").strip() + if request_id: + await Storage.set(_request_key(request_id), record, "json") await _maybe_activate_pro_license(record, allow_fallback=False) await _maybe_refresh_pro_license(record) diff --git a/flocks/server/routes/notifications.py b/flocks/server/routes/notifications.py index 709a9240f..0792abd3b 100644 --- a/flocks/server/routes/notifications.py +++ b/flocks/server/routes/notifications.py @@ -26,13 +26,11 @@ async def list_active_notifications( request: Request, locale: str | None = None, - current_version: str | None = None, ) -> list[NotificationResponse]: user = require_user(request) return await NotificationService.list_active( user_id=user.id, locale=locale, - current_version=current_version, ) diff --git a/flocks/server/routes/onboarding.py b/flocks/server/routes/onboarding.py index a94ba41c3..7ccb2a380 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": "minimax-m2.7", + "threatbook_default_model_id": "kimi-k2.7-code", "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": "minimax-m2.7", + "threatbook_default_model_id": "kimi-k2.7-code", "threatbook_api_service_id": "threatbook-io", "threatbook_mcp_name": None, "threatbook_mcp_url": None, diff --git a/flocks/server/routes/project.py b/flocks/server/routes/project.py index e03f5dd45..9143f1f0d 100644 --- a/flocks/server/routes/project.py +++ b/flocks/server/routes/project.py @@ -1,102 +1,328 @@ -""" -Project management routes +"""Project registry routes.""" -Routes for project listing, retrieval, and updates -""" +from __future__ import annotations +import os +from pathlib import Path from typing import List, Optional -from fastapi import APIRouter, HTTPException + +from fastapi import APIRouter, HTTPException, Query, Request from pydantic import BaseModel -from flocks.project.project import Project, ProjectInfo, ProjectIcon +from flocks.auth.context import AuthUser +from flocks.project.project import ( + Project, + ProjectDeletionError, + ProjectInfo, + ProjectNameConflictError, + ProjectPathConflictError, +) +from flocks.server.auth import require_user +from flocks.session.policy import SessionPolicy +from flocks.session.session import Session from flocks.utils.log import Log router = APIRouter() log = Log.create(service="routes.project") +_MANAGER_CATEGORIES = {"user", "workflow", "entity-config"} + class ProjectUpdateRequest(BaseModel): - """Project update request""" + """Editable project properties.""" + + name: str + + +class ProjectCreateRequest(BaseModel): + """Register an existing project directory.""" + + worktree: str name: Optional[str] = None - icon: Optional[ProjectIcon] = None -@router.get("/", response_model=List[ProjectInfo], summary="List all projects") -async def list_projects(): - """ - List all projects - - Get a list of projects that have been opened with Flocks. - """ +class FolderEntry(BaseModel): + """Directory row returned by the server-side folder picker.""" + + name: str + path: str + + +class FolderBrowserResponse(BaseModel): + """Directory listing constrained to configured project roots.""" + + path: str + parent: Optional[str] + roots: List[FolderEntry] + entries: List[FolderEntry] + + +async def _list_project_summaries(user: AuthUser, search: Optional[str]) -> List[ProjectInfo]: + owner_id = user.id + projects = await Project.list_visible(owner_id=owner_id) + shared_project_ids = Project.shared_project_ids() + term = search.strip().casefold() if search else None + stats = Project.get_session_stats_cache(owner_id, term or "") + if stats is None: + registered_ids = {project.id for project in projects} + counts = {project.id: 0 for project in projects} + matched_counts = {project.id: 0 for project in projects} + last_activity: dict[str, int] = {} + + for session in await Session.list_all_unfiltered(): + if not SessionPolicy.can_read( + session, + user, + shared_project_ids=shared_project_ids, + ): + continue + metadata = session.metadata if isinstance(session.metadata, dict) else {} + if metadata.get("hideFromSessionManager"): + continue + if session.parent_id or session.category not in _MANAGER_CATEGORIES: + continue + if session.project_id not in registered_ids: + continue + project_id = session.project_id + counts[project_id] = counts.get(project_id, 0) + 1 + if term is None or term in session.title.casefold(): + matched_counts[project_id] = matched_counts.get(project_id, 0) + 1 + last_activity[project_id] = max( + last_activity.get(project_id, 0), + session.time.updated, + ) + stats = { + project.id: ( + counts.get(project.id, 0), + matched_counts.get(project.id, 0), + last_activity.get(project.id), + ) + for project in projects + } + Project.set_session_stats_cache(owner_id, term or "", stats) + + return [ + project.model_copy( + update={ + "session_count": stats.get(project.id, (0, 0, None))[0], + "matched_session_count": stats.get(project.id, (0, 0, None))[1], + "last_activity_at": stats.get(project.id, (0, 0, None))[2], + }, + ) + for project in projects + ] + + +@router.get("", response_model=List[ProjectInfo], include_in_schema=False) +@router.get("/", response_model=List[ProjectInfo], summary="List projects") +async def list_projects(request: Request, search: Optional[str] = Query(None)): + """List user-registered project folders.""" + + user = require_user(request) + try: + return await _list_project_summaries(user, search) + except Exception as exc: + log.error("project.list.error", {"error": str(exc)}) + raise HTTPException(status_code=500, detail=str(exc)) from exc + + +@router.post("", response_model=ProjectInfo, include_in_schema=False) +@router.post("/", response_model=ProjectInfo, summary="Create project") +async def create_project(request: Request, payload: ProjectCreateRequest): + """Register an existing directory without writing project database rows.""" + + user = require_user(request) try: - projects = await Project.list() - return projects - except Exception as e: - log.error("project.list.error", {"error": str(e)}) - raise HTTPException(status_code=500, detail=str(e)) + return await Project.create( + owner_id=user.id, + name=payload.name, + worktree=payload.worktree, + ) + except ProjectPathConflictError as exc: + return exc.project + except ProjectNameConflictError as exc: + raise HTTPException(status_code=409, detail=str(exc)) from exc + except ValueError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + except Exception as exc: + log.error("project.create.error", {"error": str(exc)}) + raise HTTPException(status_code=500, detail=str(exc)) from exc @router.get("/current", response_model=ProjectInfo, summary="Get current project") -async def get_current_project(): - """ - Get current project - - Retrieve the currently active project that Flocks is working with. - """ - current = Project.current() - if not current: - raise HTTPException(status_code=404, detail="No current project") - - return current - - -@router.patch("/{project_id}", response_model=ProjectInfo, summary="Update project") -async def update_project(project_id: str, update: ProjectUpdateRequest): - """ - Update project - - Update project properties such as name, icon and color. - """ +async def get_current_project(request: Request, project_id: Optional[str] = Query(None, alias="projectID")): + """Return a requested project or the current runtime context.""" + + user = require_user(request) + if project_id: + project = await Project.get(project_id, owner_id=user.id) + else: + from flocks.project.instance import Instance + + project = Instance.get_project() + if project is None: + raise HTTPException(status_code=404, detail="Current project is unavailable") + return project + + +@router.get("/folders", response_model=FolderBrowserResponse, summary="Browse project folders") +async def browse_project_folders(request: Request, path: Optional[str] = Query(None)): + """List directories under configured roots without exposing file contents.""" + + require_user(request) + roots = Project.allowed_roots() + if not roots: + raise HTTPException(status_code=400, detail="No project roots are available") + try: - # Get existing project - project = await Project.get(project_id) - if not project: - raise HTTPException(status_code=404, detail=f"Project {project_id} not found") - - # Update fields - update_data = {} - if update.name is not None: - update_data["name"] = update.name - if update.icon is not None: - update_data["icon"] = update.icon - - updated_project = await Project.update(project_id, **update_data) - - log.info("project.updated", {"id": project_id}) - - return updated_project - except HTTPException: - raise - except Exception as e: - log.error("project.update.error", {"error": str(e), "id": project_id}) - raise HTTPException(status_code=400, detail=str(e)) + if path: + current = Path(path).expanduser().resolve(strict=True) + else: + home = Path.home().resolve(strict=True) + current = ( + home + if any(home == root or home.is_relative_to(root) for root in roots) + else roots[0] + ) + except (OSError, RuntimeError) as exc: + raise HTTPException(status_code=400, detail="Directory does not exist") from exc + if not current.is_dir(): + raise HTTPException(status_code=400, detail="Path must be a directory") + containing_root = next( + (root for root in roots if current == root or current.is_relative_to(root)), + None, + ) + if containing_root is None: + raise HTTPException(status_code=403, detail="Directory is outside the allowed roots") + entries: List[FolderEntry] = [] + try: + children = sorted( + (child for child in current.iterdir() if child.is_dir()), + key=lambda child: (child.name.startswith("."), child.name.casefold()), + ) + for child in children: + try: + resolved = child.resolve(strict=True) + except (OSError, RuntimeError): + continue + if not (resolved == containing_root or resolved.is_relative_to(containing_root)): + continue + if not os.access(resolved, os.R_OK | os.X_OK): + continue + entries.append(FolderEntry(name=child.name, path=str(resolved))) + except PermissionError as exc: + raise HTTPException(status_code=403, detail="Directory is not readable") from exc -@router.get("/{project_id}", response_model=ProjectInfo, summary="Get project") -async def get_project(project_id: str): - """ - Get project - - Retrieve information about a specific project. - """ + parent = current.parent + parent_value = ( + str(parent) + if current != containing_root and (parent == containing_root or parent.is_relative_to(containing_root)) + else None + ) + return FolderBrowserResponse( + path=str(current), + parent=parent_value, + roots=[FolderEntry(name=root.name or str(root), path=str(root)) for root in roots], + entries=entries, + ) + + +@router.patch("/{project_id}", response_model=ProjectInfo, summary="Rename project") +async def update_project(project_id: str, payload: ProjectUpdateRequest, request: Request): + user = require_user(request) try: - project = await Project.get(project_id) - if not project: - raise HTTPException(status_code=404, detail=f"Project {project_id} not found") - - return project - except HTTPException: - raise - except Exception as e: - log.error("project.get.error", {"error": str(e), "id": project_id}) - raise HTTPException(status_code=500, detail=str(e)) + return await Project.update(project_id, owner_id=user.id, name=payload.name) + except ProjectNameConflictError as exc: + raise HTTPException(status_code=409, detail=str(exc)) from exc + except ProjectDeletionError as exc: + raise HTTPException(status_code=409, detail=str(exc)) from exc + except ValueError as exc: + raise HTTPException(status_code=404, detail=str(exc)) from exc + + +@router.post( + "/{project_id}/share-local", + response_model=ProjectInfo, + summary="Share project locally", +) +async def share_project_local(project_id: str, request: Request): + """Share a project and all of its sessions as read-only.""" + + user = require_user(request) + try: + return await Project.set_local_shared( + project_id, + owner_id=user.id, + shared=True, + ) + except ProjectDeletionError as exc: + raise HTTPException(status_code=409, detail=str(exc)) from exc + except ValueError as exc: + raise HTTPException(status_code=404, detail=str(exc)) from exc + + +@router.post( + "/{project_id}/unshare-local", + response_model=ProjectInfo, + summary="Unshare project locally", +) +async def unshare_project_local(project_id: str, request: Request): + """Cancel local sharing for a project and all of its sessions.""" + + user = require_user(request) + try: + return await Project.set_local_shared( + project_id, + owner_id=user.id, + shared=False, + ) + except ProjectDeletionError as exc: + raise HTTPException(status_code=409, detail=str(exc)) from exc + except ValueError as exc: + raise HTTPException(status_code=404, detail=str(exc)) from exc + + +@router.delete("/{project_id}", response_model=bool, summary="Delete project") +async def delete_project(project_id: str, request: Request): + """Remove a project and its sessions while preserving project files.""" + + user = require_user(request) + try: + if await Project.get(project_id, owner_id=user.id) is None: + raise ValueError(f"Project {project_id} not found") + + project_sessions = [ + session + for session in await Session.list_all_unfiltered() + if session.project_id == project_id + ] + if any(not SessionPolicy.can_delete(session, user) for session in project_sessions): + raise HTTPException( + status_code=403, + detail="Only session owners can delete project sessions", + ) + + from flocks.server.routes.session import delete_session_for_user + + for session in project_sessions: + if await Session.get(project_id, session.id) is not None: + await delete_session_for_user(session.id, user) + + return await Project.delete(project_id, owner_id=user.id) + except ProjectDeletionError as exc: + raise HTTPException(status_code=409, detail=str(exc)) from exc + except ValueError as exc: + raise HTTPException(status_code=404, detail=str(exc)) from exc + + +@router.get("/{project_id}", response_model=ProjectInfo, summary="Get project") +async def get_project(project_id: str, request: Request): + user = require_user(request) + project = await Project.get( + project_id, + owner_id=user.id, + ) + if project is None: + raise HTTPException(status_code=404, detail=f"Project {project_id} not found") + return project diff --git a/flocks/server/routes/provider.py b/flocks/server/routes/provider.py index 79e4d7e37..9ec50d8fe 100644 --- a/flocks/server/routes/provider.py +++ b/flocks/server/routes/provider.py @@ -13,13 +13,15 @@ import time from pathlib import Path from typing import Any, Dict, List, Optional -from fastapi import APIRouter, Body, Depends, HTTPException, Query, status +from fastapi import APIRouter, Body, Depends, HTTPException, Query, Request, Response, status from pydantic import BaseModel, Field, ConfigDict +from flocks.audit import emit_audit_event +from flocks.auth.context import AuthUser from flocks.utils.log import Log from flocks.provider.provider import Provider, ModelInfo as ProviderModelInfo from flocks.security.secrets import SecretManager -from flocks.server.auth import require_admin +from flocks.server.auth import get_request_ip, get_request_user_agent, require_admin from flocks.config.config import Config from flocks.config.config_writer import ConfigWriter from flocks.storage.storage import Storage @@ -1175,9 +1177,32 @@ def _is_api_service_builtin(provider_id: str, tools: Optional[List[Any]] = None) def _set_api_service_tools_enabled(provider_id: str, enabled: bool) -> int: """Synchronize the enabled state of all tools under an API service.""" - matched_tools = _get_api_service_tool_infos(provider_id) - for tool_info in matched_tools: - tool_info.enabled = enabled + from flocks.tool.registry import ToolRegistry + + ToolRegistry.init() + with ToolRegistry._refresh_lock: + matched_tools = _get_api_service_tool_infos(provider_id) + tool_settings = ConfigWriter.list_tool_settings() + previous_states = { + tool_info.name: bool(tool_info.enabled) + for tool_info in matched_tools + } + for tool_info in matched_tools: + default_enabled = ToolRegistry.get_default_enabled(tool_info.name) + effective_enabled = enabled and ( + default_enabled if default_enabled is not None else True + ) + setting = tool_settings.get(tool_info.name) + if isinstance(setting, dict) and "enabled" in setting: + effective_enabled = enabled and bool(setting["enabled"]) + tool_info.enabled = effective_enabled + if effective_enabled and not previous_states[tool_info.name]: + ToolRegistry._reset_failure_state(tool_info.name) + if any( + previous_states[tool_info.name] != bool(tool_info.enabled) + for tool_info in matched_tools + ): + ToolRegistry._bump_revision("api_service_tool_state") if matched_tools: try: from flocks.server.routes.tool import _invalidate_tool_summary_cache @@ -1656,6 +1681,43 @@ class ProviderCredentialResponse(BaseModel): has_credential: bool +def _load_llm_provider_credentials( + provider_id: str, + *, + reveal_api_key: bool = False, +) -> ProviderCredentialResponse: + from flocks.security import get_secret_manager + + secrets = get_secret_manager() + + secret_id = f"{provider_id}_llm_key" + api_key = secrets.get(secret_id) + if not api_key: + secret_id = f"{provider_id}_api_key" + api_key = secrets.get(secret_id) + if not api_key: + api_key = _get_inline_provider_api_key(provider_id) + + base_url = None + raw_provider = ConfigWriter.get_provider_raw(provider_id) + if raw_provider: + options = raw_provider.get("options", {}) + base_url = options.get("baseURL") or options.get("base_url") + if not base_url: + base_url = secrets.get(f"{provider_id}_base_url") + + is_placeholder = _is_placeholder_api_key(api_key) + ui_api_key = None if is_placeholder else api_key + + return ProviderCredentialResponse( + secret_id=secret_id if api_key else None, + api_key=ui_api_key if reveal_api_key else None, + api_key_masked=SecretManager.mask(ui_api_key) if ui_api_key else None, + base_url=base_url, + has_credential=bool(api_key), + ) + + @router.get( "/{provider_id}/credentials", response_model=ProviderCredentialResponse, @@ -1673,51 +1735,53 @@ async def get_provider_credentials( For API services use GET /{provider_id}/service-credentials instead. """ - from flocks.security import get_secret_manager - try: - secrets = get_secret_manager() - - # LLM provider convention: _llm_key first, fall back to legacy _api_key - secret_id = f"{provider_id}_llm_key" - api_key = secrets.get(secret_id) - if not api_key: - secret_id = f"{provider_id}_api_key" - api_key = secrets.get(secret_id) - if not api_key: - api_key = _get_inline_provider_api_key(provider_id) + return _load_llm_provider_credentials(provider_id) + except Exception as e: + log.error("provider.credentials.get.error", {"provider_id": provider_id, "error": str(e)}) + raise HTTPException(status_code=500, detail=str(e)) - # base_url from flocks.json raw (not resolved) - base_url = None - raw_provider = ConfigWriter.get_provider_raw(provider_id) - if raw_provider: - options = raw_provider.get("options", {}) - base_url = options.get("baseURL") or options.get("base_url") - # Fallback: check legacy .secret.json entry (migration may not have run yet) - if not base_url: - base_url = secrets.get(f"{provider_id}_base_url") - # When the stored value is the no-auth sentinel, hide it from the - # response so the UI doesn't pre-fill the password input with the - # literal string "not-needed". We still report ``has_credential=True`` - # because a credential record DOES exist; the user just chose to leave - # the key blank for an internal/no-auth gateway. - is_placeholder = _is_placeholder_api_key(api_key) - ui_api_key = None if is_placeholder else api_key +@router.post( + "/{provider_id}/credentials/reveal", + response_model=ProviderCredentialResponse, + summary="Reveal provider credentials", + description="Reveal credential information for a registered LLM provider to an administrator.", +) +async def reveal_provider_credentials( + provider_id: str, + request: Request, + response: Response, + admin: AuthUser = Depends(require_admin), +): + response.headers["Cache-Control"] = "no-store" + response.headers["Pragma"] = "no-cache" - return ProviderCredentialResponse( - secret_id=secret_id if api_key else None, - api_key=None, - api_key_masked=( - SecretManager.mask(ui_api_key) - if ui_api_key - else None - ), - base_url=base_url, - has_credential=bool(api_key), - ) + try: + credentials = _load_llm_provider_credentials(provider_id, reveal_api_key=True) + try: + await emit_audit_event( + "provider.credentials_reveal", + { + "action": "credentials_reveal", + "actor_id": admin.id, + "actor_name": admin.username, + "user_id": admin.id, + "username": admin.username, + "provider_id": provider_id, + "secret_id": credentials.secret_id, + "ip": get_request_ip(request), + "user_agent": get_request_user_agent(request), + }, + ) + except Exception as audit_error: + log.warn( + "provider.credentials.reveal.audit_failed", + {"provider_id": provider_id, "error": str(audit_error)}, + ) + return credentials except Exception as e: - log.error("provider.credentials.get.error", {"provider_id": provider_id, "error": str(e)}) + log.error("provider.credentials.reveal.error", {"provider_id": provider_id, "error": str(e)}) raise HTTPException(status_code=500, detail=str(e)) diff --git a/flocks/server/routes/session.py b/flocks/server/routes/session.py index 3cf75b890..2f9f2acad 100644 --- a/flocks/server/routes/session.py +++ b/flocks/server/routes/session.py @@ -8,13 +8,15 @@ import asyncio import json +import os import time +from pathlib import Path from typing import List, Optional, Any, Dict, Literal, Union, Tuple from fastapi import APIRouter, HTTPException, status, Query, Request from fastapi.responses import StreamingResponse from pydantic import BaseModel, Field, ConfigDict -from flocks.auth.context import get_current_auth_user, set_current_auth_user, reset_current_auth_user +from flocks.auth.context import AuthUser, get_current_auth_user, reset_current_auth_user, set_current_auth_user from flocks.server.routes._timing import log_route_timing from flocks.audit import emit_audit_event from flocks.license import assert_license_active @@ -48,6 +50,62 @@ # extension (e.g. a PNG named report.pdf.exe whose tail would otherwise be # ".exe"). _UPLOAD_SAFE_EXTS = frozenset({"png", "jpg", "jpeg", "gif", "webp", "bmp", "pdf"}) +_UPLOAD_EXT_BY_MIME = { + "image/png": ".png", + "image/jpeg": ".jpg", + "image/jpg": ".jpg", + "image/gif": ".gif", + "image/webp": ".webp", + "image/bmp": ".bmp", + "application/pdf": ".pdf", +} + + +def _session_uploads_dir(session_id: str) -> Path: + """Return the application-owned upload directory for one session.""" + from flocks.config.config import Config + + uploads_root = (Config.get_data_path() / "uploads").resolve() + target = (uploads_root / session_id).resolve() + if target == uploads_root or not target.is_relative_to(uploads_root): + raise ValueError(f"Invalid session ID for upload path: {session_id}") + return target + + +def _materialize_data_url_part( + session_id: str, + data_url: str, + mime_hint: str, + filename_hint: Optional[str], + *, + failure_event: str = "session.prompt_queue.materialize_failed", +) -> str: + """Persist a data URL under the application data directory.""" + try: + import base64 + from flocks.utils.id import Identifier + + _header, _sep, encoded = data_url.partition(",") + if not encoded: + return data_url + raw_bytes = base64.b64decode(encoded) + uploads_root = _session_uploads_dir(session_id) + uploads_root.mkdir(parents=True, exist_ok=True) + + ext = _UPLOAD_EXT_BY_MIME.get(mime_hint, "") + if not ext and filename_hint: + _, _, tail = filename_hint.rpartition(".") + if tail.lower() in _UPLOAD_SAFE_EXTS: + ext = "." + tail.lower() + target = uploads_root / f"{Identifier.create('part')}{ext}" + target.write_bytes(raw_bytes) + return target.resolve().as_uri() + except Exception as exc: + log.warn(failure_event, { + "sessionID": session_id, + "error": str(exc), + }) + return data_url def _context_usage_cache_key(session_id: str, session: SessionModel) -> Tuple[str, int]: @@ -129,6 +187,7 @@ class SessionCreateRequest(BaseModel): model_config = ConfigDict(populate_by_name=True, extra="ignore") parentID: Optional[str] = Field(None, alias="parent_id", description="Parent session ID") + projectID: Optional[str] = Field(None, alias="project_id", description="Project ID") title: Optional[str] = Field(None, description="Session title") permission: Optional[List[PermissionRule]] = Field(None, description="Permission rules") category: Optional[str] = Field(None, description="Session category (e.g. 'user', 'workflow')") @@ -175,6 +234,7 @@ class SessionResponse(BaseModel): id: str = Field(..., description="Session ID") slug: str = Field("", description="Session slug") projectID: str = Field(..., description="Project ID") + effectiveProjectID: str = Field(..., description="Project used for session manager grouping") directory: str = Field(..., description="Working directory") parentID: Optional[str] = Field(None, description="Parent session ID") summary: Optional[Dict[str, Any]] = Field(None, description="Session summary with diffs") @@ -200,6 +260,9 @@ class SessionListItem(BaseModel): model_config = ConfigDict(populate_by_name=True, by_alias=True) id: str + projectID: str + effectiveProjectID: str + directory: str title: str time: SessionTime category: str = "user" @@ -212,19 +275,31 @@ class SessionListItem(BaseModel): isShared: bool = False -def _session_to_response(session: SessionModel) -> SessionResponse: +def _session_to_response( + session: SessionModel, + effective_project_id: Optional[str] = None, + shared_project_ids: Optional[set[str]] = None, +) -> SessionResponse: """ Convert SessionModel to SessionResponse """ current_user = get_current_auth_user() can_write = SessionPolicy.can_write(session, current_user) can_delete = SessionPolicy.can_delete(session, current_user) - is_shared = SessionPolicy.is_shared(session) + is_shared = SessionPolicy.is_shared(session, shared_project_ids) + from flocks.project.project import Project + + if effective_project_id is None: + effective_project_id = Project.effective_project_id( + current_user.id if current_user else API_TOKEN_SERVICE_USER_ID, + session.project_id, + ) return SessionResponse( id=session.id, slug=session.slug, projectID=session.project_id, + effectiveProjectID=effective_project_id, directory=session.directory, title=session.title, version=session.version, @@ -250,11 +325,25 @@ def _session_to_response(session: SessionModel) -> SessionResponse: ) -def _session_to_list_item(session: SessionModel) -> SessionListItem: +def _session_to_list_item( + session: SessionModel, + effective_project_id: Optional[str] = None, + shared_project_ids: Optional[set[str]] = None, +) -> SessionListItem: """Convert a session to the lightweight manager-list response shape.""" current_user = get_current_auth_user() + from flocks.project.project import Project + + if effective_project_id is None: + effective_project_id = Project.effective_project_id( + current_user.id if current_user else API_TOKEN_SERVICE_USER_ID, + session.project_id, + ) return SessionListItem( id=session.id, + projectID=session.project_id, + effectiveProjectID=effective_project_id, + directory=session.directory, title=session.title, time=SessionTime( created=session.time.created, @@ -269,13 +358,17 @@ def _session_to_list_item(session: SessionModel) -> SessionListItem: model_pinned=session.model_pinned, canWrite=SessionPolicy.can_write(session, current_user), canDelete=SessionPolicy.can_delete(session, current_user), - isShared=SessionPolicy.is_shared(session), + isShared=SessionPolicy.is_shared(session, shared_project_ids), ) -async def _session_to_response_with_goal(session: SessionModel) -> SessionResponse: +async def _session_to_response_with_goal( + session: SessionModel, + effective_project_id: Optional[str] = None, + shared_project_ids: Optional[set[str]] = None, +) -> SessionResponse: """Convert SessionModel to SessionResponse and attach persisted goal state.""" - response = _session_to_response(session) + response = _session_to_response(session, effective_project_id, shared_project_ids) try: from flocks.session.goal import GoalManager @@ -363,6 +456,52 @@ async def _get_session_by_id_unfiltered(session_id: str) -> Optional[SessionMode reset_current_auth_user(token) +async def _resolve_session_working_directory(session: SessionModel) -> str: + """Keep a valid historical cwd, otherwise use the request context.""" + + if session.directory: + directory = Path(session.directory).expanduser() + if directory.is_dir() and os.access(directory, os.R_OK | os.X_OK): + return str(directory.resolve()) + + from flocks.project.instance import Instance + + fallback_directory = Instance.get_directory() or os.getcwd() + fallback_path = Path(fallback_directory).expanduser() + if not fallback_path.is_dir() or not os.access(fallback_path, os.R_OK | os.X_OK): + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail="The session directory and current request directory are unavailable", + ) + fallback_directory = str(fallback_path.resolve()) + log.warn( + "session.directory.fallback", + { + "sessionID": session.id, + "stored_directory": session.directory, + "fallback_directory": fallback_directory, + }, + ) + try: + from flocks.server.routes.event import publish_event + + await publish_event( + "session.notice", + { + "sessionID": session.id, + "kind": "directory-fallback", + "storedDirectory": session.directory, + "fallbackDirectory": fallback_directory, + }, + ) + except Exception as exc: + log.debug( + "session.directory.fallback_notice_failed", + {"sessionID": session.id, "error": str(exc)}, + ) + return fallback_directory + + async def _publish_context_usage_update( event_publish_callback, session_id: str, @@ -440,6 +579,7 @@ async def list_sessions( view: Optional[Literal["list"]] = Query(None, description="Use lightweight list rows"), manager: Optional[bool] = Query(None, description="Apply session-manager visibility filters"), directory: Optional[str] = Query(None, description="Filter by project directory"), + projectID: Optional[str] = Query(None, description="Filter by effective project ID"), roots: Optional[bool] = Query(None, description="Only return root sessions (no parentID)"), start: Optional[int] = Query(None, description="Filter sessions updated on or after this timestamp"), search: Optional[str] = Query(None, description="Filter by title (case-insensitive)"), @@ -449,21 +589,47 @@ async def list_sessions( ) -> List[Union[SessionResponse, SessionListItem]]: """List all sessions with optional filters""" started_at = time.perf_counter() - _current_user = require_user(request) + current_user = require_user(request) + from flocks.project.project import ( + DEFAULT_PROJECT_ID, + Project, + TASK_SESSION_GROUP_ID, + ) list_started_at = time.perf_counter() - all_sessions = await Session.list_all() + all_sessions = await Session.list_all_unfiltered() list_elapsed_ms = (time.perf_counter() - list_started_at) * 1000 + 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 manager_categories = {"user", "workflow", "entity-config"} skip_remaining = offset or 0 for session in all_sessions: + if not SessionPolicy.can_read( + session, + current_user, + shared_project_ids=shared_project_ids, + ): + continue if _is_hidden_from_session_manager(session): continue if directory is not None and session.directory != directory: continue + effective_project_id = ( + session.project_id + if session.project_id in visible_project_ids + else TASK_SESSION_GROUP_ID + ) + requested_group_id = ( + TASK_SESSION_GROUP_ID + if projectID == DEFAULT_PROJECT_ID + else projectID + ) + if requested_group_id is not None and effective_project_id != requested_group_id: + continue if (roots or manager) and session.parent_id: continue if start is not None and session.time.updated < start: @@ -485,12 +651,16 @@ async def list_sessions( continue filtered.append(session) + effective_project_ids[session.id] = effective_project_id if limit is not None and len(filtered) >= limit: break if view == "list": - response = [_session_to_list_item(s) for s in filtered] + response = [ + _session_to_list_item(s, effective_project_ids[s.id], shared_project_ids) + for s in filtered + ] log_route_timing(log, "session.list.light.complete", started_at=started_at, extra={ "total": len(all_sessions), "count": len(response), @@ -500,11 +670,19 @@ async def list_sessions( "offset": offset, "search": bool(search), "category": category, + "projectID": projectID, "list_ms": round(list_elapsed_ms, 2), }) return response - response = [await _session_to_response_with_goal(s) for s in filtered] + response = [ + await _session_to_response_with_goal( + s, + effective_project_ids[s.id], + shared_project_ids, + ) + for s in filtered + ] log_route_timing(log, "session.list.complete", started_at=started_at, extra={ "count": len(response), "roots": roots, @@ -513,6 +691,7 @@ async def list_sessions( "offset": offset, "search": bool(search), "category": category, + "projectID": projectID, "list_ms": round(list_elapsed_ms, 2), }) return response @@ -529,19 +708,40 @@ async def create_session(http_request: Request, request: Optional[SessionCreateR """Create a new session""" current_user = require_user(http_request) await assert_license_active(feature="session_create") - import os - if request is None: request = SessionCreateRequest() - - # Use Instance context if available, otherwise use cwd + from flocks.project.instance import Instance - try: - directory = Instance.directory - project_id = Instance.project.id if hasattr(Instance, 'project') else "default" - except Exception: - directory = os.getcwd() - project_id = "default" + from flocks.project.project import ( + DEFAULT_PROJECT_ID, + Project, + TASK_SESSION_GROUP_ID, + ) + + directory = Instance.get_directory() or os.getcwd() + instance_project = Instance.get_project() + project_id = instance_project.id if instance_project else DEFAULT_PROJECT_ID + + # "default" is accepted for compatibility with older WebUI clients, but + # ordinary sessions use the current request context rather than a virtual + # project's worktree. + if request.projectID and request.projectID not in { + DEFAULT_PROJECT_ID, + TASK_SESSION_GROUP_ID, + }: + project = await Project.get(request.projectID, owner_id=current_user.id) + if not project: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Project {request.projectID} not found", + ) + if project.path_status != "available": + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=f"Project directory is {project.path_status}", + ) + project_id = project.id + directory = project.worktree # Trigger command:new hook if creating from parent (like /new command) if request.parentID: @@ -597,6 +797,7 @@ async def create_session(http_request: Request, request: Optional[SessionCreateR owner_username=None if is_api_token_client else current_user.username, **({"category": request.category} if request.category else {}), ) + Project.invalidate_session_stats() log.info("session.created", {"session_id": session.id}) try: @@ -760,12 +961,17 @@ async def update_session_todos(sessionID: str, todos: List[TodoInfo], request: R async def delete_session(sessionID: str, request: Request) -> bool: """Delete session by ID (returns true)""" current_user = require_user(request) - session = await _get_session_by_id_unfiltered(sessionID) + return await delete_session_for_user(sessionID, current_user) + + +async def delete_session_for_user(session_id: str, current_user: AuthUser) -> bool: + """Delete one session using the normal lifecycle cleanup.""" + session = await _get_session_by_id_unfiltered(session_id) if not session: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, - detail=f"Session {sessionID} not found" + detail=f"Session {session_id} not found" ) if not SessionPolicy.can_delete(session, current_user): @@ -774,38 +980,38 @@ async def delete_session(sessionID: str, request: Request) -> bool: from flocks.session.goal import GoalManager from flocks.session.interaction_queue import InteractionQueue - await _abort_session_processing(sessionID) - await InteractionQueue.clear(sessionID) - await GoalManager.clear(sessionID) - await _wait_for_sessions_idle([sessionID]) - await _abort_and_wait_descendant_sessions(session.project_id, sessionID) - await Session.delete(session.project_id, sessionID) + await _abort_session_processing(session_id) + await InteractionQueue.clear(session_id) + await GoalManager.clear(session_id) + await _wait_for_sessions_idle([session_id]) + await _abort_and_wait_descendant_sessions(session.project_id, session_id) + await Session.delete(session.project_id, session_id) + from flocks.project.project import Project + + Project.invalidate_session_stats() # Best-effort cleanup of any image/file uploads materialised for this - # session via ``_materialize_data_url_to_disk`` (see prompt_async). + # session via ``_materialize_data_url_part`` (see prompt_async). # The session DB row is gone, so the on-disk bytes are now orphaned — - # remove them to keep the workspace tidy. We deliberately swallow any + # remove them to keep application data tidy. We deliberately swallow any # filesystem errors: deletion of the session record is the contract, # the upload cleanup is incidental. try: import shutil - from flocks.workspace.manager import WorkspaceManager - - ws = WorkspaceManager.get_instance() - uploads_root = ws.resolve_workspace_path(f"uploads/{sessionID}") + uploads_root = _session_uploads_dir(session_id) if uploads_root.exists() and uploads_root.is_dir(): shutil.rmtree(uploads_root, ignore_errors=True) log.info("session.uploads.cleaned", { - "session_id": sessionID, + "session_id": session_id, "path": str(uploads_root), }) except Exception as exc: log.warn("session.uploads.cleanup_failed", { - "session_id": sessionID, + "session_id": session_id, "error": str(exc), }) - log.info("session.deleted", {"session_id": sessionID}) + log.info("session.deleted", {"session_id": session_id}) try: await emit_audit_event( "session_action", @@ -815,7 +1021,7 @@ async def delete_session(sessionID: str, request: Request) -> bool: "actor_name": current_user.username, "user_name": current_user.username, "username": current_user.username, - "session_id": sessionID, + "session_id": session_id, "owner_user_id": current_user.id, "project_id": session.project_id, }, @@ -944,11 +1150,7 @@ async def update_session( ) async def share_session_local(sessionID: str, http_request: Request) -> SessionResponse: current_user = require_user(http_request) - token = set_current_auth_user(current_user) - try: - existing = await Session.get_by_id(sessionID) - finally: - reset_current_auth_user(token) + existing = await _get_session_by_id_unfiltered(sessionID) if not existing: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, @@ -956,11 +1158,15 @@ async def share_session_local(sessionID: str, http_request: Request) -> SessionR ) _require_session_write_access(existing, current_user) metadata = _share_metadata(existing, shared=True, actor_user_id=current_user.id) - session = await Session.update( - project_id=existing.project_id, - session_id=sessionID, - metadata=metadata, - ) + token = set_current_auth_user(current_user) + try: + session = await Session.update( + project_id=existing.project_id, + session_id=sessionID, + metadata=metadata, + ) + finally: + reset_current_auth_user(token) if not session: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, @@ -977,11 +1183,7 @@ async def share_session_local(sessionID: str, http_request: Request) -> SessionR ) async def unshare_session_local(sessionID: str, http_request: Request) -> SessionResponse: current_user = require_user(http_request) - token = set_current_auth_user(current_user) - try: - existing = await Session.get_by_id(sessionID) - finally: - reset_current_auth_user(token) + existing = await _get_session_by_id_unfiltered(sessionID) if not existing: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, @@ -989,11 +1191,15 @@ async def unshare_session_local(sessionID: str, http_request: Request) -> Sessio ) _require_session_write_access(existing, current_user) metadata = _share_metadata(existing, shared=False, actor_user_id=current_user.id) - session = await Session.update( - project_id=existing.project_id, - session_id=sessionID, - metadata=metadata, - ) + token = set_current_auth_user(current_user) + try: + session = await Session.update( + project_id=existing.project_id, + session_id=sessionID, + metadata=metadata, + ) + finally: + reset_current_auth_user(token) if not session: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, @@ -1219,11 +1425,13 @@ async def summarize_session(sessionID: str, request: SummarizeRequest, http_requ if msg.role == MessageRole.USER: current_agent = msg.agent or DEFAULT_AGENT break + + working_directory = await _resolve_session_working_directory(session) async def _run_in_background(): try: await Instance.provide( - directory=session.directory, + directory=working_directory, init=instance_bootstrap, fn=lambda: _run_session_compaction( sessionID, @@ -1232,6 +1440,7 @@ async def _run_in_background(): explicit_model_id=request.modelID, auto=request.auto, event_publish_callback=publish_event, + working_directory=working_directory, ), ) except Exception as e: @@ -2004,6 +2213,7 @@ async def _on_error(error: str): model_id=model_id, agent_name=agent_name, callbacks=loop_callbacks, + working_directory=working_directory, ) if result.action == "queued": @@ -2133,7 +2343,7 @@ async def resend_session_message( detail="Session is currently generating a response", ) - working_directory = session.directory or os.getcwd() + working_directory = await _resolve_session_working_directory(session) async def _handle_resend() -> None: runtime = await _prepare_replay_runtime(sessionID, message) @@ -2228,7 +2438,7 @@ async def regenerate_session_message( detail="Session is currently generating a response", ) - working_directory = session.directory or os.getcwd() + working_directory = await _resolve_session_working_directory(session) async def _handle_regenerate() -> None: runtime = await _prepare_replay_runtime(sessionID, parent_message) @@ -2293,7 +2503,7 @@ async def send_session_message(sessionID: str, request: PromptRequest, http_requ current_user = require_user(http_request) _require_session_write_access(session, current_user) - working_directory = session.directory or os.getcwd() + working_directory = await _resolve_session_working_directory(session) log.info("session.message.send.processing", { "sessionID": sessionID, @@ -2477,6 +2687,7 @@ async def _run_session_compaction( auto: bool = False, event_publish_callback=None, focus_instruction: Optional[str] = None, + working_directory: Optional[str] = None, ) -> tuple[str, str, str]: """Execute session compaction directly without routing through the LLM loop. @@ -2494,6 +2705,8 @@ async def _run_session_compaction( session = await Session.get_by_id(session_id) if not session: raise ValueError(f"Session {session_id} not found") + if working_directory: + session = session.model_copy(update={"directory": working_directory}) await SessionRevert.cleanup(session) agent_name, provider_id, model_id = await _resolve_compaction_context( @@ -2764,51 +2977,6 @@ async def _process_session_message( # ------------------------------------------------------------------ from flocks.session.message import FilePart - def _materialize_data_url_to_disk( - data_url: str, mime_hint: str, filename_hint: Optional[str] - ) -> str: - """Decode a ``data:`` URL to ``~/.flocks/workspace/uploads//...``. - - Returns a ``file://`` URL pointing at the persisted file. On failure - the original ``data:`` URL is returned unchanged (older code paths - still cope with that, just with the now-known token-cost penalty). - """ - try: - import base64 - from flocks.workspace.manager import WorkspaceManager - - header, _, encoded = data_url.partition(",") - if not encoded: - return data_url - raw_bytes = base64.b64decode(encoded) - - ws = WorkspaceManager.get_instance() - # Use resolve_workspace_path to guard against path traversal if - # sessionID were ever user-controlled (e.g. ../../../tmp/x). - uploads_root = ws.resolve_workspace_path(f"uploads/{sessionID}") - uploads_root.mkdir(parents=True, exist_ok=True) - - ext_map = { - "image/png": ".png", "image/jpeg": ".jpg", "image/jpg": ".jpg", - "image/gif": ".gif", "image/webp": ".webp", "image/bmp": ".bmp", - "application/pdf": ".pdf", - } - ext = ext_map.get(mime_hint, "") - if not ext and filename_hint: - _, _, tail = filename_hint.rpartition(".") - if tail.lower() in _UPLOAD_SAFE_EXTS: - ext = "." + tail.lower() - unique_name = f"{Identifier.create('part')}{ext}" - target = uploads_root / unique_name - target.write_bytes(raw_bytes) - return target.resolve().as_uri() - except Exception as exc: - log.warn("session.message.file_part.materialize_failed", { - "sessionID": sessionID, - "error": str(exc), - }) - return data_url - for raw_part in request.parts or []: part_type = raw_part.get("type") if part_type == "text": @@ -2824,7 +2992,13 @@ def _materialize_data_url_to_disk( continue # Materialize ``data:`` URLs to disk before persisting the part. if url.startswith("data:"): - url = _materialize_data_url_to_disk(url, mime, raw_part.get("filename")) + url = _materialize_data_url_part( + sessionID, + url, + mime, + raw_part.get("filename"), + failure_event="session.message.file_part.materialize_failed", + ) file_part_id = raw_part.get("id") or Identifier.create("part") file_part = FilePart( id=file_part_id, @@ -2927,6 +3101,7 @@ async def _on_error(error: str): model_id=model_id, agent_name=agent_name, callbacks=loop_callbacks, + working_directory=working_directory, ) # ------------------------------------------------------------------ @@ -3277,50 +3452,6 @@ def _materialize_queued_parts(session_id: str, parts: List[Dict[str, Any]]) -> L return prepared -def _materialize_data_url_part( - session_id: str, - data_url: str, - mime_hint: str, - filename_hint: Optional[str], -) -> str: - try: - import base64 - from flocks.workspace.manager import WorkspaceManager - from flocks.utils.id import Identifier - - _header, _sep, encoded = data_url.partition(",") - if not encoded: - return data_url - raw_bytes = base64.b64decode(encoded) - ws = WorkspaceManager.get_instance() - uploads_root = ws.resolve_workspace_path(f"uploads/{session_id}") - uploads_root.mkdir(parents=True, exist_ok=True) - - ext_map = { - "image/png": ".png", - "image/jpeg": ".jpg", - "image/jpg": ".jpg", - "image/gif": ".gif", - "image/webp": ".webp", - "image/bmp": ".bmp", - "application/pdf": ".pdf", - } - ext = ext_map.get(mime_hint, "") - if not ext and filename_hint: - _, _, tail = filename_hint.rpartition(".") - if tail.lower() in _UPLOAD_SAFE_EXTS: - ext = "." + tail.lower() - target = uploads_root / f"{Identifier.create('part')}{ext}" - target.write_bytes(raw_bytes) - return target.resolve().as_uri() - except Exception as exc: - log.warn("session.prompt_queue.materialize_failed", { - "sessionID": session_id, - "error": str(exc), - }) - return data_url - - def _event_from_queued_prompt(item, working_directory: str): from flocks.input.events import UserInputEvent @@ -3621,6 +3752,7 @@ async def _run_session_control(output_event, parsed) -> bool: auto=False, event_publish_callback=publish_event, focus_instruction=focus_instruction, + working_directory=working_directory, ) return True @@ -3775,7 +3907,7 @@ async def run_prompt_queue_item_now(sessionID: str, queueID: str) -> Dict[str, A status_code=status.HTTP_404_NOT_FOUND, detail=f"Session {sessionID} not found", ) - working_directory = session.directory or os.getcwd() + working_directory = await _resolve_session_working_directory(session) try: await InteractionQueue.promote(sessionID, queueID) except QueueItemNotFoundError as exc: @@ -3817,7 +3949,7 @@ async def send_session_message_async( current_user = require_user(http_request) _require_session_write_access(session, current_user) - working_directory = session.directory or os.getcwd() + working_directory = await _resolve_session_working_directory(session) await _require_agent_usable_for_chat(request.agent) log.info("session.prompt_async.accepted", { @@ -3914,7 +4046,7 @@ async def send_session_command(sessionID: str, request: CommandRequest, http_req current_user = require_user(http_request) _require_session_write_access(session, current_user) - working_directory = session.directory or os.getcwd() + working_directory = await _resolve_session_working_directory(session) await _require_agent_usable_for_chat(request.agent) raw_arguments = request.arguments if not raw_arguments and request.arguments_json is not None: diff --git a/flocks/server/routes/tool.py b/flocks/server/routes/tool.py index e217cbf3f..281b8e813 100644 --- a/flocks/server/routes/tool.py +++ b/flocks/server/routes/tool.py @@ -639,33 +639,39 @@ def _get_effective_tool_enabled( def _set_global_tool_enabled(tool: Any, desired: bool) -> bool: """Persist and apply the global enabled state for a registry tool.""" - default = _get_default_enabled(tool.info) - # Service gate: only matters when the user is trying to enable. - # Disabling is always honoured. - service_ok = _service_allows_enable(tool.info) - new_enabled = desired and service_ok - - if desired == default: - removed = ConfigWriter.delete_tool_setting(tool.info.name) - log.info("tool.updated.reset_to_default", { - "name": tool.info.name, - "enabled": new_enabled, - "default": default, - "removed_overlay": removed, - }) - else: - ConfigWriter.set_tool_setting(tool.info.name, {"enabled": desired}) - log.info("tool.updated", { - "name": tool.info.name, - "enabled": new_enabled, - "requested": desired, - "blocked_by_service": desired and not service_ok, - "native": tool.info.native, - "store": "overlay", - }) + with ToolRegistry._refresh_lock: + previous_enabled = bool(tool.info.enabled) + default = _get_default_enabled(tool.info) + # Service gate: only matters when the user is trying to enable. + # Disabling is always honoured. + service_ok = _service_allows_enable(tool.info) + new_enabled = desired and service_ok + + if desired == default: + removed = ConfigWriter.delete_tool_setting(tool.info.name) + log.info("tool.updated.reset_to_default", { + "name": tool.info.name, + "enabled": new_enabled, + "default": default, + "removed_overlay": removed, + }) + else: + ConfigWriter.set_tool_setting(tool.info.name, {"enabled": desired}) + log.info("tool.updated", { + "name": tool.info.name, + "enabled": new_enabled, + "requested": desired, + "blocked_by_service": desired and not service_ok, + "native": tool.info.native, + "store": "overlay", + }) - tool.info.enabled = new_enabled - return new_enabled + tool.info.enabled = new_enabled + if new_enabled: + ToolRegistry._reset_failure_state(tool.info.name) + if new_enabled != previous_enabled: + ToolRegistry._bump_revision("tool_setting_update") + return new_enabled # Routes @@ -969,10 +975,16 @@ async def reset_tool_setting(tool_name: str, _admin: object = Depends(require_ad detail=f"Tool not found: {tool_name}", ) - removed = ConfigWriter.delete_tool_setting(tool_name) - default = _get_default_enabled(tool.info) - new_enabled = default and _service_allows_enable(tool.info) - tool.info.enabled = new_enabled + with ToolRegistry._refresh_lock: + removed = ConfigWriter.delete_tool_setting(tool_name) + default = _get_default_enabled(tool.info) + new_enabled = default and _service_allows_enable(tool.info) + previous_enabled = bool(tool.info.enabled) + tool.info.enabled = new_enabled + if new_enabled: + ToolRegistry._reset_failure_state(tool_name) + if new_enabled != previous_enabled: + ToolRegistry._bump_revision("tool_setting_reset") _invalidate_tool_summary_cache() log.info("tool.setting.reset", { diff --git a/flocks/server/routes/workflow.py b/flocks/server/routes/workflow.py index 1bf338498..0ac9d0429 100644 --- a/flocks/server/routes/workflow.py +++ b/flocks/server/routes/workflow.py @@ -22,11 +22,13 @@ import uuid from flocks.workflow.models import Workflow, Node, Edge +from flocks.workflow.service_port import resolve_service_port from flocks.workflow.runner import RunWorkflowResult, run_workflow from flocks.workflow.center import ( WorkflowCenterError, WorkflowNotFoundError, WorkflowNotPublishedError, + WorkflowPortUnavailableError, get_workflow_health, invoke_published_workflow, list_registry_entries, @@ -75,6 +77,7 @@ workflow_json_declares_triggers, workflow_trigger_definitions_from_json, ) +from flocks.workflow.triggers.models import TriggerType from flocks.workflow.triggers.dispatcher import evaluate_trigger_filter from flocks.workflow.triggers.runtime import default_runtime as default_trigger_runtime from flocks.workflow.triggers.compat import ( @@ -97,6 +100,12 @@ log = Log.create(service="workflow-routes") _PROGRESS_FLUSH_EVERY_STEPS = 5 +_WORKFLOW_LIST_ENRICH_CONCURRENCY = 8 +_WORKFLOW_API_HEALTH_INTERVAL_S = 5.0 +_WORKFLOW_API_HEALTH_PROBE_CONCURRENCY = 4 + +_workflow_api_health_cache: Dict[str, bool] = {} +_workflow_api_health_monitor_task: Optional[asyncio.Task[None]] = None _LEGACY_SINGLETON_TRIGGER_TYPES = frozenset({"schedule", "kafka", "syslog"}) _WEBHOOK_TRIGGER_TYPES = frozenset({"webhook", "custom_webhook"}) @@ -187,6 +196,32 @@ class WorkflowUpdateRequest(BaseModel): status: Optional[Literal["draft", "active", "archived"]] = Field(None, description="Status") +WorkflowCapabilityState = Literal["unconfigured", "starting", "running", "stopped", "error"] + + +class WorkflowCapabilityStatusResponse(BaseModel): + configured: bool + state: WorkflowCapabilityState + + +class WorkflowTriggerStatusItemResponse(BaseModel): + id: str + type: TriggerType + name: Optional[str] = None + state: WorkflowCapabilityState + rawState: Optional[str] = None + + +class WorkflowTriggerCapabilityStatusResponse(WorkflowCapabilityStatusResponse): + count: int = 0 + items: List[WorkflowTriggerStatusItemResponse] = Field(default_factory=list) + + +class WorkflowIntegrationStatusResponse(BaseModel): + api: WorkflowCapabilityStatusResponse + trigger: WorkflowTriggerCapabilityStatusResponse + + class WorkflowResponse(BaseModel): """Workflow response""" @@ -206,6 +241,27 @@ class WorkflowResponse(BaseModel): createdAt: int = Field(..., description="Created timestamp (ms)") updatedAt: int = Field(..., description="Updated timestamp (ms)") stats: Dict[str, Any] = Field(default_factory=dict, description="Statistics") + integrationStatus: Optional[WorkflowIntegrationStatusResponse] = Field( + None, + description="API publishing and trigger runtime summary", + ) + + +class WorkflowSummaryResponse(BaseModel): + """Lightweight workflow data used by card lists.""" + + id: str + name: str + nameI18n: Optional[Dict[str, str]] = None + description: Optional[str] = None + category: str = "default" + status: str = "draft" + source: Optional[str] = None + createdAt: int + updatedAt: int + nodeCount: int = 0 + stats: Dict[str, Any] = Field(default_factory=dict) + integrationStatus: WorkflowIntegrationStatusResponse class WorkflowRunRequest(BaseModel): @@ -260,6 +316,7 @@ class WorkflowCenterPublishRequest(BaseModel): None, description="Service driver. Defaults to FLOCKS_WORKFLOW_SERVICE_DRIVER or local.", ) + port: Optional[int] = Field(None, ge=1, le=65535, description="Stable host listening port") class WorkflowCenterInvokeRequest(BaseModel): @@ -329,9 +386,9 @@ def _workflow_integration_config_key(workflow_id: str) -> str: def _read_workflow_from_fs(workflow_id: str) -> Optional[Dict[str, Any]]: """Read workflow data from the filesystem. - Search order (lowest → highest priority), same roots as - resolve_global_workflow_roots / resolve_project_workflow_roots; per-id dir - is ``//`` with ``workflow.json`` inside. + Search order (lowest → highest priority) visits project-bundled roots before + global user roots; per-id dir is ``//`` with ``workflow.json`` + inside. """ return shared_read_workflow_from_fs(workflow_id) @@ -531,9 +588,8 @@ def _scan_workflow_base_dir(base_dir: Path, source: str) -> Dict[str, Dict[str, def _list_workflows_from_fs() -> List[Dict[str, Any]]: """Scan global and project workflow directories and return merged list. - Scan order matches *_all_scan_dirs()* (lowest -> highest priority): each - root from *resolve_global_workflow_roots* then each from - *resolve_project_workflow_roots(workspace)*; under each root, immediate + Scan order matches *_all_scan_dirs()* (lowest -> highest priority): project + bundle roots first, then global user roots; under each root, immediate subdirectories with *workflow.json* are workflows. Later entries override earlier ones when the workflow directory name (*id*) @@ -804,6 +860,7 @@ def _publish_for_config(service: Optional[Dict[str, Any]]) -> Dict[str, Any]: "enabled": bool(service) and status_value not in {"stopped", "unpublished"}, "status": status_value, "driver": service.get("driver") if service else None, + "port": resolve_service_port(service) if service else None, "serviceUrl": service.get("serviceUrl") if service else None, "invokeUrl": service.get("invokeUrl") if service else None, "containerName": service.get("containerName") if service else None, @@ -1309,6 +1366,22 @@ async def _get_workflow_stats(workflow_id: str) -> Dict[str, Any]: # ============================================================================= +def _filter_workflow_list_data( + all_data: List[Dict[str, Any]], + *, + category: Optional[str], + status: Optional[str], + exclude_id: Optional[str], +) -> List[Dict[str, Any]]: + return [ + data + for data in all_data + if not category or data.get("category") == category + if not status or data.get("status") == status + if not exclude_id or data.get("id") != exclude_id + ] + + @router.get("/workflow", response_model=List[WorkflowResponse]) async def list_workflows( category: Optional[str] = Query(None, description="Filter by category"), @@ -1326,26 +1399,31 @@ async def list_workflows( try: await _migrate_storage_to_filesystem() - all_data = _list_workflows_from_fs() - workflows = [] - - for data in all_data: - try: - if category and data.get("category") != category: - continue - if status and data.get("status") != status: - continue - if exclude_id and data.get("id") == exclude_id: - continue + filtered_data = _filter_workflow_list_data( + _list_workflows_from_fs(), + category=category, + status=status, + exclude_id=exclude_id, + ) - workflow_id = data["id"] - stats = await _get_workflow_stats(workflow_id) - data["stats"] = stats + semaphore = asyncio.Semaphore(_WORKFLOW_LIST_ENRICH_CONCURRENCY) - workflows.append(WorkflowResponse(**data)) + async def enrich_workflow(data: Dict[str, Any]) -> Optional[WorkflowResponse]: + try: + async with semaphore: + workflow_id = data["id"] + stats = await _get_workflow_stats(workflow_id) + enriched_data = { + **data, + "stats": stats, + } + return WorkflowResponse(**enriched_data) except Exception as e: log.warning("workflow.list.skip", {"id": data.get("id"), "error": str(e)}) - continue + return None + + enriched = await asyncio.gather(*(enrich_workflow(data) for data in filtered_data)) + workflows = [workflow for workflow in enriched if workflow is not None] workflows.sort(key=lambda w: w.updatedAt, reverse=True) @@ -1358,6 +1436,60 @@ async def list_workflows( raise HTTPException(status_code=500, detail=f"Failed to list workflows: {str(e)}") +@router.get("/workflow-summaries", response_model=List[WorkflowSummaryResponse]) +async def list_workflow_summaries( + category: Optional[str] = Query(None, description="Filter by category"), + status: Optional[str] = Query(None, description="Filter by status"), + exclude_id: Optional[str] = Query(None, alias="excludeId", description="Exclude workflow by ID"), +): + """Return lightweight workflow card data with cached integration states.""" + try: + await _migrate_storage_to_filesystem() + filtered_data = _filter_workflow_list_data( + _list_workflows_from_fs(), + category=category, + status=status, + exclude_id=exclude_id, + ) + semaphore = asyncio.Semaphore(_WORKFLOW_LIST_ENRICH_CONCURRENCY) + + async def build_summary(data: Dict[str, Any]) -> Optional[WorkflowSummaryResponse]: + try: + async with semaphore: + workflow_id = data["id"] + stats, integration_status = await asyncio.gather( + _get_workflow_stats(workflow_id), + _get_workflow_integration_status(workflow_id, data), + ) + workflow_json = data.get("workflowJson") or {} + nodes = workflow_json.get("nodes") if isinstance(workflow_json, dict) else [] + return WorkflowSummaryResponse( + id=workflow_id, + name=data.get("name") or workflow_id, + nameI18n=data.get("nameI18n"), + description=data.get("description"), + category=data.get("category") or "default", + status=data.get("status") or "draft", + source=data.get("source"), + createdAt=int(data.get("createdAt") or 0), + updatedAt=int(data.get("updatedAt") or 0), + nodeCount=len(nodes) if isinstance(nodes, list) else 0, + stats=stats, + integrationStatus=integration_status, + ) + except Exception as exc: + log.warning("workflow.summary.skip", {"id": data.get("id"), "error": str(exc)}) + return None + + enriched = await asyncio.gather(*(build_summary(data) for data in filtered_data)) + summaries = [summary for summary in enriched if summary is not None] + summaries.sort(key=lambda item: item.updatedAt, reverse=True) + return summaries + except Exception as exc: + log.error("workflow.summary.error", {"error": str(exc)}) + raise HTTPException(status_code=500, detail=f"Failed to list workflow summaries: {str(exc)}") + + @router.post("/workflow", response_model=WorkflowResponse, status_code=status.HTTP_201_CREATED) async def create_workflow(req: WorkflowCreateRequest): """ @@ -1782,6 +1914,7 @@ async def workflow_center_publish(workflow_id: str, req: Optional[WorkflowCenter workflow_id, image=req.image if req else None, driver=req.driver if req else None, + port=req.port if req else None, ) return result except WorkflowNotFoundError as e: @@ -2219,9 +2352,13 @@ async def _prepare_workflow_api_registry(workflow_id: str) -> tuple[Dict[str, An "sourceType": "main_storage", "workflowPath": str(workflow_path), "fingerprint": fp, - "publishStatus": "unpublished", + "publishStatus": existing_registry.get("publishStatus", "unpublished"), "registeredAt": existing_registry.get("registeredAt", now_ms), "updatedAt": now_ms, + "activeReleaseId": existing_registry.get("activeReleaseId"), + "serviceKey": existing_registry.get("serviceKey"), + "serviceUrl": existing_registry.get("serviceUrl"), + "servicePort": existing_registry.get("servicePort"), } await WorkflowStore.kv_put(f"{_REGISTRY_PREFIX_MAIN}{workflow_id}", registry_entry) return data, now_ms @@ -2262,6 +2399,7 @@ async def _normalize_listed_api_service(key: Any, entry: Any) -> Optional[Dict[s service["driver"] = runtime.get("driver") or service.get("driver") service["containerName"] = runtime.get("containerName") or service.get("containerName", "") service["image"] = runtime.get("image") or service.get("image") + service["port"] = resolve_service_port(runtime) or resolve_service_port(service) return service status = str(service.get("status") or "").strip().lower() @@ -2276,6 +2414,177 @@ async def _normalize_listed_api_service(key: Any, entry: Any) -> Optional[Dict[s return service +def _summarize_capability_state(raw_state: Any) -> WorkflowCapabilityState: + state = str(raw_state or "").strip().lower() + if state in {"running", "listening", "ready", "active"}: + return "running" + if state in {"failed", "error"}: + return "error" + if state in {"stopped", "unpublished", "paused", "disabled"}: + return "stopped" + return "starting" + + +async def refresh_workflow_api_health_cache() -> Dict[str, int]: + """Refresh API health outside list requests with process-wide bounded concurrency.""" + keys = await WorkflowStore.kv_list_keys(_API_SERVICE_PREFIX) + services = await asyncio.gather(*(WorkflowStore.kv_get(key) for key in keys)) + active_workflow_ids = [ + str(service.get("workflowId") or _workflow_id_from_api_service_key(key)) + for key, service in zip(keys, services) + if isinstance(service, dict) + and service + and _summarize_capability_state(service.get("status")) == "running" + ] + semaphore = asyncio.Semaphore(_WORKFLOW_API_HEALTH_PROBE_CONCURRENCY) + + async def probe(workflow_id: str) -> tuple[str, bool]: + async with semaphore: + try: + health = await get_workflow_health(workflow_id) + return workflow_id, bool(health.get("ok")) + except Exception as exc: + log.warning("workflow.api_health_monitor.probe_failed", {"id": workflow_id, "error": str(exc)}) + return workflow_id, False + + results = await asyncio.gather(*(probe(workflow_id) for workflow_id in active_workflow_ids)) + _workflow_api_health_cache.clear() + _workflow_api_health_cache.update(results) + return { + "checked": len(results), + "healthy": sum(1 for _, healthy in results if healthy), + } + + +async def run_workflow_api_health_monitor() -> None: + """Continuously maintain the health snapshot consumed by workflow cards.""" + while True: + try: + await refresh_workflow_api_health_cache() + except asyncio.CancelledError: + raise + except Exception as exc: + log.warning("workflow.api_health_monitor.refresh_failed", {"error": str(exc)}) + await asyncio.sleep(_WORKFLOW_API_HEALTH_INTERVAL_S) + + +def start_workflow_api_health_monitor() -> asyncio.Task[None]: + global _workflow_api_health_monitor_task + if _workflow_api_health_monitor_task is None or _workflow_api_health_monitor_task.done(): + _workflow_api_health_monitor_task = asyncio.create_task( + run_workflow_api_health_monitor(), + name="workflow-api-health-monitor", + ) + return _workflow_api_health_monitor_task + + +async def _get_workflow_integration_status( + workflow_id: str, + workflow_data: Dict[str, Any], +) -> WorkflowIntegrationStatusResponse: + api_summary = WorkflowCapabilityStatusResponse(configured=False, state="unconfigured") + try: + service = await WorkflowStore.kv_get(_api_service_key(workflow_id)) + if isinstance(service, dict) and service: + normalized_service = await _normalize_listed_api_service( + _api_service_key(workflow_id), + service, + ) + state = _summarize_capability_state((normalized_service or service).get("status")) + if state == "running": + cached_health = _workflow_api_health_cache.get(workflow_id) + if cached_health is None: + persisted_health = service.get("health") + if isinstance(persisted_health, dict) and isinstance(persisted_health.get("ok"), bool): + cached_health = persisted_health["ok"] + if cached_health is False: + state = "error" + elif cached_health is None: + state = "starting" + api_summary = WorkflowCapabilityStatusResponse(configured=True, state=state) + except Exception as exc: + log.warning("workflow.list.api_status_failed", {"id": workflow_id, "error": str(exc)}) + api_summary = WorkflowCapabilityStatusResponse(configured=False, state="error") + + try: + triggers = await _get_workflow_trigger_defs(workflow_id, workflow_data) + if not triggers: + trigger_summary = WorkflowTriggerCapabilityStatusResponse( + configured=False, + state="unconfigured", + ) + else: + statuses = await default_trigger_runtime.get_workflow_trigger_statuses( + workflow_id, + set_workflow_json_triggers(workflow_data.get("workflowJson") or {}, triggers), + ) + statuses_by_id = { + item.get("triggerId"): item + for item in statuses + if isinstance(item, dict) and item.get("triggerId") + } + trigger_items: List[WorkflowTriggerStatusItemResponse] = [] + for trigger in triggers: + runtime_status = statuses_by_id.get(trigger.id) + raw_state = runtime_status.get("state") if runtime_status else None + trigger_items.append( + WorkflowTriggerStatusItemResponse( + id=trigger.id, + type=trigger.type, + name=trigger.name, + state=_summarize_capability_state(raw_state) if raw_state else "error", + rawState=raw_state, + ) + ) + summarized_states = [item.state for item in trigger_items] + if "error" in summarized_states: + state = "error" + elif "stopped" in summarized_states: + state = "stopped" + elif summarized_states and all(item == "running" for item in summarized_states): + state = "running" + else: + state = "starting" + trigger_summary = WorkflowTriggerCapabilityStatusResponse( + configured=True, + state=state, + count=len(triggers), + items=trigger_items, + ) + except Exception as exc: + log.warning("workflow.list.trigger_status_failed", {"id": workflow_id, "error": str(exc)}) + trigger_summary = WorkflowTriggerCapabilityStatusResponse(configured=True, state="error") + + return WorkflowIntegrationStatusResponse(api=api_summary, trigger=trigger_summary) + + +async def _record_api_publish_failure( + workflow_id: str, + existing_service: Dict[str, Any], + error: Exception, +) -> None: + """Mark a failed restart only when no prior runtime survived the attempt.""" + if not existing_service: + return + try: + runtime = await WorkflowStore.kv_get(_runtime_key_main(workflow_id)) + if isinstance(runtime, dict) and runtime: + return + failed_service = { + **existing_service, + "status": "error", + "lastStartError": str(error), + "lastStartAttemptAt": int(time.time() * 1000), + "health": {"ok": False, "reason": "publish_failed"}, + } + await WorkflowStore.kv_put(_api_service_key(workflow_id), failed_service) + except Exception as persist_error: + log.warning( + "workflow.api.publish_failure_persist_failed", + {"id": workflow_id, "error": str(persist_error)}, + ) + + async def reconcile_published_workflow_api_services() -> Dict[str, int]: """Restart persisted workflow API services after the main server restarts.""" stats = {"checked": 0, "healthy": 0, "restarted": 0, "failed": 0, "skipped": 0} @@ -2315,6 +2624,7 @@ async def reconcile_published_workflow_api_services() -> Dict[str, int]: image=service.get("image") or None, driver=_service_driver_from_record(service), api_key=service.get("apiKey") or None, + port=resolve_service_port(service), ) service_url = active_record.get("serviceUrl", "") @@ -2329,6 +2639,7 @@ async def reconcile_published_workflow_api_services() -> Dict[str, int]: "containerName": active_record.get("containerName", ""), "driver": active_record.get("driver") or service.get("driver"), "image": active_record.get("image") or service.get("image"), + "port": resolve_service_port(active_record) or resolve_service_port(service), "restartedAt": int(time.time() * 1000), } ) @@ -2357,6 +2668,7 @@ class WorkflowServiceResponse(BaseModel): containerName: Optional[str] = None driver: Optional[Literal["local", "docker"]] = None image: Optional[str] = None + port: Optional[int] = None class KafkaConfigRequest(BaseModel): @@ -2450,6 +2762,7 @@ async def publish_workflow_as_api( Writes the workflow JSON to disk, registers it with the workflow center, starts the selected runtime, and returns the service URL and generated API key. """ + existing_service: Dict[str, Any] = {} try: data, now_ms = await _prepare_workflow_api_registry(workflow_id) @@ -2458,6 +2771,7 @@ async def publish_workflow_as_api( # enforce the key returned to callers. existing_service = await WorkflowStore.kv_get(_api_service_key(workflow_id)) or {} api_key = existing_service.get("apiKey") or (uuid.uuid4().hex + uuid.uuid4().hex) + requested_port = req.port if req and req.port is not None else resolve_service_port(existing_service) # Use center.py to publish the selected runtime. active_record = await publish_workflow( @@ -2465,6 +2779,7 @@ async def publish_workflow_as_api( image=req.image if req else None, driver=req.driver if req else None, api_key=api_key, + port=requested_port, ) service_url = active_record.get("serviceUrl", "") @@ -2472,6 +2787,7 @@ async def publish_workflow_as_api( container_name = active_record.get("containerName", "") driver = active_record.get("driver") or (req.driver if req else None) image = active_record.get("image") or (req.image if req else None) + host_port = resolve_service_port(active_record) or requested_port service_info = { "workflowId": workflow_id, @@ -2484,6 +2800,7 @@ async def publish_workflow_as_api( "containerName": container_name, "driver": driver, "image": image, + "port": host_port, } await WorkflowStore.kv_put(_api_service_key(workflow_id), service_info) @@ -2491,9 +2808,13 @@ async def publish_workflow_as_api( return service_info except WorkflowNotFoundError as e: raise HTTPException(status_code=404, detail=str(e)) + except WorkflowPortUnavailableError as e: + await _record_api_publish_failure(workflow_id, existing_service, e) + raise HTTPException(status_code=409, detail=str(e)) except HTTPException: raise except WorkflowCenterError as e: + await _record_api_publish_failure(workflow_id, existing_service, e) log.error("workflow.publish.center_error", {"id": workflow_id, "error": str(e)}) raise HTTPException(status_code=500, detail=f"发布失败: {str(e)}") except Exception as e: @@ -2536,7 +2857,8 @@ async def get_workflow_service(workflow_id: str): Returns null if not published. """ try: - return await WorkflowStore.kv_get(_api_service_key(workflow_id)) # None / null if not found + service = await WorkflowStore.kv_get(_api_service_key(workflow_id)) + return await _normalize_listed_api_service(_api_service_key(workflow_id), service) except Exception as e: log.error("workflow.service.get.error", {"id": workflow_id, "error": str(e)}) raise HTTPException(status_code=500, detail=f"Failed to get service info: {str(e)}") @@ -2559,6 +2881,12 @@ async def delete_workflow_service(workflow_id: str): await WorkflowStore.kv_remove(_api_service_key(workflow_id)) except Exception: pass + else: + registry_key = f"{_REGISTRY_PREFIX_MAIN}{workflow_id}" + registry = await WorkflowStore.kv_get(registry_key) + if isinstance(registry, dict) and "servicePort" in registry: + registry.pop("servicePort", None) + await WorkflowStore.kv_put(registry_key, registry) log.info("workflow.api.service_deleted", {"id": workflow_id}) return {"ok": True, "workflowId": workflow_id} diff --git a/flocks/session/features/memory.py b/flocks/session/features/memory.py index 7c2370bfb..418531b5b 100644 --- a/flocks/session/features/memory.py +++ b/flocks/session/features/memory.py @@ -66,7 +66,7 @@ async def initialize(self) -> bool: memory_config_dict = config.memory if hasattr(config, 'memory') and config.memory else None if not memory_config_dict: - log.warn("session.memory.no_config", {"session_id": self.session_id}) + log.info("session.memory.no_config", {"session_id": self.session_id}) memory_config = MemoryConfig(enabled=True) else: if isinstance(memory_config_dict, dict): diff --git a/flocks/session/lifecycle/retry.py b/flocks/session/lifecycle/retry.py index c220aff1a..6e754ee77 100644 --- a/flocks/session/lifecycle/retry.py +++ b/flocks/session/lifecycle/retry.py @@ -209,6 +209,8 @@ def retryable(error: Dict[str, Any]) -> Optional[str]: def is_connection_error(error: Dict[str, Any]) -> bool: """Return True for provider/model connection failures.""" data = error.get("data", {}) + if data.get("isConnectionError") is True: + return True message = data.get("message") or error.get("message", "") if not isinstance(message, str): return False diff --git a/flocks/session/policy.py b/flocks/session/policy.py index a1870fb5d..998779b70 100644 --- a/flocks/session/policy.py +++ b/flocks/session/policy.py @@ -56,13 +56,36 @@ def _has_no_owner(session: "SessionInfo") -> bool: return not session.owner_user_id and not session.owner_username @classmethod - def is_shared(cls, session: "SessionInfo") -> bool: + def is_shared( + cls, + session: "SessionInfo", + shared_project_ids: Optional[set[str]] = None, + ) -> bool: """Whether the session is explicitly shared to all local users. Ownerless sessions are a legacy / unauthenticated compatibility state, not a sharing state. The UI badge should only reflect explicit sharing. """ - return cls.is_local_shared(session) + return cls.is_local_shared(session) or cls.is_project_shared( + session, + shared_project_ids, + ) + + @staticmethod + def is_project_shared( + session: "SessionInfo", + shared_project_ids: Optional[set[str]] = None, + ) -> bool: + """Whether the session belongs to a locally shared project.""" + + if shared_project_ids is not None: + return session.project_id in shared_project_ids + try: + from flocks.project.project import Project + + return Project.is_local_shared(session.owner_user_id, session.project_id) + except Exception: + return False @staticmethod def _shared_read_user_ids(session: "SessionInfo") -> set[str]: @@ -75,17 +98,31 @@ def _shared_read_user_ids(session: "SessionInfo") -> set[str]: return {str(item) for item in raw if item} @classmethod - def is_shared_read_only(cls, session: "SessionInfo", user: Optional["AuthUser"]) -> bool: + def is_shared_read_only( + cls, + session: "SessionInfo", + user: Optional["AuthUser"], + shared_project_ids: Optional[set[str]] = None, + ) -> bool: if user is None: return False if cls.is_owner(session, user): return False - if cls.is_local_shared(session): + if cls.is_local_shared(session) or cls.is_project_shared( + session, + shared_project_ids, + ): return True return user.id in cls._shared_read_user_ids(session) @classmethod - def can_read(cls, session: "SessionInfo", user: Optional["AuthUser"] = None) -> bool: + def can_read( + cls, + session: "SessionInfo", + user: Optional["AuthUser"] = None, + *, + shared_project_ids: Optional[set[str]] = None, + ) -> bool: """ Whether the session should be visible in listings / fetch. @@ -100,7 +137,7 @@ def can_read(cls, session: "SessionInfo", user: Optional["AuthUser"] = None) -> return True if cls._has_no_owner(session) and cls.is_admin(resolved): return True - return cls.is_shared_read_only(session, resolved) + return cls.is_shared_read_only(session, resolved, shared_project_ids) @classmethod def can_write(cls, session: "SessionInfo", user: Optional["AuthUser"] = None) -> bool: diff --git a/flocks/session/prompt.py b/flocks/session/prompt.py index d36a961e3..56466f8cf 100644 --- a/flocks/session/prompt.py +++ b/flocks/session/prompt.py @@ -258,6 +258,16 @@ def environment_stable( ) -> List[str]: """Build stable workspace metadata that should stay cache-friendly.""" working_dir = directory or os.getcwd() + source_code_dir = str( + Path( + os.getenv( + "FLOCKS_REPO_ROOT", + str(Path(__file__).resolve().parents[2]), + ) + ) + .expanduser() + .resolve() + ) is_git = vcs == "git" from flocks.workspace.manager import WorkspaceManager @@ -267,7 +277,8 @@ def environment_stable( env_info = [ "Here is some useful information about the environment you are running in:", "", - f" Source code directory: {working_dir}", + f" flocks source code directory: {source_code_dir}", + f" current working directory: {working_dir}", f" Workspace outputs directory: {outputs_dir}", f" Is directory a git repo: {'yes' if is_git else 'no'}", f" Platform: {platform.system().lower()}", diff --git a/flocks/session/runner.py b/flocks/session/runner.py index e9c10b5ae..69aec5144 100644 --- a/flocks/session/runner.py +++ b/flocks/session/runner.py @@ -20,10 +20,13 @@ from typing import Optional, Dict, Any, List, Callable, Awaitable, Tuple from dataclasses import dataclass, field +import httpcore +import httpx + from flocks.utils.log import Log from flocks.utils.id import Identifier from flocks.session.session import Session, SessionInfo -from flocks.session.message import Message, MessageInfo, MessageRole +from flocks.session.message import Message, MessageInfo, MessageRole, TextPart from flocks.session.prompt import SessionPrompt from flocks.session.core.status import SessionStatus, SessionStatusRetry, SessionStatusBusy from flocks.session.core.defaults import ( @@ -107,6 +110,13 @@ def _annotate_with_provider_version(tool_info: Any, description: Optional[str]) # spurious failures in those cases. LLM_STREAM_ONGOING_CHUNK_TIMEOUT_S = 300 +_RETRYABLE_TRANSPORT_EXCEPTIONS = ( + httpx.TransportError, + httpcore.NetworkError, + httpcore.ProtocolError, + httpcore.TimeoutException, +) + _WORKFLOW_NODE_REF_RE = re.compile(r"^@@node:([^|\n]+)\|([^\n]+)\n?([\s\S]*)$") @@ -175,6 +185,18 @@ async def _iter_with_chunk_timeout( pass +def _find_retryable_transport_exception(exception: Exception) -> Optional[Exception]: + """Return a retryable HTTP transport error from an exception chain.""" + current: Optional[BaseException] = exception + seen: set[int] = set() + while current is not None and id(current) not in seen: + seen.add(id(current)) + if isinstance(current, _RETRYABLE_TRANSPORT_EXCEPTIONS): + return current + current = current.__cause__ or current.__context__ + return None + + @dataclass class ToolCall: """Tool call from LLM response.""" @@ -298,6 +320,139 @@ def _build_tool_loop_halt_message( "summarize the blocker, or answer directly instead of repeating the exact same call." ) + @staticmethod + def _build_session_error_dict( + message: str, + *, + name: str = "SessionError", + display_message: Optional[str] = None, + data: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + error_data = {"message": message} + if data: + error_data.update(data) + if display_message: + error_data["displayMessage"] = display_message + return { + "name": name, + "message": message, + "data": error_data, + } + + async def _publish_assistant_error_message( + self, + assistant_msg: MessageInfo, + *, + agent_name: str, + parent_id: str, + error_dict: Dict[str, Any], + text_part: Optional[TextPart] = None, + ) -> None: + if not self.callbacks.event_publish_callback: + return + now_ms = int(time.time() * 1000) + await self.callbacks.event_publish_callback("message.updated", { + "info": { + "id": assistant_msg.id, + "sessionID": self.session.id, + "role": "assistant", + "time": {"created": now_ms, "completed": now_ms}, + "parentID": parent_id, + "modelID": self.model_id, + "providerID": self.provider_id, + "agent": agent_name, + "mode": agent_name, + "finish": "error", + "error": error_dict, + "tokens": {"input": 0, "output": 0, "reasoning": 0, "cache": {"read": 0, "write": 0}}, + } + }) + if text_part is not None: + await self.callbacks.event_publish_callback("message.part.updated", { + "part": { + "id": text_part.id, + "messageID": text_part.messageID, + "sessionID": text_part.sessionID, + "type": "text", + "text": text_part.text, + } + }) + + async def _create_error_assistant_message( + self, + *, + last_user: MessageInfo, + agent: AgentInfo, + error_message: str, + error_dict: Dict[str, Any], + visible_text: Optional[str] = None, + ) -> MessageInfo: + text = visible_text or error_message + part_id = Identifier.ascending("part") + now_ms = int(time.time() * 1000) + assistant_msg = await Message.create( + session_id=self.session.id, + role=MessageRole.ASSISTANT, + content=text, + agent=agent.name, + model_id=self.model_id, + provider_id=self.provider_id, + parent_id=last_user.id, + finish="error", + error=error_dict, + time={"created": now_ms, "completed": now_ms}, + part_id=part_id, + ) + parts = await Message.parts(assistant_msg.id, self.session.id) + text_part = next( + ( + part for part in parts + if isinstance(part, TextPart) and part.id == part_id + ), + TextPart( + id=part_id, + sessionID=self.session.id, + messageID=assistant_msg.id, + type="text", + text=text, + ), + ) + await self._publish_assistant_error_message( + assistant_msg, + agent_name=agent.name, + parent_id=last_user.id, + error_dict=error_dict, + text_part=text_part, + ) + return assistant_msg + + async def _ensure_visible_error_text( + self, + assistant_msg: MessageInfo, + error_message: str, + ) -> Optional[TextPart]: + if not error_message: + return None + + parts = await Message.parts(assistant_msg.id, self.session.id) + for part in parts: + if getattr(part, "type", None) == "text" and str(getattr(part, "text", "") or "").strip(): + return None + + text_part = next((part for part in parts if getattr(part, "type", None) == "text"), None) + if isinstance(text_part, TextPart): + text_part.text = error_message + else: + text_part = TextPart( + id=Identifier.ascending("part"), + sessionID=self.session.id, + messageID=assistant_msg.id, + type="text", + text=error_message, + ) + stored = await Message.store_part(self.session.id, assistant_msg.id, text_part) + return stored if isinstance(stored, TextPart) else text_part + def _update_tool_loop_guard( self, result: StepResult, @@ -982,9 +1137,22 @@ async def _process_step( provider = Provider.get(self.provider_id) if not provider: error = f"Provider {self.provider_id} not found" + error_dict = self._build_session_error_dict( + error, + name="ProviderUnavailableError", + display_message="Model is unavailable. Please check the provider connection and model configuration.", + data={"providerID": self.provider_id, "modelID": self.model_id}, + ) + await self._create_error_assistant_message( + last_user=last_user, + agent=agent, + error_message=error, + error_dict=error_dict, + visible_text=error_dict["data"]["displayMessage"], + ) if self.callbacks.on_error: - await self.callbacks.on_error(error) - return StepResult(action="stop", error=error) + await self.callbacks.on_error(error_dict["data"]["displayMessage"]) + return StepResult(action="stop", error=error_dict["data"]["displayMessage"]) # Apply config-based provider options (api_key/base_url) try: @@ -997,9 +1165,22 @@ async def _process_step( if not provider.is_configured(): error = f"Provider {self.provider_id} not configured" + error_dict = self._build_session_error_dict( + error, + name="ProviderConfigurationError", + display_message="Model is unavailable. Please check the provider connection and model configuration.", + data={"providerID": self.provider_id, "modelID": self.model_id}, + ) + await self._create_error_assistant_message( + last_user=last_user, + agent=agent, + error_message=error, + error_dict=error_dict, + visible_text=error_dict["data"]["displayMessage"], + ) if self.callbacks.on_error: - await self.callbacks.on_error(error) - return StepResult(action="stop", error=error) + await self.callbacks.on_error(error_dict["data"]["displayMessage"]) + return StepResult(action="stop", error=error_dict["data"]["displayMessage"]) # Build prompts and tools tools_started_at = time.perf_counter() @@ -1301,6 +1482,14 @@ async def device_asset_prompt_factory() -> Optional[str]: error=empty_error_dict, finish="error", ) + text_part = await self._ensure_visible_error_text(assistant_msg, empty_error_msg) + await self._publish_assistant_error_message( + assistant_msg, + agent_name=agent.name, + parent_id=last_user.id, + error_dict=empty_error_dict, + text_part=text_part, + ) return StepResult(action="stop", error=empty_error_msg) tool_loop_guard = self._update_tool_loop_guard( @@ -1341,18 +1530,20 @@ async def device_asset_prompt_factory() -> Optional[str]: except Exception as e: error_attempt += 1 - log.error("runner.step.error", { + error_log_context = { "error": str(e), - "attempt": error_attempt, - }) - + "error_type": type(e).__name__, + "error_module": type(e).__module__, + "error_repr": repr(e)[:1000], + } # Convert exception to error dict for retry check error_dict = self._exception_to_error_dict(e) - + # Check if retryable retry_message = SessionRetry.retryable(error_dict) + will_retry = retry_message is not None and error_attempt <= MAX_ERROR_RETRIES - if retry_message is not None and error_attempt <= MAX_ERROR_RETRIES: + if will_retry: # Error is retryable and we have budget left delay_ms = SessionRetry.delay(error_attempt, error_dict) # Always cap the sleep to RETRY_MAX_DELAY_NO_HEADERS so a @@ -1361,8 +1552,9 @@ async def device_asset_prompt_factory() -> Optional[str]: from flocks.session.lifecycle.retry import RETRY_MAX_DELAY_NO_HEADERS delay_ms = min(delay_ms, RETRY_MAX_DELAY_NO_HEADERS) next_retry_time = int(asyncio.get_event_loop().time() * 1000) + delay_ms - - log.info("runner.step.retry", { + + log.warn("runner.step.retry", { + **error_log_context, "attempt": error_attempt, "delay_ms": delay_ms, "reason": retry_message, @@ -1388,12 +1580,15 @@ async def device_asset_prompt_factory() -> Optional[str]: # Error is not retryable, or retry budget exhausted if retry_message is not None: log.error("runner.step.max_retries_exceeded", { - "error": str(e), + **error_log_context, "attempt": error_attempt, "max_retries": MAX_ERROR_RETRIES, }) else: - log.error("runner.step.not_retryable", {"error": str(e)}) + log.error("runner.step.not_retryable", { + **error_log_context, + "attempt": error_attempt, + }) final_error_message = str(e) if SessionRetry.is_connection_error(error_dict): @@ -1410,6 +1605,14 @@ async def device_asset_prompt_factory() -> Optional[str]: error=error_dict, finish="error", ) + text_part = await self._ensure_visible_error_text(assistant_msg, final_error_message) + await self._publish_assistant_error_message( + assistant_msg, + agent_name=agent.name, + parent_id=last_user.id, + error_dict=error_dict, + text_part=text_part, + ) return StepResult(action="stop", error=final_error_message) @@ -1796,8 +1999,23 @@ def _exception_to_error_dict(self, exception: Exception) -> Dict[str, Any]: "message": str(exception), "data": { "message": str(exception), + "exceptionType": type(exception).__name__, + "exceptionModule": type(exception).__module__, } } + + transport_exception = _find_retryable_transport_exception(exception) + if transport_exception is not None: + transport_type = type(transport_exception).__name__ + error_dict["name"] = "APIError" + error_dict["data"].update({ + "message": str(exception) or f"Transport connection error ({transport_type})", + "isRetryable": True, + "isConnectionError": True, + "displayMessage": CONNECTION_ERROR_DISPLAY_MESSAGE, + "transportExceptionType": transport_type, + "transportExceptionModule": type(transport_exception).__module__, + }) # Check if it's an API error with specific attributes if hasattr(exception, 'status_code'): diff --git a/flocks/session/session.py b/flocks/session/session.py index c9807b8af..b83d898c2 100644 --- a/flocks/session/session.py +++ b/flocks/session/session.py @@ -11,15 +11,19 @@ from datetime import datetime from pydantic import BaseModel, Field, ConfigDict -# Sentinel for explicitly setting a field to None via Session.update() -_UNSET = object() - -from flocks.auth.context import get_current_auth_user +from flocks.auth.context import ( + get_current_auth_user, + reset_current_auth_user, + set_current_auth_user, +) from flocks.storage.storage import Storage from flocks.utils.log import Log from flocks.utils.id import Identifier from flocks.session.message import Message, MessageInfo, AssistantMessageInfo +# Sentinel for explicitly setting a field to None via Session.update() +_UNSET = object() + log = Log.create(service="session") @@ -160,6 +164,9 @@ def _is_owned_by_auth_user(session: SessionInfo, auth_user) -> bool: @classmethod def _sync_list_cache(cls, session: SessionInfo) -> None: """Keep the in-memory list cache aligned with session mutations.""" + from flocks.project.project import Project + + Project.invalidate_session_stats() if cls._all_sessions_cache is None: return @@ -483,6 +490,16 @@ async def list_all(cls) -> List[SessionInfo]: except Exception as e: log.error("session.list_all.error", {"error": str(e)}) return [] + + @classmethod + async def list_all_unfiltered(cls) -> List[SessionInfo]: + """List all sessions for callers that apply an explicit access policy.""" + + token = set_current_auth_user(None) + try: + return await cls.list_all() + finally: + reset_current_auth_user(token) @classmethod async def update( diff --git a/flocks/session/session_loop.py b/flocks/session/session_loop.py index 35faa5678..adb1e112e 100644 --- a/flocks/session/session_loop.py +++ b/flocks/session/session_loop.py @@ -297,6 +297,7 @@ async def run( model_id: Optional[str] = None, agent_name: Optional[str] = None, callbacks: Optional[LoopCallbacks] = None, + working_directory: Optional[str] = None, ) -> LoopResult: """ Run session loop @@ -339,6 +340,8 @@ async def run( action="error", error=f"Session {session_id} not found", ) + if working_directory: + session = session.model_copy(update={"directory": working_directory}) # Resolve model when not explicitly provided if not provider_id or not model_id: diff --git a/flocks/skill/skill.py b/flocks/skill/skill.py index db5f2e1e5..9ac60b57b 100644 --- a/flocks/skill/skill.py +++ b/flocks/skill/skill.py @@ -198,30 +198,47 @@ class Skill: Skill discovery and management. Discovers SKILL.md files from (lowest → highest priority): - - .flocks dirs (global + project-level) - - .claude dirs (global ~/.claude + project-level) - - ~/.flocks (global user-level) - - /.flocks (project-level, wins on collision) + - .claude dirs + - global built-in and project-bundled .flocks dirs + - ~/.flocks/plugins (user-level customizations, wins on collision) """ _cache: Optional[Dict[str, SkillInfo]] = None + _cache_key: Optional[tuple[str, str]] = None _cache_lock = threading.Lock() _cache_generation = 0 - _cache_loads: Dict[int, threading.Event] = {} + _cache_loads: Dict[tuple[int, tuple[str, str]], threading.Event] = {} + + @staticmethod + def _source_root() -> Path: + """Return the Flocks source installation root.""" + return Path(__file__).resolve().parents[2] + + @staticmethod + def _cache_context() -> tuple[str, str]: + """Return the normalized request context that scopes discovery.""" + current_dir = Instance.get_directory() or os.getcwd() + worktree = Instance.get_worktree() or current_dir + return ( + os.path.normcase(os.path.realpath(current_dir)), + os.path.normcase(os.path.realpath(worktree)), + ) @classmethod def _all_sync(cls) -> List[SkillInfo]: - """Discover skills once per cache generation without stale refills.""" + """Discover skills once per request context and cache generation.""" + context_key = cls._cache_context() while True: with cls._cache_lock: - if cls._cache is not None: + if cls._cache is not None and cls._cache_key == context_key: return list(cls._cache.values()) generation = cls._cache_generation - completion = cls._cache_loads.get(generation) + load_key = (generation, context_key) + completion = cls._cache_loads.get(load_key) should_discover = completion is None if completion is None: completion = threading.Event() - cls._cache_loads[generation] = completion + cls._cache_loads[load_key] = completion if not should_discover: completion.wait() @@ -231,16 +248,21 @@ def _all_sync(cls) -> List[SkillInfo]: discovered = cls._discover() except BaseException: with cls._cache_lock: - cls._cache_loads.pop(generation, None) + cls._cache_loads.pop(load_key, None) completion.set() raise with cls._cache_lock: - if generation == cls._cache_generation and cls._cache is None: + if generation == cls._cache_generation: cls._cache = discovered - cls._cache_loads.pop(generation, None) + cls._cache_key = context_key + cls._cache_loads.pop(load_key, None) completion.set() - if generation == cls._cache_generation and cls._cache is not None: + if ( + generation == cls._cache_generation + and cls._cache is not None + and cls._cache_key == context_key + ): return list(cls._cache.values()) # The cache was invalidated while discovery was running. Repeat # against the new generation instead of publishing stale data. @@ -392,11 +414,21 @@ def _scan_directory( skill_info = cls._parse_skill_md(match, source=source) if skill_info: if skill_info.name in skills: - log.warn("skill.duplicate", { - "name": skill_info.name, - "existing": skills[skill_info.name].location, - "duplicate": match, - }) + existing = skills[skill_info.name] + if existing.source != skill_info.source: + log.info("skill.override", { + "name": skill_info.name, + "selected": match, + "selected_source": skill_info.source, + "replaced": existing.location, + "replaced_source": existing.source, + }) + else: + log.warn("skill.duplicate", { + "name": skill_info.name, + "existing": existing.location, + "duplicate": match, + }) skills[skill_info.name] = skill_info log.debug("skill.found", { @@ -417,8 +449,10 @@ def _discover(cls) -> Dict[str, SkillInfo]: Scan order (lowest → highest priority): 1. .claude dirs (global ~/.claude + project-level) - 2. ~/.flocks (global user-level, overrides .claude) - 3. /.flocks (project-level, highest priority) + 2. ~/.flocks/skill[s] (global built-ins) + 3. /.flocks (bundled with the Flocks installation) + 4. /.flocks (project-level) + 5. ~/.flocks/plugins (user customizations) Source labels: "flocks" — built-in skills inside .flocks/skills/ directories @@ -442,6 +476,7 @@ def _discover(cls) -> Dict[str, SkillInfo]: "plugins/skills/**/SKILL.md", ) global_flocks = os.path.join(home_dir, ".flocks") + source_flocks = str(cls._source_root() / ".flocks") # 1) .claude directories — lowest priority global_claude = os.path.join(home_dir, ".claude") @@ -450,24 +485,38 @@ def _discover(cls) -> Dict[str, SkillInfo]: for claude_dir in cls._find_dirs_up(".claude", current_dir, worktree): cls._scan_directory(claude_dir, "skills/**/SKILL.md", skills, source="claude") - # 2) Global ~/.flocks — overrides .claude - # Built-in skills: source="flocks"; user-installed plugins: source="user" + # 2) Global built-in skills — overrides .claude if os.path.isdir(global_flocks): for pattern in builtin_patterns: cls._scan_directory(global_flocks, pattern, skills, source="flocks") + + # 3) Source installation .flocks — available regardless of request cwd. + # Keep plugin skills labelled "project" for the current Skill UI's + # built-in/custom classification. + if os.path.normpath(source_flocks) != os.path.normpath(global_flocks): + for pattern in builtin_patterns: + cls._scan_directory(source_flocks, pattern, skills, source="flocks") for pattern in plugin_patterns: - cls._scan_directory(global_flocks, pattern, skills, source="user") + cls._scan_directory(source_flocks, pattern, skills, source="project") - # 3) Project-level .flocks — highest priority + # 4) Project-level .flocks — highest priority # Built-in skills: source="flocks"; project-installed plugins: source="project" for flocks_dir in cls._find_dirs_up(".flocks", current_dir, worktree): - if os.path.normpath(flocks_dir) == os.path.normpath(global_flocks): + if os.path.normpath(flocks_dir) in { + os.path.normpath(global_flocks), + os.path.normpath(source_flocks), + }: continue for pattern in builtin_patterns: cls._scan_directory(flocks_dir, pattern, skills, source="flocks") for pattern in plugin_patterns: cls._scan_directory(flocks_dir, pattern, skills, source="project") + # 5) User-installed plugins — explicit customizations win over project bundles + if os.path.isdir(global_flocks): + for pattern in plugin_patterns: + cls._scan_directory(global_flocks, pattern, skills, source="user") + log.info("skill.discovery.complete", {"count": len(skills), "names": list(skills.keys())}) return skills @@ -543,6 +592,7 @@ def clear_cache(cls) -> None: """Clear the skill cache (for testing or forced refresh)""" with cls._cache_lock: cls._cache = None + cls._cache_key = None cls._cache_generation += 1 log.info("skill.cache.cleared") diff --git a/flocks/tool/agent/delegate_task.py b/flocks/tool/agent/delegate_task.py index 176a8f8ec..88c42d52b 100644 --- a/flocks/tool/agent/delegate_task.py +++ b/flocks/tool/agent/delegate_task.py @@ -458,9 +458,13 @@ async def delegate_task_tool( if not parent_session: return ToolResult(success=False, error="Parent session not found") + from flocks.project.instance import Instance + + runtime_directory = Instance.get_directory() or parent_session.directory + create_kwargs = dict( project_id=parent_session.project_id, - directory=parent_session.directory, + directory=runtime_directory, title=f"{description} (@{agent_to_use} subagent)", parent_id=parent_session.id, agent=agent_to_use, @@ -474,6 +478,8 @@ async def delegate_task_tool( model_pinned=bool(explicit_model), ) created = await Session.create(**create_kwargs) + if ctx.extra.get("workflow_temp_parent") is True: + ctx.extra["workflow_child_session_created"] = True await Message.create( session_id=created.id, role=MessageRole.USER, diff --git a/flocks/tool/channel/channel_message.py b/flocks/tool/channel/channel_message.py index 6a908e8e1..cfed0b452 100644 --- a/flocks/tool/channel/channel_message.py +++ b/flocks/tool/channel/channel_message.py @@ -2,7 +2,7 @@ channel_message tool — sends a message to the messaging channel bound to a given session. Looks up the SessionBinding for the given session_id to automatically resolve -the target channel (WeCom / Weixin / Feishu / DingTalk / Telegram / WhatsApp / Email) +the target channel (WeCom / Weixin / Feishu / DingTalk / Telegram / WhatsApp / Email / Slack) and chat_id, so the caller does not need to specify them manually. @@ -29,6 +29,7 @@ "telegram": ["telegram", "tg", "tele"], "whatsapp": ["whatsapp", "wa"], "email": ["email", "mail", "邮件", "imap", "smtp"], + "slack": ["slack", "sl"], } @@ -136,7 +137,7 @@ async def _http_session_send( description=( "Send a message to the messaging channel bound to a session. " "Channel types: WeCom/企业微信=wecom, Weixin/微信=weixin, Feishu=feishu, DingTalk=dingtalk, " - "Telegram=telegram, WhatsApp=whatsapp, Email/邮件=email. " + "Telegram=telegram, WhatsApp=whatsapp, Email/邮件=email, Slack=slack. " "Resolves the target channel and chat automatically from session_id. " "Use channel_type to target a specific channel when the session has multiple bindings." ), @@ -166,6 +167,7 @@ async def _http_session_send( "telegram", "whatsapp", "email", + "slack", "企微", "企业微信", "微信", @@ -175,7 +177,7 @@ async def _http_session_send( ], description=( "Target channel: wecom=企业微信, weixin=微信, feishu=飞书, dingtalk=钉钉, " - "telegram=Telegram, whatsapp=WhatsApp, or email=邮件. " + "telegram=Telegram, whatsapp=WhatsApp, email=邮件, or slack=Slack. " "Chinese aliases are accepted. " "If omitted and the session has only one binding, that channel is used automatically. " "If omitted and the session has multiple bindings, the message is sent to all of them." diff --git a/flocks/tool/channel/im_send_message.py b/flocks/tool/channel/im_send_message.py index 834218ada..83d1af039 100644 --- a/flocks/tool/channel/im_send_message.py +++ b/flocks/tool/channel/im_send_message.py @@ -23,6 +23,7 @@ "telegram": ["telegram", "tg", "tele"], "whatsapp": ["whatsapp", "wa"], "email": ["email", "mail", "邮件", "imap", "smtp"], + "slack": ["slack", "sl"], } @@ -204,10 +205,10 @@ async def _resolve_target( name="im_send_message", description=( "Resolve a messaging channel target session and optionally send a message. " - "Use this for WeCom/企业微信, Weixin/微信, Feishu, DingTalk, Telegram, WhatsApp, Email/邮件, " + "Use this for WeCom/企业微信, Weixin/微信, Feishu, DingTalk, Telegram, WhatsApp, Email/邮件, Slack, " "or custom channel sessions when the user asks to send a message through a connected channel. " "Use channel_type=wecom for 企业微信, channel_type=weixin for 微信, channel_type=telegram for Telegram, " - "channel_type=whatsapp for WhatsApp, and channel_type=email for 邮件. " + "channel_type=whatsapp for WhatsApp, channel_type=email for 邮件, and channel_type=slack for Slack. " "If session_id is omitted, it uses the current channel session when available, otherwise asks the user to pick one." ), category=ToolCategory.SYSTEM, @@ -230,7 +231,7 @@ async def _resolve_target( required=False, description=( "Optional channel filter, such as wecom=企业微信, weixin=微信, " - "feishu, dingtalk, telegram, whatsapp, email=邮件, or a custom channel id." + "feishu, dingtalk, telegram, whatsapp, email=邮件, slack, or a custom channel id." ), ), ToolParameter( diff --git a/flocks/tool/device/manage_tool.py b/flocks/tool/device/manage_tool.py index d6e8bf20d..963debab8 100644 --- a/flocks/tool/device/manage_tool.py +++ b/flocks/tool/device/manage_tool.py @@ -1,16 +1,25 @@ """Built-in device management tool for Rex. -``device_manage`` is the single system-tool entrypoint for device discovery, -non-secret config updates, and standard connectivity checks. Its +``device_manage`` is the single system-tool entrypoint for device and template +discovery, non-secret config updates, and standard connectivity checks. Its ``connectivity_test`` action reuses the existing device test path, so card state stays consistent with the ``POST /api/devices/{id}/test`` endpoint. """ from __future__ import annotations +import asyncio from typing import Any, Optional -from flocks.tool.device.intake import DeviceNotFoundError, test_device, update_device -from flocks.tool.device.models import DeviceIntegrationUpdate +from flocks.tool.device.intake import ( + DeviceNotFoundError, + create_device, + test_device, + update_device, +) +from flocks.tool.device.models import ( + DeviceIntegrationCreate, + DeviceIntegrationUpdate, +) from flocks.tool.registry import ( ParameterType, ToolCategory, @@ -28,6 +37,8 @@ name="device_manage", description=( "管理已接入安全设备。action=list 用于列出机房、设备、device_id 和工具集；" + "action=list_templates 用于列出已有设备模板、安装状态和配置字段；" + "action=create 用于从已安装模板创建设备实例，仅接受非敏感配置；" "action=update 用于写入/更新已有设备实例的非敏感配置字段；" "action=connectivity_test 用于测试指定设备连通性并更新设备卡片状态。" ), @@ -37,9 +48,37 @@ ToolParameter( name="action", type=ParameterType.STRING, - description="操作类型：list 列出设备；update 更新已有设备非敏感配置；connectivity_test 测试设备连通性。", + description=( + "操作类型：list 列出设备实例；list_templates 列出已有设备模板；" + "create 从已安装模板创建设备实例；update 更新已有设备非敏感配置；" + "connectivity_test 测试设备连通性。" + ), required=True, - enum=["list", "update", "connectivity_test"], + enum=[ + "list", + "list_templates", + "create", + "update", + "connectivity_test", + ], + ), + ToolParameter( + name="storage_key", + type=ParameterType.STRING, + description="已安装设备模板的 storage_key。action=create 时必填。", + required=False, + ), + ToolParameter( + name="device_name", + type=ParameterType.STRING, + description="新设备实例名称。action=create 时必填。", + required=False, + ), + ToolParameter( + name="group_id", + type=ParameterType.STRING, + description="目标机房 ID。action=create 时可选，默认使用默认机房。", + required=False, ), ToolParameter( name="device_id", @@ -51,15 +90,16 @@ name="fields", type=ParameterType.OBJECT, description=( - "要更新的非敏感设备配置字段，例如 {\"base_url\":\"https://device.local\"}。" - "禁止传入 api_key、secret、password、token、cookie、auth_state 等敏感字段。" + "要创建或更新的非敏感设备配置字段，例如 " + "{\"base_url\":\"https://device.local\"}。" + "禁止传入 api_key、secret、password、token、cookie 等敏感字段。" ), required=False, ), ToolParameter( name="verify_ssl", type=ParameterType.BOOLEAN, - description="是否开启 SSL 证书验证。仅 action=update 时使用。", + description="是否开启 SSL 证书验证。action=create 或 update 时使用。", required=False, ), ], @@ -67,21 +107,37 @@ async def device_manage( ctx: ToolContext, action: str, + storage_key: Optional[str] = None, + device_name: Optional[str] = None, + group_id: Optional[str] = None, device_id: Optional[str] = None, fields: Optional[dict[str, Any]] = None, verify_ssl: Optional[bool] = None, ) -> ToolResult: - """List devices, update non-secret config, or run a standard probe.""" + """List devices/templates, update non-secret config, or run a probe.""" normalized_action = (action or "").strip() if normalized_action == "list": return await _list_devices() + if normalized_action == "list_templates": + return await _list_device_templates() + if normalized_action == "create": + return await _create_device_from_template( + storage_key, + device_name, + group_id, + fields, + verify_ssl, + ) if normalized_action == "update": return await _update_device_config(ctx, device_id, fields, verify_ssl) if normalized_action == "connectivity_test": return await _connectivity_test(ctx, device_id) return ToolResult( success=False, - error="未知 action，请使用 list、update 或 connectivity_test。", + error=( + "未知 action，请使用 list、list_templates、create、update 或 " + "connectivity_test。" + ), ) @@ -104,6 +160,131 @@ async def _list_devices() -> ToolResult: return ToolResult(success=False, error=f"查询设备列表失败: {exc}") +async def _list_device_templates() -> ToolResult: + """Return the same device template index consumed by the access page.""" + try: + from flocks.tool.device.plugin_index import list_device_templates + + templates = await asyncio.to_thread(list_device_templates, refresh=False) + return ToolResult( + success=True, + output=[template.model_dump(mode="json") for template in templates], + metadata={ + "template_count": len(templates), + "installed_count": sum(template.installed for template in templates), + }, + title="设备模板列表", + ) + except Exception as exc: + log.warn("tool.device_manage.list_templates_failed", {"error": str(exc)}) + return ToolResult(success=False, error=f"查询设备模板失败: {exc}") + + +async def _create_device_from_template( + storage_key: Optional[str], + device_name: Optional[str], + group_id: Optional[str], + fields: Optional[dict[str, Any]], + verify_ssl: Optional[bool], +) -> ToolResult: + target_storage_key = (storage_key or "").strip() + name = (device_name or "").strip() + if not target_storage_key: + return ToolResult(success=False, error="action=create 时 storage_key 不能为空。") + if not name: + return ToolResult(success=False, error="action=create 时 device_name 不能为空。") + + try: + from flocks.tool.device.plugin_index import list_device_templates + + templates = await asyncio.to_thread(list_device_templates, refresh=False) + template = next( + (item for item in templates if item.storage_key == target_storage_key), + None, + ) + except Exception as exc: + log.warn("tool.device_manage.create_template_lookup_failed", {"error": str(exc)}) + return ToolResult(success=False, error=f"查询设备模板失败: {exc}") + + if template is None: + return ToolResult( + success=False, + error=( + f"未找到 storage_key={target_storage_key!r} 的设备模板，" + "请先调用 device_manage(action='list_templates')。" + ), + ) + if not template.installed: + return ToolResult( + success=False, + error=( + f"模板 {template.name!r} 尚未安装。请先在 FlockHub 安装 " + f"plugin_id={template.plugin_id!r}，然后重新查询模板。" + ), + ) + + schema = { + str(field.get("key") or "").strip(): field + for field in template.credential_schema + if str(field.get("key") or "").strip() + } + normalized_fields = { + str(key).strip(): "" if value is None else str(value) + for key, value in (fields or {}).items() + } + unknown = sorted(set(normalized_fields).difference(schema)) + if unknown: + return ToolResult( + success=False, + error="模板未声明字段：" + ", ".join(f"`{key}`" for key in unknown), + ) + + sensitive_fields = sorted( + key + for key, field in schema.items() + if field.get("storage") == "secret" + or field.get("sensitive") is True + or field.get("input_type") == "password" + or key.lower() in _SENSITIVE_FIELD_KEYS + ) + rejected = sorted(set(normalized_fields).intersection(sensitive_fields)) + if rejected: + return ToolResult( + success=False, + error=( + "拒绝写入敏感字段：" + + ", ".join(f"`{key}`" for key in rejected) + + "。请创建设备后在设备接入页面填写。" + ), + ) + + body = DeviceIntegrationCreate( + name=name, + storage_key=template.storage_key, + service_id=template.service_id, + group_id=(group_id or "").strip() or None, + verify_ssl=bool(verify_ssl) if verify_ssl is not None else False, + fields=normalized_fields, + ) + try: + created = await create_device(body) + except ValueError as exc: + return ToolResult(success=False, error=f"创建设备失败: {exc}") + except Exception as exc: + log.warn("tool.device_manage.create_failed", {"error": str(exc)}) + return ToolResult(success=False, error=f"创建设备失败: {exc}") + + output = created.model_dump(mode="json") + output["device_id"] = created.id + output["sensitive_fields_to_complete"] = sensitive_fields + return ToolResult( + success=True, + output=output, + metadata={"device_id": created.id, "sensitive_fields": sensitive_fields}, + title="设备已创建", + ) + + _SENSITIVE_FIELD_KEYS = frozenset( { "api_key", @@ -116,7 +297,6 @@ async def _list_devices() -> ToolResult: "access_token", "refresh_token", "cookie", - "auth_state", } ) @@ -144,7 +324,7 @@ def _normalize_update_fields( return {}, ( "拒绝通过 device_manage(action='update') 写入敏感字段：" + ", ".join(f"`{key}`" for key in sorted(rejected)) - + "。请在设备接入页面的配置表单中填写密钥、密码、Token、Cookie 或 auth_state。" + + "。请在设备接入页面的配置表单中填写密钥、密码、Token 或 Cookie。" ) return normalized, None diff --git a/flocks/tool/device/startup.py b/flocks/tool/device/startup.py index 82a32920a..7b1b0454b 100644 --- a/flocks/tool/device/startup.py +++ b/flocks/tool/device/startup.py @@ -7,6 +7,8 @@ """ from __future__ import annotations +import aiosqlite + from flocks.storage.storage import Storage from flocks.utils.log import Log @@ -40,6 +42,7 @@ async def _heal_stale_service_ids() -> None: from flocks.tool.device.store import storage_key_to_service_id async with Storage.connect(Storage.get_db_path()) as db: + db.row_factory = aiosqlite.Row cur = await db.execute("SELECT id, storage_key, service_id FROM device_integrations") rows = await cur.fetchall() updates: list[tuple[str, str]] = [] diff --git a/flocks/tool/file/glob.py b/flocks/tool/file/glob.py index c9f5eb8a0..3a5e9b769 100644 --- a/flocks/tool/file/glob.py +++ b/flocks/tool/file/glob.py @@ -24,6 +24,7 @@ # Constants MAX_FILES = 100 +_ripgrep_fallback_logged = False # Description matching Flocks' glob.txt @@ -195,7 +196,10 @@ async def glob_tool( 'mtime': mtime }) else: - log.warn("glob.ripgrep_not_found", {"fallback": "python_glob"}) + global _ripgrep_fallback_logged + if not _ripgrep_fallback_logged: + _ripgrep_fallback_logged = True + log.info("glob.ripgrep_not_found", {"fallback": "python_glob"}) for filepath in fallback_glob(search_path, pattern): if len(files) >= MAX_FILES: diff --git a/flocks/tool/registry.py b/flocks/tool/registry.py index d348a8387..b172a847d 100644 --- a/flocks/tool/registry.py +++ b/flocks/tool/registry.py @@ -376,6 +376,13 @@ def _normalize_param_key(name: str) -> str: return "".join(ch for ch in str(name).lower() if ch.isalnum()) +def _api_service_storage_key(provider: str) -> str: + """Resolve the single persisted key used for an API-like provider.""" + from flocks.config.api_versioning import versioned_storage_key_for + + return versioned_storage_key_for(provider) or provider + + _SCOPED_SCHEMA_ALIASES: Dict[str, Dict[str, str]] = { # SkyEye historically surfaced "威胁级别", which some callers guessed as # threat_level. Keep the schema canonical on hazard_level, but accept the @@ -516,7 +523,7 @@ async def execute(self, ctx: ToolContext, **kwargs) -> ToolResult: # Validate required parameters for required_param in schema.required: if required_param not in effective_kwargs: - log.error("tool.execute.missing_param", { + log.warn("tool.execute.missing_param", { "tool": self.info.name, "missing": required_param, "provided": list(effective_kwargs.keys()), @@ -607,7 +614,11 @@ class ToolRegistry: _dynamic_modules: Dict[str, str] = {} _dynamic_tools_by_module: Dict[str, List[str]] = {} _plugin_tool_names: List[str] = [] + # Source-level plugin failures observed by the last load. They remain + # diagnostic only because the plugin loader already isolates each source. + _plugin_load_errors: List[str] = [] _revision: int = 0 + _config_state_token: Optional[tuple[str, int, int, int]] = None _failure_state: Dict[str, Dict[str, Any]] = {} _failure_disable_threshold: int = 3 _init_lock = threading.Lock() @@ -668,22 +679,80 @@ def register(cls, tool: Tool) -> None: def revision(cls) -> int: """Return the current registry revision. - The revision is bumped when plugin or dynamic tools are reloaded so - long-lived session caches can detect toolset changes. + The revision is bumped when tool membership or enabled state changes + so long-lived session caches can detect toolset changes. """ + cls._ensure_initialized() + cls._sync_configured_enabled_states() with cls._refresh_lock: return cls._revision + @classmethod + def _current_config_state_token(cls) -> tuple[str, int, int, int]: + """Return a cross-process token for the active flocks config file.""" + from flocks.config.config_writer import ConfigWriter + + path = ConfigWriter._get_config_path() + try: + stat = path.stat() + return str(path), stat.st_ino, stat.st_mtime_ns, stat.st_size + except OSError: + return str(path), -1, -1, -1 + + @classmethod + def _sync_configured_enabled_states(cls) -> None: + """Apply config changes written by another server worker. + + Registry revisions and ``ToolInfo.enabled`` values are process-local, + while ``flocks.json`` is shared. Rebuild enabled states whenever the + atomically replaced config file changes so a request handled by a + different worker immediately observes tool/service toggles and + repeated-failure auto-disables. + """ + try: + token = cls._current_config_state_token() + except Exception: + return + + with cls._refresh_lock: + if token == cls._config_state_token: + return + + previous_states = { + name: bool(tool.info.enabled) + for name, tool in cls._tools.items() + } + for name, tool in cls._tools.items(): + default_enabled = cls._enabled_defaults.get(name) + if default_enabled is not None: + tool.info.enabled = default_enabled + cls._sync_api_service_states() + cls._apply_tool_settings() + cls._config_state_token = token + + changed_names = [ + name + for name, tool in cls._tools.items() + if previous_states[name] != bool(tool.info.enabled) + ] + for name in changed_names: + if cls._tools[name].info.enabled: + cls._reset_failure_state(name) + if changed_names: + cls._bump_revision("config_enabled_state_sync") + @classmethod def _bump_revision(cls, reason: str) -> None: """Advance the registry revision and invalidate agent prompt caches.""" - cls._revision += 1 + with cls._refresh_lock: + cls._revision += 1 + revision = cls._revision try: from flocks.agent.registry import Agent Agent.invalidate_cache() except Exception as e: log.debug("tool.revision.agent_invalidate_failed", {"error": str(e)}) - log.debug("tool.registry.revision.bumped", {"revision": cls._revision, "reason": reason}) + log.debug("tool.registry.revision.bumped", {"revision": revision, "reason": reason}) @classmethod def register_function( @@ -756,6 +825,7 @@ def _ensure_initialized(cls) -> None: def get(cls, name: str) -> Optional[Tool]: """Get a tool by name""" cls._ensure_initialized() + cls._sync_configured_enabled_states() with cls._refresh_lock: return cls._tools.get(name) @@ -763,6 +833,7 @@ def get(cls, name: str) -> Optional[Tool]: def list_tools(cls, category: Optional[ToolCategory] = None) -> List[ToolInfo]: """List all registered tools, optionally filtered by category""" cls._ensure_initialized() + cls._sync_configured_enabled_states() with cls._refresh_lock: tools = list(cls._tools.values()) if category: @@ -875,7 +946,11 @@ async def execute( if result.success: cls._reset_failure_state(tool_name) else: - disabled = cls._record_failure(tool, kwargs, result.error) + if await cls._failure_auto_disable_enabled(): + disabled = cls._record_failure(tool, kwargs, result.error) + else: + cls._reset_failure_state(tool_name) + disabled = False if disabled: result.metadata = {**(result.metadata or {}), "disabled": True, "disabled_reason": "repeated_error"} suffix = f"tool disabled after {cls._failure_disable_threshold} identical errors" @@ -1003,6 +1078,7 @@ def all_tool_ids(cls) -> List[str]: def snapshot_identity(cls) -> tuple[int, tuple[str, ...]]: """Return a revision/tool-id identity from one consistent registry state.""" cls._ensure_initialized() + cls._sync_configured_enabled_states() with cls._refresh_lock: return cls._revision, tuple(cls._tools.keys()) @@ -1052,7 +1128,7 @@ def init(cls) -> None: cls._register_builtin_tools() cls._register_dynamic_tools() cls._register_plugin_extension_point() - cls._load_plugin_tools() + cls._plugin_load_errors = cls._load_plugin_tools() cls._initialized = True log.debug("tool_registry.initialized", {"count": len(cls._tools)}) finally: @@ -1064,7 +1140,7 @@ async def init_async(cls) -> None: await asyncio.to_thread(cls.init) @classmethod - def _load_plugin_tools(cls, errors: Optional[List[str]] = None) -> None: + def _load_plugin_tools(cls, errors: Optional[List[str]] = None) -> List[str]: """Load plugin tools from both user-level and project-level plugin dirs on init. Without this, YAML/Python plugin tools only appear after an explicit @@ -1077,18 +1153,23 @@ def _load_plugin_tools(cls, errors: Optional[List[str]] = None) -> None: Tracks which tool names were added so that ``refresh_plugin_tools`` can accurately unregister stale entries (regardless of the ``ToolInfo.source`` value). + + Returns every source-level load error for diagnostics and scoped Hub + refresh decisions. """ before = set(cls._tools.keys()) + load_errors: List[str] = [] try: from flocks.plugin import PluginLoader - load_errors = PluginLoader.load_extension("TOOLS", load_entry_points=True) - if errors is not None: - errors.extend(load_errors) + reported_errors = PluginLoader.load_extension("TOOLS", load_entry_points=True) + if reported_errors: + load_errors.extend(reported_errors) except Exception as e: log.warn("tool_registry.plugin_load_failed", {"error": str(e)}) - if errors is not None: - errors.append(f"plugin loader: {e}") + load_errors.append(f"plugin loader: {e}") + if errors is not None: + errors.extend(load_errors) after = set(cls._tools.keys()) new_plugin_tools = sorted(after - before) python_tool_sources: Dict[str, Path] = {} @@ -1131,8 +1212,14 @@ def _load_plugin_tools(cls, errors: Optional[List[str]] = None) -> None: except ValueError: tool.info.native = True cls._plugin_tool_names = sorted(set(new_plugin_tools) | python_plugin_names) - if errors: - return + if errors is not None and load_errors: + return load_errors + cls._finalize_plugin_tools_load() + return load_errors + + @classmethod + def _finalize_plugin_tools_load(cls) -> None: + """Apply configuration overlays after an accepted plugin load.""" cls._bootstrap_user_api_services() # Defence-in-depth: ``register()`` is the canonical writer for # ``_enabled_defaults`` but this catches any tool that landed in @@ -1174,7 +1261,7 @@ def _bootstrap_user_api_services(cls) -> None: or not info.enabled ): continue - provider = info.provider + provider = _api_service_storage_key(info.provider) if provider in seen_providers: continue seen_providers.add(provider) @@ -1226,9 +1313,10 @@ def _sync_api_service_states(cls) -> None: disabled_count = 0 restored_count = 0 for tool in cls._tools.values(): - provider = tool.info.provider - if not provider: + configured_provider = tool.info.provider + if not configured_provider: continue + provider = _api_service_storage_key(configured_provider) svc = api_services.get(provider, {}) svc_enabled = svc.get("enabled", False) if not svc_enabled: @@ -1612,39 +1700,80 @@ def start_watcher(cls) -> None: cls._watcher.start() @classmethod - def refresh_plugin_tools(cls) -> List[str]: + def refresh_plugin_tools(cls, changed_path: Optional[Path] = None) -> List[str]: """Reload plugin tools (YAML + Python) from disk. Unregisters stale plugin tools first so that deleted files are - correctly removed from the registry. + correctly removed from the registry. When a Hub operation supplies + ``changed_path``, source failures outside that path are isolated and + do not block the operation; failures inside it still roll back. """ cls._ensure_initialized() with cls._refresh_lock: tools_before = cls._tools.copy() defaults_before = cls._enabled_defaults.copy() plugin_names_before = list(cls._plugin_tool_names) + plugin_errors_before = list(cls._plugin_load_errors) + enabled_before = { + name: tool.info.enabled + for name, tool in tools_before.items() + } + revision_before = cls._revision load_errors: List[str] = [] - try: - cls._unregister_plugin_tools() - cls._load_plugin_tools(load_errors) - except Exception: + def _restore_snapshot() -> None: + for name, enabled in enabled_before.items(): + tool = tools_before.get(name) + if tool is not None: + tool.info.enabled = enabled cls._tools.clear() cls._tools.update(tools_before) cls._enabled_defaults.clear() cls._enabled_defaults.update(defaults_before) cls._plugin_tool_names = plugin_names_before - raise - if load_errors: - cls._tools.clear() - cls._tools.update(tools_before) - cls._enabled_defaults.clear() - cls._enabled_defaults.update(defaults_before) - cls._plugin_tool_names = plugin_names_before - raise ToolRefreshError("plugin", load_errors) + cls._plugin_load_errors = plugin_errors_before + cls._revision = revision_before - cls._bump_revision("plugin_refresh") - return cls.all_tool_ids() + try: + cls._unregister_plugin_tools() + cls._load_plugin_tools(load_errors) + changed_root = changed_path.resolve() if changed_path is not None else None + fatal_errors: List[str] = [] + for error in load_errors: + if changed_root is None or error.startswith( + ( + "plugin loader:", + "extension point not found:", + "entry point scan:", + ) + ): + fatal_errors.append(error) + continue + if error.startswith("entry point "): + continue + source_text, separator, _detail = error.partition(": ") + if not separator: + fatal_errors.append(error) + continue + source_path = Path(source_text) + if not source_path.is_absolute(): + fatal_errors.append(error) + continue + source_path = source_path.resolve() + if source_path == changed_root or changed_root in source_path.parents: + fatal_errors.append(error) + if fatal_errors: + raise ToolRefreshError("plugin", fatal_errors) + + if load_errors: + cls._finalize_plugin_tools_load() + + cls._plugin_load_errors = load_errors + cls._bump_revision("plugin_refresh") + return cls.all_tool_ids() + except Exception: + _restore_snapshot() + raise @classmethod def refresh_dynamic_tools(cls) -> List[str]: @@ -1685,9 +1814,22 @@ def refresh_dynamic_tools(cls) -> List[str]: @classmethod def _reset_failure_state(cls, tool_name: str) -> None: """Reset failure tracking for a tool after success.""" - if tool_name in cls._failure_state: + with cls._refresh_lock: cls._failure_state.pop(tool_name, None) + @classmethod + async def _failure_auto_disable_enabled(cls) -> bool: + """Return the configured repeated-failure behavior, defaulting on.""" + try: + from flocks.config.config import Config + + config = await Config.get() + if config.tool_failure is not None: + return config.tool_failure.disable_on_repeated_failure + except Exception: + pass + return True + @classmethod def _should_track_failure(cls, tool: Tool) -> bool: """Track failures only for standalone custom tools. @@ -1737,23 +1879,34 @@ def _record_failure(cls, tool: Tool, params: Dict[str, Any], error: Optional[str tool_name = tool.info.name key = cls._failure_key(tool_name, params, error) - state = cls._failure_state.get(tool_name, {"key": None, "count": 0}) + with cls._refresh_lock: + state = cls._failure_state.get(tool_name, {"key": None, "count": 0}) - if state.get("key") == key: - state["count"] = state.get("count", 0) + 1 - else: - state = {"key": key, "count": 1} + if state.get("key") == key: + state["count"] = state.get("count", 0) + 1 + else: + state = {"key": key, "count": 1} - cls._failure_state[tool_name] = state + cls._failure_state[tool_name] = state - if state["count"] >= cls._failure_disable_threshold: - tool.info.enabled = False - log.warn("tool.disabled.repeated_error", { - "tool": tool_name, - "count": state["count"], - "error": error, - }) - return True + if state["count"] >= cls._failure_disable_threshold and tool.info.enabled: + tool.info.enabled = False + try: + from flocks.config.config_writer import ConfigWriter + + ConfigWriter.set_tool_setting(tool_name, {"enabled": False}) + except Exception as exc: + log.warn("tool.disabled.repeated_error.persist_failed", { + "tool": tool_name, + "error": str(exc), + }) + cls._bump_revision("failure_auto_disable") + log.warn("tool.disabled.repeated_error", { + "tool": tool_name, + "count": state["count"], + "error": error, + }) + return True return False diff --git a/flocks/tool/system/memory.py b/flocks/tool/system/memory.py index 2a041e55e..b8770b004 100644 --- a/flocks/tool/system/memory.py +++ b/flocks/tool/system/memory.py @@ -38,10 +38,12 @@ async def _get_session_memory(ctx: ToolContext) -> tuple[Optional[SessionMemory] if not session: return None, ToolResult(success=False, error="Session not found") + from flocks.project.instance import Instance + memory = SessionMemory( session_id=session.id, project_id=session.project_id, - workspace_dir=session.directory, + workspace_dir=Instance.get_directory() or session.directory, enabled=session.memory_enabled, ) diff --git a/flocks/updater/__init__.py b/flocks/updater/__init__.py index 0ba8cc329..512b9d4b6 100644 --- a/flocks/updater/__init__.py +++ b/flocks/updater/__init__.py @@ -2,8 +2,8 @@ Flocks Updater Provides self-update capability via GitHub releases. -Downloads source archives, backs up the current installation, -and replaces source files — no git binary required at runtime. +Downloads source archives and delegates source replacement to a detached +handoff — no git binary required at runtime. """ from flocks.updater.deploy import DeployMode, detect_deploy_mode @@ -13,6 +13,7 @@ check_update, get_current_version, get_latest_release, + install_or_repair_source, perform_pro_bundle_downgrade, perform_update, perform_pro_bundle_install, @@ -28,6 +29,7 @@ "check_update", "get_current_version", "get_latest_release", + "install_or_repair_source", "perform_update", "perform_pro_bundle_install", "perform_pro_bundle_downgrade", diff --git a/flocks/updater/restart_handoff.py b/flocks/updater/restart_handoff.py index 0d429a4db..f763d8970 100644 --- a/flocks/updater/restart_handoff.py +++ b/flocks/updater/restart_handoff.py @@ -1,16 +1,10 @@ -"""Restart handoff helper for the self-updater. - -The updater process owns the backend port while it is spawning the restart -command. Starting the new backend before that process has fully exited can race -with port release. This helper is spawned instead; it waits for the old backend -to exit, clears any remaining backend listener, runs post-apply upgrade tasks, -and then starts the real restart command. -""" +"""Detached restart and source-upgrade handoff helper.""" from __future__ import annotations import argparse import asyncio +import json import shutil import subprocess import time @@ -18,6 +12,7 @@ from typing import Sequence from flocks.cli import service_manager +from flocks.updater import updater as updater_module from flocks.utils.log import append_upgrade_text_log DEFAULT_PARENT_TIMEOUT_SECONDS = 20.0 @@ -27,6 +22,13 @@ DEFAULT_POLL_INTERVAL_SECONDS = 0.25 +class _NullConsole: + """Discard service-manager progress output in the detached helper.""" + + def print(self, *_args, **_kwargs) -> None: + return None + + def _record_handoff_log(message: str) -> None: append_upgrade_text_log(f"restart_handoff {message}") @@ -74,50 +76,73 @@ def _ensure_backend_port_free(backend_port: int) -> bool: def _stop_supervisor_before_restart( *, + daemon_pid: int | None = None, + backend_port: int | None = None, + service_ports: Sequence[int] = (), + force_daemon_stop: bool = False, timeout_seconds: float = SUPERVISOR_STOP_TIMEOUT_SECONDS, poll_interval_seconds: float = DEFAULT_POLL_INTERVAL_SECONDS, ) -> bool: from flocks.cli import service_control paths = service_manager.runtime_paths() - if not service_control.supervisor_is_running(paths): - return True - - try: - service_control.request_stop(paths=paths, timeout=timeout_seconds) - except Exception as exc: - _record_handoff_log(f"supervisor_stop_request_failed error={exc}") - return False + ports = {port for port in (backend_port, *service_ports) if port is not None} + control_running = service_control.supervisor_is_running(paths) + daemon_running = service_manager.pid_is_running(daemon_pid) + if control_running or daemon_running: + try: + service_control.request_stop(paths=paths, timeout=timeout_seconds) + except Exception as exc: + _record_handoff_log(f"supervisor_stop_request_failed error={exc}") + if not force_daemon_stop or daemon_pid is None: + return False + try: + service_manager._terminate_orphan_pid(daemon_pid, "daemon", _NullConsole()) + except Exception as terminate_exc: + _record_handoff_log(f"supervisor_force_stop_failed error={terminate_exc}") + return False deadline = time.monotonic() + timeout_seconds while time.monotonic() < deadline: - if not service_control.supervisor_is_running(paths): + control_stopped = not service_control.supervisor_is_running(paths) + daemon_stopped = not service_manager.pid_is_running(daemon_pid) + ports_stopped = all(not _backend_port_in_use(port) for port in ports) + if control_stopped and daemon_stopped and ports_stopped: return True time.sleep(poll_interval_seconds) - return not service_control.supervisor_is_running(paths) + return ( + not service_control.supervisor_is_running(paths) + and not service_manager.pid_is_running(daemon_pid) + and all(not _backend_port_in_use(port) for port in ports) + ) def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: parser = argparse.ArgumentParser(description="Flocks restart handoff helper") - parser.add_argument("--parent-pid", type=int, required=True) + parser.add_argument("--mode", choices=("restart", "upgrade"), default="restart") + parser.add_argument("--parent-pid", type=int) parser.add_argument("--backend-host", required=True) parser.add_argument("--backend-port", type=int, required=True) parser.add_argument("--frontend-host", required=True) parser.add_argument("--frontend-port", type=int, required=True) parser.add_argument("--backend-pid-file") parser.add_argument("--install-root", required=True) + parser.add_argument("--content-root") + parser.add_argument("--backup-path") + parser.add_argument("--was-running", action="store_true") + parser.add_argument("--daemon-pid", type=int) + parser.add_argument("--service-config-json") parser.add_argument("--uv-path", required=True) parser.add_argument("--sync-timeout", type=int, required=True) parser.add_argument("--version", required=True) parser.add_argument("--current-version", required=True) - parser.add_argument("--backup-path") parser.add_argument("--uv-default-index") parser.add_argument("--npm-registry") parser.add_argument("--pro-wheel-path") parser.add_argument("--pro-bundle-manifest-path") parser.add_argument("--bundle-sha256") parser.add_argument("--cleanup-dir") - parser.add_argument("--prepare-handover", action="store_true") + parser.add_argument("--prepare-handover", action="store_true", help=argparse.SUPPRESS) parser.add_argument("restart_argv", nargs=argparse.REMAINDER) args = parser.parse_args(argv) if args.restart_argv and args.restart_argv[0] == "--": @@ -126,23 +151,177 @@ def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: def _run_upgrade_tasks(args: argparse.Namespace) -> str | None: - from flocks.updater import updater - - return asyncio.run( - updater.run_handoff_upgrade_tasks( - install_root=Path(args.install_root), - uv_path=args.uv_path, - version=args.version, - uv_default_index=args.uv_default_index, - npm_registry=args.npm_registry, - pro_wheel_path=Path(args.pro_wheel_path) if args.pro_wheel_path else None, - pro_bundle_manifest_path=( - Path(args.pro_bundle_manifest_path) if args.pro_bundle_manifest_path else None - ), - bundle_sha256=args.bundle_sha256, - sync_timeout=args.sync_timeout, + try: + asyncio.run( + updater_module.install_or_repair_source( + install_root=Path(args.install_root), + uv_path=args.uv_path, + version=args.version, + uv_default_index=args.uv_default_index, + npm_registry=args.npm_registry, + pro_wheel_path=Path(args.pro_wheel_path) if args.pro_wheel_path else None, + pro_bundle_manifest_path=( + Path(args.pro_bundle_manifest_path) if args.pro_bundle_manifest_path else None + ), + bundle_sha256=args.bundle_sha256, + sync_timeout=args.sync_timeout, + ) ) + except RuntimeError as exc: + return str(exc) + return None + + +def _service_config_from_args(args: argparse.Namespace) -> service_manager.ServiceConfig: + """Load the captured service config without consulting the old daemon.""" + from flocks.cli.service_config import service_config_from_payload + + try: + payload = json.loads(args.service_config_json or "") + except json.JSONDecodeError as exc: + raise ValueError(f"invalid service config JSON: {exc}") from exc + if not isinstance(payload, dict): + raise ValueError("service config JSON must contain an object") + return service_config_from_payload(payload) + + +def _validate_simple_upgrade_args(args: argparse.Namespace) -> None: + """Validate all inputs before stopping or changing the active install.""" + if not args.content_root: + raise ValueError("upgrade content root is missing") + content_root = Path(args.content_root) + if not content_root.is_dir(): + raise ValueError(f"upgrade content root does not exist: {content_root}") + if not args.backup_path: + raise ValueError("upgrade backup path is missing") + backup_path = Path(args.backup_path) + if not backup_path.is_file(): + raise ValueError(f"upgrade backup does not exist: {backup_path}") + if args.was_running and not args.restart_argv: + raise ValueError("running service requires a restart runtime") + _service_config_from_args(args) + + +def _service_ports(args: argparse.Namespace) -> tuple[int, ...]: + """Return every port that must be released before source replacement.""" + config = _service_config_from_args(args) + return tuple(sorted({config.backend_port, config.frontend_port})) + + +def _stop_services_before_upgrade(args: argparse.Namespace) -> bool: + """Stop managed services and wait for the captured daemon and port to exit.""" + try: + service_manager.stop_all(_NullConsole()) + except service_manager.ServiceError as exc: + _record_handoff_log(f"service_stop_failed error={exc}") + return False + return _stop_supervisor_before_restart( + daemon_pid=args.daemon_pid, + backend_port=args.backend_port, + service_ports=_service_ports(args), + ) + + +def _apply_new_source(args: argparse.Namespace) -> None: + """Replace the active source tree with the staged source tree.""" + if not args.content_root: + raise RuntimeError("upgrade content root is missing") + content_root = Path(args.content_root) + if not content_root.is_dir(): + raise RuntimeError(f"upgrade content root does not exist: {content_root}") + updater_module._replace_install_dir(content_root, Path(args.install_root)) + + +def _build_captured_start_argv(args: argparse.Namespace) -> list[str]: + """Build ``flocks start`` directly from the pre-upgrade config snapshot.""" + if not args.restart_argv: + return [] + config = _service_config_from_args(args) + argv = [ + args.restart_argv[0], + "-m", + "flocks.cli.main", + "start", + "--host", + config.frontend_host, + "--port", + str(config.frontend_port), + ] + if config.no_browser: + argv.append("--no-browser") + if config.skip_frontend_build: + argv.append("--skip-webui-build") + if config.legacy_backend_host is not None: + argv.extend(["--server-host", config.legacy_backend_host]) + if config.legacy_backend_port is not None: + argv.extend(["--server-port", str(config.legacy_backend_port)]) + return argv + + +def _start_service_after_upgrade(args: argparse.Namespace) -> tuple[bool, str, str]: + """Synchronously restore the service only when it ran before the upgrade.""" + if not args.was_running: + return True, "", "" + start_argv = _build_captured_start_argv(args) + if not start_argv: + _record_handoff_log("missing_restart_argv") + return False, "", "missing restart runtime" + completed = subprocess.run( + start_argv, + cwd=Path(args.install_root), + capture_output=True, + text=False, + check=False, + ) + stdout = updater_module._clean_process_output(completed.stdout) + stderr = updater_module._clean_process_output(completed.stderr) + if completed.returncode == 0: + return True, stdout, stderr + _record_handoff_log( + f"restart_failed returncode={completed.returncode} stdout={stdout} stderr={stderr}" ) + return False, stdout, stderr + + +def _write_upgrade_result( + *, + args: argparse.Namespace, + phase: str, + failed_stage: str | None = None, + error: str | None = None, + backup_path: Path | None = None, + stdout: str | None = None, + stderr: str | None = None, +) -> None: + """Persist a result record without driving automatic recovery.""" + try: + config_payload = json.loads(args.service_config_json or "{}") + except json.JSONDecodeError: + config_payload = {} + payload = { + "phase": phase, + "version": args.version, + "current_version": args.current_version, + "was_running": args.was_running, + "service_config": config_payload, + } + if error: + payload["last_error"] = error + if failed_stage: + payload["failed_stage"] = failed_stage + if backup_path is not None: + payload["backup_path"] = str(backup_path) + if stdout: + payload["stdout"] = stdout + if stderr: + payload["stderr"] = stderr + try: + updater_module._write_upgrade_result_state(payload) + except Exception as exc: + try: + _record_handoff_log(f"upgrade_result_write_failed phase={phase} error={exc}") + except Exception: + pass def _report_pending_pro_bundle_install_receipt(args: argparse.Namespace) -> None: @@ -161,39 +340,86 @@ def _report_pending_pro_bundle_install_receipt(args: argparse.Namespace) -> None _record_handoff_log("install_receipt_report_skipped") -def _rollback_failed_upgrade(args: argparse.Namespace, error: str) -> None: - from flocks.updater import updater - - _record_handoff_log(f"upgrade_tasks_failed error={error}") - backup_path = Path(args.backup_path) if args.backup_path else None +def _run_simple_upgrade(args: argparse.Namespace) -> int: + """Run stop, source replacement, installation, and restart in order.""" try: - updater._rollback_failed_update( - backup_path, - Path(args.install_root), - args.current_version, + _validate_simple_upgrade_args(args) + except ValueError as exc: + error = str(exc) + _record_handoff_log(f"upgrade_validation_failed error={error}") + _write_upgrade_result( + args=args, + phase="failed", + failed_stage="validation", + error=error, ) - except Exception as exc: - _record_handoff_log(f"rollback_failed error={exc}") + _cleanup_dir(args.cleanup_dir) + return 1 + + if args.parent_pid is not None and not _wait_for_parent_exit(args.parent_pid): + error = f"parent exit timed out: {args.parent_pid}" + _record_handoff_log(error) + _write_upgrade_result(args=args, phase="failed", failed_stage="wait_parent", error=error) + return 1 + + _write_upgrade_result(args=args, phase="running") + if not _stop_services_before_upgrade(args): + error = "service stop timed out" + _record_handoff_log(error) + _write_upgrade_result(args=args, phase="failed", failed_stage="stop", error=error) + return 1 -def _prepare_upgrade_handover(args: argparse.Namespace) -> bool: - from flocks.updater import updater + backup_path = Path(args.backup_path) try: - updater._prepare_upgrade_handover(args.version) + _apply_new_source(args) except Exception as exc: - _record_handoff_log(f"prepare_handover_failed error={exc}") - return False - return True - - -def _rollback_upgrade_handover() -> None: - from flocks.updater import updater + error = str(exc) + _record_handoff_log(f"source_replace_failed error={error}") + _write_upgrade_result( + args=args, + phase="failed", + failed_stage="source_replace", + error=error, + backup_path=backup_path, + ) + return 1 try: - updater.rollback_upgrade_handover() + task_error = _run_upgrade_tasks(args) except Exception as exc: - _record_handoff_log(f"handover_rollback_failed error={exc}") + task_error = f"upgrade tasks crashed: {exc}" + if task_error is not None: + _record_handoff_log(f"install_failed error={task_error}") + _write_upgrade_result( + args=args, + phase="failed", + failed_stage="install", + error=task_error, + backup_path=backup_path, + stderr=task_error, + ) + return 1 + + _report_pending_pro_bundle_install_receipt(args) + started, stdout, stderr = _start_service_after_upgrade(args) + if not started: + error = "service restart failed" + _write_upgrade_result( + args=args, + phase="failed", + failed_stage="start", + error=error, + backup_path=backup_path, + stdout=stdout, + stderr=stderr, + ) + return 1 + + _write_upgrade_result(args=args, phase="done", backup_path=backup_path) + _cleanup_dir(args.cleanup_dir) + return 0 def _cleanup_dir(path_value: str | None) -> None: @@ -202,6 +428,86 @@ def _cleanup_dir(path_value: str | None) -> None: shutil.rmtree(Path(path_value), ignore_errors=True) +def _legacy_upgrade_page_pids(args: argparse.Namespace, page_dir: Path, pid_path: Path) -> list[int]: + """Find trusted temporary upgrade-page processes from older handoffs.""" + candidates: set[int] = set() + try: + candidates.update(service_manager.port_owner_pids(args.frontend_port)) + except Exception: + pass + try: + candidates.add(int(pid_path.read_text(encoding="utf-8").strip())) + except (OSError, ValueError): + pass + + page_dir_text = str(page_dir).lower() + matches: list[int] = [] + for pid in sorted(candidates): + try: + command_line = service_manager._process_command_line(pid).lower() + except Exception: + continue + if "http.server" in command_line and "upgrade-page" in command_line and page_dir_text in command_line: + matches.append(pid) + return matches + + +def _cleanup_legacy_upgrade_handover(args: argparse.Namespace) -> bool: + """Stop and remove upgrade-page artifacts left by old handoff protocols.""" + run_dir = updater_module._flocks_root() / "run" + state_path = run_dir / "upgrade-state.json" + pid_path = run_dir / "upgrade_server.pid" + page_dir = run_dir / "upgrade-page" + if not state_path.exists() and not pid_path.exists() and not page_dir.exists(): + return True + page_pids = _legacy_upgrade_page_pids(args, page_dir, pid_path) + + try: + for pid in page_pids: + service_manager._terminate_orphan_pid(pid, "升级临时页", _NullConsole()) + except Exception as exc: + _record_handoff_log(f"legacy_handover_stop_failed error={exc}") + return False + + remaining = _legacy_upgrade_page_pids(args, page_dir, pid_path) + if remaining: + _record_handoff_log(f"legacy_handover_still_running pids={remaining}") + return False + + state_path.unlink(missing_ok=True) + pid_path.unlink(missing_ok=True) + shutil.rmtree(page_dir, ignore_errors=True) + if page_pids: + _record_handoff_log(f"legacy_handover_cleaned pids={page_pids}") + return True + + +def _legacy_supervisor_pid(args: argparse.Namespace) -> int | None: + """Capture the daemon PID associated with an older handoff request.""" + try: + candidates = service_manager.trusted_daemon_process_pids(root=Path(args.install_root)) + except Exception as exc: + _record_handoff_log(f"legacy_daemon_scan_failed error={exc}") + return None + + port_tokens = (f"--server-port {args.backend_port}", f"--server-port={args.backend_port}") + matches: list[int] = [] + for pid in candidates: + try: + command_line = service_manager._process_command_line(pid).lower() + except Exception: + continue + if any(token in command_line for token in port_tokens): + matches.append(pid) + if len(matches) == 1: + return matches[0] + if not matches and len(candidates) == 1: + return candidates[0] + if candidates: + _record_handoff_log(f"legacy_daemon_scan_ambiguous pids={candidates}") + return None + + def _cli_subcommand(argv: Sequence[str]) -> str | None: """Return the flocks.cli.main subcommand embedded in a Python argv.""" for index, value in enumerate(argv[:-2]): @@ -210,6 +516,40 @@ def _cli_subcommand(argv: Sequence[str]) -> str | None: return None +def _restore_legacy_handoff_endpoints(args: argparse.Namespace) -> None: + """Restore endpoints lost by pre-v2026.7.8 split-service handoffs.""" + if not args.backend_pid_file or _cli_subcommand(args.restart_argv) != "serve": + return + + state_path = updater_module._flocks_root() / "run" / "upgrade-state.json" + try: + payload = json.loads(state_path.read_text(encoding="utf-8")) + except FileNotFoundError: + return + except (OSError, json.JSONDecodeError) as exc: + _record_handoff_log(f"legacy_service_config_invalid error={exc}") + return + + if not isinstance(payload, dict): + _record_handoff_log("legacy_service_config_invalid error=state is not an object") + return + + hosts = (payload.get("backend_host"), payload.get("frontend_host")) + ports = (payload.get("backend_port"), payload.get("frontend_port")) + if not all(isinstance(host, str) and host.strip() for host in hosts) or not all( + isinstance(port, int) and not isinstance(port, bool) and 0 < port <= 65535 for port in ports + ): + _record_handoff_log("legacy_service_config_invalid error=invalid host or port") + return + + args.backend_host, args.frontend_host = hosts + args.backend_port, args.frontend_port = ports + _record_handoff_log( + "legacy_service_config_restored " + f"backend={args.backend_host}:{args.backend_port} frontend={args.frontend_host}:{args.frontend_port}" + ) + + def _restart_argv_for_current_runtime(args: argparse.Namespace, restart_argv: Sequence[str]) -> list[str]: if _cli_subcommand(restart_argv) != "serve": return list(restart_argv) @@ -236,6 +576,10 @@ def _restart_argv_for_current_runtime(args: argparse.Namespace, restart_argv: Se def run(argv: Sequence[str] | None = None) -> int: args = _parse_args(argv) + if args.mode == "upgrade": + return _run_simple_upgrade(args) + + _restore_legacy_handoff_endpoints(args) restart_argv = _restart_argv_for_current_runtime(args, args.restart_argv) if not restart_argv: _record_handoff_log("missing_restart_argv") @@ -246,14 +590,36 @@ def run(argv: Sequence[str] | None = None) -> int: f"parent_pid={args.parent_pid} backend={args.backend_host}:{args.backend_port} " f"frontend={args.frontend_host}:{args.frontend_port}" ) + legacy_daemon_pid = _legacy_supervisor_pid(args) if args.prepare_handover else None - if not _wait_for_parent_exit(args.parent_pid): + if args.parent_pid is not None and not _wait_for_parent_exit(args.parent_pid): _record_handoff_log(f"parent_exit_timeout parent_pid={args.parent_pid}") _cleanup_dir(args.cleanup_dir) return 1 + supervisor_stopped = False if args.prepare_handover: - if not _prepare_upgrade_handover(args): + supervisor_stopped = _stop_supervisor_before_restart( + daemon_pid=legacy_daemon_pid, + backend_port=args.backend_port, + service_ports=(args.frontend_port,), + force_daemon_stop=True, + ) + if not supervisor_stopped: + _record_handoff_log("legacy_handover_stop_timeout") + _cleanup_dir(args.cleanup_dir) + return 1 + elif args.pro_wheel_path or args.pro_bundle_manifest_path: + supervisor_stopped = _stop_supervisor_before_restart( + backend_port=args.backend_port, + service_ports=(args.frontend_port,), + ) + if not supervisor_stopped: + _record_handoff_log("pro_handover_stop_timeout") + _cleanup_dir(args.cleanup_dir) + return 1 + if not _ensure_backend_port_free(args.backend_port): + _record_handoff_log(f"backend_port_unavailable port={args.backend_port}") _cleanup_dir(args.cleanup_dir) return 1 elif not _ensure_backend_port_free(args.backend_port): @@ -266,15 +632,19 @@ def run(argv: Sequence[str] | None = None) -> int: except Exception as exc: task_error = f"upgrade tasks crashed: {exc}" if task_error is not None: - _rollback_failed_upgrade(args, task_error) + _record_handoff_log(f"upgrade_tasks_failed error={task_error}") _cleanup_dir(args.cleanup_dir) return 1 _report_pending_pro_bundle_install_receipt(args) - if not _stop_supervisor_before_restart(): + uses_legacy_protocol = bool(args.backup_path or args.prepare_handover or args.backend_pid_file) + if uses_legacy_protocol and not _cleanup_legacy_upgrade_handover(args): + _record_handoff_log("legacy_handover_cleanup_failed") + _cleanup_dir(args.cleanup_dir) + return 1 + + if not supervisor_stopped and not _stop_supervisor_before_restart(): _record_handoff_log("supervisor_stop_timeout") - if args.prepare_handover: - _rollback_upgrade_handover() _cleanup_dir(args.cleanup_dir) return 1 @@ -286,8 +656,6 @@ def run(argv: Sequence[str] | None = None) -> int: ) except OSError as exc: _record_handoff_log(f"restart_spawn_failed error={exc}") - if args.prepare_handover: - _rollback_upgrade_handover() _cleanup_dir(args.cleanup_dir) return 1 diff --git a/flocks/updater/updater.py b/flocks/updater/updater.py index 3ebdf7791..bfc418415 100644 --- a/flocks/updater/updater.py +++ b/flocks/updater/updater.py @@ -1,10 +1,8 @@ -""" -Core updater logic +"""Core updater logic. -Checks for updates via GitHub / Gitee / GitLab Releases API, downloads -the source archive (zip / tar.gz), backs up the current installation, -extracts the new source over it, re-syncs dependencies, and restarts the -process in-place. +The active process downloads and validates source archives. A detached +handoff stops the old service, backs up and replaces the source tree, repairs +the installation, and restores the captured service state. No git binary is required at runtime — all code fetching is done via HTTP. @@ -13,13 +11,11 @@ """ import asyncio -import contextlib import importlib.util import json import os import re import shutil -import signal import subprocess import sys import tarfile @@ -40,13 +36,6 @@ _DEFAULT_REPO = "AgentFlocks/Flocks" _BACKUP_DIR = Path.home() / ".flocks" / "version" -_UPGRADE_PAGE_MARKER = "flocks-upgrade-in-progress" -_UPGRADE_PHASE_HANDOVER_PREPARING = "handover_preparing" -_UPGRADE_PHASE_TEMP_PAGE_ACTIVE = "temporary_page_active" -_UPGRADE_PHASE_CUTOVER_APPLIED = "cutover_applied" -_UPGRADE_PHASE_ROLLBACK_IN_PROGRESS = "rollback_in_progress" -_UPGRADE_PHASE_ROLLBACK_FAILED = "rollback_failed" -_STATE_FIELD_UNSET = object() _UPDATE_REGION_CN = "cn" _CN_NPM_REGISTRY = "https://registry.npmmirror.com/" _CN_UV_DEFAULT_INDEX = "https://mirrors.aliyun.com/pypi/simple" @@ -54,9 +43,8 @@ _CURL_USER_AGENT = "curl/8.7.1" _FRONTEND_DEPENDENCY_INSTALL_TIMEOUT_SECONDS = 300 _FRONTEND_BUILD_TIMEOUT_SECONDS = 300 -_DEPENDENCY_SYNC_TIMEOUT_SECONDS = 180 +_DEPENDENCY_SYNC_TIMEOUT_SECONDS = 300 _WINDOWS_DEPENDENCY_SYNC_TIMEOUT_SECONDS = 300 -_CANCELLATION_RETRY_DELAY_SECONDS = 0.1 _PRESERVE_NAMES: set[str] = { ".venv", @@ -67,6 +55,17 @@ "__pycache__", } +_ROOT_RUNTIME_NAMES: set[str] = { + ".secret.json", + "config.json", + "data", + "flocks.json", + "flocks.jsonc", + "mcp_list.json", + "run", + "storage", +} + log = Log.create(service="updater") @@ -109,6 +108,15 @@ class UpdateMirrorProfile: pip_index_url: str | None = None +@dataclass(frozen=True) +class ServiceSnapshot: + """Service state captured once before the detached upgrade starts.""" + + config: Any + daemon_pid: int | None + was_running: bool + + @dataclass(frozen=True) class _FrontendNpmCandidate: """A single npm launcher candidate for frontend rebuilds.""" @@ -118,12 +126,6 @@ class _FrontendNpmCandidate: source: str -@dataclass(frozen=True) -class _ProComponentSnapshot: - installed: bool - version: str | None = None - - # ------------------------------------------------------------------ # # Install root # ------------------------------------------------------------------ # @@ -517,18 +519,6 @@ async def build_updated_frontend( raise RuntimeError(frontend_error) -async def _await_ignoring_cancellation(awaitable): - """Finish a post-handover critical step even if the SSE client disconnects.""" - task = asyncio.create_task(awaitable) - while True: - try: - return await asyncio.shield(task) - except asyncio.CancelledError: - log.warning("updater.restart.critical_step_cancelled_ignored") - with contextlib.suppress(asyncio.CancelledError): - await asyncio.sleep(_CANCELLATION_RETRY_DELAY_SECONDS) - - def _dependency_sync_timeout_seconds() -> int: """Return the timeout budget for ``uv sync`` during self-update.""" if sys.platform == "win32": @@ -538,7 +528,7 @@ def _dependency_sync_timeout_seconds() -> int: def _build_dependency_sync_command(uv_path: str, *, uv_default_index: str | None = None) -> list[str]: """Build the ``uv sync`` command used by the self-updater.""" - cmd = [uv_path, "sync", "--frozen", "--no-python-downloads"] + cmd = [uv_path, "sync", "--no-python-downloads"] if uv_default_index: cmd.extend(["--default-index", uv_default_index]) return cmd @@ -569,10 +559,31 @@ async def _run_uv_sync(cmd: list[str]) -> tuple[int, str, str]: def _timeout_message() -> str: return f"Dependency sync timed out after {effective_timeout}s while running uv sync." + def _log_timeout(exc: subprocess.TimeoutExpired, *, retry_without_default_index: bool) -> None: + log.warning( + "updater.dependencies.sync_timeout", + { + "command": list(exc.cmd) if not isinstance(exc.cmd, str) else exc.cmd, + "timeout": effective_timeout, + "stdout": _clean_process_output(exc.stdout), + "stderr": _clean_process_output(exc.stderr), + "retry_without_default_index": retry_without_default_index, + }, + ) + try: code, _, err = await _run_uv_sync(uv_cmd) - except subprocess.TimeoutExpired: - return _timeout_message() + except subprocess.TimeoutExpired as exc: + _log_timeout(exc, retry_without_default_index=bool(uv_default_index)) + if not uv_default_index: + return _timeout_message() + uv_cmd = _build_dependency_sync_command(uv_path) + uv_default_index = None + try: + code, _, err = await _run_uv_sync(uv_cmd) + except subprocess.TimeoutExpired as fallback_exc: + _log_timeout(fallback_exc, retry_without_default_index=False) + return _timeout_message() if ( code != 0 @@ -595,7 +606,8 @@ def _timeout_message() -> str: await asyncio.sleep(2) try: code, _, err = await _run_uv_sync(uv_cmd) - except subprocess.TimeoutExpired: + except subprocess.TimeoutExpired as exc: + _log_timeout(exc, retry_without_default_index=False) return _timeout_message() if code != 0 and uv_default_index: @@ -610,7 +622,8 @@ def _timeout_message() -> str: uv_cmd = _build_dependency_sync_command(uv_path) try: code, _, err = await _run_uv_sync(uv_cmd) - except subprocess.TimeoutExpired: + except subprocess.TimeoutExpired as exc: + _log_timeout(exc, retry_without_default_index=False) return _timeout_message() if code != 0: @@ -618,7 +631,8 @@ def _timeout_message() -> str: await asyncio.sleep(3) try: code, _, err = await _run_uv_sync(uv_cmd) - except subprocess.TimeoutExpired: + except subprocess.TimeoutExpired as exc: + _log_timeout(exc, retry_without_default_index=False) return _timeout_message() if code != 0: @@ -684,197 +698,17 @@ def _upgrade_run_dir() -> Path: return _flocks_root() / "run" -def _upgrade_log_dir() -> Path: - return _flocks_root() / "logs" - - -def _upgrade_state_path() -> Path: - return _upgrade_run_dir() / "upgrade-state.json" +def _upgrade_result_path() -> Path: + return _upgrade_run_dir() / "upgrade-result.json" -def _upgrade_server_pid_path() -> Path: - return _upgrade_run_dir() / "upgrade_server.pid" - - -def _upgrade_page_dir() -> Path: - return _upgrade_run_dir() / "upgrade-page" - - -def _upgrade_page_log_path() -> Path: - return _upgrade_log_dir() / "upgrade-page.log" - - -def _read_upgrade_state() -> dict[str, Any] | None: - path = _upgrade_state_path() - if not path.exists(): - return None - try: - payload = json.loads(path.read_text(encoding="utf-8")) - except (OSError, json.JSONDecodeError): - return None - return payload if isinstance(payload, dict) else None - - -def _write_upgrade_state(payload: dict[str, Any]) -> None: - path = _upgrade_state_path() +def _write_upgrade_result_state(payload: dict[str, Any]) -> None: + """Persist diagnostics from the detached upgrade helper.""" + path = _upgrade_result_path() path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(json.dumps(payload, ensure_ascii=True, sort_keys=True), encoding="utf-8") - - -def _clear_upgrade_state() -> None: - _upgrade_state_path().unlink(missing_ok=True) - - -def _persist_upgrade_state( - payload: dict[str, Any], - *, - phase: str | None = None, - last_error: str | None | object = _STATE_FIELD_UNSET, -) -> dict[str, Any]: - if phase is not None: - payload["phase"] = phase - if last_error is not _STATE_FIELD_UNSET: - if last_error: - payload["last_error"] = last_error - else: - payload.pop("last_error", None) - _write_upgrade_state(payload) - return payload - - -def _upgrade_page_html(version: str) -> str: - return f""" - - - - - Flocks 升级中 - - - -
-
Flocks 正在升级
-

系统升级中，请稍候

-

升级期间服务会短暂重启。当前页面会在新版本恢复后自动刷新。

-
v{version}
-
- - 正在切换到新版本并恢复服务... -
-
如果页面长时间没有恢复，请稍后手动刷新一次。
-
- - - -""" - - -# ------------------------------------------------------------------ # -# Config helper -# ------------------------------------------------------------------ # + result = dict(payload) + result["updated_at"] = datetime.now(timezone.utc).isoformat() + path.write_text(json.dumps(result, ensure_ascii=True, sort_keys=True), encoding="utf-8") async def _get_updater_config(): @@ -1622,7 +1456,7 @@ def _backup_current_version( retain_count: int = 1, ) -> Path | None: """ - Compress the current install directory into ~/.flocks/version/ . + Compress the current source tree into ~/.flocks/version/ . Returns the backup path on success, None on failure. Preserved runtime/user directories are excluded. """ @@ -1633,7 +1467,9 @@ def _backup_current_version( def _filter(info: tarfile.TarInfo) -> tarfile.TarInfo | None: parts = info.name.split("/") - for index, part in enumerate(parts): + if len(parts) >= 2 and parts[0] == "flocks" and parts[1] in _ROOT_RUNTIME_NAMES: + return None + for part in parts: if part in _PRESERVE_NAMES: return None if part == "dist": @@ -1737,62 +1573,34 @@ def _venv_python_path(install_root: Path) -> Path: return install_root / ".venv" / "bin" / "python" -async def _snapshot_pro_component(install_root: Path) -> _ProComponentSnapshot: - python_path = _venv_python_path(install_root) - if not python_path.exists(): - return _ProComponentSnapshot(installed=False) - code, out, _ = await _run_async( - [ - str(python_path), - "-c", - "import importlib.metadata as m\n" - "import sys\n" - "try:\n" - " print(m.version('flockspro'))\n" - "except m.PackageNotFoundError:\n" - " sys.exit(2)\n", - ], - cwd=install_root, - timeout=30, - ) - if code == 0 and out.strip(): - return _ProComponentSnapshot(installed=True, version=out.strip()) - return _ProComponentSnapshot(installed=False) - - -async def _restore_pro_component_snapshot( - snapshot: _ProComponentSnapshot, - *, - uv_path: str, +async def install_or_repair_source( install_root: Path, - env: dict[str, str] | None, -) -> str | None: - python_path = _venv_python_path(install_root) - if not python_path.exists(): - return f"Python environment may need manual repair: missing {python_path}" - if snapshot.installed and snapshot.version: - cmd = [uv_path, "pip", "install", "--python", str(python_path), "--no-deps", f"flockspro=={snapshot.version}"] - else: - cmd = [uv_path, "pip", "uninstall", "--python", str(python_path), "-y", "flockspro"] - code, _, err = await _run_async(cmd, cwd=install_root, timeout=180, env=env) - if code != 0: - return f"Python environment may need manual repair: {err}" - return None - - -async def run_handoff_upgrade_tasks( *, - install_root: Path, - uv_path: str, version: str, + uv_path: str, uv_default_index: str | None = None, npm_registry: str | None = None, + sync_timeout: int = 300, pro_wheel_path: Path | None = None, pro_bundle_manifest_path: Path | None = None, bundle_sha256: str | None = None, - sync_timeout: int | None = None, -) -> str | None: - """Run upgrade work that must happen after the old service exits.""" +) -> None: + """Install or repair the active source tree after managed services stop.""" + install_webui_dir = install_root / "webui" + if install_webui_dir.is_dir() and (install_webui_dir / "package.json").exists(): + from flocks.cli import service_manager + + install_script = "scripts/install.ps1" if sys.platform == "win32" else "scripts/install.sh" + if service_manager.resolve_npm_executable() is None: + raise RuntimeError( + f"npm was not found. Run {install_script} to install Node.js and npm." + ) + if not service_manager.node_version_satisfies_requirement(): + raise RuntimeError( + f"Node.js {service_manager.MIN_NODE_MAJOR}+ is required. " + f"Run {install_script} to install a compatible runtime." + ) + sync_env = _build_uv_sync_env() sync_error = await _sync_project_dependencies( uv_path=uv_path, @@ -1802,7 +1610,7 @@ async def run_handoff_upgrade_tasks( env=sync_env, ) if sync_error is not None: - return sync_error + raise RuntimeError(sync_error) if pro_wheel_path is not None: python_path = _venv_python_path(install_root) @@ -1814,34 +1622,28 @@ async def run_handoff_upgrade_tasks( env=sync_env, ) if code != 0: - return f"Flocks Pro component install failed: {err}" + raise RuntimeError(f"Flocks Pro component install failed: {err}") - validation_error = await _validate_restart_runtime(install_root) - if validation_error: - return validation_error - - install_webui_dir = install_root / "webui" if install_webui_dir.is_dir() and (install_webui_dir / "package.json").exists(): frontend_error = await _build_frontend_workspace( install_webui_dir, npm_registry=npm_registry, ) if frontend_error is not None: - return frontend_error + raise RuntimeError(frontend_error) + + validation_error = await _validate_restart_runtime(install_root) + if validation_error: + raise RuntimeError(validation_error) + + _refresh_global_cli_entry(install_root) - _write_version_marker(version.lstrip("v")) if pro_bundle_manifest_path is not None and pro_bundle_manifest_path.is_file(): _write_pro_bundle_install_marker( _load_json_file(pro_bundle_manifest_path), bundle_sha256=bundle_sha256, ) - - try: - _refresh_global_cli_entry(install_root) - except Exception as exc: - log.warning("updater.refresh_cli.failed", {"error": str(exc)}) - - return None + _write_version_marker(version.lstrip("v")) def _merge_console_manifest_release_identity( @@ -1971,11 +1773,6 @@ async def _uninstall_pro_component( return None -class _NullConsole: - def print(self, *args, **kwargs) -> None: - return None - - def _current_service_config(): from flocks.cli import service_manager from flocks.cli.service_config import service_config_from_status_payload @@ -1992,6 +1789,68 @@ def _current_service_config(): ) +def _capture_service_snapshot() -> ServiceSnapshot: + """Capture restart configuration and liveness before the service is stopped.""" + from flocks.cli import service_manager + from flocks.cli.service_config import ServiceConfig, service_config_from_status_payload + from flocks.cli.service_control import read_supervisor_status + + paths = service_manager.runtime_paths() + try: + status = read_supervisor_status(paths=paths, timeout=1.0) + except Exception as exc: + backend_record = service_manager.read_runtime_record(paths.backend_pid) + frontend_record = service_manager.read_runtime_record(paths.frontend_pid) + backend_running = service_manager.runtime_record_is_running(backend_record) + frontend_running = service_manager.runtime_record_is_running(frontend_record) + if backend_running or frontend_running: + backend_host = backend_record.host if backend_record and backend_record.host else "127.0.0.1" + backend_port = backend_record.port if backend_record and backend_record.port else 5173 + frontend_host = frontend_record.host if frontend_record and frontend_record.host else backend_host + frontend_port = frontend_record.port if frontend_record and frontend_record.port else backend_port + config = ServiceConfig( + backend_host=backend_host, + backend_port=backend_port, + frontend_host=frontend_host, + frontend_port=frontend_port, + no_browser=True, + skip_frontend_build=True, + ) + return ServiceSnapshot(config=config, daemon_pid=None, was_running=True) + + try: + daemon_pids = service_manager.trusted_daemon_process_pids(root=_get_repo_root()) + except Exception: + daemon_pids = [] + if daemon_pids: + raise RuntimeError( + "Supervisor is running but its control API is unavailable; " + "cannot safely capture the service configuration." + ) from exc + return ServiceSnapshot( + config=ServiceConfig(no_browser=True, skip_frontend_build=True), + daemon_pid=None, + was_running=False, + ) + + config = service_config_from_status_payload( + status.raw, + ) + backend_state = status.backend.state.lower() + was_running = bool(status.backend.pid) or backend_state in { + "healthy", + "starting", + "restarting", + "paused", + "static", + } + return ServiceSnapshot( + config=config, + daemon_pid=status.daemon.pid, + was_running=was_running, + ) + + def _spawn_detached_process( command: list[str], *, @@ -2027,491 +1886,19 @@ def _spawn_detached_process( handle.close() -def _write_upgrade_page(version: str) -> Path: - page_dir = _upgrade_page_dir() - page_dir.mkdir(parents=True, exist_ok=True) - (page_dir / "index.html").write_text(_upgrade_page_html(version), encoding="utf-8") - return page_dir - - -def _upgrade_page_probe_urls(frontend_host: str, frontend_port: int) -> list[str]: - from flocks.cli import service_manager - - if frontend_host == "::": - return [ - f"http://[::1]:{frontend_port}", - f"http://127.0.0.1:{frontend_port}", - ] - return [f"http://{service_manager.access_host(frontend_host)}:{frontend_port}"] - - -def _wait_for_upgrade_page(config) -> None: - page_urls = _upgrade_page_probe_urls(config.frontend_host, config.frontend_port) - with httpx.Client(timeout=1.5) as client: - for _ in range(40): - for page_url in page_urls: - try: - response = client.get(page_url) - if response.status_code < 500: - return - except Exception: - pass - time.sleep(0.25) - raise RuntimeError("Upgrade page server failed to start in time") - - -def _start_upgrade_page_server(config, version: str) -> dict[str, Any]: - from flocks.cli import service_manager - - page_dir = _write_upgrade_page(version) - page_dir_resolved = page_dir.resolve() - process = _spawn_detached_process( - service_manager.resolve_python_subprocess_command(_get_repo_root()) - + [ - "-m", - "http.server", - str(config.frontend_port), - "--bind", - config.frontend_host, - "--directory", - str(page_dir_resolved), - ], - cwd=page_dir_resolved, - log_path=_upgrade_page_log_path(), - ) - _upgrade_server_pid_path().write_text(str(process.pid), encoding="utf-8") - _wait_for_upgrade_page(config) - return { - "page_dir": str(page_dir_resolved), - "page_log": str(_upgrade_page_log_path().resolve()), - "upgrade_server_pid": process.pid, - } - - -def _stop_upgrade_page_server(*, frontend_port: int | None = None) -> None: - pid_path = _upgrade_server_pid_path() - if pid_path.exists(): - try: - pid = int(pid_path.read_text(encoding="utf-8").strip()) - except (OSError, ValueError): - pid = None - pid_path.unlink(missing_ok=True) - - if pid is not None: - try: - if sys.platform == "win32": - subprocess.run(["taskkill", "/PID", str(pid), "/T", "/F"], check=False, capture_output=True) - else: - os.kill(pid, signal.SIGTERM) - except OSError: - pass - - for _ in range(40): - try: - os.kill(pid, 0) - except OSError: - break - time.sleep(0.1) - else: - if sys.platform != "win32": - try: - os.kill(pid, signal.SIGKILL) - except OSError: - pass - time.sleep(0.1) - - pid_path.unlink(missing_ok=True) - - if frontend_port is None: - return - - from flocks.cli import service_manager - - remaining = [ - pid - for pid in service_manager.port_owner_pids(frontend_port) - if _looks_like_upgrade_page_process(pid) - ] - if remaining: - log.info( - "updater.upgrade_page.port_fallback_kill", - { - "port": frontend_port, - "pids": remaining, - }, - ) - for rpid in remaining: - try: - if sys.platform == "win32": - subprocess.run(["taskkill", "/PID", str(rpid), "/T", "/F"], check=False, capture_output=True) - else: - os.kill(rpid, signal.SIGKILL) - except OSError: - pass - - if sys.platform == "win32": - wait_attempts = 40 - wait_interval = 0.25 - for _ in range(wait_attempts): - if not any(_looks_like_upgrade_page_process(pid) for pid in service_manager.port_owner_pids(frontend_port)): - return - time.sleep(wait_interval) - return - - if remaining: - time.sleep(0.3) - - -def _looks_like_upgrade_page_process(pid: int) -> bool: - """Return True only for the temporary upgrade-page http.server process.""" - try: - from flocks.cli import service_manager - - command_line = service_manager._process_command_line(pid).lower() - except Exception: - return False - if not command_line: - return False - page_dir = str(_upgrade_page_dir()).lower() - return "http.server" in command_line and "upgrade-page" in command_line and page_dir in command_line - - -def _prepare_upgrade_handover(version: str) -> dict[str, Any]: - from flocks.cli import service_manager - from flocks.cli.service_control import request_prepare_upgrade - - config = _current_service_config() - payload: dict[str, Any] = { - "version": version, - "backend_host": config.backend_host, - "backend_port": config.backend_port, - "frontend_host": config.frontend_host, - "frontend_port": config.frontend_port, - "skip_frontend_build": True, - "phase": _UPGRADE_PHASE_HANDOVER_PREPARING, - } - _persist_upgrade_state(payload, last_error=None) - - console = _NullConsole() - paths = service_manager.runtime_paths() - request_prepare_upgrade(paths=paths, timeout=30.0) - - try: - payload.update(_start_upgrade_page_server(config, version)) - _persist_upgrade_state( - payload, - phase=_UPGRADE_PHASE_TEMP_PAGE_ACTIVE, - last_error=None, - ) - except Exception: - _stop_upgrade_page_server(frontend_port=config.frontend_port) - _clear_upgrade_state() - try: - _start_frontend_with_fallback(config, console, allow_build_fallback=False) - except Exception as restart_error: - log.error("updater.frontend.restore_failed", {"error": str(restart_error)}) - raise - - return payload - - def _spawn_restart_handoff(command: list[str], *, cwd: Path) -> subprocess.Popen: - creationflags = 0 - kwargs: dict[str, object] = {} - if sys.platform == "win32": - creationflags = getattr(subprocess, "CREATE_NEW_PROCESS_GROUP", 0) | getattr(subprocess, "CREATE_NO_WINDOW", 0) - startupinfo_cls = getattr(subprocess, "STARTUPINFO", None) - if startupinfo_cls is not None: - startupinfo = startupinfo_cls() - startupinfo.dwFlags |= getattr(subprocess, "STARTF_USESHOWWINDOW", 0) - startupinfo.wShowWindow = getattr(subprocess, "SW_HIDE", 0) - kwargs["startupinfo"] = startupinfo - else: - kwargs["start_new_session"] = True - return subprocess.Popen(command, cwd=cwd, close_fds=True, creationflags=creationflags, **kwargs) - - -def _service_config_from_payload( - payload: dict[str, Any], - *, - skip_frontend_build: bool | None = None, -): - from flocks.cli.service_config import ServiceConfig, service_config_from_payload - - resolved_skip_frontend_build = ( - bool(payload.get("skip_frontend_build", True)) if skip_frontend_build is None else skip_frontend_build - ) - migrated_payload = dict(payload) - backend_port = migrated_payload.get("backend_port") - frontend_port = migrated_payload.get("frontend_port") - if isinstance(backend_port, int) and isinstance(frontend_port, int) and backend_port != frontend_port: - migrated_payload["legacy_backend_host"] = migrated_payload.get("backend_host") - migrated_payload["legacy_backend_port"] = backend_port - migrated_payload["backend_host"] = migrated_payload.get("frontend_host") or migrated_payload.get("backend_host") - migrated_payload["backend_port"] = frontend_port - migrated_payload["server_port_migration_hint"] = True - return service_config_from_payload( - migrated_payload, - default=ServiceConfig(), - no_browser=True, - skip_frontend_build=resolved_skip_frontend_build, + return _spawn_detached_process( + command, + cwd=cwd, + log_path=_flocks_root() / "logs" / "backend.log", ) def _handoff_service_config(): - payload = _read_upgrade_state() - if payload is not None: - return _service_config_from_payload(payload, skip_frontend_build=True) + """Read the current config for restart-only handoffs.""" return _current_service_config() -def _read_upgrade_server_pid() -> tuple[int | None, bool]: - pid_path = _upgrade_server_pid_path() - if not pid_path.exists(): - return None, False - try: - return int(pid_path.read_text(encoding="utf-8").strip()), True - except (OSError, ValueError): - pid_path.unlink(missing_ok=True) - return None, True - - -def _payload_frontend_port(payload: dict[str, Any] | None) -> int | None: - if not payload: - return None - value = payload.get("frontend_port") - if isinstance(value, bool): - return None - if isinstance(value, int): - return value if value > 0 else None - if isinstance(value, str) and value.isdigit(): - parsed = int(value) - return parsed if parsed > 0 else None - return None - - -def read_upgrade_runtime_state(frontend_port: int | None = None) -> dict[str, Any]: - payload = _read_upgrade_state() - upgrade_pid, pid_file_present = _read_upgrade_server_pid() - resolved_port = _payload_frontend_port(payload) or frontend_port - frontend_host = str(payload.get("frontend_host")) if payload and payload.get("frontend_host") else None - - listener_pids: list[int] = [] - if resolved_port is not None: - try: - from flocks.cli import service_manager - - listener_pids = service_manager.port_owner_pids(resolved_port) - except Exception: - listener_pids = [] - - listener_matches_pid = upgrade_pid is not None and upgrade_pid in listener_pids - return { - "payload": payload, - "payload_present": payload is not None, - "pid_file_present": pid_file_present, - "upgrade_pid": upgrade_pid, - "frontend_host": frontend_host, - "frontend_port": resolved_port, - "listener_pids": listener_pids, - "listener_matches_pid": listener_matches_pid, - "page_active": listener_matches_pid, - "has_artifacts": payload is not None or pid_file_present, - } - - -def _webui_runtime_ready(state: str) -> bool: - return state in {"healthy", "static"} - - -def _start_frontend_with_fallback(config, console, *, allow_build_fallback: bool) -> None: - from flocks.cli.service_config import with_frontend_build - from flocks.cli.service_control import request_restart_webui, request_resume_upgrade - - try: - status = request_resume_upgrade( - config, - paths=None, - timeout=180.0, - ) - if not _webui_runtime_ready(status.webui.state): - raise RuntimeError(status.webui.last_error or "WebUI restart did not become healthy") - return - except Exception: - if not allow_build_fallback or not config.skip_frontend_build: - raise - - rebuilt_config = with_frontend_build(config, skip_frontend_build=False) - result = request_restart_webui(rebuilt_config, force_frontend_build=True, paths=None, timeout=180.0) - if not _webui_runtime_ready(result.webui.state): - raise RuntimeError(result.webui.last_error or "WebUI restart did not become healthy") - - -def cleanup_orphan_upgrade_state(*, frontend_port: int | None = None) -> bool: - state = read_upgrade_runtime_state(frontend_port=frontend_port) - if not state["has_artifacts"]: - return False - - resolved_port = state["frontend_port"] - if resolved_port is None: - _stop_upgrade_page_server() - else: - _stop_upgrade_page_server(frontend_port=resolved_port) - - _clear_upgrade_state() - _upgrade_server_pid_path().unlink(missing_ok=True) - shutil.rmtree(_upgrade_page_dir(), ignore_errors=True) - return True - - -def resolve_upgrade_runtime_state( - *, - attempt_recover: bool = True, - frontend_port: int | None = None, -) -> dict[str, Any]: - state = read_upgrade_runtime_state(frontend_port=frontend_port) - if not state["has_artifacts"]: - return { - **state, - "action": "noop", - "error": None, - } - - resolved_port = state["frontend_port"] - if attempt_recover and state["payload_present"]: - try: - recover_upgrade_state() - return { - **state, - "action": "recovered", - "error": None, - } - except Exception as exc: - cleanup_orphan_upgrade_state(frontend_port=resolved_port) - return { - **state, - "action": "cleanup_after_failed_recover", - "error": str(exc), - } - - cleanup_orphan_upgrade_state(frontend_port=resolved_port) - return { - **state, - "action": "cleaned", - "error": None, - } - - -def _restore_backup_if_possible( - backup_path: Path, - install_root: Path, - previous_version: str, -) -> None: - """Restore install tree from backup without touching upgrade-page state.""" - try: - _restore_backup_archive(backup_path, install_root) - _write_version_marker(previous_version) - log.info("updater.backup.restored", {"backup": str(backup_path)}) - except Exception as exc: - log.error( - "updater.backup.restore_failed", - { - "backup": str(backup_path), - "error": str(exc), - }, - ) - - -def _restore_backup_archive(backup_path: Path, install_root: Path) -> None: - restore_dir = Path(tempfile.mkdtemp(prefix="flocks-rollback-")) - try: - content_root = _extract_archive(backup_path, restore_dir) - _replace_install_dir(content_root, install_root) - finally: - shutil.rmtree(restore_dir, ignore_errors=True) - - -def _rollback_failed_update( - backup_path: Path | None, - install_root: Path, - previous_version: str, -) -> None: - payload = _read_upgrade_state() - if not payload: - return - - _persist_upgrade_state( - payload, - phase=_UPGRADE_PHASE_ROLLBACK_IN_PROGRESS, - last_error=None, - ) - restored_backup = False - restore_error: str | None = None - if backup_path is not None: - try: - _restore_backup_archive(backup_path, install_root) - _write_version_marker(previous_version.lstrip("v")) - restored_backup = True - except Exception as exc: - restore_error = f"Failed to restore backup: {exc}" - log.error("updater.backup.restore_failed", {"error": str(exc), "backup": str(backup_path)}) - - console = _NullConsole() - config = _service_config_from_payload(payload, skip_frontend_build=True) - _stop_upgrade_page_server(frontend_port=config.frontend_port) - try: - _start_frontend_with_fallback( - config, - console, - allow_build_fallback=restored_backup, - ) - except Exception as exc: - log.error( - "updater.frontend.rollback_failed", - {"error": str(exc), "restored_backup": restored_backup, "restore_error": restore_error}, - ) - finally: - _clear_upgrade_state() - shutil.rmtree(_upgrade_page_dir(), ignore_errors=True) - - -def recover_upgrade_state() -> None: - payload = _read_upgrade_state() - if not payload: - return - - console = _NullConsole() - config = _service_config_from_payload(payload) - - _stop_upgrade_page_server(frontend_port=config.frontend_port) - try: - _start_frontend_with_fallback(config, console, allow_build_fallback=True) - except Exception as exc: - log.error("updater.frontend.resume_failed", {"error": str(exc)}) - raise - finally: - _clear_upgrade_state() - shutil.rmtree(_upgrade_page_dir(), ignore_errors=True) - - -def rollback_upgrade_handover() -> None: - payload = _read_upgrade_state() - if not payload: - return - - console = _NullConsole() - config = _service_config_from_payload(payload, skip_frontend_build=True) - - _stop_upgrade_page_server(frontend_port=config.frontend_port) - try: - _start_frontend_with_fallback(config, console, allow_build_fallback=False) - except Exception as exc: - log.error("updater.frontend.rollback_failed", {"error": str(exc)}) - finally: - _clear_upgrade_state() - shutil.rmtree(_upgrade_page_dir(), ignore_errors=True) - - def cleanup_replaced_files(root: Path | None = None) -> None: install_root = root or _get_repo_root() leftovers = sorted( @@ -2607,23 +1994,26 @@ def _safe_remove(target: Path) -> None: def _replace_install_dir( source_dir: Path, install_root: Path, + *, + _is_root: bool = True, ) -> None: """ - Overwrite *install_root* with the contents of *source_dir*, while - preserving user/runtime directories listed in ``_PRESERVE_NAMES`` - at **any** directory depth (not only the top level). + Overwrite *install_root* with the contents of *source_dir*. Build caches + are preserved at every depth; root-level configuration and runtime data + are preserved only at the installation root. """ install_root.mkdir(parents=True, exist_ok=True) source_names = {item.name for item in source_dir.iterdir()} + preserve_names = _PRESERVE_NAMES | (_ROOT_RUNTIME_NAMES if _is_root else set()) for item in source_dir.iterdir(): - if item.name in _PRESERVE_NAMES: + if item.name in preserve_names: continue target = install_root / item.name if item.is_dir() and not item.is_symlink(): if target.exists() or target.is_symlink(): if target.is_dir() and not target.is_symlink(): - _replace_install_dir(item, target) + _replace_install_dir(item, target, _is_root=False) continue _safe_remove(target) shutil.copytree(item, target, symlinks=True) @@ -2633,7 +2023,7 @@ def _replace_install_dir( shutil.copy2(item, target) for child in install_root.iterdir(): - if child.name not in source_names and child.name not in _PRESERVE_NAMES: + if child.name not in source_names and child.name not in preserve_names: _safe_remove(child) @@ -3108,11 +2498,7 @@ async def perform_pro_bundle_downgrade( current_version=current_version, ) log.info("updater.downgrade.restart_handoff_spawn", {"argv": handoff_argv}) - subprocess.Popen( - handoff_argv, - cwd=install_root, - close_fds=True, - ) + _spawn_restart_handoff(handoff_argv, cwd=install_root) os._exit(0) except Exception as exc: log.error("updater.downgrade.restart_failed", {"error": str(exc)}) @@ -3137,6 +2523,7 @@ async def perform_update( locale: str | None = None, region: str | None = None, force_console_manifest: bool = False, + wait_for_handoff: bool = False, ) -> AsyncGenerator[UpdateProgress, None]: """ Async generator that executes the upgrade steps and yields progress events. @@ -3268,27 +2655,8 @@ async def _queue_download_progress(progress: UpdateProgress) -> None: return # ------------------------------------------------------------------ # - # Step 2 – backup current version + # Step 2 – validate and extract the staged source # ------------------------------------------------------------------ # - yield UpdateProgress(stage="backing_up", message="Backing up current version...") - - backup_path = await asyncio.to_thread( - _backup_current_version, - install_root, - current_version, - ucfg.backup_retain_count, - ) - if backup_path: - yield UpdateProgress( - stage="backing_up", - message=f"Backup complete: {backup_path.name}", - ) - else: - yield UpdateProgress( - stage="backing_up", - message="Backup skipped (non-fatal, continuing upgrade)", - ) - extract_dir = tmp_dir / "extracted" extract_dir.mkdir() try: @@ -3309,8 +2677,6 @@ async def _queue_download_progress(progress: UpdateProgress) -> None: except Exception as exc: shutil.rmtree(tmp_dir, ignore_errors=True) msg = f"Failed to extract files: {exc}" - if backup_path: - msg += f"\nRestore from backup: {backup_path}" _record_update_journal(f"ERROR {msg}") yield UpdateProgress(stage="error", message=msg, success=False) return @@ -3328,66 +2694,9 @@ async def _queue_download_progress(progress: UpdateProgress) -> None: else: effective_update_version = latest_tag - # ------------------------------------------------------------------ # - # Step 3 – determine whether frontend handover is needed - # ------------------------------------------------------------------ # - staged_webui_dir = content_root / "webui" - needs_handover = ( - not skip_core_replace - and staged_webui_dir.is_dir() - and (staged_webui_dir / "package.json").exists() - ) - - # ------------------------------------------------------------------ # - # Step 4 – replace install tree - # (Frontend proxy is still alive so the SSE stream keeps flowing.) - # ------------------------------------------------------------------ # - yield UpdateProgress( - stage="applying", - message=( - f"Keeping local Flocks {_version_label(current_version)} and installing the Pro component..." - if skip_core_replace - else f"Applying v{latest_tag}..." - ), - ) - - async def _restore_after_apply_failure() -> None: - if backup_path is None: - return - await asyncio.to_thread( - _restore_backup_if_possible, - backup_path, - install_root, - current_version, - ) - - try: - if not skip_core_replace: - await asyncio.to_thread( - _replace_install_dir, - content_root, - install_root, - ) - except Exception as exc: - final_replace_error: Exception | None = exc - if final_replace_error is not None: - shutil.rmtree(tmp_dir, ignore_errors=True) - await _restore_after_apply_failure() - msg = f"Failed to replace files: {final_replace_error}" - if backup_path: - msg += f"\nRestore from backup: {backup_path}" - yield UpdateProgress(stage="error", message=msg, success=False) - return - - # ------------------------------------------------------------------ # - # Step 5 – prepare dependency sync - # ------------------------------------------------------------------ # - yield UpdateProgress(stage="syncing", message="Preparing dependency sync...") - uv_path = _find_executable("uv") if not uv_path: shutil.rmtree(tmp_dir, ignore_errors=True) - await _restore_after_apply_failure() hint = ( "Dependency sync failed: uv is required but was not found. " "Please install uv (https://docs.astral.sh/uv/) and ensure it " @@ -3409,30 +2718,32 @@ async def _restore_after_apply_failure() -> None: ) # ------------------------------------------------------------------ # - # Step 6 – run post-apply tasks or restart via handoff - # - # With restart=True, dependency sync, Pro install, frontend build, marker - # writes, and CLI refresh happen in restart_handoff after the old backend - # exits. This avoids mutating the active Python environment while the old - # service is still running. + # Step 3 – pure Pro installs may stay in process; source changes are backed + # up here and then always use the detached handoff. # ------------------------------------------------------------------ # - if not restart: - task_error = await run_handoff_upgrade_tasks( - install_root=install_root, - uv_path=uv_path, - version=effective_update_version, - uv_default_index=profile.uv_default_index, - npm_registry=profile.npm_registry, - pro_wheel_path=pro_wheel_path, - pro_bundle_manifest_path=pro_bundle_manifest_path, - bundle_sha256=bundle_sha256, - sync_timeout=sync_timeout, + if skip_core_replace: + yield UpdateProgress( + stage="applying", + message=f"Keeping local Flocks {_version_label(current_version)} and installing the Pro component...", ) - if task_error is not None: + + if skip_core_replace and not restart: + try: + await install_or_repair_source( + install_root=install_root, + uv_path=uv_path, + version=effective_update_version, + uv_default_index=profile.uv_default_index, + npm_registry=profile.npm_registry, + pro_wheel_path=pro_wheel_path, + pro_bundle_manifest_path=pro_bundle_manifest_path, + bundle_sha256=bundle_sha256, + sync_timeout=sync_timeout, + ) + except RuntimeError as exc: shutil.rmtree(tmp_dir, ignore_errors=True) - await _restore_after_apply_failure() - _record_update_journal(f"ERROR {task_error}") - yield UpdateProgress(stage="error", message=task_error, success=False) + _record_update_journal(f"ERROR {exc}") + yield UpdateProgress(stage="error", message=str(exc), success=False) return shutil.rmtree(tmp_dir, ignore_errors=True) @@ -3444,36 +2755,40 @@ async def _restore_after_apply_failure() -> None: ) return - yield UpdateProgress(stage="restarting", message="Restarting service...") - - log.info( - "updater.restart", - { - "tag": latest_tag, - "effective_version": effective_update_version, - "sources": profile.sources, - "repo": ucfg.repo, - "region": profile.region, - }, - ) - await asyncio.sleep(0.8) + if not restart: + shutil.rmtree(tmp_dir, ignore_errors=True) + message = "Source upgrades require the detached handoff; restart=False is only valid for Pro-only changes." + _record_update_journal(f"ERROR {message}") + yield UpdateProgress(stage="error", message=message, success=False) + return - if "--reload" in sys.argv: - log.info("updater.restart.reload_exit3") - sys.exit(3) + backup_path: Path | None = None + if not skip_core_replace: + yield UpdateProgress(stage="backing_up", message="Backing up current version...") + try: + backup_path = await asyncio.to_thread( + _backup_current_version, + install_root, + current_version, + ucfg.backup_retain_count, + ) + except Exception as exc: + log.error("updater.backup.failed", {"error": str(exc)}) + if backup_path is None: + shutil.rmtree(tmp_dir, ignore_errors=True) + message = "Failed to back up the current source; the update was not applied." + _record_update_journal(f"ERROR {message}") + yield UpdateProgress(stage="error", message=message, success=False) + return - try: - restart_argv = _build_restart_argv(install_root) - except Exception as exc: - log.error("updater.restart.build_argv_failed", {"error": str(exc)}) yield UpdateProgress( - stage="error", - message=f"Failed to build restart command: {exc}", - success=False, + stage="applying", + message=f"Staging v{latest_tag} for handoff...", ) - return try: + restart_argv = _build_restart_argv(install_root) + service_snapshot = _capture_service_snapshot() if not skip_core_replace else None handoff_argv = _build_restart_handoff_argv( restart_argv, install_root, @@ -3481,6 +2796,8 @@ async def _restore_after_apply_failure() -> None: sync_timeout=sync_timeout, version=effective_update_version, current_version=current_version, + content_root=content_root if not skip_core_replace else None, + service_snapshot=service_snapshot, backup_path=backup_path, uv_default_index=profile.uv_default_index, npm_registry=profile.npm_registry, @@ -3488,8 +2805,37 @@ async def _restore_after_apply_failure() -> None: pro_bundle_manifest_path=pro_bundle_manifest_path, bundle_sha256=bundle_sha256, cleanup_dir=tmp_dir, - prepare_handover=needs_handover, + wait_for_parent=not wait_for_handoff, + ) + except Exception as exc: + log.error("updater.restart.build_argv_failed", {"error": str(exc)}) + shutil.rmtree(tmp_dir, ignore_errors=True) + yield UpdateProgress( + stage="error", + message=f"Failed to prepare restart handoff: {exc}", + success=False, ) + return + + yield UpdateProgress(stage="restarting", message="Restarting service...") + + log.info( + "updater.restart", + { + "tag": latest_tag, + "effective_version": effective_update_version, + "sources": profile.sources, + "repo": ucfg.repo, + "region": profile.region, + }, + ) + await asyncio.sleep(0.8) + + if "--reload" in sys.argv: + log.info("updater.restart.reload_exit3") + sys.exit(3) + + try: log.info( "updater.restart.handoff_spawn", { @@ -3497,7 +2843,22 @@ async def _restore_after_apply_failure() -> None: "restart_argv": restart_argv, }, ) - _spawn_restart_handoff(handoff_argv, cwd=install_root) + handoff_process = _spawn_restart_handoff(handoff_argv, cwd=install_root) + if wait_for_handoff: + returncode = await asyncio.to_thread(handoff_process.wait) + if returncode != 0: + yield UpdateProgress( + stage="error", + message=f"Upgrade handoff failed with exit code {returncode}. See backend.log for details.", + success=False, + ) + return + yield UpdateProgress( + stage="done", + message=f"Upgraded to {_version_label(effective_update_version)}", + success=True, + ) + return os._exit(0) except Exception as exc: log.error("updater.restart.handoff_spawn_failed", {"error": str(exc)}) @@ -3615,6 +2976,8 @@ def _build_restart_handoff_argv( sync_timeout: int, version: str, current_version: str, + content_root: Path | None = None, + service_snapshot: ServiceSnapshot | None = None, backup_path: Path | None = None, uv_default_index: str | None = None, npm_registry: str | None = None, @@ -3622,56 +2985,91 @@ def _build_restart_handoff_argv( pro_bundle_manifest_path: Path | None = None, bundle_sha256: str | None = None, cleanup_dir: Path | None = None, - prepare_handover: bool = False, + wait_for_parent: bool = True, ) -> list[str]: """Wrap the real restart command in a helper that finishes upgrade work.""" if not restart_argv: raise ValueError("restart command is empty") - config = _handoff_service_config() - managed_restart_argv = [ - restart_argv[0], - "-m", - "flocks.cli.main", - "start", - "--no-browser", - "--skip-webui-build", - "--host", - str(config.backend_host), - "--port", - str(config.backend_port), - ] - if config.legacy_backend_host is not None: - managed_restart_argv.extend(["--server-host", str(config.legacy_backend_host)]) - if config.legacy_backend_port is not None: - managed_restart_argv.extend(["--server-port", str(config.legacy_backend_port)]) + source_upgrade = content_root is not None + if source_upgrade: + if service_snapshot is None: + raise ValueError("source upgrade requires a captured service snapshot") + if backup_path is None: + raise ValueError("source upgrade requires a backup path") + config = service_snapshot.config + else: + config = _handoff_service_config() + + from flocks.cli.service_config import service_config_payload + + if source_upgrade: + # The upgrade helper reconstructs ``flocks start`` from the serialized + # snapshot after the old source has been replaced. + managed_restart_argv = [restart_argv[0]] + else: + managed_restart_argv = [ + restart_argv[0], + "-m", + "flocks.cli.main", + "start", + "--no-browser", + "--skip-webui-build", + "--host", + str(config.frontend_host), + "--port", + str(config.frontend_port), + ] + if config.legacy_backend_host is not None: + managed_restart_argv.extend(["--server-host", str(config.legacy_backend_host)]) + if config.legacy_backend_port is not None: + managed_restart_argv.extend(["--server-port", str(config.legacy_backend_port)]) argv = [ restart_argv[0], "-m", "flocks.updater.restart_handoff", - "--parent-pid", - str(os.getpid()), - "--backend-host", - str(config.backend_host), - "--backend-port", - str(config.backend_port), - "--frontend-host", - str(config.frontend_host), - "--frontend-port", - str(config.frontend_port), - "--install-root", - str(install_root), - "--uv-path", - uv_path, - "--sync-timeout", - str(sync_timeout), - "--version", - version, - "--current-version", - current_version, ] - if backup_path is not None: - argv.extend(["--backup-path", str(backup_path)]) + if wait_for_parent: + argv.extend(["--parent-pid", str(os.getpid())]) + argv.extend( + [ + "--backend-host", + str(config.backend_host), + "--backend-port", + str(config.backend_port), + "--frontend-host", + str(config.frontend_host), + "--frontend-port", + str(config.frontend_port), + "--install-root", + str(install_root), + "--uv-path", + uv_path, + "--sync-timeout", + str(sync_timeout), + "--version", + version, + "--current-version", + current_version, + ] + ) + if source_upgrade: + argv.extend( + [ + "--mode", + "upgrade", + "--content-root", + str(content_root), + "--backup-path", + str(backup_path), + "--service-config-json", + json.dumps(service_config_payload(config), ensure_ascii=True, sort_keys=True), + ] + ) + if service_snapshot.was_running: + argv.append("--was-running") + if service_snapshot.daemon_pid is not None: + argv.extend(["--daemon-pid", str(service_snapshot.daemon_pid)]) if uv_default_index: argv.extend(["--uv-default-index", uv_default_index]) if npm_registry: @@ -3684,8 +3082,6 @@ def _build_restart_handoff_argv( argv.extend(["--bundle-sha256", bundle_sha256]) if cleanup_dir is not None: argv.extend(["--cleanup-dir", str(cleanup_dir)]) - if prepare_handover: - argv.append("--prepare-handover") argv.extend(["--", *managed_restart_argv]) return argv diff --git a/flocks/workflow/center.py b/flocks/workflow/center.py index ad05aff3b..90a4f55b0 100644 --- a/flocks/workflow/center.py +++ b/flocks/workflow/center.py @@ -30,6 +30,7 @@ from flocks.utils.log import Log from flocks.workflow.models import Workflow from flocks.workflow.requirements import resolve_python_package_index_url +from flocks.workflow.service_port import collect_service_ports, resolve_service_port from flocks.workflow.visibility import is_hidden_workflow log = Log.create(service="workflow.center") @@ -69,6 +70,10 @@ class WorkflowNotPublishedError(WorkflowCenterError): """Raised when invoking a workflow that is not published.""" +class WorkflowPortUnavailableError(WorkflowCenterError): + """Raised when a requested workflow service port cannot be used.""" + + def _key_to_string(key: Any) -> str: if isinstance(key, list): return "/".join(str(part) for part in key) @@ -100,6 +105,11 @@ def _normalize_workflow_id(path: Path) -> str: return digest[:24] +def _logical_workflow_id(path: Path) -> str: + """Return the filesystem workflow ID shared across discovery roots.""" + return path.parent.name + + def _fingerprint(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as f: @@ -156,6 +166,17 @@ def resolve_project_workflow_roots(base_dir: Optional[Path] = None) -> list[Path ] +def resolve_workflow_scan_roots( + base_dir: Optional[Path] = None, +) -> list[tuple[Path, str]]: + """Return workflow roots ordered from lowest to highest priority.""" + return [ + (root, "project") for root in resolve_project_workflow_roots(base_dir) + ] + [ + (root, "global") for root in resolve_global_workflow_roots() + ] + + def _resolve_project_workflow_root(base_dir: Optional[Path] = None) -> Path: """/.flocks/plugins/workflows/ — canonical project-level workflow storage.""" root = base_dir or Path.cwd() @@ -172,37 +193,7 @@ def _is_port_available(port: int) -> bool: return False -def _port_from_service_url(value: Any) -> Optional[int]: - if not value: - return None - try: - parsed = urlparse(str(value)) - return parsed.port - except (TypeError, ValueError): - return None - - -def _ports_from_service_record(record: Any) -> set[int]: - if not isinstance(record, dict): - return set() - - ports: set[int] = set() - for key in ("hostPort", "port"): - try: - port = int(record.get(key) or 0) - except (TypeError, ValueError): - port = 0 - if 1 <= port <= 65535: - ports.add(port) - - for key in ("serviceUrl", "invokeUrl"): - port = _port_from_service_url(record.get(key)) - if port is not None: - ports.add(port) - return ports - - -async def _reserved_service_ports() -> set[int]: +async def _reserved_service_ports(*, exclude_workflow_id: Optional[str] = None) -> set[int]: ports: set[int] = set() for prefix in (_RUNTIME_PREFIX, _API_SERVICE_PREFIX, _REGISTRY_PREFIX): try: @@ -217,7 +208,12 @@ async def _reserved_service_ports() -> set[int]: except Exception as exc: log.warning("workflow.port.read_reserved_failed", {"key": _key_to_string(key), "error": str(exc)}) continue - ports.update(_ports_from_service_record(record)) + if exclude_workflow_id and isinstance(record, dict): + record_workflow_id = str(record.get("workflowId") or "") + key_value = _key_to_string(key) + if record_workflow_id == exclude_workflow_id or key_value.endswith(f"/{exclude_workflow_id}"): + continue + ports.update(collect_service_ports(record)) return ports @@ -234,14 +230,43 @@ def _release_port_reservation(port: Optional[int]) -> None: _IN_FLIGHT_PORT_RESERVATIONS.pop(port, None) -async def _allocate_port() -> int: +def _select_publish_port( + requested_port: Optional[int], + registry: Dict[str, Any], + previous_instance: Any, +) -> tuple[Optional[int], Optional[int]]: + """Return the desired port and the port currently bound by this workflow.""" + current_port = resolve_service_port(previous_instance, keys=("hostPort", "port", "servicePort")) + configured_port = resolve_service_port(registry, keys=("servicePort", "port", "hostPort")) + desired_port = requested_port if requested_port is not None else configured_port or current_port + return desired_port, current_port + + +async def _allocate_port( + preferred_port: Optional[int] = None, + *, + workflow_id: Optional[str] = None, + allow_bound_port: Optional[int] = None, +) -> int: start = int(os.getenv("FLOCKS_WORKFLOW_SERVICE_PORT_START", str(_DEFAULT_PORT_START))) end = int(os.getenv("FLOCKS_WORKFLOW_SERVICE_PORT_END", str(_DEFAULT_PORT_END))) if start > end: raise WorkflowCenterError("Invalid workflow service port range") async with _PORT_ALLOCATION_LOCK: - reserved_ports = await _reserved_service_ports() + reserved_ports = await _reserved_service_ports(exclude_workflow_id=workflow_id) reserved_ports.update(_reserved_in_flight_ports()) + if preferred_port is not None: + if isinstance(preferred_port, bool) or not 1 <= preferred_port <= 65535: + raise WorkflowCenterError("Workflow service port must be between 1 and 65535") + bound_by_current_runtime = preferred_port == allow_bound_port + if preferred_port in reserved_ports or ( + not bound_by_current_runtime and not _is_port_available(preferred_port) + ): + raise WorkflowPortUnavailableError( + f"Workflow service port {preferred_port} is already reserved or in use" + ) + _IN_FLIGHT_PORT_RESERVATIONS[preferred_port] = time.time() + _PORT_RESERVATION_TTL_S + return preferred_port for port in range(start, end + 1): if port not in reserved_ports and _is_port_available(port): _IN_FLIGHT_PORT_RESERVATIONS[port] = time.time() + _PORT_RESERVATION_TTL_S @@ -256,6 +281,17 @@ async def _read_registry(workflow_id: str) -> Dict[str, Any]: return data +def _load_discoverable_workflow(workflow_path: Path) -> Optional[Dict[str, Any]]: + """Load and validate a workflow, returning None when intentionally hidden.""" + raw = json.loads(workflow_path.read_text(encoding="utf-8")) + meta_path = workflow_path.parent / "meta.json" + meta = json.loads(meta_path.read_text(encoding="utf-8")) if meta_path.is_file() else None + if is_hidden_workflow(raw, meta): + return None + Workflow.from_dict(raw) + return raw + + async def _scan_workflow_dir( workflow_root: Path, source_type: str, @@ -270,12 +306,9 @@ async def _scan_workflow_dir( return for workflow_path in sorted(workflow_root.glob("*/workflow.json")): try: - raw = json.loads(workflow_path.read_text(encoding="utf-8")) - meta_path = workflow_path.parent / "meta.json" - meta = json.loads(meta_path.read_text(encoding="utf-8")) if meta_path.is_file() else None - if is_hidden_workflow(raw, meta): + raw = _load_discoverable_workflow(workflow_path) + if raw is None: continue - Workflow.from_dict(raw) except Exception as exc: log.warning( "workflow.center.scan.skip_invalid", @@ -284,6 +317,7 @@ async def _scan_workflow_dir( continue workflow_id = _normalize_workflow_id(workflow_path) + logical_workflow_id = _logical_workflow_id(workflow_path) fp = _fingerprint(workflow_path) now_ms = _now_ms() existing = await WorkflowStore.kv_get(_registry_key(workflow_id)) or {} @@ -291,6 +325,7 @@ async def _scan_workflow_dir( draft_changed = bool(existing) and existing.get("fingerprint") != fp entry = { "workflowId": workflow_id, + "logicalWorkflowId": logical_workflow_id, "name": raw.get("name") or workflow_path.parent.name, "description": raw.get("description") or "", "sourceType": source_type, @@ -306,9 +341,10 @@ async def _scan_workflow_dir( "activeReleaseId": existing.get("activeReleaseId"), "serviceKey": existing.get("serviceKey"), "serviceUrl": existing.get("serviceUrl"), + "servicePort": existing.get("servicePort"), } await WorkflowStore.kv_put(_registry_key(workflow_id), entry) - by_id[workflow_id] = entry + by_id[logical_workflow_id] = entry async def scan_skill_workflows(base_dir: Optional[Path] = None) -> List[Dict[str, Any]]: @@ -316,21 +352,15 @@ async def scan_skill_workflows(base_dir: Optional[Path] = None) -> List[Dict[str Scan order (lowest → highest priority), see resolve_global_workflow_roots / resolve_project_workflow_roots: - 1. ~/.flocks/plugins/workflow/ (global legacy, sourceType="global") - 2. ~/.flocks/workflow/ (global compat, sourceType="global") - 3. ~/.flocks/plugins/workflows/ (global canonical, sourceType="global") - 4. /.flocks/plugins/workflow/ (project legacy, sourceType="project") - 5. /.flocks/workflow/ (project compat, sourceType="project") - 6. /.flocks/plugins/workflows/ (project canonical, sourceType="project") + 1. /.flocks project-bundled workflows + 2. ~/.flocks user workflows (highest priority) When two directories contain a workflow with the same ID, the later (higher-priority) entry wins. """ by_id: Dict[str, Dict[str, Any]] = {} - for path in resolve_global_workflow_roots(): - await _scan_workflow_dir(path, "global", by_id) - for path in resolve_project_workflow_roots(base_dir): - await _scan_workflow_dir(path, "project", by_id) + for path, source_type in resolve_workflow_scan_roots(base_dir): + await _scan_workflow_dir(path, source_type, by_id) entries = list(by_id.values()) entries.sort(key=lambda item: item.get("updatedAt", 0), reverse=True) @@ -376,16 +406,71 @@ def format_workflow_entries( async def list_registry_entries() -> List[Dict[str, Any]]: """List registered skill workflows.""" keys = await WorkflowStore.kv_list(_REGISTRY_PREFIX) - items: List[Dict[str, Any]] = [] + selected_entries: Dict[str, tuple[tuple[int, int, int], Dict[str, Any]]] = {} for raw_key in keys: key = _key_to_string(raw_key) entry = await WorkflowStore.kv_get(key) if entry: - items.append(entry) + logical_id = entry.get("logicalWorkflowId") + if not logical_id: + workflow_path = entry.get("workflowPath") + logical_id = ( + _logical_workflow_id(Path(str(workflow_path))) + if workflow_path + else str(entry.get("workflowId") or key) + ) + priority = _registry_entry_priority(entry) + if priority[1] >= 0 and not _registry_entry_is_discoverable(entry): + continue + # Only collapse entries that belong to the currently active project + # or user roots. Registrations from another project remain manageable. + selection_key = ( + f"logical:{logical_id}" + if priority[1] >= 0 + else f"registry:{entry.get('workflowId') or key}" + ) + selected = selected_entries.get(selection_key) + if selected is None or priority >= selected[0]: + selected_entries[selection_key] = (priority, entry) + items = [entry for _priority, entry in selected_entries.values()] items.sort(key=lambda item: item.get("updatedAt", 0), reverse=True) return items +def _registry_entry_priority(entry: Dict[str, Any]) -> tuple[int, int, int]: + """Rank registry entries using the same project-then-user scan order.""" + workflow_path = Path(str(entry.get("workflowPath") or "")) + try: + resolved_path = workflow_path.resolve() + except OSError: + resolved_path = workflow_path + + roots = [root for root, _source in resolve_workflow_scan_roots(Path.cwd())] + root_priority = -1 + for index, root in enumerate(roots): + try: + resolved_path.relative_to(root.resolve()) + except (OSError, ValueError): + continue + root_priority = index + break + + return ( + int(resolved_path.is_file()), + root_priority, + int(entry.get("updatedAt") or 0), + ) + + +def _registry_entry_is_discoverable(entry: Dict[str, Any]) -> bool: + """Return whether an active-root registry entry still passes discovery.""" + workflow_path = Path(str(entry.get("workflowPath") or "")) + try: + return _load_discoverable_workflow(workflow_path) is not None + except Exception: + return False + + def _service_release_file(workflow_id: str, release_id: str) -> Path: base = Config.get_data_path() / _SERVICE_DATA_DIR / "releases" / workflow_id base.mkdir(parents=True, exist_ok=True) @@ -737,22 +822,58 @@ async def _stop_runtime_record( } -async def _stop_existing_runtime_for_publish(workflow_id: str) -> None: - """Best-effort cleanup before starting a replacement service.""" +async def _stop_existing_runtime_for_publish(workflow_id: str) -> bool: + """Stop the current runtime and report whether an old process existed.""" runtime = await WorkflowStore.kv_get(_runtime_key(workflow_id)) if isinstance(runtime, dict) and runtime: await _stop_runtime_record(workflow_id, runtime, update_registry=False) - else: + return True + legacy_pid = await WorkflowStore.kv_get(_local_pid_key(workflow_id)) + if legacy_pid: await _stop_local_service(workflow_id) + return True + return False + + +async def _remove_runtime_records_for_release(workflow_id: str, release_id: str) -> None: + """Remove only runtime pointers created by the failed release.""" + for key in (_runtime_key(workflow_id), _active_release_key(workflow_id)): + record = await WorkflowStore.kv_get(key) + if isinstance(record, dict) and record.get("releaseId") == release_id: + await WorkflowStore.kv_remove(key) -async def publish_workflow_local(workflow_id: str, *, api_key: Optional[str] = None) -> Dict[str, Any]: +async def _restore_registry_after_publish_failure( + workflow_id: str, + previous_registry: Dict[str, Any], + *, + previous_runtime_survived: bool, +) -> None: + registry = dict(previous_registry) + if not previous_runtime_survived: + registry.update( + { + "publishStatus": "failed", + "serviceUrl": None, + "updatedAt": _now_ms(), + } + ) + await WorkflowStore.kv_put(_registry_key(workflow_id), registry) + + +async def publish_workflow_local( + workflow_id: str, + *, + api_key: Optional[str] = None, + port: Optional[int] = None, +) -> Dict[str, Any]: """Publish a workflow as a local subprocess using the current Python env. This is the default driver for development: no Docker, instant startup, uses the same .venv as the main server. """ registry = await _read_registry(workflow_id) + previous_registry = dict(registry) workflow_path = Path(str(registry["workflowPath"])) if not workflow_path.exists(): raise WorkflowCenterError(f"Workflow file not found: {workflow_path}") @@ -761,22 +882,31 @@ async def publish_workflow_local(workflow_id: str, *, api_key: Optional[str] = N Workflow.from_dict(workflow_json) release_id = uuid.uuid4().hex - now_ms = _now_ms() - registry["publishStatus"] = "publishing" - registry["updatedAt"] = now_ms - await WorkflowStore.kv_put(_registry_key(workflow_id), registry) - - release_snapshot_file = await _write_release_snapshot(workflow_id, release_id, workflow_json) - - await _stop_existing_runtime_for_publish(workflow_id) - - host_port = await _allocate_port() - service_url = _host_service_url(host_port) service_key = workflow_id runtime_api_key = api_key or _generate_api_key() + previous_runtime = await WorkflowStore.kv_get(_runtime_key(workflow_id)) or await WorkflowStore.kv_get( + _active_release_key(workflow_id) + ) + preferred_port, previous_port = _select_publish_port(port, registry, previous_runtime) + host_port: Optional[int] = None proc: asyncio.subprocess.Process | None = None + previous_runtime_stopped = False try: + host_port = await _allocate_port( + preferred_port, + workflow_id=workflow_id, + allow_bound_port=previous_port, + ) + release_snapshot_file = await _write_release_snapshot(workflow_id, release_id, workflow_json) + registry["publishStatus"] = "publishing" + registry["updatedAt"] = _now_ms() + await WorkflowStore.kv_put(_registry_key(workflow_id), registry) + + previous_runtime_stopped = await _stop_existing_runtime_for_publish(workflow_id) + if not _is_port_available(host_port): + raise WorkflowPortUnavailableError(f"Workflow service port {host_port} is still in use after restart") + service_url = _host_service_url(host_port) env = os.environ.copy() env[_SERVICE_API_KEY_ENV] = runtime_api_key proc = await asyncio.create_subprocess_exec( @@ -838,28 +968,34 @@ async def publish_workflow_local(workflow_id: str, *, api_key: Optional[str] = N } await WorkflowStore.kv_put(_active_release_key(workflow_id), active_record) await WorkflowStore.kv_put(_runtime_key(workflow_id), active_record) - _release_port_reservation(host_port) + + registry["publishStatus"] = "active" + registry["activeReleaseId"] = release_id + registry["serviceKey"] = service_key + registry["serviceUrl"] = service_url + registry["servicePort"] = host_port + registry["updatedAt"] = _now_ms() + await WorkflowStore.kv_put(_registry_key(workflow_id), registry) except Exception as exc: - _release_port_reservation(host_port) if proc is not None: try: - _signal_local_process(proc.pid, signal.SIGTERM, proc.pid) - await _wait_for_pid_exit(proc.pid, 1.0) + await _stop_local_service(workflow_id) except Exception: pass - registry["publishStatus"] = "failed" - registry["updatedAt"] = _now_ms() - await WorkflowStore.kv_put(_registry_key(workflow_id), registry) + if _pid_is_running(proc.pid): + _signal_local_process(proc.pid, signal.SIGTERM, proc.pid) + await _wait_for_pid_exit(proc.pid, 1.0) + await _remove_runtime_records_for_release(workflow_id, release_id) + await _restore_registry_after_publish_failure( + workflow_id, + previous_registry, + previous_runtime_survived=bool(previous_runtime) and not previous_runtime_stopped, + ) if isinstance(exc, WorkflowCenterError): raise raise WorkflowCenterError(str(exc)) from exc - - registry["publishStatus"] = "active" - registry["activeReleaseId"] = release_id - registry["serviceKey"] = service_key - registry["serviceUrl"] = service_url - registry["updatedAt"] = _now_ms() - await WorkflowStore.kv_put(_registry_key(workflow_id), registry) + finally: + _release_port_reservation(host_port) log.info("workflow.local.published", {"id": workflow_id, "port": host_port, "pid": proc.pid}) return active_record @@ -892,14 +1028,15 @@ async def publish_workflow( image: Optional[str] = None, driver: Optional[str] = None, api_key: Optional[str] = None, + port: Optional[int] = None, ) -> Dict[str, Any]: """Publish a workflow using the configured service driver (local or docker).""" resolved_driver = (driver or _service_driver()).strip().lower() if resolved_driver == "docker": - return await _publish_workflow_docker(workflow_id, image=image, api_key=api_key) + return await _publish_workflow_docker(workflow_id, image=image, api_key=api_key, port=port) if resolved_driver != "local": raise WorkflowCenterError(f"Unsupported workflow service driver: {resolved_driver}") - return await publish_workflow_local(workflow_id, api_key=api_key) + return await publish_workflow_local(workflow_id, api_key=api_key, port=port) async def stop_workflow_service(workflow_id: str) -> Dict[str, Any]: @@ -920,9 +1057,11 @@ async def _publish_workflow_docker( image: Optional[str] = None, *, api_key: Optional[str] = None, + port: Optional[int] = None, ) -> Dict[str, Any]: """Publish a registered workflow as a Docker service container.""" registry = await _read_registry(workflow_id) + previous_registry = dict(registry) workflow_path = Path(str(registry["workflowPath"])) if not workflow_path.exists(): raise WorkflowCenterError(f"Workflow file not found: {workflow_path}") @@ -932,10 +1071,6 @@ async def _publish_workflow_docker( release_id = uuid.uuid4().hex now_ms = _now_ms() - registry["publishStatus"] = "publishing" - registry["updatedAt"] = now_ms - await WorkflowStore.kv_put(_registry_key(workflow_id), registry) - release_snapshot_file = await _write_release_snapshot(workflow_id, release_id, workflow_json) release_runtime_dir = release_snapshot_file.parent has_requirements_snapshot = await _write_requirements_snapshot(release_runtime_dir) @@ -950,14 +1085,14 @@ async def _publish_workflow_docker( "activatedAt": None, "deactivatedAt": None, } - await WorkflowStore.kv_put(_release_key(workflow_id, release_id), release_record) - previous_runtime = await WorkflowStore.kv_get(_runtime_key(workflow_id)) or {} previous_active = await WorkflowStore.kv_get(_active_release_key(workflow_id)) or {} - previous_container_name = previous_active.get("containerName") - previous_release_id = previous_active.get("releaseId") + previous_instance = previous_runtime or previous_active + previous_container_name = previous_instance.get("containerName") + previous_release_id = previous_instance.get("releaseId") + preferred_port, previous_port = _select_publish_port(port, registry, previous_instance) - host_port = await _allocate_port() + host_port: Optional[int] = None container_name = _workflow_container_name(workflow_id, release_id) image_name = image or os.getenv("FLOCKS_WORKFLOW_SERVICE_IMAGE", _DEFAULT_IMAGE) runtime_install = os.getenv("FLOCKS_WORKFLOW_SERVICE_RUNTIME_INSTALL", "1").strip().lower() not in { @@ -981,7 +1116,7 @@ async def _publish_workflow_docker( "--name", container_name, "-p", - f"{host_port}:{_SERVICE_CONTAINER_PORT}", + "", "-v", f"{project_root}:/app:ro", "-v", @@ -1085,7 +1220,28 @@ async def _publish_workflow_docker( ] ) + previous_stopped = False + publish_started = False try: + host_port = await _allocate_port( + preferred_port, + workflow_id=workflow_id, + allow_bound_port=previous_port, + ) + port_binding_index = cmd.index("-p") + 1 + cmd[port_binding_index] = f"{host_port}:{_SERVICE_CONTAINER_PORT}" + + registry["publishStatus"] = "publishing" + registry["updatedAt"] = now_ms + await WorkflowStore.kv_put(_registry_key(workflow_id), registry) + await WorkflowStore.kv_put(_release_key(workflow_id, release_id), release_record) + publish_started = True + + if previous_instance and host_port == previous_port: + await _stop_runtime_record(workflow_id, previous_instance, update_registry=False) + previous_stopped = True + if not _is_port_available(host_port): + raise WorkflowPortUnavailableError(f"Workflow service port {host_port} is still in use after restart") stdout, _, _ = await exec_docker(cmd) container_id = stdout.strip() service_url = _host_service_url(host_port) @@ -1118,39 +1274,52 @@ async def _publish_workflow_docker( } await WorkflowStore.kv_put(_active_release_key(workflow_id), active_record) await WorkflowStore.kv_put(_runtime_key(workflow_id), active_record) - _release_port_reservation(host_port) registry["publishStatus"] = "active" registry["activeReleaseId"] = release_id registry["serviceKey"] = service_key registry["serviceUrl"] = service_url + registry["servicePort"] = host_port registry["updatedAt"] = _now_ms() await WorkflowStore.kv_put(_registry_key(workflow_id), registry) + except Exception as exc: + if publish_started: + await _stop_and_remove_container(container_name) + await _remove_runtime_records_for_release(workflow_id, release_id) + release_record["status"] = "failed" + release_record["deactivatedAt"] = _now_ms() + await WorkflowStore.kv_put(_release_key(workflow_id, release_id), release_record) + + await _restore_registry_after_publish_failure( + workflow_id, + previous_registry, + previous_runtime_survived=bool(previous_instance) and not previous_stopped, + ) + if isinstance(exc, WorkflowCenterError): + raise + raise WorkflowCenterError(str(exc)) from exc + finally: + _release_port_reservation(host_port) - if isinstance(previous_runtime, dict) and previous_runtime: + if not previous_stopped and isinstance(previous_runtime, dict) and previous_runtime: + try: await _stop_runtime_record( workflow_id, previous_runtime, update_registry=False, clear_runtime_keys=False, ) - elif previous_container_name and previous_container_name != container_name: + except Exception as exc: + log.warning("workflow.publish.previous_runtime_cleanup_failed", {"id": workflow_id, "error": str(exc)}) + elif not previous_stopped and previous_container_name and previous_container_name != container_name: + try: await _stop_and_remove_container(previous_container_name) if previous_release_id: await _mark_release_inactive(workflow_id, previous_release_id) + except Exception as exc: + log.warning("workflow.publish.previous_container_cleanup_failed", {"id": workflow_id, "error": str(exc)}) - return active_record - except Exception as exc: - _release_port_reservation(host_port) - await _stop_and_remove_container(container_name) - release_record["status"] = "failed" - release_record["deactivatedAt"] = _now_ms() - await WorkflowStore.kv_put(_release_key(workflow_id, release_id), release_record) - - registry["publishStatus"] = "failed" - registry["updatedAt"] = _now_ms() - await WorkflowStore.kv_put(_registry_key(workflow_id), registry) - raise WorkflowCenterError(str(exc)) from exc + return active_record async def _stop_workflow_service_docker(workflow_id: str) -> Dict[str, Any]: diff --git a/flocks/workflow/fs_store.py b/flocks/workflow/fs_store.py index 8f1900268..0f3471959 100644 --- a/flocks/workflow/fs_store.py +++ b/flocks/workflow/fs_store.py @@ -9,7 +9,7 @@ from flocks.utils.log import Log -from .center import resolve_global_workflow_roots, resolve_project_workflow_roots +from .center import resolve_workflow_scan_roots log = Log.create(service="workflow.fs-store") @@ -104,12 +104,7 @@ def find_workspace_root() -> Path: def workflow_scan_dirs() -> list[tuple[Path, str]]: """Return all workflow roots ordered from lowest to highest priority.""" - workspace = find_workspace_root() - return [ - (root, "global") for root in resolve_global_workflow_roots() - ] + [ - (root, "project") for root in resolve_project_workflow_roots(workspace) - ] + return resolve_workflow_scan_roots(find_workspace_root()) def read_workflow_dir( diff --git a/flocks/workflow/poller_manager.py b/flocks/workflow/poller_manager.py index c4596286c..a85aa3b48 100644 --- a/flocks/workflow/poller_manager.py +++ b/flocks/workflow/poller_manager.py @@ -28,6 +28,10 @@ from flocks.workflow.models import Workflow from flocks.workflow.runner import RunWorkflowResult, run_workflow from flocks.workflow.store import WorkflowStore +from flocks.workflow.tool_context import ( + build_workflow_tool_context, + cleanup_workflow_tool_context, +) WORKFLOW_POLLER_CONFIG_PREFIX = "workflow_poller_config/" DEFAULT_INTERVAL_SECONDS = 30 @@ -191,7 +195,18 @@ async def start_all(self) -> None: if not workflow_id: continue if isinstance(data, dict) and data.get("enabled"): - await self.restart_workflow(workflow_id) + try: + await self.restart_workflow(workflow_id, startup=True) + except Exception as exc: + self._status[workflow_id] = { + **self._base_status(workflow_id), + "state": "failed", + "error": str(exc), + } + log.warning( + "poller.start_failed", + {"workflow_id": workflow_id, "error": str(exc)}, + ) async def stop_all(self) -> None: for workflow_id in list(self._tasks.keys()): @@ -229,7 +244,12 @@ async def stop_workflow(self, workflow_id: str) -> None: self._run_cancel_events.pop(workflow_id, None) self._status[workflow_id] = current - async def restart_workflow(self, workflow_id: str) -> Dict[str, Any]: + async def restart_workflow( + self, + workflow_id: str, + *, + startup: bool = False, + ) -> Dict[str, Any]: await self.stop_workflow(workflow_id) try: stored = await WorkflowStore.get_config(workflow_id, kind="workflow_poller_config") @@ -237,8 +257,7 @@ async def restart_workflow(self, workflow_id: str) -> Dict[str, Any]: log.warning("poller.restart_read_failed", {"workflow_id": workflow_id, "error": str(exc)}) return {"workflowId": workflow_id, "state": "failed", "error": str(exc)} - config = self._normalize_config(workflow_id, stored) - if not config.get("enabled"): + if not isinstance(stored, dict) or not stored.get("enabled"): self._status[workflow_id] = { **self._base_status(workflow_id), "workflowId": workflow_id, @@ -253,11 +272,20 @@ async def restart_workflow(self, workflow_id: str) -> Dict[str, Any]: self._status[workflow_id] = { **self.get_status(workflow_id), "workflowId": workflow_id, - "state": "failed", + "state": "stopped" if startup else "failed", "error": err, } + poller_log = log.info if startup else log.warning + poller_log( + "poller.workflow_not_found_on_start" if startup else "poller.workflow_not_found", + { + "workflow_id": workflow_id, + **({"action": "stale_config_skipped"} if startup else {}), + }, + ) return self.get_status(workflow_id) + config = self._normalize_config(workflow_id, stored) workflow_json = wf_data.get("workflowJson") if not workflow_json: err = "workflow_json_missing" @@ -434,7 +462,12 @@ async def _execute_run( current["activeRuns"] = self._cleanup_done_runs(workflow_id) self._status[workflow_id] = current + tool_context = None try: + tool_context = await build_workflow_tool_context( + workflow_id=workflow_id, + action_name="trigger:schedule", + ) result = await asyncio.to_thread( run_workflow, workflow=workflow_json, @@ -445,6 +478,7 @@ async def _execute_run( execution_profile="high_frequency", cancel=cancel_event.is_set, on_step_complete=step_recorder.on_step_complete, + tool_context=tool_context, ) if not isinstance(result, RunWorkflowResult): result = RunWorkflowResult(status="failed", error="invalid_run_result") @@ -521,6 +555,7 @@ async def _execute_run( self._status[workflow_id] = current log.warning("poller.run_failed", {"workflow_id": workflow_id, "error": str(exc)}) finally: + await cleanup_workflow_tool_context(tool_context) try: await record_execution_result(workflow_id, exec_id, exec_data) except Exception as exc: diff --git a/flocks/workflow/service_port.py b/flocks/workflow/service_port.py new file mode 100644 index 000000000..c1c27c270 --- /dev/null +++ b/flocks/workflow/service_port.py @@ -0,0 +1,50 @@ +"""Shared workflow service port parsing helpers.""" + +from __future__ import annotations + +from collections.abc import Iterator, Mapping, Sequence +from typing import Any, Optional +from urllib.parse import urlparse + +SERVICE_PORT_KEYS = ("port", "servicePort", "hostPort") +SERVICE_URL_KEYS = ("serviceUrl", "invokeUrl") + + +def _valid_port(value: Any) -> Optional[int]: + if isinstance(value, bool): + return None + try: + port = int(value or 0) + except (TypeError, ValueError): + return None + return port if 1 <= port <= 65535 else None + + +def _iter_service_ports(record: Any, keys: Sequence[str]) -> Iterator[int]: + if not isinstance(record, Mapping): + return + for key in keys: + port = _valid_port(record.get(key)) + if port is not None: + yield port + for key in SERVICE_URL_KEYS: + try: + port = urlparse(str(record.get(key) or "")).port + except (TypeError, ValueError): + port = None + if port is not None: + yield port + + +def resolve_service_port( + record: Any, + *, + keys: Sequence[str] = SERVICE_PORT_KEYS, +) -> Optional[int]: + """Resolve the first valid explicit or URL-derived port in a service record.""" + return next(_iter_service_ports(record, keys), None) + + +def collect_service_ports(record: Any) -> set[int]: + """Collect every valid port represented by a service record.""" + return set(_iter_service_ports(record, SERVICE_PORT_KEYS)) diff --git a/flocks/workflow/tool_context.py b/flocks/workflow/tool_context.py index 859ddcc20..e06ad3b42 100644 --- a/flocks/workflow/tool_context.py +++ b/flocks/workflow/tool_context.py @@ -11,8 +11,11 @@ from flocks.session.message import Message, MessageRole from flocks.session.session import Session from flocks.tool import ToolContext +from flocks.utils.log import Log from flocks.workflow.fs_store import find_workspace_root +log = Log.create(service="workflow.tool_context") + async def build_workflow_tool_context( *, @@ -33,6 +36,7 @@ async def build_workflow_tool_context( effective_session_id = str(session_id or "").strip() effective_message_id = str(message_id or "").strip() effective_agent = str(agent or "").strip() + created_temp_parent = not effective_session_id workspace_dir = os.getcwd() project_id = "default" @@ -93,5 +97,56 @@ async def build_workflow_tool_context( extra={ "workspace_dir": workspace_dir, "main_session_key": effective_session_id, + "workflow_temp_parent": created_temp_parent, }, ) + + +async def cleanup_workflow_tool_context(tool_context: Optional[ToolContext]) -> bool: + """Purge an unused temporary workflow parent session. + + Parents with delegated child sessions are retained so returned task session + IDs remain valid. Caller-provided sessions are never modified. + """ + if tool_context is None: + return False + extra = getattr(tool_context, "extra", None) + if not isinstance(extra, dict) or extra.get("workflow_temp_parent") is not True: + return False + + session_id = str(getattr(tool_context, "session_id", "") or "").strip() + if not session_id: + return False + + try: + session = await Session.get_by_id(session_id) + metadata = getattr(session, "metadata", None) if session is not None else None + if session is None or not isinstance(metadata, dict) or metadata.get("workflowTempParent") is not True: + return False + + if extra.get("workflow_child_session_created") is True: + return False + + await Session.delete(session.project_id, session.id) + + # Session.delete is intentionally a soft delete. Trigger parents without + # child tasks are implementation details, so remove their residual rows + # to keep high-frequency trigger traffic storage-bounded. + from flocks.storage.storage import Storage + + await Storage.delete(f"session:{session.project_id}:{session.id}") + await Storage.delete(f"message:{session.id}") + await Storage.delete(f"todo:{session.id}") + await Storage.delete(f"goal:{session.id}") + await Storage.delete(f"session_diff:{session.id}") + await Storage.clear(prefix=f"message_diff:{session.id}:") + await Storage.clear(prefix=f"system_prompts:{session.id}:") + + Message.invalidate_cache(session.id) + return True + except Exception as exc: + log.warning( + "workflow.tool_context.cleanup_failed", + {"session_id": session_id, "error": str(exc)}, + ) + return False diff --git a/flocks/workflow/triggers/runtime.py b/flocks/workflow/triggers/runtime.py index 74cb329e2..7230bf0e2 100644 --- a/flocks/workflow/triggers/runtime.py +++ b/flocks/workflow/triggers/runtime.py @@ -18,6 +18,10 @@ from flocks.workflow.fs_store import read_workflow_dir, workflow_scan_dirs from flocks.workflow.runner import run_workflow from flocks.workflow.store import WorkflowStore +from flocks.workflow.tool_context import ( + build_workflow_tool_context, + cleanup_workflow_tool_context, +) from .compat import ( LEGACY_KAFKA_CONFIG_PREFIX, @@ -218,12 +222,18 @@ async def _execute_workflow( ) exec_id = exec_data["id"] started_at = time.time() + tool_context = None try: + tool_context = await build_workflow_tool_context( + workflow_id=workflow_id, + action_name=f"trigger:{trigger.type}", + ) result = await asyncio.to_thread( run_workflow, workflow=workflow_json, inputs=mapped_inputs, trace=False, + tool_context=tool_context, ) status_value, error_message = resolve_execution_outcome(result) exec_data.update( @@ -258,6 +268,8 @@ async def _execute_workflow( "triggerSource": mapped_inputs.get("_flocks", {}).get("trigger", {}).get("source"), } ) + finally: + await cleanup_workflow_tool_context(tool_context) await record_execution_result(workflow_id, exec_id, exec_data) return exec_data diff --git a/pyproject.toml b/pyproject.toml index 783712fde..5c58a57f4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "flocks" -version = "v2026.7.15" +version = "v2026.7.22" description = "AI-Native SecOps platform with multi-agent collaboration" authors = [ {name = "Flocks Team", email = "team@example.com"} @@ -71,6 +71,9 @@ dependencies = [ "lark-oapi>=1.3.0", # DingTalk Stream Mode (WebSocket long-connection) "dingtalk-stream>=0.20", + # Slack Socket Mode + "slack-bolt>=1.26.0", + "slack-sdk>=3.38.0", "python-socks>=2.8.1", "pytz>=2026.1.post1", # docs parser diff --git a/scripts/container-start.sh b/scripts/container-start.sh index bca3dcf1e..752570a94 100644 --- a/scripts/container-start.sh +++ b/scripts/container-start.sh @@ -5,86 +5,18 @@ set -euo pipefail ROOT_DIR="${ROOT_DIR:-/opt/flocks}" -BACKEND_HOST="${FLOCKS_BACKEND_HOST:-0.0.0.0}" -BACKEND_PORT="${FLOCKS_BACKEND_PORT:-8000}" -FRONTEND_HOST="${FLOCKS_FRONTEND_HOST:-0.0.0.0}" -FRONTEND_PORT="${FLOCKS_FRONTEND_PORT:-5173}" -FRONTEND_DIST_DIR="${FLOCKS_FRONTEND_DIST_DIR:-$ROOT_DIR/webui/dist}" -FRONTEND_PROXY_TARGET="${FLOCKS_FRONTEND_PROXY_TARGET:-http://127.0.0.1:${BACKEND_PORT}}" - -backend_pid="" -frontend_pid="" -ready_hint_pid="" +service_started=false cleanup() { - local pid - for pid in "$backend_pid" "$frontend_pid" "$ready_hint_pid"; do - if [[ -n "$pid" ]] && kill -0 "$pid" 2>/dev/null; then - kill -TERM "$pid" 2>/dev/null || true - fi - done - - for pid in "$backend_pid" "$frontend_pid" "$ready_hint_pid"; do - if [[ -n "$pid" ]]; then - wait "$pid" 2>/dev/null || true - fi - done + trap - EXIT INT TERM + if [[ "$service_started" == true ]]; then + flocks stop || true + fi } -trap cleanup EXIT +trap cleanup EXIT INT TERM cd "$ROOT_DIR" -[[ -f "$FRONTEND_DIST_DIR/index.html" ]] || { - printf '[flocks] error: WebUI 构建产物不存在: %s\n' "$FRONTEND_DIST_DIR/index.html" >&2 - exit 1 -} -[[ -f "$ROOT_DIR/scripts/serve_webui.py" ]] || { - printf '[flocks] error: 缺少前端静态服务脚本: %s\n' "$ROOT_DIR/scripts/serve_webui.py" >&2 - exit 1 -} - -printf '[flocks] publish ports to access from host: -p %s:%s -p %s:%s\n' \ - "$BACKEND_PORT" "$BACKEND_PORT" "$FRONTEND_PORT" "$FRONTEND_PORT" - -printf '[flocks] starting backend on %s:%s\n' "$BACKEND_HOST" "$BACKEND_PORT" -python -m uvicorn flocks.server.app:app \ - --host "$BACKEND_HOST" \ - --port "$BACKEND_PORT" & -backend_pid=$! - -printf '[flocks] starting frontend on %s:%s (proxy -> %s)\n' \ - "$FRONTEND_HOST" "$FRONTEND_PORT" "$FRONTEND_PROXY_TARGET" -( exec python "$ROOT_DIR/scripts/serve_webui.py" \ - --directory "$FRONTEND_DIST_DIR" \ - --host "$FRONTEND_HOST" \ - --port "$FRONTEND_PORT" \ - --proxy-target "$FRONTEND_PROXY_TARGET" ) & -frontend_pid=$! - -( - while kill -0 "$backend_pid" 2>/dev/null && kill -0 "$frontend_pid" 2>/dev/null; do - if python - "$BACKEND_PORT" <<'PY' -import sys -import urllib.request - -port = sys.argv[1] -url = f"http://127.0.0.1:{port}/api/health" - -try: - with urllib.request.urlopen(url, timeout=1) as response: - sys.exit(0 if response.status == 200 else 1) -except Exception: - sys.exit(1) -PY - then - printf '[flocks] open WebUI in your browser: http://127.0.0.1:%s\n' "$FRONTEND_PORT" - exit 0 - fi - - sleep 1 - done -) & -ready_hint_pid=$! - -wait -n "$backend_pid" "$frontend_pid" -exit $? +flocks start --no-browser --skip-webui-build +service_started=true +flocks logs --follow diff --git a/scripts/validate_flockshub.py b/scripts/validate_flockshub.py index 5aa356aaf..7068f7cdc 100644 --- a/scripts/validate_flockshub.py +++ b/scripts/validate_flockshub.py @@ -80,6 +80,30 @@ def main() -> int: if not (package_dir / entrypoint).exists(): fail(f"Missing entrypoint {entrypoint} in {manifest_rel}") + python_files = { + package_dir / entrypoint + for entrypoint in manifest.get("entrypoints", []) + if Path(entrypoint).suffix == ".py" + } + if plugin_type in {"tool", "device"}: + python_files.update( + path + for path in package_dir.rglob("*.py") + if not any( + part in {".git", "__pycache__"} + for part in path.relative_to(package_dir).parts + ) + ) + for path in sorted(python_files): + if not path.is_file(): + continue + try: + source = path.read_text(encoding="utf-8") + compile(source, str(path), "exec") + except (SyntaxError, UnicodeDecodeError) as exc: + relative_path = path.relative_to(package_dir).as_posix() + fail(f"Invalid Python source {relative_path} in {manifest_rel}: {exc}") + for path in package_dir.rglob("*"): rel = path.relative_to(package_dir).as_posix() ensure_relative(rel) diff --git a/tests/agent/test_agent_factory.py b/tests/agent/test_agent_factory.py index 30ce129e6..78f86fdac 100644 --- a/tests/agent/test_agent_factory.py +++ b/tests/agent/test_agent_factory.py @@ -15,6 +15,7 @@ import textwrap from pathlib import Path +from unittest.mock import MagicMock import pytest @@ -780,6 +781,43 @@ def test_user_plugin_agent_is_not_native( assert "custom-agent" in result assert result["custom-agent"].native is False + def test_user_plugin_agent_wins_over_project_bundle( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ): + user_agents_dir = tmp_path / "user_plugins" / "agents" + project_dir = tmp_path / "project" + project_agents_dir = project_dir / ".flocks" / "plugins" / "agents" + self._write_agent(user_agents_dir, "shared-agent") + self._write_agent(project_agents_dir, "shared-agent") + (user_agents_dir / "shared-agent" / "agent.yaml").write_text( + "name: shared-agent\ndescription: user customization\nmode: subagent\n", + encoding="utf-8", + ) + (project_agents_dir / "shared-agent" / "agent.yaml").write_text( + "name: shared-agent\ndescription: project bundle\nmode: subagent\n", + encoding="utf-8", + ) + + info_log = MagicMock() + warn_log = MagicMock() + monkeypatch.setattr(_factory_module, "_PLUGIN_AGENTS_DIR", user_agents_dir) + monkeypatch.setattr(_factory_module.log, "info", info_log) + monkeypatch.setattr(_factory_module.log, "warn", warn_log) + monkeypatch.chdir(project_dir) + + result = scan_and_load() + + assert result["shared-agent"].description == "user customization" + assert any( + call.args and call.args[0] == "agent.factory.lower_priority_skipped" + for call in info_log.call_args_list + ) + assert not any( + call.args and call.args[0] == "agent.factory.name_conflict" + and call.args[1].get("name") == "shared-agent" + for call in warn_log.call_args_list + ) + # =========================================================================== # Project-level plugin agent CRUD diff --git a/tests/browser/test_admin.py b/tests/browser/test_admin.py index 52c786431..0ee22f407 100644 --- a/tests/browser/test_admin.py +++ b/tests/browser/test_admin.py @@ -1,3 +1,7 @@ +from pathlib import Path + +import pytest + from flocks.browser import admin @@ -57,6 +61,56 @@ def test_generic_remote_debugging_message_triggers_prompt() -> None: assert admin._needs_chrome_remote_debugging_prompt(msg) +def test_current_remote_debugging_unreachable_message_triggers_prompt() -> None: + msg = ( + "Chromium-based browser remote debugging is not reachable for any detected profile: " + "/tmp/chrome: 127.0.0.1:9222 (connection refused) — for manual setup, " + "start the browser with --remote-debugging-port and a non-default --user-data-dir" + ) + assert admin._needs_chrome_remote_debugging_prompt(msg) + + +def test_local_debugging_setup_lines_use_windows_user_data_dir() -> None: + lines = "\n".join(admin._local_debugging_setup_lines("Windows")) + + assert "chrome://inspect" in lines + assert "does not reliably show it" in lines + assert "--remote-debugging-port=9222" in lines + assert '--user-data-dir="$env:USERPROFILE\\.flocks\\chrome-debug-profile"' in lines + assert "msedge.exe" in lines + assert '--user-data-dir="$env:USERPROFILE\\.flocks\\edge-debug-profile"' in lines + assert "chromium.exe" in lines + assert "brave.exe" in lines + assert "http://127.0.0.1:9222/json/version" in lines + assert "rerun `flocks browser --setup`" in lines + + +def test_local_debugging_setup_lines_list_chromium_browser_candidates() -> None: + mac_lines = "\n".join(admin._local_debugging_setup_lines("Darwin")) + linux_lines = "\n".join(admin._local_debugging_setup_lines("Linux")) + + assert "Microsoft\\ Edge.app" in mac_lines + assert "Chromium.app" in mac_lines + assert "Brave\\ Browser.app" in mac_lines + assert "microsoft-edge" in linux_lines + assert "chromium" in linux_lines + assert "brave-browser" in linux_lines + + +def test_browser_skill_docs_do_not_reintroduce_inspect_allow_setup_flow() -> None: + repo_root = Path(__file__).resolve().parents[2] + docs = [ + repo_root / ".flocks/plugins/skills/browser-use/SKILL.md", + repo_root / ".flocks/plugins/skills/browser-use/references/cdp-direct.md", + repo_root / ".flocks/plugins/skills/browser-use/references/cdp-setup.md", + ] + text = "\n".join(path.read_text(encoding="utf-8") for path in docs) + + assert "chrome://inspect/#remote-debugging" not in text + assert "edge://inspect/#remote-debugging" not in text + assert "Allow remote debugging" not in text + + def test_daemon_endpoint_names_discovers_default_and_named_sessions(tmp_path, monkeypatch) -> None: monkeypatch.setattr(admin.ipc, "IS_WINDOWS", False) monkeypatch.setenv("FLOCKS_ROOT", str(tmp_path)) @@ -189,6 +243,40 @@ def fake_popen(*args, **kwargs): assert len(spawned) == 2 +def test_ensure_daemon_prints_manual_guidance_without_blind_retry(tmp_path, monkeypatch, capsys) -> None: + spawned = [] + restarted = [] + + class FakeProcess: + def poll(self): + return 1 + + def fake_popen(*args, **kwargs): + spawned.append((args, kwargs)) + return FakeProcess() + + monkeypatch.setenv("FLOCKS_ROOT", str(tmp_path)) + monkeypatch.setattr(admin.ipc, "endpoint_reachable", lambda name: False) + monkeypatch.setattr(admin, "daemon_alive", lambda name: False) + monkeypatch.setattr("subprocess.Popen", fake_popen) + monkeypatch.setattr( + admin, + "_log_tail", + lambda name: "Chromium-based browser remote debugging is not reachable for any detected profile", + ) + monkeypatch.setattr(admin, "restart_daemon", lambda name=None: restarted.append(name)) + monkeypatch.setattr(admin, "_local_debugging_setup_lines", lambda: ["debug browser instructions"]) + + with pytest.raises(RuntimeError): + admin.ensure_daemon(wait=1.0, name="manual-session") + + err = capsys.readouterr().err + assert "local Chromium-based browser remote debugging is not reachable" in err + assert "debug browser instructions" in err + assert len(spawned) == 1 + assert restarted == ["manual-session"] + + def test_run_doctor_prints_active_browser_connections_and_active_pages(monkeypatch, capsys) -> None: monkeypatch.setattr(admin, "_version", lambda: "0.1.0") monkeypatch.setattr(admin, "_install_mode", lambda: "git") @@ -275,19 +363,21 @@ def test_run_doctor_suggests_setup_when_target_exists_but_daemon_missing(monkeyp def test_run_setup_uses_generic_missing_browser_wording(monkeypatch, capsys) -> None: monkeypatch.setattr(admin, "daemon_alive", lambda: False) monkeypatch.setattr(admin, "_chrome_running", lambda: False) + monkeypatch.setattr(admin, "_local_debugging_setup_lines", lambda: ["debug browser instructions"]) assert admin.run_setup() == 1 out = capsys.readouterr().out - assert "no Chrome/Chromium/Edge process detected" in out + assert "no Chrome/Chromium/Edge/Brave process detected" in out + assert "start a Chromium-based browser with remote debugging" in out + assert "debug browser instructions" in out -def test_run_setup_uses_generic_remote_debugging_wording(monkeypatch, capsys) -> None: +def test_run_setup_prints_remote_debugging_guidance_without_blind_retry(monkeypatch, capsys) -> None: monkeypatch.setattr(admin, "daemon_alive", lambda: False) monkeypatch.setattr(admin, "_chrome_running", lambda: True) monkeypatch.setattr(admin, "_is_local_chrome_mode", lambda env=None: True) - open_calls = [] - monkeypatch.setattr(admin, "_open_browser_inspect", lambda: open_calls.append(True)) + monkeypatch.setattr(admin, "_local_debugging_setup_lines", lambda: ["debug browser instructions"]) calls = {"count": 0} @@ -301,13 +391,13 @@ def fake_ensure_daemon(*args, **kwargs): monkeypatch.setattr(admin, "ensure_daemon", fake_ensure_daemon) - assert admin.run_setup() == 0 + assert admin.run_setup() == 1 out = capsys.readouterr().out - assert "browser remote debugging is not enabled on the current profile." in out - assert "opening your browser's inspect page" in out - assert "if the browser shows the profile picker" in out - assert open_calls == [True] + assert "Chromium-based browser remote debugging is not reachable for the current profile." in out + assert "debug browser instructions" in out + assert "opening your browser's inspect page" not in out + assert calls["count"] == 1 def test_run_setup_restarts_stale_existing_local_daemon(monkeypatch, capsys) -> None: @@ -325,7 +415,7 @@ def test_run_setup_restarts_stale_existing_local_daemon(monkeypatch, capsys) -> assert "browser connection is stale; restarting" in out assert "daemon is up." in out assert restarted == [None] - assert ensure_calls == [{"wait": admin._SETUP_ATTACH_WAIT, "_open_inspect": False}] + assert ensure_calls == [{"wait": admin._SETUP_ATTACH_WAIT, "_show_debugging_guidance": False}] def test_run_setup_retries_at_most_once(monkeypatch, capsys) -> None: @@ -345,8 +435,8 @@ def fake_ensure_daemon(**kwargs): out = capsys.readouterr() assert "retrying once" in out.out assert ensure_calls == [ - {"wait": admin._SETUP_ATTACH_WAIT, "_open_inspect": False}, - {"wait": admin._SETUP_RETRY_WAIT, "_open_inspect": False}, + {"wait": admin._SETUP_ATTACH_WAIT, "_show_debugging_guidance": False}, + {"wait": admin._SETUP_RETRY_WAIT, "_show_debugging_guidance": False}, ] @@ -362,7 +452,7 @@ def test_run_setup_allows_explicit_remote_cdp_without_local_browser(monkeypatch, out = capsys.readouterr().out assert "attaching via BU_CDP_WS" in out assert "daemon is up." in out - assert ensure_calls == [{"wait": admin._SETUP_ATTACH_WAIT, "_open_inspect": False}] + assert ensure_calls == [{"wait": admin._SETUP_ATTACH_WAIT, "_show_debugging_guidance": False}] def test_run_setup_restarts_existing_daemon_for_explicit_remote_cdp(monkeypatch, capsys) -> None: @@ -383,7 +473,7 @@ def test_run_setup_restarts_existing_daemon_for_explicit_remote_cdp(monkeypatch, assert "restarting to attach via BU_CDP_URL" in out assert "daemon is up." in out assert restarted == [None] - assert ensure_calls == [{"wait": admin._SETUP_ATTACH_WAIT, "_open_inspect": False}] + assert ensure_calls == [{"wait": admin._SETUP_ATTACH_WAIT, "_show_debugging_guidance": False}] def test_run_doctor_uses_generic_browser_wording_when_missing(monkeypatch, capsys) -> None: @@ -398,8 +488,24 @@ def test_run_doctor_uses_generic_browser_wording_when_missing(monkeypatch, capsy out = capsys.readouterr().out assert "[FAIL] browser running" in out - assert "start Chrome, Chromium, or Edge and rerun `flocks browser --setup`" in out - assert "next action start Chrome/Chromium/Edge or provide BU_CDP_URL/BU_CDP_WS" in out + assert "start Chrome, Chromium, Edge, or Brave and rerun `flocks browser --setup`" in out + assert "next action start Chrome/Chromium/Edge/Brave or provide BU_CDP_URL/BU_CDP_WS" in out + + +def test_run_doctor_points_to_debug_browser_instructions_when_daemon_missing(monkeypatch, capsys) -> None: + monkeypatch.setattr(admin, "_version", lambda: "0.1.0") + monkeypatch.setattr(admin, "_install_mode", lambda: "git") + monkeypatch.setattr(admin, "_chrome_running", lambda: True) + monkeypatch.setattr(admin, "daemon_alive", lambda: False) + monkeypatch.setattr(admin, "browser_connections", lambda: []) + monkeypatch.setattr(admin, "_latest_release_tag", lambda: "0.1.0") + + assert admin.run_doctor() == 1 + + out = capsys.readouterr().out + assert "follow its debug-browser instructions" in out + assert "remote debugging is not reachable" in out + assert "inspect-page prompt" not in out def test_run_doctor_accepts_explicit_remote_cdp_without_local_browser(monkeypatch, capsys) -> None: @@ -530,8 +636,22 @@ def test_chrome_running_on_windows_detects_chromium_process(monkeypatch) -> None assert admin._chrome_running() +def test_chrome_running_on_windows_detects_brave_process(monkeypatch) -> None: + monkeypatch.setattr("platform.system", lambda: "Windows") + monkeypatch.setattr("subprocess.check_output", lambda *args, **kwargs: b"brave.exe\r\n") + + assert admin._chrome_running() + + def test_chrome_running_on_non_windows_matches_text_output(monkeypatch) -> None: monkeypatch.setattr("platform.system", lambda: "Darwin") monkeypatch.setattr("subprocess.check_output", lambda *args, **kwargs: "Google Chrome\n") assert admin._chrome_running() + + +def test_chrome_running_on_non_windows_detects_brave_process(monkeypatch) -> None: + monkeypatch.setattr("platform.system", lambda: "Darwin") + monkeypatch.setattr("subprocess.check_output", lambda *args, **kwargs: "Brave Browser\n") + + assert admin._chrome_running() diff --git a/tests/browser/test_daemon.py b/tests/browser/test_daemon.py index b791cb869..b872a506a 100644 --- a/tests/browser/test_daemon.py +++ b/tests/browser/test_daemon.py @@ -163,6 +163,12 @@ def test_is_real_page_accepts_normal_https_pages() -> None: def test_profile_dirs_only_returns_paths_for_requested_os() -> None: home = Path.home() / "profile-test-home" local_app_data = home / "AppData/Local" + flocks_debug_profiles = [ + home / ".flocks/chrome-debug-profile", + home / ".flocks/edge-debug-profile", + home / ".flocks/chromium-debug-profile", + home / ".flocks/brave-debug-profile", + ] mac_profiles = daemon.profile_dirs(system="Darwin", home=home, environ={}) linux_profiles = daemon.profile_dirs(system="Linux", home=home, environ={}) @@ -172,9 +178,12 @@ def test_profile_dirs_only_returns_paths_for_requested_os() -> None: environ={"LOCALAPPDATA": str(local_app_data)}, ) - assert all("Library" in path.parts and "Application Support" in path.parts for path in mac_profiles) - assert all(".config" in path.parts or ".var" in path.parts for path in linux_profiles) - assert all(path.is_relative_to(local_app_data) for path in windows_profiles) + assert mac_profiles[:4] == flocks_debug_profiles + assert linux_profiles[:4] == flocks_debug_profiles + assert windows_profiles[:4] == flocks_debug_profiles + assert all("Library" in path.parts and "Application Support" in path.parts for path in mac_profiles[4:]) + assert all(".config" in path.parts or ".var" in path.parts for path in linux_profiles[4:]) + assert all(path.is_relative_to(local_app_data) for path in windows_profiles[4:]) def test_get_ws_url_skips_unreachable_profile_and_uses_next_candidate(tmp_path, monkeypatch) -> None: @@ -302,6 +311,35 @@ def fake_urlopen(url: str, timeout: float): assert daemon.get_ws_url() == "ws://127.0.0.1:9222/devtools/browser/current" +def test_get_ws_url_uses_devtools_active_port_path_when_version_omits_websocket_url( + tmp_path, monkeypatch +) -> None: + profile = tmp_path / "chrome" + profile.mkdir() + (profile / "DevToolsActivePort").write_text( + "9222\n/devtools/browser/current\n", + encoding="utf-8", + ) + + class FakeHttpResponse: + def __enter__(self): + return self + + def __exit__(self, *_args) -> None: + pass + + def read(self) -> bytes: + return b'{"Browser":"Chrome/126.0","Protocol-Version":"1.3"}' + + monkeypatch.delenv("BU_CDP_WS", raising=False) + monkeypatch.delenv("BU_CDP_URL", raising=False) + monkeypatch.setattr(daemon, "profile_dirs", lambda: [profile]) + monkeypatch.setattr(daemon.urllib.request, "urlopen", lambda _url, timeout: FakeHttpResponse()) + monkeypatch.setattr(daemon.time, "sleep", lambda _seconds: pytest.fail("missing field must not be retried")) + + assert daemon.get_ws_url() == "ws://127.0.0.1:9222/devtools/browser/current" + + def test_load_env_uses_shared_loader_for_existing_files(tmp_path, monkeypatch) -> None: workspace = tmp_path / "workspace" workspace.mkdir() diff --git a/tests/channel/test_channel.py b/tests/channel/test_channel.py index ca43b75fb..34a7830cf 100644 --- a/tests/channel/test_channel.py +++ b/tests/channel/test_channel.py @@ -32,6 +32,7 @@ ChatType, DeliveryResult, InboundMessage, + NonRetryableChannelError, OutboundContext, ) from flocks.config.config import ChannelAccountConfig, ChannelConfig, ConfigInfo @@ -598,6 +599,106 @@ async def fake_deliver(ctx, session_id=None): assert update_mock.await_args.args == ("channel", "session_old") assert update_mock.await_args.kwargs["status"] == "archived" + @pytest.mark.asyncio + async def test_new_command_with_args_sends_args_to_new_session(self, monkeypatch): + from flocks.channel.inbound.dispatcher import InboundDispatcher + from flocks.channel.inbound.session_binding import SessionBinding + + dispatcher = InboundDispatcher() + dispatcher._trigger_command_hook = AsyncMock() + append_mock = AsyncMock() + run_mock = AsyncMock() + monkeypatch.setattr(dispatcher, "_append_user_message", append_mock) + monkeypatch.setattr(dispatcher, "_run_agent", run_mock) + dispatcher.binding_service.rebind = AsyncMock( + return_value=SessionBinding( + channel_id="slack", + account_id="T1", + chat_id="C123", + chat_type=ChatType.CHANNEL, + thread_id="171.1", + session_id="session_new", + agent_id="rex", + created_at=0, + last_message_at=0, + ) + ) + binding = SessionBinding( + channel_id="slack", + account_id="T1", + chat_id="C123", + chat_type=ChatType.CHANNEL, + thread_id="171.1", + session_id="session_old", + agent_id="rex", + created_at=0, + last_message_at=0, + ) + msg = InboundMessage( + channel_id="slack", + account_id="T1", + message_id="171.1", + sender_id="U123", + chat_id="C123", + chat_type=ChatType.CHANNEL, + text="<@UBOT> /new 叫我uuuu", + mention_text="/new 叫我uuuu", + thread_id="171.1", + mentioned=True, + ) + + delivered: list[str] = [] + + async def fake_deliver(ctx, session_id=None): + delivered.append(ctx.text) + + monkeypatch.setattr( + "flocks.channel.outbound.deliver.OutboundDelivery.deliver", + fake_deliver, + ) + monkeypatch.setattr( + "flocks.session.session.Session.get_by_id", + AsyncMock( + return_value=SimpleNamespace( + id="session_old", + project_id="channel", + directory="/tmp/project", + agent="rex", + provider="anthropic", + model="claude-sonnet-4-20250514", + model_pinned=True, + ) + ), + ) + monkeypatch.setattr( + "flocks.session.session.Session.create", + AsyncMock(return_value=SimpleNamespace(id="session_new", agent="rex")), + ) + monkeypatch.setattr( + "flocks.session.session.Session.update", + AsyncMock(return_value=None), + ) + monkeypatch.setattr( + "flocks.channel.inbound.dispatcher._resolve_session_model", + AsyncMock(return_value={"providerID": "anthropic", "modelID": "claude"}), + ) + + handled = await dispatcher._handle_feishu_native_command( + binding=binding, + msg=msg, + channel_config=ChannelConfig(enabled=True), + user_text="/new 叫我uuuu", + scope_override=None, + ) + + assert handled is True + assert "已开始全新对话。" in delivered[0] + append_mock.assert_awaited_once() + assert append_mock.await_args.args[:3] == ("session_new", "叫我uuuu", msg) + assert append_mock.await_args.kwargs["agent"] == "rex" + run_mock.assert_awaited_once() + assert run_mock.await_args.args[0].session_id == "session_new" + @pytest.mark.asyncio async def test_reset_alias_matches_new_semantics(self, monkeypatch): from flocks.channel.inbound.dispatcher import InboundDispatcher @@ -1083,6 +1184,23 @@ def test_dm_allowlist_passes_listed(self): cfg = ChannelConfig(dm_policy="allowlist", allow_from=["u1"]) assert _check_allowlist(self._make_msg(sender_id="u1"), cfg) is True + def test_slack_dm_allow_from_without_policy_blocks_unlisted(self): + from flocks.channel.inbound.dispatcher import _check_allowlist + cfg = ChannelConfig(allow_from=["U_ALLOWED"]) + msg = self._make_msg(channel_id="slack", sender_id="U_BLOCKED") + assert _check_allowlist(msg, cfg) is False + + def test_slack_dm_allow_from_without_policy_passes_listed(self): + from flocks.channel.inbound.dispatcher import _check_allowlist + cfg = ChannelConfig(allow_from=["U_ALLOWED"]) + msg = self._make_msg(channel_id="slack", sender_id="U_ALLOWED") + assert _check_allowlist(msg, cfg) is True + + def test_non_slack_dm_allow_from_without_policy_stays_open(self): + from flocks.channel.inbound.dispatcher import _check_allowlist + cfg = ChannelConfig(allow_from=["u2"]) + assert _check_allowlist(self._make_msg(channel_id="test", sender_id="u1"), cfg) is True + def test_dm_allowlist_empty_blocks_all(self): from flocks.channel.inbound.dispatcher import _check_allowlist cfg = ChannelConfig(dm_policy="allowlist", allow_from=None) @@ -1407,6 +1525,70 @@ def test_record_error_updates_status(self): assert plugin.status.last_error == "fail" assert plugin.status.error_count == 1 + async def test_run_with_reconnect_stops_on_non_retryable_error(self): + from flocks.channel.gateway.manager import GatewayManager + + manager = GatewayManager() + plugin = _StubChannel() + attempts = 0 + abort_event = asyncio.Event() + + async def _fail_start(config, on_message, abort_event=None): + nonlocal attempts + attempts += 1 + raise NonRetryableChannelError("bad credentials") + + plugin.start = _fail_start # type: ignore[method-assign] + + await manager._run_with_reconnect( + channel_id="stub", + plugin=plugin, + config={}, + on_message=AsyncMock(), + abort_event=abort_event, + ) + + assert attempts == 1 + assert plugin.status.connected is False + assert plugin.status.last_error == "bad credentials" + assert plugin.status.error_count == 1 + + async def test_run_with_reconnect_does_not_pre_mark_self_managed_connection(self): + from flocks.channel.gateway.manager import GatewayManager + + class _SelfManagedChannel(_StubChannel): + def capabilities(self) -> ChannelCapabilities: + return ChannelCapabilities( + chat_types=[ChatType.DIRECT], + self_managed_connection=True, + ) + + manager = GatewayManager() + plugin = _SelfManagedChannel() + started = asyncio.Event() + abort_event = asyncio.Event() + + async def _start(config, on_message, abort_event=None): + started.set() + await asyncio.sleep(0.75) + abort_event.set() + raise RuntimeError("handshake failed") + + plugin.start = _start # type: ignore[method-assign] + + await manager._run_with_reconnect( + channel_id="stub", + plugin=plugin, + config={}, + on_message=AsyncMock(), + abort_event=abort_event, + ) + + assert started.is_set() + assert plugin.status.connected is False + assert plugin.status.last_error == "handshake failed" + assert plugin.status.error_count == 1 + async def test_stop_all_drains_cancelled_tasks(self, monkeypatch): from flocks.channel.gateway.manager import GatewayManager from flocks.channel.registry import ChannelRegistry @@ -1551,6 +1733,34 @@ async def fake_telegram(msg, config): assert result is not None assert captured["msg"].channel_id == "telegram" + @pytest.mark.asyncio + async def test_dispatch_to_slack_downloader(self, monkeypatch): + from flocks.channel.inbound import dispatcher as dispatch_mod + + captured = {} + + async def fake_slack(msg, config): + captured["msg"] = msg + captured["config"] = config + return SimpleNamespace( + filename="slack.png", mime="image/png", + url="file:///tmp/slack.png", source={"channel": "slack"}, + ) + + import flocks.channel.builtin.slack.inbound_media as slack_inb + monkeypatch.setattr(slack_inb, "download_inbound_media", fake_slack) + + result = await dispatch_mod._download_channel_media( + InboundMessage( + channel_id="slack", account_id="T1", message_id="m", + sender_id="u", media_url="https://files.slack.com/x", + ), + {"botToken": "xoxb-test"}, + ) + assert result is not None + assert captured["msg"].channel_id == "slack" + assert captured["config"] == {"botToken": "xoxb-test"} + @pytest.mark.asyncio async def test_unknown_channel_returns_none(self): from flocks.channel.inbound import dispatcher as dispatch_mod @@ -1573,6 +1783,7 @@ class TestIsPlaceholderText: 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 diff --git a/tests/channel/test_e2e_file_roundtrip.py b/tests/channel/test_e2e_file_roundtrip.py index f35441146..ab4a71e6b 100644 --- a/tests/channel/test_e2e_file_roundtrip.py +++ b/tests/channel/test_e2e_file_roundtrip.py @@ -342,8 +342,8 @@ async def fake_get_http_client(): # _download_file to use the server directly. async def fake_get_file_path(*, bot_token, api_base, file_id, timeout): return server.file_id_to_path[file_id], file_id - async def fake_download(*, bot_token, file_path, max_bytes, timeout): - return server.get(f"https://api.telegram.org/file/bot{bot_token}/{file_path}")._body + async def fake_download(*, download_base, file_path, max_bytes, timeout): + return server.get(f"{download_base.rstrip('/')}/{file_path}")._body monkeypatch.setattr(inbound_media, "_get_file_path", fake_get_file_path) monkeypatch.setattr(inbound_media, "_download_file", fake_download) diff --git a/tests/channel/test_slack.py b/tests/channel/test_slack.py new file mode 100644 index 000000000..1d5bdf4bc --- /dev/null +++ b/tests/channel/test_slack.py @@ -0,0 +1,906 @@ +from __future__ import annotations + +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from flocks.channel.base import ChatType, NonRetryableChannelError, OutboundContext +import flocks.channel.builtin.slack.channel as slack_mod +import flocks.channel.builtin.slack.inbound_media as slack_media_mod +from flocks.channel.builtin.slack.channel import SlackChannel +from flocks.channel.builtin.slack.format import markdown_to_slack_mrkdwn +from flocks.channel.builtin.slack.inbound import build_inbound_message, slack_thread_cache_key +from flocks.channel.builtin.slack.manifest import build_slack_app_manifest + + +def test_plugin_exports_slack_channel(): + plugin = SlackChannel() + assert plugin.meta().id == "slack" + assert plugin.meta().label == "Slack" + assert "sl" in plugin.meta().aliases + assert plugin.capabilities().media is True + assert plugin.capabilities().rich_text is True + assert plugin.capabilities().self_managed_connection is True + + +def test_slack_manifest_matches_socket_mode_setup_needs(): + manifest = build_slack_app_manifest() + + assert manifest["settings"]["socket_mode_enabled"] is True + scopes = set(manifest["oauth_config"]["scopes"]["bot"]) + assert { + "app_mentions:read", + "channels:history", + "channels:read", + "chat:write", + "files:read", + "groups:history", + "groups:read", + "im:history", + "im:read", + "im:write", + "mpim:history", + "mpim:read", + "users:read", + }.issubset(scopes) + events = set(manifest["settings"]["event_subscriptions"]["bot_events"]) + assert { + "app_mention", + "message.channels", + "message.groups", + "message.im", + "message.mpim", + }.issubset(events) + + +def test_validate_config_requires_tokens(monkeypatch): + monkeypatch.setattr(slack_mod, "SLACK_AVAILABLE", True) + plugin = SlackChannel() + + assert "botToken" in (plugin.validate_config({}) or "") + assert "appToken" in (plugin.validate_config({"botToken": "xoxb-1"}) or "") + assert plugin.validate_config({"botToken": "xoxb-1", "appToken": "xapp-1"}) is None + + +def test_validate_config_rejects_non_slack_token_prefixes(monkeypatch): + monkeypatch.setattr(slack_mod, "SLACK_AVAILABLE", True) + plugin = SlackChannel() + + assert "xoxb-" in ( + plugin.validate_config({"botToken": "xoxp-user", "appToken": "xapp-1"}) or "" + ) + assert "xapp-" in ( + plugin.validate_config({"botToken": "xoxb-1", "appToken": "xoxb-wrong"}) or "" + ) + assert plugin.validate_config( + {"botToken": "{secret:slack_bot}", "appToken": "{secret:slack_app}"} + ) is None + + +def test_validate_config_reports_missing_dependency(monkeypatch): + monkeypatch.setattr(slack_mod, "SLACK_AVAILABLE", False) + plugin = SlackChannel() + + assert "slack-bolt" in (plugin.validate_config({"botToken": "x", "appToken": "y"}) or "") + + +def test_markdown_to_slack_mrkdwn_preserves_code_and_links(): + text = "# Title\n\n**bold** and [OpenAI](https://openai.com)\n\n```py\n**not bold**\n```" + + rendered = markdown_to_slack_mrkdwn(text) + + assert "*Title*" in rendered + assert "*bold*" in rendered + assert "" in rendered + assert "```py\n**not bold**\n```" in rendered + + +def test_markdown_to_slack_mrkdwn_keeps_escaped_special_mentions(): + rendered = markdown_to_slack_mrkdwn("safe <!channel> and &") + + assert "<!channel>" in rendered + assert "" not in rendered + assert "&" in rendered + + +def test_markdown_to_slack_mrkdwn_escapes_raw_slack_special_mentions(): + rendered = markdown_to_slack_mrkdwn( + "notify <@U123>; compare 1 < 2 & 3 > 2" + ) + + assert "" not in rendered + assert "" not in rendered + assert "" not in rendered + assert "<@U123>" not in rendered + assert "<!here>" in rendered + assert "<!channel>" in rendered + assert "<!everyone>" in rendered + assert "<@U123>" in rendered + assert "1 < 2 & 3 > 2" in rendered + + +def test_markdown_to_slack_mrkdwn_escapes_unsafe_link_url_delimiters(): + rendered = markdown_to_slack_mrkdwn("[x](https://example.com/a|)") + + assert rendered == "" + assert "|" not in rendered + + +def test_build_inbound_direct_message(): + msg = build_inbound_message( + { + "channel": "D123", + "channel_type": "im", + "user": "U123", + "ts": "171.1", + "text": "hello", + "team": "T1", + }, + bot_user_id="UBOT", + config={}, + known_thread_ids=set(), + ) + + assert msg is not None + assert msg.channel_id == "slack" + assert msg.account_id == "T1" + assert msg.chat_type == ChatType.DIRECT + assert msg.chat_id == "D123" + assert msg.message_id == "171.1" + assert msg.text == "hello" + assert msg.mentioned is False + assert msg.mention_text == "" + + +def test_build_inbound_group_mention_strips_bot_mention(): + msg = build_inbound_message( + { + "channel": "C123", + "channel_type": "channel", + "user": "U123", + "ts": "171.2", + "text": "<@UBOT> summarize this", + }, + bot_user_id="UBOT", + config={}, + known_thread_ids=set(), + ) + + assert msg is not None + assert msg.chat_type == ChatType.CHANNEL + assert msg.mentioned is True + assert msg.mention_text == "summarize this" + + +def test_build_inbound_extracts_rich_text_blocks(): + msg = build_inbound_message( + { + "channel": "C123", + "channel_type": "channel", + "user": "U123", + "ts": "171.22", + "blocks": [ + { + "type": "rich_text", + "elements": [ + { + "type": "rich_text_section", + "elements": [ + {"type": "user", "user_id": "UBOT"}, + {"type": "text", "text": " summarize "}, + {"type": "link", "text": "this", "url": "https://example.com"}, + ], + } + ], + } + ], + }, + bot_user_id="UBOT", + config={}, + known_thread_ids=set(), + ) + + assert msg is not None + assert "<@UBOT> summarize this (https://example.com)" in msg.text + assert msg.mentioned is True + assert "summarize this" in msg.mention_text + + +def test_build_inbound_merges_blocks_when_plain_text_exists(): + msg = build_inbound_message( + { + "channel": "C123", + "channel_type": "channel", + "user": "U123", + "ts": "171.24", + "text": "<@UBOT> please inspect", + "blocks": [ + { + "type": "rich_text", + "elements": [ + { + "type": "rich_text_section", + "elements": [ + {"type": "user", "user_id": "UBOT"}, + {"type": "text", "text": " please inspect"}, + ], + }, + { + "type": "rich_text_quote", + "elements": [ + { + "type": "rich_text_section", + "elements": [ + {"type": "text", "text": "quoted outage context"}, + ], + } + ], + }, + ], + } + ], + }, + bot_user_id="UBOT", + config={}, + known_thread_ids=set(), + ) + + assert msg is not None + assert "please inspect" in msg.text + assert "quoted outage context" in msg.text + assert "quoted outage context" in msg.mention_text + + +def test_build_inbound_appends_attachments_and_file_names(): + msg = build_inbound_message( + { + "channel": "C123", + "channel_type": "channel", + "user": "U123", + "ts": "171.23", + "text": "<@UBOT> review alert", + "attachments": [ + { + "title": "Alert context", + "fields": [ + {"title": "Severity", "value": "high"}, + ], + } + ], + "files": [ + {"name": "screenshot.png"}, + ], + }, + bot_user_id="UBOT", + config={}, + known_thread_ids=set(), + ) + + assert msg is not None + assert "review alert" in msg.text + assert "Alert context" in msg.text + assert "Severity: high" in msg.text + assert "[文件消息: screenshot.png]" in msg.text + + +def test_build_inbound_image_file_sets_media_url_and_mime(): + msg = build_inbound_message( + { + "channel": "D123", + "channel_type": "im", + "user": "U123", + "ts": "171.25", + "team": "T1", + "files": [ + { + "id": "F123", + "name": "screenshot.png", + "mimetype": "image/png", + "url_private_download": "https://files.slack.com/files-pri/T1-F123/download/screenshot.png", + }, + ], + }, + bot_user_id="UBOT", + config={}, + known_thread_ids=set(), + ) + + assert msg is not None + assert msg.text == "[图片消息: screenshot.png]" + assert msg.media_url == "https://files.slack.com/files-pri/T1-F123/download/screenshot.png" + assert msg.media_mime == "image/png" + + +def test_build_inbound_file_stub_uses_slack_file_uri(): + msg = build_inbound_message( + { + "channel": "C123", + "channel_type": "channel", + "user": "U123", + "ts": "171.26", + "text": "<@UBOT> inspect", + "files": [ + { + "id": "F_STUB", + "name": "stub.pdf", + "mimetype": "application/pdf", + "file_access": "check_file_info", + }, + ], + }, + bot_user_id="UBOT", + config={}, + known_thread_ids=set(), + ) + + assert msg is not None + assert msg.media_url == "slack://file/F_STUB" + assert msg.media_mime == "application/pdf" + + +def test_build_inbound_uses_socket_mode_body_team_id_for_stable_dm_sessions(): + from flocks.channel.inbound.session_binding import _resolve_session_key + + event_without_team = { + "channel": "D123", + "channel_type": "im", + "user": "U123", + "ts": "171.27", + "text": "hello", + } + event_with_team = { + **event_without_team, + "ts": "171.28", + "team": "T1", + } + + first = build_inbound_message( + event_without_team, + bot_user_id="UBOT", + config={}, + known_thread_ids=set(), + body={"team_id": "T1"}, + ) + second = build_inbound_message( + event_with_team, + bot_user_id="UBOT", + config={}, + known_thread_ids=set(), + ) + + assert first is not None + assert second is not None + assert first.account_id == second.account_id == "T1" + assert _resolve_session_key(first) == _resolve_session_key(second) + + +def test_build_inbound_group_dm_uses_group_chat_type(): + msg = build_inbound_message( + { + "channel": "G123", + "channel_type": "mpim", + "user": "U123", + "ts": "171.21", + "text": "<@UBOT> summarize this", + }, + bot_user_id="UBOT", + config={}, + known_thread_ids=set(), + ) + + assert msg is not None + assert msg.chat_type == ChatType.GROUP + assert msg.mentioned is True + + +def test_build_inbound_thread_reply_to_known_bot_thread_triggers_without_mention(): + msg = build_inbound_message( + { + "channel": "C123", + "channel_type": "channel", + "user": "U123", + "ts": "171.3", + "thread_ts": "171.2", + "text": "continue", + }, + bot_user_id="UBOT", + config={}, + known_thread_ids={slack_thread_cache_key("default", "C123", "171.2")}, + ) + + assert msg is not None + assert msg.mentioned is True + assert msg.mention_text == "continue" + assert msg.thread_id == "171.2" + + +def test_build_inbound_same_thread_ts_in_other_channel_does_not_trigger(): + msg = build_inbound_message( + { + "channel": "C999", + "channel_type": "channel", + "user": "U123", + "ts": "171.3", + "thread_ts": "171.2", + "text": "same timestamp, different channel", + "team": "T1", + }, + bot_user_id="UBOT", + config={}, + known_thread_ids={slack_thread_cache_key("T1", "C123", "171.2")}, + ) + + assert msg is not None + assert msg.mentioned is False + + +def test_build_inbound_thread_reply_to_parent_bot_user_triggers_after_restart(): + msg = build_inbound_message( + { + "channel": "C123", + "channel_type": "channel", + "user": "U123", + "ts": "171.4", + "thread_ts": "171.2", + "parent_user_id": "UBOT", + "text": "continue after restart", + }, + bot_user_id="UBOT", + config={}, + known_thread_ids=set(), + ) + + assert msg is not None + assert msg.mentioned is True + assert msg.mention_text == "continue after restart" + assert msg.thread_id == "171.2" + + +def test_build_inbound_allowed_bot_message_without_user_uses_bot_id(): + msg = build_inbound_message( + { + "channel": "C123", + "channel_type": "channel", + "bot_id": "B_OTHER", + "subtype": "bot_message", + "ts": "171.5", + "text": "automation says hi", + }, + bot_user_id="UBOT", + config={"allowBots": "all"}, + known_thread_ids=set(), + ) + + assert msg is not None + assert msg.sender_id == "B_OTHER" + assert msg.text == "automation says hi" + + +def test_build_inbound_bot_message_mentions_policy_requires_bot_mention(): + blocked = build_inbound_message( + { + "channel": "C123", + "channel_type": "channel", + "bot_id": "B_OTHER", + "subtype": "bot_message", + "ts": "171.6", + "text": "automation says hi", + }, + bot_user_id="UBOT", + config={"allowBots": "mentions"}, + known_thread_ids=set(), + ) + allowed = build_inbound_message( + { + "channel": "C123", + "channel_type": "channel", + "bot_id": "B_OTHER", + "subtype": "bot_message", + "ts": "171.7", + "text": "<@UBOT> automation says hi", + }, + bot_user_id="UBOT", + config={"allowBots": "mentions"}, + known_thread_ids=set(), + ) + + assert blocked is None + assert allowed is not None + assert allowed.sender_id == "B_OTHER" + assert allowed.mention_text == "automation says hi" + + +def test_build_inbound_ignores_own_and_edit_messages(): + own = build_inbound_message( + {"channel": "C1", "user": "UBOT", "ts": "1", "text": "echo"}, + bot_user_id="UBOT", + config={}, + known_thread_ids=set(), + ) + edited = build_inbound_message( + {"channel": "C1", "user": "U1", "ts": "2", "text": "x", "subtype": "message_changed"}, + bot_user_id="UBOT", + config={}, + known_thread_ids=set(), + ) + + assert own is None + assert edited is None + + +@pytest.mark.asyncio +async def test_send_text_posts_thread_reply_and_remembers_thread(): + plugin = SlackChannel() + plugin._config = {"replyInThread": True, "replyBroadcast": True} + fake_client = SimpleNamespace( + chat_postMessage=AsyncMock(return_value={"ts": "172.1"}) + ) + plugin._app = SimpleNamespace(client=fake_client) + + result = await plugin.send_text( + OutboundContext( + channel_id="slack", + to="slack:C123", + text="hello", + reply_to_id="171.1", + ) + ) + + assert result.success is True + assert result.message_id == "172.1" + fake_client.chat_postMessage.assert_awaited_once_with( + channel="C123", + text="hello", + mrkdwn=True, + thread_ts="171.1", + reply_broadcast=True, + ) + assert slack_thread_cache_key("default", "C123", "171.1") in plugin._known_thread_ids + assert slack_thread_cache_key("default", "C123", "172.1") not in plugin._known_thread_ids + + +@pytest.mark.asyncio +async def test_known_thread_roots_are_persisted_and_restored(monkeypatch): + stored: dict[str, list[str]] = {} + + class FakeStorage: + @staticmethod + async def get(key): + return stored.get(key) + + @staticmethod + async def set(key, value, value_type="json"): + stored[key] = value + + monkeypatch.setattr(slack_mod, "Storage", FakeStorage) + + plugin = SlackChannel() + plugin._remember_thread("T1", "C123", "171.2") + await plugin._persist_known_threads() + + restarted = SlackChannel() + await restarted._load_known_threads() + msg = build_inbound_message( + { + "channel": "C123", + "channel_type": "channel", + "user": "U123", + "ts": "171.3", + "thread_ts": "171.2", + "parent_user_id": "U123", + "text": "continue after process restart", + "team": "T1", + }, + bot_user_id="UBOT", + config={}, + known_thread_ids=set(restarted._known_thread_ids.keys()), + ) + + assert msg is not None + assert msg.mentioned is True + assert msg.mention_text == "continue after process restart" + + +@pytest.mark.asyncio +async def test_connect_socket_mode_marks_connected_after_connect(monkeypatch): + plugin = SlackChannel() + plugin._config = {"socketConnectTimeoutSeconds": 1} + plugin._app = SimpleNamespace() + plugin._verify_app_token = AsyncMock() + + class FakeHandler: + def __init__(self, app, token): + self.app = app + self.token = token + self.connected = False + + async def connect_async(self): + self.connected = True + + async def close_async(self): + pass + + monkeypatch.setattr(slack_mod, "AsyncSocketModeHandler", FakeHandler) + + await plugin._connect_socket_mode("xapp-ok") + + assert plugin.status.connected is True + assert plugin._handler.connected is True + + +@pytest.mark.asyncio +async def test_connect_socket_mode_failure_marks_disconnected(monkeypatch): + plugin = SlackChannel() + plugin._config = {"socketConnectTimeoutSeconds": 1} + plugin._app = SimpleNamespace() + plugin._verify_app_token = AsyncMock() + + class FakeHandler: + def __init__(self, app, token): + pass + + async def connect_async(self): + raise RuntimeError("invalid app token") + + monkeypatch.setattr(slack_mod, "AsyncSocketModeHandler", FakeHandler) + + with pytest.raises(NonRetryableChannelError): + await plugin._connect_socket_mode("xapp-bad") + + assert plugin.status.connected is False + assert "Slack App Token" in (plugin.status.last_error or "") + assert "xapp-" in (plugin.status.last_error or "") + + +@pytest.mark.asyncio +async def test_connect_socket_mode_slack_response_error_is_actionable(monkeypatch): + plugin = SlackChannel() + plugin._config = {"socketConnectTimeoutSeconds": 1} + plugin._app = SimpleNamespace() + plugin._verify_app_token = AsyncMock() + + class FakeSlackError(Exception): + def __init__(self): + super().__init__("Slack API error") + self.response = SimpleNamespace(data={"error": "missing_scope"}) + + class FakeHandler: + def __init__(self, app, token): + pass + + async def connect_async(self): + raise FakeSlackError() + + monkeypatch.setattr(slack_mod, "AsyncSocketModeHandler", FakeHandler) + + with pytest.raises(NonRetryableChannelError): + await plugin._connect_socket_mode("xapp-missing-scope") + + assert "connections:write" in (plugin.status.last_error or "") + + +@pytest.mark.asyncio +async def test_connect_socket_mode_timeout_marks_disconnected(monkeypatch): + plugin = SlackChannel() + plugin._config = {"socketConnectTimeoutSeconds": 1} + plugin._app = SimpleNamespace() + plugin._verify_app_token = AsyncMock() + + class FakeHandler: + def __init__(self, app, token): + pass + + async def connect_async(self): + await asyncio.sleep(2) + + monkeypatch.setattr(slack_mod, "AsyncSocketModeHandler", FakeHandler) + + with pytest.raises(TimeoutError): + await plugin._connect_socket_mode("xapp-slow") + + assert plugin.status.connected is False + assert "timed out" in (plugin.status.last_error or "") + + +@pytest.mark.asyncio +async def test_connect_socket_mode_preflights_app_token_before_handler(monkeypatch): + plugin = SlackChannel() + plugin._config = {"socketConnectTimeoutSeconds": 1} + plugin._app = SimpleNamespace() + + class FakeSlackError(Exception): + def __init__(self): + super().__init__("Slack API error") + self.response = SimpleNamespace(data={"error": "invalid_auth"}) + + class FakeWebClient: + def __init__(self, token): + self.token = token + + async def apps_connections_open(self, *, app_token): + assert app_token == "xapp-bad" + raise FakeSlackError() + + class FakeHandler: + def __init__(self, app, token): + raise AssertionError("handler should not start for invalid app token") + + monkeypatch.setattr(slack_mod, "AsyncWebClient", FakeWebClient) + monkeypatch.setattr(slack_mod, "AsyncSocketModeHandler", FakeHandler) + + with pytest.raises(NonRetryableChannelError): + await plugin._connect_socket_mode("xapp-bad") + + assert "Slack App Token" in (plugin.status.last_error or "") + assert "xapp-" in (plugin.status.last_error or "") + + +@pytest.mark.asyncio +async def test_start_invalid_config_raises_non_retryable(monkeypatch): + monkeypatch.setattr(slack_mod, "SLACK_AVAILABLE", True) + plugin = SlackChannel() + + with pytest.raises(NonRetryableChannelError): + await plugin.start({}, AsyncMock()) + + assert plugin.status.connected is False + assert "botToken" in (plugin.status.last_error or "") + + +@pytest.mark.asyncio +async def test_start_rejects_user_token_auth_identity(monkeypatch): + monkeypatch.setattr(slack_mod, "SLACK_AVAILABLE", True) + + class FakeClient: + async def auth_test(self): + return { + "user_id": "U_HUMAN", + "team_id": "T1", + } + + class FakeApp: + def __init__(self, token): + self.token = token + self.client = FakeClient() + + monkeypatch.setattr(slack_mod, "AsyncApp", FakeApp) + plugin = SlackChannel() + + with pytest.raises(NonRetryableChannelError): + await plugin.start({"botToken": "xoxb-looks-valid", "appToken": "xapp-1"}, AsyncMock()) + + assert plugin.status.connected is False + assert "Bot User OAuth Token" in (plugin.status.last_error or "") + assert "xoxp-" in (plugin.status.last_error or "") + + +@pytest.mark.asyncio +async def test_handle_slack_event_dispatches_inbound_message(): + plugin = SlackChannel() + plugin._config = {} + plugin._bot_user_id = "UBOT" + dispatched = [] + + async def on_message(msg): + dispatched.append(msg) + + plugin._on_message = on_message + + await plugin._handle_slack_event( + { + "channel": "C123", + "channel_type": "channel", + "user": "U123", + "ts": "171.2", + "text": "<@UBOT> ping", + } + ) + + assert len(dispatched) == 1 + assert dispatched[0].mention_text == "ping" + + +@pytest.mark.asyncio +async def test_handle_slack_event_passes_socket_body_to_inbound_parser(): + plugin = SlackChannel() + plugin._config = {} + plugin._bot_user_id = "UBOT" + dispatched = [] + + async def on_message(msg): + dispatched.append(msg) + + plugin._on_message = on_message + + await plugin._handle_slack_event( + { + "channel": "D123", + "channel_type": "im", + "user": "U123", + "ts": "171.29", + "text": "hello", + }, + body={"authorizations": [{"team_id": "T_BODY"}]}, + ) + + assert len(dispatched) == 1 + assert dispatched[0].account_id == "T_BODY" + + +@pytest.mark.asyncio +async def test_slack_inbound_media_downloads_private_file_with_bot_token(monkeypatch, tmp_path): + monkeypatch.setenv("HOME", str(tmp_path)) + captured = {} + + class FakeResponse: + headers = { + "content-type": "image/png", + "content-disposition": 'attachment; filename="from-slack.png"', + } + + def raise_for_status(self): + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return False + + async def aiter_bytes(self, size): + yield b"png-bytes" + + class FakeClient: + def __init__(self, *args, **kwargs): + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return False + + def stream(self, method, url, headers): + captured["method"] = method + captured["url"] = url + captured["headers"] = headers + return FakeResponse() + + monkeypatch.setattr(slack_media_mod.httpx, "AsyncClient", FakeClient) + + msg = build_inbound_message( + { + "channel": "D123", + "channel_type": "im", + "user": "U123", + "ts": "171.30", + "team": "T1", + "files": [ + { + "id": "F123", + "name": "screenshot.png", + "mimetype": "image/png", + "url_private_download": "https://files.slack.com/files-pri/T1-F123/download/screenshot.png", + } + ], + }, + bot_user_id="UBOT", + config={}, + known_thread_ids=set(), + ) + + assert msg is not None + media = await slack_media_mod.download_inbound_media( + msg, + {"botToken": "xoxb-test"}, + ) + + assert media is not None + assert captured["headers"] == {"Authorization": "Bearer xoxb-test"} + assert captured["url"].startswith("https://files.slack.com/") + assert media.filename == "screenshot.png" + assert media.mime == "image/png" + assert media.url.startswith("file://") + assert media.source["channel"] == "slack" diff --git a/tests/channel/test_weixin_log_levels.py b/tests/channel/test_weixin_log_levels.py new file mode 100644 index 000000000..cc083af66 --- /dev/null +++ b/tests/channel/test_weixin_log_levels.py @@ -0,0 +1,47 @@ +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +import flocks.channel.builtin.weixin.channel as weixin_module + + +@pytest.mark.asyncio +async def test_group_policy_note_logs_info(monkeypatch): + sessions = [MagicMock(), MagicMock()] + info_log = MagicMock() + warn_log = MagicMock() + warning_log = MagicMock() + + monkeypatch.setattr(weixin_module, "AIOHTTP_AVAILABLE", True) + monkeypatch.setattr(weixin_module, "CRYPTO_AVAILABLE", True) + monkeypatch.setattr( + weixin_module, + "aiohttp", + SimpleNamespace( + ClientTimeout=MagicMock(return_value=MagicMock()), + ClientSession=MagicMock(side_effect=sessions), + ), + ) + monkeypatch.setattr(weixin_module.ilink, "make_ssl_connector", MagicMock(return_value=None)) + monkeypatch.setattr(weixin_module, "ContextTokenStore", MagicMock(return_value=MagicMock())) + monkeypatch.setattr(weixin_module, "MessageDedup", MagicMock(return_value=MagicMock())) + monkeypatch.setattr(weixin_module, "MediaCache", MagicMock(return_value=MagicMock())) + monkeypatch.setattr(weixin_module.log, "info", info_log) + monkeypatch.setattr(weixin_module.log, "warn", warn_log) + monkeypatch.setattr(weixin_module.log, "warning", warning_log) + + channel = weixin_module.WeixinChannel() + monkeypatch.setattr(channel, "_poll_loop", AsyncMock(return_value=None)) + monkeypatch.setattr(channel, "_close_sessions", AsyncMock(return_value=None)) + + await channel.start( + {"token": "token", "accountId": "account", "groupPolicy": "all"}, + AsyncMock(), + ) + + assert any(call.args and call.args[0] == "weixin.group_policy.note" for call in info_log.call_args_list) + assert not any( + call.args and call.args[0] == "weixin.group_policy.note" + for call in warn_log.call_args_list + warning_log.call_args_list + ) diff --git a/tests/cli/test_doctor_command.py b/tests/cli/test_doctor_command.py index dea132de8..b9f21513a 100644 --- a/tests/cli/test_doctor_command.py +++ b/tests/cli/test_doctor_command.py @@ -12,17 +12,16 @@ async def _noop_log_init(**_: object) -> None: return None -def test_doctor_runs_source_installer_from_source_root(monkeypatch, tmp_path) -> None: +def test_doctor_runs_core_installer_from_source_root(monkeypatch, tmp_path) -> None: monkeypatch.setenv("FLOCKS_CONFIG_DIR", str(tmp_path)) monkeypatch.setattr(cli_main.Log, "init", _noop_log_init) - calls: list[tuple[list[str], object, bool]] = [] - - def fake_run(command, *, cwd, check, env): - _ = env - calls.append((command, cwd, check)) - - monkeypatch.setattr(doctor_cmd.subprocess, "run", fake_run) + calls: list[object] = [] + monkeypatch.setattr( + doctor_cmd, + "_run_core_install", + lambda source_root, **_kwargs: calls.append(source_root), + ) monkeypatch.setattr( "flocks.cli.service_manager.build_status_lines", lambda: [ @@ -35,16 +34,30 @@ def fake_run(command, *, cwd, check, env): assert result.exit_code == 0, result.stdout assert "Flocks source directory:" in result.stdout - assert "scripts/install.sh" in result.stdout + assert "Repairing Flocks core installation" in result.stdout assert "安装正常" in result.stdout assert "运行状态正常" in result.stdout assert len(calls) == 1 - command, cwd, check = calls[0] - assert command[0] == "bash" - assert command[1].endswith("scripts/install.sh") - assert cwd == doctor_cmd._find_source_root() - assert check is True + assert calls == [doctor_cmd._find_source_root()] + + +def test_doctor_uses_shared_core_install_entry(monkeypatch, tmp_path) -> None: + monkeypatch.setenv("FLOCKS_CONFIG_DIR", str(tmp_path)) + monkeypatch.setattr(cli_main.Log, "init", _noop_log_init) + calls: list[object] = [] + + monkeypatch.setattr( + doctor_cmd, + "_run_core_install", + lambda source_root, **_kwargs: calls.append(source_root), + ) + monkeypatch.setattr(doctor_cmd, "_print_service_diagnosis", lambda: None) + + result = runner.invoke(cli_main.app, ["doctor"]) + + assert result.exit_code == 0, result.stdout + assert calls == [doctor_cmd._find_source_root()] def test_doctor_uses_cn_environment_for_zh_install_profile(monkeypatch, tmp_path) -> None: @@ -54,14 +67,11 @@ def test_doctor_uses_cn_environment_for_zh_install_profile(monkeypatch, tmp_path monkeypatch.setattr(cli_main.Log, "init", _noop_log_init) captured: dict[str, object] = {} - - def fake_run(command, *, cwd, check, env): - captured["command"] = command - captured["cwd"] = cwd - captured["check"] = check - captured["env"] = env - - monkeypatch.setattr(doctor_cmd.subprocess, "run", fake_run) + monkeypatch.setattr( + doctor_cmd, + "_run_core_install", + lambda source_root, *, env: captured.update(source_root=source_root, env=env), + ) monkeypatch.setattr("flocks.cli.service_manager.build_status_lines", lambda: ["[flocks] 后端未运行", "[flocks] WebUI 未运行"]) result = runner.invoke(cli_main.app, ["doctor"]) @@ -99,7 +109,7 @@ def fail_run(*_args, **_kwargs): result = runner.invoke(cli_main.app, ["doctor"]) assert result.exit_code == 0, result.stdout - assert "scripts/install.ps1" in result.stdout + assert "Repairing Flocks core installation" in result.stdout assert "installer will continue in this console" in result.stdout assert captured["cwd"] == doctor_cmd._find_source_root() assert captured["close_fds"] is True @@ -121,16 +131,14 @@ def test_doctor_windows_handoff_child_runs_installer_synchronously(monkeypatch, captured: dict[str, object] = {} - def fake_run(command, *, cwd, check, env): - captured["command"] = command - captured["cwd"] = cwd - captured["check"] = check - captured["env"] = env - def fail_popen(*_args, **_kwargs): raise AssertionError("handoff child should run the installer directly") - monkeypatch.setattr(doctor_cmd.subprocess, "run", fake_run) + monkeypatch.setattr( + doctor_cmd, + "_run_core_install", + lambda source_root, *, env: captured.update(source_root=source_root, env=env), + ) monkeypatch.setattr(doctor_cmd.subprocess, "Popen", fail_popen) monkeypatch.setattr( "flocks.cli.service_manager.build_status_lines", @@ -143,11 +151,7 @@ def fail_popen(*_args, **_kwargs): result = runner.invoke(cli_main.app, ["doctor"]) assert result.exit_code == 0, result.stdout - command = captured["command"] - assert isinstance(command, list) - assert command[-1].endswith("scripts/install.ps1") - assert captured["cwd"] == doctor_cmd._find_source_root() - assert captured["check"] is True + assert captured["source_root"] == doctor_cmd._find_source_root() assert "安装正常" in result.stdout assert "运行状态正常" in result.stdout @@ -175,18 +179,3 @@ def test_service_status_is_healthy_accepts_legacy_backend_and_webui() -> None: ] ) assert not doctor_cmd._service_status_is_healthy(["[flocks] 后端运行中: PID=111", "[flocks] WebUI 未运行"]) - - -def test_doctor_builds_windows_install_command(monkeypatch) -> None: - monkeypatch.setattr(doctor_cmd.shutil, "which", lambda name: None) - - command = doctor_cmd._build_source_install_command(doctor_cmd._find_source_root() / "scripts" / "install.ps1") - - assert command == [ - "powershell", - "-NoProfile", - "-ExecutionPolicy", - "Bypass", - "-File", - str(doctor_cmd._find_source_root() / "scripts" / "install.ps1"), - ] diff --git a/tests/cli/test_service_manager.py b/tests/cli/test_service_manager.py index d75c1cac0..42f3736ad 100644 --- a/tests/cli/test_service_manager.py +++ b/tests/cli/test_service_manager.py @@ -15,7 +15,6 @@ from tests.helpers.service_supervisor import ( SleeperProcessAdapter, make_short_runtime_root, - wait_for_process_exit, ) @@ -30,7 +29,6 @@ def print(self, *args, **kwargs) -> None: @pytest.fixture(autouse=True) def _skip_backend_webui_dist_check(monkeypatch) -> None: monkeypatch.setattr(service_manager, "_ensure_webui_dist", lambda *_args, **_kwargs: None) - monkeypatch.setattr(service_manager, "_resolve_upgrade_runtime", lambda *_args, **_kwargs: {"action": "noop", "error": None}) def _make_runtime_paths(tmp_path: Path) -> service_manager.RuntimePaths: @@ -973,27 +971,22 @@ def test_start_all_starts_supervisor_when_control_api_is_down(monkeypatch) -> No assert call_order == ["ensure_runtime_dirs", "_start_all_without_stop"] -def test_start_all_resolves_upgrade_runtime_before_supervisor_status(monkeypatch) -> None: +def test_start_all_checks_supervisor_before_starting(monkeypatch) -> None: events: list[str] = [] console = DummyConsole() paths = _make_runtime_paths(Path("/tmp/flocks-test")) - def resolve_upgrade_runtime(_console, *, frontend_port: int, attempt_recover: bool) -> dict[str, object]: - events.append(f"upgrade:{frontend_port}:{attempt_recover}") - return {"action": "cleaned", "error": None} - def supervisor_running(_paths) -> bool: events.append("supervisor") return False monkeypatch.setattr(service_manager, "ensure_runtime_dirs", lambda: paths) - monkeypatch.setattr(service_manager, "_resolve_upgrade_runtime", resolve_upgrade_runtime) monkeypatch.setattr(service_manager, "supervisor_is_running", supervisor_running) monkeypatch.setattr(service_manager, "_start_all_without_stop", lambda _config, _console: events.append("start")) service_manager.start_all(service_manager.ServiceConfig(frontend_port=5173), console) - assert events == ["upgrade:5173:False", "supervisor", "start"] + assert events == ["supervisor", "start"] def test_start_all_does_not_duplicate_running_supervisor(monkeypatch) -> None: @@ -1487,6 +1480,10 @@ def test_build_frontend_env_sets_portal_defaults_when_env_missing(monkeypatch) - env["__VITE_ADDITIONAL_SERVER_ALLOWED_HOSTS"] == service_manager.DEFAULT_VITE_ADDITIONAL_SERVER_ALLOWED_HOSTS ) + allowed_hosts = env["__VITE_ADDITIONAL_SERVER_ALLOWED_HOSTS"].split(",") + assert "127.0.0.1" in allowed_hosts + assert "localhost" in allowed_hosts + assert "portalflocks.threatbook.cn" in allowed_hosts def test_build_frontend_env_allows_overriding_portal_defaults(monkeypatch) -> None: @@ -1613,47 +1610,6 @@ def test_supervisor_rejects_static_webui_stop_control_api(monkeypatch, tmp_path: assert exc_info.value.response.status_code == 409 -@pytest.mark.skipif(sys.platform == "win32", reason="uses the Unix domain socket control API") -def test_supervisor_upgrade_prepare_control_api_pauses_real_child_restart(monkeypatch, tmp_path: Path) -> None: - del tmp_path - short_root = make_short_runtime_root("flocks-supervisor-") - paths = _make_runtime_paths(short_root) - paths.run_dir.mkdir(parents=True) - paths.log_dir.mkdir(parents=True) - monkeypatch.setattr(service_manager, "ensure_runtime_dirs", lambda: paths) - backend_adapter = SleeperProcessAdapter() - daemon = service_supervisor.SupervisorDaemon( - service_manager.ServiceConfig(backend_port=9995, frontend_port=9996), - backend_adapter=backend_adapter, - ) - daemon._start_control_server() - - try: - daemon.restart_all(reason="test startup") - backend_process = daemon.backend.process - assert backend_process is not None - assert daemon.webui.process is None - assert daemon.webui.state == "static" - - status = service_control.request_prepare_upgrade(paths=paths) - - wait_for_process_exit(backend_process) - assert status.backend.paused is True - assert status.webui.paused is True - assert daemon.backend.process is None - assert backend_process.pid in backend_adapter.stopped - - daemon.tick() - - assert len(backend_adapter.started) == 1 - assert daemon.backend.process is None - assert daemon.status_payload()["backend"]["paused"] is True - finally: - daemon.shutdown_children() - daemon._stop_control_server() - shutil.rmtree(short_root, ignore_errors=True) - - def test_build_webui_dist_tolerates_windows_node_assertion_after_build(monkeypatch, tmp_path: Path) -> None: webui_dir = tmp_path / "webui" webui_dist = webui_dir / "dist" diff --git a/tests/cli/test_session_runner_project_scope.py b/tests/cli/test_session_runner_project_scope.py new file mode 100644 index 000000000..01b0a1a6b --- /dev/null +++ b/tests/cli/test_session_runner_project_scope.py @@ -0,0 +1,48 @@ +from io import StringIO +from pathlib import Path + +import pytest +from rich.console import Console + +from flocks.cli.session_runner import CLISessionRunner +from flocks.session.session import Session, SessionInfo, SessionTime + + +def _session(session_id: str, directory: Path, *, updated: int) -> SessionInfo: + return SessionInfo( + id=session_id, + projectID="default", + directory=str(directory), + title=session_id, + time=SessionTime(created=updated, updated=updated), + ) + + +@pytest.mark.asyncio +async def test_continue_uses_latest_session_from_current_worktree( + monkeypatch, + tmp_path: Path, +) -> None: + project_a = tmp_path / "project-a" + project_b = tmp_path / "project-b" + project_a.mkdir() + project_b.mkdir() + session_a = _session("ses-a", project_a, updated=2_000) + session_b = _session("ses-b", project_b, updated=1_000) + + async def fake_list(project_id: str): + assert project_id == "default" + return [session_a, session_b] + + monkeypatch.setattr(Session, "list", fake_list) + runner = CLISessionRunner( + console=Console(file=StringIO()), + directory=project_b, + ) + + resumed = await runner._get_or_create_session( + project_id="default", + continue_session=True, + ) + + assert resumed.id == session_b.id diff --git a/tests/cli/test_stats_command.py b/tests/cli/test_stats_command.py index 14b9ef561..3f88cf183 100644 --- a/tests/cli/test_stats_command.py +++ b/tests/cli/test_stats_command.py @@ -20,6 +20,33 @@ def _build_session(session_id: str = "ses_stats_test") -> SessionInfo: ) +@pytest.mark.asyncio +async def test_current_project_stats_only_include_sessions_from_current_worktree( + monkeypatch, + tmp_path, +) -> None: + current_project = tmp_path / "current" + other_project = tmp_path / "other" + current_project.mkdir() + other_project.mkdir() + current = _build_session("ses-current") + current.directory = str(current_project) + current.project_id = "default" + other = _build_session("ses-other") + other.directory = str(other_project) + other.project_id = "default" + + async def fake_get_all_sessions(): + return [other, current] + + monkeypatch.chdir(current_project) + monkeypatch.setattr(stats_cmd, "_get_all_sessions", fake_get_all_sessions) + + sessions = await stats_cmd._resolve_project_sessions("") + + assert [session.id for session in sessions] == [current.id] + + @pytest.mark.asyncio async def test_aggregate_stats_uses_usage_records(monkeypatch) -> None: session = _build_session() diff --git a/tests/cli/test_update_command.py b/tests/cli/test_update_command.py index 76b5e971e..5dd268b15 100644 --- a/tests/cli/test_update_command.py +++ b/tests/cli/test_update_command.py @@ -7,7 +7,6 @@ import flocks.cli.commands.update as update_cmd import flocks.cli.main as cli_main -import flocks.cli.service_manager as service_manager import flocks.updater as updater_pkg from flocks.updater.models import UpdateProgress, VersionInfo @@ -18,10 +17,10 @@ async def _noop_log_init(**_: object) -> None: return None -def test_updater_package_exports_build_updated_frontend() -> None: +def test_updater_package_exports_shared_installer() -> None: from flocks.updater import updater as updater_module - assert updater_pkg.build_updated_frontend is updater_module.build_updated_frontend + assert updater_pkg.install_or_repair_source is updater_module.install_or_repair_source def test_update_cli_accepts_force_option(monkeypatch, tmp_path) -> None: @@ -92,8 +91,6 @@ def test_update_prompts_for_cn_mirror_before_upgrade_confirmation(monkeypatch, t check_regions: list[str | None] = [] confirm_prompts: list[str] = [] captured: dict[str, object] = {} - stop_calls: list[str] = [] - build_calls: list[str | None] = [] answers = iter([True, True]) async def fake_check_update(*, locale: str | None = None, region: str | None = None) -> VersionInfo: @@ -123,6 +120,7 @@ async def fake_perform_update( restart: bool = True, locale: str | None = None, region: str | None = None, + wait_for_handoff: bool = False, ): captured["latest_tag"] = latest_tag captured["zipball_url"] = zipball_url @@ -131,6 +129,7 @@ async def fake_perform_update( captured["bundle_format"] = bundle_format captured["perform_region"] = region captured["restart"] = restart + captured["wait_for_handoff"] = wait_for_handoff async for step in _fake_progress(): yield step @@ -138,18 +137,10 @@ def fake_confirm(prompt: str, default: bool = False) -> bool: confirm_prompts.append(prompt) return next(answers) - def fake_stop_all(console) -> None: - stop_calls.append("stop") - - async def fake_build_updated_frontend(*, locale: str | None = None, region: str | None = None) -> None: - build_calls.append(region) - monkeypatch.setattr(updater_pkg, "check_update", fake_check_update) monkeypatch.setattr(updater_pkg, "perform_update", fake_perform_update) - monkeypatch.setattr(updater_pkg, "build_updated_frontend", fake_build_updated_frontend) monkeypatch.setattr(updater_pkg, "detect_deploy_mode", lambda: "source") monkeypatch.setattr(update_cmd.typer, "confirm", fake_confirm) - monkeypatch.setattr(service_manager, "stop_all", fake_stop_all) import asyncio @@ -157,8 +148,6 @@ async def fake_build_updated_frontend(*, locale: str | None = None, region: str assert check_regions == ["cn"] assert confirm_prompts == ["\n是否使用中国镜像进行升级？", "\n是否立即升级？"] - assert stop_calls == ["stop"] - assert build_calls == ["cn"] assert captured == { "latest_tag": "2026.4.2", "zipball_url": "https://gitee.example.com/flocks.zip", @@ -166,13 +155,17 @@ async def fake_build_updated_frontend(*, locale: str | None = None, region: str "bundle_sha256": None, "bundle_format": None, "perform_region": "cn", - "restart": False, + "restart": True, + "wait_for_handoff": True, } assert "已切换为中国镜像源" not in output.getvalue() async def _fake_progress(): yield UpdateProgress(stage="fetching", message="fetching") + yield UpdateProgress(stage="backing_up", message="backing up") + yield UpdateProgress(stage="applying", message="applying") + yield UpdateProgress(stage="restarting", message="restarting") yield UpdateProgress(stage="done", message="done", success=True) @@ -197,8 +190,6 @@ async def fake_check_update(*, locale: str | None = None, region: str | None = N ) captured: dict[str, object] = {} - stop_calls: list[str] = [] - build_calls: list[str | None] = [] async def fake_perform_update( latest_tag: str, @@ -210,6 +201,7 @@ async def fake_perform_update( restart: bool = True, locale: str | None = None, region: str | None = None, + wait_for_handoff: bool = False, ): captured["latest_tag"] = latest_tag captured["zipball_url"] = zipball_url @@ -218,20 +210,13 @@ async def fake_perform_update( captured["bundle_format"] = bundle_format captured["perform_region"] = region captured["restart"] = restart + captured["wait_for_handoff"] = wait_for_handoff async for step in _fake_progress(): yield step - def fake_stop_all(console) -> None: - stop_calls.append("stop") - - async def fake_build_updated_frontend(*, locale: str | None = None, region: str | None = None) -> None: - build_calls.append(region) - monkeypatch.setattr(updater_pkg, "check_update", fake_check_update) monkeypatch.setattr(updater_pkg, "perform_update", fake_perform_update) - monkeypatch.setattr(updater_pkg, "build_updated_frontend", fake_build_updated_frontend) monkeypatch.setattr(updater_pkg, "detect_deploy_mode", lambda: "source") - monkeypatch.setattr(service_manager, "stop_all", fake_stop_all) import asyncio @@ -245,15 +230,16 @@ async def fake_build_updated_frontend(*, locale: str | None = None, region: str "bundle_format": None, "check_region": "cn", "perform_region": "cn", - "restart": False, + "restart": True, + "wait_for_handoff": True, } - assert stop_calls == ["stop"] - assert build_calls == ["cn"] assert "强制重新安装 v2026.4.2" in output.getvalue() + assert "[3/3] 应用新版本... ✓" in output.getvalue() + assert "重启服务" not in output.getvalue() assert "升级完成" in output.getvalue() -def test_update_executes_flocks_stop_before_upgrade(monkeypatch, tmp_path) -> None: +def test_update_delegates_stop_install_build_and_restart_to_handoff(monkeypatch, tmp_path) -> None: output = StringIO() monkeypatch.setenv("FLOCKS_CONFIG_DIR", str(tmp_path)) monkeypatch.delenv("FLOCKS_INSTALL_LANGUAGE", raising=False) @@ -288,8 +274,9 @@ async def fake_perform_update( restart: bool = True, locale: str | None = None, region: str | None = None, + wait_for_handoff: bool = False, ): - events.append("perform_update") + events.append(f"perform_update:{wait_for_handoff}") async for step in _fake_progress(): yield step @@ -297,29 +284,21 @@ def fake_confirm(prompt: str, default: bool = False) -> bool: confirm_prompts.append(prompt) return next(answers) - def fake_stop_all(console) -> None: - events.append("stop") - - async def fake_build_updated_frontend(*, locale: str | None = None, region: str | None = None) -> None: - events.append("build") - monkeypatch.setattr(updater_pkg, "check_update", fake_check_update) monkeypatch.setattr(updater_pkg, "perform_update", fake_perform_update) - monkeypatch.setattr(updater_pkg, "build_updated_frontend", fake_build_updated_frontend) monkeypatch.setattr(updater_pkg, "detect_deploy_mode", lambda: "source") monkeypatch.setattr(update_cmd.typer, "confirm", fake_confirm) - monkeypatch.setattr(service_manager, "stop_all", fake_stop_all) import asyncio asyncio.run(update_cmd._update(check=False, yes=False, force=False, region=None)) assert confirm_prompts == ["\n是否使用中国镜像进行升级？", "\n是否立即升级？"] - assert events == ["stop", "perform_update", "build"] - assert "已执行 flocks stop" in output.getvalue() + assert events == ["perform_update:True"] + assert "已执行 flocks stop" not in output.getvalue() -def test_update_reports_frontend_build_failure_after_common_upgrade(monkeypatch, tmp_path) -> None: +def test_update_reports_handoff_preparation_failure(monkeypatch, tmp_path) -> None: output = StringIO() monkeypatch.setenv("FLOCKS_CONFIG_DIR", str(tmp_path)) monkeypatch.delenv("FLOCKS_INSTALL_LANGUAGE", raising=False) @@ -350,21 +329,13 @@ async def fake_perform_update( restart: bool = True, locale: str | None = None, region: str | None = None, + wait_for_handoff: bool = False, ): - async for step in _fake_progress(): - yield step - - def fake_stop_all(console) -> None: - return None - - async def fake_build_updated_frontend(*, locale: str | None = None, region: str | None = None) -> None: - raise RuntimeError("npm run build failed") + yield UpdateProgress(stage="error", message="handoff preparation failed", success=False) monkeypatch.setattr(updater_pkg, "check_update", fake_check_update) monkeypatch.setattr(updater_pkg, "perform_update", fake_perform_update) - monkeypatch.setattr(updater_pkg, "build_updated_frontend", fake_build_updated_frontend) monkeypatch.setattr(updater_pkg, "detect_deploy_mode", lambda: "source") - monkeypatch.setattr(service_manager, "stop_all", fake_stop_all) import asyncio @@ -372,4 +343,4 @@ async def fake_build_updated_frontend(*, locale: str | None = None, region: str asyncio.run(update_cmd._update(check=False, yes=True, force=False, region=None)) assert excinfo.value.exit_code == 1 - assert "前端构建失败" in output.getvalue() + assert "handoff preparation failed" in output.getvalue() diff --git a/tests/config/test_config_writer.py b/tests/config/test_config_writer.py index d66640115..1a136f4cf 100644 --- a/tests/config/test_config_writer.py +++ b/tests/config/test_config_writer.py @@ -187,13 +187,13 @@ def test_build_provider_config(self, temp_project): "deepseek", npm="@ai-sdk/openai-compatible", base_url="https://api.deepseek.com/v1", - models={"deepseek-chat": {"name": "DeepSeek Chat"}}, + models={"deepseek-v4-flash": {"name": "DeepSeek V4 Flash"}}, ) assert config["npm"] == "@ai-sdk/openai-compatible" assert config["options"]["apiKey"] == "{secret:deepseek_llm_key}" assert config["options"]["baseURL"] == "https://api.deepseek.com/v1" - assert config["models"]["deepseek-chat"]["name"] == "DeepSeek Chat" + assert config["models"]["deepseek-v4-flash"]["name"] == "DeepSeek V4 Flash" def test_build_provider_config_no_url(self, temp_project): from flocks.config.config_writer import ConfigWriter diff --git a/tests/docker/test_container_start.py b/tests/docker/test_container_start.py new file mode 100644 index 000000000..900067300 --- /dev/null +++ b/tests/docker/test_container_start.py @@ -0,0 +1,42 @@ +import os +import subprocess +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[2] +CONTAINER_START = REPO_ROOT / "scripts" / "container-start.sh" + + +def test_container_start_uses_supervised_service_commands(tmp_path: Path) -> None: + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + call_log = tmp_path / "calls.log" + flocks = bin_dir / "flocks" + flocks.write_text( + '#!/usr/bin/env bash\nprintf \'%s\\n\' "$*" >> "$CALL_LOG"\n', + encoding="utf-8", + ) + flocks.chmod(0o755) + env = os.environ.copy() + env.update( + { + "CALL_LOG": str(call_log), + "PATH": f"{bin_dir}:{env['PATH']}", + "ROOT_DIR": str(tmp_path), + } + ) + + completed = subprocess.run( + ["bash", str(CONTAINER_START)], + check=False, + capture_output=True, + env=env, + text=True, + ) + + assert completed.returncode == 0 + assert call_log.read_text(encoding="utf-8").splitlines() == [ + "start --no-browser --skip-webui-build", + "logs --follow", + "stop", + ] diff --git a/tests/docker/test_dockerfile_runtime_requirements.py b/tests/docker/test_dockerfile_runtime_requirements.py index bbb61e1c8..e2340771a 100644 --- a/tests/docker/test_dockerfile_runtime_requirements.py +++ b/tests/docker/test_dockerfile_runtime_requirements.py @@ -1,8 +1,42 @@ +import re from pathlib import Path +import pytest + +from flocks.cli.service_config import build_service_config + REPO_ROOT = Path(__file__).resolve().parents[2] DOCKERFILE = REPO_ROOT / "docker" / "Dockerfile" +PORT_ENV_NAMES = ( + "FLOCKS_PORT", + "FLOCKS_PUBLIC_PORT", + "FLOCKS_WEBUI_PORT", + "FLOCKS_FRONTEND_PORT", + "FLOCKS_SERVER_PORT", + "FLOCKS_BACKEND_PORT", +) + + +def _docker_port_env() -> dict[str, str]: + dockerfile = DOCKERFILE.read_text(encoding="utf-8") + values = {} + for name in PORT_ENV_NAMES: + match = re.search(rf"\b{re.escape(name)}=([^\s\\]+)", dockerfile) + if match: + values[name] = match.group(1) + return values + + +def _resolve_docker_port(monkeypatch: pytest.MonkeyPatch, runtime_env: dict[str, str]) -> int: + for name in PORT_ENV_NAMES: + monkeypatch.delenv(name, raising=False) + for name, value in _docker_port_env().items(): + monkeypatch.setenv(name, value) + for name, value in runtime_env.items(): + monkeypatch.setenv(name, value) + config = build_service_config(default_server_host="127.0.0.1", default_server_port=8000) + return config.backend_port def test_runtime_image_no_longer_bundles_system_chromium() -> None: @@ -10,3 +44,32 @@ def test_runtime_image_no_longer_bundles_system_chromium() -> None: assert "AGENT_BROWSER_EXECUTABLE_PATH=/usr/bin/chromium" not in dockerfile assert " chromium \\" not in dockerfile + + +def test_runtime_image_uses_unified_supervised_service_port() -> None: + dockerfile = DOCKERFILE.read_text(encoding="utf-8") + + assert "FLOCKS_BACKEND_HOST=0.0.0.0" in dockerfile + assert "FLOCKS_BACKEND_PORT=5173" in dockerfile + assert "FLOCKS_HOST=0.0.0.0" not in dockerfile + assert "FLOCKS_PORT=5173" not in dockerfile + assert "EXPOSE 5173" in dockerfile + assert "EXPOSE 8000" not in dockerfile + assert 'ENTRYPOINT ["/usr/bin/tini", "-g", "--"]' in dockerfile + + +@pytest.mark.parametrize( + ("runtime_env", "expected_port"), + [ + ({}, 5173), + ({"FLOCKS_BACKEND_PORT": "6184"}, 6184), + ({"FLOCKS_FRONTEND_PORT": "6184"}, 6184), + ({"FLOCKS_PORT": "6184"}, 6184), + ], +) +def test_runtime_port_overrides_remain_backward_compatible( + monkeypatch: pytest.MonkeyPatch, + runtime_env: dict[str, str], + expected_port: int, +) -> None: + assert _resolve_docker_port(monkeypatch, runtime_env) == expected_port diff --git a/tests/hub/test_bundled_tools.py b/tests/hub/test_bundled_tools.py index e2fe2c483..93cec8f5d 100644 --- a/tests/hub/test_bundled_tools.py +++ b/tests/hub/test_bundled_tools.py @@ -134,7 +134,10 @@ def _write_bundled_tool( @pytest.mark.asyncio -async def test_tool_runtime_refresh_clears_device_template_cache(monkeypatch): +async def test_tool_runtime_refresh_scopes_errors_and_clears_device_template_cache( + monkeypatch, + tmp_path, +): from flocks.config import api_versioning from flocks.tool.device import plugin_index from flocks.tool.registry import ToolRegistry @@ -144,7 +147,9 @@ async def test_tool_runtime_refresh_clears_device_template_cache(monkeypatch): monkeypatch.setattr( ToolRegistry, "refresh_plugin_tools", - classmethod(lambda cls: calls.append("refresh")), + classmethod( + lambda cls, changed_path=None: calls.append(f"refresh:{changed_path}") + ), ) monkeypatch.setattr( api_versioning, @@ -153,15 +158,17 @@ async def test_tool_runtime_refresh_clears_device_template_cache(monkeypatch): ) plugin_index._template_cache = [] - await _refresh_runtime("device") + changed_path = tmp_path / "device" + await _refresh_runtime("device", changed_path) assert plugin_index._template_cache is None - assert calls == ["init", "refresh", "discover:True"] + assert calls == ["init", f"refresh:{changed_path}", "discover:True"] @pytest.mark.asyncio async def test_uninstall_missing_tool_record_clears_device_template_cache(isolated_hub, monkeypatch): calls: list[str] = [] + missing_path = isolated_hub["home"] / ".flocks" / "plugins" / "tools" / "api" / "ghost_tool" local.save_installed_record( InstalledPluginRecord( id="ghost_tool", @@ -169,7 +176,7 @@ async def test_uninstall_missing_tool_record_clears_device_template_cache(isolat version="1.0", source="bundled", installedAt=1, - installPath=str(isolated_hub["home"] / ".flocks" / "plugins" / "tools" / "api" / "ghost_tool"), + installPath=str(missing_path), ) ) monkeypatch.setattr("flocks.hub.installer.clear_catalog_caches", lambda: calls.append("catalog")) @@ -178,8 +185,8 @@ async def test_uninstall_missing_tool_record_clears_device_template_cache(isolat lambda plugin_type: calls.append(f"device:{plugin_type}"), ) - async def fake_refresh_runtime(plugin_type): - calls.append(f"refresh:{plugin_type}") + async def fake_refresh_runtime(plugin_type, changed_path=None): + calls.append(f"refresh:{plugin_type}:{changed_path}") monkeypatch.setattr("flocks.hub.installer._refresh_runtime", fake_refresh_runtime) @@ -187,7 +194,37 @@ async def fake_refresh_runtime(plugin_type): assert removed is True assert local.get_record("tool", "ghost_tool") is None - assert calls == ["catalog", "device:tool", "refresh:tool"] + assert calls == ["catalog", "device:tool", f"refresh:tool:{missing_path}"] + + +@pytest.mark.asyncio +async def test_uninstall_missing_tool_record_without_path_uses_canonical_path( + isolated_hub, + monkeypatch, +): + local.save_installed_record( + InstalledPluginRecord( + id="ghost_tool", + type="tool", + version="1.0", + source="bundled", + installedAt=1, + installPath=None, + ) + ) + refresh_paths: list[Path | None] = [] + + async def fake_refresh_runtime(_plugin_type, changed_path=None): + refresh_paths.append(changed_path) + + monkeypatch.setattr("flocks.hub.installer._refresh_runtime", fake_refresh_runtime) + + removed = await uninstall_plugin("tool", "ghost_tool") + + assert removed is True + assert refresh_paths == [ + isolated_hub["home"] / ".flocks" / "plugins" / "tools" / "ghost_tool" + ] class TestBundledToolRoots: @@ -777,6 +814,23 @@ def test_path_escape_still_rejected(self, tmp_path): with pytest.raises(ValueError, match="escapes"): validate_package(package, self._manifest()) + def test_python_syntax_error_is_rejected_before_install(self, tmp_path): + package = tmp_path / "demo" + package.mkdir() + (package / "broken.handler.py").write_text("def broken(:\n", encoding="utf-8") + + with pytest.raises(ValueError, match=r"Invalid Python source.*broken\.handler\.py"): + validate_package(package, self._manifest()) + + def test_nested_python_syntax_error_is_rejected_before_install(self, tmp_path): + package = tmp_path / "demo" + scripts = package / "scripts" + scripts.mkdir(parents=True) + (scripts / "broken.py").write_text("def broken(:\n", encoding="utf-8") + + with pytest.raises(ValueError, match=r"Invalid Python source.*scripts/broken\.py"): + validate_package(package, self._manifest()) + # --------------------------------------------------------------------------- # API-service summary name disambiguation diff --git a/tests/hub/test_hub_catalog.py b/tests/hub/test_hub_catalog.py index f26ec863e..bdb479b2e 100644 --- a/tests/hub/test_hub_catalog.py +++ b/tests/hub/test_hub_catalog.py @@ -9,7 +9,7 @@ from flocks.hub import local from flocks.hub.catalog import list_catalog, load_manifest, load_taxonomy from flocks.hub.files import file_tree, read_file_content -from flocks.hub.installer import install_plugin, uninstall_plugin +from flocks.hub.installer import install_plugin, uninstall_plugin, update_plugin from flocks.plugin.loader import PluginLoader @@ -114,6 +114,162 @@ def test_soc_workspace_component_exposes_chinese_name(): assert entries["soc-workspace"].nameCn == "SOC 工作区场景套件" +async def test_chaitin_safeline_repair_release_replaces_broken_handler( + isolated_hub_env, + monkeypatch: pytest.MonkeyPatch, +): + async def noop_refresh(_plugin_type, _changed_path=None): + return None + + monkeypatch.setattr("flocks.hub.installer._refresh_runtime", noop_refresh) + plugin_id = "chaitin_safeline_waf_v1_0_0" + install_dir = ( + isolated_hub_env["home"] + / ".flocks" + / "plugins" + / "tools" + / "device" + / plugin_id + ) + install_dir.mkdir(parents=True) + legacy_handler = install_dir / "chaitin_leichi_waf.handler.py" + legacy_handler.write_text("def broken(:\n", encoding="utf-8") + + entries = {entry.id: entry for entry in list_catalog(plugin_type="device")} + assert entries[plugin_id].state == "updateAvailable" + assert entries[plugin_id].version == "1.0.1" + assert entries[plugin_id].installedVersion is None + + record = await update_plugin("device", plugin_id) + + repaired_handler = install_dir / "chaitin_safeline_waf.handler.py" + assert record.version == "1.0.1" + assert local.get_record("device", plugin_id) == record + assert not legacy_handler.exists() + compile(repaired_handler.read_text(encoding="utf-8"), str(repaired_handler), "exec") + + +async def test_hub_install_runtime_failure_rolls_back_new_plugin( + isolated_hub_env, + monkeypatch: pytest.MonkeyPatch, +): + async def fail_refresh(_plugin_type, _changed_path=None): + raise RuntimeError("import failed") + + monkeypatch.setattr("flocks.hub.installer._refresh_runtime", fail_refresh) + install_dir = ( + isolated_hub_env["home"] + / ".flocks" + / "plugins" + / "tools" + / "python" + / "soc_workspace_query" + ) + + with pytest.raises(RuntimeError, match="import failed"): + await install_plugin("tool", "soc_workspace_query") + + assert not install_dir.exists() + assert local.get_record("tool", "soc_workspace_query") is None + + +async def test_hub_update_runtime_failure_restores_previous_plugin( + isolated_hub_env, + monkeypatch: pytest.MonkeyPatch, +): + install_dir = ( + isolated_hub_env["home"] + / ".flocks" + / "plugins" + / "tools" + / "python" + / "soc_workspace_query" + ) + install_dir.mkdir(parents=True) + old_handler = install_dir / "old_tool.py" + old_handler.write_text("OLD = True\n", encoding="utf-8") + previous_record = local.make_record( + plugin_type="tool", + plugin_id="soc_workspace_query", + version="0.9.0", + source="bundled:old", + install_path=install_dir, + scope="global", + ) + local.save_installed_record(previous_record) + refresh_calls = 0 + + async def fail_then_recover(_plugin_type, _changed_path=None): + nonlocal refresh_calls + refresh_calls += 1 + if refresh_calls == 1: + raise RuntimeError("import failed") + + monkeypatch.setattr("flocks.hub.installer._refresh_runtime", fail_then_recover) + + with pytest.raises(RuntimeError, match="import failed"): + await update_plugin("tool", "soc_workspace_query") + + assert refresh_calls == 2 + assert old_handler.read_text(encoding="utf-8") == "OLD = True\n" + assert not (install_dir / "soc_workspace_query.py").exists() + assert local.get_record("tool", "soc_workspace_query") == previous_record + + +async def test_hub_webui_runtime_failure_rolls_back_package_and_access_contracts( + isolated_hub_env, + monkeypatch: pytest.MonkeyPatch, +): + async def fail_refresh(_plugin_type, _changed_path=None): + raise RuntimeError("reconcile failed") + + monkeypatch.setattr("flocks.hub.installer._refresh_runtime", fail_refresh) + _patch_webui_bundle_build(monkeypatch) + home_plugins = isolated_hub_env["home"] / ".flocks" / "plugins" + + with pytest.raises(RuntimeError, match="reconcile failed"): + await install_plugin("webui", "soc_ui") + + assert not (home_plugins / "contracts" / "webui" / "soc_ui").exists() + assert not (home_plugins / "contracts" / "access" / "soc_ui").exists() + assert local.get_record("webui", "soc_ui") is None + + +def test_catalog_ignores_stale_record_when_legacy_install_is_inferred( + isolated_hub_env, +): + plugin_id = "chaitin_safeline_waf_v1_0_0" + legacy_dir = ( + isolated_hub_env["home"] + / ".flocks" + / "plugins" + / "tools" + / "device" + / plugin_id + ) + legacy_dir.mkdir(parents=True) + (legacy_dir / "legacy.yaml").write_text( + "name: legacy\nhandler:\n type: http\n", + encoding="utf-8", + ) + stale_record = local.make_record( + plugin_type="device", + plugin_id=plugin_id, + version="1.0.1", + source="bundled:stale", + install_path=isolated_hub_env["home"] / "missing", + scope="global", + ) + local.save_installed_record(stale_record) + + entries = {entry.id: entry for entry in list_catalog(plugin_type="device")} + + assert entries[plugin_id].state == "updateAvailable" + assert entries[plugin_id].installedVersion is None + assert entries[plugin_id].installPath == str(legacy_dir) + assert local.get_record("device", plugin_id) is None + + def test_pentest_agents_are_listed_in_agent_catalog(): entries = list_catalog(plugin_type="agent") ids = {entry.id for entry in entries} @@ -214,7 +370,7 @@ async def test_hub_installs_pentest_subagent(isolated_hub_env): async def test_hub_installs_soc_webui_package(isolated_hub_env, monkeypatch: pytest.MonkeyPatch): - async def noop_refresh(_plugin_type): + async def noop_refresh(_plugin_type, _changed_path=None): return None monkeypatch.setattr("flocks.hub.installer._refresh_runtime", noop_refresh) @@ -246,7 +402,7 @@ async def test_hub_webui_install_fails_when_bundle_build_fails( isolated_hub_env, monkeypatch: pytest.MonkeyPatch, ): - async def noop_refresh(_plugin_type): + async def noop_refresh(_plugin_type, _changed_path=None): return None def fail_build(_self, _page_id: str): @@ -269,7 +425,7 @@ async def test_hub_installed_soc_webui_registers_alert_access_contract( isolated_hub_env, monkeypatch: pytest.MonkeyPatch, ): - async def noop_refresh(_plugin_type): + async def noop_refresh(_plugin_type, _changed_path=None): return None monkeypatch.setattr("flocks.hub.installer._refresh_runtime", noop_refresh) @@ -296,7 +452,7 @@ async def test_hub_installed_soc_webui_serves_alert_access_operation( isolated_hub_env, monkeypatch: pytest.MonkeyPatch, ): - async def noop_refresh(_plugin_type): + async def noop_refresh(_plugin_type, _changed_path=None): return None monkeypatch.setattr("flocks.hub.installer._refresh_runtime", noop_refresh) @@ -392,7 +548,7 @@ async def noop_refresh(_plugin_type): async def test_hub_installs_soc_workspace_component_children(isolated_hub_env, monkeypatch: pytest.MonkeyPatch): - async def noop_refresh(_plugin_type): + async def noop_refresh(_plugin_type, _changed_path=None): return None monkeypatch.setattr("flocks.hub.installer._refresh_runtime", noop_refresh) @@ -421,7 +577,7 @@ async def noop_refresh(_plugin_type): async def test_hub_component_uninstall_preserves_existing_children(isolated_hub_env, monkeypatch: pytest.MonkeyPatch): - async def noop_refresh(_plugin_type): + async def noop_refresh(_plugin_type, _changed_path=None): return None monkeypatch.setattr("flocks.hub.installer._refresh_runtime", noop_refresh) @@ -448,7 +604,7 @@ async def noop_refresh(_plugin_type): async def test_hub_component_adopts_existing_soc_workspace_tool(isolated_hub_env, monkeypatch: pytest.MonkeyPatch): - async def noop_refresh(_plugin_type): + async def noop_refresh(_plugin_type, _changed_path=None): return None monkeypatch.setattr("flocks.hub.installer._refresh_runtime", noop_refresh) @@ -473,7 +629,7 @@ async def noop_refresh(_plugin_type): async def test_hub_component_uninstall_cleans_adoptable_legacy_tool_record(isolated_hub_env, monkeypatch: pytest.MonkeyPatch): - async def noop_refresh(_plugin_type): + async def noop_refresh(_plugin_type, _changed_path=None): return None monkeypatch.setattr("flocks.hub.installer._refresh_runtime", noop_refresh) @@ -496,7 +652,7 @@ async def noop_refresh(_plugin_type): async def test_hub_component_uninstall_cleans_legacy_unrecorded_webui(isolated_hub_env, monkeypatch: pytest.MonkeyPatch): - async def noop_refresh(_plugin_type): + async def noop_refresh(_plugin_type, _changed_path=None): return None monkeypatch.setattr("flocks.hub.installer._refresh_runtime", noop_refresh) @@ -521,7 +677,7 @@ async def noop_refresh(_plugin_type): async def test_hub_component_install_failure_rolls_back_children(isolated_hub_env, monkeypatch: pytest.MonkeyPatch): - async def fail_tool_refresh(plugin_type): + async def fail_tool_refresh(plugin_type, _changed_path=None): if plugin_type == "tool": raise RuntimeError("tool refresh failed") return None @@ -548,7 +704,7 @@ async def fail_tool_refresh(plugin_type): async def test_hub_component_uninstall_cleans_orphan_children(isolated_hub_env, monkeypatch: pytest.MonkeyPatch): - async def noop_refresh(_plugin_type): + async def noop_refresh(_plugin_type, _changed_path=None): return None monkeypatch.setattr("flocks.hub.installer._refresh_runtime", noop_refresh) @@ -582,7 +738,7 @@ async def noop_refresh(_plugin_type): async def test_hub_uninstalls_python_tool_without_record(isolated_hub_env, monkeypatch: pytest.MonkeyPatch): - async def noop_refresh(_plugin_type): + async def noop_refresh(_plugin_type, _changed_path=None): return None monkeypatch.setattr("flocks.hub.installer._refresh_runtime", noop_refresh) @@ -758,7 +914,7 @@ def test_hub_component_install_stream_reports_child_progress(isolated_hub_env, m from flocks.server.auth import require_admin from flocks.server.routes.hub import router - async def noop_refresh(_plugin_type): + async def noop_refresh(_plugin_type, _changed_path=None): return None monkeypatch.setattr("flocks.hub.installer._refresh_runtime", noop_refresh) diff --git a/tests/ingest/test_kafka_manager.py b/tests/ingest/test_kafka_manager.py index 31de29d02..111cdf34e 100644 --- a/tests/ingest/test_kafka_manager.py +++ b/tests/ingest/test_kafka_manager.py @@ -17,6 +17,7 @@ import json import sys from types import SimpleNamespace +from unittest.mock import AsyncMock import pytest @@ -25,6 +26,105 @@ from flocks.workflow.triggers.models import TriggerDefinition +@pytest.fixture(autouse=True) +def trigger_tool_context(monkeypatch: pytest.MonkeyPatch) -> SimpleNamespace: + context = object() + builder = AsyncMock(return_value=context) + cleanup = AsyncMock() + monkeypatch.setattr(kafka_manager, "build_workflow_tool_context", builder) + monkeypatch.setattr(kafka_manager, "cleanup_workflow_tool_context", cleanup) + return SimpleNamespace(context=context, builder=builder, cleanup=cleanup) + + +@pytest.mark.asyncio +async def test_start_all_skips_stale_workflow_config_as_info( + monkeypatch: pytest.MonkeyPatch, +) -> None: + workflow_id = "removed-workflow" + config = { + "enabled": True, + "inputBroker": "", + "inputTopic": "", + } + info_events: list[tuple[str, dict]] = [] + warning_events: list[str] = [] + + async def _list_configs(*, kind: str): # noqa: ANN202 + assert kind == "workflow_kafka_config" + return [(workflow_id, config)] + + async def _get_config(_workflow_id: str, *, kind: str) -> dict: + assert kind == "workflow_kafka_config" + return config + + monkeypatch.setattr(kafka_manager.WorkflowStore, "list_configs", _list_configs) + monkeypatch.setattr(kafka_manager.WorkflowStore, "get_config", _get_config) + monkeypatch.setattr(kafka_manager, "read_workflow_from_fs", lambda _workflow_id: None) + monkeypatch.setattr( + kafka_manager.log, + "info", + lambda event, data=None: info_events.append((event, data)), + ) + monkeypatch.setattr( + kafka_manager.log, + "warning", + lambda event, _data=None: warning_events.append(event), + ) + + manager = kafka_manager.KafkaManager() + await manager.start_all() + + assert manager.get_consumer_status(workflow_id) == { + "state": "stopped", + "error": "workflow_not_found", + } + assert info_events == [ + ( + "kafka.workflow_not_found_on_start", + {"workflow_id": workflow_id, "action": "stale_config_skipped"}, + ) + ] + assert warning_events == [] + + +@pytest.mark.asyncio +async def test_start_all_continues_after_one_workflow_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + restarted: list[str] = [] + warning_events: list[tuple[str, dict]] = [] + + async def _list_configs(*, kind: str): # noqa: ANN202 + assert kind == "workflow_kafka_config" + return [ + ("broken-workflow", {"enabled": True}), + ("healthy-workflow", {"enabled": True}), + ] + + async def _restart(workflow_id: str, *, startup: bool = False) -> dict: + assert startup is True + restarted.append(workflow_id) + if workflow_id == "broken-workflow": + raise ValueError("invalid config") + return {"state": "running", "error": None} + + monkeypatch.setattr(kafka_manager.WorkflowStore, "list_configs", _list_configs) + monkeypatch.setattr(kafka_manager.log, "warning", lambda event, data: warning_events.append((event, data))) + + manager = kafka_manager.KafkaManager() + monkeypatch.setattr(manager, "restart_workflow", _restart) + + await manager.start_all() + + assert restarted == ["broken-workflow", "healthy-workflow"] + assert warning_events == [ + ( + "kafka.start_failed", + {"workflow_id": "broken-workflow", "error": "invalid config"}, + ) + ] + + @pytest.mark.asyncio async def test_worker_pool_bounds_in_flight_dispatches(monkeypatch: pytest.MonkeyPatch) -> None: """The fixed worker pool must cap concurrent ``_trigger_workflow`` calls.""" @@ -61,7 +161,7 @@ async def _fake_trigger(workflow_id, workflow_json, msg, input_key, producer=Non manager._abort_events[workflow_id] = abort workers = [ asyncio.create_task( - manager._worker_loop(workflow_id, {}, trigger, {}, queue, abort, "topic-a"), + manager._worker_loop(workflow_id, {}, trigger, {}, queue, abort, "topic-a"), name=f"test-worker-{i}", ) for i in range(pool_size) @@ -83,8 +183,7 @@ async def _fake_trigger(workflow_id, workflow_json, msg, input_key, producer=Non assert completed == burst_size, f"expected {burst_size} dispatches, got {completed}" assert max_in_flight <= pool_size, ( - f"in-flight dispatches exceeded worker pool size: " - f"max_in_flight={max_in_flight}, pool_size={pool_size}" + f"in-flight dispatches exceeded worker pool size: max_in_flight={max_in_flight}, pool_size={pool_size}" ) @@ -190,6 +289,11 @@ async def _fake_get_config(workflow_id: str, *, kind: str) -> dict: return {"enabled": True, "inputBroker": "", "inputTopic": ""} monkeypatch.setattr(kafka_manager.WorkflowStore, "get_config", _fake_get_config) + monkeypatch.setattr( + kafka_manager, + "read_workflow_from_fs", + lambda _workflow_id: {"workflowJson": {}}, + ) status = await manager.restart_workflow("wf-no-broker") assert status["state"] == "failed" @@ -405,6 +509,7 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 @pytest.mark.asyncio async def test_trigger_workflow_applies_mapping_and_filter( monkeypatch: pytest.MonkeyPatch, + trigger_tool_context: SimpleNamespace, ) -> None: manager = kafka_manager.KafkaManager() captured_run_kwargs: dict = {} @@ -456,6 +561,12 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 assert captured_run_kwargs["inputs"]["order_id"] == 7 assert captured_run_kwargs["inputs"]["region"] == "cn" assert captured_run_kwargs["inputs"]["pipeline"] == "orders" + assert captured_run_kwargs["tool_context"] is trigger_tool_context.context + trigger_tool_context.builder.assert_awaited_once_with( + workflow_id="wf-orders", + action_name="trigger:kafka", + ) + trigger_tool_context.cleanup.assert_awaited_once_with(trigger_tool_context.context) assert recorded_exec_data["triggerId"] == "kafka-orders" assert recorded_exec_data["triggerSource"] == "orders-topic" diff --git a/tests/ingest/test_syslog_manager_backpressure.py b/tests/ingest/test_syslog_manager_backpressure.py index a042e2fc5..5309e4be8 100644 --- a/tests/ingest/test_syslog_manager_backpressure.py +++ b/tests/ingest/test_syslog_manager_backpressure.py @@ -18,6 +18,7 @@ import asyncio from types import SimpleNamespace +from unittest.mock import AsyncMock import pytest @@ -26,6 +27,16 @@ from flocks.workflow.triggers.models import TriggerDefinition +@pytest.fixture(autouse=True) +def trigger_tool_context(monkeypatch: pytest.MonkeyPatch) -> SimpleNamespace: + context = object() + builder = AsyncMock(return_value=context) + cleanup = AsyncMock() + monkeypatch.setattr(syslog_manager, "build_workflow_tool_context", builder) + monkeypatch.setattr(syslog_manager, "cleanup_workflow_tool_context", cleanup) + return SimpleNamespace(context=context, builder=builder, cleanup=cleanup) + + @pytest.mark.asyncio async def test_worker_pool_bounds_in_flight_dispatches(monkeypatch: pytest.MonkeyPatch) -> None: """The fixed worker pool must cap concurrent ``_trigger_workflow`` calls. @@ -96,8 +107,7 @@ async def _fake_trigger(workflow_id, workflow_json, msg, input_key, **kwargs): assert completed == burst_size, f"expected {burst_size} dispatches, got {completed}" assert max_in_flight <= pool_size, ( - f"in-flight dispatches exceeded worker pool size: " - f"max_in_flight={max_in_flight}, pool_size={pool_size}" + f"in-flight dispatches exceeded worker pool size: max_in_flight={max_in_flight}, pool_size={pool_size}" ) @@ -167,6 +177,7 @@ async def _noop_trigger(*args, **kwargs): # noqa: ANN001, D401 @pytest.mark.asyncio async def test_trigger_workflow_applies_mapping_and_filter( monkeypatch: pytest.MonkeyPatch, + trigger_tool_context: SimpleNamespace, ) -> None: manager = syslog_manager.SyslogManager() captured_run_kwargs: dict = {} @@ -240,7 +251,13 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 assert captured_run_kwargs["inputs"]["pipeline"] == "syslog" 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["on_step_complete"]) + trigger_tool_context.builder.assert_awaited_once_with( + workflow_id="wf-syslog", + action_name="trigger:syslog", + ) + trigger_tool_context.cleanup.assert_awaited_once_with(trigger_tool_context.context) assert recorded_steps[0][0] == "exec-syslog" assert recorded_steps[0][1] == 1 assert recorded_steps[0][2]["node_id"] == "receive_alert" diff --git a/tests/ingest/test_syslog_manager_bind_failure.py b/tests/ingest/test_syslog_manager_bind_failure.py index ecd621fcb..75809307d 100644 --- a/tests/ingest/test_syslog_manager_bind_failure.py +++ b/tests/ingest/test_syslog_manager_bind_failure.py @@ -14,6 +14,7 @@ import asyncio import socket +from unittest.mock import AsyncMock, MagicMock import pytest @@ -28,6 +29,84 @@ def _find_busy_udp_port() -> tuple[socket.socket, int]: return sock, port +@pytest.mark.asyncio +async def test_start_all_logs_stale_workflow_config_as_info( + monkeypatch: pytest.MonkeyPatch, +) -> None: + workflow_id = "removed-workflow" + config = {"workflowId": workflow_id, "enabled": True} + info_log = MagicMock() + warning_log = MagicMock() + + monkeypatch.setattr( + syslog_manager.WorkflowStore, + "list_configs", + AsyncMock(return_value=[(workflow_id, config)]), + ) + monkeypatch.setattr( + syslog_manager.WorkflowStore, + "get_config", + AsyncMock(return_value=config), + ) + monkeypatch.setattr(syslog_manager, "read_workflow_from_fs", lambda _workflow_id: None) + monkeypatch.setattr(syslog_manager.log, "info", info_log) + monkeypatch.setattr(syslog_manager.log, "warning", warning_log) + + manager = syslog_manager.SyslogManager() + await manager.start_all() + + assert manager.get_listener_status(workflow_id) == { + "state": "stopped", + "error": "workflow_not_found", + } + info_log.assert_called_once_with( + "syslog.workflow_not_found_on_start", + { + "workflow_id": workflow_id, + "action": "stale_config_skipped", + }, + ) + warning_log.assert_not_called() + + +@pytest.mark.asyncio +async def test_start_all_continues_after_one_workflow_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + restarted: list[str] = [] + warning_log = MagicMock() + + monkeypatch.setattr( + syslog_manager.WorkflowStore, + "list_configs", + AsyncMock( + return_value=[ + ("broken-workflow", {"enabled": True}), + ("healthy-workflow", {"enabled": True}), + ] + ), + ) + + async def _restart(workflow_id: str, *, startup: bool = False) -> dict: + assert startup is True + restarted.append(workflow_id) + if workflow_id == "broken-workflow": + raise ValueError("invalid config") + return {"state": "listening", "error": None} + + manager = syslog_manager.SyslogManager() + monkeypatch.setattr(manager, "restart_workflow", _restart) + monkeypatch.setattr(syslog_manager.log, "warning", warning_log) + + await manager.start_all() + + assert restarted == ["broken-workflow", "healthy-workflow"] + warning_log.assert_called_once_with( + "syslog.start_failed", + {"workflow_id": "broken-workflow", "error": "invalid config"}, + ) + + @pytest.mark.asyncio async def test_restart_workflow_reports_failure_on_port_conflict( monkeypatch: pytest.MonkeyPatch, @@ -46,11 +125,6 @@ async def test_restart_workflow_reports_failure_on_port_conflict( "inputKey": "syslog_message", } - async def _fake_storage_read(key: str): # noqa: ANN001 - if key == syslog_manager.SyslogManager._config_key(workflow_id): - return config - return None - def _fake_read_workflow_from_fs(wid: str): # noqa: ANN001 return { "id": wid, @@ -62,7 +136,11 @@ def _fake_read_workflow_from_fs(wid: str): # noqa: ANN001 } # Patch the *module-level* names ``manager.py`` looks up at call time. - monkeypatch.setattr(syslog_manager.Storage, "read", _fake_storage_read) + monkeypatch.setattr( + syslog_manager.WorkflowStore, + "get_config", + AsyncMock(return_value=config), + ) monkeypatch.setattr(syslog_manager, "read_workflow_from_fs", _fake_read_workflow_from_fs) manager = syslog_manager.SyslogManager() @@ -95,10 +173,11 @@ async def test_restart_workflow_returns_stopped_when_disabled( "inputKey": "syslog_message", } - async def _fake_storage_read(key: str): # noqa: ANN001 - return config - - monkeypatch.setattr(syslog_manager.Storage, "read", _fake_storage_read) + monkeypatch.setattr( + syslog_manager.WorkflowStore, + "get_config", + AsyncMock(return_value=config), + ) manager = syslog_manager.SyslogManager() status = await manager.restart_workflow(workflow_id) diff --git a/tests/mcp/test_mcp_client.py b/tests/mcp/test_mcp_client.py index 6cb6bbfd6..3cd12fc25 100644 --- a/tests/mcp/test_mcp_client.py +++ b/tests/mcp/test_mcp_client.py @@ -295,3 +295,65 @@ async def fake_remote(startup_future): assert cancelled.is_set() is True assert client._owner_task is None assert client._command_queue is None + + +class TestMcpClientOwnerFailure: + @pytest.mark.asyncio + async def test_call_tool_fails_when_owner_exits_during_request(self): + client = McpClient( + name="demo", + server_type="remote", + url="https://example.com/mcp", + ) + client._connected = True + client._owner_loop = asyncio.get_running_loop() + client._command_queue = asyncio.Queue() + + async def failing_owner() -> None: + await client._command_queue.get() + raise RuntimeError("HTTP 504 from https://example.com/mcp") + + owner_task = asyncio.create_task(failing_owner()) + owner_task.add_done_callback(client._handle_owner_task_done) + client._owner_task = owner_task + + with pytest.raises(RuntimeError, match="HTTP 504"): + await asyncio.wait_for( + client.call_tool("demo_tool", {"value": "x"}), + timeout=1.0, + ) + + @pytest.mark.asyncio + async def test_call_tool_cancels_response_when_caller_is_cancelled(self): + client = McpClient( + name="demo", + server_type="remote", + url="https://example.com/mcp", + ) + client._connected = True + client._owner_loop = asyncio.get_running_loop() + client._command_queue = asyncio.Queue() + command_seen = asyncio.get_running_loop().create_future() + + async def blocked_owner() -> None: + command = await client._command_queue.get() + command_seen.set_result(command) + await asyncio.Future() + + owner_task = asyncio.create_task(blocked_owner()) + client._owner_task = owner_task + + try: + call_task = asyncio.create_task(client.call_tool("demo_tool", {})) + command = await asyncio.wait_for(command_seen, timeout=1.0) + call_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await call_task + + assert command.response is not None + assert command.response.cancelled() is True + finally: + owner_task.cancel() + with pytest.raises(asyncio.CancelledError): + await owner_task diff --git a/tests/mcp/test_mcp_client_sse.py b/tests/mcp/test_mcp_client_sse.py index 479a32a44..ed4e9487d 100644 --- a/tests/mcp/test_mcp_client_sse.py +++ b/tests/mcp/test_mcp_client_sse.py @@ -5,9 +5,12 @@ from types import MethodType, SimpleNamespace import pytest +from mcp.shared.exceptions import McpError +from mcp.types import ErrorData, METHOD_NOT_FOUND import flocks.mcp.client as mcp_client_module from flocks.mcp.client import McpClient, _extract_root_cause +from flocks.mcp.errors import is_method_not_found_error def _make_session_class( @@ -374,3 +377,21 @@ class MockRequest: result = _extract_root_cause(exc) assert "401" in result assert "secret" not in result + + def test_method_not_found_detection_uses_structured_error_code(self): + exc = McpError(ErrorData(code=METHOD_NOT_FOUND, message="Unsupported operation")) + assert is_method_not_found_error(exc) is True + + def test_method_not_found_detection_checks_every_exception_group_child(self): + method_error = McpError( + ErrorData(code=METHOD_NOT_FOUND, message="Unsupported operation") + ) + exc = ExceptionGroup( + "request failed", + [RuntimeError("transport closed"), method_error], + ) + assert is_method_not_found_error(exc) is True + + def test_method_not_found_detection_does_not_match_unrelated_phrase(self): + exc = RuntimeError("cache method not found during local lookup") + assert is_method_not_found_error(exc) is False diff --git a/tests/mcp/test_mcp_server.py b/tests/mcp/test_mcp_server.py index b89a6b26d..7e3f6c098 100644 --- a/tests/mcp/test_mcp_server.py +++ b/tests/mcp/test_mcp_server.py @@ -4,7 +4,10 @@ from types import SimpleNamespace import pytest +from mcp.shared.exceptions import McpError +from mcp.types import ErrorData, METHOD_NOT_FOUND +from flocks.mcp import server as mcp_server_module from flocks.mcp.server import McpServerManager from flocks.mcp.types import McpStatus @@ -81,6 +84,70 @@ def __init__(self, **kwargs) -> None: assert manager._status["legacy-demo"].status == McpStatus.CONNECTED +@pytest.mark.asyncio +async def test_connect_treats_missing_resources_method_as_optional(monkeypatch: pytest.MonkeyPatch): + class ToolsOnlyClient(_FakeMcpClient): + async def list_resources(self) -> list: + raise McpError( + ErrorData(code=METHOD_NOT_FOUND, message="Unsupported operation") + ) + + info_events: list[str] = [] + warn_events: list[str] = [] + monkeypatch.setattr(mcp_server_module, "McpClient", ToolsOnlyClient) + monkeypatch.setattr( + mcp_server_module.log, + "info", + lambda event, _data=None: info_events.append(event), + ) + monkeypatch.setattr( + mcp_server_module.log, + "warn", + lambda event, _data=None: warn_events.append(event), + ) + + manager = McpServerManager() + await manager._connect_and_register( + "tools-only", + {"type": "local", "command": ["python", "-m", "demo"]}, + ) + + assert manager._resources_cache["tools-only"] == [] + assert manager._status["tools-only"].status == McpStatus.CONNECTED + assert "mcp.resources_unsupported" in info_events + assert "mcp.resources_unavailable" not in warn_events + + +@pytest.mark.asyncio +async def test_connect_logs_resource_discovery_failure_once(monkeypatch: pytest.MonkeyPatch): + class BrokenResourcesClient(_FakeMcpClient): + async def list_resources(self) -> list: + raise RuntimeError("resource endpoint unavailable") + + warn_events: list[str] = [] + error_events: list[str] = [] + monkeypatch.setattr(mcp_server_module, "McpClient", BrokenResourcesClient) + monkeypatch.setattr( + mcp_server_module.log, + "warn", + lambda event, _data=None: warn_events.append(event), + ) + monkeypatch.setattr( + mcp_server_module.log, + "error", + lambda event, _data=None: error_events.append(event), + ) + + manager = McpServerManager() + await manager._connect_and_register( + "broken-resources", + {"type": "local", "command": ["python", "-m", "demo"]}, + ) + + assert warn_events == ["mcp.resources_unavailable"] + assert error_events == [] + + @pytest.mark.asyncio async def test_init_is_serialized_for_concurrent_callers(monkeypatch: pytest.MonkeyPatch): manager = McpServerManager() diff --git a/tests/project/test_project.py b/tests/project/test_project.py new file mode 100644 index 000000000..bd9e5e7a2 --- /dev/null +++ b/tests/project/test_project.py @@ -0,0 +1,204 @@ +import asyncio +import json +import uuid +from unittest.mock import AsyncMock, patch + +import pytest + +from flocks.project.project import ( + DEFAULT_PROJECT_ID, + Project, + ProjectDeletionError, + ProjectNameConflictError, + ProjectPathConflictError, + TASK_SESSION_GROUP_ID, +) +from flocks.storage.storage import Storage + + +@pytest.fixture +def project_root(tmp_path, monkeypatch): + root = tmp_path / ".flocks" + monkeypatch.setenv("FLOCKS_ROOT", str(root)) + monkeypatch.setenv("FLOCKS_PROJECT_ROOTS", str(tmp_path)) + return tmp_path + + +@pytest.mark.asyncio +async def test_list_projects_uses_json_registry_without_virtual_default(project_root): + labs = project_root / "labs" + labs.mkdir() + created = await Project.create(owner_id="user-1", name="Labs", worktree=str(labs)) + + with patch.object(Storage, "list_entries", new=AsyncMock()) as list_entries: + projects = await Project.list(owner_id="user-1") + + assert [project.id for project in projects] == [created.id] + assert projects[0].worktree == str(labs.resolve()) + assert projects[0].is_default is False + assert await Project.get(DEFAULT_PROJECT_ID, owner_id="user-1") is None + list_entries.assert_not_awaited() + + registry = json.loads(Project.registry_path("user-1").read_text(encoding="utf-8")) + assert "defaultWorktree" not in registry + assert not (project_root / "workspace").exists() + assert registry["projects"][0]["id"].startswith("prj_") + uuid.UUID(registry["projects"][0]["id"].removeprefix("prj_")) + + +@pytest.mark.asyncio +async def test_create_rejects_duplicate_canonical_directory(project_root): + labs = project_root / "labs" + labs.mkdir() + existing = await Project.create(owner_id="user-1", name="Labs", worktree=str(labs)) + + with pytest.raises(ProjectPathConflictError) as exc_info: + await Project.create(owner_id="user-1", name="Other", worktree=str(labs / ".")) + + assert exc_info.value.project.id == existing.id + assert len((await Project.list(owner_id="user-1"))) == 1 + + +@pytest.mark.asyncio +async def test_create_rejects_duplicate_name_across_directories(project_root): + first = project_root / "first" + second = project_root / "second" + first.mkdir() + second.mkdir() + await Project.create(owner_id="user-1", name="Labs", worktree=str(first)) + + with pytest.raises(ProjectNameConflictError): + await Project.create(owner_id="user-1", name=" labs ", worktree=str(second)) + + +@pytest.mark.asyncio +async def test_create_project_creates_a_missing_worktree(project_root): + worktree = project_root / "new" / "nested-project" + + project = await Project.create( + owner_id="user-1", + name="Nested Project", + worktree=str(worktree), + ) + + assert worktree.is_dir() + assert project.worktree == str(worktree.resolve()) + + +@pytest.mark.asyncio +async def test_create_project_does_not_create_outside_allowed_roots(project_root, monkeypatch): + allowed_root = project_root / "allowed" + allowed_root.mkdir() + outside = project_root / "outside" / "project" + monkeypatch.setenv("FLOCKS_PROJECT_ROOTS", str(allowed_root)) + + with pytest.raises(ValueError, match="outside the allowed roots"): + await Project.create( + owner_id="user-1", + name="Outside", + worktree=str(outside), + ) + + assert not outside.exists() + + +@pytest.mark.asyncio +async def test_update_renames_registered_project(project_root): + labs = project_root / "labs" + labs.mkdir() + created = await Project.create(owner_id="user-1", name="Labs", worktree=str(labs)) + + updated = await Project.update(created.id, owner_id="user-1", name=" Security ") + + assert updated.name == "Security" + assert (await Project.get(created.id, owner_id="user-1")).name == "Security" + + +@pytest.mark.asyncio +async def test_delete_only_removes_registry_entry(project_root): + labs = project_root / "labs" + labs.mkdir() + created = await Project.create(owner_id="user-1", name="Labs", worktree=str(labs)) + + with patch.object(Storage, "delete", new=AsyncMock()) as storage_delete: + result = await Project.delete(created.id, owner_id="user-1") + + assert result is True + assert labs.exists() + assert await Project.get(created.id, owner_id="user-1") is None + storage_delete.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_delete_rejects_default_project(project_root): + await Project.ensure_registry("user-1") + + with pytest.raises(ProjectDeletionError): + await Project.delete(DEFAULT_PROJECT_ID, owner_id="user-1") + + +@pytest.mark.asyncio +async def test_effective_project_maps_legacy_and_removed_projects_to_tasks(project_root): + labs = project_root / "labs" + labs.mkdir() + created = await Project.create(owner_id="user-1", name="Labs", worktree=str(labs)) + + assert Project.effective_project_id("user-1", created.id) == created.id + assert Project.effective_project_id("user-1", "old-git-project") == TASK_SESSION_GROUP_ID + + await Project.delete(created.id, owner_id="user-1") + assert Project.effective_project_id("user-1", created.id) == TASK_SESSION_GROUP_ID + + +@pytest.mark.asyncio +async def test_from_directory_never_writes_project_storage(project_root): + with patch.object(Storage, "write", new=AsyncMock()) as storage_write: + result = await Project.from_directory(str(project_root)) + + assert result["project"].id == DEFAULT_PROJECT_ID + storage_write.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_concurrent_project_creates_keep_every_registry_entry(project_root): + worktrees = [] + for index in range(8): + worktree = project_root / f"project-{index}" + worktree.mkdir() + worktrees.append(worktree) + + created = await asyncio.gather(*( + Project.create( + owner_id="user-1", + name=f"Project {index}", + worktree=str(worktree), + ) + for index, worktree in enumerate(worktrees) + )) + + projects = await Project.list(owner_id="user-1") + assert {project.id for project in projects} == {project.id for project in created} + + +@pytest.mark.asyncio +async def test_registry_writes_do_not_create_backup(project_root): + labs = project_root / "labs" + labs.mkdir() + created = await Project.create(owner_id="user-1", name="Labs", worktree=str(labs)) + await Project.update(created.id, owner_id="user-1", name="Renamed Labs") + + registry_path = Project.registry_path("user-1") + assert registry_path.exists() + assert not registry_path.with_suffix(".json.bak").exists() + + +@pytest.mark.asyncio +async def test_corrupt_registry_returns_no_projects(project_root): + await Project.ensure_registry("user-1") + registry_path = Project.registry_path("user-1") + registry_path.write_text("{broken", encoding="utf-8") + + projects = await Project.list(owner_id="user-1") + + assert projects == [] + assert registry_path.read_text(encoding="utf-8") == "{broken" diff --git a/tests/provider/test_anthropic_interleaved.py b/tests/provider/test_anthropic_interleaved.py index 573d8b3e7..8241e235b 100644 --- a/tests/provider/test_anthropic_interleaved.py +++ b/tests/provider/test_anthropic_interleaved.py @@ -118,7 +118,7 @@ def test_anthropic_formatter_preserves_unsigned_thinking_for_deepseek_compatible formatted = AnthropicProvider._format_messages_anthropic( [message], base_url="https://api.deepseek.com/anthropic", - model_id="deepseek-chat", + model_id="deepseek-v4-flash", ) assert formatted[0]["content"][0] == { diff --git a/tests/provider/test_api_service_management.py b/tests/provider/test_api_service_management.py index 4647fb8df..fbb1d9ebf 100644 --- a/tests/provider/test_api_service_management.py +++ b/tests/provider/test_api_service_management.py @@ -585,6 +585,58 @@ def test_bootstrap_adds_enabled_true_when_key_missing(self): assert written["apiKey"] == "{secret:key}" assert written["base_url"] == "https://api.example.com" + def test_bootstrap_writes_unique_versioned_storage_key(self): + tools = {"tdp_tool": _make_api_tool("tdp_tool", "tdp_api")} + with patch( + "flocks.config.api_versioning.versioned_storage_key_for", + return_value="tdp_api_v3_3_10", + ): + mock_set = self._run_bootstrap(tools, get_raw_side_effect=lambda _: None) + + mock_set.assert_called_once_with("tdp_api_v3_3_10", {"enabled": True}) + + def test_bootstrap_then_sync_reads_the_same_versioned_storage_key(self): + tool = _make_api_tool("tdp_tool", "tdp_api") + tools = {"tdp_tool": tool} + services = { + "tdp_api": {"apiKey": "legacy-copy"}, + "tdp_api_v3_3_10": {"apiKey": "versioned-copy"}, + } + + def get_raw(storage_key: str): + value = services.get(storage_key) + return value.copy() if value is not None else None + + def set_service(storage_key: str, value: dict): + services[storage_key] = value.copy() + + with ( + patch.object(ToolRegistry, "_tools", tools), + patch.object(ToolRegistry, "_enabled_defaults", {"tdp_tool": True}), + patch( + "flocks.config.api_versioning.versioned_storage_key_for", + return_value="tdp_api_v3_3_10", + ), + patch( + "flocks.config.config_writer.ConfigWriter.get_api_service_raw", + side_effect=get_raw, + ), + patch( + "flocks.config.config_writer.ConfigWriter.set_api_service", + side_effect=set_service, + ), + patch( + "flocks.config.config_writer.ConfigWriter.list_api_services_raw", + side_effect=lambda: services, + ), + ): + ToolRegistry._bootstrap_user_api_services() + ToolRegistry._sync_api_service_states() + + assert services["tdp_api_v3_3_10"]["enabled"] is True + assert "enabled" not in services["tdp_api"] + assert tool.info.enabled is True + def test_bootstrap_does_not_overwrite_explicit_disabled(self): """When api_services..enabled is False, bootstrap leaves it untouched.""" existing = {"enabled": False, "apiKey": "{secret:key}"} diff --git a/tests/provider/test_chinese_providers.py b/tests/provider/test_chinese_providers.py index 92a7b1b5e..af892d611 100644 --- a/tests/provider/test_chinese_providers.py +++ b/tests/provider/test_chinese_providers.py @@ -141,19 +141,10 @@ def test_deepseek_catalog(self): models = get_provider_model_definitions("deepseek") assert {m.id for m in models} == { - "deepseek-chat", - "deepseek-reasoner", "deepseek-v4-flash", "deepseek-v4-pro", } - r1 = next(m for m in models if m.id == "deepseek-reasoner") - assert r1.capabilities.supports_reasoning is True - assert r1.capabilities.interleaved["field"] == "reasoning_content" - assert r1.capabilities.interleaved["placeholder"] == " " - assert r1.pricing.currency == "CNY" - assert r1.pricing.output == 16.0 - v4_flash = next(m for m in models if m.id == "deepseek-v4-flash") assert v4_flash.capabilities.supports_reasoning is True assert v4_flash.capabilities.interleaved["field"] == "reasoning_content" @@ -253,6 +244,7 @@ def test_threatbook_cn_llm_catalog(self): assert get_provider_default_url("threatbook-cn-llm") == "https://llm.threatbook.cn/v1" models = get_provider_model_definitions("threatbook-cn-llm") assert {m.id for m in models} == { + "kimi-k2.7-code", "minimax-m3", "minimax-m2.7", "minimax-m2.5", @@ -263,6 +255,19 @@ def test_threatbook_cn_llm_catalog(self): "deepseek-v4-flash", } + 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" + assert kimi_code.pricing.currency == "CNY" + assert kimi_code.pricing.cache_read == 1.3 + assert kimi_code.pricing.input == 6.5 + assert kimi_code.pricing.output == 27.0 + assert kimi_code.limits.context_window == 256000 + assert kimi_code.limits.max_input_tokens == 224000 + assert kimi_code.limits.max_output_tokens == 16000 + qwen = next(m for m in models if m.id == "qwen3.6-plus") assert qwen.capabilities.supports_vision is True @@ -297,6 +302,7 @@ def test_threatbook_io_llm_catalog(self): assert get_provider_default_url("threatbook-io-llm") == "https://llm.threatbook.io/v1" models = get_provider_model_definitions("threatbook-io-llm") assert {m.id for m in models} == { + "kimi-k2.7-code", "minimax-m3", "minimax-m2.7", "minimax-m2.5", @@ -306,6 +312,19 @@ def test_threatbook_io_llm_catalog(self): "deepseek-v4-flash", } + 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" + assert kimi_code.pricing.currency == "CNY" + assert kimi_code.pricing.cache_read == 1.3 + assert kimi_code.pricing.input == 6.5 + assert kimi_code.pricing.output == 27.0 + assert kimi_code.limits.context_window == 256000 + assert kimi_code.limits.max_input_tokens == 224000 + assert kimi_code.limits.max_output_tokens == 16000 + qwen = next(m for m in models if m.id == "qwen3.6-plus") assert qwen.capabilities.supports_vision is True diff --git a/tests/provider/test_reasoning_replay.py b/tests/provider/test_reasoning_replay.py index 15f901154..261446f6b 100644 --- a/tests/provider/test_reasoning_replay.py +++ b/tests/provider/test_reasoning_replay.py @@ -48,7 +48,7 @@ def test_prepare_reasoning_uses_placeholder_for_strict_echo_provider(): prepared = prepare_reasoning_for_replay( provider_id="deepseek", - model_id="deepseek-reasoner", + model_id="deepseek-v4-pro", message=message, interleaved={ "field": "reasoning_content", @@ -78,7 +78,7 @@ def test_prepare_reasoning_upgrades_whitespace_only_reasoning_content(): prepared = prepare_reasoning_for_replay( provider_id="deepseek", - model_id="deepseek-reasoner", + model_id="deepseek-v4-pro", message=message, interleaved={ "field": "reasoning_content", @@ -132,7 +132,7 @@ def test_prepare_reasoning_drops_details_when_target_uses_reasoning_content(): prepared = prepare_reasoning_for_replay( provider_id="deepseek", - model_id="deepseek-reasoner", + model_id="deepseek-v4-pro", message=message, interleaved={ "field": "reasoning_content", diff --git a/tests/provider/test_thinking_params.py b/tests/provider/test_thinking_params.py index 6d838d8fe..db82214f1 100644 --- a/tests/provider/test_thinking_params.py +++ b/tests/provider/test_thinking_params.py @@ -247,7 +247,7 @@ def test_glm5_emits_official_thinking_payload_on_every_provider( ("threatbook-io-llm", "minimax-m2.7", "reasoning_details", {"reasoning_split": True}), ("threatbook-io-llm", "minimax-m3", "reasoning_details", {"reasoning_split": True}), ("minimax", "minimax-m2.5", "reasoning_details", {"reasoning_split": True}), - ("deepseek", "deepseek-reasoner", "reasoning_content", DEEPSEEK_THINKING_EXTRA_BODY), + ("deepseek", "deepseek-v4-pro", "reasoning_content", DEEPSEEK_THINKING_EXTRA_BODY), ("stepfun", "step-3.5-flash", "reasoning_content", {"enable_thinking": True}), ], ) @@ -314,26 +314,6 @@ def test_no_shape_registry(self) -> None: "interleaved emits extra_body inline" ) - def test_deepseek_v3_is_not_auto_thinking_model(self) -> None: - """``deepseek-chat`` (V3) must not inherit thinking params from a - broad ``deepseek`` substring. - """ - catalog = model_catalog.get_raw_catalog() - v3_entry = catalog.get("deepseek", {}).get("models", {}).get("deepseek-chat") - assert v3_entry is not None, "deepseek-chat missing from catalog" - assert v3_entry.get("capabilities", {}).get("interleaved") is None, ( - "deepseek-chat now declares interleaved in catalog — remove the " - "series-token assertion and let the catalog coverage test pin it" - ) - - options = provider_options.build_provider_options( - "deepseek", "deepseek-chat", resolve_max_tokens=False, - ) - assert "extra_body" not in options, ( - "deepseek-chat does not declare interleaved in catalog and should " - f"not be auto-enabled by a broad deepseek token. options={options!r}" - ) - def test_explicit_reasoning_toggle_propagates(self) -> None: """``reasoning_enabled=False`` should produce ``enable_thinking: false`` on a generic_chat transport, mirroring the old token-matching branch's @@ -488,7 +468,6 @@ def test_series_token_fallback_emits_expected_extra_body( @pytest.mark.parametrize( "provider_id,model_id,expected_extra_body", [ - ("deepseek", "deepseek-reasoner", DEEPSEEK_THINKING_EXTRA_BODY), ("deepseek", "deepseek-v4-flash", DEEPSEEK_THINKING_EXTRA_BODY), ("minimax", "minimax-m3", {"reasoning_split": True}), ("stepfun", "step-3.5-flash", {"enable_thinking": True}), diff --git a/tests/server/conftest.py b/tests/server/conftest.py index 4e1b3f819..2880ab582 100644 --- a/tests/server/conftest.py +++ b/tests/server/conftest.py @@ -20,11 +20,18 @@ async def _server_isolated_env(tmp_path: Path, monkeypatch): """ data_dir = tmp_path / "flocks_data" data_dir.mkdir(parents=True, exist_ok=True) + flocks_root = tmp_path / ".flocks" + workspace_dir = flocks_root / "workspace" + workspace_dir.mkdir(parents=True, exist_ok=True) monkeypatch.setenv("FLOCKS_DATA_DIR", str(data_dir)) + monkeypatch.setenv("FLOCKS_ROOT", str(flocks_root)) + monkeypatch.setenv("FLOCKS_WORKSPACE_DIR", str(workspace_dir)) + monkeypatch.setenv("FLOCKS_PROJECT_ROOTS", str(tmp_path)) from flocks.config.config import Config from flocks.storage.storage import Storage from flocks.project.instance import Instance, _state_manager + from flocks.project.project import Project Config._global_config = None Config._cached_config = None @@ -38,6 +45,7 @@ async def _server_isolated_env(tmp_path: Path, monkeypatch): original_cache = dict(Instance._cache) Instance._cache.clear() + Project.invalidate_session_stats() for key in list(_state_manager._states.keys()): _state_manager._states.pop(key, None) _state_manager._disposers.pop(key, None) @@ -69,6 +77,7 @@ async def _server_isolated_env(tmp_path: Path, monkeypatch): AuthService._has_users_cached = False Instance._cache.clear() Instance._cache.update(original_cache) + Project.invalidate_session_stats() for key in list(_state_manager._states.keys()): _state_manager._states.pop(key, None) _state_manager._disposers.pop(key, None) diff --git a/tests/server/routes/test_console_upgrade_routes.py b/tests/server/routes/test_console_upgrade_routes.py index d4dc0b55e..b7c0c94df 100644 --- a/tests/server/routes/test_console_upgrade_routes.py +++ b/tests/server/routes/test_console_upgrade_routes.py @@ -1522,7 +1522,7 @@ async def test_auto_activate_reinstalls_when_existing_pro_marker_is_not_target_b async def _fake_perform_pro_bundle_install(*args, **kwargs): assert args == () - assert kwargs["restart"] is False + assert kwargs["restart"] is True marker_state["payload"] = { "release_id": "rel_20260605", "bundle_release_id": "rel_20260605", @@ -1627,7 +1627,7 @@ async def test_auto_activate_installs_pro_bundle_when_core_version_is_latest( async def _fake_perform_pro_bundle_install(*args, **kwargs): nonlocal installed assert args == () - assert kwargs["restart"] is False + assert kwargs["restart"] is True yield UpdateProgress(stage="syncing", message="Installing Flocks Pro component...", success=None) installed = True yield UpdateProgress(stage="done", message="Flocks Pro component installed from v2026.5.9", success=True) diff --git a/tests/server/routes/test_global_mutation_auth.py b/tests/server/routes/test_global_mutation_auth.py index 2c2bd8d69..40cdb9f67 100644 --- a/tests/server/routes/test_global_mutation_auth.py +++ b/tests/server/routes/test_global_mutation_auth.py @@ -40,6 +40,7 @@ def test_provider_global_configuration_credentials_and_tests_require_admin(): client.post("/api/provider/example/test"), client.put("/api/provider/example", json={}), client.get("/api/provider/example/credentials"), + client.post("/api/provider/example/credentials/reveal"), client.post("/api/provider/example/credentials", json={"api_key": "secret"}), client.delete("/api/provider/example/credentials"), client.get("/api/provider/example/service-credentials"), @@ -53,7 +54,9 @@ def test_provider_global_configuration_credentials_and_tests_require_admin(): assert [response.status_code for response in responses] == [403] * len(responses) -def test_admin_credential_routes_return_only_masked_secrets(monkeypatch: pytest.MonkeyPatch): +def test_admin_reveal_returns_raw_llm_key_with_audit_and_no_store( + monkeypatch: pytest.MonkeyPatch, +): from flocks.server.routes import provider as provider_routes raw_llm_key = "sk-raw-llm-secret-1234" @@ -63,7 +66,9 @@ def test_admin_credential_routes_return_only_masked_secrets(monkeypatch: pytest. "llm-example_llm_key": raw_llm_key, "service-example_api_key": raw_service_key, }.get(key) + audit = AsyncMock() monkeypatch.setattr("flocks.security.get_secret_manager", lambda: secrets) + monkeypatch.setattr(provider_routes, "emit_audit_event", audit) monkeypatch.setattr( provider_routes.ConfigWriter, "get_provider_raw", @@ -100,20 +105,33 @@ def test_admin_credential_routes_return_only_masked_secrets(monkeypatch: pytest. client = _user_client(provider_routes.router, prefix="/api/provider", role="admin") llm_response = client.get("/api/provider/llm-example/credentials") + reveal_response = client.post("/api/provider/llm-example/credentials/reveal") service_response = client.get("/api/provider/service-example/service-credentials") assert llm_response.status_code == 200 + assert reveal_response.status_code == 200 assert service_response.status_code == 200 llm_payload = llm_response.json() + reveal_payload = reveal_response.json() service_payload = service_response.json() assert llm_payload["api_key"] is None assert llm_payload["api_key_masked"] + assert reveal_payload["api_key"] == raw_llm_key + assert reveal_response.headers["Cache-Control"] == "no-store" + assert reveal_response.headers["Pragma"] == "no-cache" assert service_payload["api_key"] is None assert service_payload["api_key_masked"] assert service_payload["fields"]["api_key"] != raw_service_key assert service_payload["fields"]["base_url"] == "https://service.example.test" assert raw_llm_key not in llm_response.text + assert raw_llm_key in reveal_response.text assert raw_service_key not in service_response.text + audit.assert_awaited_once() + event_type, audit_payload = audit.await_args.args + assert event_type == "provider.credentials_reveal" + assert audit_payload["provider_id"] == "llm-example" + assert audit_payload["username"] == "admin-test" + assert raw_llm_key not in repr(audit_payload) def test_admin_provider_update_never_writes_raw_api_key_to_storage(monkeypatch: pytest.MonkeyPatch): diff --git a/tests/server/routes/test_notifications_routes.py b/tests/server/routes/test_notifications_routes.py index 1245f8295..a0400e801 100644 --- a/tests/server/routes/test_notifications_routes.py +++ b/tests/server/routes/test_notifications_routes.py @@ -25,7 +25,7 @@ async def test_notifications_require_browser_login(client: AsyncClient): async def test_active_notifications_and_dismiss_forever(client: AsyncClient): response = await client.get( "/api/notifications/active", - params={"locale": "zh-CN", "current_version": "2026.04.27"}, + params={"locale": "zh-CN"}, ) assert response.status_code == 200, response.text items = response.json() @@ -37,14 +37,14 @@ async def test_active_notifications_and_dismiss_forever(client: AsyncClient): response = await client.get( "/api/notifications/active", - params={"locale": "zh-CN", "current_version": "2026.04.27"}, + params={"locale": "zh-CN"}, ) assert response.status_code == 200, response.text assert response.json() == [] response = await client.get( "/api/notifications/active", - params={"locale": "zh-CN", "current_version": "2026.05.01"}, + params={"locale": "zh-CN"}, ) assert response.status_code == 200, response.text assert response.json() == [] diff --git a/tests/server/routes/test_onboarding_routes.py b/tests/server/routes/test_onboarding_routes.py index f2b812a48..53675b541 100644 --- a/tests/server/routes/test_onboarding_routes.py +++ b/tests/server/routes/test_onboarding_routes.py @@ -198,6 +198,14 @@ async def fake_test_mcp(region: str, api_key: str): class TestOnboardingApplyRoutes: + def test_threatbook_region_presets_use_kimi_k27_code(self): + assert onboarding_routes.ONBOARDING_REGION_PRESETS["cn"]["threatbook_default_model_id"] == ( + "kimi-k2.7-code" + ) + assert onboarding_routes.ONBOARDING_REGION_PRESETS["global"]["threatbook_default_model_id"] == ( + "kimi-k2.7-code" + ) + def test_ensure_threatbook_mcp_config_uses_explicit_secret_reference( self, monkeypatch: pytest.MonkeyPatch ): @@ -277,7 +285,7 @@ async def fake_connect_mcp(name: str): return True async def fake_set_default_model(model_type, body): - calls.append(("default_model", body.provider_id)) + calls.append(("default_model", body.provider_id, body.model_id)) return {"provider_id": body.provider_id, "model_id": body.model_id} monkeypatch.setattr(onboarding_routes, "_validate_onboarding_request", fake_validate) @@ -302,13 +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 ("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") in calls + assert ("default_model", "threatbook-cn-llm", "kimi-k2.7-code") in calls @pytest.mark.asyncio async def test_apply_returns_400_when_validation_fails( diff --git a/tests/server/routes/test_project_routes.py b/tests/server/routes/test_project_routes.py new file mode 100644 index 000000000..f147ebf13 --- /dev/null +++ b/tests/server/routes/test_project_routes.py @@ -0,0 +1,308 @@ +from pathlib import Path + +import pytest +from fastapi import status +from httpx import AsyncClient + +from flocks.auth.context import AuthUser +from flocks.project.project import Project +from flocks.session.session import Session +from flocks.storage.storage import Storage + + +@pytest.mark.asyncio +async def test_project_crud_uses_registry_without_project_database_rows( + client: AsyncClient, + tmp_path: Path, +): + worktree = tmp_path / "labs" + worktree.mkdir() + + create_response = await client.post( + "/api/project", + json={"name": "Labs", "worktree": str(worktree)}, + ) + assert create_response.status_code == status.HTTP_200_OK + project = create_response.json() + assert project["id"].startswith("prj_") + assert await Storage.list_keys(prefix="project/") == [] + + list_response = await client.get("/api/project") + assert list_response.status_code == status.HTTP_200_OK + projects = list_response.json() + assert [item["id"] for item in projects] == [project["id"]] + assert all(item["isDefault"] is False for item in projects) + + rename_response = await client.patch( + f"/api/project/{project['id']}", + json={"name": "Security Labs"}, + ) + assert rename_response.status_code == status.HTTP_200_OK + assert rename_response.json()["name"] == "Security Labs" + + retained_file = worktree / "keep.txt" + retained_file.write_text("project files stay", encoding="utf-8") + session_response = await client.post( + "/api/session", + json={"title": "Delete with project", "projectID": project["id"]}, + ) + assert session_response.status_code == status.HTTP_200_OK + session_id = session_response.json()["id"] + + delete_response = await client.delete(f"/api/project/{project['id']}") + assert delete_response.status_code == status.HTTP_200_OK + assert worktree.exists() + assert retained_file.read_text(encoding="utf-8") == "project files stay" + assert (await client.get("/api/project")).json() == [] + assert (await client.get(f"/api/session/{session_id}")).status_code == status.HTTP_404_NOT_FOUND + + +@pytest.mark.asyncio +async def test_create_project_creates_missing_directory( + client: AsyncClient, + tmp_path: Path, +): + worktree = tmp_path / "workspace" / "created-by-project" + + response = await client.post( + "/api/project", + json={"name": "Created Project", "worktree": str(worktree)}, + ) + + assert response.status_code == status.HTTP_200_OK + assert response.json()["worktree"] == str(worktree.resolve()) + assert worktree.is_dir() + + +@pytest.mark.asyncio +async def test_duplicate_project_directory_returns_existing_project( + client: AsyncClient, + tmp_path: Path, +): + worktree = tmp_path / "same" + worktree.mkdir() + + first = await client.post( + "/api/project", + json={"name": "First", "worktree": str(worktree)}, + ) + second = await client.post( + "/api/project", + json={"name": "Second", "worktree": str(worktree / ".")}, + ) + + assert first.status_code == status.HTTP_200_OK + assert second.status_code == status.HTTP_200_OK + assert second.json()["id"] == first.json()["id"] + assert len((await client.get("/api/project")).json()) == 1 + + +@pytest.mark.asyncio +async def test_folder_browser_starts_at_home_when_path_is_omitted( + client: AsyncClient, + monkeypatch: pytest.MonkeyPatch, +): + home = Path.home().resolve() + monkeypatch.setenv("FLOCKS_PROJECT_ROOTS", str(home)) + + response = await client.get("/api/project/folders") + + assert response.status_code == status.HTTP_200_OK + assert response.json()["path"] == str(home) + + +@pytest.mark.asyncio +async def test_folder_browser_lists_directories_and_blocks_outside_root( + client: AsyncClient, + tmp_path: Path, +): + parent = tmp_path / "parent" + child = parent / "child" + child.mkdir(parents=True) + + response = await client.get("/api/project/folders", params={"path": str(parent)}) + + assert response.status_code == status.HTTP_200_OK + payload = response.json() + assert payload["path"] == str(parent.resolve()) + assert {entry["name"] for entry in payload["entries"]} == {"child"} + + outside = await client.get("/api/project/folders", params={"path": str(Path.home())}) + assert outside.status_code == status.HTTP_403_FORBIDDEN + + escape_link = parent / "escape" + escape_link.symlink_to(Path.home(), target_is_directory=True) + escaped_listing = await client.get( + "/api/project/folders", + params={"path": str(escape_link)}, + ) + assert escaped_listing.status_code == status.HTTP_403_FORBIDDEN + + +@pytest.mark.asyncio +async def test_project_local_sharing_includes_existing_and_future_sessions( + client: AsyncClient, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +): + from flocks.server.routes import project as project_routes + from flocks.server.routes import session as session_routes + + owner = AuthUser(id="usr_owner", username="owner", role="member", status="active") + viewer = AuthUser(id="usr_viewer", username="viewer", role="member", status="active") + worktree = tmp_path / "shared-project" + worktree.mkdir() + project = await Project.create( + owner_id=owner.id, + name="Shared Project", + worktree=str(worktree), + ) + existing = await Session.create( + project_id=project.id, + directory=str(worktree), + title="Existing session", + owner_user_id=owner.id, + owner_username=owner.username, + ) + + monkeypatch.setattr(project_routes, "require_user", lambda _request: owner) + share_response = await client.post(f"/api/project/{project.id}/share-local") + + assert share_response.status_code == status.HTTP_200_OK + assert share_response.json()["isShared"] is True + + future = await Session.create( + project_id=project.id, + directory=str(worktree), + title="Future session", + owner_user_id=owner.id, + owner_username=owner.username, + ) + + monkeypatch.setattr(project_routes, "require_user", lambda _request: viewer) + project_list = (await client.get("/api/project")).json() + shared_project = next(item for item in project_list if item["id"] == project.id) + assert shared_project["isShared"] is True + assert shared_project["canWrite"] is False + assert shared_project["canDelete"] is False + + non_owner_share = await client.post(f"/api/project/{project.id}/share-local") + assert non_owner_share.status_code == status.HTTP_404_NOT_FOUND + + monkeypatch.setattr(session_routes, "require_user", lambda _request: viewer) + session_list = ( + await client.get( + "/api/session", + params={"view": "list", "manager": True, "projectID": project.id}, + ) + ).json() + assert {item["id"] for item in session_list} == {existing.id, future.id} + assert all(item["effectiveProjectID"] == project.id for item in session_list) + assert all(item["isShared"] is True for item in session_list) + assert all(item["canWrite"] is False for item in session_list) + + monkeypatch.setattr(project_routes, "require_user", lambda _request: owner) + unshare_response = await client.post(f"/api/project/{project.id}/unshare-local") + assert unshare_response.status_code == status.HTTP_200_OK + assert unshare_response.json()["isShared"] is False + + monkeypatch.setattr(project_routes, "require_user", lambda _request: viewer) + assert project.id not in {item["id"] for item in (await client.get("/api/project")).json()} + monkeypatch.setattr(session_routes, "require_user", lambda _request: viewer) + assert ( + await client.get( + "/api/session", + params={"view": "list", "manager": True, "projectID": project.id}, + ) + ).json() == [] + + +@pytest.mark.asyncio +async def test_project_counts_and_session_lists_are_user_scoped( + client: AsyncClient, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +): + from flocks.server.routes import project as project_routes + from flocks.server.routes import session as session_routes + + alice = AuthUser(id="usr_alice", username="alice", role="member", status="active") + bob = AuthUser(id="usr_bob", username="bob", role="member", status="active") + + alice_worktree = tmp_path / "alice-project" + alice_worktree.mkdir() + monkeypatch.setattr(project_routes, "require_user", lambda _request: alice) + project_response = await client.post( + "/api/project", + json={"name": "Alice Project", "worktree": str(alice_worktree)}, + ) + alice_project = project_response.json() + + monkeypatch.setattr(session_routes, "require_user", lambda _request: alice) + alice_session = await client.post( + "/api/session", + json={"title": "Alice session", "projectID": alice_project["id"]}, + ) + assert alice_session.status_code == status.HTTP_200_OK + + monkeypatch.setattr(session_routes, "require_user", lambda _request: bob) + bob_session = await client.post("/api/session", json={"title": "Bob session"}) + assert bob_session.status_code == status.HTTP_200_OK + + monkeypatch.setattr(session_routes, "require_user", lambda _request: alice) + projects = (await client.get("/api/project")).json() + assert projects[0]["id"] == alice_project["id"] + assert projects[0]["sessionCount"] == 1 + + second_alice_session = await client.post( + "/api/session", + json={"title": "Second Alice session", "projectID": alice_project["id"]}, + ) + assert second_alice_session.status_code == status.HTTP_200_OK + refreshed_projects = (await client.get("/api/project")).json() + assert refreshed_projects[0]["sessionCount"] == 2 + + sessions = ( + await client.get( + "/api/session", + params={"view": "list", "manager": "true", "roots": "true"}, + ) + ).json() + assert {item["title"] for item in sessions} == { + "Alice session", + "Second Alice session", + } + + +@pytest.mark.asyncio +async def test_session_title_update_invalidates_project_search_stats( + client: AsyncClient, + tmp_path: Path, +): + worktree = tmp_path / "search-project" + worktree.mkdir() + project = ( + await client.post( + "/api/project", + json={"name": "Search Project", "worktree": str(worktree)}, + ) + ).json() + created = await client.post( + "/api/session", + json={"title": "Initial title", "projectID": project["id"]}, + ) + assert created.status_code == status.HTTP_200_OK + + initial = await client.get("/api/project", params={"search": "triage"}) + assert initial.status_code == status.HTTP_200_OK + assert initial.json()[0]["matchedSessionCount"] == 0 + + updated = await client.patch( + f"/api/session/{created.json()['id']}", + json={"title": "Triage findings"}, + ) + assert updated.status_code == status.HTTP_200_OK + + refreshed = await client.get("/api/project", params={"search": "triage"}) + assert refreshed.status_code == status.HTTP_200_OK + assert refreshed.json()[0]["matchedSessionCount"] == 1 diff --git a/tests/server/routes/test_provider_model_bootstrap.py b/tests/server/routes/test_provider_model_bootstrap.py index aaee5869c..5ab24fae7 100644 --- a/tests/server/routes/test_provider_model_bootstrap.py +++ b/tests/server/routes/test_provider_model_bootstrap.py @@ -31,6 +31,8 @@ async def test_set_provider_credentials_bootstraps_kimi_k26_from_catalog( raw = ConfigWriter.get_provider_raw("threatbook-cn-llm") assert raw is not None + assert "kimi-k2.7-code" in raw["models"] + assert raw["models"]["kimi-k2.7-code"]["name"] == "kimi-k2.7-code" assert "kimi-k2.6" in raw["models"] assert raw["models"]["kimi-k2.6"]["name"] == "kimi-k2.6" fake_secrets.set.assert_called_once_with("threatbook-cn-llm_llm_key", "tb-key") diff --git a/tests/server/routes/test_remaining_routes.py b/tests/server/routes/test_remaining_routes.py index ce681b707..173c4afdf 100644 --- a/tests/server/routes/test_remaining_routes.py +++ b/tests/server/routes/test_remaining_routes.py @@ -5,10 +5,11 @@ from __future__ import annotations import asyncio +import json from pathlib import Path import pytest -from fastapi import status +from fastapi import HTTPException, status from httpx import AsyncClient from unittest.mock import AsyncMock @@ -508,6 +509,148 @@ async def test_config_has_expected_top_level_keys(self, client: AsyncClient): f"No expected keys found. Got: {list(data.keys())}" ) + @pytest.mark.asyncio + async def test_null_allow_from_deletes_existing_channel_field( + self, + tmp_path, + monkeypatch, + ): + """channels..allowFrom=null removes the persisted allowlist.""" + from flocks.server.routes import config as config_routes + + config_file = tmp_path / "config" / "flocks.json" + config_file.parent.mkdir(parents=True) + config_file.write_text( + json.dumps( + { + "channels": { + "slack": { + "enabled": False, + "allowFrom": ["U123"], + } + } + } + ), + encoding="utf-8", + ) + monkeypatch.setattr(config_routes.Config, "get_config_file", lambda: config_file) + + async def fake_get_config(): + return {} + + monkeypatch.setattr(config_routes, "get_config", fake_get_config) + + await config_routes.update_config( + {"channels": {"slack": {"enabled": False, "allowFrom": None}}} + ) + + saved = json.loads(config_file.read_text(encoding="utf-8")) + assert "allowFrom" not in saved["channels"]["slack"] + assert saved["channels"]["slack"]["dmPolicy"] == "open" + + @pytest.mark.asyncio + async def test_slack_allow_from_save_sets_dm_policy_allowlist( + self, + tmp_path, + monkeypatch, + ): + """Slack allowFrom must also enable DM allowlist enforcement.""" + from flocks.server.routes import config as config_routes + + config_file = tmp_path / "config" / "flocks.json" + config_file.parent.mkdir(parents=True) + config_file.write_text("{}", encoding="utf-8") + monkeypatch.setattr(config_routes.Config, "get_config_file", lambda: config_file) + + async def fake_get_config(): + return {} + + monkeypatch.setattr(config_routes, "get_config", fake_get_config) + + await config_routes.update_config( + {"channels": {"slack": {"enabled": False, "allowFrom": ["U123"]}}} + ) + + saved = json.loads(config_file.read_text(encoding="utf-8")) + assert saved["channels"]["slack"]["allowFrom"] == ["U123"] + assert saved["channels"]["slack"]["dmPolicy"] == "allowlist" + + @pytest.mark.asyncio + async def test_invalid_patch_does_not_delete_channel_allow_from( + self, + tmp_path, + monkeypatch, + ): + """A failing PATCH must not relax persisted channel allowlists.""" + from flocks.server.routes import config as config_routes + + config_file = tmp_path / "config" / "flocks.json" + config_file.parent.mkdir(parents=True) + config_file.write_text( + json.dumps( + { + "share": "manual", + "channels": { + "slack": { + "enabled": False, + "allowFrom": ["U123"], + } + }, + } + ), + encoding="utf-8", + ) + monkeypatch.setattr(config_routes.Config, "get_config_file", lambda: config_file) + + with pytest.raises(HTTPException) as exc: + await config_routes.update_config( + { + "share": "invalid-mode", + "channels": {"slack": {"allowFrom": None}}, + } + ) + + assert exc.value.status_code == status.HTTP_400_BAD_REQUEST + saved = json.loads(config_file.read_text(encoding="utf-8")) + assert saved["channels"]["slack"]["allowFrom"] == ["U123"] + + @pytest.mark.asyncio + async def test_tool_failure_preference_defaults_on_and_updates_only_its_section( + self, + client: AsyncClient, + tmp_path, + monkeypatch, + ): + from flocks.config.config import Config + from flocks.config.config_writer import ConfigWriter + + monkeypatch.setenv("FLOCKS_CONFIG_DIR", str(tmp_path / "config")) + Config._global_config = None + Config._cached_config = None + ConfigWriter._write_raw({"theme": "dark"}) + + resp = await client.get("/api/config/tool-failure") + assert resp.status_code == status.HTTP_200_OK + assert resp.json() == {"disableOnRepeatedFailure": True} + + resp = await client.patch( + "/api/config/tool-failure", + json={"disableOnRepeatedFailure": False}, + ) + assert resp.status_code == status.HTTP_200_OK + assert resp.json() == {"disableOnRepeatedFailure": False} + assert ConfigWriter._read_raw() == { + "theme": "dark", + "toolFailure": {"disableOnRepeatedFailure": False}, + } + + resp = await client.get("/api/config/tool-failure") + assert resp.status_code == status.HTTP_200_OK + assert resp.json() == {"disableOnRepeatedFailure": False} + + resp = await client.patch("/api/config/tool-failure", json={}) + assert resp.status_code == status.HTTP_422_UNPROCESSABLE_CONTENT + @pytest.mark.asyncio async def test_ui_display_defaults_and_updates( self, diff --git a/tests/server/routes/test_session_routes.py b/tests/server/routes/test_session_routes.py index f46b9d153..05d0470f1 100644 --- a/tests/server/routes/test_session_routes.py +++ b/tests/server/routes/test_session_routes.py @@ -35,6 +35,37 @@ # CRUD # =========================================================================== + +@pytest.mark.asyncio +async def test_missing_session_directory_uses_cwd_and_publishes_notice( + tmp_path, + monkeypatch: pytest.MonkeyPatch, +): + from flocks.server.routes import event as event_routes + from flocks.server.routes import session as session_routes + + monkeypatch.chdir(tmp_path) + publish_event = AsyncMock() + monkeypatch.setattr(event_routes, "publish_event", publish_event) + + working_directory = await session_routes._resolve_session_working_directory( + SimpleNamespace( + id="ses_missing_directory", + directory=str(tmp_path / "missing"), + ), + ) + + assert working_directory == str(tmp_path) + publish_event.assert_awaited_once_with( + "session.notice", + { + "sessionID": "ses_missing_directory", + "kind": "directory-fallback", + "storedDirectory": str(tmp_path / "missing"), + "fallbackDirectory": str(tmp_path), + }, + ) + class TestSessionCRUD: """Basic create / read / update / delete for sessions.""" @@ -48,6 +79,23 @@ async def test_create_session_minimal(self, client: AsyncClient): assert "projectID" in data assert "directory" in data + @pytest.mark.asyncio + async def test_create_ordinary_session_uses_process_cwd( + self, + client: AsyncClient, + tmp_path, + monkeypatch: pytest.MonkeyPatch, + ): + """A task session must not resolve through a virtual default project.""" + process_cwd = tmp_path / "server-cwd" + process_cwd.mkdir() + monkeypatch.chdir(process_cwd) + + response = await client.post("/api/session", json={}) + + assert response.status_code == status.HTTP_200_OK + assert response.json()["directory"] == str(process_cwd) + @pytest.mark.asyncio async def test_create_session_with_title(self, client: AsyncClient): """Created session reflects the provided title.""" @@ -235,6 +283,9 @@ async def test_list_sessions_light_manager_filters_and_omits_heavy_fields(self, row = next(item for item in data if item["id"] == user_id) assert set(row) == { "id", + "projectID", + "effectiveProjectID", + "directory", "title", "time", "category", @@ -246,9 +297,110 @@ async def test_list_sessions_light_manager_filters_and_omits_heavy_fields(self, "canDelete", "isShared", } + assert row["projectID"] + assert row["directory"] assert "goal" not in row assert "summary" not in row + @pytest.mark.asyncio + async def test_create_session_in_user_managed_project( + self, + client: AsyncClient, + tmp_path, + monkeypatch: pytest.MonkeyPatch, + ): + """Session creation can target a user-managed project.""" + worktree = tmp_path / "labs" + worktree.mkdir() + monkeypatch.setenv("FLOCKS_PROJECT_ROOTS", str(tmp_path)) + project_resp = await client.post( + "/api/project", + json={"name": "Labs", "worktree": str(worktree)}, + ) + assert project_resp.status_code == status.HTTP_200_OK + project = project_resp.json() + + session_resp = await client.post( + "/api/session", + json={"title": "Project Session", "projectID": project["id"]}, + ) + assert session_resp.status_code == status.HTTP_200_OK + assert session_resp.json()["projectID"] == project["id"] + + list_resp = await client.get( + "/api/session", + params={"view": "list", "manager": "true", "roots": "true", "limit": "100"}, + ) + row = next(item for item in list_resp.json() if item["id"] == session_resp.json()["id"]) + assert row["projectID"] == project["id"] + assert row["effectiveProjectID"] == project["id"] + assert row["directory"] == project["worktree"] + + @pytest.mark.asyncio + async def test_legacy_session_is_grouped_under_tasks_without_rewrite( + self, + client: AsyncClient, + tmp_path, + ): + """Legacy project IDs are projected to Tasks while storage stays unchanged.""" + legacy = await Session.create( + project_id="legacy-git-project", + directory=str(tmp_path), + title="Legacy Session", + ) + + response = await client.get( + "/api/session", + params={ + "view": "list", + "manager": "true", + "roots": "true", + "projectID": "tasks", + }, + ) + + assert response.status_code == status.HTTP_200_OK + row = next(item for item in response.json() if item["id"] == legacy.id) + assert row["projectID"] == "legacy-git-project" + assert row["effectiveProjectID"] == "tasks" + stored = await Session.get("legacy-git-project", legacy.id) + assert stored is not None + assert stored.project_id == "legacy-git-project" + + @pytest.mark.asyncio + async def test_deleted_project_sessions_are_removed( + self, + client: AsyncClient, + tmp_path, + ): + """Deleting a project removes its sessions but preserves its directory.""" + worktree = tmp_path / "removable-project" + worktree.mkdir() + project_response = await client.post( + "/api/project", + json={"name": "Removable", "worktree": str(worktree)}, + ) + project = project_response.json() + session_response = await client.post( + "/api/session", + json={"title": "Keep Me", "projectID": project["id"]}, + ) + session_id = session_response.json()["id"] + + delete_response = await client.delete(f"/api/project/{project['id']}") + assert delete_response.status_code == status.HTTP_200_OK + assert worktree.exists() + + session_get_response = await client.get(f"/api/session/{session_id}") + assert session_get_response.status_code == status.HTTP_404_NOT_FOUND + + tasks_response = await client.get( + "/api/session", + params={"view": "list", "manager": "true", "projectID": "tasks"}, + ) + assert all(item["id"] != session_id for item in tasks_response.json()) + assert await Session.get(project["id"], session_id) is None + @pytest.mark.asyncio async def test_get_session(self, client: AsyncClient, session_id: str): """GET /api/session/{id} returns the specific session.""" @@ -803,6 +955,7 @@ async def test_owner_can_share_and_unshare_session( create_resp = await client.post("/api/session", json={"title": "share-session"}) assert create_resp.status_code == status.HTTP_200_OK session_id = create_resp.json()["id"] + assert await Session.get_by_id(session_id) is not None share_resp = await client.post(f"/api/session/{session_id}/share-local") assert share_resp.status_code == status.HTTP_200_OK diff --git a/tests/server/routes/test_skill_routes.py b/tests/server/routes/test_skill_routes.py new file mode 100644 index 000000000..2148dbf7e --- /dev/null +++ b/tests/server/routes/test_skill_routes.py @@ -0,0 +1,31 @@ +from pathlib import Path + +import pytest +from httpx import AsyncClient + +from flocks.skill.skill import Skill + + +@pytest.mark.asyncio +async def test_ordinary_request_keeps_bundled_skills_visible( + client: AsyncClient, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Bundled skills stay visible when an ordinary request starts elsewhere.""" + ordinary_cwd = tmp_path / "ordinary-cwd" + ordinary_cwd.mkdir() + monkeypatch.chdir(ordinary_cwd) + Skill.clear_cache() + + response = await client.get("/api/skills") + + assert response.status_code == 200 + source_skills_root = Path(Skill._source_root()) / ".flocks" + bundled_skills = [ + skill + for skill in response.json() + if Path(skill["location"]).is_relative_to(source_skills_root) + ] + assert bundled_skills + assert any(skill["source"] == "project" for skill in bundled_skills) diff --git a/tests/server/routes/test_tool_routes.py b/tests/server/routes/test_tool_routes.py index faa087fa8..036e8f68c 100644 --- a/tests/server/routes/test_tool_routes.py +++ b/tests/server/routes/test_tool_routes.py @@ -26,6 +26,7 @@ def _temporary_tool(tool: Tool) -> Iterator[None]: ToolRegistry.init() existing = ToolRegistry._tools.get(tool.info.name) + existing_default = ToolRegistry._enabled_defaults.get(tool.info.name) ToolRegistry.register(tool) _clear_tool_summary_cache() try: @@ -36,6 +37,10 @@ def _temporary_tool(tool: Tool) -> Iterator[None]: ToolRegistry._tools[tool.info.name] = existing else: ToolRegistry._tools.pop(tool.info.name, None) + if existing_default is not None: + ToolRegistry._enabled_defaults[tool.info.name] = existing_default + else: + ToolRegistry._enabled_defaults.pop(tool.info.name, None) _clear_tool_summary_cache() @@ -405,6 +410,67 @@ async def handler(ctx) -> ToolResult: class TestToolListPageRoute: + @pytest.mark.asyncio + async def test_auto_disable_syncs_cached_page_across_worker_state(self, client: AsyncClient): + from flocks.config.config_writer import ConfigWriter + + async def handler(_ctx, probe_id: str) -> ToolResult: + return ToolResult(success=False, error=f"repeated failure for {probe_id}") + + tool = Tool( + info=ToolInfo( + name="page_auto_disable_cache_tool", + description="NeedleAutoDisableCache", + category=ToolCategory.CUSTOM, + source="plugin_py", + ), + handler=handler, + ) + + with _temporary_tool(tool): + initial = await client.get( + "/api/tools/page", + params={"q": "needleautodisablecache", "limit": 25}, + ) + for _ in range(ToolRegistry._failure_disable_threshold): + result = await client.post( + f"/api/tools/{tool.info.name}/test", + json={"params": {"probe_id": "same-input"}}, + ) + setting = ConfigWriter.get_tool_setting(tool.info.name) + # Simulate a second worker whose in-memory registry and list cache + # still contain the pre-disable state. The persisted config token + # must make the next page request repair both immediately. + tool.info.enabled = True + refreshed = await client.get( + "/api/tools/page", + params={"q": "needleautodisablecache", "limit": 25}, + ) + disabled_after_refresh = tool.info.enabled + + # The reverse transition must also propagate: another worker can + # manually restore the default by deleting the disable overlay. + removed = ConfigWriter.delete_tool_setting(tool.info.name) + tool.info.enabled = False + reenabled = await client.get( + "/api/tools/page", + params={"q": "needleautodisablecache", "limit": 25}, + ) + + assert initial.status_code == 200, initial.text + assert initial.json()["items"][0]["enabled"] is True + assert result.status_code == 200, result.text + assert result.json()["metadata"]["disabled"] is True + assert setting == {"enabled": False} + assert disabled_after_refresh is False + assert refreshed.status_code == 200, refreshed.text + assert refreshed.json()["items"][0]["enabled"] is False + assert removed is True + assert tool.info.enabled is True + assert tool.info.name not in ToolRegistry._failure_state + assert reenabled.status_code == 200, reenabled.text + assert reenabled.json()["items"][0]["enabled"] is True + @pytest.mark.asyncio async def test_list_page_searches_server_side_and_omits_parameter_payload(self, client: AsyncClient): async def handler(ctx, text: str) -> ToolResult: diff --git a/tests/server/routes/test_workflow_health_monitor.py b/tests/server/routes/test_workflow_health_monitor.py new file mode 100644 index 000000000..40f6a7c59 --- /dev/null +++ b/tests/server/routes/test_workflow_health_monitor.py @@ -0,0 +1,209 @@ +from __future__ import annotations + +import asyncio +from types import SimpleNamespace +from typing import Any + +import pytest + +from flocks.server.routes import workflow as workflow_routes + + +@pytest.fixture(autouse=True) +def clear_workflow_health_cache() -> None: + workflow_routes._workflow_api_health_cache.clear() + + +@pytest.mark.asyncio +async def test_workflow_api_health_cache_has_process_wide_probe_limit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + workflow_ids = [f"wf-{index}" for index in range(7)] + active = 0 + peak = 0 + + async def fake_list_keys(_prefix: str) -> list[str]: + return [workflow_routes._api_service_key(workflow_id) for workflow_id in workflow_ids] + + async def fake_kv_get(key: Any, *_args: Any, **_kwargs: Any) -> dict[str, Any]: + workflow_id = str(key).removeprefix(workflow_routes._API_SERVICE_PREFIX) + return {"workflowId": workflow_id, "status": "running"} + + async def fake_health(_workflow_id: str) -> dict[str, Any]: + nonlocal active, peak + active += 1 + peak = max(peak, active) + await asyncio.sleep(0.01) + active -= 1 + return {"ok": True} + + monkeypatch.setattr(workflow_routes.WorkflowStore, "kv_list_keys", fake_list_keys) + monkeypatch.setattr(workflow_routes.WorkflowStore, "kv_get", fake_kv_get) + monkeypatch.setattr(workflow_routes, "get_workflow_health", fake_health) + + result = await workflow_routes.refresh_workflow_api_health_cache() + + assert result == {"checked": 7, "healthy": 7} + assert peak == workflow_routes._WORKFLOW_API_HEALTH_PROBE_CONCURRENCY + assert workflow_routes._workflow_api_health_cache == dict.fromkeys(workflow_ids, True) + + +@pytest.mark.asyncio +async def test_workflow_integration_status_summarizes_running_capabilities( + monkeypatch: pytest.MonkeyPatch, +) -> None: + workflow_id = "wf-running" + workflow_data = { + "id": workflow_id, + "workflowJson": { + "start": "n1", + "nodes": [{"id": "n1", "type": "python"}], + "edges": [], + "triggers": [ + {"id": "schedule-default", "type": "schedule", "enabled": True}, + {"id": "webhook-default", "type": "webhook", "enabled": True}, + ], + }, + } + + async def fake_kv_get(key: Any, *_args: Any, **_kwargs: Any) -> Any: + if str(key) == workflow_routes._api_service_key(workflow_id): + return {"workflowId": workflow_id, "status": "running"} + if str(key) == workflow_routes._runtime_key_main(workflow_id): + return {"status": "running"} + return None + + async def fake_trigger_statuses(_workflow_id: str, _workflow_json: dict[str, Any]) -> list[dict[str, Any]]: + return [ + {"triggerId": "webhook-default", "state": "ready"}, + {"triggerId": "schedule-default", "state": "running"}, + ] + + monkeypatch.setattr(workflow_routes.WorkflowStore, "kv_get", fake_kv_get) + workflow_routes._workflow_api_health_cache[workflow_id] = True + monkeypatch.setattr( + workflow_routes, + "default_trigger_runtime", + SimpleNamespace(get_workflow_trigger_statuses=fake_trigger_statuses), + ) + + status = await workflow_routes._get_workflow_integration_status(workflow_id, workflow_data) + + assert status.model_dump() == { + "api": {"configured": True, "state": "running"}, + "trigger": { + "configured": True, + "state": "running", + "count": 2, + "items": [ + { + "id": "schedule-default", + "type": "schedule", + "name": None, + "state": "running", + "rawState": "running", + }, + { + "id": "webhook-default", + "type": "webhook", + "name": None, + "state": "running", + "rawState": "ready", + }, + ], + }, + } + + +@pytest.mark.asyncio +async def test_workflow_integration_status_marks_unhealthy_api_as_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + workflow_id = "wf-unhealthy" + workflow_data = { + "id": workflow_id, + "workflowJson": { + "start": "n1", + "nodes": [{"id": "n1", "type": "python"}], + "edges": [], + "triggers": [], + }, + } + + async def fake_kv_get(key: Any, *_args: Any, **_kwargs: Any) -> Any: + if str(key) == workflow_routes._api_service_key(workflow_id): + return {"workflowId": workflow_id, "status": "running"} + if str(key) == workflow_routes._runtime_key_main(workflow_id): + return {"status": "running"} + return None + + monkeypatch.setattr(workflow_routes.WorkflowStore, "kv_get", fake_kv_get) + workflow_routes._workflow_api_health_cache[workflow_id] = False + + status = await workflow_routes._get_workflow_integration_status(workflow_id, workflow_data) + + assert status.api.model_dump() == {"configured": True, "state": "error"} + + +@pytest.mark.asyncio +async def test_workflow_integration_status_marks_stopped_and_failed_red_states( + monkeypatch: pytest.MonkeyPatch, +) -> None: + workflow_id = "wf-stopped" + workflow_data = { + "id": workflow_id, + "workflowJson": { + "start": "n1", + "nodes": [{"id": "n1", "type": "python"}], + "edges": [], + "triggers": [ + {"id": "schedule-default", "type": "schedule", "enabled": False}, + {"id": "kafka-default", "type": "kafka", "enabled": True}, + ], + }, + } + + async def fake_kv_get(key: Any, *_args: Any, **_kwargs: Any) -> Any: + if str(key) == workflow_routes._api_service_key(workflow_id): + return {"workflowId": workflow_id, "status": "stopped", "stoppedAt": 1} + return None + + async def fake_trigger_statuses(_workflow_id: str, _workflow_json: dict[str, Any]) -> list[dict[str, Any]]: + return [ + {"triggerId": "schedule-default", "state": "stopped"}, + {"triggerId": "kafka-default", "state": "failed"}, + ] + + monkeypatch.setattr(workflow_routes.WorkflowStore, "kv_get", fake_kv_get) + monkeypatch.setattr( + workflow_routes, + "default_trigger_runtime", + SimpleNamespace(get_workflow_trigger_statuses=fake_trigger_statuses), + ) + + status = await workflow_routes._get_workflow_integration_status(workflow_id, workflow_data) + + assert status.model_dump() == { + "api": {"configured": True, "state": "stopped"}, + "trigger": { + "configured": True, + "state": "error", + "count": 2, + "items": [ + { + "id": "schedule-default", + "type": "schedule", + "name": None, + "state": "stopped", + "rawState": "stopped", + }, + { + "id": "kafka-default", + "type": "kafka", + "name": None, + "state": "error", + "rawState": "failed", + }, + ], + }, + } diff --git a/tests/server/routes/test_workflow_publish_api.py b/tests/server/routes/test_workflow_publish_api.py index a715239f4..483298457 100644 --- a/tests/server/routes/test_workflow_publish_api.py +++ b/tests/server/routes/test_workflow_publish_api.py @@ -39,7 +39,14 @@ async def test_publish_workflow_as_api_reuses_key_for_runtime( async def fake_read(key: Any, *_args: Any, **_kwargs: Any) -> Any: if str(key) == workflow_routes._api_service_key(workflow_id): - return {"apiKey": existing_key} + return {"apiKey": existing_key, "serviceUrl": "http://127.0.0.1:19007"} + if str(key) == f"{workflow_routes._REGISTRY_PREFIX_MAIN}{workflow_id}": + return { + "workflowId": workflow_id, + "publishStatus": "active", + "servicePort": 19007, + "serviceUrl": "http://127.0.0.1:19007", + } return None async def fake_write(key: Any, value: Any) -> None: @@ -50,6 +57,7 @@ async def fake_publish_workflow( image: str | None = None, driver: str | None = None, api_key: str | None = None, + port: int | None = None, ) -> dict[str, Any]: publish_calls.append( { @@ -57,6 +65,7 @@ async def fake_publish_workflow( "image": image, "driver": driver, "api_key": api_key, + "port": port, } ) return { @@ -64,6 +73,7 @@ async def fake_publish_workflow( "containerName": "local-wf-1", "driver": driver or "local", "apiKey": api_key, + "hostPort": port, } monkeypatch.setattr(workflow_routes.WorkflowStore, "kv_get", fake_read) @@ -72,7 +82,7 @@ async def fake_publish_workflow( result = await workflow_routes.publish_workflow_as_api( workflow_id, - workflow_routes.WorkflowCenterPublishRequest(driver="local"), + workflow_routes.WorkflowCenterPublishRequest(driver="local", port=25000), ) assert publish_calls == [ @@ -81,10 +91,105 @@ async def fake_publish_workflow( "image": None, "driver": "local", "api_key": existing_key, + "port": 25000, } ] assert result["apiKey"] == existing_key + assert result["port"] == 25000 assert writes[workflow_routes._api_service_key(workflow_id)]["apiKey"] == existing_key + assert writes[workflow_routes._api_service_key(workflow_id)]["port"] == 25000 + registry = writes[f"{workflow_routes._REGISTRY_PREFIX_MAIN}{workflow_id}"] + assert registry["publishStatus"] == "active" + assert registry["servicePort"] == 19007 + + +@pytest.mark.asyncio +async def test_publish_workflow_as_api_returns_conflict_for_unavailable_port( + monkeypatch: pytest.MonkeyPatch, +) -> None: + async def fake_prepare_registry(_workflow_id: str) -> tuple[dict[str, Any], int]: + return {"name": "Demo Workflow"}, 123 + + workflow_id = "wf-1" + service = { + "workflowId": workflow_id, + "status": "running", + "apiKey": "existing-key", + "port": 25000, + } + writes: list[tuple[Any, Any]] = [] + + async def fake_read(key: Any, *_args: Any, **_kwargs: Any) -> dict[str, Any]: + if str(key) == workflow_routes._api_service_key(workflow_id): + return service + if str(key) == workflow_routes._runtime_key_main(workflow_id): + return {"workflowId": workflow_id, "hostPort": 25000} + return {} + + async def fake_write(key: Any, value: Any) -> None: + writes.append((key, value)) + + async def fake_publish_workflow(*_args: Any, **_kwargs: Any) -> dict[str, Any]: + raise workflow_routes.WorkflowPortUnavailableError( + "Workflow service port 25000 is already reserved or in use" + ) + + monkeypatch.setattr(workflow_routes, "_prepare_workflow_api_registry", fake_prepare_registry) + monkeypatch.setattr(workflow_routes.WorkflowStore, "kv_get", fake_read) + monkeypatch.setattr(workflow_routes.WorkflowStore, "kv_put", fake_write) + monkeypatch.setattr(workflow_routes, "publish_workflow", fake_publish_workflow) + + with pytest.raises(workflow_routes.HTTPException) as exc_info: + await workflow_routes.publish_workflow_as_api( + workflow_id, + workflow_routes.WorkflowCenterPublishRequest(driver="local", port=25000), + ) + + assert exc_info.value.status_code == 409 + assert "25000" in str(exc_info.value.detail) + assert writes == [] + + +@pytest.mark.asyncio +async def test_publish_workflow_as_api_marks_failed_restart_when_runtime_was_stopped( + monkeypatch: pytest.MonkeyPatch, +) -> None: + workflow_id = "wf-failed-restart" + service_key = workflow_routes._api_service_key(workflow_id) + store: dict[str, Any] = { + service_key: { + "workflowId": workflow_id, + "status": "running", + "apiKey": "existing-key", + "port": 25000, + } + } + + async def fake_prepare_registry(_workflow_id: str) -> tuple[dict[str, Any], int]: + return {"name": "Demo Workflow"}, 123 + + async def fake_read(key: Any, *_args: Any, **_kwargs: Any) -> Any: + return store.get(str(key)) + + async def fake_write(key: Any, value: Any) -> None: + store[str(key)] = value + + async def fake_publish_workflow(*_args: Any, **_kwargs: Any) -> dict[str, Any]: + raise workflow_routes.WorkflowCenterError("replacement failed health check") + + monkeypatch.setattr(workflow_routes, "_prepare_workflow_api_registry", fake_prepare_registry) + monkeypatch.setattr(workflow_routes.WorkflowStore, "kv_get", fake_read) + monkeypatch.setattr(workflow_routes.WorkflowStore, "kv_put", fake_write) + monkeypatch.setattr(workflow_routes, "publish_workflow", fake_publish_workflow) + + with pytest.raises(workflow_routes.HTTPException) as exc_info: + await workflow_routes.publish_workflow_as_api(workflow_id) + + assert exc_info.value.status_code == 500 + assert store[service_key]["status"] == "error" + assert store[service_key]["port"] == 25000 + assert store[service_key]["health"] == {"ok": False, "reason": "publish_failed"} + assert "replacement failed" in store[service_key]["lastStartError"] @pytest.mark.asyncio @@ -131,6 +236,7 @@ async def fake_publish_workflow( image: str | None = None, driver: str | None = None, api_key: str | None = None, + port: int | None = None, ) -> dict[str, Any]: publish_calls.append( { @@ -138,14 +244,16 @@ async def fake_publish_workflow( "image": image, "driver": driver, "api_key": api_key, + "port": port, } ) return { - "serviceUrl": "http://127.0.0.1:19001", + "serviceUrl": "http://127.0.0.1:19000", "containerName": "flocks-wf-wf-1-rel-1", "driver": driver, "image": image, "apiKey": api_key, + "hostPort": port, } monkeypatch.setattr(workflow_routes.WorkflowStore, "kv_list_keys", fake_list_keys) @@ -165,12 +273,14 @@ async def fake_publish_workflow( "image": "custom-image:latest", "driver": "docker", "api_key": existing_key, + "port": 19000, } ] assert store[service_key]["status"] == "running" assert store[service_key]["apiKey"] == existing_key - assert store[service_key]["serviceUrl"] == "http://127.0.0.1:19001" - assert store[service_key]["invokeUrl"] == "http://127.0.0.1:19001/invoke" + assert store[service_key]["serviceUrl"] == "http://127.0.0.1:19000" + assert store[service_key]["invokeUrl"] == "http://127.0.0.1:19000/invoke" + assert store[service_key]["port"] == 19000 @pytest.mark.asyncio @@ -211,6 +321,7 @@ async def fake_publish_workflow( image: str | None = None, driver: str | None = None, api_key: str | None = None, + port: int | None = None, ) -> dict[str, Any]: publish_calls.append(requested_id) return { @@ -275,7 +386,7 @@ async def fake_health(requested_id: str) -> dict[str, Any]: @pytest.mark.asyncio -async def test_get_workflow_service_does_not_probe_runtime_health( +async def test_get_workflow_service_normalizes_stale_state_without_health_probe( monkeypatch: pytest.MonkeyPatch, ) -> None: workflow_id = "wf-service-read" @@ -288,8 +399,11 @@ async def test_get_workflow_service_does_not_probe_runtime_health( health_calls: list[str] = [] async def fake_read(key: Any, *_args: Any, **_kwargs: Any) -> Any: - assert str(key) == workflow_routes._api_service_key(workflow_id) - return service + if str(key) == workflow_routes._api_service_key(workflow_id): + return service + if str(key) == workflow_routes._runtime_key_main(workflow_id): + return None + raise AssertionError(f"Unexpected storage key: {key}") async def fake_write(key: Any, value: Any) -> None: writes.append((key, value)) @@ -304,7 +418,9 @@ async def fake_health(requested_id: str) -> dict[str, Any]: result = await workflow_routes.get_workflow_service(workflow_id) - assert result is service + assert result["status"] == "stopped" + assert result["health"] == {"ok": False, "stale": True, "reason": "missing_runtime"} + assert service["status"] == "running" assert health_calls == [] assert writes == [] @@ -354,3 +470,42 @@ async def fake_write(key: Any, value: Any) -> None: assert "stoppedAt" not in result[0] assert store[service_key]["status"] == "running" assert writes == [] + + +@pytest.mark.asyncio +async def test_delete_workflow_service_clears_persisted_port(monkeypatch: pytest.MonkeyPatch) -> None: + workflow_id = "wf-delete-service" + service_key = workflow_routes._api_service_key(workflow_id) + registry_key = f"{workflow_routes._REGISTRY_PREFIX_MAIN}{workflow_id}" + store: dict[str, Any] = { + service_key: {"workflowId": workflow_id, "status": "stopped", "port": 25000}, + registry_key: { + "workflowId": workflow_id, + "publishStatus": "stopped", + "servicePort": 25000, + }, + } + + async def fake_read(key: Any, *_args: Any, **_kwargs: Any) -> Any: + return store.get(str(key)) + + async def fake_write(key: Any, value: Any) -> None: + store[str(key)] = value + + async def fake_remove(key: Any) -> bool: + store.pop(str(key), None) + return True + + async def fake_stop(_workflow_id: str) -> dict[str, Any]: + return {"status": "stopped"} + + monkeypatch.setattr(workflow_routes.WorkflowStore, "kv_get", fake_read) + monkeypatch.setattr(workflow_routes.WorkflowStore, "kv_put", fake_write) + monkeypatch.setattr(workflow_routes.WorkflowStore, "kv_remove", fake_remove) + monkeypatch.setattr(workflow_routes, "stop_workflow_service", fake_stop) + + result = await workflow_routes.delete_workflow_service(workflow_id) + + assert result == {"ok": True, "workflowId": workflow_id} + assert service_key not in store + assert "servicePort" not in store[registry_key] diff --git a/tests/server/routes/test_workflow_summary_routes.py b/tests/server/routes/test_workflow_summary_routes.py new file mode 100644 index 000000000..d8d3bb373 --- /dev/null +++ b/tests/server/routes/test_workflow_summary_routes.py @@ -0,0 +1,138 @@ +from __future__ import annotations + +import asyncio +from typing import Any, Awaitable, Callable + +import pytest + +from flocks.server.routes import workflow as workflow_routes + + +def _workflow_data(workflow_id: str) -> dict[str, Any]: + return { + "id": workflow_id, + "name": workflow_id, + "category": "default", + "workflowJson": {"start": "n1", "nodes": [], "edges": []}, + "status": "active", + "source": "global", + "createdAt": 1, + "updatedAt": 2, + } + + +async def _noop_migrate() -> None: + return None + + +async def _zero_stats(_workflow_id: str) -> dict[str, Any]: + return {"callCount": 0} + + +async def _unconfigured_status( + _workflow_id: str, + _workflow_data: dict[str, Any], +) -> workflow_routes.WorkflowIntegrationStatusResponse: + return workflow_routes.WorkflowIntegrationStatusResponse( + api=workflow_routes.WorkflowCapabilityStatusResponse(configured=False, state="unconfigured"), + trigger=workflow_routes.WorkflowTriggerCapabilityStatusResponse( + configured=False, + state="unconfigured", + ), + ) + + +def _patch_list_dependencies( + monkeypatch: pytest.MonkeyPatch, + workflows: list[dict[str, Any]], + *, + stats: Callable[[str], Awaitable[dict[str, Any]]] = _zero_stats, +) -> None: + monkeypatch.setattr(workflow_routes, "_migrate_storage_to_filesystem", _noop_migrate) + monkeypatch.setattr(workflow_routes, "_list_workflows_from_fs", lambda: workflows) + monkeypatch.setattr(workflow_routes, "_get_workflow_stats", stats) + monkeypatch.setattr(workflow_routes, "_get_workflow_integration_status", _unconfigured_status) + + +@pytest.mark.asyncio +async def test_list_workflow_summaries_includes_integration_status_without_full_definition( + monkeypatch: pytest.MonkeyPatch, +) -> None: + workflow = _workflow_data("wf-list") + workflow["name"] = "Listed Workflow" + _patch_list_dependencies(monkeypatch, [workflow]) + + async def running_status( + _workflow_id: str, + _workflow_data: dict[str, Any], + ) -> workflow_routes.WorkflowIntegrationStatusResponse: + return workflow_routes.WorkflowIntegrationStatusResponse( + api=workflow_routes.WorkflowCapabilityStatusResponse(configured=True, state="running"), + trigger=workflow_routes.WorkflowTriggerCapabilityStatusResponse( + configured=False, + state="unconfigured", + ), + ) + + monkeypatch.setattr(workflow_routes, "_get_workflow_integration_status", running_status) + + items = await workflow_routes.list_workflow_summaries(category=None, status=None, exclude_id=None) + + assert items[0].integrationStatus.model_dump() == { + "api": {"configured": True, "state": "running"}, + "trigger": {"configured": False, "state": "unconfigured", "count": 0, "items": []}, + } + assert items[0].nodeCount == 0 + assert not hasattr(items[0], "workflowJson") + + +@pytest.mark.asyncio +async def test_list_workflow_summaries_enriches_multiple_items_concurrently( + monkeypatch: pytest.MonkeyPatch, +) -> None: + both_started = asyncio.Event() + release = asyncio.Event() + started = 0 + + async def blocking_stats(_workflow_id: str) -> dict[str, Any]: + nonlocal started + started += 1 + if started == 2: + both_started.set() + await release.wait() + return {"callCount": 0} + + _patch_list_dependencies( + monkeypatch, + [_workflow_data("wf-one"), _workflow_data("wf-two")], + stats=blocking_stats, + ) + + list_task = asyncio.create_task( + workflow_routes.list_workflow_summaries(category=None, status=None, exclude_id=None) + ) + await asyncio.wait_for(both_started.wait(), timeout=1) + release.set() + + items = await list_task + + assert {item.id for item in items} == {"wf-one", "wf-two"} + + +@pytest.mark.asyncio +async def test_full_workflow_list_does_not_load_integration_status( + monkeypatch: pytest.MonkeyPatch, +) -> None: + workflow = _workflow_data("wf-full") + workflow["name"] = "Full Workflow" + _patch_list_dependencies(monkeypatch, [workflow]) + + async def unexpected_integration_status(*_args: Any, **_kwargs: Any) -> None: + raise AssertionError("full workflow list must not load card integration status") + + monkeypatch.setattr(workflow_routes, "_get_workflow_integration_status", unexpected_integration_status) + + items = await workflow_routes.list_workflows(category=None, status=None, exclude_id=None) + + assert items[0].id == "wf-full" + assert items[0].integrationStatus is None diff --git a/tests/server/test_app_errors.py b/tests/server/test_app_errors.py index 673cd89a1..71e97b879 100644 --- a/tests/server/test_app_errors.py +++ b/tests/server/test_app_errors.py @@ -1,6 +1,8 @@ import json +from unittest.mock import MagicMock import pytest +from starlette.exceptions import HTTPException from starlette.requests import Request from flocks.server import app as app_module @@ -31,3 +33,43 @@ async def test_general_exception_response_does_not_expose_traceback(): "error": "InternalServerError", "message": "Internal server error", } + + +@pytest.mark.asyncio +async def test_http_4xx_logs_warning(monkeypatch): + warning = MagicMock() + error = MagicMock() + monkeypatch.setattr(app_module.log, "warn", warning) + monkeypatch.setattr(app_module.log, "error", error) + + response = await app_module.http_exception_handler( + _request("/missing"), + HTTPException(status_code=404, detail="Not found"), + ) + + assert response.status_code == 404 + warning.assert_called_once_with( + "http.error", + {"path": "/missing", "status": 404, "detail": "Not found"}, + ) + error.assert_not_called() + + +@pytest.mark.asyncio +async def test_http_5xx_logs_error(monkeypatch): + warning = MagicMock() + error = MagicMock() + monkeypatch.setattr(app_module.log, "warn", warning) + monkeypatch.setattr(app_module.log, "error", error) + + response = await app_module.http_exception_handler( + _request("/unavailable"), + HTTPException(status_code=503, detail="Unavailable"), + ) + + assert response.status_code == 503 + error.assert_called_once_with( + "http.error", + {"path": "/unavailable", "status": 503, "detail": "Unavailable"}, + ) + warning.assert_not_called() diff --git a/tests/server/test_asgi_middleware.py b/tests/server/test_asgi_middleware.py index 227545326..3644f8557 100644 --- a/tests/server/test_asgi_middleware.py +++ b/tests/server/test_asgi_middleware.py @@ -382,3 +382,40 @@ async def send(_message): await asyncio.gather(*requests, return_exceptions=True) assert len(cleared_tokens) == 11 + + +@pytest.mark.asyncio +async def test_instance_context_without_directory_uses_process_cwd( + monkeypatch: pytest.MonkeyPatch, + tmp_path, +) -> None: + """Ordinary requests keep the pre-project-management cwd behavior.""" + from flocks.project.instance import Instance + + process_cwd = tmp_path / "server-cwd" + process_cwd.mkdir() + monkeypatch.chdir(process_cwd) + provided_directories: list[str] = [] + + async def provide(*, directory, init, fn): + provided_directories.append(directory) + return await fn() + + async def endpoint(_scope, _receive, _send) -> None: + return None + + async def receive(): + return {"type": "http.request", "body": b"", "more_body": False} + + async def send(_message): + return None + + monkeypatch.setattr(Instance, "provide", provide) + + await server_app._InstanceContextMiddleware(endpoint)( + _http_scope("/api/skill"), + receive, + send, + ) + + assert provided_directories == [str(process_cwd)] diff --git a/tests/server/test_input_dispatcher.py b/tests/server/test_input_dispatcher.py index ff1c8bae6..17a11d8f0 100644 --- a/tests/server/test_input_dispatcher.py +++ b/tests/server/test_input_dispatcher.py @@ -357,13 +357,13 @@ def test_materialize_queued_data_url_returns_readable_file_uri(self, monkeypatch from flocks.server.routes import session as session_routes from flocks.session.utils.file_extractor import read_file_part_bytes - class FakeWorkspace: - def resolve_workspace_path(self, rel_path: str): - return tmp_path / rel_path - monkeypatch.setattr( "flocks.workspace.manager.WorkspaceManager.get_instance", - lambda: FakeWorkspace(), + lambda: (_ for _ in ()).throw(PermissionError("workspace access denied")), + ) + monkeypatch.setattr( + "flocks.config.config.Config.get_data_path", + staticmethod(lambda: tmp_path), ) data_url = "data:image/png;base64," + base64.b64encode(b"png-bytes").decode() @@ -377,6 +377,17 @@ def resolve_workspace_path(self, rel_path: str): assert url.startswith("file://") assert read_file_part_bytes(url) == b"png-bytes" + def test_session_upload_path_rejects_traversal(self, monkeypatch, tmp_path): + from flocks.server.routes import session as session_routes + + monkeypatch.setattr( + "flocks.config.config.Config.get_data_path", + staticmethod(lambda: tmp_path), + ) + + with pytest.raises(ValueError, match="Invalid session ID"): + session_routes._session_uploads_dir("../outside") + @pytest.mark.asyncio async def test_prompt_async_queues_when_session_running_without_creating_message(self, monkeypatch): from flocks.server.routes import session as session_routes diff --git a/tests/server/test_lifespan.py b/tests/server/test_lifespan.py index de6677698..14e8741f4 100644 --- a/tests/server/test_lifespan.py +++ b/tests/server/test_lifespan.py @@ -19,7 +19,7 @@ def warn(self, *_args, **_kwargs) -> None: @pytest.mark.asyncio -async def test_lifespan_cleans_leftovers_before_recovering_upgrade_state( +async def test_lifespan_cleans_replaced_files_without_upgrade_recovery( monkeypatch: pytest.MonkeyPatch, ) -> None: events: list[str] = [] @@ -27,9 +27,6 @@ async def test_lifespan_cleans_leftovers_before_recovering_upgrade_state( def cleanup_replaced_files() -> None: events.append("cleanup_replaced_files") - def recover_upgrade_state() -> None: - events.append("recover_upgrade_state") - async def fake_storage_init() -> None: return None @@ -56,7 +53,6 @@ async def fake_async_noop(*_args, **_kwargs) -> None: "flocks.updater.updater", types.SimpleNamespace( cleanup_replaced_files=cleanup_replaced_files, - recover_upgrade_state=recover_upgrade_state, ), ) monkeypatch.setitem( @@ -149,4 +145,4 @@ async def fake_async_noop(*_args, **_kwargs) -> None: async with app_module.lifespan(SimpleNamespace()): pass - assert events == ["cleanup_replaced_files", "recover_upgrade_state"] + assert events == ["cleanup_replaced_files"] diff --git a/tests/server/test_tool_setting_routes.py b/tests/server/test_tool_setting_routes.py index 023360fe9..8ee244019 100644 --- a/tests/server/test_tool_setting_routes.py +++ b/tests/server/test_tool_setting_routes.py @@ -34,7 +34,10 @@ # ─── Fixtures ──────────────────────────────────────────────────────────────── -def _stub_api_tool(name: str, *, enabled: bool, provider: str = "onesec_api") -> Tool: +_TEST_SERVICE_ID = "test_api_service" + + +def _stub_api_tool(name: str, *, enabled: bool, provider: str = _TEST_SERVICE_ID) -> Tool: async def handler(ctx: ToolContext, value: str = "ok") -> ToolResult: return ToolResult(success=True, output=value) @@ -45,6 +48,7 @@ async def handler(ctx: ToolContext, value: str = "ok") -> ToolResult: category=ToolCategory.CUSTOM, enabled=enabled, provider=provider, + source="api", ), handler=handler, ) @@ -58,7 +62,7 @@ def tool_client(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): overlay/service-gate combinations without touching real plugin YAML. """ config_dir = tmp_path / ".flocks" / "config" - config_dir.mkdir(parents=True) + config_dir.mkdir(parents=True, exist_ok=True) monkeypatch.setenv("FLOCKS_CONFIG_DIR", str(config_dir)) from flocks.config.config import Config @@ -68,6 +72,9 @@ def tool_client(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): saved_tools = dict(ToolRegistry._tools) saved_defaults = dict(ToolRegistry._enabled_defaults) + saved_failure_state = dict(ToolRegistry._failure_state) + saved_revision = ToolRegistry._revision + saved_config_state_token = ToolRegistry._config_state_token saved_initialized = ToolRegistry._initialized enabled_tool = _stub_api_tool("onesec_dns_test", enabled=True) @@ -80,6 +87,7 @@ def tool_client(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): enabled_tool.info.name: True, disabled_tool.info.name: False, } + ToolRegistry._failure_state = {} # Skip plugin discovery — our stub registry is enough. ToolRegistry._initialized = True @@ -100,10 +108,13 @@ def tool_client(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): ToolRegistry._tools = saved_tools ToolRegistry._enabled_defaults = saved_defaults + ToolRegistry._failure_state = saved_failure_state + ToolRegistry._revision = saved_revision + ToolRegistry._config_state_token = saved_config_state_token ToolRegistry._initialized = saved_initialized -def _set_service(*, enabled: bool, sid: str = "onesec_api") -> None: +def _set_service(*, enabled: bool, sid: str = _TEST_SERVICE_ID) -> None: from flocks.config.config_writer import ConfigWriter ConfigWriter.set_api_service(sid, { "apiKey": "{secret:test_key}", @@ -192,6 +203,30 @@ def test_reload_path_refreshes_enabled_default(self, tool_client): class TestUpdateTool: + def test_manual_reenable_clears_repeated_failure_count(self, tool_client): + client, enabled_tool, _ = tool_client + _set_service(enabled=True) + revision_before = ToolRegistry.revision() + enabled_tool.info.enabled = False + ToolRegistry._failure_state[enabled_tool.info.name] = { + "key": "same-failure", + "count": ToolRegistry._failure_disable_threshold, + } + ToolRegistry._failure_state[enabled_tool.info.name] = { + "key": "same-failure", + "count": ToolRegistry._failure_disable_threshold, + } + + res = client.patch( + f"/api/tools/{enabled_tool.info.name}", + json={"enabled": True}, + ) + + assert res.status_code == 200 + assert res.json()["enabled"] is True + assert enabled_tool.info.name not in ToolRegistry._failure_state + assert ToolRegistry.revision() == revision_before + 1 + def test_overlay_persisted_when_differs_from_default(self, tool_client): client, _, disabled_tool = tool_client _set_service(enabled=True) @@ -312,7 +347,75 @@ def test_device_disable_writes_false_override_only( delete_override.assert_not_awaited() +class TestApiServiceToolSync: + def test_service_enable_does_not_resurrect_disabled_tool_overlay( + self, + tool_client, + monkeypatch: pytest.MonkeyPatch, + ): + client, enabled_tool, _ = tool_client + _set_service(enabled=True) + client.patch(f"/api/tools/{enabled_tool.info.name}", json={"enabled": False}) + assert enabled_tool.info.enabled is False + + from flocks.server.routes import provider as provider_routes + + monkeypatch.setattr( + provider_routes, + "_get_api_service_tool_infos", + lambda _provider_id: [enabled_tool.info, tool_client[2].info], + ) + + matched = provider_routes._set_api_service_tools_enabled(_TEST_SERVICE_ID, True) + + assert matched == 2 + assert enabled_tool.info.enabled is False + assert _read_settings() == {enabled_tool.info.name: {"enabled": False}} + + def test_service_state_change_bumps_tool_revision( + self, + tool_client, + monkeypatch: pytest.MonkeyPatch, + ): + _, enabled_tool, disabled_tool = tool_client + _set_service(enabled=True) + revision_before = ToolRegistry.revision() + enabled_tool.info.enabled = False + + from flocks.server.routes import provider as provider_routes + + monkeypatch.setattr( + provider_routes, + "_get_api_service_tool_infos", + lambda _provider_id: [enabled_tool.info, disabled_tool.info], + ) + + provider_routes._set_api_service_tools_enabled(_TEST_SERVICE_ID, True) + + assert enabled_tool.info.enabled is True + assert disabled_tool.info.enabled is False + assert enabled_tool.info.name not in ToolRegistry._failure_state + assert ToolRegistry.revision() == revision_before + 1 + + class TestResetToolSetting: + def test_reset_reenable_clears_failure_count_and_bumps_revision(self, tool_client): + client, enabled_tool, _ = tool_client + _set_service(enabled=True) + client.patch(f"/api/tools/{enabled_tool.info.name}", json={"enabled": False}) + ToolRegistry._failure_state[enabled_tool.info.name] = { + "key": "same-failure", + "count": ToolRegistry._failure_disable_threshold, + } + revision_before = ToolRegistry.revision() + + res = client.post(f"/api/tools/{enabled_tool.info.name}/reset") + + assert res.status_code == 200 + assert res.json()["enabled"] is True + assert enabled_tool.info.name not in ToolRegistry._failure_state + assert ToolRegistry.revision() == revision_before + 1 + def test_reset_restores_default_and_removes_overlay(self, tool_client): client, _, disabled_tool = tool_client _set_service(enabled=True) diff --git a/tests/session/test_prompt_tokens.py b/tests/session/test_prompt_tokens.py index 66acdbc7e..fddedf832 100644 --- a/tests/session/test_prompt_tokens.py +++ b/tests/session/test_prompt_tokens.py @@ -196,7 +196,15 @@ async def test_returns_list(self): async def test_includes_working_directory(self): result = await SystemPrompt.environment("/my/work/dir") combined = "\n".join(result) - assert "/my/work/dir" in combined + assert "current working directory: /my/work/dir" in combined + + def test_distinguishes_source_code_and_working_directories(self, monkeypatch): + monkeypatch.setenv("FLOCKS_REPO_ROOT", "/opt/flocks") + + combined = "\n".join(SystemPrompt.environment_stable("/workspace/project")) + + assert "flocks source code directory: /opt/flocks" in combined + assert "current working directory: /workspace/project" in combined @pytest.mark.asyncio async def test_includes_date_info(self): diff --git a/tests/session/test_runner_step.py b/tests/session/test_runner_step.py index 9d67e6565..6109a1aaa 100644 --- a/tests/session/test_runner_step.py +++ b/tests/session/test_runner_step.py @@ -10,6 +10,8 @@ - SessionRunner construction and abort behavior (from existing tests) """ +import httpcore +import httpx import pytest from types import SimpleNamespace from unittest.mock import MagicMock, patch, AsyncMock @@ -270,6 +272,51 @@ def test_incomplete_chunked_read_exception_is_retryable_connection_error(self): assert result["data"]["isRetryable"] is True assert result["data"]["displayMessage"] == runner_mod.CONNECTION_ERROR_DISPLAY_MESSAGE + @pytest.mark.parametrize( + "exc", + [ + httpx.ReadError("", request=httpx.Request("GET", "https://example.test")), + httpx.RemoteProtocolError( + "", + request=httpx.Request("GET", "https://example.test"), + ), + httpcore.ReadError(), + httpcore.RemoteProtocolError(), + ], + ) + def test_empty_transport_exception_is_retryable_connection_error(self, exc): + runner = _make_runner() + + result = runner._exception_to_error_dict(exc) + + assert result["name"] == "APIError" + assert result["data"]["isRetryable"] is True + assert result["data"]["isConnectionError"] is True + assert result["data"]["displayMessage"] == runner_mod.CONNECTION_ERROR_DISPLAY_MESSAGE + assert result["data"]["exceptionType"] == type(exc).__name__ + + def test_wrapped_transport_exception_is_retryable_connection_error(self): + runner = _make_runner() + transport_error = httpcore.ReadError() + exc = RuntimeError("stream failed") + exc.__cause__ = transport_error + + result = runner._exception_to_error_dict(exc) + + assert result["name"] == "APIError" + assert result["data"]["isRetryable"] is True + assert result["data"]["isConnectionError"] is True + assert result["data"]["exceptionType"] == "RuntimeError" + assert result["data"]["transportExceptionType"] == "ReadError" + + def test_empty_non_transport_exception_is_not_retryable(self): + runner = _make_runner() + + result = runner._exception_to_error_dict(RuntimeError()) + + assert result["name"] == "RuntimeError" + assert result["data"].get("isRetryable") is not True + def test_exception_with_status_code_429(self): runner = _make_runner() exc = Exception("Rate limited") @@ -2101,7 +2148,7 @@ async def test_to_chat_messages_skips_reasoning_only_aborted_assistant(monkeypat def test_provider_capability_key_includes_interleaved_policy(monkeypatch): runner = _make_runner("ses_runner_interleaved_capability_key") runner.provider_id = "deepseek" - runner.model_id = "deepseek-reasoner" + runner.model_id = "deepseek-v4-pro" monkeypatch.setattr(SessionRunner, "_model_supports_vision", lambda self: False) monkeypatch.setattr( @@ -2249,6 +2296,8 @@ async def test_process_step_limits_connection_error_retries(monkeypatch): provider.is_configured.return_value = True assistant_msg = SimpleNamespace(id="msg_assistant_connection_error") update_mock = AsyncMock(return_value=None) + warn_log = MagicMock() + error_log = MagicMock() call_count = 0 async def fake_call_llm(*_args, **_kwargs): @@ -2271,6 +2320,8 @@ async def fake_call_llm(*_args, **_kwargs): monkeypatch.setattr(runner_mod.Message, "update", update_mock) monkeypatch.setattr(runner_mod.SessionRetry, "sleep", AsyncMock(return_value=None)) monkeypatch.setattr(runner, "_call_llm", fake_call_llm) + monkeypatch.setattr(runner_mod.log, "warn", warn_log) + monkeypatch.setattr(runner_mod.log, "error", error_log) result = await runner._process_step([last_user], last_user) @@ -2284,6 +2335,21 @@ async def fake_call_llm(*_args, **_kwargs): assert final_update["error"]["data"]["message"] == "Connection error." assert final_update["error"]["data"]["displayMessage"] == runner_mod.CONNECTION_ERROR_DISPLAY_MESSAGE + retry_logs = [ + call for call in warn_log.call_args_list + if call.args and call.args[0] == "runner.step.retry" + ] + max_retry_logs = [ + call for call in error_log.call_args_list + if call.args and call.args[0] == "runner.step.max_retries_exceeded" + ] + assert len(retry_logs) == 3 + assert len(max_retry_logs) == 1 + assert not any( + call.args and call.args[0] == "runner.step.error" + for call in [*warn_log.call_args_list, *error_log.call_args_list] + ) + @pytest.mark.asyncio async def test_process_step_marks_aborted_llm_message_as_error(monkeypatch): @@ -2419,6 +2485,167 @@ async def chat_stream(self, **kwargs): # noqa: ANN003 run_after_mock.assert_not_awaited() +@pytest.mark.asyncio +async def test_process_step_persists_visible_error_when_provider_missing(monkeypatch): + runner = _make_runner("ses_runner_missing_provider_error") + user = await Message.create( + runner.session.id, + MessageRole.USER, + "hello", + agent="rex", + model={"providerID": "missing-provider", "modelID": "missing-model"}, + ) + messages = await Message.list(runner.session.id) + agent = SimpleNamespace(name="rex", tools=[], steps=5, prompt=None) + events = [] + callback_errors = [] + + async def publish_event(event_name, payload): + events.append((event_name, payload)) + + async def on_error(error): + callback_errors.append(error) + + monkeypatch.setattr(runner_mod.Agent, "get", AsyncMock(return_value=agent)) + monkeypatch.setattr(runner_mod.Provider, "get", staticmethod(lambda _provider_id: None)) + + runner.provider_id = "missing-provider" + runner.model_id = "missing-model" + runner.callbacks = RunnerCallbacks( + on_error=on_error, + event_publish_callback=publish_event, + ) + + result = await runner._process_step(messages, user) + messages_with_parts = await Message.list_with_parts(runner.session.id) + assistant = next(item for item in messages_with_parts if item.info.role == MessageRole.ASSISTANT) + visible_text_parts = [ + part for part in assistant.parts + if getattr(part, "type", None) == "text" and getattr(part, "text", "").strip() + ] + + assert result.action == "stop" + assert result.error == runner_mod.CONNECTION_ERROR_DISPLAY_MESSAGE + assert callback_errors == [runner_mod.CONNECTION_ERROR_DISPLAY_MESSAGE] + assert assistant.info.finish == "error" + assert assistant.info.error["name"] == "ProviderUnavailableError" + assert visible_text_parts + assert visible_text_parts[0].text == runner_mod.CONNECTION_ERROR_DISPLAY_MESSAGE + assert any(event_name == "message.part.updated" for event_name, _payload in events) + + +@pytest.mark.asyncio +async def test_process_step_persists_visible_error_when_provider_not_configured(monkeypatch): + runner = _make_runner("ses_runner_unconfigured_provider_error") + user = await Message.create( + runner.session.id, + MessageRole.USER, + "hello", + agent="rex", + model={"providerID": "unconfigured-provider", "modelID": "unconfigured-model"}, + ) + messages = await Message.list(runner.session.id) + agent = SimpleNamespace(name="rex", tools=[], steps=5, prompt=None) + events = [] + callback_errors = [] + + class UnconfiguredProvider: + def is_configured(self): + return False + + async def publish_event(event_name, payload): + events.append((event_name, payload)) + + async def on_error(error): + callback_errors.append(error) + + monkeypatch.setattr(runner_mod.Agent, "get", AsyncMock(return_value=agent)) + monkeypatch.setattr(runner_mod.Provider, "get", staticmethod(lambda _provider_id: UnconfiguredProvider())) + monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) + + runner.provider_id = "unconfigured-provider" + runner.model_id = "unconfigured-model" + runner.callbacks = RunnerCallbacks( + on_error=on_error, + event_publish_callback=publish_event, + ) + + result = await runner._process_step(messages, user) + messages_with_parts = await Message.list_with_parts(runner.session.id) + assistant = next(item for item in messages_with_parts if item.info.role == MessageRole.ASSISTANT) + visible_text_parts = [ + part for part in assistant.parts + if getattr(part, "type", None) == "text" and getattr(part, "text", "").strip() + ] + + assert result.action == "stop" + assert result.error == runner_mod.CONNECTION_ERROR_DISPLAY_MESSAGE + assert callback_errors == [runner_mod.CONNECTION_ERROR_DISPLAY_MESSAGE] + assert assistant.info.finish == "error" + assert assistant.info.error["name"] == "ProviderConfigurationError" + assert visible_text_parts + assert visible_text_parts[0].text == runner_mod.CONNECTION_ERROR_DISPLAY_MESSAGE + assert any(event_name == "message.part.updated" for event_name, _payload in events) + + +@pytest.mark.asyncio +async def test_process_step_persists_visible_error_when_model_returns_empty_stream(monkeypatch): + runner = _make_runner("ses_runner_empty_stream_error") + user = await Message.create( + runner.session.id, + MessageRole.USER, + "hello", + agent="rex", + model={"providerID": "empty-provider", "modelID": "empty-model"}, + ) + messages = await Message.list(runner.session.id) + agent = SimpleNamespace(name="rex", tools=[], steps=5, prompt=None, model=None) + events = [] + + class EmptyProvider: + def is_configured(self): + return True + + async def chat_stream(self, **kwargs): + del kwargs + if False: + yield None + + async def publish_event(event_name, payload): + events.append((event_name, payload)) + + monkeypatch.setattr(runner_mod.Agent, "get", AsyncMock(return_value=agent)) + monkeypatch.setattr(runner_mod.Provider, "get", staticmethod(lambda _provider_id: EmptyProvider())) + monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) + monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompts", AsyncMock(return_value=[])) + monkeypatch.setattr(SessionRunner, "_build_callable_tool_schema", AsyncMock(return_value=[])) + monkeypatch.setattr(runner_mod.SessionRetry, "sleep", AsyncMock(return_value=None)) + + runner.provider_id = "empty-provider" + runner.model_id = "empty-model" + runner.callbacks = RunnerCallbacks(event_publish_callback=publish_event) + + result = await runner._process_step(messages, user) + messages_with_parts = await Message.list_with_parts(runner.session.id) + assistant = next(item for item in messages_with_parts if item.info.role == MessageRole.ASSISTANT) + visible_text = "\n".join( + getattr(part, "text", "") + for part in assistant.parts + if getattr(part, "type", None) == "text" and getattr(part, "text", "").strip() + ) + + assert result.action == "stop" + assert "returned an empty response" in result.error + assert assistant.info.finish == "error" + assert assistant.info.error["name"] == "EmptyResponseError" + assert "returned an empty response" in visible_text + assert any( + event_name == "message.part.updated" + and "returned an empty response" in _payload["part"]["text"] + for event_name, _payload in events + ) + + @pytest.mark.asyncio async def test_process_step_uses_loaded_tool_schema_names_for_prompt_guidance(monkeypatch): runner = _make_runner("ses_runner_prompt_guidance_tool_names") @@ -2634,6 +2861,55 @@ async def test_process_step_empty_retry_records_usage_per_attempt(monkeypatch): record_mock.assert_any_await(second_usage, message_id=assistant_msg.id) +@pytest.mark.asyncio +async def test_process_step_retries_empty_transport_exception(monkeypatch): + runner = _make_runner("ses_runner_transport_retry") + runner.callbacks = RunnerCallbacks(on_error=AsyncMock()) + + last_user = UserMessageInfo( + id="msg_user_transport_retry", + sessionID=runner.session.id, + role="user", + time={"created": 1_000}, + agent="rex", + model={"providerID": "openai", "modelID": "gpt-5"}, + ) + agent = SimpleNamespace(name="rex", steps=None, mode="primary", prompt="", tools=[]) + provider = MagicMock() + provider.is_configured.return_value = True + assistant_msg = SimpleNamespace(id="msg_assistant_transport_retry") + call_llm_mock = AsyncMock( + side_effect=[ + httpcore.ReadError(), + StepResult(action="stop", content="recovered"), + ] + ) + sleep_mock = AsyncMock(return_value=None) + + monkeypatch.setattr(runner_mod.Agent, "get", AsyncMock(return_value=agent)) + monkeypatch.setattr(runner_mod.Provider, "get", lambda provider_id: provider) + monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) + monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompts", AsyncMock(return_value=[])) + monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) + monkeypatch.setattr( + runner, + "_to_chat_messages", + AsyncMock(return_value=[SimpleNamespace(role="user", content="hi")]), + ) + monkeypatch.setattr(runner_mod.Message, "get_text_content", AsyncMock(return_value="hi")) + monkeypatch.setattr(runner_mod.Message, "create", AsyncMock(return_value=assistant_msg)) + monkeypatch.setattr(runner_mod.Message, "update", AsyncMock(return_value=None)) + monkeypatch.setattr(runner, "_call_llm", call_llm_mock) + monkeypatch.setattr(runner_mod.SessionRetry, "sleep", sleep_mock) + + result = await runner._process_step([last_user], last_user) + + assert result.content == "recovered" + assert call_llm_mock.await_count == 2 + sleep_mock.assert_awaited_once() + runner.callbacks.on_error.assert_not_awaited() + + @pytest.mark.asyncio async def test_process_step_uses_default_max_steps_when_agent_steps_missing(monkeypatch): runner = _make_runner("ses_runner_default_max_steps") diff --git a/tests/session/test_session_loop_working_directory.py b/tests/session/test_session_loop_working_directory.py new file mode 100644 index 000000000..4c2cc5f3a --- /dev/null +++ b/tests/session/test_session_loop_working_directory.py @@ -0,0 +1,42 @@ +from unittest.mock import AsyncMock + +import pytest + +from flocks.bus.bus import Bus +from flocks.session.message import Message +from flocks.session.session import Session, SessionInfo +from flocks.session.session_loop import LoopResult, SessionLoop + + +@pytest.mark.asyncio +async def test_run_uses_runtime_working_directory(monkeypatch: pytest.MonkeyPatch): + session = SessionInfo( + id="ses_runtime_directory", + projectID="legacy-project", + directory="/missing/original", + title="Legacy session", + ) + run_loop = AsyncMock(return_value=LoopResult(action="stop")) + + monkeypatch.setattr(Session, "get_by_id", AsyncMock(return_value=session)) + monkeypatch.setattr(Session, "touch", AsyncMock()) + monkeypatch.setattr(Message, "list", AsyncMock(return_value=[])) + monkeypatch.setattr(SessionLoop, "_run_loop", run_loop) + monkeypatch.setattr( + "flocks.session.orphan_tools.abort_orphan_running_parts", + AsyncMock(), + ) + monkeypatch.setattr(Bus, "publish", AsyncMock()) + + result = await SessionLoop.run( + session.id, + provider_id="test-provider", + model_id="test-model", + working_directory="/available/default", + ) + + assert result.action == "stop" + loop_context = run_loop.await_args.args[0] + assert loop_context.session.directory == "/available/default" + assert loop_context.session_ctx.directory == "/available/default" + assert session.directory == "/missing/original" diff --git a/tests/session/test_session_memory_log_levels.py b/tests/session/test_session_memory_log_levels.py new file mode 100644 index 000000000..5a0571cf1 --- /dev/null +++ b/tests/session/test_session_memory_log_levels.py @@ -0,0 +1,40 @@ +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +import flocks.session.features.memory as memory_module + + +@pytest.mark.asyncio +async def test_missing_memory_config_logs_info(monkeypatch, tmp_path): + manager = SimpleNamespace(initialize=AsyncMock(return_value=None)) + info_log = MagicMock() + warn_log = MagicMock() + + monkeypatch.setattr( + memory_module.Config, + "get", + AsyncMock(return_value=SimpleNamespace(memory=None)), + ) + monkeypatch.setattr( + memory_module.MemoryManager, + "get_instance", + lambda **_kwargs: manager, + ) + monkeypatch.setattr(memory_module.log, "info", info_log) + monkeypatch.setattr(memory_module.log, "warn", warn_log) + + memory = memory_module.SessionMemory( + session_id="session-no-memory-config", + project_id="project", + workspace_dir=str(tmp_path), + enabled=True, + ) + + assert await memory.initialize() is True + info_log.assert_any_call( + "session.memory.no_config", + {"session_id": "session-no-memory-config"}, + ) + assert not any(call.args and call.args[0] == "session.memory.no_config" for call in warn_log.call_args_list) diff --git a/tests/skill/test_device_integration_guide_skill.py b/tests/skill/test_device_integration_guide_skill.py new file mode 100644 index 000000000..e619f8462 --- /dev/null +++ b/tests/skill/test_device_integration_guide_skill.py @@ -0,0 +1,47 @@ +from pathlib import Path + + +PROJECT_ROOT = Path(__file__).resolve().parents[2] +SKILL_FILE = ( + PROJECT_ROOT + / ".flocks" + / "plugins" + / "skills" + / "device-integration-guide" + / "SKILL.md" +) + + +def test_device_integration_guide_queries_templates_before_custom_creation() -> None: + content = SKILL_FILE.read_text(encoding="utf-8") + + assert 'device_manage(action="list_templates")' in content + assert "设备实例为空不代表模板不存在" in content + assert "没有匹配模板：进入自定义" in content + + +def test_device_integration_guide_routes_uninstalled_template_to_flockhub() -> None: + content = SKILL_FILE.read_text(encoding="utf-8") + + assert "installed=false" in content + assert "FlockHub 安装返回的 `plugin_id`" in content + assert "不要创建自定义模板" in content + + +def test_device_integration_guide_creates_from_installed_template() -> None: + content = SKILL_FILE.read_text(encoding="utf-8") + + assert 'device_manage(action="create")' in content + assert 'device_manage(\n action="create"' not in content + assert "不能用 `update` 代替创建" in content + assert "只有 `list_templates` 已返回匹配模板且 `installed=true`" in content + assert "sensitive_fields_to_complete" in content + assert "同一轮会话已经拿到成功返回的 `device_id`" in content + + +def test_device_integration_guide_does_not_embed_page_json_protocol() -> None: + content = SKILL_FILE.read_text(encoding="utf-8") + + assert "```json" not in content + assert '"storage_key":""' not in content + assert "一键回填" not in content diff --git a/tests/skill/test_skill.py b/tests/skill/test_skill.py index e704e785e..045253ab0 100644 --- a/tests/skill/test_skill.py +++ b/tests/skill/test_skill.py @@ -10,6 +10,7 @@ from pathlib import Path from unittest.mock import patch +import flocks.skill.skill as skill_module from flocks.skill.skill import Skill, SkillInfo, SkillRequires, SkillInstallSpec @@ -367,6 +368,71 @@ async def test_discover_project_source_for_project_plugins(tmp_path): assert skill.source == "project", f"expected 'project', got {skill.source!r}" +@pytest.mark.asyncio +async def test_discover_source_root_plugins_from_unrelated_cwd(tmp_path): + """Bundled source-root skills remain visible outside the source checkout.""" + source_root = tmp_path / "source" + ordinary_cwd = tmp_path / "ordinary-cwd" + ordinary_cwd.mkdir() + bundled_skill = source_root / ".flocks" / "plugins" / "skills" / "bundled-skill" + _write_skill_md(bundled_skill / "SKILL.md", "bundled-skill") + + module_file = source_root / "flocks" / "skill" / "skill.py" + with ( + patch("os.path.expanduser", return_value=str(tmp_path / "home")), + patch("flocks.skill.skill.__file__", str(module_file)), + patch( + "flocks.skill.skill.Instance.get_directory", + return_value=str(ordinary_cwd), + ), + patch( + "flocks.skill.skill.Instance.get_worktree", + return_value=str(ordinary_cwd), + ), + ): + skills = await Skill.all() + + skill = next((item for item in skills if item.name == "bundled-skill"), None) + assert skill is not None, "source-root bundled skill not discovered" + assert skill.source == "project" + + +@pytest.mark.asyncio +async def test_skill_cache_does_not_leak_across_project_contexts(tmp_path): + """An ordinary-session discovery must not poison a project lookup.""" + fake_home = tmp_path / "home" + ordinary_cwd = tmp_path / "ordinary-cwd" + project_dir = tmp_path / "project" + source_root = tmp_path / "source" + ordinary_cwd.mkdir() + project_dir.mkdir() + _write_skill_md( + project_dir / ".flocks" / "plugins" / "skills" / "project-only" / "SKILL.md", + "project-only", + ) + + context = {"directory": str(ordinary_cwd)} + module_file = source_root / "flocks" / "skill" / "skill.py" + with ( + patch("os.path.expanduser", return_value=str(fake_home)), + patch("flocks.skill.skill.__file__", str(module_file)), + patch( + "flocks.skill.skill.Instance.get_directory", + side_effect=lambda: context["directory"], + ), + patch( + "flocks.skill.skill.Instance.get_worktree", + side_effect=lambda: context["directory"], + ), + ): + ordinary_skills = await Skill.all() + context["directory"] = str(project_dir) + project_skills = await Skill.all() + + assert "project-only" not in {skill.name for skill in ordinary_skills} + assert "project-only" in {skill.name for skill in project_skills} + + @pytest.mark.asyncio async def test_discover_flocks_source_for_builtin_skills(tmp_path): """Skills under .flocks/skills/ (not plugins/) must get source='flocks'.""" @@ -392,8 +458,8 @@ async def test_discover_flocks_source_for_builtin_skills(tmp_path): @pytest.mark.asyncio -async def test_discover_project_overrides_global(tmp_path): - """Project-level plugin skill must override global skill of the same name.""" +async def test_discover_user_plugin_overrides_project_bundle(tmp_path): + """User plugin skills must override project-bundled skills of the same name.""" fake_home = tmp_path / "home" project_dir = tmp_path / "myproject" project_dir.mkdir() @@ -409,14 +475,24 @@ async def test_discover_project_overrides_global(tmp_path): patch("os.path.expanduser", return_value=str(fake_home)), patch("flocks.skill.skill.Instance.get_directory", return_value=str(project_dir)), patch("flocks.skill.skill.Instance.get_worktree", return_value=str(project_dir)), + patch.object(skill_module.log, "info") as info_log, + patch.object(skill_module.log, "warn") as warn_log, ): Skill.clear_cache() skills = await Skill.all() skill = next((s for s in skills if s.name == "shared-skill"), None) assert skill is not None - assert skill.source == "project", "project-level skill should override global" - assert "Project version" in skill.description + assert skill.source == "user", "user-level skill should override project bundle" + assert "Global version" in skill.description + assert any( + call.args and call.args[0] == "skill.override" + for call in info_log.call_args_list + ) + assert not any( + call.args and call.args[0] == "skill.duplicate" + for call in warn_log.call_args_list + ) @pytest.mark.asyncio diff --git a/tests/tool/test_channel_message.py b/tests/tool/test_channel_message.py index ed815ef57..fe772981e 100644 --- a/tests/tool/test_channel_message.py +++ b/tests/tool/test_channel_message.py @@ -26,6 +26,11 @@ def test_channel_message_normalizes_wecom_aliases() -> None: assert _normalize_channel_type("wxwork") == "wecom" +def test_channel_message_normalizes_slack_aliases() -> None: + assert _normalize_channel_type("slack") == "slack" + assert _normalize_channel_type("sl") == "slack" + + def test_channel_message_normalizes_telegram_whatsapp_email_aliases() -> None: assert _normalize_channel_type("telegram") == "telegram" assert _normalize_channel_type("tg") == "telegram" @@ -52,6 +57,7 @@ def test_channel_message_schema_includes_builtin_channels() -> None: "telegram", "whatsapp", "email", + "slack", "邮件", ): assert value in channel_enum diff --git a/tests/tool/test_delegate_task_compat.py b/tests/tool/test_delegate_task_compat.py index 047ba6d07..1229aa1f9 100644 --- a/tests/tool/test_delegate_task_compat.py +++ b/tests/tool/test_delegate_task_compat.py @@ -23,6 +23,8 @@ def test_delegate_task_schema_allows_omitting_optional_fields(self): @pytest.mark.asyncio async def test_delegate_task_derives_description_and_ignores_blank_skills(self): + ctx = _make_ctx() + ctx.extra["workflow_temp_parent"] = True parent_session = SimpleNamespace( id="test-session", project_id="proj", @@ -47,7 +49,7 @@ async def test_delegate_task_derives_description_and_ignores_blank_skills(self): ): result = await ToolRegistry.execute( "delegate_task", - ctx=_make_ctx(), + ctx=ctx, subagent_type="asset-survey", prompt="Investigate threatbook.cn assets", load_skills=["", " "], @@ -56,6 +58,7 @@ async def test_delegate_task_derives_description_and_ignores_blank_skills(self): assert result.success is True assert result.title == "Investigate threatbook.cn assets" assert result.metadata["sessionId"] == "ses-child" + assert ctx.extra["workflow_child_session_created"] is True skill_get.assert_not_awaited() permissions = create_session.await_args.kwargs["permission"] denied_permissions = {rule.permission for rule in permissions if rule.action == "deny"} diff --git a/tests/tool/test_device_manage_tool.py b/tests/tool/test_device_manage_tool.py index 8e7c82861..01ed78407 100644 --- a/tests/tool/test_device_manage_tool.py +++ b/tests/tool/test_device_manage_tool.py @@ -7,7 +7,7 @@ from flocks.tool.device.intake import DeviceNotFoundError from flocks.tool.device.manage_tool import device_manage -from flocks.tool.device.models import DeviceIntegration, DeviceTestResult +from flocks.tool.device.models import DeviceIntegration, DeviceTemplate, DeviceTestResult from flocks.tool.registry import ToolContext, ToolRegistry @@ -41,19 +41,43 @@ def make_device(**overrides) -> DeviceIntegration: return DeviceIntegration(**data) +def make_template(**overrides) -> DeviceTemplate: + data = { + "plugin_id": "360-waf", + "storage_key": "360_waf_v5_5", + "service_id": "360_waf", + "name": "360 WAF", + "credential_schema": [], + "installed": True, + "state": "installed", + "source": "project", + } + data.update(overrides) + return DeviceTemplate(**data) + + def test_device_manage_is_registered(): tools = {tool.name for tool in ToolRegistry.list_tools()} assert "device_manage" in tools -def test_device_manage_schema_includes_update_action(): +def test_device_manage_schema_includes_template_discovery_action(): tool = ToolRegistry.get("device_manage") assert tool is not None action_param = next(param for param in tool.info.parameters if param.name == "action") - assert action_param.enum == ["list", "update", "connectivity_test"] + assert action_param.enum == [ + "list", + "list_templates", + "create", + "update", + "connectivity_test", + ] assert {param.name for param in tool.info.parameters} >= { "device_id", + "device_name", + "storage_key", + "group_id", "fields", "verify_ssl", } @@ -71,6 +95,228 @@ async def test_device_manage_list_returns_device_inventory(): assert "dev-1" in result.output +@pytest.mark.asyncio +async def test_device_manage_list_templates_returns_existing_template_metadata(): + templates = [ + make_template( + plugin_id="tdp", + storage_key="tdp_v3_3_10", + service_id="tdp", + name="TDP", + version="3.3.10", + vendor="threatbook", + credential_schema=[ + { + "key": "base_url", + "label": "Base URL", + "required": True, + "secret": False, + }, + { + "key": "api_key", + "label": "API Key", + "required": True, + "secret": True, + }, + ], + tool_count=2, + ) + ] + with patch( + "flocks.tool.device.plugin_index.list_device_templates", + return_value=templates, + ) as mocked_list: + result = await device_manage(make_ctx(), action="list_templates") + + mocked_list.assert_called_once_with(refresh=False) + assert result.success is True + assert result.output == [templates[0].model_dump(mode="json")] + assert result.metadata == { + "template_count": 1, + "installed_count": 1, + } + + +@pytest.mark.asyncio +async def test_device_manage_create_uses_installed_template_identity(): + template = make_template( + credential_schema=[ + { + "key": "base_url", + "label": "Base URL", + "storage": "config", + "required": True, + }, + { + "key": "username", + "label": "Username", + "storage": "config", + "required": False, + }, + { + "key": "auth_state", + "label": "Auth State", + "storage": "config", + "required": False, + }, + { + "key": "password", + "label": "Password", + "storage": "secret", + "required": True, + }, + ], + ) + created = make_device( + id="dev-360", + name="360 WAF", + storage_key=template.storage_key, + service_id=template.service_id, + fields={"base_url": "https://192.168.1.100"}, + fields_set={"base_url": True}, + ) + with ( + patch( + "flocks.tool.device.plugin_index.list_device_templates", + return_value=[template], + ), + patch( + "flocks.tool.device.manage_tool.create_device", + AsyncMock(return_value=created), + ) as mocked_create, + ): + result = await device_manage( + make_ctx(), + action="create", + storage_key="360_waf_v5_5", + device_name="360 WAF", + fields={ + "base_url": "https://192.168.1.100", + "auth_state": "ready", + }, + verify_ssl=False, + ) + + mocked_create.assert_awaited_once() + body = mocked_create.await_args.args[0] + assert body.name == "360 WAF" + assert body.storage_key == "360_waf_v5_5" + assert body.service_id == "360_waf" + assert body.fields == { + "base_url": "https://192.168.1.100", + "auth_state": "ready", + } + assert result.success is True + assert result.output["device_id"] == "dev-360" + assert result.metadata["sensitive_fields"] == ["password"] + + +@pytest.mark.asyncio +async def test_device_manage_create_rejects_uninstalled_template(): + template = make_template( + installed=False, + state="available", + source="bundled", + ) + with patch( + "flocks.tool.device.plugin_index.list_device_templates", + return_value=[template], + ): + result = await device_manage( + make_ctx(), + action="create", + storage_key=template.storage_key, + device_name="360 WAF", + ) + + assert result.success is False + assert "尚未安装" in (result.error or "") + assert "360-waf" in (result.error or "") + + +@pytest.mark.asyncio +async def test_device_manage_create_rejects_secret_and_unknown_fields(): + template = make_template( + credential_schema=[ + {"key": "base_url", "storage": "config", "required": True}, + {"key": "password", "storage": "secret", "required": True}, + ], + ) + with patch( + "flocks.tool.device.plugin_index.list_device_templates", + return_value=[template], + ): + secret_result = await device_manage( + make_ctx(), + action="create", + storage_key=template.storage_key, + device_name="360 WAF", + fields={ + "base_url": "https://192.168.1.100", + "password": "do-not-store", + }, + ) + unknown_result = await device_manage( + make_ctx(), + action="create", + storage_key=template.storage_key, + device_name="360 WAF", + fields={ + "base_url": "https://192.168.1.100", + "account": "admin", + }, + ) + + assert secret_result.success is False + assert "敏感字段" in (secret_result.error or "") + assert "password" in (secret_result.error or "") + assert unknown_result.success is False + assert "模板未声明字段" in (unknown_result.error or "") + assert "account" in (unknown_result.error or "") + + +@pytest.mark.asyncio +async def test_device_manage_create_leaves_missing_fields_for_page_completion(): + template = make_template( + credential_schema=[ + { + "key": "base_url", + "storage": "config", + "required": True, + "default_value": "https://default.local", + }, + {"key": "password", "storage": "secret", "required": True}, + ], + ) + created = make_device( + name="360 WAF", + storage_key=template.storage_key, + service_id=template.service_id, + fields={}, + fields_set={}, + ) + with ( + patch( + "flocks.tool.device.plugin_index.list_device_templates", + return_value=[template], + ), + patch( + "flocks.tool.device.manage_tool.create_device", + AsyncMock(return_value=created), + ) as mocked_create, + ): + result = await device_manage( + make_ctx(), + action="create", + storage_key=template.storage_key, + device_name="360 WAF", + ) + + assert result.success is True + assert mocked_create.await_args.args[0].fields == {} + assert result.output["sensitive_fields_to_complete"] == ["password"] + + @pytest.mark.asyncio async def test_device_manage_update_updates_existing_device_non_secret_config(): updated_device = make_device(verify_ssl=True) @@ -82,7 +328,11 @@ async def test_device_manage_update_updates_existing_device_non_secret_config(): make_ctx(), action="update", device_id="dev-1", - fields={"base_url": "https://device.local", "port": 443}, + fields={ + "base_url": "https://device.local", + "port": 443, + "auth_state": "ready", + }, verify_ssl=True, ) @@ -92,11 +342,12 @@ async def test_device_manage_update_updates_existing_device_non_secret_config(): assert update_body.fields == { "base_url": "https://device.local", "port": "443", + "auth_state": "ready", } assert update_body.verify_ssl is True assert result.success is True assert result.output["device_id"] == "dev-1" - assert result.output["updated_fields"] == ["base_url", "port"] + assert result.output["updated_fields"] == ["auth_state", "base_url", "port"] assert result.metadata["verify_ssl"] is True diff --git a/tests/tool/test_device_startup_sync.py b/tests/tool/test_device_startup_sync.py index 37986aa48..0481e22e5 100644 --- a/tests/tool/test_device_startup_sync.py +++ b/tests/tool/test_device_startup_sync.py @@ -15,6 +15,7 @@ """ from __future__ import annotations +import sqlite3 from pathlib import Path from typing import List @@ -24,6 +25,33 @@ from flocks.tool.device import startup +@pytest.mark.asyncio +async def test_heal_stale_service_ids_reads_named_sqlite_rows(tmp_path, monkeypatch): + db_path = tmp_path / "devices.db" + with sqlite3.connect(db_path) as db: + db.execute( + "CREATE TABLE device_integrations (id TEXT, storage_key TEXT, service_id TEXT)" + ) + db.execute( + "INSERT INTO device_integrations VALUES (?, ?, ?)", + ("dev-1", "onesec_api_v2_8_2", "stale_service_id"), + ) + + monkeypatch.setattr(startup.Storage, "get_db_path", staticmethod(lambda: db_path)) + monkeypatch.setattr( + "flocks.tool.device.store.storage_key_to_service_id", + lambda storage_key: "onesec_api" if storage_key == "onesec_api_v2_8_2" else storage_key, + ) + + await startup._heal_stale_service_ids() + + with sqlite3.connect(db_path) as db: + row = db.execute( + "SELECT service_id FROM device_integrations WHERE id = ?", ("dev-1",) + ).fetchone() + assert row == ("onesec_api",) + + @pytest.fixture def isolated_plugins(monkeypatch, tmp_path): """Point plugin discovery at an empty tmp HOME so production plugins diff --git a/tests/tool/test_failure_auto_disable_config.py b/tests/tool/test_failure_auto_disable_config.py new file mode 100644 index 000000000..b54c654cb --- /dev/null +++ b/tests/tool/test_failure_auto_disable_config.py @@ -0,0 +1,130 @@ +from __future__ import annotations + +from concurrent.futures import ThreadPoolExecutor +from unittest.mock import Mock + +import pytest + +from flocks.config.config import Config, ConfigInfo, ToolFailureConfig +from flocks.config.config_writer import ConfigWriter +from flocks.tool.registry import ( + Tool, + ToolCategory, + ToolContext, + ToolInfo, + ToolRegistry, + ToolResult, +) + + +def _failing_tool(name: str = "failing_custom_tool") -> Tool: + async def handler(_ctx: ToolContext, **_kwargs) -> ToolResult: + return ToolResult(success=False, error="synthetic repeated failure") + + return Tool( + info=ToolInfo( + name=name, + description="Always fails for repeated-failure tests", + category=ToolCategory.CUSTOM, + source="plugin_py", + ), + handler=handler, + ) + + +@pytest.fixture +def isolated_failure_tracking(monkeypatch: pytest.MonkeyPatch): + tool = _failing_tool() + monkeypatch.setattr(ToolRegistry, "_initialized", True) + monkeypatch.setattr(ToolRegistry, "_tools", {tool.info.name: tool}) + monkeypatch.setattr(ToolRegistry, "_failure_state", {}) + monkeypatch.setattr(ToolRegistry, "_revision", 41) + monkeypatch.setattr(Config, "_cached_config", ConfigInfo()) + monkeypatch.setattr(ConfigWriter, "set_tool_setting", lambda _name, _setting: None) + return tool + + +def test_tool_failure_config_defaults_to_enabled() -> None: + section = ToolFailureConfig() + + assert section.disable_on_repeated_failure is True + + +def test_tool_failure_config_accepts_camel_and_snake_case() -> None: + camel = ConfigInfo.model_validate( + {"toolFailure": {"disableOnRepeatedFailure": False}} + ) + snake = ConfigInfo.model_validate( + {"tool_failure": {"disable_on_repeated_failure": False}} + ) + + assert camel.tool_failure is not None + assert camel.tool_failure.disable_on_repeated_failure is False + assert snake.tool_failure is not None + assert snake.tool_failure.disable_on_repeated_failure is False + + +@pytest.mark.asyncio +async def test_repeated_failures_disable_tool_by_default( + isolated_failure_tracking: Tool, +) -> None: + tool = isolated_failure_tracking + + for _ in range(ToolRegistry._failure_disable_threshold): + result = await ToolRegistry.execute(tool.info.name, query="same") + + assert tool.info.enabled is False + assert result.metadata == { + "disabled": True, + "disabled_reason": "repeated_error", + } + assert ToolRegistry.revision() == 42 + + +def test_concurrent_failures_commit_one_disable_transition( + isolated_failure_tracking: Tool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + tool = isolated_failure_tracking + params = {"query": "same"} + error = "synthetic repeated failure" + ToolRegistry._failure_state[tool.info.name] = { + "key": ToolRegistry._failure_key(tool.info.name, params, error), + "count": ToolRegistry._failure_disable_threshold - 1, + } + persist_setting = Mock() + monkeypatch.setattr(ConfigWriter, "set_tool_setting", persist_setting) + + with ThreadPoolExecutor(max_workers=8) as executor: + transitions = list(executor.map( + lambda _: ToolRegistry._record_failure(tool, params, error), + range(8), + )) + + assert transitions.count(True) == 1 + assert tool.info.enabled is False + assert ToolRegistry.revision() == 42 + persist_setting.assert_called_once_with(tool.info.name, {"enabled": False}) + + +@pytest.mark.asyncio +async def test_config_can_turn_off_repeated_failure_auto_disable( + isolated_failure_tracking: Tool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + tool = isolated_failure_tracking + monkeypatch.setattr( + Config, + "_cached_config", + ConfigInfo.model_validate( + {"toolFailure": {"disableOnRepeatedFailure": False}} + ), + ) + + for _ in range(ToolRegistry._failure_disable_threshold + 1): + result = await ToolRegistry.execute(tool.info.name, query="same") + + assert result.success is False + assert "disabled" not in result.metadata + assert tool.info.enabled is True + assert ToolRegistry._failure_state == {} diff --git a/tests/tool/test_huaweicloud_waf_huorong_device_plugins.py b/tests/tool/test_huaweicloud_waf_huorong_device_plugins.py new file mode 100644 index 000000000..e3a8cc8ee --- /dev/null +++ b/tests/tool/test_huaweicloud_waf_huorong_device_plugins.py @@ -0,0 +1,181 @@ +from __future__ import annotations + +import ast +import inspect +import shutil +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest +import yaml + +from flocks.tool.registry import ToolContext, ToolResult +from flocks.tool.tool_loader import yaml_to_tool + + +_ROOT = Path(__file__).resolve().parents[2] +_DEVICE_ROOT = _ROOT / ".flocks" / "flockshub" / "plugins" / "tools" / "device" +_HUAWEI_DIR = _DEVICE_ROOT / "huaweicloud_waf_v39" +_HUORONG_DIR = _DEVICE_ROOT / "huorong_edr_v1_0" + + +def _installed_tool( + source_dir: Path, + yaml_name: str, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +): + project_root = tmp_path / "project" + install_dir = project_root / ".flocks" / "plugins" / "tools" / "device" / source_dir.name + shutil.copytree(source_dir, install_dir) + monkeypatch.chdir(project_root) + yaml_path = install_dir / yaml_name + raw = yaml.safe_load(yaml_path.read_text(encoding="utf-8")) + return raw, yaml_to_tool(raw, yaml_path) + + +def _script_globals(tool) -> dict[str, Any]: + return inspect.getclosurevars(tool.handler).nonlocals["fn"].__globals__ + + +def _ctx() -> ToolContext: + return ToolContext(session_id="device-plugin-test", message_id="device-plugin-test") + + +def test_huawei_manifests_separate_dispatch_action_from_api_payload(): + event = yaml.safe_load((_HUAWEI_DIR / "hw_waf_event.yaml").read_text(encoding="utf-8")) + event_properties = event["inputSchema"]["properties"] + + assert "event_list" in event_properties["action"]["enum"] + assert event_properties["event_action"]["enum"] == ["block", "log"] + + policy = yaml.safe_load((_HUAWEI_DIR / "hw_waf_policy.yaml").read_text(encoding="utf-8")) + policy_properties = policy["inputSchema"]["properties"] + + assert "policy_list" in policy_properties["action"]["enum"] + assert policy_properties["rule_action"]["type"] == "object" + + +@pytest.mark.parametrize( + "handler_path", + [_HUAWEI_DIR / "hw_waf.handler.py", _HUORONG_DIR / "huorong.handler.py"], +) +def test_tool_result_calls_use_supported_fields(handler_path: Path): + module = ast.parse(handler_path.read_text(encoding="utf-8"), filename=str(handler_path)) + unsupported: list[tuple[int, str]] = [] + for node in ast.walk(module): + if not isinstance(node, ast.Call) or not isinstance(node.func, ast.Name) or node.func.id != "ToolResult": + continue + unsupported.extend( + (node.lineno, keyword.arg) + for keyword in node.keywords + if keyword.arg is not None and keyword.arg not in ToolResult.model_fields + ) + + assert unsupported == [] + + +@pytest.mark.parametrize( + "yaml_name", + [ + "hw_waf_host.yaml", + "hw_waf_policy.yaml", + "hw_waf_event.yaml", + "hw_waf_overview.yaml", + ], +) +@pytest.mark.asyncio +async def test_huawei_handlers_execute_through_script_loader( + yaml_name: str, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +): + _, tool = _installed_tool(_HUAWEI_DIR, yaml_name, tmp_path, monkeypatch) + globals_ = _script_globals(tool) + monkeypatch.setitem(globals_, "_load_config", lambda *_: SimpleNamespace(project_id="project-1")) + + result = await tool.handler(_ctx(), action="unsupported") + + assert result.success is False + assert result.error == "Unknown action: unsupported" + + +@pytest.mark.asyncio +async def test_huawei_event_action_maps_to_api_filter( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +): + _, tool = _installed_tool(_HUAWEI_DIR, "hw_waf_event.yaml", tmp_path, monkeypatch) + globals_ = _script_globals(tool) + calls: list[dict[str, Any]] = [] + + async def fake_request(cfg, method, path, query=None, body=None): + calls.append({"method": method, "path": path, "query": query, "body": body}) + return ToolResult(success=True, output={}) + + monkeypatch.setitem(globals_, "_load_config", lambda *_: SimpleNamespace(project_id="project-1")) + monkeypatch.setitem(globals_, "_request", fake_request) + + result = await tool.handler( + _ctx(), + action="event_list", + event_action="block", + page=1, + ) + + assert result.success is True + assert calls[0]["query"] == {"page": 1, "action": "block"} + + +@pytest.mark.asyncio +async def test_huawei_rule_action_maps_to_api_body( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +): + _, tool = _installed_tool(_HUAWEI_DIR, "hw_waf_policy.yaml", tmp_path, monkeypatch) + globals_ = _script_globals(tool) + calls: list[dict[str, Any]] = [] + + async def fake_request(cfg, method, path, query=None, body=None): + calls.append({"method": method, "path": path, "query": query, "body": body}) + return ToolResult(success=True, output={}) + + monkeypatch.setitem(globals_, "_load_config", lambda *_: SimpleNamespace(project_id="project-1")) + monkeypatch.setitem(globals_, "_request", fake_request) + rule_action = {"category": "block"} + + result = await tool.handler( + _ctx(), + action="cc_rule_create", + policy_id="policy-1", + url="/login", + rule_action=rule_action, + ) + + assert result.success is True + assert calls[0]["body"] == {"url": "/login", "action": rule_action} + + +@pytest.mark.parametrize( + "yaml_name", + ["huorong_group.yaml", "huorong_clnts.yaml", "huorong_task.yaml"], +) +@pytest.mark.asyncio +async def test_huorong_handlers_execute_through_script_loader( + yaml_name: str, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +): + _, tool = _installed_tool(_HUORONG_DIR, yaml_name, tmp_path, monkeypatch) + globals_ = _script_globals(tool) + monkeypatch.setitem( + globals_, + "_resolve_runtime_config", + lambda: ("https://huorong.example.com", 30, "secret-id", "secret-key", False), + ) + + result = await tool.handler(_ctx(), action="unsupported") + + assert result.success is False + assert result.error == "Unknown action: unsupported" diff --git a/tests/tool/test_im_send_message.py b/tests/tool/test_im_send_message.py index 42e41a674..f582d3d41 100644 --- a/tests/tool/test_im_send_message.py +++ b/tests/tool/test_im_send_message.py @@ -51,6 +51,11 @@ def test_im_send_message_normalizes_wecom_aliases() -> None: assert _normalize_channel_type("wxwork") == "wecom" +def test_im_send_message_normalizes_slack_aliases() -> None: + assert _normalize_channel_type("slack") == "slack" + assert _normalize_channel_type("sl") == "slack" + + def test_im_send_message_normalizes_telegram_whatsapp_email_aliases() -> None: assert _normalize_channel_type("telegram") == "telegram" assert _normalize_channel_type("tg") == "telegram" @@ -68,7 +73,7 @@ def test_im_send_message_schema_mentions_extended_builtin_channels() -> None: assert schema is not None channel_description = schema.properties["channel_type"]["description"] tool_description = ToolRegistry.get("im_send_message").info.description - for channel in ("telegram", "whatsapp", "email"): + for channel in ("telegram", "whatsapp", "email", "slack"): assert channel in channel_description.lower() assert channel in tool_description.lower() diff --git a/tests/tool/test_sangfor_edr_handler.py b/tests/tool/test_sangfor_edr_handler.py new file mode 100644 index 000000000..3eb286367 --- /dev/null +++ b/tests/tool/test_sangfor_edr_handler.py @@ -0,0 +1,183 @@ +import importlib.util +from pathlib import Path + + +_HANDLER_PATH = ( + Path(__file__).resolve().parents[2] + / ".flocks" + / "plugins" + / "tools" + / "device" + / "sangfor_edr_webcli" + / "sangfor_edr.handler.py" +) + + +def _load_handler(): + spec = importlib.util.spec_from_file_location("_test_sangfor_edr_handler", _HANDLER_PATH) + assert spec and spec.loader + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _cfg(handler, state_path: Path): + return handler.RuntimeConfig( + base_url="https://edr.example.com", + auth_state_path=state_path, + username="admin", + password="secret", + login_path="/ui/login.php", + index_path="/ui/#/index", + timeout=5, + auto_ocr_code=True, + max_captcha_retry=1, + username_selector="", + password_selector="", + captcha_selector="", + agreement_selector="", + submit_selector="", + ) + + +def test_validate_missing_state_does_not_start_daemon(tmp_path, monkeypatch): + handler = _load_handler() + monkeypatch.setattr(handler, "_ensure_browser_daemon", lambda: (_ for _ in ()).throw(AssertionError())) + + result = handler._validate_auth_state(_cfg(handler, tmp_path / "missing.json")) + + assert result["reason"] == "auth_state_not_found" + + +def test_validate_existing_state_ensures_daemon(tmp_path, monkeypatch): + handler = _load_handler() + state_path = tmp_path / "auth-state.json" + state_path.write_text('{"cookies": [], "origins": []}', encoding="utf-8") + calls = [] + monkeypatch.setattr(handler, "_ensure_browser_daemon", lambda: calls.append("ensure")) + monkeypatch.setattr(handler.helpers, "load_state", lambda *args, **kwargs: {}) + monkeypatch.setattr(handler.helpers, "wait_for_load", lambda *args, **kwargs: True) + monkeypatch.setattr(handler, "_is_logged_in", lambda cfg: True) + + result = handler._validate_auth_state(_cfg(handler, state_path)) + + assert result["valid"] is True + assert calls == ["ensure"] + + +def test_validate_reports_daemon_failure_separately(tmp_path, monkeypatch): + handler = _load_handler() + state_path = tmp_path / "auth-state.json" + state_path.write_text('{"cookies": [], "origins": []}', encoding="utf-8") + monkeypatch.setattr( + handler, + "_ensure_browser_daemon", + lambda: (_ for _ in ()).throw(RuntimeError("browser daemon did not start")), + ) + + result = handler._validate_auth_state(_cfg(handler, state_path)) + + assert result["status"] == "browser_daemon_not_ready" + assert result["reason"] == "auth_state_load_failed_browser_daemon_not_ready" + assert "flocks browser --setup" in result["next_action"] + + +def test_refresh_distinguishes_page_open_failure(tmp_path, monkeypatch): + handler = _load_handler() + cfg = _cfg(handler, tmp_path / "auth-state.json") + monkeypatch.setattr(handler, "_ensure_browser_daemon", lambda: None) + monkeypatch.setattr( + handler, + "_open_page", + lambda url: (_ for _ in ()).throw(RuntimeError("EDR connection refused")), + ) + + result = handler._refresh_auth_state_with_cdp_login(cfg) + + assert result["status"] == "browser_login_page_open_failed" + assert result["reason"] == "browser_login_page_open_failed" + assert result["login_url"] == "https://edr.example.com/ui/login.php" + + +def test_refresh_falls_back_to_manual_login_when_form_is_missing(tmp_path, monkeypatch): + handler = _load_handler() + cfg = _cfg(handler, tmp_path / "auth-state.json") + monkeypatch.setattr(handler, "_ensure_browser_daemon", lambda: None) + monkeypatch.setattr(handler, "_open_page", lambda url: None) + monkeypatch.setattr(handler, "_wait_for_login_form_ready", lambda cfg: (_ for _ in ()).throw(RuntimeError("missing form"))) + monkeypatch.setattr(handler, "_page_text", lambda: "请输入动态口令") + + result = handler._refresh_auth_state_with_cdp_login(cfg) + + assert result["status"] == "manual_login_required" + assert result["reason"] == "mfa_required" + assert result["browser_left_open"] is True + + +def test_complete_manual_login_saves_state(tmp_path, monkeypatch): + handler = _load_handler() + cfg = _cfg(handler, tmp_path / "auth-state.json") + monkeypatch.setattr(handler, "_ensure_browser_daemon", lambda: None) + monkeypatch.setattr(handler.helpers, "page_info", lambda: {"url": "https://edr.example.com/ui/#/index"}) + monkeypatch.setattr(handler, "_is_logged_in", lambda cfg: True) + monkeypatch.setattr(handler, "_save_auth_state", lambda cfg: {"cookies": 1}) + monkeypatch.setattr(handler, "_save_captured_login_token", lambda: True) + + result = handler._complete_manual_login(cfg) + + assert result["status"] == "manual_login_captured_auth_state" + assert result["saved"] == {"cookies": 1} + assert result["token_saved"] is True + + +def test_complete_manual_login_keeps_waiting_when_not_logged_in(tmp_path, monkeypatch): + handler = _load_handler() + cfg = _cfg(handler, tmp_path / "auth-state.json") + monkeypatch.setattr(handler, "_ensure_browser_daemon", lambda: None) + monkeypatch.setattr(handler.helpers, "page_info", lambda: {"url": "https://edr.example.com/ui/login.php"}) + monkeypatch.setattr(handler, "_is_logged_in", lambda cfg: False) + + result = handler._complete_manual_login(cfg) + + assert result["status"] == "manual_login_required" + assert result["reason"] == "manual_login_not_completed" + + +def test_missing_credentials_opens_browser_for_manual_login(tmp_path, monkeypatch): + handler = _load_handler() + cfg = _cfg(handler, tmp_path / "auth-state.json") + cfg.username = "" + opened = [] + monkeypatch.setattr(handler, "_ensure_browser_daemon", lambda: None) + monkeypatch.setattr(handler, "_open_page", lambda url: opened.append(url)) + + result = handler._refresh_auth_state_with_cdp_login(cfg) + + assert opened == ["https://edr.example.com/ui/login.php"] + assert result["reason"] == "missing_cdp_login_credentials" + assert result["missing"] == ["username"] + + +def test_captured_login_token_is_saved_as_secret(monkeypatch): + handler = _load_handler() + saved = {} + monkeypatch.setattr(handler.helpers, "js", lambda script: "token-value") + monkeypatch.setattr( + handler, + "_get_secret_manager", + lambda: type("Secrets", (), {"set": lambda self, key, value: saved.update({key: value})})(), + ) + + assert handler._save_captured_login_token(timeout=0.1) is True + assert saved == {"sangfor_edr_token": "token-value"} + + +def test_login_token_capture_hooks_fetch_and_xhr(monkeypatch): + handler = _load_handler() + scripts = [] + monkeypatch.setattr(handler.helpers, "js", lambda script: scripts.append(script) or True) + + assert handler._install_login_token_capture() is True + assert "/launch_login.php" in scripts[0] + assert "window.fetch" in scripts[0] + assert "XMLHttpRequest.prototype" in scripts[0] diff --git a/tests/tool/test_skyeye_api_tools.py b/tests/tool/test_skyeye_api_tools.py index 97c71bf07..9929919b8 100644 --- a/tests/tool/test_skyeye_api_tools.py +++ b/tests/tool/test_skyeye_api_tools.py @@ -1,3 +1,5 @@ +import base64 +import gzip from pathlib import Path from unittest.mock import MagicMock, patch @@ -207,6 +209,45 @@ async def test_skyeye_alarm_list_accepts_legacy_threat_level_alias(): assert "threat_level" not in request_kwargs["params"] +@pytest.mark.asyncio +async def test_skyeye_alarm_list_encodes_ip_filters_without_affecting_other_params(): + tool = _load_tool("skyeye_alarm_list.yaml") + mock_secret_manager = MagicMock() + mock_secret_manager.get.side_effect = lambda key: { + "skyeye_api_key": "login-key-123", + }.get(key) + fake_session = _FakeSession([ + _FakeResponse(json_payload={"access_token": "token-123", "status": 200}), + _FakeResponse(text_payload=''), + _FakeResponse(json_payload={"data": {"items": [], "total": 0}, "status": 1000, "message": "ok"}), + ]) + + with ( + patch("flocks.security.get_secret_manager", return_value=mock_secret_manager), + patch( + "flocks.config.config_writer.ConfigWriter.get_api_service_raw", + return_value={ + "apiKey": "{secret:skyeye_api_key}", + "base_url": "https://skyeye.local", + }, + ), + patch("aiohttp.ClientSession", return_value=fake_session), + ): + result = await tool.handler( + ToolContext(session_id="test", message_id="test"), + alarm_sip="10.105.249.255", + attack_sip="10.61.116.134", + limit=3, + ) + + assert result.success is True + _, _, request_kwargs = fake_session.calls[2] + params = request_kwargs["params"] + assert gzip.decompress(base64.b64decode(params["alarm_sip"])).decode() == "10.105.249.255" + assert gzip.decompress(base64.b64decode(params["attack_sip"])).decode() == "10.61.116.134" + assert params["limit"] == 3 + + @pytest.mark.asyncio async def test_skyeye_alarm_list_still_rejects_unknown_level_fields(): tool = _load_tool("skyeye_alarm_list.yaml") diff --git a/tests/tool/test_tdp_api_tools.py b/tests/tool/test_tdp_api_tools.py index 232ae407e..4d80f5988 100644 --- a/tests/tool/test_tdp_api_tools.py +++ b/tests/tool/test_tdp_api_tools.py @@ -186,8 +186,8 @@ async def test_tdp_dashboard_status_uses_non_deprecated_threat_topic_path(): @pytest.mark.asyncio -async def test_tdp_incident_list_can_use_secret_manager_credentials_and_default_filters(): - tool = _load_tool("tdp_incident_list.yaml") +async def test_tdp_threat_intelligent_aggregation_can_use_secret_manager_credentials_and_default_filters(): + tool = _load_tool("tdp_threat_intelligent_aggregation.yaml") mock_secret_manager = MagicMock() mock_secret_manager.get.side_effect = lambda key: { "tdp_api_key": "secret-api", @@ -212,7 +212,7 @@ async def test_tdp_incident_list_can_use_secret_manager_credentials_and_default_ assert method == "POST" assert request_url == "https://tdp.internal/api/v1/incident/search" condition = request_kwargs["json"]["condition"] - assert condition["duration"] == {"begin_duration": 0, "end_duration": 24} + assert condition["duration"] == {"begin_duration": 0} assert condition["time_from"] < condition["time_to"] @@ -252,8 +252,8 @@ async def test_tdp_machine_asset_list_can_switch_to_web_app_framework_action(): @pytest.mark.asyncio -async def test_tdp_threat_inbound_attack_uses_severity_distribution_endpoint(): - tool = _load_tool("tdp_threat_inbound_attack.yaml") +async def test_tdp_threat_external_attack_uses_severity_distribution_endpoint(): + tool = _load_tool("tdp_threat_external_attack.yaml") fake_session = _FakeSession([_FakeResponse(json_payload={"response_code": 0, "data": [{"key": "4", "value": 2}]})]) with ( diff --git a/tests/tool/test_tdp_skyeye_api_plugins.py b/tests/tool/test_tdp_skyeye_api_plugins.py index d3cf573f6..19013757c 100644 --- a/tests/tool/test_tdp_skyeye_api_plugins.py +++ b/tests/tool/test_tdp_skyeye_api_plugins.py @@ -28,23 +28,23 @@ def _built_json_payload(mock_run): return api_name, path, body_builder(body) -async def test_tdp_incident_timeline_requires_incident_id(): +async def test_tdp_attack_timeline_requires_incident_id(): module = _load_module("test_tdp_handler_incident", _TDP_HANDLER) - result = await module.incident_list(_ctx(), action="timeline") + result = await module.incident_list(_ctx(), action="attack_timeline") assert result.success is False assert "incident_id" in result.error -async def test_tdp_incident_timeline_defaults_show_attack_true(): +async def test_tdp_attack_timeline_defaults_show_attack_true(): module = _load_module("test_tdp_handler_incident_timeline_show_attack", _TDP_HANDLER) mock_result = ToolResult(success=True, output={"status": 200}) with patch.object(module, "_run_action_json_tool", AsyncMock(return_value=mock_result)) as mock_run: result = await module.incident_list( _ctx(), - action="timeline", + action="attack_timeline", incident_id="incident-1", time_from=1700000000, time_to=1700003600, @@ -56,30 +56,30 @@ async def test_tdp_incident_timeline_defaults_show_attack_true(): assert body["show_attack"] is True -async def test_tdp_incident_alert_search_requires_page(): +async def test_tdp_attack_alert_list_requires_page(): module = _load_module("test_tdp_handler_incident_alert_page", _TDP_HANDLER) - result = await module.incident_list(_ctx(), action="alert_search", alert_ids=["alert-1"]) + result = await module.incident_list(_ctx(), action="attack_alert_list", alert_ids=["alert-1"]) assert result.success is False assert "page" in result.error -async def test_tdp_alert_host_events_requires_asset_machine(): +async def test_tdp_host_threat_list_requires_asset_machine(): module = _load_module("test_tdp_handler_alert_host", _TDP_HANDLER) - result = await module.threat_host_list(_ctx(), action="events", condition={}) + result = await module.threat_host_list(_ctx(), action="host_threat_list", condition={}) assert result.success is False assert "condition.asset_machine" in result.error -async def test_tdp_alert_host_summary_uses_all_threat_types_defaults(): - module = _load_module("test_tdp_handler_alert_host_summary_defaults", _TDP_HANDLER) +async def test_tdp_alert_host_list_uses_all_threat_types_defaults(): + module = _load_module("test_tdp_handler_alert_host_list_defaults", _TDP_HANDLER) mock_result = ToolResult(success=True, output={"status": 200}) with patch.object(module, "_run_json_tool", AsyncMock(return_value=mock_result)) as mock_run: - result = await module.threat_host_list(_ctx(), action="summary") + result = await module.threat_host_list(_ctx(), action="alert_host_list") assert result.success is True mock_run.assert_awaited_once() @@ -304,14 +304,14 @@ async def test_tdp_policy_resolve_host_requires_assets_machine_status_and_sub_st assert "sub_status" in result_missing_sub_status.error -async def test_tdp_incident_alert_search_maps_explicit_params_to_condition(): +async def test_tdp_attack_alert_list_maps_explicit_params_to_condition(): module = _load_module("test_tdp_handler_incident_mapping", _TDP_HANDLER) mock_result = ToolResult(success=True, output={"status": 200}) with patch.object(module, "_run_action_json_tool", AsyncMock(return_value=mock_result)) as mock_run: result = await module.incident_list( _ctx(), - action="alert_search", + action="attack_alert_list", alert_ids=["alert-1"], include_risk=True, include_action=False, @@ -325,8 +325,8 @@ async def test_tdp_incident_alert_search_maps_explicit_params_to_condition(): default_action = mock_run.await_args.kwargs["default_action"] action = mock_run.await_args.kwargs["action"] body = mock_run.await_args.kwargs["body"] - assert default_action == "search" - assert action == "alert_search" + assert default_action == "incident_search" + assert action == "attack_alert_list" assert body["condition"]["id"] == ["alert-1"] assert body["condition"]["include_risk"] is True assert body["condition"]["include_action"] is False @@ -752,7 +752,7 @@ async def test_tdp_medium_query_tools_map_semantic_filters(): mock_run.reset_mock() result = await module.incident_list( _ctx(), - action="search", + action="incident_search", severity=[4], phase=["exploit"], result=["success"], @@ -821,13 +821,50 @@ async def test_skyeye_alarm_list_forwards_extended_filters(): assert params["limit"] == 10 -def test_tdp_incident_yaml_loads_with_provider(): - yaml_path = _WORKSPACE_ROOT / ".flocks/plugins/tools/device/tdp_v3_3_10/tdp_incident_list.yaml" +def test_tdp_threat_tool_names_match_document_feature_groups(): + tool_dir = _WORKSPACE_ROOT / ".flocks/plugins/tools/device/tdp_v3_3_10" + expected_names = [ + "tdp_threat_alert_host", + "tdp_threat_monitor_list", + "tdp_threat_external_attack", + "tdp_threat_intelligent_aggregation", + ] + for expected_name in expected_names: + yaml_path = tool_dir / f"{expected_name}.yaml" + raw = _read_yaml_raw(yaml_path) + tool = yaml_to_tool(raw, yaml_path) + assert tool.info.name == expected_name + assert tool.info.provider == "tdp_api_v3_3_10" + + +def test_tdp_threat_action_names_describe_returned_data(): + tool_dir = _WORKSPACE_ROOT / ".flocks/plugins/tools/device/tdp_v3_3_10" + alert_host = _read_yaml_raw(tool_dir / "tdp_threat_alert_host.yaml") + intelligent_aggregation = _read_yaml_raw(tool_dir / "tdp_threat_intelligent_aggregation.yaml") + + assert alert_host["inputSchema"]["properties"]["action"]["enum"] == [ + "alert_host_list", + "host_threat_list", + ] + assert intelligent_aggregation["inputSchema"]["properties"]["action"]["enum"] == [ + "incident_search", + "top_attacked_entity", + "attack_success", + "attack_timeline", + "attack_alert_list", + "attack_result_distribution", + "attacker_ip_list", + "attacker_ip_detail", + ] + + +def test_tdp_threat_intelligent_aggregation_promotes_semantic_fields(): + yaml_path = ( + _WORKSPACE_ROOT + / ".flocks/plugins/tools/device/tdp_v3_3_10/tdp_threat_intelligent_aggregation.yaml" + ) raw = _read_yaml_raw(yaml_path) - tool = yaml_to_tool(raw, yaml_path) - assert tool.info.name == "tdp_incident_list" - assert tool.info.provider == "tdp_api_v3_3_10" assert "body" not in raw["inputSchema"]["properties"] assert "condition" in raw["inputSchema"]["properties"] assert "severity" in raw["inputSchema"]["properties"] @@ -837,14 +874,14 @@ def test_tdp_incident_yaml_loads_with_provider(): def test_tdp_query_yaml_promotes_semantic_top_level_fields(): expected_fields = { - "tdp_host_threat_list.yaml": {"severity", "threat_characters", "keyword", "cur_page"}, + "tdp_threat_alert_host.yaml": {"severity", "threat_characters", "keyword", "cur_page"}, "tdp_vulnerability_list.yaml": {"severity", "status", "keyword", "sort_by"}, "tdp_interface_list.yaml": {"host", "methods", "privacy_tags", "keyword"}, "tdp_interface_risk_list.yaml": {"api_risk_type", "keyword", "sort_by"}, "tdp_login_weakpwd_list.yaml": {"data", "result", "app_class", "keyword"}, "tdp_assets_domain_list.yaml": {"domain_name_or_ip", "has_login_api", "second_level_domain"}, "tdp_privacy_diagram.yaml": {"itag", "methods", "fuzzy_url_host"}, - "tdp_threat_inbound_attack.yaml": {"severity", "result_list", "keyword"}, + "tdp_threat_external_attack.yaml": {"severity", "result_list", "keyword"}, "tdp_threat_monitor_list.yaml": {"time_from", "time_to", "sql", "assets_group", "cur_page"}, "tdp_machine_asset_list.yaml": {"time_from", "time_to", "service", "service_class", "application", "keyword"}, "tdp_mdr_alert_list.yaml": {"section_list", "threat_severity", "keyword"}, diff --git a/tests/tool/test_tool_refresh_outcomes.py b/tests/tool/test_tool_refresh_outcomes.py index 12f826f47..4efd7544b 100644 --- a/tests/tool/test_tool_refresh_outcomes.py +++ b/tests/tool/test_tool_refresh_outcomes.py @@ -2,6 +2,7 @@ from collections.abc import Iterator from contextlib import contextmanager +from pathlib import Path import threading import pytest @@ -34,6 +35,8 @@ def _isolated_registry() -> Iterator[None]: "tools": ToolRegistry._tools, "defaults": ToolRegistry._enabled_defaults, "plugin_names": ToolRegistry._plugin_tool_names, + "plugin_errors": ToolRegistry._plugin_load_errors, + "revision": ToolRegistry._revision, "dynamic_modules": ToolRegistry._dynamic_modules, "dynamic_tools": ToolRegistry._dynamic_tools_by_module, } @@ -42,6 +45,8 @@ def _isolated_registry() -> Iterator[None]: ToolRegistry._tools = {} ToolRegistry._enabled_defaults = {} ToolRegistry._plugin_tool_names = [] + ToolRegistry._plugin_load_errors = [] + ToolRegistry._revision = 0 ToolRegistry._dynamic_modules = {} ToolRegistry._dynamic_tools_by_module = {} yield @@ -50,11 +55,13 @@ def _isolated_registry() -> Iterator[None]: ToolRegistry._tools = state["tools"] ToolRegistry._enabled_defaults = state["defaults"] ToolRegistry._plugin_tool_names = state["plugin_names"] + ToolRegistry._plugin_load_errors = state["plugin_errors"] + ToolRegistry._revision = state["revision"] ToolRegistry._dynamic_modules = state["dynamic_modules"] ToolRegistry._dynamic_tools_by_module = state["dynamic_tools"] -def test_plugin_refresh_restores_last_known_good_on_load_error( +def test_plugin_refresh_restores_last_known_good_on_loader_failure( monkeypatch: pytest.MonkeyPatch, ): with _isolated_registry(): @@ -66,17 +73,257 @@ def _failed_load(cls, errors): replacement = _tool("partial_replacement", "plugin_py") cls._tools[replacement.info.name] = replacement cls._plugin_tool_names = [replacement.info.name] - errors.append("broken plugin") + errors.append("plugin loader: unavailable") monkeypatch.setattr(ToolRegistry, "_load_plugin_tools", classmethod(_failed_load)) - with pytest.raises(ToolRefreshError, match="broken plugin"): + with pytest.raises(ToolRefreshError, match="plugin loader: unavailable"): ToolRegistry.refresh_plugin_tools() assert ToolRegistry._tools == {previous.info.name: previous} assert ToolRegistry._plugin_tool_names == [previous.info.name] +def test_plugin_refresh_accepts_later_unrelated_syntax_error( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +): + from flocks.plugin import PluginLoader + + with _isolated_registry(): + project_dir = tmp_path / "project" + plugin_root = tmp_path / "plugins" + plugin_dir = plugin_root / "tools" / "device" / "legacy_device" + project_dir.mkdir() + plugin_dir.mkdir(parents=True) + + monkeypatch.chdir(project_dir) + monkeypatch.setattr(PluginLoader, "_plugin_root", plugin_root) + monkeypatch.setattr( + ToolRegistry, + "_finalize_plugin_tools_load", + classmethod(lambda cls: None), + ) + monkeypatch.setattr(ToolRegistry, "_bump_revision", lambda _reason: None) + extension_points_before = PluginLoader._extension_points.copy() + try: + ToolRegistry._register_plugin_extension_point() + startup_errors = ToolRegistry._load_plugin_tools() + assert startup_errors == [] + ToolRegistry._plugin_load_errors = startup_errors + + (plugin_dir / "broken.handler.py").write_text( + "def broken(:\n", + encoding="utf-8", + ) + + (plugin_root / "tools" / "newly_installed.py").write_text( + "async def handle(ctx):\n" + " return None\n\n" + "TOOLS = [{\n" + " 'name': 'newly_installed',\n" + " 'description': 'newly installed tool',\n" + " 'handler': handle,\n" + "}]\n", + encoding="utf-8", + ) + + refreshed = ToolRegistry.refresh_plugin_tools( + changed_path=plugin_root / "tools" / "newly_installed.py", + ) + + assert "newly_installed" in refreshed + assert ToolRegistry.get("newly_installed") is not None + assert any("SyntaxError" in error for error in ToolRegistry._plugin_load_errors) + finally: + PluginLoader.clear_extension_points() + PluginLoader._extension_points.update(extension_points_before) + + +def test_registry_init_records_plugin_load_errors( + monkeypatch: pytest.MonkeyPatch, +): + with _isolated_registry(): + ToolRegistry._initialized = False + monkeypatch.setattr( + ToolRegistry, + "_register_builtin_tools", + classmethod(lambda cls: None), + ) + monkeypatch.setattr( + ToolRegistry, + "_register_dynamic_tools", + classmethod(lambda cls, errors=None: None), + ) + monkeypatch.setattr( + ToolRegistry, + "_register_plugin_extension_point", + classmethod(lambda cls: None), + ) + monkeypatch.setattr( + ToolRegistry, + "_load_plugin_tools", + classmethod(lambda cls: ["broken legacy plugin"]), + ) + + ToolRegistry.init() + + assert ToolRegistry._plugin_load_errors == ["broken legacy plugin"] + + +def test_plugin_refresh_commits_valid_sources_alongside_new_source_error( + monkeypatch: pytest.MonkeyPatch, +): + with _isolated_registry(): + previous = _tool("previous_plugin", "plugin_py") + ToolRegistry._tools[previous.info.name] = previous + ToolRegistry._plugin_tool_names = [previous.info.name] + ToolRegistry._plugin_load_errors = ["/plugins/legacy.py: broken legacy plugin"] + + def _load_with_new_error(cls, errors): + replacement = _tool("partial_replacement", "plugin_py") + cls._tools[replacement.info.name] = replacement + cls._plugin_tool_names = [replacement.info.name] + errors.extend( + [ + "/plugins/legacy.py: broken legacy plugin", + "/plugins/unrelated.py: new broken plugin", + "entry point legacy_package: RuntimeError: broken", + ] + ) + + monkeypatch.setattr( + ToolRegistry, + "_load_plugin_tools", + classmethod(_load_with_new_error), + ) + + refreshed = ToolRegistry.refresh_plugin_tools( + changed_path=Path("/plugins/newly_installed.py"), + ) + + assert refreshed == ["partial_replacement"] + assert set(ToolRegistry._tools) == {"partial_replacement"} + assert ToolRegistry._plugin_tool_names == ["partial_replacement"] + assert ToolRegistry._plugin_load_errors == [ + "/plugins/legacy.py: broken legacy plugin", + "/plugins/unrelated.py: new broken plugin", + "entry point legacy_package: RuntimeError: broken", + ] + + +def test_plugin_refresh_rejects_error_inside_changed_path( + monkeypatch: pytest.MonkeyPatch, +): + with _isolated_registry(): + previous = _tool("previous_plugin", "plugin_py") + ToolRegistry._tools[previous.info.name] = previous + ToolRegistry._plugin_tool_names = [previous.info.name] + + def _load_target_error(cls, errors): + errors.append("/plugins/new/broken.py: ImportError: missing dependency") + + monkeypatch.setattr( + ToolRegistry, + "_load_plugin_tools", + classmethod(_load_target_error), + ) + + with pytest.raises(ToolRefreshError, match="missing dependency"): + ToolRegistry.refresh_plugin_tools(changed_path=Path("/plugins/new")) + + assert ToolRegistry._tools == {previous.info.name: previous} + assert ToolRegistry._plugin_tool_names == [previous.info.name] + + +def test_plugin_refresh_restores_snapshot_when_finalize_fails( + monkeypatch: pytest.MonkeyPatch, +): + with _isolated_registry(): + builtin = _tool("builtin", "builtin") + previous = _tool("previous_plugin", "plugin_py") + ToolRegistry._tools = { + builtin.info.name: builtin, + previous.info.name: previous, + } + ToolRegistry._enabled_defaults = { + builtin.info.name: True, + previous.info.name: True, + } + ToolRegistry._plugin_tool_names = [previous.info.name] + ToolRegistry._plugin_load_errors = ["old error"] + ToolRegistry._revision = 7 + + def _load_with_unrelated_error(cls, errors): + replacement = _tool("replacement", "plugin_py") + cls._tools[replacement.info.name] = replacement + cls._enabled_defaults[replacement.info.name] = True + cls._plugin_tool_names = [replacement.info.name] + errors.append("/plugins/unrelated.py: SyntaxError: broken") + + def _failed_finalize(cls): + cls._tools[builtin.info.name].info.enabled = False + raise RuntimeError("finalize failed") + + monkeypatch.setattr( + ToolRegistry, + "_load_plugin_tools", + classmethod(_load_with_unrelated_error), + ) + monkeypatch.setattr( + ToolRegistry, + "_finalize_plugin_tools_load", + classmethod(_failed_finalize), + ) + + with pytest.raises(RuntimeError, match="finalize failed"): + ToolRegistry.refresh_plugin_tools( + changed_path=Path("/plugins/target.py"), + ) + + assert ToolRegistry._tools == { + builtin.info.name: builtin, + previous.info.name: previous, + } + assert builtin.info.enabled is True + assert ToolRegistry._enabled_defaults == { + builtin.info.name: True, + previous.info.name: True, + } + assert ToolRegistry._plugin_tool_names == [previous.info.name] + assert ToolRegistry._plugin_load_errors == ["old error"] + assert ToolRegistry._revision == 7 + + +def test_plugin_refresh_rejects_entry_point_scan_failure( + monkeypatch: pytest.MonkeyPatch, +): + with _isolated_registry(): + previous = _tool("entrypoint_old", "plugin_py") + ToolRegistry._tools[previous.info.name] = previous + ToolRegistry._plugin_tool_names = [previous.info.name] + + def _load_with_scan_failure(cls, errors): + replacement = _tool("hub_target", "plugin_py") + cls._tools[replacement.info.name] = replacement + cls._plugin_tool_names = [replacement.info.name] + errors.append("entry point scan: RuntimeError: metadata unavailable") + + monkeypatch.setattr( + ToolRegistry, + "_load_plugin_tools", + classmethod(_load_with_scan_failure), + ) + + with pytest.raises(ToolRefreshError, match="metadata unavailable"): + ToolRegistry.refresh_plugin_tools( + changed_path=Path("/plugins/hub_target.py"), + ) + + assert ToolRegistry._tools == {previous.info.name: previous} + assert ToolRegistry._plugin_tool_names == [previous.info.name] + + def test_dynamic_refresh_restores_last_known_good_on_load_error( monkeypatch: pytest.MonkeyPatch, ): @@ -121,7 +368,7 @@ def test_refresh_transactions_are_serialized_across_stages( def _failed_plugin_load(cls, errors): plugin_load_entered.set() assert allow_plugin_failure.wait(timeout=2) - errors.append("broken plugin") + errors.append("plugin loader: unavailable") def _successful_dynamic_load(cls, _errors): dynamic_load_entered.set() @@ -199,7 +446,7 @@ def test_registry_readers_do_not_observe_an_in_progress_refresh( def _failed_plugin_load(cls, errors): load_entered.set() assert allow_failure.wait(timeout=2) - errors.append("broken plugin") + errors.append("plugin loader: unavailable") monkeypatch.setattr( ToolRegistry, diff --git a/tests/tool/test_tools.py b/tests/tool/test_tools.py index e63e97cf7..895cf0f80 100644 --- a/tests/tool/test_tools.py +++ b/tests/tool/test_tools.py @@ -20,6 +20,7 @@ import uuid from pathlib import Path from typing import Dict, Any, List +from unittest.mock import MagicMock # Import the tool system from flocks.tool import ( @@ -35,6 +36,8 @@ ParameterType, ) from flocks.tool.code import bash as bash_module +import flocks.tool.file.glob as glob_module +import flocks.tool.registry as registry_module import flocks.tool.system.question as question_module @@ -579,6 +582,33 @@ async def test_glob_no_matches(self, tool_context, temp_dir): assert result.success assert "No files found" in result.output + @pytest.mark.asyncio + async def test_python_fallback_logs_info_once(self, monkeypatch, tool_context, temp_dir): + monkeypatch.setattr(glob_module, "_ripgrep_fallback_logged", False) + monkeypatch.setattr(glob_module, "find_ripgrep", lambda: None) + info_log = MagicMock() + warn_log = MagicMock() + monkeypatch.setattr(glob_module.log, "info", info_log) + monkeypatch.setattr(glob_module.log, "warn", warn_log) + + for _ in range(2): + result = await glob_module.glob_tool( + tool_context, + pattern="*.txt", + path=temp_dir, + ) + assert result.success + + fallback_calls = [ + call for call in info_log.call_args_list + if call.args and call.args[0] == "glob.ripgrep_not_found" + ] + assert len(fallback_calls) == 1 + assert not any( + call.args and call.args[0] == "glob.ripgrep_not_found" + for call in warn_log.call_args_list + ) + # ============================================================================= # P1 Tools Tests @@ -1030,8 +1060,13 @@ class TestErrorHandling: """Test error handling across tools""" @pytest.mark.asyncio - async def test_missing_required_parameter(self, tool_context): + async def test_missing_required_parameter(self, monkeypatch, tool_context): """Test error when required parameter is missing""" + warn_log = MagicMock() + error_log = MagicMock() + monkeypatch.setattr(registry_module.log, "warn", warn_log) + monkeypatch.setattr(registry_module.log, "error", error_log) + result = await ToolRegistry.execute( "read", ctx=tool_context @@ -1040,6 +1075,18 @@ async def test_missing_required_parameter(self, tool_context): assert not result.success assert "required" in result.error.lower() or "missing" in result.error.lower() + warn_log.assert_called_once_with( + "tool.execute.missing_param", + { + "tool": "read", + "missing": "filePath", + "provided": [], + }, + ) + assert not any( + call.args and call.args[0] == "tool.execute.missing_param" + for call in error_log.call_args_list + ) @pytest.mark.asyncio async def test_nonexistent_tool(self, tool_context): diff --git a/tests/updater/test_restart_handoff.py b/tests/updater/test_restart_handoff.py index e9df3cbb7..b044d3baa 100644 --- a/tests/updater/test_restart_handoff.py +++ b/tests/updater/test_restart_handoff.py @@ -1,4 +1,6 @@ +import json import shutil +import subprocess import sys from pathlib import Path from types import SimpleNamespace @@ -6,11 +8,17 @@ import pytest from flocks.cli import service_manager +from flocks.cli.service_config import service_config_payload from flocks.updater import restart_handoff -from tests.helpers.service_supervisor import make_short_runtime_root, start_supervisor, stop_supervisor, wait_for_supervisor +from tests.helpers.service_supervisor import ( + make_short_runtime_root, + start_supervisor, + stop_supervisor, + wait_for_supervisor, +) -def _handoff_args(tmp_path: Path, restart_argv: list[str], *, prepare_handover: bool = False) -> list[str]: +def _handoff_args(tmp_path: Path, restart_argv: list[str]) -> list[str]: args = [ "--parent-pid", "1234", @@ -33,11 +41,423 @@ def _handoff_args(tmp_path: Path, restart_argv: list[str], *, prepare_handover: "--current-version", "2026.3.31", ] - if prepare_handover: - args.append("--prepare-handover") return [*args, "--", *restart_argv] +def _simple_upgrade_handoff_args(tmp_path: Path) -> list[str]: + content_root = tmp_path / "staged" + content_root.mkdir() + backup_path = tmp_path / "backup.tar.gz" + backup_path.write_text("backup", encoding="utf-8") + config = service_manager.ServiceConfig( + backend_host="10.0.0.8", + backend_port=5273, + frontend_host="10.0.0.8", + frontend_port=5273, + legacy_backend_host="0.0.0.0", + legacy_backend_port=9000, + no_browser=True, + skip_frontend_build=True, + ) + return [ + "--mode", + "upgrade", + "--parent-pid", + "1234", + "--backend-host", + config.backend_host, + "--backend-port", + str(config.backend_port), + "--frontend-host", + config.frontend_host, + "--frontend-port", + str(config.frontend_port), + "--install-root", + str(tmp_path / "install"), + "--content-root", + str(content_root), + "--backup-path", + str(backup_path), + "--was-running", + "--daemon-pid", + "2468", + "--service-config-json", + json.dumps(service_config_payload(config)), + "--uv-path", + "uv", + "--sync-timeout", + "300", + "--version", + "2026.4.1", + "--current-version", + "2026.3.31", + "--", + "/install/.venv/bin/python", + ] + + +def _v2026_7_1_handoff_args(tmp_path: Path, restart_argv: list[str]) -> list[str]: + """Build the handoff protocol emitted by the v2026.7.1 updater.""" + args = _handoff_args(tmp_path, restart_argv) + args[args.index("--install-root") : args.index("--install-root")] = [ + "--backend-pid-file", + str(tmp_path / "backend.pid"), + ] + separator_index = args.index("--") + args[separator_index:separator_index] = [ + "--backup-path", + str(tmp_path / "backup.tar.gz"), + ] + return args + + +def _v2026_7_15_handoff_args(tmp_path: Path, restart_argv: list[str]) -> list[str]: + """Build the handoff protocol emitted by the v2026.7.15 updater.""" + args = _handoff_args(tmp_path, restart_argv) + args[args.index("--backend-port") + 1] = "5173" + separator_index = args.index("--") + args[separator_index:separator_index] = [ + "--backup-path", + str(tmp_path / "backup.tar.gz"), + "--prepare-handover", + ] + return args + + +def test_current_version_upgrade_handoff_stops_replaces_installs_and_restarts( + monkeypatch, + tmp_path: Path, +) -> None: + events: list[str] = [] + cleanup_dir = tmp_path / "cleanup" + cleanup_dir.mkdir() + args = _simple_upgrade_handoff_args(tmp_path) + args[args.index("--") : args.index("--")] = ["--cleanup-dir", str(cleanup_dir)] + + monkeypatch.setattr(restart_handoff, "_record_handoff_log", lambda _message: None) + monkeypatch.setattr( + restart_handoff, + "_wait_for_parent_exit", + lambda parent_pid: events.append(f"wait-parent:{parent_pid}") or True, + ) + monkeypatch.setattr( + restart_handoff, + "_stop_services_before_upgrade", + lambda args: events.append(f"stop:{args.daemon_pid}") or True, + ) + monkeypatch.setattr( + restart_handoff, + "_apply_new_source", + lambda args: events.append("replace"), + ) + monkeypatch.setattr( + restart_handoff, + "_run_upgrade_tasks", + lambda args: events.append("install") or None, + ) + monkeypatch.setattr( + restart_handoff, + "_start_service_after_upgrade", + lambda args: events.append("start") or (True, "", ""), + ) + monkeypatch.setattr(restart_handoff, "_write_upgrade_result", lambda **_kwargs: None) + + assert restart_handoff.run(args) == 0 + assert events == [ + "wait-parent:1234", + "stop:2468", + "replace", + "install", + "start", + ] + assert not cleanup_dir.exists() + + +def test_current_version_upgrade_handoff_can_run_without_waiting_for_parent( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + events: list[str] = [] + args = _simple_upgrade_handoff_args(tmp_path) + parent_index = args.index("--parent-pid") + del args[parent_index : parent_index + 2] + + monkeypatch.setattr(restart_handoff, "_record_handoff_log", lambda _message: None) + monkeypatch.setattr( + restart_handoff, + "_wait_for_parent_exit", + lambda _parent_pid: pytest.fail("parent wait must be skipped"), + ) + monkeypatch.setattr( + restart_handoff, + "_stop_services_before_upgrade", + lambda _args: events.append("stop") or True, + ) + monkeypatch.setattr(restart_handoff, "_apply_new_source", lambda _args: events.append("replace")) + monkeypatch.setattr( + restart_handoff, + "_run_upgrade_tasks", + lambda _args: events.append("install") or None, + ) + monkeypatch.setattr( + restart_handoff, + "_start_service_after_upgrade", + lambda _args: events.append("start") or (True, "", ""), + ) + monkeypatch.setattr(restart_handoff, "_write_upgrade_result", lambda **_kwargs: None) + + assert restart_handoff.run(args) == 0 + assert events == ["stop", "replace", "install", "start"] + + +@pytest.mark.parametrize( + ("host", "port"), + [ + ("127.0.0.1", 5173), + ("0.0.0.0", 5273), + ("10.20.30.40", 9527), + ("::", 6173), + ("2001:db8::20", 7173), + ], +) +def test_upgrade_start_argv_uses_captured_host_port_without_control_api(host: str, port: int) -> None: + config = service_manager.ServiceConfig( + backend_host=host, + backend_port=port, + frontend_host=host, + frontend_port=port, + legacy_backend_host="0.0.0.0", + legacy_backend_port=9000, + no_browser=True, + skip_frontend_build=True, + ) + args = SimpleNamespace( + restart_argv=["/install/.venv/bin/python"], + service_config_json=json.dumps(service_config_payload(config)), + ) + + assert restart_handoff._build_captured_start_argv(args) == [ + "/install/.venv/bin/python", + "-m", + "flocks.cli.main", + "start", + "--host", + host, + "--port", + str(port), + "--no-browser", + "--skip-webui-build", + "--server-host", + "0.0.0.0", + "--server-port", + "9000", + ] + + +@pytest.mark.parametrize( + ("public_host", "public_port"), + [ + ("0.0.0.0", 5173), + ("10.20.30.40", 9527), + ("::", 6173), + ("2001:db8::20", 7173), + ], +) +def test_upgrade_start_argv_preserves_distinct_legacy_public_endpoint( + public_host: str, + public_port: int, +) -> None: + config = service_manager.ServiceConfig( + backend_host="127.0.0.1", + backend_port=8000, + frontend_host=public_host, + frontend_port=public_port, + legacy_backend_host="127.0.0.1", + legacy_backend_port=8000, + no_browser=True, + skip_frontend_build=True, + ) + args = SimpleNamespace( + restart_argv=["/install/.venv/bin/python"], + service_config_json=json.dumps(service_config_payload(config)), + ) + + assert restart_handoff._build_captured_start_argv(args) == [ + "/install/.venv/bin/python", + "-m", + "flocks.cli.main", + "start", + "--host", + public_host, + "--port", + str(public_port), + "--no-browser", + "--skip-webui-build", + "--server-host", + "127.0.0.1", + "--server-port", + "8000", + ] + + +def test_upgrade_wait_ports_exclude_legacy_cleanup_port(tmp_path: Path) -> None: + args = restart_handoff._parse_args(_simple_upgrade_handoff_args(tmp_path)) + + assert restart_handoff._service_ports(args) == (5273,) + + +def test_upgrade_install_failure_keeps_backup_and_temp_without_restart_or_rollback( + monkeypatch, + tmp_path: Path, +) -> None: + events: list[str] = [] + results: list[dict[str, object]] = [] + cleanup_dir = tmp_path / "cleanup" + cleanup_dir.mkdir() + backup_path = tmp_path / "backup.tar.gz" + backup_path.write_text("backup", encoding="utf-8") + args = _simple_upgrade_handoff_args(tmp_path) + args[args.index("--") : args.index("--")] = ["--cleanup-dir", str(cleanup_dir)] + + monkeypatch.setattr(restart_handoff, "_wait_for_parent_exit", lambda _pid: True) + monkeypatch.setattr(restart_handoff, "_stop_services_before_upgrade", lambda _args: True) + monkeypatch.setattr( + restart_handoff, + "_apply_new_source", + lambda _args: events.append("replace"), + ) + monkeypatch.setattr(restart_handoff, "_run_upgrade_tasks", lambda _args: "npm build failed") + monkeypatch.setattr( + restart_handoff, + "_start_service_after_upgrade", + lambda _args: events.append("start") or (True, "", ""), + ) + monkeypatch.setattr(restart_handoff, "_record_handoff_log", lambda message: events.append(message)) + monkeypatch.setattr(restart_handoff, "_write_upgrade_result", lambda **kwargs: results.append(kwargs)) + + assert restart_handoff.run(args) == 1 + assert events[0] == "replace" + assert "start" not in events + assert not any("rollback" in event for event in events) + assert cleanup_dir.exists() + assert backup_path.exists() + assert results[-1]["failed_stage"] == "install" + assert results[-1]["backup_path"] == backup_path + + +def test_upgrade_start_failure_records_process_output_and_keeps_temp( + monkeypatch, + tmp_path: Path, +) -> None: + results: list[dict[str, object]] = [] + cleanup_dir = tmp_path / "cleanup" + cleanup_dir.mkdir() + backup_path = tmp_path / "backup.tar.gz" + backup_path.write_text("backup", encoding="utf-8") + args = _simple_upgrade_handoff_args(tmp_path) + args[args.index("--") : args.index("--")] = ["--cleanup-dir", str(cleanup_dir)] + + monkeypatch.setattr(restart_handoff, "_wait_for_parent_exit", lambda _pid: True) + monkeypatch.setattr(restart_handoff, "_stop_services_before_upgrade", lambda _args: True) + monkeypatch.setattr(restart_handoff, "_apply_new_source", lambda _args: None) + monkeypatch.setattr(restart_handoff, "_run_upgrade_tasks", lambda _args: None) + monkeypatch.setattr(restart_handoff, "_report_pending_pro_bundle_install_receipt", lambda _args: None) + monkeypatch.setattr( + restart_handoff, + "_start_service_after_upgrade", + lambda _args: (False, "start stdout", "start stderr"), + ) + monkeypatch.setattr(restart_handoff, "_write_upgrade_result", lambda **kwargs: results.append(kwargs)) + + assert restart_handoff.run(args) == 1 + assert cleanup_dir.exists() + assert results[-1]["failed_stage"] == "start" + assert results[-1]["stdout"] == "start stdout" + assert results[-1]["stderr"] == "start stderr" + + +def test_upgrade_writes_running_after_parent_exits(monkeypatch, tmp_path: Path) -> None: + events: list[str] = [] + args = _simple_upgrade_handoff_args(tmp_path) + + monkeypatch.setattr( + restart_handoff, + "_wait_for_parent_exit", + lambda _pid: events.append("parent-exited") or True, + ) + monkeypatch.setattr( + restart_handoff, + "_write_upgrade_result", + lambda **kwargs: events.append(f"write:{kwargs['phase']}"), + ) + monkeypatch.setattr( + restart_handoff, + "_stop_services_before_upgrade", + lambda _args: events.append("stop") or False, + ) + monkeypatch.setattr(restart_handoff, "_record_handoff_log", lambda _message: None) + + assert restart_handoff.run(args) == 1 + assert events[:3] == ["parent-exited", "write:running", "stop"] + + +def test_upgrade_result_write_failure_is_non_fatal(monkeypatch, tmp_path: Path) -> None: + events: list[str] = [] + args = restart_handoff._parse_args(_simple_upgrade_handoff_args(tmp_path)) + + monkeypatch.setattr( + restart_handoff.updater_module, + "_write_upgrade_result_state", + lambda _payload: (_ for _ in ()).throw(OSError("disk unavailable")), + ) + monkeypatch.setattr( + restart_handoff, + "_record_handoff_log", + lambda message: events.append(message), + ) + + restart_handoff._write_upgrade_result(args=args, phase="running") + + assert events == ["upgrade_result_write_failed phase=running error=disk unavailable"] + + +def test_upgrade_that_was_stopped_does_not_start_service(monkeypatch, tmp_path: Path) -> None: + args = _simple_upgrade_handoff_args(tmp_path) + args.remove("--was-running") + parsed = restart_handoff._parse_args(args) + monkeypatch.setattr( + restart_handoff.subprocess, + "run", + lambda *_args, **_kwargs: pytest.fail("stopped service must remain stopped"), + ) + + assert restart_handoff._start_service_after_upgrade(parsed) == (True, "", "") + + +def test_upgrade_start_decodes_invalid_windows_output_without_crashing(monkeypatch, tmp_path: Path) -> None: + parsed = restart_handoff._parse_args(_simple_upgrade_handoff_args(tmp_path)) + logs: list[str] = [] + + def fake_run(command, **kwargs): + assert kwargs["text"] is False + return subprocess.CompletedProcess( + command, + 1, + stdout=b"start stdout \x80", + stderr=b"start stderr \x81", + ) + + monkeypatch.setattr(restart_handoff.subprocess, "run", fake_run) + monkeypatch.setattr(restart_handoff, "_record_handoff_log", logs.append) + + assert restart_handoff._start_service_after_upgrade(parsed) == ( + False, + "start stdout �", + "start stderr �", + ) + assert logs == ["restart_failed returncode=1 stdout=start stdout � stderr=start stderr �"] + + def test_run_waits_for_parent_and_backend_port_before_spawning( monkeypatch, tmp_path: Path, @@ -75,8 +495,9 @@ def test_run_waits_for_parent_and_backend_port_before_spawning( monkeypatch.setattr( restart_handoff.subprocess, "Popen", - lambda argv, cwd=None, close_fds=False: events.append(f"spawn:{list(argv)}:{cwd}:{close_fds}") - or SimpleNamespace(pid=4321), + lambda argv, cwd=None, close_fds=False: ( + events.append(f"spawn:{list(argv)}:{cwd}:{close_fds}") or SimpleNamespace(pid=4321) + ), ) monkeypatch.setattr(restart_handoff, "_run_upgrade_tasks", lambda args: events.append("tasks") or None) monkeypatch.setattr( @@ -123,8 +544,9 @@ def test_run_keeps_current_start_restart_argv(monkeypatch, tmp_path: Path) -> No monkeypatch.setattr( restart_handoff.subprocess, "Popen", - lambda argv, cwd=None, close_fds=False: events.append(f"spawn:{list(argv)}:{cwd}:{close_fds}") - or SimpleNamespace(pid=4321), + lambda argv, cwd=None, close_fds=False: ( + events.append(f"spawn:{list(argv)}:{cwd}:{close_fds}") or SimpleNamespace(pid=4321) + ), ) code = restart_handoff.run(_handoff_args(tmp_path, restart_argv)) @@ -137,7 +559,7 @@ def test_run_accepts_legacy_backend_pid_file_argument(monkeypatch, tmp_path: Pat events: list[str] = [] restart_argv = ["python.exe", "-m", "flocks.cli.main", "start"] args = _handoff_args(tmp_path, restart_argv) - args[args.index("--install-root"):args.index("--install-root")] = [ + args[args.index("--install-root") : args.index("--install-root")] = [ "--backend-pid-file", str(tmp_path / "backend.pid"), ] @@ -146,12 +568,14 @@ def test_run_accepts_legacy_backend_pid_file_argument(monkeypatch, tmp_path: Pat monkeypatch.setattr(restart_handoff, "_wait_for_parent_exit", lambda parent_pid: True) monkeypatch.setattr(restart_handoff, "_ensure_backend_port_free", lambda backend_port: True) monkeypatch.setattr(restart_handoff, "_run_upgrade_tasks", lambda args: None) + monkeypatch.setattr(restart_handoff, "_cleanup_legacy_upgrade_handover", lambda args: True) monkeypatch.setattr(restart_handoff, "_stop_supervisor_before_restart", lambda: True) monkeypatch.setattr( restart_handoff.subprocess, "Popen", - lambda argv, cwd=None, close_fds=False: events.append(f"spawn:{list(argv)}:{cwd}:{close_fds}") - or SimpleNamespace(pid=4321), + lambda argv, cwd=None, close_fds=False: ( + events.append(f"spawn:{list(argv)}:{cwd}:{close_fds}") or SimpleNamespace(pid=4321) + ), ) code = restart_handoff.run(args) @@ -160,14 +584,122 @@ def test_run_accepts_legacy_backend_pid_file_argument(monkeypatch, tmp_path: Pat assert f"spawn:{restart_argv}:{tmp_path}:True" in events -def test_run_prepares_handover_after_parent_exit_without_waiting_for_page_port( +def test_v2026_7_1_upgrade_handoff_runs_tasks_and_restarts(monkeypatch, tmp_path: Path) -> None: + events: list[str] = [] + flocks_root = tmp_path / "flocks-root" + upgrade_state = flocks_root / "run" / "upgrade-state.json" + upgrade_state.parent.mkdir(parents=True) + upgrade_state.write_text( + json.dumps( + { + "backend_host": "127.0.0.1", + "backend_port": 8000, + "frontend_host": "0.0.0.0", + "frontend_port": 8888, + } + ), + encoding="utf-8", + ) + restart_argv = [ + "python.exe", + "-m", + "flocks.cli.main", + "serve", + "--host", + "127.0.0.1", + "--port", + "8000", + ] + expected_restart_argv = [ + "python.exe", + "-m", + "flocks.cli.main", + "start", + "--no-browser", + "--skip-webui-build", + "--host", + "0.0.0.0", + "--port", + "8888", + "--server-host", + "127.0.0.1", + "--server-port", + "8000", + ] + args = _v2026_7_1_handoff_args(tmp_path, restart_argv) + + monkeypatch.setattr(restart_handoff.updater_module, "_flocks_root", lambda: flocks_root) + monkeypatch.setattr(restart_handoff, "_record_handoff_log", lambda _message: None) + monkeypatch.setattr( + restart_handoff, + "_wait_for_parent_exit", + lambda parent_pid: events.append(f"wait-parent:{parent_pid}") or True, + ) + monkeypatch.setattr( + restart_handoff, + "_ensure_backend_port_free", + lambda backend_port: events.append(f"free-port:{backend_port}") or True, + ) + monkeypatch.setattr( + restart_handoff, + "_run_upgrade_tasks", + lambda _args: events.append("install") or None, + ) + monkeypatch.setattr( + restart_handoff, + "_cleanup_legacy_upgrade_handover", + lambda _args: events.append("cleanup-handover") or True, + ) + monkeypatch.setattr( + restart_handoff, + "_stop_supervisor_before_restart", + lambda: events.append("stop-supervisor") or True, + ) + monkeypatch.setattr( + restart_handoff.subprocess, + "Popen", + lambda argv, cwd=None, close_fds=False: ( + events.append(f"spawn:{list(argv)}:{cwd}:{close_fds}") or SimpleNamespace(pid=4321) + ), + ) + + assert restart_handoff.run(args) == 0 + assert events == [ + "wait-parent:1234", + "free-port:8000", + "install", + "cleanup-handover", + "stop-supervisor", + f"spawn:{expected_restart_argv}:{tmp_path}:True", + ] + + +def test_v2026_7_15_upgrade_handoff_stops_before_tasks_and_restarts( monkeypatch, tmp_path: Path, ) -> None: events: list[str] = [] - restart_argv = ["python.exe", "-m", "flocks.cli.main", "start"] + restart_argv = [ + "python.exe", + "-m", + "flocks.cli.main", + "start", + "--no-browser", + "--skip-webui-build", + "--host", + "0.0.0.0", + "--port", + "5173", + "--server-host", + "127.0.0.1", + "--server-port", + "8000", + ] + args = _v2026_7_15_handoff_args(tmp_path, restart_argv) + args[args.index("--backend-host") + 1] = "0.0.0.0" + args[args.index("--frontend-host") + 1] = "0.0.0.0" - monkeypatch.setattr(restart_handoff, "_record_handoff_log", lambda message: events.append(f"log:{message}")) + monkeypatch.setattr(restart_handoff, "_record_handoff_log", lambda _message: None) monkeypatch.setattr( restart_handoff, "_wait_for_parent_exit", @@ -175,8 +707,91 @@ def test_run_prepares_handover_after_parent_exit_without_waiting_for_page_port( ) monkeypatch.setattr( restart_handoff, - "_prepare_upgrade_handover", - lambda args: events.append(f"prepare:{args.version}") or True, + "_ensure_backend_port_free", + lambda _backend_port: pytest.fail("legacy handover must stop the supervisor first"), + ) + monkeypatch.setattr( + restart_handoff, + "_stop_supervisor_before_restart", + lambda **kwargs: events.append(f"stop-supervisor:{kwargs}") or True, + ) + monkeypatch.setattr(restart_handoff, "_legacy_supervisor_pid", lambda _args: 2468) + monkeypatch.setattr( + restart_handoff, + "_run_upgrade_tasks", + lambda _args: events.append("install") or None, + ) + monkeypatch.setattr( + restart_handoff, + "_cleanup_legacy_upgrade_handover", + lambda _args: events.append("cleanup-handover") or True, + ) + monkeypatch.setattr( + restart_handoff.subprocess, + "Popen", + lambda argv, cwd=None, close_fds=False: ( + events.append(f"spawn:{list(argv)}:{cwd}:{close_fds}") or SimpleNamespace(pid=4321) + ), + ) + + assert restart_handoff.run(args) == 0 + assert events == [ + "wait-parent:1234", + ( + "stop-supervisor:{'daemon_pid': 2468, 'backend_port': 5173, " + "'service_ports': (5173,), 'force_daemon_stop': True}" + ), + "install", + "cleanup-handover", + f"spawn:{restart_argv}:{tmp_path}:True", + ] + + +def test_cleanup_legacy_upgrade_handover_stops_trusted_page(monkeypatch, tmp_path: Path) -> None: + run_dir = tmp_path / "run" + page_dir = run_dir / "upgrade-page" + page_dir.mkdir(parents=True) + state_path = run_dir / "upgrade-state.json" + state_path.write_text("{}", encoding="utf-8") + pid_path = run_dir / "upgrade_server.pid" + pid_path.write_text("2468", encoding="utf-8") + alive = {2468} + + monkeypatch.setattr(restart_handoff.updater_module, "_flocks_root", lambda: tmp_path) + monkeypatch.setattr( + restart_handoff.service_manager, + "port_owner_pids", + lambda _port: sorted(alive), + ) + monkeypatch.setattr( + restart_handoff.service_manager, + "_process_command_line", + lambda pid: f"python -m http.server --directory {page_dir}" if pid in alive else "", + ) + monkeypatch.setattr( + restart_handoff.service_manager, + "_terminate_orphan_pid", + lambda pid, _label, _console: alive.discard(pid), + ) + + assert restart_handoff._cleanup_legacy_upgrade_handover(SimpleNamespace(frontend_port=5173)) + assert not state_path.exists() + assert not pid_path.exists() + assert not page_dir.exists() + + +def test_restart_only_waits_for_port_after_parent_exit( + monkeypatch, + tmp_path: Path, +) -> None: + events: list[str] = [] + restart_argv = ["python.exe", "-m", "flocks.cli.main", "start"] + + monkeypatch.setattr(restart_handoff, "_record_handoff_log", lambda message: events.append(f"log:{message}")) + monkeypatch.setattr( + restart_handoff, + "_wait_for_parent_exit", + lambda parent_pid: events.append(f"wait-parent:{parent_pid}") or True, ) monkeypatch.setattr( restart_handoff, @@ -192,16 +807,17 @@ def test_run_prepares_handover_after_parent_exit_without_waiting_for_page_port( monkeypatch.setattr( restart_handoff.subprocess, "Popen", - lambda argv, cwd=None, close_fds=False: events.append(f"spawn:{list(argv)}:{cwd}:{close_fds}") - or SimpleNamespace(pid=4321), + lambda argv, cwd=None, close_fds=False: ( + events.append(f"spawn:{list(argv)}:{cwd}:{close_fds}") or SimpleNamespace(pid=4321) + ), ) - code = restart_handoff.run(_handoff_args(tmp_path, restart_argv, prepare_handover=True)) + code = restart_handoff.run(_handoff_args(tmp_path, restart_argv)) assert code == 0 assert events[1:] == [ "wait-parent:1234", - "prepare:2026.4.1", + "free-port:8000", "tasks", "stop-supervisor", f"spawn:{restart_argv}:{tmp_path}:True", @@ -223,7 +839,8 @@ def test_run_reports_pending_install_receipt_after_pro_bundle_tasks( monkeypatch.setattr(restart_handoff, "_record_handoff_log", lambda message: events.append(f"log:{message}")) monkeypatch.setattr(restart_handoff, "_wait_for_parent_exit", lambda parent_pid: True) - monkeypatch.setattr(restart_handoff, "_ensure_backend_port_free", lambda backend_port, backend_pid_file: True) + monkeypatch.setattr(restart_handoff, "_stop_supervisor_before_restart", lambda **_kwargs: True) + monkeypatch.setattr(restart_handoff, "_ensure_backend_port_free", lambda backend_port: True) monkeypatch.setattr(restart_handoff, "_run_upgrade_tasks", lambda args: events.append("tasks") or None) monkeypatch.setattr( restart_handoff, @@ -235,12 +852,66 @@ def test_run_reports_pending_install_receipt_after_pro_bundle_tasks( "Popen", lambda argv, cwd=None, close_fds=False: events.append(f"spawn:{list(argv)}") or SimpleNamespace(pid=4321), ) - monkeypatch.setattr(restart_handoff, "_record_backend_runtime_if_direct_serve", lambda *_args, **_kwargs: None) - code = restart_handoff.run(args) assert code == 0 - assert events[1:4] == ["tasks", "receipt", f"spawn:{restart_argv}"] + assert events.index("tasks") < events.index("receipt") + assert any(event.startswith("spawn:") for event in events) + + +def test_pro_bundle_handoff_stops_supervisor_before_port_check_and_install( + monkeypatch, + tmp_path: Path, +) -> None: + events: list[str] = [] + restart_argv = ["python.exe", "-m", "flocks.cli.main", "start"] + pro_wheel = tmp_path / "flockspro.whl" + manifest = tmp_path / "manifest.json" + args = _handoff_args(tmp_path, restart_argv) + separator_index = args.index("--") + args[separator_index:separator_index] = [ + "--pro-wheel-path", + str(pro_wheel), + "--pro-bundle-manifest-path", + str(manifest), + ] + + monkeypatch.setattr(restart_handoff, "_record_handoff_log", lambda _message: None) + monkeypatch.setattr( + restart_handoff, + "_wait_for_parent_exit", + lambda parent_pid: events.append(f"wait-parent:{parent_pid}") or True, + ) + monkeypatch.setattr( + restart_handoff, + "_stop_supervisor_before_restart", + lambda **kwargs: events.append(f"stop-supervisor:{kwargs}") or True, + ) + monkeypatch.setattr( + restart_handoff, + "_ensure_backend_port_free", + lambda backend_port: events.append(f"free-port:{backend_port}") or True, + ) + monkeypatch.setattr( + restart_handoff, + "_run_upgrade_tasks", + lambda _args: events.append("install") or None, + ) + monkeypatch.setattr(restart_handoff, "_report_pending_pro_bundle_install_receipt", lambda _args: None) + monkeypatch.setattr( + restart_handoff.subprocess, + "Popen", + lambda _argv, **_kwargs: events.append("spawn") or SimpleNamespace(pid=4321), + ) + + assert restart_handoff.run(args) == 0 + assert events == [ + "wait-parent:1234", + "stop-supervisor:{'backend_port': 8000, 'service_ports': (5173,)}", + "free-port:8000", + "install", + "spawn", + ] def test_run_does_not_spawn_when_parent_exit_times_out(monkeypatch, tmp_path: Path) -> None: @@ -258,7 +929,10 @@ def test_run_does_not_spawn_when_parent_exit_times_out(monkeypatch, tmp_path: Pa code = restart_handoff.run(_handoff_args(tmp_path, ["python.exe", "-m", "flocks.cli.main", "start"])) assert code == 1 - assert events == ["log:started parent_pid=1234 backend=127.0.0.1:8000 frontend=127.0.0.1:5173", "log:parent_exit_timeout parent_pid=1234"] + assert events == [ + "log:started parent_pid=1234 backend=127.0.0.1:8000 frontend=127.0.0.1:5173", + "log:parent_exit_timeout parent_pid=1234", + ] def test_run_does_not_spawn_when_upgrade_tasks_fail(monkeypatch, tmp_path: Path) -> None: @@ -269,7 +943,6 @@ def test_run_does_not_spawn_when_upgrade_tasks_fail(monkeypatch, tmp_path: Path) monkeypatch.setattr(restart_handoff, "_wait_for_parent_exit", lambda parent_pid: True) monkeypatch.setattr(restart_handoff, "_ensure_backend_port_free", lambda backend_port: True) monkeypatch.setattr(restart_handoff, "_run_upgrade_tasks", lambda args: "sync failed") - monkeypatch.setattr(restart_handoff, "_rollback_failed_upgrade", lambda args, error: events.append(f"rollback:{error}")) monkeypatch.setattr( restart_handoff.subprocess, "Popen", @@ -279,11 +952,11 @@ def test_run_does_not_spawn_when_upgrade_tasks_fail(monkeypatch, tmp_path: Path) code = restart_handoff.run(_handoff_args(tmp_path, restart_argv)) assert code == 1 - assert "rollback:sync failed" in events + assert "log:upgrade_tasks_failed error=sync failed" in events assert "spawn" not in events -def test_run_rolls_back_and_cleans_up_when_upgrade_tasks_crash(monkeypatch, tmp_path: Path) -> None: +def test_run_logs_and_cleans_up_when_upgrade_tasks_crash(monkeypatch, tmp_path: Path) -> None: events: list[str] = [] cleanup_dir = tmp_path / "cleanup" cleanup_dir.mkdir() @@ -300,7 +973,6 @@ def crash(_args): monkeypatch.setattr(restart_handoff, "_wait_for_parent_exit", lambda parent_pid: True) monkeypatch.setattr(restart_handoff, "_ensure_backend_port_free", lambda backend_port: True) monkeypatch.setattr(restart_handoff, "_run_upgrade_tasks", crash) - monkeypatch.setattr(restart_handoff, "_rollback_failed_upgrade", lambda args, error: events.append(f"rollback:{error}")) monkeypatch.setattr( restart_handoff.subprocess, "Popen", @@ -310,7 +982,7 @@ def crash(_args): code = restart_handoff.run(args) assert code == 1 - assert "rollback:upgrade tasks crashed: boom" in events + assert "log:upgrade_tasks_failed error=upgrade tasks crashed: boom" in events assert not cleanup_dir.exists() assert "spawn" not in events @@ -337,50 +1009,78 @@ def test_run_does_not_spawn_when_supervisor_stop_fails(monkeypatch, tmp_path: Pa assert "spawn" not in events -def test_run_rolls_back_prepared_handover_when_supervisor_stop_fails(monkeypatch, tmp_path: Path) -> None: +def test_restart_only_does_not_rollback_when_supervisor_stop_fails(monkeypatch, tmp_path: Path) -> None: events: list[str] = [] restart_argv = ["python.exe", "-m", "flocks.cli.main", "start"] monkeypatch.setattr(restart_handoff, "_record_handoff_log", lambda message: events.append(f"log:{message}")) monkeypatch.setattr(restart_handoff, "_wait_for_parent_exit", lambda parent_pid: True) - monkeypatch.setattr(restart_handoff, "_prepare_upgrade_handover", lambda args: events.append("prepare") or True) + monkeypatch.setattr(restart_handoff, "_ensure_backend_port_free", lambda _port: True) monkeypatch.setattr(restart_handoff, "_run_upgrade_tasks", lambda args: None) monkeypatch.setattr(restart_handoff, "_stop_supervisor_before_restart", lambda: False) - monkeypatch.setattr(restart_handoff, "_rollback_upgrade_handover", lambda: events.append("rollback-handover")) monkeypatch.setattr( restart_handoff.subprocess, "Popen", lambda *_args, **_kwargs: events.append("spawn"), ) - code = restart_handoff.run(_handoff_args(tmp_path, restart_argv, prepare_handover=True)) + code = restart_handoff.run(_handoff_args(tmp_path, restart_argv)) assert code == 1 - assert "rollback-handover" in events + assert not any("rollback" in event for event in events) assert "spawn" not in events -def test_run_rolls_back_prepared_handover_when_restart_spawn_fails(monkeypatch, tmp_path: Path) -> None: +def test_restart_only_does_not_rollback_when_restart_spawn_fails(monkeypatch, tmp_path: Path) -> None: events: list[str] = [] restart_argv = ["python.exe", "-m", "flocks.cli.main", "start"] monkeypatch.setattr(restart_handoff, "_record_handoff_log", lambda message: events.append(f"log:{message}")) monkeypatch.setattr(restart_handoff, "_wait_for_parent_exit", lambda parent_pid: True) - monkeypatch.setattr(restart_handoff, "_prepare_upgrade_handover", lambda args: events.append("prepare") or True) + monkeypatch.setattr(restart_handoff, "_ensure_backend_port_free", lambda _port: True) monkeypatch.setattr(restart_handoff, "_run_upgrade_tasks", lambda args: None) monkeypatch.setattr(restart_handoff, "_stop_supervisor_before_restart", lambda: True) - monkeypatch.setattr(restart_handoff, "_rollback_upgrade_handover", lambda: events.append("rollback-handover")) monkeypatch.setattr( restart_handoff.subprocess, "Popen", lambda *_args, **_kwargs: (_ for _ in ()).throw(OSError("spawn failed")), ) - code = restart_handoff.run(_handoff_args(tmp_path, restart_argv, prepare_handover=True)) + code = restart_handoff.run(_handoff_args(tmp_path, restart_argv)) assert code == 1 assert "log:restart_spawn_failed error=spawn failed" in events - assert "rollback-handover" in events + assert not any("rollback" in event for event in events) + + +def test_stop_supervisor_requests_stop_when_health_probe_is_temporarily_unavailable(monkeypatch) -> None: + from flocks.cli import service_control + + events: list[str] = [] + daemon_running = True + + monkeypatch.setattr(service_control, "supervisor_is_running", lambda _paths: False) + monkeypatch.setattr( + restart_handoff.service_manager, + "pid_is_running", + lambda _pid: daemon_running, + ) + monkeypatch.setattr(restart_handoff, "_backend_port_in_use", lambda _port: False) + + def request_stop(*, paths, timeout): + del paths, timeout + nonlocal daemon_running + events.append("request-stop") + daemon_running = False + + monkeypatch.setattr(service_control, "request_stop", request_stop) + + assert restart_handoff._stop_supervisor_before_restart( + daemon_pid=2468, + timeout_seconds=0.01, + poll_interval_seconds=0.001, + ) + assert events == ["request-stop"] @pytest.mark.skipif(sys.platform == "win32", reason="uses the Unix domain socket control API") diff --git a/tests/updater/test_restart_handoff_process.py b/tests/updater/test_restart_handoff_process.py new file mode 100644 index 000000000..69c302a43 --- /dev/null +++ b/tests/updater/test_restart_handoff_process.py @@ -0,0 +1,281 @@ +import json +import os +import shutil +import socket +import subprocess +import sys +import time +from pathlib import Path + +import pytest + +from flocks.cli import service_manager +from flocks.cli.service_config import service_config_payload + + +pytestmark = pytest.mark.skipif(sys.platform == "win32", reason="uses Unix executable shims and symlinks") + +REPO_ROOT = Path(__file__).resolve().parents[2] +LOCALE_HANDOFF_ARGS = [ + pytest.param([], id="english-upgrade"), + pytest.param( + [ + "--uv-default-index", + "https://mirrors.aliyun.com/pypi/simple", + "--npm-registry", + "https://registry.npmmirror.com/", + ], + id="chinese-upgrade", + ), +] + + +def _free_port() -> int: + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def _prepare_isolated_runtime(tmp_path: Path) -> tuple[Path, Path, dict[str, str]]: + install_root = tmp_path / "install" + install_root.mkdir() + (install_root / "pyproject.toml").write_text( + '[project]\nname = "handoff-process-test"\nversion = "0.0.0"\nrequires-python = ">=3.12"\ndependencies = []\n', + encoding="utf-8", + ) + runtime_python = install_root / ".venv" / "bin" / "python" + runtime_python.parent.mkdir(parents=True) + runtime_python.symlink_to(sys.executable) + + uv_path = shutil.which("uv") + if uv_path is None: + pytest.skip("uv is required for the real handoff process test") + + isolated_home = tmp_path / "home" + isolated_home.mkdir() + env = os.environ.copy() + env.update( + { + "FLOCKS_ROOT": str(tmp_path / "flocks-root"), + "HOME": str(isolated_home), + "UV_CACHE_DIR": str(tmp_path / "uv-cache"), + "UV_PYTHON_DOWNLOADS": "never", + } + ) + env.pop("PYTHONPATH", None) + return install_root, Path(uv_path), env + + +def _parent_process() -> subprocess.Popen[str]: + return subprocess.Popen( + [sys.executable, "-c", "import time; time.sleep(0.2)"], + text=True, + ) + + +def _base_handoff_command( + *, + parent_pid: int, + install_root: Path, + uv_path: Path, + backend_port: int, + frontend_port: int, + current_version: str, +) -> list[str]: + return [ + sys.executable, + "-m", + "flocks.updater.restart_handoff", + "--parent-pid", + str(parent_pid), + "--backend-host", + "127.0.0.1", + "--backend-port", + str(backend_port), + "--frontend-host", + "127.0.0.1", + "--frontend-port", + str(frontend_port), + "--install-root", + str(install_root), + "--uv-path", + str(uv_path), + "--sync-timeout", + "10", + "--version", + "2026.7.15", + "--current-version", + current_version, + ] + + +def _restart_marker_command(install_root: Path, marker: Path, value: str) -> list[str]: + runtime_python = install_root / ".venv" / "bin" / "python" + script = f"from pathlib import Path; Path({str(marker)!r}).write_text({value!r}, encoding='utf-8')" + return [str(runtime_python), "-c", script] + + +def _run_handoff(command: list[str], *, env: dict[str, str]) -> subprocess.CompletedProcess[str]: + return subprocess.run( + command, + cwd=REPO_ROOT, + env=env, + capture_output=True, + text=True, + timeout=20, + check=False, + ) + + +def _wait_for_marker(marker: Path, expected: str) -> None: + deadline = time.monotonic() + 5 + while time.monotonic() < deadline: + if marker.is_file() and marker.read_text(encoding="utf-8") == expected: + return + time.sleep(0.05) + raise AssertionError(f"restart marker was not written: {marker}") + + +@pytest.mark.parametrize("locale_args", LOCALE_HANDOFF_ARGS) +def test_v2026_7_1_real_handoff_command_restarts_with_current_code( + tmp_path: Path, + locale_args: list[str], +) -> None: + install_root, fake_uv, env = _prepare_isolated_runtime(tmp_path) + marker = tmp_path / "v2026.7.1-restarted" + backup_path = tmp_path / "backup.tar.gz" + backup_path.write_text("backup", encoding="utf-8") + parent = _parent_process() + command = _base_handoff_command( + parent_pid=parent.pid, + install_root=install_root, + uv_path=fake_uv, + backend_port=_free_port(), + frontend_port=_free_port(), + current_version="2026.7.1", + ) + install_index = command.index("--install-root") + command[install_index:install_index] = [ + "--backend-pid-file", + str(tmp_path / "backend.pid"), + ] + command.extend( + [ + *locale_args, + "--backup-path", + str(backup_path), + "--", + *_restart_marker_command(install_root, marker, "v2026.7.1"), + ] + ) + + completed = _run_handoff(command, env=env) + parent.wait(timeout=5) + + assert completed.returncode == 0, completed.stderr + _wait_for_marker(marker, "v2026.7.1") + + +@pytest.mark.parametrize("locale_args", LOCALE_HANDOFF_ARGS) +def test_v2026_7_15_real_handoff_command_restarts_with_current_code( + tmp_path: Path, + locale_args: list[str], +) -> None: + install_root, fake_uv, env = _prepare_isolated_runtime(tmp_path) + marker = tmp_path / "v2026.7.15-restarted" + backup_path = tmp_path / "backup.tar.gz" + backup_path.write_text("backup", encoding="utf-8") + public_port = _free_port() + parent = _parent_process() + command = _base_handoff_command( + parent_pid=parent.pid, + install_root=install_root, + uv_path=fake_uv, + backend_port=public_port, + frontend_port=public_port, + current_version="2026.7.15", + ) + command.extend( + [ + *locale_args, + "--backup-path", + str(backup_path), + "--prepare-handover", + "--", + *_restart_marker_command(install_root, marker, "v2026.7.15"), + ] + ) + + completed = _run_handoff(command, env=env) + parent.wait(timeout=5) + + assert completed.returncode == 0, completed.stderr + _wait_for_marker(marker, "v2026.7.15") + + +@pytest.mark.parametrize("locale_args", LOCALE_HANDOFF_ARGS) +def test_current_real_upgrade_handoff_command_restarts_with_replaced_source( + tmp_path: Path, + locale_args: list[str], +) -> None: + install_root, fake_uv, env = _prepare_isolated_runtime(tmp_path) + marker = tmp_path / "current-restarted" + env["FLOCKS_RESTART_MARKER"] = str(marker) + + content_root = tmp_path / "staged" + fake_cli = content_root / "flocks" / "cli" + fake_cli.mkdir(parents=True) + (content_root / "flocks" / "__init__.py").write_text("", encoding="utf-8") + shutil.copy2(install_root / "pyproject.toml", content_root / "pyproject.toml") + (fake_cli / "__init__.py").write_text("", encoding="utf-8") + (fake_cli / "main.py").write_text( + "import os\n" + "from pathlib import Path\n" + "Path(os.environ['FLOCKS_RESTART_MARKER']).write_text('current', encoding='utf-8')\n", + encoding="utf-8", + ) + backup_path = tmp_path / "backup.tar.gz" + backup_path.write_text("backup", encoding="utf-8") + + backend_port = _free_port() + frontend_port = _free_port() + config = service_manager.ServiceConfig( + backend_host="127.0.0.1", + backend_port=backend_port, + frontend_host="127.0.0.1", + frontend_port=frontend_port, + no_browser=True, + skip_frontend_build=True, + ) + parent = _parent_process() + command = _base_handoff_command( + parent_pid=parent.pid, + install_root=install_root, + uv_path=fake_uv, + backend_port=backend_port, + frontend_port=frontend_port, + current_version="2026.7.15", + ) + command.extend( + [ + *locale_args, + "--mode", + "upgrade", + "--content-root", + str(content_root), + "--backup-path", + str(backup_path), + "--was-running", + "--service-config-json", + json.dumps(service_config_payload(config)), + "--", + str(install_root / ".venv" / "bin" / "python"), + ] + ) + + completed = _run_handoff(command, env=env) + parent.wait(timeout=5) + + assert completed.returncode == 0, completed.stderr + assert marker.read_text(encoding="utf-8") == "current" + assert (install_root / "flocks" / "cli" / "main.py").is_file() diff --git a/tests/updater/test_updater.py b/tests/updater/test_updater.py index cbeb4b043..58dec2b10 100644 --- a/tests/updater/test_updater.py +++ b/tests/updater/test_updater.py @@ -11,6 +11,7 @@ import pytest from flocks.cli import service_control, service_manager +from flocks.cli.service_config import service_config_payload from flocks.updater import updater from tests.helpers.service_supervisor import ( make_short_runtime_root, @@ -73,6 +74,87 @@ def test_current_service_config_requires_supervisor_control_api(monkeypatch: pyt updater._current_service_config() +def test_capture_service_snapshot_preserves_complete_supervisor_config( + monkeypatch: pytest.MonkeyPatch, +) -> None: + config = service_manager.ServiceConfig( + backend_host="2001:db8::20", + backend_port=9527, + frontend_host="2001:db8::20", + frontend_port=9527, + legacy_backend_host="0.0.0.0", + legacy_backend_port=9000, + server_port_migration_hint=True, + no_browser=False, + skip_frontend_build=False, + ) + payload = { + "daemon": {"pid": 2468, "state": "running"}, + "backend": { + "pid": 3579, + "host": config.backend_host, + "port": config.backend_port, + "state": "healthy", + "health": "healthy", + }, + "webui": { + "host": config.frontend_host, + "port": config.frontend_port, + "state": "static", + }, + "config": service_config_payload(config), + } + status = service_control.parse_supervisor_status(payload) + monkeypatch.setattr(service_control, "read_supervisor_status", lambda **_kwargs: status) + + snapshot = updater._capture_service_snapshot() + + assert snapshot.config == config + assert snapshot.daemon_pid == 2468 + assert snapshot.was_running is True + + +def test_spawn_restart_handoff_redirects_output_to_backend_log( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + capfd: pytest.CaptureFixture[str], +) -> None: + flocks_root = tmp_path / "flocks-root" + command = [ + sys.executable, + "-c", + "import sys; print('handoff stdout'); print('handoff stderr', file=sys.stderr)", + ] + monkeypatch.setenv("FLOCKS_ROOT", str(flocks_root)) + + process = updater._spawn_restart_handoff(command, cwd=tmp_path) + + assert process.wait(timeout=10) == 0 + captured = capfd.readouterr() + assert captured.out == "" + assert captured.err == "" + backend_output = (flocks_root / "logs" / "backend.log").read_text(encoding="utf-8") + assert "handoff stdout" in backend_output + assert "handoff stderr" in backend_output + + +def test_capture_service_snapshot_allows_stopped_service_without_control_api( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + service_control, + "read_supervisor_status", + lambda **_kwargs: (_ for _ in ()).throw(RuntimeError("control down")), + ) + monkeypatch.setattr(service_manager, "read_runtime_record", lambda _path: None) + monkeypatch.setattr(service_manager, "trusted_daemon_process_pids", lambda **_kwargs: []) + + snapshot = updater._capture_service_snapshot() + + assert snapshot.daemon_pid is None + assert snapshot.was_running is False + + def test_run_handles_none_process_output(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: def fake_run(*args, **kwargs): return subprocess.CompletedProcess(args=args[0], returncode=0, stdout=None, stderr=None) @@ -107,41 +189,6 @@ def fake_run(*args, **kwargs): assert stderr == "failed�output" -@pytest.mark.asyncio -async def test_await_ignoring_cancellation_backs_off_on_repeated_cancellation( - monkeypatch: pytest.MonkeyPatch, -) -> None: - shield_calls = 0 - sleep_delays: list[float] = [] - - async def fake_shield(task): - nonlocal shield_calls - shield_calls += 1 - if shield_calls <= 2: - raise updater.asyncio.CancelledError - return await task - - async def fake_sleep(delay: float) -> None: - sleep_delays.append(delay) - if len(sleep_delays) == 1: - raise updater.asyncio.CancelledError - - async def critical_step() -> str: - return "done" - - monkeypatch.setattr(updater.asyncio, "shield", fake_shield) - monkeypatch.setattr(updater.asyncio, "sleep", fake_sleep) - - result = await updater._await_ignoring_cancellation(critical_step()) - - assert result == "done" - assert shield_calls == 3 - assert sleep_delays == [ - updater._CANCELLATION_RETRY_DELAY_SECONDS, - updater._CANCELLATION_RETRY_DELAY_SECONDS, - ] - - def test_get_current_version_prefers_higher_marker_version( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, @@ -516,7 +563,6 @@ def test_build_dependency_sync_command_installs_project_on_windows( assert updater._build_dependency_sync_command("uv", uv_default_index="https://mirror.example/simple") == [ "uv", "sync", - "--frozen", "--no-python-downloads", "--default-index", "https://mirror.example/simple", @@ -528,7 +574,7 @@ def test_build_dependency_sync_command_keeps_project_install_on_non_windows( ) -> None: monkeypatch.setattr(updater.sys, "platform", "linux") - assert updater._build_dependency_sync_command("uv") == ["uv", "sync", "--frozen", "--no-python-downloads"] + assert updater._build_dependency_sync_command("uv") == ["uv", "sync", "--no-python-downloads"] def test_wheel_build_config_does_not_force_include_runtime_or_build_outputs() -> None: @@ -563,21 +609,6 @@ def test_build_frontend_subprocess_env_prepends_bundled_node_on_windows( assert result["PATH"].split(os.pathsep)[0] == str(node_home) -def test_upgrade_page_html_contains_marker_and_version() -> None: - html = updater._upgrade_page_html("2026.3.31.1") - - assert "flocks-upgrade-in-progress" in html - assert "v2026.3.31.1" in html - assert "window.location.reload()" in html - - -def test_upgrade_page_probe_urls_support_ipv6_loopback_fallback() -> None: - assert updater._upgrade_page_probe_urls("::", 5173) == [ - "http://[::1]:5173", - "http://127.0.0.1:5173", - ] - - def test_resolve_update_mirror_profile_uses_cn_defaults_for_zh_locale() -> None: profile = updater._resolve_update_mirror_profile( ["github", "gitee", "gitlab"], @@ -823,10 +854,9 @@ def test_build_restart_handoff_argv_rewrites_serve_to_managed_start( sync_timeout=300, version="2026.4.1", current_version="2026.3.31", - prepare_handover=True, ) - assert "--prepare-handover" in argv[: argv.index("--")] + assert "--mode" not in argv assert argv[argv.index("--") + 1 :] == [ "python", "-m", @@ -845,6 +875,29 @@ def test_build_restart_handoff_argv_rewrites_serve_to_managed_start( ] +def test_build_restart_handoff_argv_can_skip_parent_wait( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + monkeypatch.setattr( + updater, + "_handoff_service_config", + lambda: service_manager.ServiceConfig(), + ) + + argv = updater._build_restart_handoff_argv( + ["python"], + tmp_path, + uv_path="uv", + sync_timeout=300, + version="2026.4.1", + current_version="2026.3.31", + wait_for_parent=False, + ) + + assert "--parent-pid" not in argv + + def test_refresh_global_cli_entry_creates_symlink_on_unix( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, @@ -992,6 +1045,70 @@ async def fail_run_async(*_args, **_kwargs): assert await updater._validate_restart_runtime(tmp_path) is None +@pytest.mark.asyncio +async def test_shared_installer_runs_core_steps_in_order( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + webui_dir = tmp_path / "webui" + webui_dir.mkdir() + (webui_dir / "package.json").write_text("{}", encoding="utf-8") + events: list[str] = [] + + async def fake_sync(**_kwargs) -> None: + events.append("sync") + return None + + async def fake_build(*_args, **_kwargs) -> None: + events.append("build") + return None + + async def fake_validate(_root: Path) -> None: + events.append("validate") + return None + + monkeypatch.setattr(service_manager, "resolve_npm_executable", lambda: "/usr/bin/npm") + monkeypatch.setattr(service_manager, "node_version_satisfies_requirement", lambda: True) + monkeypatch.setattr(updater, "_sync_project_dependencies", fake_sync) + monkeypatch.setattr(updater, "_build_frontend_workspace", fake_build) + monkeypatch.setattr(updater, "_validate_restart_runtime", fake_validate) + monkeypatch.setattr(updater, "_refresh_global_cli_entry", lambda _root: events.append("refresh-cli")) + monkeypatch.setattr(updater, "_write_version_marker", lambda version: events.append(f"marker:{version}")) + + await updater.install_or_repair_source( + tmp_path, + version="v2026.7.2", + uv_path="/usr/bin/uv", + ) + + assert events == ["sync", "build", "validate", "refresh-cli", "marker:2026.7.2"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("platform", "script"), + [("darwin", "scripts/install.sh"), ("win32", "scripts/install.ps1")], +) +async def test_shared_installer_reports_missing_npm_with_platform_installer_hint( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + platform: str, + script: str, +) -> None: + webui_dir = tmp_path / "webui" + webui_dir.mkdir() + (webui_dir / "package.json").write_text("{}", encoding="utf-8") + monkeypatch.setattr(updater.sys, "platform", platform) + monkeypatch.setattr(service_manager, "resolve_npm_executable", lambda: None) + + with pytest.raises(RuntimeError, match=script): + await updater.install_or_repair_source( + tmp_path, + version="2026.7.2", + uv_path="uv", + ) + + def test_rmtree_onerror_retries_before_logging_skip(monkeypatch: pytest.MonkeyPatch) -> None: attempts: list[str] = [] warnings: list[tuple[str, dict[str, str]]] = [] @@ -1063,495 +1180,6 @@ def test_safe_remove_renames_locked_directory_on_windows( assert (leftovers[0] / "dist" / "index.html").exists() -@pytest.mark.skipif(sys.platform == "win32", reason="uses the Unix domain socket control API") -def test_prepare_upgrade_handover_writes_state_and_stops_frontend_with_real_control_api( - monkeypatch: pytest.MonkeyPatch, -) -> None: - short_root = make_short_runtime_root("flocks-updater-") - monkeypatch.setenv("FLOCKS_ROOT", str(short_root)) - paths = service_manager.runtime_paths() - config = service_manager.ServiceConfig( - backend_host="127.0.0.1", - backend_port=9995, - frontend_host="127.0.0.1", - frontend_port=9996, - ) - daemon, thread = start_supervisor(config) - wait_for_supervisor(paths, running=True) - - monkeypatch.setattr( - updater, - "_start_upgrade_page_server", - lambda _config, _version: { - "upgrade_server_pid": 321, - "page_dir": str(short_root / "page"), - "page_log": str(short_root / "logs" / "upgrade.log"), - }, - ) - - try: - payload = updater._prepare_upgrade_handover("2026.3.31.1") - - status = service_control.read_supervisor_status(paths) - assert status.backend.paused is True - assert status.webui.paused is True - assert payload["upgrade_server_pid"] == 321 - assert payload["backend_port"] == 9995 - assert payload["frontend_port"] == 9996 - assert updater._read_upgrade_state()["version"] == "2026.3.31.1" - finally: - stop_supervisor(daemon, thread) - shutil.rmtree(short_root, ignore_errors=True) - - -@pytest.mark.skipif(sys.platform == "win32", reason="uses the Unix domain socket control API") -def test_prepare_upgrade_handover_restores_frontend_when_upgrade_page_fails_with_real_control_api( - monkeypatch: pytest.MonkeyPatch, -) -> None: - short_root = make_short_runtime_root("flocks-updater-") - monkeypatch.setenv("FLOCKS_ROOT", str(short_root)) - paths = service_manager.runtime_paths() - config = service_manager.ServiceConfig( - backend_host="127.0.0.1", - backend_port=9995, - frontend_host="127.0.0.1", - frontend_port=9996, - ) - daemon, thread = start_supervisor(config) - wait_for_supervisor(paths, running=True) - calls: list[str] = [] - - monkeypatch.setattr(updater, "_stop_upgrade_page_server", lambda **_kw: calls.append("stop_page")) - monkeypatch.setattr( - updater, - "_start_upgrade_page_server", - lambda _config, _version: (_ for _ in ()).throw(RuntimeError("page failed")), - ) - - try: - with pytest.raises(RuntimeError, match="page failed"): - updater._prepare_upgrade_handover("2026.3.31.1") - - status = service_control.read_supervisor_status(paths) - assert calls == ["stop_page"] - assert status.backend.paused is False - assert status.backend.pid is not None - assert status.webui.paused is False - assert status.webui.pid is None - assert status.webui.state == "static" - assert updater._read_upgrade_state() is None - finally: - stop_supervisor(daemon, thread) - shutil.rmtree(short_root, ignore_errors=True) - - -@pytest.mark.skipif(sys.platform == "win32", reason="uses the Unix domain socket control API") -def test_rollback_failed_update_resumes_backend_when_handoff_tasks_fail( - monkeypatch: pytest.MonkeyPatch, -) -> None: - short_root = make_short_runtime_root("flocks-updater-") - monkeypatch.setenv("FLOCKS_ROOT", str(short_root)) - paths = service_manager.runtime_paths() - config = service_manager.ServiceConfig( - backend_host="127.0.0.1", - backend_port=9995, - frontend_host="127.0.0.1", - frontend_port=9996, - ) - daemon, thread = start_supervisor(config) - wait_for_supervisor(paths, running=True) - monkeypatch.setattr(updater, "_stop_upgrade_page_server", lambda **_kw: None) - - try: - updater._write_upgrade_state( - { - "version": "2026.4.1", - "backend_host": "127.0.0.1", - "backend_port": 9995, - "frontend_host": "127.0.0.1", - "frontend_port": 9996, - "skip_frontend_build": True, - } - ) - old_backend = daemon.backend.process - assert old_backend is not None - service_control.request_prepare_upgrade(paths=paths) - wait_for_process_exit(old_backend) - - updater._rollback_failed_update(None, short_root / "install", "2026.3.31") - - status = service_control.read_supervisor_status(paths) - assert status.backend.paused is False - assert status.webui.paused is False - assert status.backend.pid is not None - assert status.backend.pid != old_backend.pid - assert status.webui.pid is None - assert status.webui.state == "static" - assert updater._read_upgrade_state() is None - finally: - stop_supervisor(daemon, thread) - shutil.rmtree(short_root, ignore_errors=True) - - -def test_recover_upgrade_state_restarts_frontend_and_clears_marker( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - monkeypatch.setenv("FLOCKS_ROOT", str(tmp_path / ".flocks")) - started: list[tuple[int, bool | None]] = [] - stopped: list[str] = [] - - monkeypatch.setattr(updater, "_stop_upgrade_page_server", lambda **kw: stopped.append("stop")) - monkeypatch.setattr( - service_control, - "request_resume_upgrade", - lambda config, **_kwargs: started.append((config.frontend_port, config.skip_frontend_build)) - or _webui_control_status(), - ) - updater._write_upgrade_state( - { - "version": "2026.3.31.1", - "backend_host": "127.0.0.1", - "backend_port": 8000, - "frontend_host": "127.0.0.1", - "frontend_port": 5173, - "skip_frontend_build": True, - } - ) - - updater.recover_upgrade_state() - - assert stopped == ["stop"] - assert started == [(5173, True)] - assert updater._read_upgrade_state() is None - - -def test_recover_upgrade_state_retries_frontend_with_build_when_dist_is_missing( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - monkeypatch.setenv("FLOCKS_ROOT", str(tmp_path / ".flocks")) - starts: list[tuple[str, bool | None, bool | None]] = [] - - monkeypatch.setattr(updater, "_stop_upgrade_page_server", lambda **kw: None) - - results = iter([ - _webui_control_payload("degraded", "missing dist"), - _webui_control_payload(), - ]) - - def fake_resume_upgrade(config, **_kwargs): - starts.append(("resume", config.skip_frontend_build, None)) - return service_control.parse_supervisor_status(next(results)) - - def fake_restart_webui(config, *, force_frontend_build=False, **_kwargs): - starts.append(("restart_webui", config.skip_frontend_build, force_frontend_build or None)) - return service_control.parse_supervisor_status(next(results)) - - monkeypatch.setattr(service_control, "request_resume_upgrade", fake_resume_upgrade) - monkeypatch.setattr(service_control, "request_restart_webui", fake_restart_webui) - updater._write_upgrade_state( - { - "version": "2026.3.31.1", - "backend_host": "127.0.0.1", - "backend_port": 8000, - "frontend_host": "127.0.0.1", - "frontend_port": 5173, - "skip_frontend_build": True, - } - ) - - updater.recover_upgrade_state() - - assert starts == [("resume", True, None), ("restart_webui", False, True)] - assert updater._read_upgrade_state() is None - - -def test_recover_upgrade_state_restart_failure_clears_state_without_restarting_page( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - monkeypatch.setenv("FLOCKS_ROOT", str(tmp_path / ".flocks")) - starts: list[tuple[str, bool | None, bool | None]] = [] - - monkeypatch.setattr(updater, "_stop_upgrade_page_server", lambda **kw: None) - - def fake_resume_upgrade(config, **_kwargs): - starts.append(("resume", config.skip_frontend_build, None)) - return _webui_control_status("degraded", "still broken") - - def fake_restart_webui(config, *, force_frontend_build=False, **_kwargs): - starts.append(("restart_webui", config.skip_frontend_build, force_frontend_build or None)) - return _webui_control_status("degraded", "still broken") - - monkeypatch.setattr(service_control, "request_resume_upgrade", fake_resume_upgrade) - monkeypatch.setattr(service_control, "request_restart_webui", fake_restart_webui) - updater._write_upgrade_state( - { - "version": "2026.3.31.1", - "backend_host": "127.0.0.1", - "backend_port": 8000, - "frontend_host": "127.0.0.1", - "frontend_port": 5173, - "skip_frontend_build": True, - } - ) - - with pytest.raises(RuntimeError, match="still broken"): - updater.recover_upgrade_state() - - assert starts == [("resume", True, None), ("restart_webui", False, True)] - assert updater._read_upgrade_state() is None - - -def test_start_upgrade_page_server_binds_configured_frontend_host( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - flocks_root = tmp_path / ".flocks" - (flocks_root / "run").mkdir(parents=True) - monkeypatch.setenv("FLOCKS_ROOT", str(flocks_root)) - - page_dir = tmp_path / "page" - page_dir.mkdir() - captured: dict[str, object] = {} - - monkeypatch.setattr(updater, "_write_upgrade_page", lambda _version: page_dir) - monkeypatch.setattr(updater, "_wait_for_upgrade_page", lambda config: captured.setdefault("wait_host", config.frontend_host)) - monkeypatch.setattr( - service_manager, - "resolve_python_subprocess_command", - lambda _root=None: ["/env/bin/python"], - ) - monkeypatch.setattr( - updater, - "_spawn_detached_process", - lambda command, *, cwd, log_path: captured.update({ - "command": command, - "cwd": cwd, - "log_path": log_path, - }) or SimpleNamespace(pid=4321), - ) - - config = service_manager.ServiceConfig(frontend_host="0.0.0.0", frontend_port=5173) - payload = updater._start_upgrade_page_server(config, "2026.4.1") - - assert payload["upgrade_server_pid"] == 4321 - assert captured["command"] == [ - "/env/bin/python", - "-m", - "http.server", - "5173", - "--bind", - "0.0.0.0", - "--directory", - str(page_dir.resolve()), - ] - assert captured["wait_host"] == "0.0.0.0" - - -def test_stop_upgrade_page_server_does_not_kill_unified_flocks_service( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - flocks_root = tmp_path / ".flocks" - monkeypatch.setenv("FLOCKS_ROOT", str(flocks_root)) - killed: list[int] = [] - - monkeypatch.setattr(service_manager, "port_owner_pids", lambda _port: [111]) - monkeypatch.setattr( - service_manager, - "_process_command_line", - lambda _pid: "/env/bin/python -m flocks.cli.main serve --host 127.0.0.1 --port 5173", - ) - monkeypatch.setattr(updater.os, "kill", lambda pid, _sig: killed.append(pid)) - - updater._stop_upgrade_page_server(frontend_port=5173) - - assert killed == [] - - -def test_stop_upgrade_page_server_kills_only_upgrade_page_process( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - flocks_root = tmp_path / ".flocks" - page_dir = flocks_root / "run" / "upgrade-page" - monkeypatch.setenv("FLOCKS_ROOT", str(flocks_root)) - killed: list[int] = [] - - def fake_command_line(pid: int) -> str: - if pid == 222: - return f"/env/bin/python -m http.server 5173 --directory {page_dir}" - return "/env/bin/python -m flocks.cli.main serve --host 127.0.0.1 --port 5173" - - monkeypatch.setattr(service_manager, "port_owner_pids", lambda _port: [111, 222]) - monkeypatch.setattr(service_manager, "_process_command_line", fake_command_line) - monkeypatch.setattr(updater.os, "kill", lambda pid, _sig: killed.append(pid)) - monkeypatch.setattr(updater.time, "sleep", lambda _seconds: None) - - updater._stop_upgrade_page_server(frontend_port=5173) - - assert killed == [222] - - -def test_wait_for_upgrade_page_uses_access_host_for_local_probe( - monkeypatch: pytest.MonkeyPatch, -) -> None: - requested_urls: list[str] = [] - - class _FakeClient: - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - def get(self, url): - requested_urls.append(url) - return SimpleNamespace(status_code=200) - - monkeypatch.setattr(updater.httpx, "Client", lambda timeout: _FakeClient()) - monkeypatch.setattr(service_manager, "access_host", lambda host: "127.0.0.1" if host == "0.0.0.0" else host) - monkeypatch.setattr(updater.time, "sleep", lambda _seconds: None) - - updater._wait_for_upgrade_page(service_manager.ServiceConfig(frontend_host="0.0.0.0", frontend_port=5173)) - - assert requested_urls == ["http://127.0.0.1:5173"] - - -def test_wait_for_upgrade_page_falls_back_from_ipv6_to_ipv4_probe( - monkeypatch: pytest.MonkeyPatch, -) -> None: - requested_urls: list[str] = [] - - class _FakeClient: - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - def get(self, url): - requested_urls.append(url) - if url == "http://[::1]:5173": - raise OSError("ipv6 unavailable") - return SimpleNamespace(status_code=200) - - monkeypatch.setattr(updater.httpx, "Client", lambda timeout: _FakeClient()) - monkeypatch.setattr(updater.time, "sleep", lambda _seconds: None) - - updater._wait_for_upgrade_page(service_manager.ServiceConfig(frontend_host="::", frontend_port=5173)) - - assert requested_urls == [ - "http://[::1]:5173", - "http://127.0.0.1:5173", - ] - - -def test_rollback_failed_update_restores_backup_and_rebuilds_frontend_if_needed( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - monkeypatch.setenv("FLOCKS_ROOT", str(tmp_path / ".flocks")) - events: list[str] = [] - - monkeypatch.setattr(updater, "_restore_backup_archive", lambda backup, root: events.append(f"restore:{backup.name}:{root.name}")) - monkeypatch.setattr(updater, "_write_version_marker", lambda version: events.append(f"marker:{version}")) - monkeypatch.setattr(updater, "_stop_upgrade_page_server", lambda **kw: events.append("stop_page")) - monkeypatch.setattr(updater.shutil, "rmtree", lambda path, ignore_errors=True: events.append(f"rmtree:{Path(path).name}")) - - results = iter([ - _webui_control_payload("degraded", "missing dist"), - _webui_control_payload(), - ]) - - def fake_resume_upgrade(config, **_kwargs) -> service_control.SupervisorStatus: - events.append(f"resume:{config.skip_frontend_build}") - return service_control.parse_supervisor_status(next(results)) - - def fake_restart_webui(config, *, force_frontend_build=False, **_kwargs) -> service_control.SupervisorStatus: - events.append(f"restart_webui:{config.skip_frontend_build}:{force_frontend_build or None}") - return service_control.parse_supervisor_status(next(results)) - - monkeypatch.setattr(service_control, "request_resume_upgrade", fake_resume_upgrade) - monkeypatch.setattr(service_control, "request_restart_webui", fake_restart_webui) - updater._write_upgrade_state( - { - "version": "2026.4.1", - "backend_host": "127.0.0.1", - "backend_port": 8000, - "frontend_host": "127.0.0.1", - "frontend_port": 5173, - "skip_frontend_build": True, - } - ) - - backup_path = tmp_path / "backup.tar.gz" - backup_path.write_text("backup", encoding="utf-8") - updater._rollback_failed_update(backup_path, tmp_path / "install", "2026.3.31") - - assert events == [ - "restore:backup.tar.gz:install", - "marker:2026.3.31", - "stop_page", - "resume:True", - "restart_webui:False:True", - "rmtree:upgrade-page", - ] - assert updater._read_upgrade_state() is None - - -def test_rollback_failed_update_clears_state_when_restore_and_frontend_both_fail( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - monkeypatch.setenv("FLOCKS_ROOT", str(tmp_path / ".flocks")) - events: list[str] = [] - - monkeypatch.setattr( - updater, - "_restore_backup_archive", - lambda *_args: (_ for _ in ()).throw(RuntimeError("backup broken")), - ) - monkeypatch.setattr(updater, "_write_version_marker", lambda version: events.append(f"marker:{version}")) - monkeypatch.setattr(updater, "_stop_upgrade_page_server", lambda **kw: events.append("stop_page")) - monkeypatch.setattr(updater.shutil, "rmtree", lambda path, ignore_errors=True: events.append(f"rmtree:{Path(path).name}")) - - def fake_resume_upgrade(config, **_kwargs) -> service_control.SupervisorStatus: - events.append(f"resume:{config.skip_frontend_build}") - return _webui_control_status("degraded", "frontend still broken") - - def fake_restart_webui(config, **_kwargs) -> service_control.SupervisorStatus: - events.append(f"restart_webui:{config.skip_frontend_build}") - return _webui_control_status("degraded", "frontend still broken") - - monkeypatch.setattr(service_control, "request_resume_upgrade", fake_resume_upgrade) - monkeypatch.setattr(service_control, "request_restart_webui", fake_restart_webui) - updater._write_upgrade_state( - { - "version": "2026.4.1", - "backend_host": "127.0.0.1", - "backend_port": 8000, - "frontend_host": "127.0.0.1", - "frontend_port": 5173, - "skip_frontend_build": True, - "phase": "cutover_applied", - } - ) - - backup_path = tmp_path / "backup.tar.gz" - backup_path.write_text("backup", encoding="utf-8") - updater._rollback_failed_update(backup_path, tmp_path / "install", "2026.3.31") - - assert events == [ - "stop_page", - "resume:True", - "rmtree:upgrade-page", - ] - assert updater._read_upgrade_state() is None - - def test_backup_current_version_excludes_all_dist_directories( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, @@ -1566,6 +1194,12 @@ def test_backup_current_version_excludes_all_dist_directories( (webui_dist / "index.html").write_text("ok", encoding="utf-8") (git_dir / "HEAD").write_text("ref: refs/heads/main", encoding="utf-8") (install_root / "flocks.json").write_text('{"keep": true}', encoding="utf-8") + runtime_data = install_root / "data" + runtime_data.mkdir() + (runtime_data / "flocks.db").write_text("runtime", encoding="utf-8") + nested_source_config = install_root / "webui" / "src" / "locales" / "en-US" / "config.json" + nested_source_config.parent.mkdir(parents=True) + nested_source_config.write_text('{"source": true}', encoding="utf-8") (other_dist / "ignored.txt").write_text("nope", encoding="utf-8") backup_dir = tmp_path / "backups" @@ -1577,7 +1211,9 @@ def test_backup_current_version_excludes_all_dist_directories( names = tar.getnames() assert "flocks/webui/dist/index.html" not in names - assert "flocks/flocks.json" in names + assert "flocks/flocks.json" not in names + assert "flocks/data/flocks.db" not in names + assert "flocks/webui/src/locales/en-US/config.json" in names assert "flocks/.git/HEAD" not in names assert "flocks/dist/ignored.txt" not in names @@ -1689,6 +1325,8 @@ def test_replace_install_dir_copies_dot_flocks_plugins_from_source( (source_dir / "flocks.json").write_text('{"version": "new"}', encoding="utf-8") (install_root / "flocks.json").write_text('{"keep": true}', encoding="utf-8") + (install_root / "run").mkdir() + (install_root / "run" / "service.pid").write_text("2468", encoding="utf-8") (install_root / ".git").mkdir() (install_root / ".git" / "HEAD").write_text("ref: refs/heads/main", encoding="utf-8") @@ -1697,14 +1335,37 @@ def test_replace_install_dir_copies_dot_flocks_plugins_from_source( assert (inst_plugins / "fofa" / "_provider.yaml").read_text(encoding="utf-8") == "version: new" assert (inst_plugins / "new_release_plugin" / "tool.yaml").read_text(encoding="utf-8") == "name: new" assert not (inst_plugins / "obsolete_plugin").exists() - assert (install_root / "flocks.json").read_text(encoding="utf-8") == '{"version": "new"}' + assert (install_root / "flocks.json").read_text(encoding="utf-8") == '{"keep": true}' + assert (install_root / "run" / "service.pid").read_text(encoding="utf-8") == "2468" assert (install_root / ".git" / "HEAD").read_text(encoding="utf-8") == "ref: refs/heads/main" @pytest.mark.asyncio -async def test_perform_update_schedules_handoff_with_deferred_handover( +@pytest.mark.parametrize( + ("locale", "expected_sources", "expected_mirror_args", "wait_for_handoff"), + [ + pytest.param("en-US", ["github", "gitee"], [], False, id="english-detached-upgrade"), + pytest.param( + "zh-CN", + ["gitee", "github"], + [ + "--uv-default-index", + "https://mirrors.aliyun.com/pypi/simple", + "--npm-registry", + "https://registry.npmmirror.com/", + ], + True, + id="chinese-waited-upgrade", + ), + ], +) +async def test_perform_update_only_stages_source_and_schedules_upgrade_handoff( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, + locale: str, + expected_sources: list[str], + expected_mirror_args: list[str], + wait_for_handoff: bool, ) -> None: archive_path = tmp_path / "flocks.zip" archive_path.write_text("archive", encoding="utf-8") @@ -1720,11 +1381,13 @@ async def test_perform_update_schedules_handoff_with_deferred_handover( events: list[str] = [] popen_calls: list[list[str]] = [] + download_sources: list[str] = [] + progress_stages: list[str] = [] async def fake_get_updater_config(): return SimpleNamespace( archive_format="zip", - sources=["github"], + sources=["github", "gitee"], repo="AgentFlocks/Flocks", token=None, gitee_token=None, @@ -1733,7 +1396,8 @@ async def fake_get_updater_config(): gitee_repo=None, ) - async def fake_download_with_fallback(**_kwargs): + async def fake_download_with_fallback(**kwargs): + download_sources.extend(kwargs["sources"]) return archive_path async def fake_sleep(_seconds) -> None: @@ -1747,15 +1411,30 @@ async def fake_sleep(_seconds) -> None: monkeypatch.setattr(updater, "_get_repo_root", lambda: install_root) monkeypatch.setattr(updater, "get_current_version", lambda: "2026.3.31") monkeypatch.setattr(updater, "_download_with_fallback", fake_download_with_fallback) - monkeypatch.setattr(updater, "_backup_current_version", lambda *_args, **_kwargs: tmp_path / "backup.tar.gz") + monkeypatch.setattr( + updater, + "_backup_current_version", + lambda *_args, **_kwargs: events.append("backup") or tmp_path / "backup.tar.gz", + ) monkeypatch.setattr(updater, "_extract_archive", lambda *_args, **_kwargs: staged_root) monkeypatch.setattr( updater, "_find_executable", lambda name: "/usr/bin/npm" if name in {"npm", "npm.cmd"} else "/usr/bin/uv", ) - monkeypatch.setattr(updater, "_prepare_upgrade_handover", lambda _version: events.append("handover") or {}) - monkeypatch.setattr(updater, "_handoff_service_config", lambda: service_manager.ServiceConfig()) + config = service_manager.ServiceConfig( + backend_host="10.20.30.40", + backend_port=9527, + frontend_host="10.20.30.40", + frontend_port=9527, + legacy_backend_host="0.0.0.0", + legacy_backend_port=9000, + ) + monkeypatch.setattr( + updater, + "_capture_service_snapshot", + lambda: updater.ServiceSnapshot(config=config, daemon_pid=2468, was_running=True), + ) monkeypatch.setattr( updater, "_replace_install_dir", @@ -1764,43 +1443,111 @@ async def fake_sleep(_seconds) -> None: ) monkeypatch.setattr(updater, "_build_restart_argv", lambda install_root=None: ["/usr/bin/python3", "-m", "flocks.cli.main", "start"]) monkeypatch.setattr(updater.asyncio, "sleep", fake_sleep) - monkeypatch.setattr(updater, "_rollback_failed_update", lambda *_args: events.append("rollback")) - monkeypatch.setattr(updater, "rollback_upgrade_handover", lambda *_args: events.append("rollback_handover")) monkeypatch.setattr( updater, "_spawn_restart_handoff", - lambda argv, **_kwargs: popen_calls.append(list(argv)) or SimpleNamespace(pid=4321), + lambda argv, **_kwargs: popen_calls.append(list(argv)) + or SimpleNamespace(pid=4321, wait=lambda: events.append("wait") or 0), ) monkeypatch.setattr(updater.os, "_exit", lambda code: (_ for _ in ()).throw(SystemExit(code))) - with pytest.raises(SystemExit, match="0"): - async for _step in updater.perform_update("2026.4.1"): - pass - - assert events[:2] == ["replace", "sleep"] - assert "handover" not in events + if wait_for_handoff: + async for step in updater.perform_update( + "2026.4.1", + locale=locale, + wait_for_handoff=True, + ): + progress_stages.append(step.stage) + else: + with pytest.raises(SystemExit, match="0"): + async for step in updater.perform_update("2026.4.1", locale=locale): + progress_stages.append(step.stage) + + expected_events = ["backup", "sleep"] + expected_stages = ["fetching", "backing_up", "applying", "restarting"] + if wait_for_handoff: + expected_events.append("wait") + expected_stages.append("done") + + assert events == expected_events + assert progress_stages == expected_stages + assert download_sources == expected_sources assert len(popen_calls) == 1 handoff_argv = popen_calls[0] assert handoff_argv[:3] == ["/usr/bin/python3", "-m", "flocks.updater.restart_handoff"] assert "--uv-path" in handoff_argv assert "--version" in handoff_argv - assert "--prepare-handover" in handoff_argv[: handoff_argv.index("--")] - assert handoff_argv[handoff_argv.index("--") + 1 :] == [ - "/usr/bin/python3", - "-m", - "flocks.cli.main", - "start", - "--no-browser", - "--skip-webui-build", - "--host", - "127.0.0.1", - "--port", - "5173", - "--server-host", - "127.0.0.1", - "--server-port", - "8000", - ] + assert "--mode" in handoff_argv + assert handoff_argv[handoff_argv.index("--mode") + 1] == "upgrade" + assert handoff_argv[handoff_argv.index("--content-root") + 1] == str(staged_root) + assert handoff_argv[handoff_argv.index("--backup-path") + 1] == str(tmp_path / "backup.tar.gz") + assert handoff_argv[handoff_argv.index("--daemon-pid") + 1] == "2468" + assert "--was-running" in handoff_argv + assert "--prepare-handover" not in handoff_argv + assert ("--parent-pid" in handoff_argv) is not wait_for_handoff + if expected_mirror_args: + mirror_index = handoff_argv.index("--uv-default-index") + assert handoff_argv[mirror_index : mirror_index + len(expected_mirror_args)] == expected_mirror_args + else: + assert "--uv-default-index" not in handoff_argv + assert "--npm-registry" not in handoff_argv + assert handoff_argv[handoff_argv.index("--") + 1 :] == ["/usr/bin/python3"] + + +@pytest.mark.asyncio +async def test_perform_update_aborts_before_handoff_when_backup_raises( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + archive_path = tmp_path / "flocks.zip" + archive_path.write_text("archive", encoding="utf-8") + update_dir = tmp_path / "update" + update_dir.mkdir() + staged_root = update_dir / "staged" + staged_root.mkdir() + install_root = tmp_path / "install-root" + install_root.mkdir() + handoff_spawned = False + + async def fake_get_updater_config(): + return SimpleNamespace( + archive_format="zip", + sources=["github"], + repo="AgentFlocks/Flocks", + token=None, + gitee_token=None, + backup_retain_count=3, + base_url=None, + gitee_repo=None, + ) + + async def fake_download_with_fallback(**_kwargs): + return archive_path + + def fail_backup(*_args, **_kwargs): + raise OSError("backup directory is read-only") + + def record_handoff(*_args, **_kwargs): + nonlocal handoff_spawned + handoff_spawned = True + + monkeypatch.setattr(updater, "_get_updater_config", fake_get_updater_config) + monkeypatch.setattr(updater, "_get_repo_root", lambda: install_root) + monkeypatch.setattr(updater, "get_current_version", lambda: "2026.3.31") + monkeypatch.setattr(updater.tempfile, "mkdtemp", lambda **_kwargs: str(update_dir)) + monkeypatch.setattr(updater, "_download_with_fallback", fake_download_with_fallback) + monkeypatch.setattr(updater, "_extract_archive", lambda *_args, **_kwargs: staged_root) + monkeypatch.setattr(updater, "_find_executable", lambda _name: "/usr/bin/uv") + monkeypatch.setattr(updater, "_backup_current_version", fail_backup) + monkeypatch.setattr(updater, "_spawn_restart_handoff", record_handoff) + monkeypatch.setattr(updater, "_record_update_journal", lambda _message: None) + + progresses = [step async for step in updater.perform_update("2026.4.1")] + + assert progresses[-1].stage == "error" + assert progresses[-1].message == "Failed to back up the current source; the update was not applied." + assert handoff_spawned is False + assert not update_dir.exists() @pytest.mark.asyncio @@ -1873,10 +1620,13 @@ async def fake_sleep(_seconds) -> None: monkeypatch.setattr(updater.asyncio, "sleep", fake_sleep) monkeypatch.setattr( updater, - "_prepare_upgrade_handover", - lambda _version: (_ for _ in ()).throw(RuntimeError("handover boom")), + "_capture_service_snapshot", + lambda: updater.ServiceSnapshot( + config=service_manager.ServiceConfig(), + daemon_pid=None, + was_running=False, + ), ) - monkeypatch.setattr(updater, "_restore_backup_if_possible", lambda *_args: events.append("restore")) monkeypatch.setattr( updater, "_spawn_restart_handoff", @@ -1888,13 +1638,14 @@ async def fake_sleep(_seconds) -> None: async for _step in updater.perform_update("2026.4.1"): pass - assert events == ["replace"] + assert events == [] assert len(popen_calls) == 1 - assert "--prepare-handover" in popen_calls[0][: popen_calls[0].index("--")] + assert "--mode" in popen_calls[0] + assert "--prepare-handover" not in popen_calls[0] @pytest.mark.asyncio -async def test_perform_update_uses_cn_mirror_profile_for_sources_and_dependency_commands( +async def test_perform_update_rejects_source_update_without_handoff( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: @@ -1958,25 +1709,14 @@ async def fake_run_async(cmd, cwd=None, timeout=None, env=None): ) ] - assert progresses[-1].stage == "done" + assert progresses[-1].stage == "error" + assert "detached handoff" in progresses[-1].message assert captured["sources"] == ["gitee", "github"] - assert run_calls == [ - ( - [ - "/usr/bin/uv", - "sync", - "--frozen", - "--no-python-downloads", - "--default-index", - "https://mirrors.aliyun.com/pypi/simple", - ], - None, - ), - ] + assert run_calls == [] @pytest.mark.asyncio -async def test_perform_update_retries_cn_uv_sync_with_default_source( +async def test_dependency_sync_retries_cn_mirror_with_default_source( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: @@ -2005,7 +1745,6 @@ async def fake_run_async(cmd, cwd=None, timeout=None, env=None): if cmd == [ "/usr/bin/uv", "sync", - "--frozen", "--no-python-downloads", "--default-index", "https://mirrors.aliyun.com/pypi/simple", @@ -2032,24 +1771,74 @@ async def fake_sleep(_seconds) -> None: monkeypatch.setattr(updater, "_write_version_marker", lambda _v: None) monkeypatch.setattr(updater.asyncio, "sleep", fake_sleep) - progresses = [step async for step in updater.perform_update("2026.4.1", restart=False, locale="zh-CN")] + error = await updater._sync_project_dependencies( + uv_path="/usr/bin/uv", + install_root=install_root, + uv_default_index="https://mirrors.aliyun.com/pypi/simple", + env=None, + ) - assert progresses[-1].stage == "done" + assert error is None assert run_calls == [ [ "/usr/bin/uv", "sync", - "--frozen", "--no-python-downloads", "--default-index", "https://mirrors.aliyun.com/pypi/simple", ], - ["/usr/bin/uv", "sync", "--frozen", "--no-python-downloads"], + ["/usr/bin/uv", "sync", "--no-python-downloads"], ] @pytest.mark.asyncio -async def test_perform_update_prefers_bundled_npm_for_windows_frontend_rebuild( +async def test_dependency_sync_retries_default_source_after_timeout_and_logs_output( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + calls: list[list[str]] = [] + warnings: list[tuple[str, dict[str, object]]] = [] + + async def fake_run_async(cmd, **_kwargs): + calls.append(list(cmd)) + if "--default-index" in cmd: + raise subprocess.TimeoutExpired( + cmd=cmd, + timeout=300, + output=b"mirror stdout", + stderr=b"mirror stderr", + ) + return 0, "", "" + + monkeypatch.setattr(updater, "_run_async", fake_run_async) + monkeypatch.setattr(updater.log, "warning", lambda event, payload: warnings.append((event, payload))) + + error = await updater._sync_project_dependencies( + uv_path="uv", + install_root=tmp_path, + uv_default_index="https://mirrors.aliyun.com/pypi/simple", + sync_timeout=300, + ) + + assert error is None + assert calls == [ + [ + "uv", + "sync", + "--no-python-downloads", + "--default-index", + "https://mirrors.aliyun.com/pypi/simple", + ], + ["uv", "sync", "--no-python-downloads"], + ] + timeout_payload = next(payload for event, payload in warnings if event == "updater.dependencies.sync_timeout") + assert timeout_payload["stdout"] == "mirror stdout" + assert timeout_payload["stderr"] == "mirror stderr" + assert timeout_payload["retry_without_default_index"] is True + + +@pytest.mark.asyncio +async def test_frontend_build_prefers_bundled_npm_on_windows( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: @@ -2112,9 +1901,13 @@ async def fake_run_async(cmd, cwd=None, timeout=None, env=None): monkeypatch.delenv("FLOCKS_INSTALL_ROOT", raising=False) monkeypatch.setenv("PATH", "/usr/bin:/bin") - progresses = [step async for step in updater.perform_update("2026.4.1", restart=False, locale="zh-CN")] + shutil.copytree(staged_webui, install_root / "webui") + frontend_error = await updater._build_frontend_workspace( + install_root / "webui", + npm_registry="https://registry.npmmirror.com/", + ) - assert progresses[-1].stage == "done" + assert frontend_error is None frontend_calls = [ call for call in run_calls if call[0][0] == str(bundled_npm) ] @@ -2133,7 +1926,7 @@ async def fake_run_async(cmd, cwd=None, timeout=None, env=None): @pytest.mark.asyncio -async def test_perform_update_retries_windows_frontend_with_system_npm_after_bundled_build_failure( +async def test_frontend_build_retries_system_npm_after_bundled_build_failure( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: @@ -2224,9 +2017,13 @@ def fake_find(name: str) -> str | None: monkeypatch.delenv("FLOCKS_INSTALL_ROOT", raising=False) monkeypatch.setenv("PATH", "/usr/bin:/bin") - progresses = [step async for step in updater.perform_update("2026.4.1", restart=False, locale="zh-CN")] + shutil.copytree(staged_webui, install_webui) + frontend_error = await updater._build_frontend_workspace( + install_webui, + npm_registry="https://registry.npmmirror.com/", + ) - assert progresses[-1].stage == "done" + assert frontend_error is None frontend_calls = [ call for call in run_calls if call[0][0] in {str(bundled_npm), system_npm} ] @@ -2364,7 +2161,7 @@ async def fake_run_async(cmd, cwd=None, timeout=None, env=None): @pytest.mark.asyncio -async def test_perform_update_retries_windows_frontend_with_full_timeout_after_bundled_install_and_ci_timeout( +async def test_frontend_build_retries_system_npm_after_bundled_timeouts( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: @@ -2454,9 +2251,13 @@ def fake_find(name: str) -> str | None: monkeypatch.delenv("FLOCKS_INSTALL_ROOT", raising=False) monkeypatch.setenv("PATH", "/usr/bin:/bin") - progresses = [step async for step in updater.perform_update("2026.4.1", restart=False, locale="zh-CN")] + shutil.copytree(staged_webui, install_webui) + frontend_error = await updater._build_frontend_workspace( + install_webui, + npm_registry="https://registry.npmmirror.com/", + ) - assert progresses[-1].stage == "done" + assert frontend_error is None frontend_calls = [ call for call in run_calls if call[0][0] in {str(bundled_npm), system_npm} ] @@ -2520,7 +2321,7 @@ async def fake_download_with_fallback(**_kwargs): @pytest.mark.asyncio -async def test_perform_update_retries_uv_sync_on_first_failure( +async def test_dependency_sync_retries_first_failure( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: @@ -2574,16 +2375,18 @@ async def fake_sleep(_s): monkeypatch.setattr(updater, "_write_version_marker", lambda _v: None) monkeypatch.setattr(updater.asyncio, "sleep", fake_sleep) - progresses = [ - step async for step in updater.perform_update("2026.4.1", restart=False) - ] + error = await updater._sync_project_dependencies( + uv_path="/usr/bin/uv", + install_root=install_root, + env=None, + ) - assert progresses[-1].stage == "done" + assert error is None assert call_count == 2 @pytest.mark.asyncio -async def test_perform_update_syncs_windows_venv_in_place( +async def test_dependency_sync_updates_windows_venv_in_place( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: @@ -2633,16 +2436,20 @@ async def fake_run_async(cmd, cwd=None, timeout=None, env=None): monkeypatch.setattr(updater, "_write_version_marker", lambda _v: None) monkeypatch.setattr(updater, "_refresh_global_cli_entry", lambda _root: None) - progresses = [step async for step in updater.perform_update("2026.4.1", restart=False)] + error = await updater._sync_project_dependencies( + uv_path=r"C:\tools\uv.exe", + install_root=install_root, + env=None, + ) - assert progresses[-1].stage == "done" - assert sync_calls == [([r"C:\tools\uv.exe", "sync", "--frozen", "--no-python-downloads"], install_root)] + assert error is None + assert sync_calls == [([r"C:\tools\uv.exe", "sync", "--no-python-downloads"], install_root)] assert (install_root / ".venv" / "Scripts" / "python.exe").read_text(encoding="utf-8") == "old" assert not (install_root / ".venv.flocks_backup").exists() @pytest.mark.asyncio -async def test_perform_update_rolls_back_when_windows_uv_sync_times_out( +async def test_dependency_sync_reports_windows_timeout_without_rollback( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: @@ -2699,31 +2506,32 @@ async def fake_run_async(cmd, cwd=None, timeout=None, env=None): ) monkeypatch.setattr(updater, "_build_uv_sync_env", lambda: None) monkeypatch.setattr(updater, "_replace_install_dir", lambda *_a, **_kw: None) - monkeypatch.setattr(updater, "_restore_backup_if_possible", lambda *_a: events.append("restore")) - progresses = [step async for step in updater.perform_update("2026.4.1", restart=False)] + error = await updater._sync_project_dependencies( + uv_path=r"C:\tools\uv.exe", + install_root=install_root, + env=None, + ) - assert progresses[-1].stage == "error" expected_timeout = updater._dependency_sync_timeout_seconds() - assert progresses[-1].message == ( + assert error == ( f"Dependency sync timed out after {expected_timeout}s while running uv sync." ) - assert events == ["restore"] + assert events == [] assert (install_root / ".venv" / "Scripts" / "python.exe").read_text(encoding="utf-8") == "new" assert not (install_root / ".venv.flocks_failed").exists() assert not (install_root / ".venv.flocks_backup").exists() @pytest.mark.asyncio -async def test_perform_update_fails_after_uv_sync_retry_exhausted( +async def test_dependency_sync_fails_after_retry_exhausted( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: - """When uv sync fails twice, the updater should give up and rollback.""" + """When uv sync fails twice, the updater should return the final error.""" archive_path = tmp_path / "flocks.tar.gz" archive_path.write_text("archive", encoding="utf-8") staged_root = tmp_path / "staged" - rolled_back = False async def fake_get_updater_config(): return SimpleNamespace( @@ -2745,10 +2553,6 @@ async def fake_run_async(cmd, cwd=None, timeout=None, env=None): async def fake_download(**_kw): return archive_path - def fake_restore(*_args): - nonlocal rolled_back - rolled_back = True - monkeypatch.setattr(updater, "_get_updater_config", fake_get_updater_config) monkeypatch.setattr(updater, "_get_repo_root", lambda: tmp_path / "install-root") monkeypatch.setattr(updater, "get_current_version", lambda: "2026.3.31") @@ -2758,25 +2562,25 @@ def fake_restore(*_args): monkeypatch.setattr(updater, "_run_async", fake_run_async) monkeypatch.setattr(updater, "_find_executable", lambda _name: "/usr/bin/uv") monkeypatch.setattr(updater, "_build_uv_sync_env", lambda: None) + async def fake_sleep(_s): pass monkeypatch.setattr(updater, "_replace_install_dir", lambda *_a, **_kw: None) - monkeypatch.setattr(updater, "_restore_backup_if_possible", fake_restore) monkeypatch.setattr(updater.asyncio, "sleep", fake_sleep) - progresses = [ - step async for step in updater.perform_update("2026.4.1", restart=False) - ] + error = await updater._sync_project_dependencies( + uv_path="/usr/bin/uv", + install_root=tmp_path / "install-root", + env=None, + ) - error_events = [p for p in progresses if p.stage == "error"] - assert len(error_events) == 1 - assert "resolution failed" in error_events[0].message - assert rolled_back + assert error is not None + assert "resolution failed" in error @pytest.mark.asyncio -async def test_perform_update_repairs_broken_uv_managed_python_cache_before_retrying_sync( +async def test_dependency_sync_repairs_broken_uv_managed_python_cache_before_retry( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: @@ -2841,18 +2645,20 @@ async def fake_sleep(seconds): monkeypatch.setattr(updater, "_repair_windows_uv_managed_python_install", lambda text: repaired_messages.append(text) or Path(r"C:\Users\worker\AppData\Roaming\uv\python\cpython-3.12-windows-x86_64-none")) monkeypatch.setattr(updater.asyncio, "sleep", fake_sleep) - progresses = [ - step async for step in updater.perform_update("2026.4.1", restart=False) - ] + error = await updater._sync_project_dependencies( + uv_path=r"C:\Users\worker\AppData\Local\Programs\Flocks\tools\uv\uv.exe", + install_root=install_root, + env=None, + ) - assert progresses[-1].stage == "done" + assert error is None assert call_count == 2 assert sleep_calls == [2] assert len(repaired_messages) == 1 @pytest.mark.asyncio -async def test_perform_update_rolls_back_when_replace_fails_on_windows_locked_file( +async def test_perform_update_without_handoff_never_attempts_source_replace( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: @@ -2903,7 +2709,6 @@ async def fake_run_async(cmd, cwd=None, timeout=None, env=None): "_find_executable", lambda name: "/usr/bin/npm" if name in {"npm", "npm.cmd"} else "/usr/bin/uv", ) - monkeypatch.setattr(updater, "_prepare_upgrade_handover", lambda _version: events.append("handover") or {}) monkeypatch.setattr( updater, "_replace_install_dir", @@ -2913,18 +2718,16 @@ async def fake_run_async(cmd, cwd=None, timeout=None, env=None): ) ), ) - monkeypatch.setattr(updater, "_restore_backup_if_possible", lambda *_args: events.append("restore")) - - progresses = [step async for step in updater.perform_update("2026.4.1")] + progresses = [step async for step in updater.perform_update("2026.4.1", restart=False)] assert progresses[-1].stage == "error" - assert "WinError 5" in progresses[-1].message - assert events == ["restore"] + assert "detached handoff" in progresses[-1].message + assert events == [] assert "handover" not in events @pytest.mark.asyncio -async def test_perform_update_reports_windows_file_lock_without_stopping_current_backend( +async def test_perform_update_without_handoff_does_not_touch_locked_windows_files( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: @@ -2983,11 +2786,16 @@ def fake_replace_install_dir(*_args, **_kwargs): "_find_executable", lambda name: "/usr/bin/npm" if name in {"npm", "npm.cmd"} else "/usr/bin/uv", ) - monkeypatch.setattr(updater, "_prepare_upgrade_handover", lambda _version: events.append("handover")) - monkeypatch.setattr(updater, "_handoff_service_config", lambda: service_manager.ServiceConfig()) + monkeypatch.setattr( + updater, + "_capture_service_snapshot", + lambda: updater.ServiceSnapshot( + config=service_manager.ServiceConfig(), + daemon_pid=None, + was_running=False, + ), + ) monkeypatch.setattr(updater, "_replace_install_dir", fake_replace_install_dir) - monkeypatch.setattr(updater, "_rollback_failed_update", lambda *_args: events.append("rollback")) - monkeypatch.setattr(updater, "_restore_backup_if_possible", lambda *_args: events.append("restore")) monkeypatch.setattr(updater, "_build_restart_argv", lambda install_root=None: [r"C:\tool\python.exe", "-m", "flocks.cli.main", "start"]) monkeypatch.setattr( updater, @@ -2996,17 +2804,18 @@ def fake_replace_install_dir(*_args, **_kwargs): ) monkeypatch.setattr(updater.os, "_exit", lambda code: (_ for _ in ()).throw(SystemExit(code))) - progresses = [step async for step in updater.perform_update("2026.4.1")] + progresses = [step async for step in updater.perform_update("2026.4.1", restart=False)] assert progresses[-1].stage == "error" - assert "WinError 32" in progresses[-1].message - assert events == ["replace-1", "restore"] + assert "detached handoff" in progresses[-1].message + assert replace_attempts["count"] == 0 + assert events == [] assert "handover" not in events assert "popen" not in events @pytest.mark.asyncio -async def test_perform_update_reports_frontend_dependency_install_timeout( +async def test_frontend_workspace_reports_dependency_install_timeout( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: @@ -3064,23 +2873,19 @@ async def fake_run_async(cmd, cwd=None, timeout=None, env=None): or shutil.copytree(staged_webui, install_webui, dirs_exist_ok=True), ) monkeypatch.setattr(updater, "_write_version_marker", lambda _v: None) - monkeypatch.setattr(updater, "_restore_backup_if_possible", lambda *_args: events.append("restore")) - progresses = [step async for step in updater.perform_update("2026.4.1", restart=False)] + shutil.copytree(staged_webui, install_webui) + frontend_error = await updater._build_frontend_workspace(install_webui) - assert progresses[-1].stage == "error" - assert progresses[-1].message == "Frontend dependency install timed out after 300s while running npm ci." + assert frontend_error == "Frontend dependency install timed out after 300s while running npm ci." assert events == [ - "replace", - "/usr/bin/uv sync --frozen --no-python-downloads", "/usr/bin/npm install", "/usr/bin/npm ci", - "restore", ] @pytest.mark.asyncio -async def test_perform_update_rolls_back_handover_when_current_frontend_build_fails( +async def test_frontend_workspace_reports_build_failure_without_rollback( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: @@ -3144,13 +2949,12 @@ async def fake_run_async(cmd, cwd=None, timeout=None, env=None): lambda *_args, **_kwargs: events.append("replace") or shutil.copytree(staged_webui, install_webui, dirs_exist_ok=True), ) - monkeypatch.setattr(updater, "_restore_backup_if_possible", lambda *_args: events.append("restore")) - progresses = [step async for step in updater.perform_update("2026.4.1", restart=False)] + shutil.copytree(staged_webui, install_webui) + frontend_error = await updater._build_frontend_workspace(install_webui) - assert progresses[-1].stage == "error" - assert progresses[-1].message == "Frontend build failed: boom" - assert events == ["replace", "uv-sync", "npm-install", "npm-build", "restore"] + assert frontend_error == "Frontend build failed: boom" + assert events == ["npm-install", "npm-build"] @pytest.mark.asyncio @@ -3204,7 +3008,6 @@ async def fake_run_async(cmd, cwd=None, timeout=None, env=None): "_find_executable", lambda name: "/usr/bin/npm" if name in {"npm", "npm.cmd"} else "/usr/bin/uv", ) - monkeypatch.setattr(updater, "_prepare_upgrade_handover", lambda _version: events.append("handover")) monkeypatch.setattr(updater, "_replace_install_dir", lambda *_args, **_kwargs: None) monkeypatch.setattr(updater, "_write_version_marker", lambda _v: None) @@ -3215,7 +3018,7 @@ async def fake_run_async(cmd, cwd=None, timeout=None, env=None): await gen.aclose() assert "handover" not in events - assert updater._read_upgrade_state() is None + assert not updater._upgrade_result_path().exists() @pytest.mark.asyncio @@ -3271,7 +3074,6 @@ async def fake_run_async(cmd, cwd=None, timeout=None, env=None): monkeypatch.setattr(updater, "_write_version_marker", lambda _v: None) monkeypatch.setattr(updater, "_refresh_global_cli_entry", lambda _root: None) monkeypatch.setattr(updater, "_build_restart_argv", lambda install_root=None: [r"C:\tool\python.exe", "-m", "flocks.cli.main", "start"]) - monkeypatch.setattr(updater, "_prepare_upgrade_handover", lambda _version: events.append("handover")) monkeypatch.setattr(updater, "_handoff_service_config", lambda: service_manager.ServiceConfig()) monkeypatch.setattr( updater, @@ -3291,29 +3093,15 @@ async def fake_run_async(cmd, cwd=None, timeout=None, env=None): assert handoff_argv[:3] == [r"C:\tool\python.exe", "-m", "flocks.updater.restart_handoff"] assert "--parent-pid" in handoff_argv assert "--backend-port" in handoff_argv - assert "--prepare-handover" in handoff_argv[: handoff_argv.index("--")] - assert handoff_argv[handoff_argv.index("--") + 1 :] == [ - r"C:\tool\python.exe", - "-m", - "flocks.cli.main", - "start", - "--no-browser", - "--skip-webui-build", - "--host", - "127.0.0.1", - "--port", - "5173", - "--server-host", - "127.0.0.1", - "--server-port", - "8000", - ] + assert handoff_argv[handoff_argv.index("--mode") + 1] == "upgrade" + assert "--prepare-handover" not in handoff_argv + assert handoff_argv[handoff_argv.index("--") + 1 :] == [r"C:\tool\python.exe"] assert events == [] assert "execv" not in events @pytest.mark.asyncio -async def test_perform_update_stops_when_windows_restart_runtime_validation_fails( +async def test_validate_restart_runtime_reports_missing_windows_python( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: @@ -3364,15 +3152,13 @@ async def fake_validate_restart_runtime(_install_root: Path) -> str | None: ) monkeypatch.setattr(updater, "_replace_install_dir", lambda *_args, **_kwargs: None) monkeypatch.setattr(updater, "_write_version_marker", lambda _v: events.append("marker")) - monkeypatch.setattr(updater, "_restore_backup_if_possible", lambda *_args: events.append("restore")) monkeypatch.setattr(updater, "_validate_restart_runtime", fake_validate_restart_runtime) monkeypatch.setattr(updater.subprocess, "Popen", lambda *_args, **_kwargs: events.append("popen")) - progresses = [step async for step in updater.perform_update("2026.4.1", restart=False)] + validation_error = await updater._validate_restart_runtime(tmp_path / "install-root") - assert progresses[-1].stage == "error" - assert progresses[-1].message == "Restart runtime is missing: .venv/Scripts/python.exe" - assert events == ["restore"] + assert validation_error == "Restart runtime is missing: .venv/Scripts/python.exe" + assert events == [] @pytest.mark.asyncio @@ -3426,12 +3212,12 @@ async def fake_run_async(cmd, cwd=None, timeout=None, env=None): FileNotFoundError("python.exe not found"), ), ) - monkeypatch.setattr(updater, "_prepare_upgrade_handover", lambda _version: events.append("handover")) progresses = [step async for step in updater.perform_update("2026.4.1")] assert progresses[-1].stage == "error" - assert "Failed to build restart command" in progresses[-1].message + assert "restarting" not in [step.stage for step in progresses] + assert "Failed to prepare restart handoff" in progresses[-1].message assert "handover" not in events @@ -3488,9 +3274,7 @@ async def fake_run_async(cmd, cwd=None, timeout=None, env=None): monkeypatch.setattr(updater, "_write_version_marker", lambda _v: None) monkeypatch.setattr(updater, "_refresh_global_cli_entry", lambda _root: None) monkeypatch.setattr(updater, "_build_restart_argv", lambda install_root=None: [r"C:\tool\python.exe", "-m", "flocks.cli.main"]) - monkeypatch.setattr(updater, "_prepare_upgrade_handover", lambda _version: events.append("handover")) monkeypatch.setattr(updater, "_handoff_service_config", lambda: service_manager.ServiceConfig()) - monkeypatch.setattr(updater, "rollback_upgrade_handover", lambda: events.append("rollback_handover")) monkeypatch.setattr( updater, "_spawn_restart_handoff", diff --git a/tests/updater/test_updater_console_manifest_bundle.py b/tests/updater/test_updater_console_manifest_bundle.py index 118c75f55..c43ed1c58 100644 --- a/tests/updater/test_updater_console_manifest_bundle.py +++ b/tests/updater/test_updater_console_manifest_bundle.py @@ -1,7 +1,8 @@ from __future__ import annotations -import zipfile import json +import sys +import zipfile from datetime import UTC, datetime, timedelta from types import SimpleNamespace @@ -527,6 +528,7 @@ async def _after_uninstall(): monkeypatch.setattr(updater, "_is_pro_component_installed", lambda: True) monkeypatch.setattr(updater, "_find_executable", lambda name: "/usr/bin/uv" if name == "uv" else None) monkeypatch.setattr(updater, "_uninstall_pro_component", _fake_uninstall_pro_component) + monkeypatch.setattr(updater, "_write_version_marker", lambda _version: None) monkeypatch.setattr(updater, "_refresh_global_cli_entry", lambda _install_root: None) progresses = [ @@ -545,6 +547,68 @@ async def _after_uninstall(): assert len(archived_marker) == 1 +@pytest.mark.asyncio +async def test_perform_pro_bundle_downgrade_uses_restart_handoff_spawn( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + from flocks.updater import deploy as deploy_mod + + flocks_root = tmp_path / "flocks-root" + run_dir = flocks_root / "run" + run_dir.mkdir(parents=True) + (run_dir / "pro-bundle-installed.json").write_text( + json.dumps( + { + "bundle_version": "v2026.6.24", + "core_version": "v2026.6.21", + "flockspro_component_version": "v2026.6.24-pro", + "installed_at": "2026-06-24T08:00:00+00:00", + } + ), + encoding="utf-8", + ) + install_root = tmp_path / "install-root" + install_root.mkdir() + handoff_argv = [sys.executable, "-m", "flocks.updater.restart_handoff"] + spawn_calls: list[tuple[list[str], Path]] = [] + + async def fake_uninstall_pro_component(*, uv_path, install_root, env): + return None + + async def fake_sleep(_seconds: float) -> None: + return None + + monkeypatch.setenv("FLOCKS_ROOT", str(flocks_root)) + monkeypatch.setattr(deploy_mod, "detect_deploy_mode", lambda: "source") + monkeypatch.setattr(updater, "_get_repo_root", lambda: install_root) + monkeypatch.setattr(updater, "get_current_version", lambda: "2026.6.24") + monkeypatch.setattr(updater, "_is_pro_component_installed", lambda: True) + monkeypatch.setattr(updater, "_find_executable", lambda name: "/usr/bin/uv" if name == "uv" else None) + monkeypatch.setattr(updater, "_uninstall_pro_component", fake_uninstall_pro_component) + monkeypatch.setattr(updater, "_write_version_marker", lambda _version: None) + monkeypatch.setattr(updater, "_build_restart_argv", lambda _install_root: [sys.executable]) + monkeypatch.setattr(updater, "_build_restart_handoff_argv", lambda *_args, **_kwargs: handoff_argv) + monkeypatch.setattr(updater.asyncio, "sleep", fake_sleep) + monkeypatch.setattr( + updater, + "_spawn_restart_handoff", + lambda argv, *, cwd: spawn_calls.append((argv, cwd)) or SimpleNamespace(pid=4321), + ) + monkeypatch.setattr( + updater.subprocess, + "Popen", + lambda *_args, **_kwargs: pytest.fail("downgrade must use the restart handoff spawn helper"), + ) + monkeypatch.setattr(updater.os, "_exit", lambda code: (_ for _ in ()).throw(SystemExit(code))) + + with pytest.raises(SystemExit, match="0"): + async for _step in updater.perform_pro_bundle_downgrade(): + pass + + assert spawn_calls == [(handoff_argv, install_root)] + + @pytest.mark.asyncio async def test_load_console_session_token_falls_back_to_shared_session( monkeypatch: pytest.MonkeyPatch, @@ -811,7 +875,7 @@ def stream(self, method, url, headers=None): @pytest.mark.asyncio -async def test_perform_pro_bundle_install_replaces_core_and_installs_wheel( +async def test_core_pro_bundle_requires_source_upgrade_handoff( monkeypatch: pytest.MonkeyPatch, tmp_path, ) -> None: @@ -884,20 +948,12 @@ async def _fake_run_async(cmd, **_kwargs): monkeypatch.setattr(updater, "_run_async", _fake_run_async) progresses = [step async for step in updater.perform_pro_bundle_install(restart=False)] - assert progresses[-1].stage == "done" - assert (install_root / "new_core.py").is_file() - assert not (install_root / "old_core.py").exists() - assert any(cmd[:2] == ["/usr/bin/uv", "sync"] for cmd in captured) - pip_installs = [cmd for cmd in captured if cmd[:3] == ["/usr/bin/uv", "pip", "install"]] - assert pip_installs - assert "--no-deps" in pip_installs[-1] - assert str(wheel.name) in pip_installs[-1][-1] - marker = tmp_path / "flocks-root" / "run" / "pro-bundle-installed.json" - assert marker.is_file() - marker_payload = __import__("json").loads(marker.read_text(encoding="utf-8")) - assert version_writes == ["2026.5.10"] - assert marker_payload["bundle_version"] == "v2026.5.11" - assert marker_payload["core_version"] == "v2026.5.10" + assert progresses[-1].stage == "error" + assert "detached handoff" in progresses[-1].message + assert not (install_root / "new_core.py").exists() + assert (install_root / "old_core.py").exists() + assert captured == [] + assert version_writes == [] @pytest.mark.asyncio @@ -977,6 +1033,8 @@ async def _fake_run_async(cmd, **_kwargs): progresses = [step async for step in updater.perform_pro_bundle_install(restart=False)] assert progresses[-1].stage == "done" + assert "backing_up" not in [step.stage for step in progresses] + assert "syncing" not in [step.stage for step in progresses] assert any("Keeping local Flocks v2026.6.18" in step.message for step in progresses) assert (install_root / "current_core.py").is_file() assert not (install_root / "older_core.py").exists() diff --git a/tests/workflow/test_fs_store.py b/tests/workflow/test_fs_store.py index 711f072ee..f5c7b4c31 100644 --- a/tests/workflow/test_fs_store.py +++ b/tests/workflow/test_fs_store.py @@ -53,6 +53,45 @@ def test_read_workflow_from_fs_refreshes_cached_workspace_root( assert fs_store.find_workspace_root() == second_workspace +def test_read_workflow_from_fs_prefers_user_workflow_over_project_bundle( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +): + workflow_id = "shared-workflow" + project_root = tmp_path / "project-workflows" + user_root = tmp_path / "user-workflows" + + for root, name in ( + (project_root, "project bundle"), + (user_root, "user customization"), + ): + workflow_dir = root / workflow_id + workflow_dir.mkdir(parents=True) + (workflow_dir / "workflow.json").write_text( + json.dumps( + { + "name": name, + "start": "n1", + "nodes": [{"id": "n1", "type": "python", "code": "pass"}], + "edges": [], + } + ), + encoding="utf-8", + ) + + monkeypatch.setattr( + fs_store, + "resolve_workflow_scan_roots", + lambda _workspace: [(project_root, "project"), (user_root, "global")], + ) + + workflow = fs_store.read_workflow_from_fs(workflow_id) + + assert workflow is not None + assert workflow["source"] == "global" + assert workflow["workflowJson"]["name"] == "user customization" + + def test_read_workflow_dir_uses_latest_file_mtime_when_meta_is_stale( tmp_path: Path, ): diff --git a/tests/workflow/test_poller_manager.py b/tests/workflow/test_poller_manager.py index 4b7f915df..37efd76e9 100644 --- a/tests/workflow/test_poller_manager.py +++ b/tests/workflow/test_poller_manager.py @@ -4,14 +4,76 @@ import threading from types import SimpleNamespace from typing import Any +from unittest.mock import AsyncMock import pytest +from flocks.tool import ToolContext from flocks.workflow import poller_manager from flocks.workflow import execution_store from flocks.workflow.runner import RunWorkflowResult +@pytest.fixture(autouse=True) +def trigger_tool_context(monkeypatch: pytest.MonkeyPatch) -> SimpleNamespace: + context = ToolContext( + session_id="schedule-parent", + message_id="schedule-message", + agent="rex", + ) + builder = AsyncMock(return_value=context) + cleanup = AsyncMock() + monkeypatch.setattr(poller_manager, "build_workflow_tool_context", builder) + monkeypatch.setattr(poller_manager, "cleanup_workflow_tool_context", cleanup) + return SimpleNamespace(context=context, builder=builder, cleanup=cleanup) + + +@pytest.mark.asyncio +async def test_start_all_skips_stale_workflow_config_as_info( + monkeypatch: pytest.MonkeyPatch, +) -> None: + workflow_id = "removed-workflow" + config = {"enabled": True, "intervalSeconds": "not-a-number"} + info_events: list[tuple[str, dict]] = [] + warning_events: list[str] = [] + + async def _list_configs(*, kind: str): # noqa: ANN202 + assert kind == "workflow_poller_config" + return [(workflow_id, config)] + + async def _get_config(_workflow_id: str, *, kind: str) -> dict[str, Any]: + assert kind == "workflow_poller_config" + return config + + monkeypatch.setattr(poller_manager.WorkflowStore, "list_configs", _list_configs) + monkeypatch.setattr(poller_manager.WorkflowStore, "get_config", _get_config) + monkeypatch.setattr(poller_manager, "read_workflow_from_fs", lambda _workflow_id: None) + monkeypatch.setattr( + poller_manager.log, + "info", + lambda event, data=None: info_events.append((event, data)), + ) + monkeypatch.setattr( + poller_manager.log, + "warning", + lambda event, _data=None: warning_events.append(event), + ) + + manager = poller_manager.WorkflowPollerManager() + await manager.start_all() + + status = manager.get_status(workflow_id) + assert status["state"] == "stopped" + assert status["error"] == "workflow_not_found" + assert info_events == [ + ( + "poller.workflow_not_found_on_start", + {"workflow_id": workflow_id, "action": "stale_config_skipped"}, + ) + ] + assert warning_events == [] + + @pytest.mark.asyncio async def test_restart_disabled_config_reports_stopped(monkeypatch: pytest.MonkeyPatch) -> None: manager = poller_manager.WorkflowPollerManager() @@ -42,7 +104,10 @@ async def _fake_get_config(_workflow_id: str, *, kind: str) -> dict[str, Any]: @pytest.mark.asyncio -async def test_run_once_injects_dynamic_inputs_and_summary(monkeypatch: pytest.MonkeyPatch) -> None: +async def test_run_once_injects_dynamic_inputs_and_summary( + monkeypatch: pytest.MonkeyPatch, + trigger_tool_context: SimpleNamespace, +) -> None: manager = poller_manager.WorkflowPollerManager() captured_inputs: dict[str, Any] = {} @@ -63,6 +128,7 @@ def _fake_run_workflow( # noqa: ANN001 on_step_complete, run_id: str, execution_profile: str, + tool_context: ToolContext, ): captured_inputs.update(inputs) assert workflow["start"] == "n1" @@ -71,6 +137,7 @@ def _fake_run_workflow( # noqa: ANN001 assert trace is False assert run_id == "exec-wf-run-once" assert execution_profile == "high_frequency" + assert tool_context is trigger_tool_context.context assert cancel() is False return RunWorkflowResult( status="success", @@ -130,6 +197,11 @@ def _fake_run_workflow( # noqa: ANN001 assert captured_inputs["input_date"] assert captured_inputs["_trigger"] == "poller" assert captured_inputs["_poller_run_id"].startswith("poller-") + trigger_tool_context.builder.assert_awaited_once_with( + workflow_id="wf-run-once", + action_name="trigger:schedule", + ) + trigger_tool_context.cleanup.assert_awaited_once_with(trigger_tool_context.context) @pytest.mark.asyncio @@ -193,6 +265,7 @@ def _fake_run_workflow( # noqa: ANN001 on_step_complete, run_id: str, execution_profile: str, + tool_context: ToolContext, ): assert workflow["start"] == "n1" assert workflow["nodes"][0]["id"] == "n1" @@ -284,6 +357,7 @@ def _fake_run_workflow( # noqa: ANN001 on_step_complete, run_id: str, execution_profile: str, + tool_context: ToolContext, ): _ = workflow, inputs, timeout_s, trace, cancel, run_id _ = on_step_complete @@ -380,6 +454,7 @@ def _fake_run_workflow( # noqa: ANN001 on_step_complete, run_id: str, execution_profile: str, + tool_context: ToolContext, ): _ = workflow, inputs, timeout_s, trace, cancel, run_id _ = on_step_complete @@ -411,22 +486,42 @@ def _fake_run_workflow( # noqa: ANN001 async def test_start_all_only_restarts_enabled_configs(monkeypatch: pytest.MonkeyPatch) -> None: manager = poller_manager.WorkflowPollerManager() restarted: list[str] = [] + warning_events: list[tuple[str, dict[str, str]]] = [] async def _fake_list_configs(*, kind: str) -> list[tuple[str, dict[str, Any]]]: return [ + ("wf-broken", {"enabled": True}), ("wf-enabled", {"enabled": True}), ("wf-disabled", {"enabled": False}), ] - async def _fake_restart(workflow_id: str) -> dict[str, Any]: + async def _fake_restart( + workflow_id: str, + *, + startup: bool = False, + ) -> dict[str, Any]: + assert startup is True restarted.append(workflow_id) + if workflow_id == "wf-broken": + raise ValueError("invalid config") return {"workflowId": workflow_id, "state": "running"} monkeypatch.setattr(poller_manager.WorkflowStore, "list_configs", _fake_list_configs) monkeypatch.setattr(manager, "restart_workflow", _fake_restart) + monkeypatch.setattr( + poller_manager.log, + "warning", + lambda event, data: warning_events.append((event, data)), + ) await manager.start_all() - assert restarted == ["wf-enabled"] + assert restarted == ["wf-broken", "wf-enabled"] + assert warning_events == [ + ( + "poller.start_failed", + {"workflow_id": "wf-broken", "error": "invalid config"}, + ) + ] @pytest.mark.asyncio diff --git a/tests/workflow/test_stream_alert_triage_persistence.py b/tests/workflow/test_stream_alert_triage_persistence.py index 3c8c05720..722f844e4 100644 --- a/tests/workflow/test_stream_alert_triage_persistence.py +++ b/tests/workflow/test_stream_alert_triage_persistence.py @@ -6,6 +6,7 @@ import os import re import sqlite3 +import threading import time from concurrent.futures import ThreadPoolExecutor from pathlib import Path @@ -27,6 +28,11 @@ def _concurrent_triage_code() -> str: return next(node["code"] for node in workflow["nodes"] if node["id"] == "concurrent_triage") +def _load_dedup_code() -> str: + workflow = json.loads(WORKFLOW_PATH.read_text(encoding="utf-8")) + return next(node["code"] for node in workflow["nodes"] if node["id"] == "load_dedup_file") + + def _load_functions(*names: str) -> dict[str, object]: tree = ast.parse(_concurrent_triage_code()) body = [ @@ -46,6 +52,62 @@ def _load_functions(*names: str) -> dict[str, object]: return namespace +def test_load_node_defaults_to_safe_single_alert_concurrency(tmp_path: Path) -> None: + input_path = tmp_path / "dedup_result_001.jsonl" + input_path.write_text( + json.dumps({"dedup_key": "first", "is_duplicate": False}) + "\n", + encoding="utf-8", + ) + namespace: dict[str, object] = { + "inputs": {"input_path": str(input_path)}, + "outputs": {}, + } + + exec(compile(_load_dedup_code(), str(WORKFLOW_PATH), "exec"), namespace) + + assert namespace["outputs"]["concurrency"] == 1 + + +def test_llm_calls_share_the_run_concurrency_budget() -> None: + functions = _load_functions("_ask_llm") + active = 0 + peak = 0 + lock = threading.Lock() + + class FakeLLM: + def ask(self, _prompt: str, **_kwargs: object) -> str: + nonlocal active, peak + with lock: + active += 1 + peak = max(peak, active) + try: + time.sleep(0.02) + return "ok" + finally: + with lock: + active -= 1 + + functions.update( + { + "llm": FakeLLM(), + "LLM_CALL_TIMEOUT_S": 120.0, + "LLM_CALL_MAX_RETRIES": 1, + "_llm_slots": threading.BoundedSemaphore(5), + } + ) + ask_llm = functions["_ask_llm"] + + with ThreadPoolExecutor(max_workers=20) as pool: + results = list(pool.map(ask_llm, [f"prompt-{i}" for i in range(40)])) + + assert results == ["ok"] * 40 + assert peak == 5 + assert re.search( + r"_llm_slots\s*=\s*threading\.BoundedSemaphore\(concurrency\)", + _concurrent_triage_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"] diff --git a/tests/workflow/test_trigger_runtime.py b/tests/workflow/test_trigger_runtime.py index 023ce4f32..dabf9bb7e 100644 --- a/tests/workflow/test_trigger_runtime.py +++ b/tests/workflow/test_trigger_runtime.py @@ -2,10 +2,69 @@ import asyncio from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock import pytest +from flocks.tool import ToolContext from flocks.workflow.triggers import runtime as runtime_module +from flocks.workflow.triggers.models import TriggerDefinition + + +@pytest.mark.asyncio +async def test_trigger_execution_builds_tool_context_for_workflow_tools( + monkeypatch: pytest.MonkeyPatch, +) -> None: + tool_context = ToolContext( + session_id="trigger-parent", + message_id="trigger-message", + agent="rex", + ) + build_context = AsyncMock(return_value=tool_context) + cleanup_context = AsyncMock() + + def _fake_run_workflow(**kwargs): # noqa: ANN003 + missing_context = kwargs.get("tool_context") is None + return SimpleNamespace( + status="FAILED" if missing_context else "SUCCEEDED", + outputs={}, + error="Parent session not found" if missing_context else None, + history=[], + last_node_id="notify", + steps=1, + ) + + monkeypatch.setattr( + runtime_module, + "build_workflow_tool_context", + build_context, + ) + monkeypatch.setattr(runtime_module, "cleanup_workflow_tool_context", cleanup_context) + monkeypatch.setattr(runtime_module, "run_workflow", Mock(side_effect=_fake_run_workflow)) + monkeypatch.setattr( + runtime_module, + "create_execution_record", + AsyncMock(return_value={"id": "exec-1"}), + ) + monkeypatch.setattr(runtime_module, "record_execution_result", AsyncMock()) + + trigger = TriggerDefinition.model_validate({"id": "webhook-trigger", "type": "custom_webhook"}) + runtime = runtime_module.TriggerRuntime() + + result = await runtime._execute_workflow( # noqa: SLF001 + workflow_id="wf-trigger", + workflow_json={"start": "notify", "nodes": [], "edges": []}, + trigger=trigger, + mapped_inputs={"message": "hello"}, + ) + + assert result["status"] == "success" + build_context.assert_awaited_once_with( + workflow_id="wf-trigger", + action_name="trigger:custom_webhook", + ) + assert runtime_module.run_workflow.call_args.kwargs["tool_context"] is tool_context + cleanup_context.assert_awaited_once_with(tool_context) @pytest.mark.asyncio diff --git a/tests/workflow/test_workflow_center.py b/tests/workflow/test_workflow_center.py index 7aa12cc8e..6031882d9 100644 --- a/tests/workflow/test_workflow_center.py +++ b/tests/workflow/test_workflow_center.py @@ -163,12 +163,14 @@ async def test_publish_invoke_stop_workflow_service( docker_calls = [] json_post_calls = [] + port_requests = [] async def fake_exec_docker(args, allow_failure=False, **_kwargs): docker_calls.append((args, allow_failure)) return ("container-abc\n", "", 0) - async def fake_allocate_port() -> int: + async def fake_allocate_port(*args, **kwargs) -> int: + port_requests.append((args, kwargs)) return 19123 async def fake_wait_docker_service_healthy(*_args, **_kwargs) -> bool: @@ -187,6 +189,8 @@ def fake_json_post(*args, **kwargs): assert published["status"] == "active" assert published["hostPort"] == 19123 assert len(published["apiKey"]) == 64 + registry = await center.WorkflowStore.kv_get(center._registry_key(workflow_id)) + assert registry["servicePort"] == 19123 invoked = await center.invoke_published_workflow(workflow_id, inputs={"k": "v"}) assert invoked["status"] == "SUCCEEDED" @@ -196,6 +200,13 @@ def fake_json_post(*args, **kwargs): stopped = await center.stop_workflow_service(workflow_id) assert stopped["status"] == "stopped" assert stopped["stopped"] is True + registry = await center.WorkflowStore.kv_get(center._registry_key(workflow_id)) + assert registry["servicePort"] == 19123 + + republished = await center.publish_workflow(workflow_id) + assert republished["hostPort"] == 19123 + assert port_requests[-1][0][0] == 19123 + await center.stop_workflow_service(workflow_id) assert any(call[0][:3] == ["run", "-d", "--name"] for call in docker_calls) run_call = next(call for call in docker_calls if call[0][:3] == ["run", "-d", "--name"]) diff --git a/tests/workflow/test_workflow_center_lifecycle.py b/tests/workflow/test_workflow_center_lifecycle.py index 43d11ab83..9b90ee305 100644 --- a/tests/workflow/test_workflow_center_lifecycle.py +++ b/tests/workflow/test_workflow_center_lifecycle.py @@ -139,7 +139,7 @@ async def test_allocate_port_skips_reserved_service_records(monkeypatch) -> None }, "workflow_registry/wf-registry": { "workflowId": "wf-registry", - "serviceUrl": "http://127.0.0.1:19002", + "servicePort": 19002, "publishStatus": "active", }, } @@ -179,6 +179,220 @@ async def fake_list_keys(_prefix): center._IN_FLIGHT_PORT_RESERVATIONS.clear() +@pytest.mark.asyncio +async def test_allocate_preferred_port_reuses_current_workflow_reservation(monkeypatch) -> None: + store: dict[str, Any] = { + "workflow_api_service/wf-current": { + "workflowId": "wf-current", + "port": 25000, + "status": "stopped", + } + } + + async def fake_list_keys(prefix): + return [key for key in store if key.startswith(str(prefix))] + + async def fake_read(key): + return store.get(str(key)) + + center._IN_FLIGHT_PORT_RESERVATIONS.clear() + monkeypatch.setenv("FLOCKS_WORKFLOW_SERVICE_PORT_START", "19000") + monkeypatch.setenv("FLOCKS_WORKFLOW_SERVICE_PORT_END", "19001") + monkeypatch.setattr(center.WorkflowStore, "kv_list_keys", fake_list_keys) + monkeypatch.setattr(center.WorkflowStore, "kv_get", fake_read) + monkeypatch.setattr(center, "_is_port_available", lambda _port: True) + + try: + assert await center._allocate_port(25000, workflow_id="wf-current") == 25000 + finally: + center._IN_FLIGHT_PORT_RESERVATIONS.clear() + + +@pytest.mark.asyncio +async def test_allocate_preferred_port_rejects_other_workflow_reservation(monkeypatch) -> None: + store: dict[str, Any] = { + "workflow_api_service/wf-other": { + "workflowId": "wf-other", + "port": 25000, + "status": "stopped", + } + } + + async def fake_list_keys(prefix): + return [key for key in store if key.startswith(str(prefix))] + + async def fake_read(key): + return store.get(str(key)) + + center._IN_FLIGHT_PORT_RESERVATIONS.clear() + monkeypatch.setattr(center.WorkflowStore, "kv_list_keys", fake_list_keys) + monkeypatch.setattr(center.WorkflowStore, "kv_get", fake_read) + monkeypatch.setattr(center, "_is_port_available", lambda _port: True) + + with pytest.raises(center.WorkflowPortUnavailableError, match="25000"): + await center._allocate_port(25000, workflow_id="wf-current") + + +@pytest.mark.asyncio +async def test_allocate_preferred_port_allows_port_bound_by_current_runtime(monkeypatch) -> None: + async def fake_list_keys(_prefix): + return [] + + center._IN_FLIGHT_PORT_RESERVATIONS.clear() + monkeypatch.setattr(center.WorkflowStore, "kv_list_keys", fake_list_keys) + monkeypatch.setattr(center, "_is_port_available", lambda _port: False) + + try: + assert ( + await center._allocate_port( + 25000, + workflow_id="wf-current", + allow_bound_port=25000, + ) + == 25000 + ) + finally: + center._IN_FLIGHT_PORT_RESERVATIONS.clear() + + +@pytest.mark.asyncio +async def test_local_port_conflict_keeps_previous_runtime_and_registry(monkeypatch, tmp_path) -> None: + workflow_id = "wf-local-conflict" + workflow_path = tmp_path / "workflow.json" + workflow_path.write_text( + json.dumps( + { + "id": workflow_id, + "start": "n1", + "nodes": [{"id": "n1", "type": "python", "code": "outputs['ok'] = True"}], + "edges": [], + } + ), + encoding="utf-8", + ) + registry = { + "workflowId": workflow_id, + "workflowPath": str(workflow_path), + "publishStatus": "active", + "servicePort": 25000, + "serviceUrl": "http://127.0.0.1:25000", + } + runtime = { + "workflowId": workflow_id, + "releaseId": "rel-old", + "driver": "local", + "hostPort": 25000, + "serviceUrl": "http://127.0.0.1:25000", + } + store: dict[str, Any] = { + f"workflow_registry/{workflow_id}": registry, + f"workflow_runtime/{workflow_id}": runtime, + } + stop_calls: list[str] = [] + + async def fake_read(key): + return store.get(str(key)) + + async def fake_write(key, value): + store[str(key)] = value + + async def reject_port(*_args, **_kwargs): + raise center.WorkflowPortUnavailableError("port conflict") + + async def fake_stop(_workflow_id): + stop_calls.append(_workflow_id) + return True + + monkeypatch.setattr(center.WorkflowStore, "kv_get", fake_read) + monkeypatch.setattr(center.WorkflowStore, "kv_put", fake_write) + monkeypatch.setattr(center, "_allocate_port", reject_port) + monkeypatch.setattr(center, "_stop_existing_runtime_for_publish", fake_stop) + + with pytest.raises(center.WorkflowPortUnavailableError, match="port conflict"): + await center.publish_workflow_local(workflow_id, port=25000) + + assert stop_calls == [] + assert store[f"workflow_runtime/{workflow_id}"] is runtime + assert store[f"workflow_registry/{workflow_id}"] == registry + + +@pytest.mark.asyncio +async def test_docker_port_conflict_keeps_previous_runtime_and_registry(monkeypatch, tmp_path) -> None: + workflow_id = "wf-docker-conflict" + workflow_path = tmp_path / "workflow.json" + workflow_path.write_text( + json.dumps( + { + "id": workflow_id, + "start": "n1", + "nodes": [{"id": "n1", "type": "python", "code": "outputs['ok'] = True"}], + "edges": [], + } + ), + encoding="utf-8", + ) + release_dir = tmp_path / "release" + release_dir.mkdir() + release_snapshot = release_dir / "workflow.json" + registry = { + "workflowId": workflow_id, + "workflowPath": str(workflow_path), + "publishStatus": "active", + "servicePort": 25000, + "serviceUrl": "http://127.0.0.1:25000", + } + runtime = { + "workflowId": workflow_id, + "releaseId": "rel-old", + "driver": "docker", + "containerName": "old-container", + "hostPort": 25000, + "serviceUrl": "http://127.0.0.1:25000", + } + store: dict[str, Any] = { + f"workflow_registry/{workflow_id}": registry, + f"workflow_runtime/{workflow_id}": runtime, + } + docker_calls: list[Any] = [] + + async def fake_read(key): + return store.get(str(key)) + + async def fake_write(key, value): + store[str(key)] = value + + async def fake_snapshot(*_args, **_kwargs): + return release_snapshot + + async def no_requirements(*_args, **_kwargs): + return False + + async def reject_port(*_args, **_kwargs): + raise center.WorkflowPortUnavailableError("port conflict") + + async def fake_exec_docker(*args, **kwargs): + docker_calls.append((args, kwargs)) + return ("", "", 0) + + monkeypatch.setattr(center.WorkflowStore, "kv_get", fake_read) + monkeypatch.setattr(center.WorkflowStore, "kv_put", fake_write) + monkeypatch.setattr(center, "_write_release_snapshot", fake_snapshot) + monkeypatch.setattr(center, "_write_requirements_snapshot", no_requirements) + monkeypatch.setattr(center, "_service_cache_dir", lambda _name: tmp_path) + monkeypatch.setattr(center.Config, "get_config_path", lambda: tmp_path / "missing-config") + monkeypatch.setattr(center, "resolve_python_package_index_url", lambda: None) + monkeypatch.setattr(center, "_allocate_port", reject_port) + monkeypatch.setattr(center, "exec_docker", fake_exec_docker) + + with pytest.raises(center.WorkflowPortUnavailableError, match="port conflict"): + await center._publish_workflow_docker(workflow_id, port=26000) + + assert docker_calls == [] + assert store[f"workflow_runtime/{workflow_id}"] is runtime + assert store[f"workflow_registry/{workflow_id}"] == registry + assert not any(key.startswith(f"workflow_release/{workflow_id}/") for key in store) + + @pytest.mark.asyncio async def test_publish_workflow_local_releases_reserved_port_on_spawn_failure( monkeypatch, @@ -212,12 +426,12 @@ async def fake_write(key, value): store[str(key)] = value async def fake_stop_existing_runtime_for_publish(_workflow_id): - return None + return False async def fake_write_release_snapshot(_workflow_id, _release_id, _workflow_json): return workflow_path - async def fake_allocate_port(): + async def fake_allocate_port(*_args, **_kwargs): center._IN_FLIGHT_PORT_RESERVATIONS[19000] = 9999999999.0 return 19000 @@ -230,6 +444,7 @@ async def fake_create_subprocess_exec(*_args, **_kwargs): monkeypatch.setattr(center, "_stop_existing_runtime_for_publish", fake_stop_existing_runtime_for_publish) monkeypatch.setattr(center, "_write_release_snapshot", fake_write_release_snapshot) monkeypatch.setattr(center, "_allocate_port", fake_allocate_port) + monkeypatch.setattr(center, "_is_port_available", lambda _port: True) monkeypatch.setattr(center.asyncio, "create_subprocess_exec", fake_create_subprocess_exec) try: diff --git a/tests/workflow/test_workflow_paths.py b/tests/workflow/test_workflow_paths.py index ff31e0237..2194b5bfb 100644 --- a/tests/workflow/test_workflow_paths.py +++ b/tests/workflow/test_workflow_paths.py @@ -162,5 +162,99 @@ async def test_new_canonical_path_wins_over_legacy( results = await center.scan_skill_workflows(tmp_path) # The new canonical path (plugins/workflows/) has higher priority and wins - names = [r["name"] for r in results] - assert "shared-wf-new" in names + assert [r["name"] for r in results] == ["shared-wf-new"] + + +@pytest.mark.asyncio +async def test_scan_and_registry_prefer_user_workflow_over_project_bundle( + tmp_path: Path, + isolated_storage, + monkeypatch: pytest.MonkeyPatch, +) -> None: + project_root = tmp_path / "project" + user_root = tmp_path / "user" + for root, name in ( + (project_root, "project bundle"), + (user_root, "user customization"), + ): + workflow_dir = root / "shared-workflow" + workflow_dir.mkdir(parents=True) + (workflow_dir / "workflow.json").write_text( + json.dumps(_workflow_payload(name)), + encoding="utf-8", + ) + + monkeypatch.setattr( + center, + "resolve_project_workflow_roots", + lambda _base: [project_root], + ) + monkeypatch.setattr(center, "resolve_global_workflow_roots", lambda: [user_root]) + + scanned = await center.scan_skill_workflows(tmp_path) + registered = await center.list_registry_entries() + + assert len(scanned) == 1 + assert scanned[0]["name"] == "user customization" + assert scanned[0]["sourceType"] == "global" + matching_registry_entries = [ + entry + for entry in registered + if entry.get("logicalWorkflowId") == "shared-workflow" + and Path(str(entry.get("workflowPath"))).is_relative_to(tmp_path) + ] + assert len(matching_registry_entries) == 1 + assert matching_registry_entries[0]["name"] == "user customization" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("invalid_user_workflow", ["invalid_json", "hidden"]) +async def test_registry_falls_back_when_user_workflow_is_no_longer_discoverable( + tmp_path: Path, + isolated_storage, + monkeypatch: pytest.MonkeyPatch, + invalid_user_workflow: str, +) -> None: + project_root = tmp_path / "project" + user_root = tmp_path / "user" + project_dir = project_root / "shared-workflow" + user_dir = user_root / "shared-workflow" + project_dir.mkdir(parents=True) + user_dir.mkdir(parents=True) + (project_dir / "workflow.json").write_text( + json.dumps(_workflow_payload("project bundle")), + encoding="utf-8", + ) + user_workflow_path = user_dir / "workflow.json" + user_workflow_path.write_text( + json.dumps(_workflow_payload("user customization")), + encoding="utf-8", + ) + + monkeypatch.setattr( + center, + "resolve_project_workflow_roots", + lambda _base: [project_root], + ) + monkeypatch.setattr(center, "resolve_global_workflow_roots", lambda: [user_root]) + + await center.scan_skill_workflows(tmp_path) + if invalid_user_workflow == "invalid_json": + user_workflow_path.write_text("{", encoding="utf-8") + else: + (user_dir / "meta.json").write_text( + json.dumps({"hidden": True}), + encoding="utf-8", + ) + + scanned = await center.scan_skill_workflows(tmp_path) + registered = await center.list_registry_entries() + matching_registry_entries = [ + entry + for entry in registered + if entry.get("logicalWorkflowId") == "shared-workflow" + and Path(str(entry.get("workflowPath"))).is_relative_to(tmp_path) + ] + + assert [entry["name"] for entry in scanned] == ["project bundle"] + assert [entry["name"] for entry in matching_registry_entries] == ["project bundle"] diff --git a/tests/workflow/test_workflow_tool_context.py b/tests/workflow/test_workflow_tool_context.py index 9e0ee8fa4..cd5a00c01 100644 --- a/tests/workflow/test_workflow_tool_context.py +++ b/tests/workflow/test_workflow_tool_context.py @@ -7,7 +7,10 @@ from flocks.session.message import Message, MessageRole from flocks.session.session import Session from flocks.storage.storage import Storage -from flocks.workflow.tool_context import build_workflow_tool_context +from flocks.workflow.tool_context import ( + build_workflow_tool_context, + cleanup_workflow_tool_context, +) @pytest.fixture @@ -86,3 +89,58 @@ async def test_build_workflow_tool_context_reuses_existing_parent_session( assert message is not None assert message.role == MessageRole.USER assert message.agent == "rex-junior" + + assert await cleanup_workflow_tool_context(tool_context) is False + assert await Session.get_by_id(parent.id) is not None + + +@pytest.mark.asyncio +async def test_cleanup_workflow_tool_context_purges_unused_temp_parent( + tmp_path: Path, + isolated_storage, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.chdir(tmp_path) + tool_context = await build_workflow_tool_context( + workflow_id="wf-cleanup", + action_name="trigger:kafka", + ) + session_id = tool_context.session_id + message_id = tool_context.message_id + + cleaned = await cleanup_workflow_tool_context(tool_context) + + assert cleaned is True + assert await Session.get_by_id(session_id) is None + assert await Message.get(session_id, message_id) is None + assert await Storage.get(f"session_callable_tools:{session_id}") is None + + +@pytest.mark.asyncio +async def test_cleanup_workflow_tool_context_preserves_parent_with_child_session( + tmp_path: Path, + isolated_storage, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.chdir(tmp_path) + tool_context = await build_workflow_tool_context( + workflow_id="wf-delegated", + action_name="trigger:schedule", + ) + parent = await Session.get_by_id(tool_context.session_id) + assert parent is not None + child = await Session.create( + project_id=parent.project_id, + directory=parent.directory, + title="Delegated child", + parent_id=parent.id, + agent="rex-junior", + category="task", + ) + tool_context.extra["workflow_child_session_created"] = True + + cleaned = await cleanup_workflow_tool_context(tool_context) + + assert cleaned is False + assert await Session.get_by_id(parent.id) is not None + assert await Session.get_by_id(child.id) is not None diff --git a/tui/flocks/cli/cmd/tui/context/local.tsx b/tui/flocks/cli/cmd/tui/context/local.tsx index b0d569018..0cd77695b 100644 --- a/tui/flocks/cli/cmd/tui/context/local.tsx +++ b/tui/flocks/cli/cmd/tui/context/local.tsx @@ -12,6 +12,7 @@ import { Provider } from "@/provider/provider" import { useArgs } from "./args" import { useSDK } from "./sdk" import { RGBA } from "@opentui/core" +import { selectFallbackModel } from "./model-selection" export const { use: useLocal, provider: LocalProvider } = createSimpleContext({ name: "Local", @@ -174,16 +175,7 @@ export const { use: useLocal, provider: LocalProvider } = createSimpleContext({ } } - const provider = sync.data.provider[0] - if (!provider) return undefined - const defaultModel = sync.data.provider_default[provider.id] - const firstModel = Object.values(provider.models)[0] - const model = defaultModel ?? firstModel?.id - if (!model) return undefined - return { - providerID: provider.id, - modelID: model, - } + return selectFallbackModel(sync.data.provider, sync.data.provider_default) }) const currentModel = createMemo(() => { diff --git a/tui/flocks/cli/cmd/tui/context/model-selection.test.ts b/tui/flocks/cli/cmd/tui/context/model-selection.test.ts new file mode 100644 index 000000000..3830009e2 --- /dev/null +++ b/tui/flocks/cli/cmd/tui/context/model-selection.test.ts @@ -0,0 +1,61 @@ +import { describe, expect, test } from "bun:test" +import { selectFallbackModel } from "./model-selection" + +describe("selectFallbackModel", () => { + test("skips an empty provider when a later provider has a model", () => { + expect( + selectFallbackModel( + [ + { id: "openai", models: {} }, + { + id: "threatbook-cn-llm", + models: { + "minimax-m3": { id: "minimax-m3" }, + }, + }, + ], + { "threatbook-cn-llm": "minimax-m3" }, + ), + ).toEqual({ + providerID: "threatbook-cn-llm", + modelID: "minimax-m3", + }) + }) + + test("uses the provider default when it is available", () => { + expect( + selectFallbackModel( + [ + { + id: "openai", + models: { + "gpt-4.1": { id: "gpt-4.1" }, + "gpt-5": { id: "gpt-5" }, + }, + }, + ], + { openai: "gpt-5" }, + ), + ).toEqual({ providerID: "openai", modelID: "gpt-5" }) + }) + + test("uses the first model when the provider default is unavailable", () => { + expect( + selectFallbackModel( + [ + { + id: "openai", + models: { + "gpt-4.1": { id: "gpt-4.1" }, + }, + }, + ], + { openai: "removed-model" }, + ), + ).toEqual({ providerID: "openai", modelID: "gpt-4.1" }) + }) + + test("returns undefined when no provider has a model", () => { + expect(selectFallbackModel([{ id: "openai", models: {} }], {})).toBeUndefined() + }) +}) diff --git a/tui/flocks/cli/cmd/tui/context/model-selection.ts b/tui/flocks/cli/cmd/tui/context/model-selection.ts new file mode 100644 index 000000000..1c8ac86cc --- /dev/null +++ b/tui/flocks/cli/cmd/tui/context/model-selection.ts @@ -0,0 +1,21 @@ +type Model = { + id: string +} + +type Provider = { + id: string + models: Record +} + +export function selectFallbackModel(providers: Provider[], defaults: Record) { + const provider = providers.find((item) => Object.keys(item.models).length > 0) + if (!provider) return undefined + const defaultModel = defaults[provider.id] + const firstModel = Object.values(provider.models)[0] + const modelID = defaultModel && provider.models[defaultModel] ? defaultModel : firstModel?.id + if (!modelID) return undefined + return { + providerID: provider.id, + modelID, + } +} diff --git a/uv.lock b/uv.lock index a09013247..a5584aa89 100644 --- a/uv.lock +++ b/uv.lock @@ -553,7 +553,7 @@ wheels = [ [[package]] name = "flocks" -version = "2026.7.15" +version = "2026.7.22" source = { editable = "." } dependencies = [ { name = "aiofiles" }, @@ -598,6 +598,8 @@ dependencies = [ { name = "pyyaml" }, { name = "requests" }, { name = "rich" }, + { name = "slack-bolt" }, + { name = "slack-sdk" }, { name = "sqlalchemy" }, { name = "sse-starlette" }, { name = "starlette" }, @@ -667,6 +669,8 @@ requires-dist = [ { name = "pyyaml", specifier = ">=6.0.1" }, { name = "requests", specifier = ">=2.32.5" }, { name = "rich", specifier = ">=13.7.0" }, + { name = "slack-bolt", specifier = ">=1.26.0" }, + { name = "slack-sdk", specifier = ">=3.38.0" }, { name = "sqlalchemy", specifier = ">=2.0.25" }, { name = "sse-starlette", specifier = ">=1.8.2" }, { name = "starlette", specifier = ">=1.0.1" }, @@ -2137,6 +2141,27 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050, upload-time = "2024-12-04T17:35:26.475Z" }, ] +[[package]] +name = "slack-bolt" +version = "1.29.0" +source = { registry = "https://mirrors.aliyun.com/pypi/simple" } +dependencies = [ + { name = "slack-sdk" }, +] +sdist = { url = "https://mirrors.aliyun.com/pypi/packages/3c/2c/7434f5d8c52eafb1e702012e9f291b013bd05985a5b32bde568b2279a28a/slack_bolt-1.29.0.tar.gz", hash = "sha256:b6271ba0a9b71e319c86b40632e6cb6240aacd0433773615b76b890b9a574762" } +wheels = [ + { url = "https://mirrors.aliyun.com/pypi/packages/2f/cc/e5f78a80a1775a4cbb2cc41e8ef433602d5e755ff0d829bd43145015740a/slack_bolt-1.29.0-py2.py3-none-any.whl", hash = "sha256:1835b66b778158f3af0da77603aa18d7dfd82fd9b9a985e25c752f95050ab826" }, +] + +[[package]] +name = "slack-sdk" +version = "3.43.0" +source = { registry = "https://mirrors.aliyun.com/pypi/simple" } +sdist = { url = "https://mirrors.aliyun.com/pypi/packages/13/75/a4964eb771a0c74d79ee7a3bee6fb5d9718909dd1b675e80d62a6a0ad90a/slack_sdk-3.43.0.tar.gz", hash = "sha256:0553152e46c4259eb69f7464cdadc35ba4802ca10f9f5a849c92cf03d6c2ba07" } +wheels = [ + { url = "https://mirrors.aliyun.com/pypi/packages/e4/55/42141b8338d46323d5b3c6095201b044c670c20f898643b322ea9b1543a1/slack_sdk-3.43.0-py2.py3-none-any.whl", hash = "sha256:4b6557c65577fc172f685af218b811f9f3b4909e24cddd839ada09565f10c585" }, +] + [[package]] name = "smmap" version = "5.0.3" diff --git a/webui/public/channel-dingtalk-transparent.png b/webui/public/channel-dingtalk-transparent.png new file mode 100644 index 000000000..7fe7032ee Binary files /dev/null and b/webui/public/channel-dingtalk-transparent.png differ diff --git a/webui/public/channel-slack.png b/webui/public/channel-slack.png new file mode 100644 index 000000000..1cc40f368 Binary files /dev/null and b/webui/public/channel-slack.png differ diff --git a/webui/public/channel-weixin-transparent.png b/webui/public/channel-weixin-transparent.png new file mode 100644 index 000000000..bf2d8399a Binary files /dev/null and b/webui/public/channel-weixin-transparent.png differ diff --git a/webui/src/api/notifications.ts b/webui/src/api/notifications.ts index 6c8927989..98ac94aa0 100644 --- a/webui/src/api/notifications.ts +++ b/webui/src/api/notifications.ts @@ -32,14 +32,10 @@ export interface NotificationAckStatus { acknowledged: boolean; } -export const getActiveNotifications = async ( - locale?: string, - currentVersion?: string | null, -): Promise => { +export const getActiveNotifications = async (locale?: string): Promise => { const response = await client.get('/api/notifications/active', { params: { ...(locale ? { locale } : {}), - ...(currentVersion ? { current_version: currentVersion } : {}), }, }); return response.data; diff --git a/webui/src/api/provider.test.ts b/webui/src/api/provider.test.ts index 1e2cda5ce..2a63e4f57 100644 --- a/webui/src/api/provider.test.ts +++ b/webui/src/api/provider.test.ts @@ -136,3 +136,19 @@ describe('providerAPI.listApiServices', () => { expect(response.data[0].name).toBe('Forced Service'); }); }); + +describe('providerAPI.revealCredentials', () => { + afterEach(() => { + vi.resetModules(); + mockPost.mockReset(); + }); + + it('uses the explicit credential reveal endpoint', async () => { + mockPost.mockResolvedValue({ data: { api_key: 'sk-revealed', has_credential: true } }); + + const { providerAPI } = await import('./provider'); + await providerAPI.revealCredentials('openai'); + + expect(mockPost).toHaveBeenCalledWith('/api/provider/openai/credentials/reveal'); + }); +}); diff --git a/webui/src/api/provider.ts b/webui/src/api/provider.ts index a35235f42..e0f3062ef 100644 --- a/webui/src/api/provider.ts +++ b/webui/src/api/provider.ts @@ -111,6 +111,9 @@ export const providerAPI = { getCredentials: (id: string) => client.get(`/api/provider/${id}/credentials`), + revealCredentials: (id: string) => + client.post(`/api/provider/${id}/credentials/reveal`), + setCredentials: (id: string, credentials: ProviderCredentialInput) => client.post<{ success: boolean; message: string }>(`/api/provider/${id}/credentials`, credentials), diff --git a/webui/src/api/session.ts b/webui/src/api/session.ts index eee2df3a0..a421cb219 100644 --- a/webui/src/api/session.ts +++ b/webui/src/api/session.ts @@ -77,6 +77,7 @@ export interface SessionListParams { limit?: number; offset?: number; directory?: string; + projectID?: string; roots?: boolean; start?: number; search?: string; @@ -125,7 +126,7 @@ export const sessionApi = { /** * 创建会话 */ - create: async (data?: { title?: string; parentID?: string }) => { + create: async (data?: { title?: string; parentID?: string; projectID?: string }) => { const response = await client.post('/api/session', data || {}); return response.data; }, diff --git a/webui/src/api/stats.test.ts b/webui/src/api/stats.test.ts index 7407e93ba..0b49842b3 100644 --- a/webui/src/api/stats.test.ts +++ b/webui/src/api/stats.test.ts @@ -97,5 +97,40 @@ describe('statsApi.getSystemStats', () => { const { statsApi } = await import('./stats'); const result = await statsApi.getSystemStats(); expect(result.skills.total).toBe(0); + expect(result.system.status).toBe('warning'); + expect(result.system.message).toBe('partial'); + }); + + it('reports an authentication-specific error when all stats resources fail with 401', async () => { + mockGet.mockImplementation((url: string) => { + if (url === '/api/stats/summary') return Promise.reject(new Error('legacy backend')); + if (url === '/api/health') return Promise.resolve({ data: { status: 'healthy' } }); + return Promise.reject({ + response: { status: 401, data: { message: 'unauthorized' } }, + }); + }); + + const { statsApi } = await import('./stats'); + const result = await statsApi.getSystemStats(); + + expect(result.system.status).toBe('error'); + expect(result.system.message).toBe('authExpired'); + }); + + it('does not issue legacy fallback requests when the summary endpoint returns 401', async () => { + mockGet.mockImplementation((url: string) => { + if (url === '/api/stats/summary') { + return Promise.reject({ response: { status: 401 } }); + } + return Promise.resolve({ data: [] }); + }); + + const { statsApi } = await import('./stats'); + const result = await statsApi.getSystemStats(); + + expect(result.system.status).toBe('error'); + expect(result.system.message).toBe('authExpired'); + expect(mockGet).toHaveBeenCalledTimes(1); + expect(mockGet).toHaveBeenCalledWith('/api/stats/summary'); }); }); diff --git a/webui/src/api/stats.ts b/webui/src/api/stats.ts index 89ad2130f..b6a8197e7 100644 --- a/webui/src/api/stats.ts +++ b/webui/src/api/stats.ts @@ -26,6 +26,10 @@ export interface SystemStats { }; } +type StatsStatus = SystemStats['system']['status']; + +const LEGACY_RESOURCE_ENDPOINT_COUNT = 6; + function shouldCountForAgentPage(agent: any): boolean { if (!agent || typeof agent !== 'object') return false; if (agent.mode === 'primary') return true; @@ -46,15 +50,75 @@ function isSystemStats(value: any): value is SystemStats { ); } +type StatsEndpointFailure = { + endpoint: string; + error: unknown; +}; + +type StatsFailureCode = + | 'authExpired' + | 'forbidden' + | 'notFound' + | 'partial' + | 'network' + | 'unavailable'; + +function getErrorStatus(error: unknown): number | undefined { + return (error as { response?: { status?: number } } | undefined)?.response?.status; +} + +function classifyStatsFailure(failures: StatsEndpointFailure[]): StatsFailureCode { + if (failures.some((failure) => getErrorStatus(failure.error) === 401)) { + return 'authExpired'; + } + + if (failures.some((failure) => getErrorStatus(failure.error) === 403)) { + return 'forbidden'; + } + + if (failures.some((failure) => getErrorStatus(failure.error) === 404)) { + return 'notFound'; + } + + if (failures.length > 0) { + return 'network'; + } + + return 'unavailable'; +} + +function emptyStats(status: StatsStatus, message: StatsFailureCode): SystemStats { + return { + tasks: { week: 0, scheduledActive: 0 }, + agents: { total: 0 }, + workflows: { total: 0 }, + skills: { total: 0 }, + tools: { total: 0 }, + models: { total: 0 }, + system: { status, message }, + }; +} + async function getSystemStatsLegacy(): Promise { + const resourceFailures: StatsEndpointFailure[] = []; + const healthFailures: StatsEndpointFailure[] = []; + const getWithFallback = async (endpoint: string, fallbackData: unknown, failures: StatsEndpointFailure[] = resourceFailures) => { + try { + return await apiClient.get(endpoint); + } catch (error) { + failures.push({ endpoint, error }); + return { data: fallbackData }; + } + }; + const [taskDash, agents, workflows, skills, tools, providers, health] = await Promise.all([ - apiClient.get('/api/task-system/dashboard').catch(() => ({ data: {} })), - apiClient.get('/api/agent').catch(() => ({ data: [] })), - apiClient.get('/api/workflow').catch(() => ({ data: [] })), - apiClient.get('/api/skills').catch(() => ({ data: [] })), - apiClient.get('/api/tools').catch(() => ({ data: [] })), - apiClient.get('/api/provider').catch(() => ({ data: { all: [] } })), - apiClient.get('/api/health').catch(() => ({ data: { status: 'error' } })), + getWithFallback('/api/task-system/dashboard', {}), + getWithFallback('/api/agent', []), + getWithFallback('/api/workflow', []), + getWithFallback('/api/skills', []), + getWithFallback('/api/tools', []), + getWithFallback('/api/provider', { all: [] }), + getWithFallback('/api/health', { status: 'error' }, healthFailures), ]); const dash = taskDash.data || {}; @@ -72,6 +136,19 @@ async function getSystemStatsLegacy(): Promise { const totalModels = providerAll .filter((p: any) => connectedSet.has(p.id)) .reduce((sum: number, p: any) => sum + Object.keys(p.models ?? {}).length, 0); + const allResourceEndpointsFailed = resourceFailures.length === LEGACY_RESOURCE_ENDPOINT_COUNT; + const hasPartialResourceFailure = resourceFailures.length > 0 && !allResourceEndpointsFailed; + const healthIsHealthy = health.data.status === 'healthy'; + const systemStatus: StatsStatus = !healthIsHealthy || allResourceEndpointsFailed + ? 'error' + : hasPartialResourceFailure + ? 'warning' + : 'healthy'; + const systemMessage: string = systemStatus === 'healthy' + ? 'healthy' + : hasPartialResourceFailure + ? 'partial' + : classifyStatsFailure(allResourceEndpointsFailed ? resourceFailures : healthFailures); return { tasks: { @@ -84,8 +161,8 @@ async function getSystemStatsLegacy(): Promise { tools: { total: toolList.length }, models: { total: totalModels }, system: { - status: health.data.status === 'healthy' ? 'healthy' : 'error', - message: health.data.status === 'healthy' ? '所有服务运行正常' : '部分服务异常', + status: systemStatus, + message: systemMessage, }, }; } @@ -99,19 +176,16 @@ export const statsApi = { } return await getSystemStatsLegacy(); } catch (error) { + const summaryStatus = getErrorStatus(error); + if (summaryStatus === 401 || summaryStatus === 403) { + return emptyStats('error', classifyStatsFailure([{ endpoint: '/api/stats/summary', error }])); + } + try { return await getSystemStatsLegacy(); } catch (fallbackError) { console.error('Failed to fetch system stats:', fallbackError || error); - return { - tasks: { week: 0, scheduledActive: 0 }, - agents: { total: 0 }, - workflows: { total: 0 }, - skills: { total: 0 }, - tools: { total: 0 }, - models: { total: 0 }, - system: { status: 'error', message: '无法连接到后端服务' }, - }; + return emptyStats('error', classifyStatsFailure([{ endpoint: '/api/stats/summary', error: fallbackError || error }])); } } }, diff --git a/webui/src/api/toolFailureConfig.ts b/webui/src/api/toolFailureConfig.ts new file mode 100644 index 000000000..e45333cba --- /dev/null +++ b/webui/src/api/toolFailureConfig.ts @@ -0,0 +1,19 @@ +import client from './client'; + +export interface ToolFailureConfig { + disableOnRepeatedFailure: boolean; +} + +export const toolFailureConfigApi = { + async get(): Promise { + const response = await client.get('/api/config/tool-failure'); + return response.data; + }, + + async update(disableOnRepeatedFailure: boolean): Promise { + const response = await client.patch('/api/config/tool-failure', { + disableOnRepeatedFailure, + }); + return response.data; + }, +}; diff --git a/webui/src/api/workflow.ts b/webui/src/api/workflow.ts index b081435f9..dc6ed9230 100644 --- a/webui/src/api/workflow.ts +++ b/webui/src/api/workflow.ts @@ -171,6 +171,61 @@ export interface WorkflowJSON { metadata?: WorkflowMetadata; } +export type WorkflowCapabilityState = + | 'unconfigured' + | 'starting' + | 'running' + | 'stopped' + | 'error'; + +export interface WorkflowCapabilityStatus { + configured: boolean; + state: WorkflowCapabilityState; + count?: number; +} + +export interface WorkflowTriggerStatusSummary { + id: string; + type: WorkflowTriggerType; + name?: string; + state: WorkflowCapabilityState; + rawState?: string; +} + +export interface WorkflowTriggerCapabilityStatus extends WorkflowCapabilityStatus { + items?: WorkflowTriggerStatusSummary[]; +} + +export interface WorkflowIntegrationStatus { + api: WorkflowCapabilityStatus; + trigger: WorkflowTriggerCapabilityStatus; +} + +export interface WorkflowStats { + callCount: number; + successCount: number; + errorCount: number; + totalRuntime: number; + avgRuntime: number; + thumbsUp: number; + thumbsDown: number; +} + +export interface WorkflowSummary { + id: string; + name: string; + nameI18n?: Record; + description?: string; + category: string; + status: 'draft' | 'active' | 'archived'; + source?: 'project' | 'global'; + createdAt: number; + updatedAt: number; + nodeCount: number; + integrationStatus?: WorkflowIntegrationStatus | null; + stats: WorkflowStats; +} + export interface Workflow { id: string; name: string; @@ -185,15 +240,8 @@ export interface Workflow { createdBy?: string; createdAt: number; updatedAt: number; - stats: { - callCount: number; - successCount: number; - errorCount: number; - totalRuntime: number; - avgRuntime: number; - thumbsUp: number; - thumbsDown: number; - }; + integrationStatus?: WorkflowIntegrationStatus | null; + stats: WorkflowStats; } export interface WorkflowExecutionStep { @@ -262,10 +310,24 @@ export interface WorkflowService { publishedAt: number; containerName?: string; driver?: 'local' | 'docker'; + port?: number; } export type WorkflowServiceDriver = 'local' | 'docker'; +export function resolveWorkflowServicePort( + service?: Pick | null, +): number | undefined { + if (service?.port && service.port >= 1 && service.port <= 65535) return service.port; + if (!service?.serviceUrl) return undefined; + try { + const port = Number(new URL(service.serviceUrl).port); + return port >= 1 && port <= 65535 ? port : undefined; + } catch { + return undefined; + } +} + export interface WorkflowIntegrationConfig { version: number; kind: string; @@ -379,6 +441,9 @@ export interface WorkflowPollerStatus { export const workflowAPI = { list: (params?: { category?: string; status?: string; excludeId?: string }) => client.get('/api/workflow', { params }), + + listSummaries: (params?: { category?: string; status?: string; excludeId?: string }) => + client.get('/api/workflow-summaries', { params }), get: (id: string) => client.get(`/api/workflow/${id}`), @@ -444,7 +509,7 @@ export const workflowAPI = { export: (id: string) => client.get(`/api/workflow/${id}/export`), - publish: (id: string, data?: { driver?: WorkflowServiceDriver }) => + publish: (id: string, data?: { driver?: WorkflowServiceDriver; port?: number }) => client.post(`/api/workflow/${id}/publish`, data, { timeout: 600000 }), unpublish: (id: string) => diff --git a/webui/src/components/common/CommandDropdown.test.ts b/webui/src/components/common/CommandDropdown.test.ts new file mode 100644 index 000000000..ab1b57a98 --- /dev/null +++ b/webui/src/components/common/CommandDropdown.test.ts @@ -0,0 +1,17 @@ +import { describe, expect, it } from 'vitest'; +import { parseSlashCommand } from './CommandDropdown'; + +describe('parseSlashCommand', () => { + it('parses normal slash commands', () => { + expect(parseSlashCommand('/help')).toEqual({ command: 'help', args: '' }); + expect(parseSlashCommand('/plan build the thing')).toEqual({ + command: 'plan', + args: 'build the thing', + }); + }); + + it('does not treat absolute filesystem paths as commands', () => { + expect(parseSlashCommand('/tmp/workspace/workflow.md')).toBeNull(); + expect(parseSlashCommand('/tmp/rex_integration_guide.md\n\nuse this')).toBeNull(); + }); +}); diff --git a/webui/src/components/common/CommandDropdown.tsx b/webui/src/components/common/CommandDropdown.tsx index a19334716..05e9b94f0 100644 --- a/webui/src/components/common/CommandDropdown.tsx +++ b/webui/src/components/common/CommandDropdown.tsx @@ -89,16 +89,21 @@ export default function CommandDropdown({ * 从输入文本中解析 slash 命令的名称和参数 * 例如 "/bug describe issue" → { command: "bug", args: "describe issue" } */ +export function isSlashCommandName(command: string): boolean { + return command.length > 0 && !/[\\/]/.test(command); +} + export function parseSlashCommand(text: string): { command: string; args: string } | null { const trimmed = text.trim(); if (!trimmed.startsWith('/')) return null; const withoutSlash = trimmed.slice(1); const spaceIndex = withoutSlash.indexOf(' '); - if (spaceIndex === -1) { - return { command: withoutSlash, args: '' }; + const command = spaceIndex === -1 ? withoutSlash : withoutSlash.slice(0, spaceIndex); + if (!isSlashCommandName(command)) { + return null; } return { - command: withoutSlash.slice(0, spaceIndex), - args: withoutSlash.slice(spaceIndex + 1).trim(), + command, + args: spaceIndex === -1 ? '' : withoutSlash.slice(spaceIndex + 1).trim(), }; } diff --git a/webui/src/components/common/OnboardingModal.test.tsx b/webui/src/components/common/OnboardingModal.test.tsx index cb04b6210..723d060fa 100644 --- a/webui/src/components/common/OnboardingModal.test.tsx +++ b/webui/src/components/common/OnboardingModal.test.tsx @@ -121,7 +121,7 @@ describe('OnboardingModal', () => { { id: 'qwen3-max', name: 'Qwen 3 Max' }, ]), makeProvider('openai-compatible', 'OpenAI Compatible', []), - makeProvider('deepseek', 'DeepSeek', [{ id: 'deepseek-chat', name: 'DeepSeek V3.2' }]), + makeProvider('deepseek', 'DeepSeek', [{ id: 'deepseek-v4-flash', name: 'DeepSeek V4 Flash' }]), ], }, }); diff --git a/webui/src/components/common/SessionChat.test.ts b/webui/src/components/common/SessionChat.test.ts index 0551d8ab4..5dcd06532 100644 --- a/webui/src/components/common/SessionChat.test.ts +++ b/webui/src/components/common/SessionChat.test.ts @@ -1,12 +1,11 @@ import React from 'react'; -import { act, render, screen, waitFor } from '@testing-library/react'; +import { act, fireEvent, render, screen, waitFor } from '@testing-library/react'; import userEvent from '@testing-library/user-event'; import { beforeEach, describe, expect, it, vi } from 'vitest'; import type { Message } from '@/types'; import { - areChatMessagePartsRenderEqual, areChatTimelineItemsRenderEqual, buildInstructionDisplayText, buildChatTimelineItems, @@ -35,6 +34,7 @@ import { shouldRefetchFinishedMessage, truncateToolDisplayText, } from './SessionChat'; +import { areChatMessagePartsRenderEqual } from './sessionChatRenderEquality'; const clientGetMock = vi.fn(); const clientPostMock = vi.fn(); @@ -435,6 +435,16 @@ describe('getMessageBubbleClassName', () => { // The message column owns the available width, so the inner bubble only // controls intrinsic sizing (`w-auto` vs `w-full`). Tests here therefore // assert width semantics, not legacy max-width literals. + it('allows every bubble variant to shrink within the message column', () => { + for (const compact of [false, true]) { + for (const isUser of [false, true]) { + const className = getMessageBubbleClassName({ compact, isUser, isEditing: false }); + expect(className).toContain('min-w-0'); + expect(className).toContain('max-w-full'); + } + } + }); + it('keeps non-editing user bubbles auto-sized in full layout', () => { const className = getMessageBubbleClassName({ compact: false, @@ -443,7 +453,7 @@ describe('getMessageBubbleClassName', () => { }); expect(className).toContain('w-auto'); - expect(className).not.toContain('w-full'); + expect(className.split(' ')).not.toContain('w-full'); }); it('expands editing user bubbles to full width in full layout', () => { @@ -1353,6 +1363,35 @@ describe('SessionChat error rendering', () => { expect(screen.getByText('Connection error.')).toBeInTheDocument(); expect(container.querySelectorAll('.animate-bounce')).toHaveLength(0); }); + + it('renders assistant error messages when the only part is blank text', () => { + useSessionMessagesMock.mockReturnValue({ + messages: [ + makeMessage({ + id: 'assistant-error-with-blank-text', + role: 'assistant', + parts: [{ id: 'blank-text', type: 'text', text: '' }] as Message['parts'], + finish: 'error', + error: { + name: 'EmptyResponseError', + data: { message: 'Model returned an empty response.' }, + } as any, + }), + ], + loading: false, + refetch: vi.fn(), + addMessage: vi.fn(), + updateMessage: vi.fn(), + updateMessagePart: vi.fn(), + replaceMessageText: vi.fn(), + truncateAfterMessage: vi.fn(), + }); + + const { container } = render(React.createElement(SessionChat, { sessionId: 'sess-1' })); + + expect(screen.getByText('Model returned an empty response.')).toBeInTheDocument(); + expect(container.querySelectorAll('.animate-bounce')).toHaveLength(0); + }); }); describe('SessionChat intermediate process collapse', () => { @@ -2252,6 +2291,32 @@ describe('SessionChat agent mentions', () => { }); }); +describe('SessionChat slash command routing', () => { + it('sends absolute filesystem paths as normal prompts instead of slash commands', async () => { + render(React.createElement(SessionChat, { + sessionId: 'sess-1', + })); + + const text = '/tmp/stream_alert_denoise/rex_integration_guide.md\n\nuse this file'; + const textarea = screen.getByPlaceholderText('请输入消息') as HTMLTextAreaElement; + fireEvent.change(textarea, { target: { value: text } }); + fireEvent.keyDown(textarea, { key: 'Enter', code: 'Enter' }); + + await waitFor(() => { + expect(clientPostMock).toHaveBeenCalledWith( + '/api/session/sess-1/prompt_async', + expect.objectContaining({ + parts: [{ type: 'text', text }], + }), + ); + }); + expect(clientPostMock).not.toHaveBeenCalledWith( + '/api/session/sess-1/command', + expect.anything(), + ); + }); +}); + describe('truncateToolDisplayText', () => { it('returns short text unchanged', () => { expect(truncateToolDisplayText('bash')).toBe('bash'); @@ -2937,6 +3002,37 @@ describe('streaming activity helpers', () => { }); describe('SessionChat fallback polling', () => { + it('reconciles pending questions while the session is busy', async () => { + vi.useFakeTimers(); + try { + render(React.createElement(SessionChat, { + sessionId: 'sess-1', + live: true, + })); + + await act(async () => { + await Promise.resolve(); + }); + pendingQuestionsHookMock.fetchPendingQuestions.mockClear(); + + act(() => { + useSSEOptionsRef.current.onEvent({ + type: 'session.status', + properties: { sessionID: 'sess-1', status: { type: 'busy' } }, + }); + }); + pendingQuestionsHookMock.fetchPendingQuestions.mockClear(); + + await act(async () => { + await vi.advanceTimersByTimeAsync(2_000); + }); + + expect(pendingQuestionsHookMock.fetchPendingQuestions).toHaveBeenCalledWith('sess-1'); + } finally { + vi.useRealTimers(); + } + }); + it('does not finish streaming while fetched messages still contain a running tool', async () => { vi.useFakeTimers(); const refetch = vi.fn(); diff --git a/webui/src/components/common/SessionChat.tsx b/webui/src/components/common/SessionChat.tsx index c52d08b8a..90a94f391 100644 --- a/webui/src/components/common/SessionChat.tsx +++ b/webui/src/components/common/SessionChat.tsx @@ -23,7 +23,7 @@ import { useTranslation } from 'react-i18next'; import LoadingSpinner from './LoadingSpinner'; import { QuestionTool, type QuestionItem } from './QuestionTool'; import DelegateTaskCard, { isDelegateTool, shouldRenderDelegateTaskCard } from './DelegateTaskCard'; -import CommandDropdown, { parseSlashCommand } from './CommandDropdown'; +import CommandDropdown, { isSlashCommandName, parseSlashCommand } from './CommandDropdown'; import ImageLightbox from './ImageLightbox'; import { useSessionMessages } from '@/hooks/useSessions'; import { useSSE, type SSEConnectionStatus } from '@/hooks/useSSE'; @@ -34,6 +34,7 @@ import type { Command } from '@/api/skill'; import type { Agent } from '@/api/agent'; import { useToast } from './Toast'; import { buildRunWorkflowHeaderSummary } from './toolStageSummary'; +import { areChatMessagePartsRenderEqual } from './sessionChatRenderEquality'; import { workspaceAPI } from '@/api/workspace'; import { formatSmartTime } from '@/utils/time'; import { getAgentDisplayDescription } from '@/utils/agentDisplay'; @@ -821,7 +822,7 @@ export function getMessageBubbleClassName({ ? (isEditing ? 'w-full max-w-full' : 'max-w-full') : 'w-full max-w-full'; - return `${widthClass} px-4 py-3 rounded-[20px] text-sm break-words shadow-sm ${ + return `${widthClass} min-w-0 px-4 py-3 rounded-[20px] text-sm break-words shadow-sm ${ isUser ? 'bg-sky-50 border border-sky-100 text-zinc-900 dark:border-zinc-700 dark:bg-zinc-800 dark:text-zinc-50 dark:shadow-none' : 'bg-white border border-zinc-200/90 dark:border-zinc-800 dark:bg-zinc-900 dark:text-zinc-100 dark:shadow-none' @@ -832,7 +833,7 @@ export function getMessageBubbleClassName({ ? (isEditing ? 'w-full' : 'w-auto') : 'w-full'; - return `${widthClass} px-5 py-4 rounded-[24px] text-sm break-words shadow-sm ${ + return `${widthClass} min-w-0 max-w-full px-5 py-4 rounded-[24px] text-sm break-words shadow-sm ${ isUser ? 'bg-sky-50 border border-sky-100 text-zinc-900 dark:border-zinc-700 dark:bg-zinc-800 dark:text-zinc-50 dark:shadow-none' : 'bg-white border border-zinc-200/90 dark:border-zinc-800 dark:bg-zinc-900 dark:text-zinc-100 dark:shadow-none' @@ -1079,83 +1080,6 @@ export function getUserAvatarSpacerClassName(_compact: boolean): string { return 'h-0'; } -function areToolStatesRenderEqual( - prevState?: ToolState, - nextState?: ToolState, -): boolean { - if (prevState === nextState) return true; - if ( - prevState?.status !== nextState?.status || - prevState?.title !== nextState?.title || - prevState?.error !== nextState?.error || - prevState?.time?.start !== nextState?.time?.start || - prevState?.time?.end !== nextState?.time?.end - ) { - return false; - } - - return ( - JSON.stringify(prevState?.input) === JSON.stringify(nextState?.input) - && JSON.stringify(prevState?.output) === JSON.stringify(nextState?.output) - && JSON.stringify(prevState?.metadata) === JSON.stringify(nextState?.metadata) - ); -} - -function areLegacyToolPayloadsRenderEqual( - prevPayload?: MessagePart['toolCall'] | MessagePart['toolResult'], - nextPayload?: MessagePart['toolCall'] | MessagePart['toolResult'], -): boolean { - if (prevPayload === nextPayload) return true; - return JSON.stringify(prevPayload) === JSON.stringify(nextPayload); -} - -export function areChatMessagePartsRenderEqual( - prevParts?: MessagePart[], - nextParts?: MessagePart[], -): boolean { - if (prevParts === nextParts) return true; - if ((prevParts?.length ?? 0) !== (nextParts?.length ?? 0)) return false; - - const total = prevParts?.length ?? 0; - for (let i = 0; i < total; i++) { - const prevPart = prevParts?.[i]; - const nextPart = nextParts?.[i]; - - if (prevPart === nextPart) continue; - if (!prevPart || !nextPart) return false; - - if ( - prevPart.id !== nextPart.id || - prevPart.type !== nextPart.type || - prevPart.text !== nextPart.text || - prevPart.thinking !== nextPart.thinking || - prevPart.synthetic !== nextPart.synthetic || - prevPart.ignored !== nextPart.ignored || - prevPart.tool !== nextPart.tool || - prevPart.callID !== nextPart.callID || - prevPart.mime !== nextPart.mime || - prevPart.filename !== nextPart.filename || - prevPart.url !== nextPart.url || - prevPart.image?.url !== nextPart.image?.url || - prevPart.image?.alt !== nextPart.image?.alt - ) { - return false; - } - - if (!areToolStatesRenderEqual(prevPart.state, nextPart.state)) { - return false; - } - if (!areLegacyToolPayloadsRenderEqual(prevPart.toolCall, nextPart.toolCall)) { - return false; - } - if (!areLegacyToolPayloadsRenderEqual(prevPart.toolResult, nextPart.toolResult)) { - return false; - } - } - - return true; -} - // ============================================================================ // Main component // ============================================================================ @@ -1163,6 +1087,7 @@ export function areChatMessagePartsRenderEqual( const ABORT_SSE_SETTLE_DELAY = 2000; const SCROLL_BOTTOM_THRESHOLD_PX = 80; const FALLBACK_POLL_MS = 5_000; +const PENDING_QUESTION_RECONCILE_MS = 2_000; const WORKSPACE_UPLOAD_DEST = 'uploads'; const FILE_INPUT_ACCEPT_DOCS = '.txt,.md,.json,.yaml,.yml,.xml,.csv,.pdf,.doc,.docx,.html,.htm,.ppt,.pptx,.xls,.xlsx'; const FILE_INPUT_ACCEPT_ALL = `${FILE_INPUT_ACCEPT_DOCS},${FILE_INPUT_ACCEPT_IMAGES}`; @@ -2198,6 +2123,19 @@ export default function SessionChat({ checkStatus(); }, [sessionId, loading, messages, fetchPendingQuestions]); + // A remote proxy or a reconnect boundary can delay or lose the one-shot + // question.asked event. Reconcile only while the session is active; the + // hook ignores responses made stale by a newer SSE update. + useEffect(() => { + if (!sessionId || (!isStreaming && !sending)) return; + const timer = setInterval(() => { + fetchPendingQuestions(sessionId).catch((err) => { + console.warn('[SessionChat] Failed to reconcile pending questions:', err); + }); + }, PENDING_QUESTION_RECONCILE_MS); + return () => clearInterval(timer); + }, [sessionId, isStreaming, sending, fetchPendingQuestions]); + // Refetch when page becomes visible again useEffect(() => { if (!sessionId) return; @@ -2205,11 +2143,14 @@ export default function SessionChat({ if (document.visibilityState === 'visible') { refetch(); fetchPromptQueue(); + fetchPendingQuestions(sessionId).catch((err) => { + console.warn('[SessionChat] Failed to recover pending questions after visibility change:', err); + }); } }; document.addEventListener('visibilitychange', handler); return () => document.removeEventListener('visibilitychange', handler); - }, [sessionId, refetch, fetchPromptQueue]); + }, [sessionId, refetch, fetchPromptQueue, fetchPendingQuestions]); // Backup refetch when compaction ends — covers SSE reconnect scenarios // where the session.status event may have been missed. @@ -3546,15 +3487,20 @@ export default function SessionChat({ const cursor = e.target.selectionStart ?? val.length; const mention = mentionAgents.length > 0 ? findMentionTrigger(val, cursor) : null; const trimmed = val.trimStart(); + const slashQuery = trimmed.startsWith('/') ? trimmed.slice(1) : ''; if (mention && !trimmed.startsWith('/')) { setMentionRange({ start: mention.start, end: mention.end }); setMentionQuery(mention.query); setSelectedMentionIndex(0); setShowCommandDropdown(false); - } else if (trimmed.startsWith('/') && !trimmed.includes(' ') && successfulAttachments.length === 0) { + } else if ( + trimmed.startsWith('/') && + !trimmed.includes(' ') && + (slashQuery === '' || isSlashCommandName(slashQuery)) && + successfulAttachments.length === 0 + ) { void loadCommandsIfNeeded(); - const q = trimmed.slice(1); - setCommandQuery(q); + setCommandQuery(slashQuery); setSelectedCommandIndex(0); setShowCommandDropdown(true); setMentionRange(null); @@ -3955,6 +3901,12 @@ function ChatMessageBubbleInner({ const iconButtonClass = 'group/action relative inline-flex h-6 w-6 items-center justify-center rounded-full border border-gray-200/80 bg-white/80 text-gray-400 transition-colors duration-150 hover:border-gray-300 hover:text-gray-700 disabled:opacity-40 disabled:cursor-not-allowed dark:border-zinc-800 dark:bg-zinc-900/80 dark:text-zinc-500 dark:hover:border-zinc-700 dark:hover:text-zinc-200'; const tooltipClass = 'pointer-events-none absolute bottom-full left-1/2 z-10 mb-1.5 -translate-x-1/2 whitespace-nowrap rounded-md bg-gray-900 px-2 py-1 text-[11px] font-medium text-white opacity-0 shadow-sm transition-opacity duration-150 group-hover/action:opacity-100'; const messageErrorText = isUser ? '' : getMessageErrorText(message); + const hasOnlyBlankTextParts = parts.length > 0 && parts.every((part) => + part.type === 'text' && !String(part.text || '').trim() + ); + const shouldRenderAssistantErrorState = !isUser && !!messageErrorText && ( + parts.length === 0 || hasOnlyBlankTextParts + ); const avatarSize = compact ? 'w-7 h-7 text-xs' : 'w-8 h-8 text-sm'; const avatar = isUser ? ( @@ -3972,7 +3924,7 @@ function ChatMessageBubbleInner({
{/* Empty / loading state */} - {parts.length === 0 && ( + {(parts.length === 0 || shouldRenderAssistantErrorState) && ( isUser ? (
diff --git a/webui/src/components/common/StreamingMarkdown.test.tsx b/webui/src/components/common/StreamingMarkdown.test.tsx index 27af58ed8..08c2437ed 100644 --- a/webui/src/components/common/StreamingMarkdown.test.tsx +++ b/webui/src/components/common/StreamingMarkdown.test.tsx @@ -311,6 +311,15 @@ describe('useStreamingContent', () => { }); describe('StreamingMarkdown', () => { + it('constrains rendered Markdown to its message container', () => { + const { container } = render( + , + ); + + expect(container.firstElementChild).toHaveClass('w-full', 'min-w-0', 'max-w-full'); + expect(container.querySelector('pre')).not.toBeNull(); + }); + it('preserves single newlines as visible line breaks', () => { const { container } = render( , diff --git a/webui/src/components/common/StreamingMarkdown.tsx b/webui/src/components/common/StreamingMarkdown.tsx index 8f6e4d83d..3691267eb 100644 --- a/webui/src/components/common/StreamingMarkdown.tsx +++ b/webui/src/components/common/StreamingMarkdown.tsx @@ -248,7 +248,7 @@ export interface StreamingMarkdownProps { */ const MarkdownContent = memo(function MarkdownContent({ content }: { content: string }) { return ( -
+
{ { id: 'qwen3-max', name: 'Qwen 3 Max' }, ]), makeProvider('openai-compatible', 'OpenAI Compatible', []), - makeProvider('deepseek', 'DeepSeek', [{ id: 'deepseek-chat', name: 'DeepSeek V3.2' }]), + makeProvider('deepseek', 'DeepSeek', [{ id: 'deepseek-v4-flash', name: 'DeepSeek V4 Flash' }]), ], }, }); @@ -481,8 +481,11 @@ describe('Layout onboarding entry', () => { expect(await screen.findByText('admin.roleMember')).toBeInTheDocument(); expect(await screen.findByText('v2026.6.21')).toBeInTheDocument(); - expect(screen.queryByRole('link', { name: 'Flocks' })).not.toBeInTheDocument(); - await waitFor(() => expect(checkUpdate).toHaveBeenCalledWith('zh-CN', 'flockspro')); + expect(screen.queryByRole('link', { name: 'Flocks Pro' })).not.toBeInTheDocument(); + expect(checkUpdate).not.toHaveBeenCalled(); + await waitFor(() => { + expect(getActiveNotifications).toHaveBeenCalledWith('zh-CN'); + }); const sidebarShell = container.querySelector('aside > div'); const logoRow = sidebarShell?.firstElementChild as HTMLElement | null; @@ -537,7 +540,7 @@ describe('Layout onboarding entry', () => { await user.click(screen.getByRole('button', { name: 'admin settings' })); - expect(screen.getByRole('link', { name: 'Flocks' })).toHaveAttribute('href', '/settings/flockspro'); + expect(screen.getByRole('link', { name: 'Flocks Pro' })).toHaveAttribute('href', '/settings/flockspro'); expect(screen.getByRole('button', { name: 'checkUpdate' })).toBeInTheDocument(); expect(screen.getByRole('link', { name: 'settings' })).toHaveAttribute('href', '/settings/preferences'); @@ -608,7 +611,7 @@ describe('Layout onboarding entry', () => { await user.click(screen.getByRole('button', { name: 'admin settings' })); - expect(screen.getByRole('link', { name: 'Flocks' })).toHaveAttribute('href', '/settings/flockspro'); + expect(screen.getByRole('link', { name: 'Flocks Pro' })).toHaveAttribute('href', '/settings/flockspro'); expect(screen.getByRole('link', { name: 'settings' })).toHaveAttribute('href', '/settings/preferences'); await user.click(screen.getByRole('button', { name: 'logout' })); @@ -845,6 +848,33 @@ describe('Layout onboarding entry', () => { expect(await screen.findByText('Token 免费期已延长')).toBeInTheDocument(); expect(screen.getByText('Flocks v2026.04.28 更新内容')).toBeInTheDocument(); }); + + it('loads backend notifications without waiting for the update check', async () => { + localStorage.setItem('flocks_onboarding_dismissed', 'true'); + const updateCheck = deferred<{ + has_update: boolean; + latest_version: null; + current_version: string; + error: null; + }>(); + checkUpdate.mockReturnValue(updateCheck.promise); + + renderHomeWithLayout(); + + await waitFor(() => { + expect(getActiveNotifications).toHaveBeenCalledWith('zh-CN'); + }); + expect(getNotificationAckStatus).not.toHaveBeenCalled(); + + await act(async () => { + updateCheck.resolve({ + has_update: false, + latest_version: null, + current_version: '0.2.0', + error: null, + }); + }); + }); }); describe('Layout WebUI contract pages navigation', () => { diff --git a/webui/src/components/layout/Layout.tsx b/webui/src/components/layout/Layout.tsx index f4758e1ca..f361c9a41 100644 --- a/webui/src/components/layout/Layout.tsx +++ b/webui/src/components/layout/Layout.tsx @@ -198,7 +198,7 @@ export default function Layout() { const { t: tWebUIContractPage } = useTranslation('webuiContractPage'); const { t: tAuth } = useTranslation('auth'); const toast = useToast(); - const { productName } = useProductName(); + const { productName, proProductName } = useProductName(); const [hasUpdate, setHasUpdate] = useState(false); const [latestVersion, setLatestVersion] = useState(null); const [currentVersion, setCurrentVersion] = useState(null); @@ -217,6 +217,8 @@ export default function Layout() { const [flocksproStatusReady, setFlocksproStatusReady] = useState(false); const [flocksproVersion, setFlocksproVersion] = useState(null); const canManageUpdates = user?.role === 'admin'; + const notificationGateReady = flocksproStatusReady + && (!canManageUpdates || hasCompletedUpdateCheck); const canCreateWorkspaceCustomPage = user?.role === 'admin'; const { pages: webuiContractPages, workspaces: webuiContractWorkspaces = [] } = useWebUIContractPages(); const [openWorkspaceMenuId, setOpenWorkspaceMenuId] = useState(null); @@ -253,7 +255,7 @@ export default function Layout() { }, [handleOpenOnboarding]); const refreshUpdateStatus = useCallback(async (bypassMinGap = false) => { - if (!flocksproStatusReady) return; + if (!flocksproStatusReady || !canManageUpdates) return; const now = Date.now(); if (checkingUpdateRef.current) return; @@ -304,7 +306,7 @@ export default function Layout() { }, [canManageUpdates, flocksproStatusReady, i18n.language, isFlocksproActive]); useEffect(() => { - if (!flocksproStatusReady) return undefined; + if (!flocksproStatusReady || !canManageUpdates) return undefined; const initialCheckTimerId = window.setTimeout(() => { refreshUpdateStatus(true); @@ -335,7 +337,7 @@ export default function Layout() { document.removeEventListener('visibilitychange', handleVisibilityChange); window.removeEventListener('focus', handleWindowFocus); }; - }, [flocksproStatusReady, refreshUpdateStatus]); + }, [canManageUpdates, flocksproStatusReady, refreshUpdateStatus]); useEffect(() => { let cancelled = false; @@ -397,16 +399,14 @@ export default function Layout() { lastNotificationFetchKeyRef.current = null; return; } - if (!hasCompletedUpdateCheck) return; - - const fetchKey = `${user.id}:${i18n.language}:${currentVersion ?? 'pending-version'}`; + const fetchKey = `${user.id}:${i18n.language}`; if (lastNotificationFetchKeyRef.current === fetchKey) return; const previousFetchKey = lastNotificationFetchKeyRef.current; lastNotificationFetchKeyRef.current = fetchKey; setBackendNotificationsReady(false); let cancelled = false; - void getActiveNotifications(i18n.language, currentVersion) + void getActiveNotifications(i18n.language) .then((items) => { if (cancelled) return; setNotifications((prev) => { @@ -431,7 +431,7 @@ export default function Layout() { return () => { cancelled = true; }; - }, [currentVersion, hasCompletedUpdateCheck, i18n.language, user?.id]); + }, [i18n.language, user?.id]); useEffect(() => { if (!user?.id) { @@ -439,7 +439,7 @@ export default function Layout() { setUpdateNotificationReady(false); return; } - if (!hasCompletedUpdateCheck) return; + if (!notificationGateReady) return; setUpdateNotificationReady(false); const notification = buildUpdateNotification(updateInfo, i18n.language); @@ -465,7 +465,7 @@ export default function Layout() { return () => { cancelled = true; }; - }, [hasCompletedUpdateCheck, i18n.language, updateInfo, user?.id]); + }, [i18n.language, notificationGateReady, updateInfo, user?.id]); const allNotifications = updateNotification ? [...notifications, updateNotification].sort((a, b) => a.priority - b.priority) @@ -914,7 +914,7 @@ export default function Layout() { className="flex items-center gap-2 px-3 py-2 text-sm font-medium text-zinc-700 transition-colors hover:bg-zinc-50 hover:text-zinc-950 dark:text-zinc-200 dark:hover:bg-zinc-800 dark:hover:text-zinc-50" > - {productName} + {proProductName} )} + + + + + + + {t('slack.openSlackApps')} + {t('slack.openSlackAppsDesc')} + + + + +
+
    +
  1. + 1 + {t('slack.stepCreateApp')} +
    2. +
  2. + 2 + {t('slack.stepInstall')} +
    3. +
  3. + 3 + {t('slack.stepTokens')} +
    4. +
  4. + 4 + {t('slack.stepSaveEnable')} +
    5. +
+

+ * + {t('slack.manifestHint')} +

+
+ + +
+ + set('botToken', v || undefined)} + placeholder="xoxb-..." + /> + + + set('appToken', v || undefined)} + placeholder="xapp-..." + /> + + + set('homeChannel', v || undefined)} + placeholder="C0123456789" + /> + + + set('homeChannelName', v || undefined)} + placeholder="general" + /> + +
+ +
+ + set('defaultAgent', v || undefined)} + placeholder={t('slack.optional')} + /> + + + set('allowBots', v as 'none' | 'mentions' | 'all')} + options={[ + { value: 'none', label: t('slack.allowBotsNone') }, + { value: 'mentions', label: t('slack.allowBotsMentions') }, + { value: 'all', label: t('slack.allowBotsAll') }, + ]} + /> + +
+ + ); +} + // ============================================================================ // Email Config Panel // ============================================================================ @@ -2791,8 +3070,8 @@ export default function ChannelPage() { configs[ch.id] = { ...defaultDingTalkConfig(), ...saved }; } else if (ch.id === 'telegram') { configs[ch.id] = { ...defaultTelegramConfig(), ...saved }; - } else if (ch.id === 'email') { - configs[ch.id] = { ...defaultEmailConfig(), ...saved }; + } else if (ch.id === 'slack') { + configs[ch.id] = { ...defaultSlackConfig(), ...saved }; } else if (ch.id === 'email') { configs[ch.id] = { ...defaultEmailConfig(), ...saved }; } else if (ch.id === 'whatsapp') { @@ -2863,7 +3142,7 @@ export default function ChannelPage() { const updatedChannels = { ...(fullConfig.channels ?? {}), ...Object.fromEntries( - Object.entries(channelConfigs).map(([id, cfg]) => [id, stripEmpty(cfg)]) + Object.entries(channelConfigs).map(([id, cfg]) => [id, stripChannelConfigForSave(id, cfg)]) ), }; @@ -3167,6 +3446,12 @@ export default function ChannelPage() { onRefresh={fetchAll} /> )} + {selectedId === 'slack' && ( + handleChannelConfigChange('slack', cfg)} + /> + )} {selectedId === 'email' && ( ): Record { } return result; } + +function stripChannelConfigForSave(channelId: string, cfg: Record): Record { + const result = stripEmpty(cfg); + + if (channelId === 'slack') { + const allowFrom = Array.isArray(cfg.allowFrom) ? cfg.allowFrom : []; + result.dmPolicy = allowFrom.length > 0 ? 'allowlist' : 'open'; + if (cfg.allowFrom === undefined) { + result.allowFrom = null; + } + } + + return result; +} diff --git a/webui/src/pages/FlocksproUpgrade/index.test.tsx b/webui/src/pages/FlocksproUpgrade/index.test.tsx index 200ccdb4b..147d47cf3 100644 --- a/webui/src/pages/FlocksproUpgrade/index.test.tsx +++ b/webui/src/pages/FlocksproUpgrade/index.test.tsx @@ -43,7 +43,10 @@ vi.mock('react-i18next', () => ({ return options.defaultValue; } if (key === 'upgrade.installedTitle') { - return `Flocks Pro ${options?.version || ''}`; + return `${options?.productName || ''} ${options?.version || ''}`; + } + if (key === 'upgrade.title') { + return `Upgrade to ${options?.productName || ''}`; } return key; }, @@ -109,6 +112,25 @@ describe('FlocksproUpgradePage', () => { expect(await screen.findByRole('button', { name: 'upgrade.startUpgrade' })).toBeInTheDocument(); }); + it('uses the Pro product name when no custom display name is configured', async () => { + consoleUpgradeApi.listRequests.mockResolvedValue([]); + consoleUpgradeApi.getProPackageStatus.mockResolvedValue({ + installed: false, + runtime_importable: false, + install_marker_present: false, + pro_enabled: false, + license_status: 'uninstalled', + }); + + render( + + + , + ); + + expect(await screen.findByText('Upgrade to Flocks Pro')).toBeInTheDocument(); + }); + it('shows the installed bundle version instead of the Pro component version', async () => { consoleUpgradeApi.listRequests.mockResolvedValue([ { diff --git a/webui/src/pages/FlocksproUpgrade/index.tsx b/webui/src/pages/FlocksproUpgrade/index.tsx index 89eb4198f..9559194c4 100644 --- a/webui/src/pages/FlocksproUpgrade/index.tsx +++ b/webui/src/pages/FlocksproUpgrade/index.tsx @@ -303,7 +303,7 @@ function requestDurationDays(item: UpgradeRequestStatus): number | null { export default function FlocksproUpgradePage() { const { t } = useTranslation('flockspro'); - const { productName } = useProductName(); + const { proProductName } = useProductName(); const [searchParams, setSearchParams] = useSearchParams(); const [consoleLoginStatus, setConsoleLoginStatus] = useState(null); const [consoleLoginLoading, setConsoleLoginLoading] = useState(false); @@ -993,8 +993,8 @@ export default function FlocksproUpgradePage() { return (
} /> @@ -1052,13 +1052,13 @@ export default function FlocksproUpgradePage() {

{isProLoaded - ? t('upgrade.installedTitle', { productName, version: proVersion }) - : t('upgrade.title', { productName })} + ? t('upgrade.installedTitle', { productName: proProductName, version: proVersion }) + : t('upgrade.title', { productName: proProductName })}

{isProLoaded - ? t('upgrade.installedDescription', { productName }) - : t('upgrade.description', { productName })} + ? t('upgrade.installedDescription', { productName: proProductName }) + : t('upgrade.description', { productName: proProductName })}

{isProLoaded ? ( @@ -1231,7 +1231,7 @@ export default function FlocksproUpgradePage() { > {currentLicenseInvalid && (
- {t('upgrade.revokedOrExpiredHint', { productName })} + {t('upgrade.revokedOrExpiredHint', { productName: proProductName })}
)}
-

{t('upgrade.applyDialogTitle', { productName })}

+

{t('upgrade.applyDialogTitle', { productName: proProductName })}

{t('upgrade.productLabel')}
@@ -1492,7 +1492,7 @@ export default function FlocksproUpgradePage() {

{t('upgrade.downgradeDialogTitle')}

-

{t('upgrade.downgradeDialogDescription', { productName })}

+

{t('upgrade.downgradeDialogDescription', { productName: proProductName })}

@@ -1543,8 +1543,8 @@ export default function FlocksproUpgradePage() { ? t('upgrade.waitingDowngradeRestart') : t('upgrade.waitingRestart') : updateMode === 'downgrade' - ? t('upgrade.downgradingHint', { productName }) - : t('upgrade.installingHint', { productName })} + ? t('upgrade.downgradingHint', { productName: proProductName }) + : t('upgrade.installingHint', { productName: proProductName })}
{upgradeSteps.length > 0 && (
@@ -1564,7 +1564,7 @@ export default function FlocksproUpgradePage() { )}
- {t(`upgrade.stageLabels.${step.stage}`, { defaultValue: step.stage, productName })} + {t(`upgrade.stageLabels.${step.stage}`, { defaultValue: step.stage, productName: proProductName })}
{step.message} diff --git a/webui/src/pages/Home/index.test.tsx b/webui/src/pages/Home/index.test.tsx index f5141e841..97b1a0adf 100644 --- a/webui/src/pages/Home/index.test.tsx +++ b/webui/src/pages/Home/index.test.tsx @@ -4,11 +4,12 @@ import userEvent from '@testing-library/user-event'; import { MemoryRouter } from 'react-router-dom'; import Home from './index'; -const { createMock, navigateMock, toastErrorMock, useAuthMock } = vi.hoisted(() => ({ +const { createMock, navigateMock, toastErrorMock, useAuthMock, useStatsMock } = vi.hoisted(() => ({ createMock: vi.fn(), navigateMock: vi.fn(), toastErrorMock: vi.fn(), useAuthMock: vi.fn(), + useStatsMock: vi.fn(), })); vi.mock('react-router-dom', async () => { @@ -26,11 +27,7 @@ vi.mock('@/api/session', () => ({ })); vi.mock('@/hooks/useStats', () => ({ - useStats: () => ({ - stats: null, - loading: false, - error: null, - }), + useStats: () => useStatsMock(), })); vi.mock('@/components/common/Toast', () => ({ @@ -54,6 +51,11 @@ describe('Home create WebUI contract page entry', () => { beforeEach(() => { vi.clearAllMocks(); createMock.mockResolvedValue({ id: 'session-webui-contract-1' }); + useStatsMock.mockReturnValue({ + stats: null, + loading: false, + error: null, + }); useAuthMock.mockReturnValue({ user: { id: 'user-1', @@ -76,7 +78,9 @@ describe('Home create WebUI contract page entry', () => { await user.click(screen.getByRole('button', { name: 'createWebUIContractPage' })); await waitFor(() => { - expect(createMock).toHaveBeenCalledWith({ title: 'createWebUIContractPageSessionTitle' }); + expect(createMock).toHaveBeenCalledWith({ + title: 'createWebUIContractPageSessionTitle', + }); }); expect(navigateMock).toHaveBeenCalledWith( @@ -104,4 +108,21 @@ describe('Home create WebUI contract page entry', () => { expect(screen.queryByRole('button', { name: 'createWebUIContractPage' })).not.toBeInTheDocument(); expect(createMock).not.toHaveBeenCalled(); }); + + it('shows the concrete stats load failure instead of a backend-not-running hint', () => { + useStatsMock.mockReturnValue({ + stats: null, + loading: false, + error: new Error('authExpired'), + }); + + render( + + + , + ); + + expect(screen.getByText('stats.loadErrorHint.authExpired')).toBeInTheDocument(); + expect(screen.queryByText(/backend is running/i)).not.toBeInTheDocument(); + }); }); diff --git a/webui/src/pages/Home/index.tsx b/webui/src/pages/Home/index.tsx index 27d4198ae..f2d6e0c79 100644 --- a/webui/src/pages/Home/index.tsx +++ b/webui/src/pages/Home/index.tsx @@ -40,6 +40,7 @@ export default function Home() { const canCreateWebUIContractPage = user?.role === 'admin'; const [isRepoMenuOpen, setIsRepoMenuOpen] = useState(false); const [creatingWebUIContractPageSession, setCreatingWebUIContractPageSession] = useState(false); + const statsErrorHint = error ? t(`stats.loadErrorHint.${error.message}`, { defaultValue: t('stats.loadErrorHint.unavailable') }) : ''; const handleCreateWebUIContractPage = useCallback(async () => { if (creatingWebUIContractPageSession) return; @@ -205,7 +206,7 @@ export default function Home() {
{t('stats.abnormal')} - Please ensure the {productName} backend is running + {statsErrorHint}
)} diff --git a/webui/src/pages/Model/index.test.tsx b/webui/src/pages/Model/index.test.tsx index 1515819c6..cbd42e1d7 100644 --- a/webui/src/pages/Model/index.test.tsx +++ b/webui/src/pages/Model/index.test.tsx @@ -21,6 +21,7 @@ const mocks = vi.hoisted(() => ({ catalogList: vi.fn(), createProvider: vi.fn(), getCredentials: vi.fn(), + revealCredentials: vi.fn(), setCredentials: vi.fn(), testCredentials: vi.fn(), })); @@ -113,6 +114,7 @@ vi.mock('@/components/common/EntitySheet', () => ({ vi.mock('@/api/provider', () => ({ providerAPI: { getCredentials: mocks.getCredentials, + revealCredentials: mocks.revealCredentials, setCredentials: mocks.setCredentials, testCredentials: mocks.testCredentials, }, @@ -277,6 +279,15 @@ describe('ModelPage configure provider dialog', () => { has_credential: true, }, }); + mocks.revealCredentials.mockResolvedValue({ + data: { + secret_id: 'openai_llm_key', + api_key: 'sk-existing-secret-1234', + api_key_masked: 'sk-***1234', + base_url: 'https://old.example.com/v1', + has_credential: true, + }, + }); mocks.setCredentials.mockResolvedValue({ data: { success: true } }); mocks.testCredentials.mockResolvedValue({ data: { @@ -294,15 +305,18 @@ describe('ModelPage configure provider dialog', () => { renderWithRouter(); await user.click(await screen.findByTitle('Configure')); expect(await screen.findByTestId('entity-sheet')).toBeInTheDocument(); + expect(mocks.revealCredentials).toHaveBeenCalledWith('openai'); } - it('preserves an existing secret when saving a new base URL with a blank key', async () => { + it('displays and preserves the existing API key when saving a new base URL', async () => { const user = userEvent.setup(); await openConfigureDialog(user); - const apiKeyInput = screen.getByPlaceholderText('Leave blank to keep the existing API key'); - expect(apiKeyInput).toHaveValue(''); + const apiKeyInput = screen.getByDisplayValue('sk-existing-secret-1234'); + expect(apiKeyInput).toHaveAttribute('type', 'password'); expect(apiKeyInput).not.toHaveValue('sk-***1234'); + await user.click(screen.getByTitle('form.show')); + expect(apiKeyInput).toHaveAttribute('type', 'text'); const baseUrlInput = screen.getByPlaceholderText('https://api.example.com/v1'); await user.clear(baseUrlInput); @@ -317,7 +331,7 @@ describe('ModelPage configure provider dialog', () => { expect(payload).not.toHaveProperty('api_key'); }); - it('persists a base URL change without a key before testing the connection', async () => { + it('persists the displayed API key and a base URL change before testing the connection', async () => { const user = userEvent.setup(); await openConfigureDialog(user); mocks.setCredentials.mockClear(); @@ -338,4 +352,19 @@ describe('ModelPage configure provider dialog', () => { expect(mocks.testCredentials).toHaveBeenCalledWith('openai', 'gpt-4o'); }); }); + + it('submits a replacement API key when the displayed value changes', async () => { + const user = userEvent.setup(); + await openConfigureDialog(user); + + const apiKeyInput = screen.getByDisplayValue('sk-existing-secret-1234'); + await user.clear(apiKeyInput); + await user.type(apiKeyInput, 'sk-replacement-secret'); + await user.click(screen.getByRole('button', { name: 'Save' })); + + await waitFor(() => expect(mocks.setCredentials).toHaveBeenCalled()); + expect(mocks.setCredentials.mock.calls.at(-1)?.[1]).toEqual(expect.objectContaining({ + api_key: 'sk-replacement-secret', + })); + }); }); diff --git a/webui/src/pages/Model/index.tsx b/webui/src/pages/Model/index.tsx index 3a8e1b4c8..b34dc23f1 100644 --- a/webui/src/pages/Model/index.tsx +++ b/webui/src/pages/Model/index.tsx @@ -505,7 +505,15 @@ export default function ModelPage() { onConfigure={async () => { await handleSelectProvider(provider); if (selectedProviderRef.current?.id === provider.id) { - setShowConfigDialog(true); + try { + const res = await providerAPI.revealCredentials(provider.id); + if (selectedProviderRef.current?.id === provider.id) { + setCredentials(res.data); + setShowConfigDialog(true); + } + } catch (err: any) { + toast.error(t('configFailed'), err.message); + } } }} /> @@ -554,7 +562,10 @@ export default function ModelPage() { provider={selectedProvider} existingCredentials={credentials} models={providerModels} - onClose={() => setShowConfigDialog(false)} + onClose={() => { + setShowConfigDialog(false); + setCredentials(null); + }} onConfigured={async () => { if (selectedProvider) { const res = await providerAPI.getCredentials(selectedProvider.id).catch(() => ({ data: null })); @@ -2179,7 +2190,11 @@ function getDefaultReasoningToggleValue(providerId: string, modelId: string): bo function allowsBuiltInVisionToggle(modelId: string): boolean { const lowered = modelId.toLowerCase(); - return lowered.includes('qwen3.6-plus') || lowered.includes('kimi-k2.6'); + return ( + lowered.includes('qwen3.6-plus') + || lowered.includes('kimi-k2.6') + || lowered.includes('kimi-k2.7-code') + ); } // ==================== Configure Dialog ==================== @@ -2192,10 +2207,7 @@ function ConfigureProviderDialog({ provider, existingCredentials, models, onClos const toast = useToast(); const { t } = useTranslation('model'); const hasExisting = existingCredentials?.has_credential ?? false; - // Credential reads intentionally never return the raw key. Keep the input - // empty and use hasExisting to express "leave blank to preserve"; in - // particular, never seed this field with api_key_masked. - const existingKey = ''; + const existingKey = existingCredentials?.api_key ?? ''; const existingBaseUrl = existingCredentials?.base_url ?? ''; const [baseUrl, setBaseUrl] = useState(existingCredentials?.base_url ?? ''); @@ -2273,6 +2285,7 @@ function ConfigureProviderDialog({ provider, existingCredentials, models, onClos const handleSubmit = async () => { const nextApiKey = apiKey.trim(); + const apiKeyChanged = nextApiKey !== existingKey.trim(); if (!nextApiKey && !hasExisting && !providerAllowsEmptyApiKey(provider.id)) { toast.warning('Please enter API Key'); return; @@ -2285,7 +2298,7 @@ function ConfigureProviderDialog({ provider, existingCredentials, models, onClos ? (providerName.trim() || undefined) : undefined, }; - if (nextApiKey) payload.api_key = nextApiKey; + if (nextApiKey && apiKeyChanged) payload.api_key = nextApiKey; await providerAPI.setCredentials(provider.id, payload); // Sync catalog model list: add newly selected, delete deselected @@ -2326,6 +2339,7 @@ function ConfigureProviderDialog({ provider, existingCredentials, models, onClos ); const nextApiKey = apiKey.trim(); + const apiKeyChanged = nextApiKey !== existingKey.trim(); if (!nextApiKey && !hasExisting && !providerAllowsEmptyApiKey(provider.id)) { toast.warning('Please enter API Key first'); return; @@ -2341,7 +2355,7 @@ function ConfigureProviderDialog({ provider, existingCredentials, models, onClos ? (providerName.trim() || undefined) : undefined, }; - if (nextApiKey) payload.api_key = nextApiKey; + if (nextApiKey && apiKeyChanged) payload.api_key = nextApiKey; await providerAPI.setCredentials(provider.id, payload); } catch (err: any) { toast.error(t('deleteFailed'), err.message); diff --git a/webui/src/pages/Session/index.test.tsx b/webui/src/pages/Session/index.test.tsx index 0f44eb568..0b2b857ed 100644 --- a/webui/src/pages/Session/index.test.tsx +++ b/webui/src/pages/Session/index.test.tsx @@ -1,10 +1,10 @@ import React from 'react'; import { beforeEach, describe, expect, it, vi } from 'vitest'; -import { act, render, screen, waitFor } from '@testing-library/react'; +import { act, fireEvent, render, screen, waitFor, within } from '@testing-library/react'; import userEvent from '@testing-library/user-event'; import { MemoryRouter, useNavigate } from 'react-router-dom'; import { __resetChatModelResourcesForTesting } from '@/hooks/useChatModelResources'; -import SessionPage, { getVisibleSessionGroupItems, groupSessionsByDate } from './index'; +import SessionPage from './index'; const { client, @@ -23,6 +23,9 @@ const { toast, } = vi.hoisted(() => ({ client: { + delete: vi.fn(), + get: vi.fn(), + patch: vi.fn(), post: vi.fn(), }, sessionApi: { @@ -92,6 +95,10 @@ vi.mock('@/components/common/Toast', () => ({ useToast: () => toast, })); +vi.mock('@/contexts/AuthContext', () => ({ + useAuth: () => ({ user: { id: 'user-1', username: 'admin', role: 'admin' } }), +})); + vi.mock('@/components/common/LoadingSpinner', () => ({ default: () =>
loading-spinner
, })); @@ -205,6 +212,7 @@ const session = { id: 'session-1', slug: 'session-1', projectID: 'project-1', + effectiveProjectID: 'default', directory: '/tmp/project', title: 'Original Session', version: '1.0.0', @@ -270,61 +278,6 @@ function deferred() { return { promise, resolve, reject }; } -describe('session sidebar grouping helpers', () => { - it('groups sessions by updated date and applies search filtering', () => { - const now = new Date(2026, 6, 9, 12, 0, 0); - const makeSession = (id: string, title: string, updated: number) => ({ - ...session, - id, - title, - time: { ...session.time, updated }, - }); - - const groups = groupSessionsByDate([ - makeSession('today', 'Today Investigation', new Date(2026, 6, 9, 9).getTime()), - makeSession('yesterday', 'Yesterday Work', new Date(2026, 6, 8, 9).getTime()), - makeSession('earlier', 'Old Investigation', new Date(2026, 5, 1, 9).getTime()), - ], 'investigation', now); - - expect(groups.map((group) => [group.key, group.items.map((item) => item.id)])).toEqual([ - ['today', ['today']], - ['earlier', ['earlier']], - ]); - }); - - it('limits collapsed older groups and reports hidden count', () => { - const group = { - key: 'thisWeek' as const, - labelKey: 'groupThisWeek', - items: Array.from({ length: 7 }, (_, index) => ({ - ...session, - id: `session-${index}`, - title: `Session ${index}`, - })), - }; - - expect(getVisibleSessionGroupItems({ - group, - expanded: false, - searching: false, - })).toMatchObject({ - visibleItems: group.items.slice(0, 5), - hiddenCount: 2, - limit: 5, - }); - - expect(getVisibleSessionGroupItems({ - group, - expanded: false, - searching: true, - })).toMatchObject({ - visibleItems: group.items, - hiddenCount: 0, - limit: Infinity, - }); - }); -}); - describe('SessionPage session actions menu', () => { beforeEach(() => { vi.clearAllMocks(); @@ -358,6 +311,16 @@ describe('SessionPage session actions menu', () => { }); defaultModelAPI.getResolved.mockResolvedValue({ data: { provider_id: '', model_id: '' } }); modelV2API.listDefinitions.mockResolvedValue({ data: { models: [] } }); + client.get.mockResolvedValue({ + data: [{ + id: 'default', + worktree: '/tmp/project', + name: '默认', + isDefault: true, + pathStatus: 'available', + sessionCount: 1, + }], + }); hubAPI.catalog.mockResolvedValue({ data: [{ id: 'soc-workspace', type: 'component', state: 'installed' }], }); @@ -365,6 +328,7 @@ describe('SessionPage session actions menu', () => { hubAPI.installStream.mockResolvedValue(undefined); sessionApi.update.mockResolvedValue({ ...session, title: 'Renamed Session' }); + client.patch.mockResolvedValue({ data: { id: 'prj_project2', worktree: '/tmp/labs', name: 'Renamed Project' } }); client.post.mockResolvedValue({ data: secondSession }); sessionApi.get.mockResolvedValue(session); sessionApi.getMessages.mockResolvedValue([ @@ -383,13 +347,725 @@ describe('SessionPage session actions menu', () => { vi.stubGlobal('confirm', vi.fn(() => true)); }); + it('shows default sessions under tasks without a default project row', async () => { + const user = userEvent.setup(); + renderSessionPage(); + + const tasksHeading = await screen.findByText('tasksSection'); + const projectsHeading = screen.getByText('projectsSection'); + const tasksSection = tasksHeading.closest('section'); + const projectsSection = projectsHeading.closest('section'); + expect(tasksSection).not.toBeNull(); + expect(projectsSection).not.toBeNull(); + expect(tasksSection?.parentElement).toBe(projectsSection?.parentElement); + expect(projectsSection).not.toContainElement(tasksHeading); + expect(useSessions).toHaveBeenLastCalledWith('', { + projectIds: ['tasks'], + pageSize: 20, + }); + expect(screen.queryByText('defaultProjectName')).not.toBeInTheDocument(); + expect(screen.getByText('Original Session')).toBeInTheDocument(); + + await user.click(screen.getByRole('button', { name: 'toggleTasks' })); + expect(screen.queryByText('Original Session')).not.toBeInTheDocument(); + + await user.click(screen.getByRole('button', { name: 'selectTasks' })); + expect(screen.getByText('Original Session')).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('collapses the projects section and restores it after remounting', async () => { + const user = userEvent.setup(); + client.get.mockResolvedValue({ + data: [ + { id: 'default', worktree: '/tmp/project', name: '默认', isDefault: true }, + { id: 'prj_labs', worktree: '/tmp/labs', name: 'Labs', isDefault: false }, + ], + }); + const firstRender = renderSessionPage(); + + await screen.findByText('Labs'); + await user.click(screen.getByRole('button', { name: 'toggleProjects' })); + expect(screen.queryByText('Labs')).not.toBeInTheDocument(); + + firstRender.unmount(); + renderSessionPage(); + + await screen.findByText('projectsSection'); + expect(screen.queryByText('Labs')).not.toBeInTheDocument(); + }); + + it('restores collapsed projects after the session page remounts', async () => { + const user = userEvent.setup(); + client.get.mockResolvedValue({ + data: [ + { id: 'default', worktree: '/tmp/project', name: '默认', isDefault: true }, + { id: 'prj_labs', worktree: '/tmp/labs', name: 'Labs', isDefault: false }, + ], + }); + useSessions.mockReturnValue({ + sessions: [{ + ...session, + projectID: 'prj_labs', + effectiveProjectID: 'prj_labs', + directory: '/tmp/labs', + }], + loading: false, + error: null, + refetch: refetchSessions, + updateSessionTitle, + removeSession, + removeSessions, + addSession, + }); + const firstRender = renderSessionPage(); + + await screen.findByText('Labs'); + await user.click(screen.getByRole('button', { name: 'selectProject' })); + expect(screen.queryByText('Original Session')).not.toBeInTheDocument(); + + firstRender.unmount(); + renderSessionPage(); + + await screen.findByText('Labs'); + expect(screen.queryByText('Original Session')).not.toBeInTheDocument(); + }); + + it('uses six sessions per page when multiple projects exist', async () => { + client.get.mockResolvedValue({ + data: [ + { + id: 'default', + worktree: '/tmp/project', + name: '默认', + isDefault: true, + pathStatus: 'available', + sessionCount: 1, + }, + { + id: 'prj_labs', + worktree: '/tmp/labs', + name: 'Labs', + isDefault: false, + pathStatus: 'available', + sessionCount: 0, + }, + ], + }); + + renderSessionPage(); + + await screen.findByText('Labs'); + expect(useSessions).toHaveBeenLastCalledWith('', { + projectIds: ['tasks', 'prj_labs'], + pageSize: 6, + }); + }); + + it('toggles project sessions when clicking the selected project row', async () => { + const user = userEvent.setup(); + client.get.mockResolvedValue({ + data: [ + { id: 'default', worktree: '/tmp/project', name: '默认', isDefault: true }, + { id: 'prj_labs', worktree: '/tmp/labs', name: 'Labs', isDefault: false }, + ], + }); + useSessions.mockReturnValue({ + sessions: [{ + ...session, + projectID: 'prj_labs', + effectiveProjectID: 'prj_labs', + directory: '/tmp/labs', + }], + loading: false, + error: null, + refetch: refetchSessions, + updateSessionTitle, + removeSession, + removeSessions, + addSession, + }); + + renderSessionPage(); + + await screen.findByText('Labs'); + const projectRow = screen.getByRole('button', { name: 'selectProject' }); + expect(screen.queryByRole('button', { name: 'toggleProject' })).not.toBeInTheDocument(); + expect(projectRow).toHaveAttribute('aria-expanded', 'true'); + + await user.click(projectRow); + expect(screen.queryByText('Original Session')).not.toBeInTheDocument(); + expect(projectRow).toHaveAttribute('aria-expanded', 'false'); + + await user.click(projectRow); + expect(screen.getByText('Original Session')).toBeInTheDocument(); + expect(projectRow).toHaveAttribute('aria-expanded', 'true'); + + await user.click(projectRow); + expect(screen.queryByText('Original Session')).not.toBeInTheDocument(); + }); + + it('renders project sessions with the same card outline as task sessions', async () => { + client.get.mockResolvedValue({ + data: [ + { id: 'default', worktree: '/tmp/project', name: '默认', isDefault: true }, + { id: 'prj_labs', worktree: '/tmp/labs', name: 'Labs', isDefault: false }, + ], + }); + useSessions.mockReturnValue({ + sessions: [{ + ...session, + projectID: 'prj_labs', + effectiveProjectID: 'prj_labs', + directory: '/tmp/labs', + }], + loading: false, + error: null, + refetch: refetchSessions, + updateSessionTitle, + removeSession, + removeSessions, + addSession, + }); + + renderSessionPage(); + + const sessionTitle = await screen.findByText('Original Session'); + const sessionCard = sessionTitle.closest('[class*="cursor-pointer"]'); + expect(sessionCard).not.toBeNull(); + expect(sessionCard).toHaveClass('border-gray-100'); + }); + + it('groups legacy sessions by the effective project returned by the backend', async () => { + client.get.mockResolvedValue({ + data: [{ + id: 'default', + worktree: '/tmp/project', + name: '默认', + isDefault: true, + pathStatus: 'available', + sessionCount: 2, + }], + }); + useSessions.mockReturnValue({ + sessions: [ + session, + { + ...secondSession, + projectID: 'old-project-id', + effectiveProjectID: 'default', + directory: '/tmp/project', + title: 'Legacy Session', + }, + ], + loading: false, + error: null, + refetch: refetchSessions, + updateSessionTitle, + removeSession, + removeSessions, + addSession, + }); + + renderSessionPage(); + + await screen.findByText('tasksSection'); + expect(screen.queryByText('defaultProjectName')).not.toBeInTheDocument(); + expect(screen.getByText('Original Session')).toBeInTheDocument(); + expect(screen.getByText('Legacy Session')).toBeInTheDocument(); + }); + + it('creates a user-managed project from the sidebar', async () => { + const user = userEvent.setup(); + let projectRows = [ + { id: 'default', worktree: '/tmp/project', name: '默认', isDefault: true }, + ]; + client.get.mockImplementation(() => Promise.resolve({ data: projectRows })); + client.post.mockImplementation((url: string, payload: Record) => { + if (url === '/api/project') { + const created = { id: 'prj_project2', worktree: payload.worktree as string, name: payload.name as string }; + projectRows = [projectRows[0], created]; + return Promise.resolve({ data: created }); + } + return Promise.resolve({ data: secondSession }); + }); + + renderSessionPage(); + + await user.click(await screen.findByRole('button', { name: 'projectDialog.createTitle' })); + const nameInput = screen.getByLabelText('projectDialog.nameLabel'); + await user.clear(nameInput); + await user.type(nameInput, 'Labs'); + const folderInput = screen.getByLabelText('projectDialog.folderLabel'); + await user.clear(folderInput); + await user.type(folderInput, '/tmp/labs'); + await user.click(screen.getByRole('button', { name: 'save' })); + + await waitFor(() => { + expect(client.post).toHaveBeenCalledWith('/api/project', { name: 'Labs', worktree: '/tmp/labs' }); + expect(screen.getByText('Labs')).toBeInTheDocument(); + }); + }); + + it('uses the current browser path when saving without selecting the folder', async () => { + const user = userEvent.setup(); + const defaultProject = { + id: 'default', + worktree: '/tmp/project', + name: '默认', + isDefault: true, + }; + client.get.mockImplementation((url: string) => Promise.resolve({ + data: url === '/api/project/folders' + ? { + path: '/home/test-user', + parent: null, + roots: [], + entries: [], + } + : [defaultProject], + })); + client.post.mockResolvedValue({ + data: { id: 'prj_home', name: 'test-user', worktree: '/home/test-user' }, + }); + + renderSessionPage(); + + await user.click(await screen.findByRole('button', { name: 'projectDialog.createTitle' })); + expect(screen.getByLabelText('projectDialog.nameLabel')).toHaveValue(''); + expect(screen.getByLabelText('projectDialog.folderLabel')).toHaveValue(''); + expect(screen.getByRole('button', { name: 'cancel' })).toBeEnabled(); + + await user.click(screen.getByRole('button', { name: 'projectDialog.chooseFolder' })); + + expect(await screen.findByText('/home/test-user')).toBeInTheDocument(); + expect(client.get).toHaveBeenCalledWith('/api/project/folders', { + params: { path: undefined }, + }); + expect(screen.getByLabelText('projectDialog.folderLabel')).toHaveValue('/home/test-user'); + + await user.click(screen.getByRole('button', { name: 'save' })); + + await waitFor(() => { + expect(client.post).toHaveBeenCalledWith('/api/project', { + name: 'test-user', + worktree: '/home/test-user', + }); + }); + }); + + it('keeps the folder input and folder browser in sync', async () => { + const user = userEvent.setup(); + client.get.mockImplementation((url: string, config?: { params?: { path?: string } }) => { + if (url !== '/api/project/folders') { + return Promise.resolve({ + data: [{ id: 'default', worktree: '/tmp/project', name: '默认', isDefault: true }], + }); + } + const path = config?.params?.path; + if (path === '/home/test-user/labs') { + return Promise.resolve({ + data: { path, parent: '/home/test-user', roots: [], entries: [] }, + }); + } + return Promise.resolve({ + data: { + path: '/home/test-user', + parent: null, + roots: [], + entries: [{ name: 'labs', path: '/home/test-user/labs' }], + }, + }); + }); + + renderSessionPage(); + + await user.click(await screen.findByRole('button', { name: 'projectDialog.createTitle' })); + const folderInput = screen.getByLabelText('projectDialog.folderLabel'); + await user.click(screen.getByRole('button', { name: 'projectDialog.chooseFolder' })); + await waitFor(() => expect(folderInput).toHaveValue('/home/test-user')); + + await user.type(folderInput, '/'); + await waitFor(() => { + expect(client.get).toHaveBeenCalledWith('/api/project/folders', { + params: { path: '/home/test-user/' }, + }); + }); + expect(folderInput).toHaveValue('/home/test-user/'); + + await user.clear(folderInput); + await user.type(folderInput, '/home/test-user/labs'); + await waitFor(() => { + expect(client.get).toHaveBeenCalledWith('/api/project/folders', { + params: { path: '/home/test-user/labs' }, + }); + expect(screen.getByText('/home/test-user/labs')).toBeInTheDocument(); + }); + }); + + it('does not submit when Enter is pressed before choosing a project folder', async () => { + const user = userEvent.setup(); + client.post.mockResolvedValue({ + data: { id: 'prj_labs', name: 'Labs', worktree: '/tmp/labs' }, + }); + renderSessionPage(); + + await user.click(await screen.findByRole('button', { name: 'projectDialog.createTitle' })); + const nameInput = screen.getByLabelText('projectDialog.nameLabel'); + await user.type(nameInput, 'Labs'); + fireEvent.keyDown(nameInput, { key: 'Enter' }); + + expect(client.post).not.toHaveBeenCalled(); + expect(toast.error).not.toHaveBeenCalled(); + + const folderInput = screen.getByLabelText('projectDialog.folderLabel'); + await user.type(folderInput, '/tmp/labs'); + fireEvent.keyDown(nameInput, { key: 'Enter', isComposing: true }); + expect(client.post).not.toHaveBeenCalled(); + + fireEvent.keyDown(nameInput, { key: 'Enter' }); + await waitFor(() => { + expect(client.post).toHaveBeenCalledWith('/api/project', { + name: 'Labs', + worktree: '/tmp/labs', + }); + }); + }); + + it('shows the backend detail when project creation fails', async () => { + const user = userEvent.setup(); + client.post.mockRejectedValue({ + message: 'Request failed with status code 400', + response: { data: { detail: 'Project directory does not exist' } }, + }); + + renderSessionPage(); + + await user.click(await screen.findByRole('button', { name: 'projectDialog.createTitle' })); + const nameInput = screen.getByLabelText('projectDialog.nameLabel'); + await user.clear(nameInput); + await user.type(nameInput, 'Missing project'); + const folderInput = screen.getByLabelText('projectDialog.folderLabel'); + await user.clear(folderInput); + await user.type(folderInput, '/tmp/missing-project'); + await user.click(screen.getByRole('button', { name: 'save' })); + + await waitFor(() => { + expect(toast.error).toHaveBeenCalledWith( + 'projectDialog.saveFailed', + 'Project directory does not exist', + ); + }); + }); + + it('submits project creation only once when save is clicked twice quickly', async () => { + const user = userEvent.setup(); + let resolveCreate: ((value: { data: Record }) => void) | undefined; + client.post.mockImplementation(() => new Promise((resolve) => { + resolveCreate = resolve; + })); + + renderSessionPage(); + + await user.click(await screen.findByRole('button', { name: 'projectDialog.createTitle' })); + const nameInput = screen.getByLabelText('projectDialog.nameLabel'); + await user.clear(nameInput); + await user.type(nameInput, 'Labs'); + const folderInput = screen.getByLabelText('projectDialog.folderLabel'); + await user.clear(folderInput); + await user.type(folderInput, '/tmp/labs'); + const saveButton = screen.getByRole('button', { name: 'save' }); + + act(() => { + saveButton.click(); + saveButton.click(); + }); + + expect(client.post).toHaveBeenCalledTimes(1); + resolveCreate?.({ data: { id: 'prj_project2', name: 'Labs', worktree: '/tmp/labs' } }); + await waitFor(() => expect(screen.queryByRole('button', { name: 'save' })).not.toBeInTheDocument()); + }); + + it('keeps a newly created empty project visible while search is active', async () => { + const user = userEvent.setup(); + const currentProject = { id: 'default', worktree: '/tmp/project', name: '默认', isDefault: true }; + client.get.mockResolvedValue({ data: [currentProject] }); + client.post.mockImplementation((url: string, payload: Record) => { + if (url === '/api/project') { + const created = { id: 'prj_project2', worktree: payload.worktree as string, name: payload.name as string }; + return Promise.resolve({ data: created }); + } + return Promise.resolve({ data: secondSession }); + }); + + renderSessionPage(); + + await user.type(screen.getByPlaceholderText('filterConversations'), 'nothing matches'); + await user.click(await screen.findByRole('button', { name: 'projectDialog.createTitle' })); + const nameInput = screen.getByLabelText('projectDialog.nameLabel'); + await user.clear(nameInput); + await user.type(nameInput, 'Labs'); + const folderInput = screen.getByLabelText('projectDialog.folderLabel'); + await user.clear(folderInput); + await user.type(folderInput, '/tmp/labs'); + await user.click(screen.getByRole('button', { name: 'save' })); + + expect(await screen.findByText('Labs')).toBeInTheDocument(); + expect(screen.getByText('noProjectSessions')).toBeInTheDocument(); + }); + + it('renames a project from the sidebar', async () => { + const user = userEvent.setup(); + const defaultProject = { id: 'default', worktree: '/tmp/project', name: '默认', isDefault: true }; + const projectList = [defaultProject, { + id: 'prj_project2', + worktree: '/tmp/labs', + name: 'Labs', + sessionCount: 8, + lastActivityAt: 10_000, + }]; + client.get + .mockResolvedValueOnce({ data: projectList }) + .mockResolvedValue({ + data: projectList.map((project) => ( + project.id === 'prj_project2' ? { ...project, name: 'Renamed Project' } : project + )), + }); + client.patch.mockResolvedValue({ + data: { id: 'prj_project2', worktree: '/tmp/labs', name: 'Renamed Project' }, + }); + + renderSessionPage(); + + const projectLabel = await screen.findByText('Labs'); + const projectRow = projectLabel.closest('[class*="group/project"]'); + expect(projectRow).not.toBeNull(); + await user.click(within(projectRow as HTMLElement).getByRole('button', { name: 'projectActions' })); + await user.click(within(projectRow as HTMLElement).getByRole('menuitem', { name: 'projectDialog.renameAction' })); + const input = screen.getByLabelText('projectDialog.nameLabel'); + await user.clear(input); + await user.type(input, 'Renamed Project'); + await user.click(screen.getByRole('button', { name: 'save' })); + + await waitFor(() => { + expect(client.patch).toHaveBeenCalledWith('/api/project/prj_project2', { name: 'Renamed Project' }); + }); + const renamedProject = await screen.findByText('Renamed Project'); + expect(renamedProject.closest('[class*="group/project"]')).toHaveTextContent('8'); + }); + + it('shares and unshares a project from the sidebar', async () => { + const user = userEvent.setup(); + const defaultProject = { id: 'default', worktree: '/tmp/project', name: '默认', isDefault: true }; + let projectRows = [ + defaultProject, + { + id: 'prj_project2', worktree: '/tmp/labs', name: 'Labs', + canWrite: true, canDelete: true, isShared: false, + }, + ]; + client.get.mockImplementation(() => Promise.resolve({ data: projectRows })); + client.post.mockImplementation((url: string) => { + if (url === '/api/project/prj_project2/share-local') { + projectRows = projectRows.map((project) => ( + project.id === 'prj_project2' ? { ...project, isShared: true } : project + )); + } + if (url === '/api/project/prj_project2/unshare-local') { + projectRows = projectRows.map((project) => ( + project.id === 'prj_project2' ? { ...project, isShared: false } : project + )); + } + return Promise.resolve({ data: true }); + }); + + renderSessionPage(); + + let projectRow = (await screen.findByText('Labs')).closest('[class*="group/project"]'); + expect(projectRow).not.toBeNull(); + await user.click(within(projectRow as HTMLElement).getByRole('button', { name: 'projectActions' })); + await user.click(within(projectRow as HTMLElement).getByRole('menuitem', { name: 'shareAction' })); + await waitFor(() => { + expect(client.post).toHaveBeenCalledWith('/api/project/prj_project2/share-local'); + expect(screen.getByText('sharedTag')).toBeInTheDocument(); + }); + + projectRow = screen.getByText('Labs').closest('[class*="group/project"]'); + await user.click(within(projectRow as HTMLElement).getByRole('button', { name: 'projectActions' })); + await user.click(within(projectRow as HTMLElement).getByRole('menuitem', { name: 'unshareAction' })); + await waitFor(() => { + expect(client.post).toHaveBeenCalledWith('/api/project/prj_project2/unshare-local'); + expect(screen.queryByText('sharedTag')).not.toBeInTheDocument(); + }); + }); + + it('keeps a shared project read-only for non-owners', async () => { + const user = userEvent.setup(); + client.get.mockResolvedValue({ + data: [ + { id: 'default', worktree: '/tmp/project', name: '默认', isDefault: true }, + { + id: 'prj_shared', worktree: '/tmp/shared', name: 'Shared Labs', + canWrite: false, canDelete: false, isShared: true, + }, + ], + }); + + renderSessionPage(); + + const projectRow = (await screen.findByText('Shared Labs')).closest('[class*="group/project"]'); + expect(projectRow).not.toBeNull(); + 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(); + expect(within(projectRow as HTMLElement).queryByRole('menuitem', { name: 'unshareAction' })).not.toBeInTheDocument(); + expect(within(projectRow as HTMLElement).queryByRole('menuitem', { name: 'projectDialog.renameAction' })).not.toBeInTheDocument(); + expect(within(projectRow as HTMLElement).queryByRole('menuitem', { name: 'projectDialog.deleteAction' })).not.toBeInTheDocument(); + }); + + it('ignores stale project results after the search changes', async () => { + const initialProject = { id: 'default', worktree: '/tmp/project', name: '默认', isDefault: true }; + const olderSearch = deferred<{ data: Array> }>(); + const latestSearch = deferred<{ data: Array> }>(); + client.get.mockImplementation((_url: string, config?: { params?: { search?: string } }) => { + const query = config?.params?.search; + if (query === 'a') return olderSearch.promise; + if (query === 'ab') return latestSearch.promise; + return Promise.resolve({ data: [initialProject] }); + }); + + renderSessionPage(); + await screen.findByText('tasksSection'); + const searchInput = screen.getByPlaceholderText('filterConversations'); + fireEvent.change(searchInput, { target: { value: 'a' } }); + fireEvent.change(searchInput, { target: { value: 'ab' } }); + + await act(async () => { + latestSearch.resolve({ + data: [initialProject, { + id: 'prj_latest', + worktree: '/tmp/latest', + name: 'Latest result', + isDefault: false, + matchedSessionCount: 1, + }], + }); + await latestSearch.promise; + }); + expect(await screen.findByText('Latest result')).toBeInTheDocument(); + + await act(async () => { + olderSearch.resolve({ + data: [initialProject, { + id: 'prj_stale', + worktree: '/tmp/stale', + name: 'Stale result', + isDefault: false, + matchedSessionCount: 1, + }], + }); + await olderSearch.promise; + }); + expect(screen.getByText('Latest result')).toBeInTheDocument(); + expect(screen.queryByText('Stale result')).not.toBeInTheDocument(); + }); + + it('does not render project actions for the default task group', async () => { + renderSessionPage(); + + await screen.findByText('tasksSection'); + expect(screen.queryByText('defaultProjectName')).not.toBeInTheDocument(); + expect(screen.queryByRole('button', { name: 'projectActions' })).not.toBeInTheDocument(); + expect(client.patch).not.toHaveBeenCalled(); + }); + + it('deletes an empty user-managed project after confirmation', async () => { + const user = userEvent.setup(); + const currentProject = { id: 'default', worktree: '/tmp/project', name: '默认', isDefault: true }; + let projectRows = [ + currentProject, + { id: 'prj_project2', worktree: '/tmp/labs', name: 'Labs' }, + ]; + client.get.mockImplementation(() => Promise.resolve({ data: projectRows })); + client.delete.mockImplementation(() => { + projectRows = [currentProject]; + return Promise.resolve({ data: true }); + }); + + renderSessionPage(); + + const projectLabel = await screen.findByText('Labs'); + const projectRow = projectLabel.closest('[class*="group/project"]'); + expect(projectRow).not.toBeNull(); + await user.click(within(projectRow as HTMLElement).getByRole('button', { name: 'projectActions' })); + await user.click(within(projectRow as HTMLElement).getByRole('menuitem', { name: 'projectDialog.deleteAction' })); + await user.click(screen.getByRole('button', { name: 'projectDialog.confirmDelete' })); + + await waitFor(() => { + expect(client.delete).toHaveBeenCalledWith('/api/project/prj_project2'); + expect(screen.queryByText('Labs')).not.toBeInTheDocument(); + }); + }); + + 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(); + await user.click(within(projectRow as HTMLElement).getByRole('button', { name: 'createSessionInProject' })); + + 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 () => { const user = userEvent.setup(); renderSessionPage(); + await screen.findByText('Original Session'); await user.click(screen.getByRole('button', { name: 'moreActions' })); + const menu = document.querySelector('[data-session-menu-portal]'); + expect(menu).toHaveClass('min-w-28', 'rounded-md', 'p-1'); + expect(menu).not.toHaveClass('w-36', 'rounded-lg'); expect(screen.getByRole('button', { name: 'rename' })).toBeInTheDocument(); expect(screen.getByRole('button', { name: 'downloadJson' })).toBeInTheDocument(); expect(screen.getByRole('button', { name: 'deleteAction' })).toBeInTheDocument(); @@ -400,6 +1076,7 @@ describe('SessionPage session actions menu', () => { renderSessionPage(); + await screen.findByText('Original Session'); await user.click(screen.getByRole('button', { name: 'moreActions' })); await user.click(screen.getByRole('button', { name: 'rename' })); @@ -476,6 +1153,7 @@ describe('SessionPage session actions menu', () => { renderSessionPage(); + await screen.findByText('Original Session'); await user.click(screen.getByRole('button', { name: 'moreActions' })); await user.click(screen.getByRole('button', { name: 'downloadJson' })); @@ -517,6 +1195,7 @@ describe('SessionPage session actions menu', () => { renderSessionPage(); + await screen.findByText('Original Session'); await user.click(screen.getByRole('button', { name: 'moreActions' })); await user.click(screen.getByRole('button', { name: 'deleteAction' })); diff --git a/webui/src/pages/Session/index.tsx b/webui/src/pages/Session/index.tsx index bf3a4a124..f9b27aba7 100644 --- a/webui/src/pages/Session/index.tsx +++ b/webui/src/pages/Session/index.tsx @@ -1,10 +1,11 @@ import { memo, useState, useEffect, useMemo, useCallback, useRef, type RefObject } from 'react'; import { - MessageSquare, Plus, Trash2, - ChevronDown, Sparkles, Shield, Search, AlertTriangle, + Plus, Trash2, + ChevronDown, ChevronRight, Sparkles, Shield, Search, AlertTriangle, PanelLeftClose, PanelLeft, Bot, Loader2, Workflow as WorkflowIcon, Settings2, CheckSquare, MoreHorizontal, PencilLine, Download, Share2, Cpu, Info, + FolderGit2, FolderPlus, FolderOpen, Copy, ArrowUp, HardDrive, } from 'lucide-react'; import { useTranslation } from 'react-i18next'; import i18n from '@/i18n'; @@ -31,6 +32,7 @@ import { buildPromptParts, type ImagePartData } from '@/utils/imageUpload'; import { getAgentDisplayDescription, getAgentDisplayName, isAgentUsableInChat } from '@/utils/agentDisplay'; import { formatSessionDate } from '@/utils/time'; import type { ModelDefinitionV2, Session } from '@/types'; +import { useAuth } from '@/contexts/AuthContext'; function sanitizeSessionExportName(value: string) { const trimmed = value.trim(); @@ -46,7 +48,44 @@ 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 TASK_SESSION_GROUP_ID = 'tasks'; type AgentSourceFilter = 'all' | 'builtin' | 'custom'; +type ProjectSummary = { + id: string; + worktree: string; + vcs?: string | null; + name?: string | null; + isDefault?: boolean; + pathStatus?: 'available' | 'missing' | 'unreadable'; + sessionCount?: number; + matchedSessionCount?: number; + lastActivityAt?: number | null; + ownerUserID?: string | null; + canWrite?: boolean; + canDelete?: boolean; + isShared?: boolean; +}; +type ProjectDialogMode = 'create' | 'rename'; +type ProjectSessionGroup = { + id: string; + label: string; + worktree: string; + sessions: Session[]; + sessionCount: number; + pathStatus: 'available' | 'missing' | 'unreadable'; + canWrite: boolean; + canDelete: boolean; + isShared: boolean; +}; +type FolderEntry = { name: string; path: string }; +type FolderBrowserResponse = { + path: string; + parent?: string | null; + roots: FolderEntry[]; + entries: FolderEntry[]; +}; +const MULTI_PROJECT_SESSION_PAGE_SIZE = 6; +const SINGLE_PROJECT_SESSION_PAGE_SIZE = 20; type ChatModelOption = { key: string; providerID: string; @@ -123,82 +162,22 @@ function makeModelKey(providerID: string, modelID: string): string { return `${providerID}::${modelID}`; } -export type SessionGroupKey = 'today' | 'yesterday' | 'thisWeek' | 'lastWeek' | 'earlier'; - -export interface SessionGroup { - key: SessionGroupKey; - labelKey: string; - items: Session[]; +function getPathBasename(path: string): string { + const normalized = path.replace(/[\\/]+$/g, ''); + const parts = normalized.split(/[\\/]/).filter(Boolean); + return parts[parts.length - 1] || path || 'Project'; } -const SESSION_GROUP_DEFAULT_LIMIT: Record = { - today: Infinity, - yesterday: Infinity, - thisWeek: 5, - lastWeek: 5, - earlier: 5, -}; - -const SESSION_GROUP_DEFS: Array> = [ - { key: 'today', labelKey: 'groupToday' }, - { key: 'yesterday', labelKey: 'groupYesterday' }, - { key: 'thisWeek', labelKey: 'groupThisWeek' }, - { key: 'lastWeek', labelKey: 'groupLastWeek' }, - { key: 'earlier', labelKey: 'groupEarlier' }, -]; - -export function groupSessionsByDate( - sessions: Session[], - searchQuery: string, - now = new Date(), -): SessionGroup[] { - const todayStart = new Date(now.getFullYear(), now.getMonth(), now.getDate()).getTime(); - const yesterdayStart = todayStart - 86400000; - const dayOfWeek = now.getDay() === 0 ? 7 : now.getDay(); - const thisWeekStart = todayStart - (dayOfWeek - 1) * 86400000; - const lastWeekStart = thisWeekStart - 7 * 86400000; - - const q = searchQuery.toLowerCase().trim(); - const filtered = q ? sessions.filter(s => s.title.toLowerCase().includes(q)) : sessions; - const buckets: SessionGroup[] = SESSION_GROUP_DEFS.map((group) => ({ - ...group, - items: [], - })); - - for (const session of filtered) { - const ts = session.time?.updated ?? 0; - if (ts >= todayStart) buckets[0].items.push(session); - else if (ts >= yesterdayStart) buckets[1].items.push(session); - else if (ts >= thisWeekStart) buckets[2].items.push(session); - else if (ts >= lastWeekStart) buckets[3].items.push(session); - else buckets[4].items.push(session); - } - - return buckets.filter(group => group.items.length > 0); -} - -export function getVisibleSessionGroupItems({ - group, - expanded, - searching, -}: { - group: SessionGroup; - expanded: boolean; - searching: boolean; -}): { visibleItems: Session[]; hiddenCount: number; limit: number } { - const limit = searching ? Infinity : SESSION_GROUP_DEFAULT_LIMIT[group.key]; - const visibleItems = (searching || expanded || group.items.length <= limit) - ? group.items - : group.items.slice(0, limit); - return { - visibleItems, - hiddenCount: group.items.length - visibleItems.length, - limit, - }; +function getProjectLabel(project?: ProjectSummary, fallbackDirectory?: string): string { + const explicitName = project?.name?.trim(); + if (explicitName) return explicitName; + const directory = project?.worktree || fallbackDirectory || ''; + return getPathBasename(directory); } interface SessionSidebarItemProps { session: Session; + nested?: boolean; selected: boolean; selectMode: boolean; checked: boolean; @@ -218,6 +197,7 @@ interface SessionSidebarItemProps { function SessionSidebarItemInner({ session, + nested = false, selected, selectMode, checked, @@ -237,12 +217,16 @@ function SessionSidebarItemInner({ return (
onSelect(session.id)} - className={`group relative mx-2 mb-1 px-3 py-2.5 rounded-xl border cursor-pointer transition-all duration-150 ${ + className={`group relative mb-1 cursor-pointer border px-3 transition-all duration-150 ${ + nested ? 'ml-7 mr-2 rounded-lg py-2' : 'mx-2 rounded-xl py-2.5' + } ${ !selectMode && selected - ? 'bg-gray-100 border-gray-300 shadow-sm dark:border-zinc-700 dark:bg-zinc-900 dark:shadow-none' + ? 'border-gray-300 bg-gray-100 shadow-sm dark:border-zinc-700 dark:bg-zinc-900 dark:shadow-none' : selectMode && checked ? 'bg-blue-50 border-blue-200 dark:border-blue-500/40 dark:bg-blue-950/30' - : 'border-gray-100 hover:border-gray-200 hover:bg-gray-50 hover:shadow-sm dark:border-transparent dark:hover:border-zinc-800 dark:hover:bg-zinc-900 dark:hover:shadow-none' + : nested + ? 'border-gray-100 hover:border-gray-200 hover:bg-gray-50 dark:border-zinc-800 dark:hover:bg-zinc-900' + : 'border-gray-100 hover:border-gray-200 hover:bg-gray-50 hover:shadow-sm dark:border-transparent dark:hover:border-zinc-800 dark:hover:bg-zinc-900 dark:hover:shadow-none' }`} >
@@ -322,6 +306,7 @@ function SessionSidebarItemInner({ const SessionSidebarItem = memo(SessionSidebarItemInner, (prev, next) => ( prev.session.id === next.session.id && + prev.nested === next.nested && prev.session.title === next.session.title && prev.session.category === next.session.category && prev.session.isShared === next.session.isShared && @@ -338,6 +323,10 @@ const SessionSidebarItem = memo(SessionSidebarItemInner, (prev, next) => ( export default function SessionPage() { const { t, i18n } = useTranslation('session'); + const { user } = useAuth(); + const activeProjectStorageKey = `flocks:sessions:active-project:${user?.id ?? 'anonymous'}`; + const collapsedProjectsStorageKey = `flocks:sessions:collapsed-projects:${user?.id ?? 'anonymous'}`; + const projectsSectionCollapsedStorageKey = `flocks:sessions:projects-section-collapsed:${user?.id ?? 'anonymous'}`; const location = useLocation(); const navigate = useNavigate(); const [searchParams, setSearchParams] = useSearchParams(); @@ -355,6 +344,34 @@ export default function SessionPage() { const [pendingInitialDisplayText, setPendingInitialDisplayText] = useState(null); const [selectMode, setSelectMode] = useState(false); const [checkedIds, setCheckedIds] = useState>(new Set()); + const [projects, setProjects] = useState([]); + const [selectedProjectId, setSelectedProjectId] = useState(null); + const [collapsedProjectIds, setCollapsedProjectIds] = useState>(() => { + try { + const stored = JSON.parse(window.localStorage.getItem(collapsedProjectsStorageKey) ?? '[]'); + return new Set(Array.isArray(stored) ? stored.filter((id): id is string => typeof id === 'string') : []); + } catch { + return new Set(); + } + }); + const [projectsSectionCollapsed, setProjectsSectionCollapsed] = useState(() => { + try { + return window.localStorage.getItem(projectsSectionCollapsedStorageKey) === 'true'; + } catch { + return false; + } + }); + const [projectDialogMode, setProjectDialogMode] = useState(null); + const [openProjectMenuId, setOpenProjectMenuId] = useState(null); + const [projectPendingDelete, setProjectPendingDelete] = useState(null); + const [projectDeleting, setProjectDeleting] = useState(false); + const [editingProjectId, setEditingProjectId] = useState(null); + const [projectNameValue, setProjectNameValue] = useState(''); + const [projectWorktreeValue, setProjectWorktreeValue] = useState(''); + const [projectNameManuallyEdited, setProjectNameManuallyEdited] = useState(false); + const [folderBrowser, setFolderBrowser] = useState(null); + const [folderBrowserLoading, setFolderBrowserLoading] = useState(false); + const [projectSubmitting, setProjectSubmitting] = useState(false); const [batchDeleting, setBatchDeleting] = useState(false); const [openMenuSessionId, setOpenMenuSessionId] = useState(null); const [menuAnchor, setMenuAnchor] = useState<{ top: number; right: number } | null>(null); @@ -369,9 +386,17 @@ export default function SessionPage() { const [selectorTooltip, setSelectorTooltip] = useState(null); const renameInputRef = useRef(null); const renameSubmitInFlightRef = useRef(false); + const projectSubmitInFlightRef = useRef(false); + const folderBrowserRequestIdRef = useRef(0); + const folderBrowserInputPathRef = useRef(null); const sessionUpdateRefetchTimerRef = useRef(null); + const projectListRequestSeqRef = useRef(0); const toast = useToast(); + const sessionProjectIds = useMemo( + () => [TASK_SESSION_GROUP_ID, ...projects.map((project) => project.id)], + [projects], + ); const { sessions, loading: loadingSessions, @@ -380,10 +405,15 @@ export default function SessionPage() { removeSession, removeSessions, addSession, - hasMore: hasMoreSessions, - loadingMore: loadingMoreSessions, + hasMoreByProject = {}, + loadingMoreProjectIds = new Set(), loadMore: loadMoreSessions, - } = useSessions(searchQuery); + } = useSessions(searchQuery, { + projectIds: sessionProjectIds, + pageSize: projects.length >= 1 + ? MULTI_PROJECT_SESSION_PAGE_SIZE + : SINGLE_PROJECT_SESSION_PAGE_SIZE, + }); const { agents, loading: loadingAgents } = useAgents(); const { providers, loading: loadingProviders } = useProviders(); const { @@ -498,44 +528,171 @@ export default function SessionPage() { } = useResolvedDefaultModel(chatModelOptions.length > 0 && !hasPinnedModelOption); const effectiveSupportsVision = selectedModelOption?.supportsVision ?? supportsVision; - const [expandedGroups, setExpandedGroups] = useState>(new Set()); - const toggleGroupExpand = useCallback((key: string) => { - setExpandedGroups(prev => { + const toggleProjectCollapsed = useCallback((projectId: string) => { + setCollapsedProjectIds(prev => { const next = new Set(prev); - next.has(key) ? next.delete(key) : next.add(key); + next.has(projectId) ? next.delete(projectId) : next.add(projectId); return next; }); }, []); - const groupedSessions = useMemo(() => { - return groupSessionsByDate(sessions, searchQuery); - }, [sessions, searchQuery]); - const visibleSessionGroups = useMemo(() => { - const searching = searchQuery.trim().length > 0; - return groupedSessions.map((group) => ({ - ...group, - ...getVisibleSessionGroupItems({ - group, - expanded: expandedGroups.has(group.key), - searching, - }), - expanded: expandedGroups.has(group.key), - searching, - })); - }, [expandedGroups, groupedSessions, searchQuery]); + useEffect(() => { + try { + window.localStorage.setItem( + collapsedProjectsStorageKey, + JSON.stringify(Array.from(collapsedProjectIds)), + ); + } catch { + // Project collapse persistence must never block the session manager. + } + }, [collapsedProjectIds, collapsedProjectsStorageKey]); + + useEffect(() => { + try { + window.localStorage.setItem( + projectsSectionCollapsedStorageKey, + String(projectsSectionCollapsed), + ); + } catch { + // Section collapse persistence must never block the session manager. + } + }, [projectsSectionCollapsed, projectsSectionCollapsedStorageKey]); + + const projectSessionGroups = useMemo(() => { + const registeredIds = new Set(projects.map((project) => project.id)); + const sessionsByProject = new Map(); + sessions.forEach((session) => { + const responseProjectId = session.effectiveProjectID || session.projectID; + if (!registeredIds.has(responseProjectId)) return; + const groupSessions = sessionsByProject.get(responseProjectId) ?? []; + groupSessions.push(session); + sessionsByProject.set(responseProjectId, groupSessions); + }); + + return projects.map((project) => { + const projectSessions = sessionsByProject.get(project.id) ?? []; + return { + id: project.id, + label: getProjectLabel(project), + worktree: project.worktree, + sessions: projectSessions, + sessionCount: searchQuery.trim() + ? project.matchedSessionCount ?? projectSessions.length + : project.sessionCount ?? projectSessions.length, + pathStatus: project.pathStatus ?? 'available', + canWrite: project.canWrite ?? true, + canDelete: project.canDelete ?? true, + isShared: project.isShared ?? false, + }; + }) + .filter((group) => { + if (!searchQuery.trim()) return true; + const project = projects.find((item) => item.id === group.id); + return (project?.matchedSessionCount ?? group.sessions.length) > 0 || group.id === selectedProjectId; + }) + .sort((a, b) => { + const aLatest = projects.find((project) => project.id === a.id)?.lastActivityAt ?? a.sessions[0]?.time?.updated ?? 0; + const bLatest = projects.find((project) => project.id === b.id)?.lastActivityAt ?? b.sessions[0]?.time?.updated ?? 0; + if (aLatest !== bLatest) return bLatest - aLatest; + return a.label.localeCompare(b.label); + }); + }, [projects, searchQuery, selectedProjectId, sessions]); + + const managedProjectSessionGroups = projectSessionGroups; + const taskSessionGroup = useMemo( + () => { + const registeredIds = new Set(projects.map((project) => project.id)); + const taskSessions = sessions.filter((session) => ( + !registeredIds.has(session.effectiveProjectID || session.projectID) + )); + return { + id: TASK_SESSION_GROUP_ID, + label: t('tasksSection'), + worktree: '', + sessions: taskSessions, + sessionCount: taskSessions.length, + pathStatus: 'available' as const, + }; + }, + [projects, sessions, t], + ); + const taskGroupCollapsed = collapsedProjectIds.has(TASK_SESSION_GROUP_ID); + const taskGroupSelected = selectedProjectId === TASK_SESSION_GROUP_ID; + const hasMoreTaskSessions = hasMoreByProject[TASK_SESSION_GROUP_ID] ?? false; + + const selectedProjectIDForCreate = selectedProjectId && selectedProjectId !== TASK_SESSION_GROUP_ID + ? selectedProjectId + : null; + + useEffect(() => { + const selectableProjectIds = new Set([ + TASK_SESSION_GROUP_ID, + ...projectSessionGroups.map((group) => group.id), + ]); + if (selectedProjectId && selectableProjectIds.has(selectedProjectId)) { + return; + } + + let storedProjectId: string | null = null; + try { + storedProjectId = window.localStorage.getItem(activeProjectStorageKey); + } catch { + storedProjectId = null; + } + const fallbackProjectId = storedProjectId && selectableProjectIds.has(storedProjectId) + ? storedProjectId + : TASK_SESSION_GROUP_ID; + setSelectedProjectId(fallbackProjectId); + }, [activeProjectStorageKey, projectSessionGroups, selectedProjectId]); + + useEffect(() => { + if ( + !selectedProjectId + || (selectedProjectId !== TASK_SESSION_GROUP_ID + && !projects.some((project) => project.id === selectedProjectId)) + ) return; + try { + window.localStorage.setItem(activeProjectStorageKey, selectedProjectId); + } catch { + // Active project persistence must never block the session manager. + } + }, [activeProjectStorageKey, projects, selectedProjectId]); // Handle SSE events for session-level updates (title changes, etc.) const handleChatError = useCallback((msg: string) => { toast.error(t('chat.error', 'Error'), msg); }, [toast, t]); + const fetchProjects = useCallback(async (ensureProject?: ProjectSummary, query = '') => { + const requestSeq = ++projectListRequestSeqRef.current; + const listResult = await client.get('/api/project', { + params: { search: query.trim() || undefined }, + }); + if (requestSeq !== projectListRequestSeqRef.current) return; + const nextProjects = Array.isArray(listResult.data) + ? listResult.data.filter((project: ProjectSummary) => ( + project.id !== TASK_SESSION_GROUP_ID && !project.isDefault + )) + : []; + setProjects((currentProjects) => { + if (!ensureProject?.id || nextProjects.some((project) => project.id === ensureProject.id)) { + return nextProjects; + } + const currentProject = currentProjects.find((project) => project.id === ensureProject.id); + return [{ ...currentProject, ...ensureProject }, ...nextProjects]; + }); + }, []); + const scheduleSessionListRefetch = useCallback(() => { if (sessionUpdateRefetchTimerRef.current !== null) return; sessionUpdateRefetchTimerRef.current = window.setTimeout(() => { sessionUpdateRefetchTimerRef.current = null; - void refetchSessions(); + void Promise.all([ + refetchSessions(), + fetchProjects(undefined, searchQuery), + ]); }, SESSION_UPDATE_REFETCH_DEBOUNCE_MS); - }, [refetchSessions]); + }, [fetchProjects, refetchSessions, searchQuery]); useEffect(() => () => { if (sessionUpdateRefetchTimerRef.current !== null) { @@ -545,6 +702,15 @@ export default function SessionPage() { }, []); const handleSSEEvent = useCallback((event: SSEChatEvent) => { + if (event.type === 'session.notice' && event.properties?.kind === 'directory-fallback') { + toast.warning( + t('projectDirectoryFallbackTitle'), + t('projectDirectoryFallbackDescription', { + path: event.properties.fallbackDirectory, + }), + ); + return; + } if (event.type === 'session.updated' && event.properties?.id) { if (event.properties?.title) { // Instant local title update so the sidebar reflects the change immediately. @@ -555,7 +721,18 @@ export default function SessionPage() { // those bursts don't turn into a request/re-render storm. scheduleSessionListRefetch(); } - }, [scheduleSessionListRefetch, updateSessionTitle]); + }, [scheduleSessionListRefetch, t, toast, updateSessionTitle]); + + useEffect(() => { + void fetchProjects(undefined, searchQuery); + }, [fetchProjects, searchQuery]); + + useEffect(() => { + if (!openProjectMenuId) return; + const closeMenu = () => setOpenProjectMenuId(null); + window.addEventListener('click', closeMenu); + return () => window.removeEventListener('click', closeMenu); + }, [openProjectMenuId]); // Keep the selected session in sync with URL query params (e.g. onboarding // or other in-app navigation to `/sessions?session=...`). Clear the params @@ -722,13 +899,25 @@ export default function SessionPage() { setRenameValue(''); }, [selectMode]); - const handleCreateSession = useCallback(async () => { + const handleCreateSession = useCallback(async (projectIdOverride?: string) => { if (creating) return; + const targetGroupId = projectIdOverride ?? selectedProjectId ?? TASK_SESSION_GROUP_ID; + const projectID = targetGroupId === TASK_SESSION_GROUP_ID ? null : targetGroupId; setCreating(true); try { - const response = await client.post('/api/session', { title: 'New Session' }); + const response = await client.post('/api/session', { + title: 'New Session', + ...(projectID ? { projectID } : {}), + }); addSession(response.data); + await fetchProjects(undefined, searchQuery); setSelectedSessionFallback(response.data); + setSelectedProjectId(targetGroupId); + setCollapsedProjectIds(prev => { + const next = new Set(prev); + next.delete(targetGroupId); + return next; + }); setSelectedAgent('rex'); setSelectedModelKey(null); setSelectedSessionId(response.data.id); @@ -737,7 +926,11 @@ export default function SessionPage() { } finally { setCreating(false); } - }, [creating, addSession, toast, t]); + }, [creating, selectedProjectId, addSession, fetchProjects, searchQuery, toast, t]); + + const handleCreateSessionInProject = useCallback((projectId: string) => { + void handleCreateSession(projectId); + }, [handleCreateSession]); const handleSelectModel = useCallback(async (option: ChatModelOption) => { setSelectedModelKey(option.key); @@ -764,10 +957,14 @@ export default function SessionPage() { options?: PromptDisplayOptions, ) => { try { - const response = await client.post('/api/session', { title: 'New Session' }); + const response = await client.post('/api/session', { + title: 'New Session', + ...(selectedProjectIDForCreate ? { projectID: selectedProjectIDForCreate } : {}), + }); const newSessionId = response.data.id; addSession(response.data); + await fetchProjects(undefined, searchQuery); setSelectedSessionFallback(response.data); setSelectedModelKey(null); setSelectedSessionId(newSessionId); @@ -785,7 +982,7 @@ export default function SessionPage() { } catch (err: any) { toast.error(t('createFailed'), err.message); } - }, [addSession, selectedAgent, toast, t]); + }, [addSession, fetchProjects, searchQuery, selectedAgent, selectedProjectIDForCreate, toast, t]); const handleSuiteInstallProgress = useCallback((progress: HubInstallProgressEvent) => { setSuiteInstallProgress(current => applySuiteInstallProgressEvent(current, progress)); @@ -861,10 +1058,11 @@ export default function SessionPage() { // No need to refetchSessions — removeSession already keeps the list accurate. if (selectedSessionId === sessionId) setSelectedSessionId(null); removeSession(sessionId); + await fetchProjects(undefined, searchQuery); } catch (err: any) { toast.error(t('deleteFailed'), err.message); } - }, [selectedSessionId, removeSession, toast, t]); + }, [fetchProjects, removeSession, searchQuery, selectedSessionId, toast, t]); const handleStartRename = useCallback((sessionId: string, currentTitle: string) => { setOpenMenuSessionId(null); @@ -908,6 +1106,212 @@ export default function SessionPage() { } }, [renameValue, sessions, t, toast, updateSessionTitle]); + const handleOpenCreateProject = useCallback(() => { + setProjectDialogMode('create'); + setEditingProjectId(null); + setProjectWorktreeValue(''); + setProjectNameValue(''); + setProjectNameManuallyEdited(false); + setFolderBrowser(null); + }, []); + + const handleOpenRenameProject = useCallback((project: ProjectSessionGroup) => { + const persistedName = projects.find((item) => item.id === project.id)?.name?.trim(); + setOpenProjectMenuId(null); + setProjectDialogMode('rename'); + setEditingProjectId(project.id); + setProjectNameValue(persistedName && persistedName !== getPathBasename(project.worktree) + ? persistedName + : project.label); + setProjectWorktreeValue(project.worktree); + setProjectNameManuallyEdited(true); + setFolderBrowser(null); + }, [projects]); + + const handleCloseProjectDialog = useCallback(() => { + if (projectSubmitting) return; + setProjectDialogMode(null); + setEditingProjectId(null); + setProjectNameValue(''); + setProjectWorktreeValue(''); + setProjectNameManuallyEdited(false); + setFolderBrowser(null); + }, [projectSubmitting]); + + const loadFolderBrowser = useCallback(async ( + path?: string, + options?: { preserveInput?: boolean; silent?: boolean }, + ) => { + const requestId = ++folderBrowserRequestIdRef.current; + setFolderBrowserLoading(true); + try { + const response = await client.get('/api/project/folders', { + params: { path: path?.trim() || undefined }, + }); + if (requestId !== folderBrowserRequestIdRef.current) return; + const browser = response.data as FolderBrowserResponse; + setFolderBrowser(browser); + if (options?.preserveInput) { + folderBrowserInputPathRef.current = path?.trim() || browser.path; + } else { + folderBrowserInputPathRef.current = browser.path; + setProjectWorktreeValue(browser.path); + } + } catch (err: any) { + if (requestId === folderBrowserRequestIdRef.current && !options?.silent) { + toast.error(t('projectDialog.folderBrowseFailed'), err.message); + } + } finally { + if (requestId === folderBrowserRequestIdRef.current) { + setFolderBrowserLoading(false); + } + } + }, [t, toast]); + + useEffect(() => { + if (!folderBrowser) return undefined; + const path = projectWorktreeValue.trim(); + if ( + !path + || path === folderBrowser.path + || path === folderBrowserInputPathRef.current + ) return undefined; + + const timer = window.setTimeout(() => { + void loadFolderBrowser(path, { preserveInput: true, silent: true }); + }, 300); + return () => window.clearTimeout(timer); + }, [folderBrowser, loadFolderBrowser, projectWorktreeValue]); + + const handleSelectProjectFolder = useCallback((path: string) => { + setProjectWorktreeValue(path); + if (!projectNameManuallyEdited) { + setProjectNameValue(getPathBasename(path)); + } + setFolderBrowser(null); + }, [projectNameManuallyEdited]); + + const handleCopyProjectPath = useCallback(async (project: ProjectSessionGroup) => { + setOpenProjectMenuId(null); + try { + await navigator.clipboard.writeText(project.worktree); + toast.success(t('projectDialog.pathCopied')); + } catch (err: any) { + toast.error(t('projectDialog.copyPathFailed'), err.message); + } + }, [t, toast]); + + const handleShareProject = useCallback(async ( + project: ProjectSessionGroup, + nextShared: boolean, + ) => { + setOpenProjectMenuId(null); + try { + const action = nextShared ? 'share-local' : 'unshare-local'; + await client.post(`/api/project/${project.id}/${action}`); + toast.success(t(nextShared ? 'shareEnabled' : 'shareDisabled')); + await Promise.all([ + fetchProjects(undefined, searchQuery), + refetchSessions(), + ]); + } catch (err: any) { + toast.error(t('shareUpdateFailed'), err.message); + } + }, [fetchProjects, refetchSessions, searchQuery, t, toast]); + + const handleOpenDeleteProject = useCallback((project: ProjectSessionGroup) => { + setOpenProjectMenuId(null); + setProjectPendingDelete(project); + }, []); + + const handleDeleteProject = useCallback(async () => { + if (!projectPendingDelete || projectDeleting) return; + + setProjectDeleting(true); + try { + await client.delete(`/api/project/${projectPendingDelete.id}`); + setProjects((current) => current.filter((project) => project.id !== projectPendingDelete.id)); + setCollapsedProjectIds((current) => { + const next = new Set(current); + next.delete(projectPendingDelete.id); + return next; + }); + setSelectedProjectId((current) => ( + current === projectPendingDelete.id ? TASK_SESSION_GROUP_ID : current + )); + setProjectPendingDelete(null); + await Promise.all([ + fetchProjects(undefined, searchQuery), + refetchSessions(), + ]); + toast.success(t('projectDialog.deleteSuccess')); + } catch (err: any) { + toast.error(t('projectDialog.deleteFailed'), err.message); + } finally { + setProjectDeleting(false); + } + }, [fetchProjects, projectDeleting, projectPendingDelete, refetchSessions, searchQuery, t, toast]); + + const handleSubmitProject = useCallback(async () => { + if (!projectDialogMode || projectSubmitInFlightRef.current) return; + const nextWorktree = projectWorktreeValue.trim() || folderBrowser?.path.trim() || ''; + const nextName = projectNameValue.trim() || ( + projectDialogMode === 'create' && !projectNameManuallyEdited && nextWorktree + ? getPathBasename(nextWorktree) + : '' + ); + if (!nextName) { + toast.error(t('projectDialog.saveFailed'), t('projectDialog.nameEmpty')); + return; + } + if (projectDialogMode === 'create' && !nextWorktree) { + toast.error(t('projectDialog.saveFailed'), t('projectDialog.folderEmpty')); + return; + } + + projectSubmitInFlightRef.current = true; + setProjectSubmitting(true); + try { + if (projectDialogMode === 'create') { + const response = await client.post('/api/project', { + name: nextName, + worktree: nextWorktree, + }); + const createdProject = response.data as ProjectSummary; + setProjects(prev => ( + prev.some((project) => project.id === createdProject.id) + ? prev.map((project) => (project.id === createdProject.id ? createdProject : project)) + : [createdProject, ...prev] + )); + setSelectedProjectId(createdProject.id); + setCollapsedProjectIds(prev => { + const next = new Set(prev); + next.delete(createdProject.id); + return next; + }); + await fetchProjects(createdProject); + } else if (editingProjectId) { + const response = await client.patch(`/api/project/${editingProjectId}`, { name: nextName }); + await fetchProjects(response.data as ProjectSummary); + } + setProjectDialogMode(null); + setEditingProjectId(null); + setProjectNameValue(''); + setProjectWorktreeValue(''); + setProjectNameManuallyEdited(false); + setFolderBrowser(null); + } catch (err: any) { + const detail = err?.response?.data?.detail; + toast.error( + t('projectDialog.saveFailed'), + typeof detail === 'string' ? detail : err?.message, + ); + } finally { + projectSubmitInFlightRef.current = false; + setProjectSubmitting(false); + } + }, [editingProjectId, fetchProjects, folderBrowser?.path, projectDialogMode, projectNameManuallyEdited, projectNameValue, projectWorktreeValue, t, toast]); + const handleDownloadSession = useCallback(async (sessionId: string, title: string) => { setOpenMenuSessionId(null); setDownloadingSessionId(sessionId); @@ -1011,6 +1415,7 @@ export default function SessionPage() { })); if (succeeded.length > 0) { removeSessions(succeeded); + await fetchProjects(undefined, searchQuery); if (selectedSessionId && succeeded.includes(selectedSessionId)) { setSelectedSessionId(null); } @@ -1023,7 +1428,101 @@ export default function SessionPage() { setSelectMode(false); } setBatchDeleting(false); - }, [checkedIds, batchDeleting, removeSessions, selectedSessionId, toast, t]); + }, [batchDeleting, checkedIds, fetchProjects, removeSessions, searchQuery, selectedSessionId, toast, t]); + + const renderSessionListItem = (session: Session) => ( +
selectMode ? handleToggleCheck(session.id) : setSelectedSessionId(session.id)} + className={`group relative mx-2 mb-1 px-3 py-2.5 rounded-xl border cursor-pointer transition-all duration-150 ${ + !selectMode && selectedSessionId === session.id + ? 'bg-gray-100 border-gray-300 shadow-sm dark:border-zinc-700 dark:bg-zinc-900 dark:shadow-none' + : selectMode && checkedIds.has(session.id) + ? 'bg-blue-50 border-blue-200 dark:border-blue-500/40 dark:bg-blue-950/30' + : 'border-gray-100 hover:border-gray-200 hover:bg-gray-50 hover:shadow-sm dark:border-transparent dark:hover:border-zinc-800 dark:hover:bg-zinc-900 dark:hover:shadow-none' + }`} + > +
+ {selectMode && ( + handleToggleCheck(session.id)} + onClick={(e) => e.stopPropagation()} + className="flex-shrink-0 w-3.5 h-3.5 accent-blue-500 cursor-pointer rounded" + /> + )} + {session.category === 'workflow' && ( + + + + )} + {session.category === 'entity-config' && ( + + + + )} + {renamingSessionId === session.id ? ( + setRenameValue(e.target.value)} + onClick={(e) => e.stopPropagation()} + onBlur={() => void handleSubmitRename(session.id)} + onKeyDown={(e) => { + if (e.key === 'Enter') { e.preventDefault(); void handleSubmitRename(session.id); } + if (e.key === 'Escape') { e.preventDefault(); handleCancelRename(); } + }} + placeholder={t('renamePlaceholder')} + disabled={renameSubmitting} + className="w-full min-w-0 rounded border border-blue-300 bg-white px-1.5 py-0.5 text-sm text-gray-900 outline-none focus:border-blue-400 dark:border-blue-500/50 dark:bg-zinc-950 dark:text-zinc-100" + aria-label={t('rename')} + data-session-rename-input + /> + ) : ( +

+ {session.title} + {session.isShared && ( + + {t('sharedTag')} + + )} +

+ )} +
+ {session.time?.updated && renamingSessionId !== session.id && ( +

+ {formatSessionDate(session.time.updated)} +

+ )} + + {!selectMode && ( +
+ +
+ )} +
+ ); if (loadingSessions) { return ( @@ -1048,7 +1547,7 @@ export default function SessionPage() { ? : } - )} - {!searching && expanded && items.length > limit && ( +
+ {!projectsSectionCollapsed && ( +
+ {managedProjectSessionGroups.map((group) => { + const collapsed = collapsedProjectIds.has(group.id); + const isSelectedProject = selectedProjectId === group.id; + const hasMoreProjectSessions = ( + group.sessions.length < group.sessionCount + || hasMoreByProject[group.id] + ); + const persistedProject = projects.find((project) => project.id === group.id); + return ( +
+
+ + + {persistedProject && ( + + )} + + {group.sessionCount} + +
+ {persistedProject && openProjectMenuId === group.id && ( +
event.stopPropagation()} + > + + {group.canWrite && ( + + )} + {group.canWrite && ( + + )} + {group.canDelete && ( +
+ )} + {group.canDelete && ( + + )} +
+ )} + {!collapsed && ( +
+ {group.sessions.length > 0 ? ( + group.sessions.map((session) => ( + + )) + ) : ( +
+ {t('noProjectSessions')} +
+ )} + {hasMoreProjectSessions && ( +
+ +
+ )} +
+ )} +
+ ); + })} +
)} -
- )) - )} - {hasMoreSessions && ( -
- -
- )} + + {taskSessionGroup && ( +
+
+ + + + + {taskSessionGroup.sessionCount} + +
+ {!taskGroupCollapsed && ( +
+ {taskSessionGroup.sessions.map((session) => ( + + ))} + {hasMoreTaskSessions && ( +
+ +
+ )} +
+ )} +
+ )} +
{/* Bottom：批量操作栏 / 批量选择入口 */} @@ -1457,6 +2181,196 @@ export default function SessionPage() { )}
+ {projectDialogMode && ( +
+
+
+

+ {projectDialogMode === 'create' ? t('projectDialog.createTitle') : t('projectDialog.renameTitle')} +

+
+
+ + { + setProjectNameValue(event.target.value); + setProjectNameManuallyEdited(true); + }} + onKeyDown={(event) => { + if (event.key === 'Enter') { + event.preventDefault(); + if ( + event.nativeEvent.isComposing + || (projectDialogMode === 'create' + && !projectWorktreeValue.trim() + && !folderBrowser?.path.trim()) + ) { + return; + } + void handleSubmitProject(); + } + if (event.key === 'Escape') { + event.preventDefault(); + handleCloseProjectDialog(); + } + }} + disabled={projectSubmitting} + placeholder={t('projectDialog.namePlaceholder')} + className="w-full rounded-lg border border-zinc-200 bg-white px-3 py-2 text-sm text-zinc-900 outline-none transition-colors placeholder:text-zinc-400 focus:border-blue-400 dark:border-zinc-800 dark:bg-zinc-900 dark:text-zinc-100 dark:placeholder:text-zinc-600 dark:focus:border-blue-500" + autoFocus + /> + {projectDialogMode === 'create' && ( + <> + +
+ { + folderBrowserRequestIdRef.current += 1; + folderBrowserInputPathRef.current = null; + setFolderBrowserLoading(false); + setProjectWorktreeValue(event.target.value); + if (!projectNameManuallyEdited) { + setProjectNameValue(getPathBasename(event.target.value)); + } + }} + disabled={projectSubmitting} + placeholder={t('projectDialog.folderPlaceholder')} + className="min-w-0 flex-1 rounded-lg border border-zinc-200 bg-white px-3 py-2 font-mono text-xs text-zinc-900 outline-none transition-colors placeholder:text-zinc-400 focus:border-blue-400 dark:border-zinc-800 dark:bg-zinc-900 dark:text-zinc-100 dark:placeholder:text-zinc-600 dark:focus:border-blue-500" + /> + +
+ {folderBrowser && ( +
+
+ + + {folderBrowser.path} + + +
+
+ {folderBrowser.roots.map((root) => ( + + ))} +
+
+ {folderBrowser.entries.length > 0 ? folderBrowser.entries.map((entry) => ( + + )) : ( +
+ {t('projectDialog.noSubfolders')} +
+ )} +
+
+ )} + + )} +
+
+ + +
+
+
+ )} + + {projectPendingDelete && ( +
+
+
+

+ {t('projectDialog.deleteTitle')} +

+
+
+ {t('projectDialog.deleteDescription', { project: projectPendingDelete.label })} +
+
+ + +
+
+
+ )} + {selectorTooltip && (
e.stopPropagation()} > -
+
+ ); +} + function PreferencesPanel() { const { t, i18n } = useTranslation('nav'); const { theme, setTheme } = useContext(ThemeContext); @@ -168,26 +201,58 @@ function PreferencesPanel() { uploadProductFavicon, resetProductFavicon, } = useProductName(); - const toast = useToast(); + const { error: showToastError, success: showToastSuccess } = useToast(); const language = i18n.language?.toLowerCase().startsWith('zh') ? 'zh-CN' : 'en-US'; const faviconInputRef = useRef(null); const [displayNameDraft, setDisplayNameDraft] = useState(configuredDisplayName ?? ''); const [savingDisplayName, setSavingDisplayName] = useState(false); const [savingFavicon, setSavingFavicon] = useState(false); + const [toolFailureAutoDisable, setToolFailureAutoDisable] = useState(true); + const [loadingToolFailure, setLoadingToolFailure] = useState(true); + const [savingToolFailure, setSavingToolFailure] = useState(false); const normalizedDisplayName = displayNameDraft.trim(); const displayNameChanged = normalizedDisplayName !== (configuredDisplayName ?? ''); + const toolFailureSettingLoadFailedMessage = t('toolFailureSettingLoadFailed'); useEffect(() => { setDisplayNameDraft(configuredDisplayName ?? ''); }, [configuredDisplayName]); + useEffect(() => { + let cancelled = false; + setLoadingToolFailure(true); + toolFailureConfigApi.get() + .then((config) => { + if (!cancelled) { + setToolFailureAutoDisable(config.disableOnRepeatedFailure); + } + }) + .catch((err: any) => { + if (!cancelled) { + showToastError( + toolFailureSettingLoadFailedMessage, + err?.response?.data?.detail || err?.message, + ); + } + }) + .finally(() => { + if (!cancelled) { + setLoadingToolFailure(false); + } + }); + + return () => { + cancelled = true; + }; + }, [showToastError, toolFailureSettingLoadFailedMessage]); + const handleSaveDisplayName = async () => { setSavingDisplayName(true); try { await updateProductName(normalizedDisplayName || null); - toast.success(t('displayNameSaved')); + showToastSuccess(t('displayNameSaved')); } catch (err: any) { - toast.error(t('displayNameSaveFailed'), err?.response?.data?.detail || err?.message); + showToastError(t('displayNameSaveFailed'), err?.response?.data?.detail || err?.message); } finally { setSavingDisplayName(false); } @@ -197,9 +262,9 @@ function PreferencesPanel() { setSavingDisplayName(true); try { await updateProductName(null); - toast.success(t('displayNameSaved')); + showToastSuccess(t('displayNameSaved')); } catch (err: any) { - toast.error(t('displayNameSaveFailed'), err?.response?.data?.detail || err?.message); + showToastError(t('displayNameSaveFailed'), err?.response?.data?.detail || err?.message); } finally { setSavingDisplayName(false); } @@ -213,9 +278,9 @@ function PreferencesPanel() { setSavingFavicon(true); try { await uploadProductFavicon(file); - toast.success(t('faviconSaved')); + showToastSuccess(t('faviconSaved')); } catch (err: any) { - toast.error(t('faviconSaveFailed'), err?.response?.data?.detail || err?.message); + showToastError(t('faviconSaveFailed'), err?.response?.data?.detail || err?.message); } finally { setSavingFavicon(false); } @@ -225,14 +290,31 @@ function PreferencesPanel() { setSavingFavicon(true); try { await resetProductFavicon(); - toast.success(t('faviconSaved')); + showToastSuccess(t('faviconSaved')); } catch (err: any) { - toast.error(t('faviconSaveFailed'), err?.response?.data?.detail || err?.message); + showToastError(t('faviconSaveFailed'), err?.response?.data?.detail || err?.message); } finally { setSavingFavicon(false); } }; + const handleToolFailureAutoDisableChange = async () => { + const nextValue = !toolFailureAutoDisable; + setSavingToolFailure(true); + try { + const config = await toolFailureConfigApi.update(nextValue); + setToolFailureAutoDisable(config.disableOnRepeatedFailure); + showToastSuccess(t('toolFailureSettingSaved')); + } catch (err: any) { + showToastError( + t('toolFailureSettingSaveFailed'), + err?.response?.data?.detail || err?.message, + ); + } finally { + setSavingToolFailure(false); + } + }; + return (
@@ -371,6 +453,19 @@ function PreferencesPanel() {
+ + + void handleToolFailureAutoDisableChange()} + /> +
); @@ -395,7 +490,7 @@ export default function SettingsPage() { const navigate = useNavigate(); const { t } = useTranslation('nav'); const { user } = useAuth(); - const { productName } = useProductName(); + const { proProductName } = useProductName(); const isAdmin = user?.role === 'admin'; const sectionId = params.sectionId; const [flocksproCapabilityReady, setFlocksproCapabilityReady] = useState(false); @@ -455,11 +550,11 @@ export default function SettingsPage() { { id: 'account', name: t('accountManagement'), icon: UserCog }, { id: 'system-logs', name: t('systemLog'), icon: ScrollText }, { id: 'audit-logs', name: t('auditLogs'), icon: ShieldCheck, adminOnly: true, requiresFlockspro: true }, - { id: 'flockspro', name: productName, icon: ArrowUpCircle, adminOnly: true }, + { id: 'flockspro', name: proProductName, icon: ArrowUpCircle, adminOnly: true }, ], }, ], - [productName, t], + [proProductName, t], ); const visibleGroups = groups diff --git a/webui/src/pages/Task/ServicesSection.test.tsx b/webui/src/pages/Task/ServicesSection.test.tsx index 78c9e21d6..634f21395 100644 --- a/webui/src/pages/Task/ServicesSection.test.tsx +++ b/webui/src/pages/Task/ServicesSection.test.tsx @@ -11,7 +11,8 @@ const mocks = vi.hoisted(() => ({ toastError: vi.fn(), })); -vi.mock('@/api/workflow', () => ({ +vi.mock('@/api/workflow', async (importOriginal) => ({ + ...(await importOriginal()), workflowAPI: { listServices: mocks.listServices, publish: mocks.publish, @@ -113,7 +114,7 @@ describe('ServicesSection', () => { await user.click(restartButton); await waitFor(() => { - expect(mocks.publish).toHaveBeenCalledWith('wf-1', { driver: 'docker' }); + expect(mocks.publish).toHaveBeenCalledWith('wf-1', { driver: 'docker', port: 19000 }); }); expect(screen.getByText('服务驱动')).toBeInTheDocument(); expect(screen.getByText('Docker')).toBeInTheDocument(); diff --git a/webui/src/pages/Task/ServicesSection.tsx b/webui/src/pages/Task/ServicesSection.tsx index 18a89a25e..32094b658 100644 --- a/webui/src/pages/Task/ServicesSection.tsx +++ b/webui/src/pages/Task/ServicesSection.tsx @@ -1,7 +1,7 @@ import { useState, useEffect, useCallback } from 'react'; import { useTranslation } from 'react-i18next'; import { Globe, StopCircle, Loader2, RefreshCw, ExternalLink } from 'lucide-react'; -import { workflowAPI, WorkflowService } from '@/api/workflow'; +import { resolveWorkflowServicePort, workflowAPI, WorkflowService } from '@/api/workflow'; import { useToast } from '@/components/common/Toast'; import EmptyState from '@/components/common/EmptyState'; import CopyButton from '@/components/common/CopyButton'; @@ -164,7 +164,7 @@ export default function ServicesSection() { setRestartingIds(prev => new Set(prev).add(service.workflowId)); try { const publishRequest = service.driver === 'local' || service.driver === 'docker' - ? { driver: service.driver } + ? { driver: service.driver, port: resolveWorkflowServicePort(service) } : undefined; await workflowAPI.publish(service.workflowId, publishRequest); await fetchServices(); diff --git a/webui/src/pages/Task/TaskSheet.test.tsx b/webui/src/pages/Task/TaskSheet.test.tsx index 958073eec..268290f14 100644 --- a/webui/src/pages/Task/TaskSheet.test.tsx +++ b/webui/src/pages/Task/TaskSheet.test.tsx @@ -100,7 +100,7 @@ vi.mock('@/api/agent', () => ({ vi.mock('@/api/workflow', () => ({ workflowAPI: { - list: mocks.workflowList, + listSummaries: mocks.workflowList, }, })); diff --git a/webui/src/pages/Task/TaskSheet.tsx b/webui/src/pages/Task/TaskSheet.tsx index 43b523827..9a8116350 100644 --- a/webui/src/pages/Task/TaskSheet.tsx +++ b/webui/src/pages/Task/TaskSheet.tsx @@ -21,7 +21,7 @@ import PillGroup from '@/components/common/PillGroup'; import { useTaskExecutionsByScheduler } from '@/hooks/useTasks'; import { describeCron, CRON_PRESETS, formatDuration, formatTime } from './helpers'; import { agentAPI, Agent } from '@/api/agent'; -import { workflowAPI, Workflow } from '@/api/workflow'; +import { workflowAPI, WorkflowSummary } from '@/api/workflow'; import { getAgentDisplayDescription } from '@/utils/agentDisplay'; import { getWorkflowDisplayName } from '@/utils/workflowDisplay'; import { StatusBadge } from './components'; @@ -106,7 +106,7 @@ export default function TaskSheet({ task, defaultScheduleKind = 'recurring', onC const [loading, setLoading] = useState(false); const [agents, setAgents] = useState([]); - const [workflows, setWorkflows] = useState([]); + const [workflows, setWorkflows] = useState([]); const guideGroups = useMemo(() => { const scope = isEdit ? 'taskSheet.edit' : 'taskSheet.create'; const name = formData.title || task?.title || t('taskSheet.entityType'); @@ -132,7 +132,7 @@ export default function TaskSheet({ task, defaultScheduleKind = 'recurring', onC useEffect(() => { agentAPI.list().then((res) => setAgents(res.data.filter((a) => !a.hidden))).catch(() => {}); - workflowAPI.list({ status: 'active' }).then((res) => setWorkflows(res.data)).catch(() => {}); + workflowAPI.listSummaries({ status: 'active' }).then((res) => setWorkflows(res.data)).catch(() => {}); }, []); const isImmediate = formData.scheduleKind === 'immediate'; @@ -332,7 +332,7 @@ interface TaskFormContentProps { isRunOnce: boolean; effectiveCron: string; agents: Agent[]; - workflows: Workflow[]; + workflows: WorkflowSummary[]; } function TaskFormContent({ diff --git a/webui/src/pages/Tool/ToolDetailDrawer.test.tsx b/webui/src/pages/Tool/ToolDetailDrawer.test.tsx index 518e6998d..7da6bdc61 100644 --- a/webui/src/pages/Tool/ToolDetailDrawer.test.tsx +++ b/webui/src/pages/Tool/ToolDetailDrawer.test.tsx @@ -1,6 +1,6 @@ import React from 'react'; import { describe, it, expect, vi } from 'vitest'; -import { render, screen } from '@testing-library/react'; +import { fireEvent, render, screen, waitFor } from '@testing-library/react'; import { ToolDetailDrawer } from './index'; vi.mock('react-i18next', () => ({ @@ -25,6 +25,7 @@ vi.mock('react-i18next', () => ({ 'toolDetail.no': '否', 'toolDetail.close': '关闭', 'toolDetail.testTool': '测试工具', + 'toolDetail.runTest': '执行测试', 'toolDetail.enabled': '启用', 'toolDetail.disabled': '禁用', }; @@ -49,6 +50,46 @@ vi.mock('@/api/tool', () => ({ })); describe('ToolDetailDrawer', () => { + const enabledTool = { + name: 'failing_custom_tool', + description: 'Always fails', + source: 'plugin_py', + category: 'custom', + enabled: true, + parameters: [], + } as any; + + const drawerProps = { + tool: enabledTool, + testParams: '{}', + testResult: null, + testing: false, + onClose: vi.fn(), + onTestParamsChange: vi.fn(), + onTest: vi.fn(), + }; + + it('disables test execution as soon as the result reports auto-disable', async () => { + const { rerender } = render(); + fireEvent.click(screen.getByRole('button', { name: '测试' })); + expect(screen.getByRole('button', { name: '执行测试' })).toBeEnabled(); + + rerender( + , + ); + + await waitFor(() => { + expect(screen.getByRole('button', { name: '执行测试' })).toBeDisabled(); + }); + }); + it('wraps long descriptions and parameter text without clipping', () => { const longDescription = 'OneSEC DNS grouped tool. dns_search_blocked_queries_by_super_long_keyword_with_no_breaks_and_more_details_to_force_wrapping'; const longParamDescription = 'DNS 分组动作名 dns_search_blocked_queries_by_super_long_keyword_with_no_breaks_and_more_details_to_force_wrapping'; diff --git a/webui/src/pages/Tool/ToolPageAutoDisable.test.tsx b/webui/src/pages/Tool/ToolPageAutoDisable.test.tsx new file mode 100644 index 000000000..aad856925 --- /dev/null +++ b/webui/src/pages/Tool/ToolPageAutoDisable.test.tsx @@ -0,0 +1,108 @@ +import React from 'react'; +import { fireEvent, render, screen, waitFor } from '@testing-library/react'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +import ToolPage from './index'; + +const { + getToolMock, + listFixturesMock, + refetchMock, + reloadMock, + testToolMock, +} = vi.hoisted(() => ({ + getToolMock: vi.fn(), + listFixturesMock: vi.fn(), + refetchMock: vi.fn(), + reloadMock: vi.fn(), + testToolMock: vi.fn(), +})); + +const enabledTool = { + name: 'failing_custom_tool', + description: 'Always fails', + source: 'plugin_py', + category: 'custom', + enabled: true, + parameters: [], +}; + +vi.mock('react-i18next', () => ({ + useTranslation: () => ({ + t: (key: string, options?: Record) => + typeof options?.defaultValue === 'string' ? options.defaultValue : key, + i18n: { language: 'zh-CN' }, + }), +})); + +vi.mock('@/components/common/Toast', () => ({ + useToast: () => ({ + error: vi.fn(), + success: vi.fn(), + warning: vi.fn(), + info: vi.fn(), + }), +})); + +vi.mock('@/hooks/useTools', () => ({ + useToolPage: () => ({ + tools: [enabledTool], + total: 1, + facets: { + category: { custom: 1 }, + source: { plugin_py: 1 }, + source_groups: {}, + source_name: { Flocks: 1 }, + enabled: { true: 1 }, + }, + loading: false, + error: null, + initialized: true, + refetch: refetchMock, + reload: reloadMock, + }), +})); + +vi.mock('@/api/tool', () => ({ + canDirectlyTestTool: vi.fn(() => true), + toolAPI: { + get: getToolMock, + listFixtures: listFixturesMock, + test: testToolMock, + }, +})); + +vi.mock('@/api/provider', () => ({ + providerAPI: { + listApiServices: vi.fn(() => Promise.resolve({ data: [] })), + }, +})); + +describe('ToolPage auto-disable synchronization', () => { + beforeEach(() => { + vi.clearAllMocks(); + getToolMock.mockResolvedValue({ data: enabledTool }); + listFixturesMock.mockResolvedValue({ data: [] }); + reloadMock.mockResolvedValue(undefined); + testToolMock.mockResolvedValue({ + data: { + success: false, + error: 'synthetic repeated failure', + metadata: { disabled: true, disabled_reason: 'repeated_error' }, + }, + }); + }); + + it('reloads list state without refreshing plugins after an auto-disable result', async () => { + render(); + + fireEvent.click(screen.getByRole('button', { name: enabledTool.name })); + await screen.findByText(enabledTool.description); + fireEvent.click(screen.getByRole('button', { name: 'toolDetail.tabTest' })); + fireEvent.click(screen.getByRole('button', { name: 'toolDetail.runTest' })); + + await waitFor(() => expect(reloadMock).toHaveBeenCalledTimes(1)); + expect(refetchMock).not.toHaveBeenCalled(); + expect(screen.getByRole('button', { name: 'toolDetail.runTest' })).toBeDisabled(); + }); +}); diff --git a/webui/src/pages/Tool/index.tsx b/webui/src/pages/Tool/index.tsx index cd0c44400..145828e34 100644 --- a/webui/src/pages/Tool/index.tsx +++ b/webui/src/pages/Tool/index.tsx @@ -227,6 +227,7 @@ export default function ToolPage() { error, initialized: toolPageInitialized, refetch, + reload: reloadToolPage, } = useToolPage(toolPageParams); const [hasLoadedToolPage, setHasLoadedToolPage] = useState(false); @@ -420,7 +421,21 @@ export default function ToolPage() { setTestResult(null); const params = JSON.parse(testParams); const response = await toolAPI.test(selectedTool.name, params); - setTestResult(response.data); + const result = response.data; + setTestResult(result); + if (result.metadata?.disabled === true) { + setSelectedTool((current) => ( + current?.name === selectedTool.name + ? { ...current, enabled: false } + : current + )); + void reloadToolPage().catch((err: unknown) => { + toast.error( + t('alert.refreshFailedTitle'), + extractErrorMessage(err, t('alert.unknownError')), + ); + }); + } } catch (err: any) { setTestResult({ success: false, error: err.message }); } finally { @@ -3538,6 +3553,12 @@ export function ToolDetailDrawer({ setEnabledDefault(tool.enabled_default ?? tool.enabled); }, [tool.enabled, tool.enabled_customized, tool.enabled_default]); + useEffect(() => { + if (testResult?.metadata?.disabled === true) { + setEnabled(false); + } + }, [testResult]); + useEffect(() => { setFixturesLoading(true); toolAPI.listFixtures(tool.name) @@ -3799,14 +3820,14 @@ export function ToolDetailDrawer({
- {!tool.enabled && ( + {!enabled && (

{t('toolDetail.disabledNote')}

)} {canDirectTest ? null : ( diff --git a/webui/src/pages/Workflow/index.test.tsx b/webui/src/pages/Workflow/index.test.tsx index 9b0c06551..7da7fcc4a 100644 --- a/webui/src/pages/Workflow/index.test.tsx +++ b/webui/src/pages/Workflow/index.test.tsx @@ -1,7 +1,8 @@ import type { ReactNode } from 'react'; import { describe, it, expect, vi, beforeEach } from 'vitest'; -import { render, screen, within } from '@testing-library/react'; +import { render, screen, waitFor, within } from '@testing-library/react'; import userEvent from '@testing-library/user-event'; +import type { WorkflowSummary } from '@/api/workflow'; import WorkflowPage from './index'; const { mockNavigate, mockUseWorkflows, mockLanguage } = vi.hoisted(() => ({ @@ -23,8 +24,19 @@ vi.mock('react-i18next', () => ({ createWorkflow: '创建工作流', 'section.custom': '自定义工作流', 'section.builtin': '内置工作流', - 'status.draft': '草稿', 'stats.nodes': '节点', + 'integration.aria': '发布与触发状态', + 'integration.api': 'API', + 'integration.trigger': 'Trigger', + 'integration.moreStatuses': '更多状态', + 'integration.triggerType.syslog': 'Syslog', + 'integration.triggerType.kafka': 'Kafka', + 'integration.triggerType.schedule': 'Schedule', + 'integration.triggerType.webhook': 'Webhook', + 'integration.state.enabled': '启用', + 'integration.state.disabled': '关闭', + 'integration.detailState.unconfigured': '未配置', + 'integration.detailState.error': '异常', noDescription: '无描述', }; return translations[key] ?? key; @@ -61,16 +73,16 @@ vi.mock('@/components/common/EmptyState', () => ({ ), })); -function makeWorkflow(overrides: Partial = {}) { +function makeWorkflow(overrides: Partial = {}): WorkflowSummary { return { id: 'wf-1', name: '默认工作流', category: 'default', - workflowJson: { start: 'node-1', nodes: [{ id: 'node-1' }], edges: [] }, status: 'draft' as const, source: 'project' as const, createdAt: Date.now(), updatedAt: Date.now(), + nodeCount: 1, stats: { callCount: 0, successCount: 0, @@ -193,4 +205,115 @@ describe('WorkflowPage', () => { expect(screen.queryByRole('region', { name: '自定义工作流' })).not.toBeInTheDocument(); expect(screen.getByRole('region', { name: '内置工作流' })).toBeInTheDocument(); }); + + it('展示 API 与 Trigger 的聚合运行状态', () => { + mockUseWorkflows.mockReturnValue({ + workflows: [makeWorkflow({ + integrationStatus: { + api: { configured: true, state: 'running' }, + trigger: { + configured: true, + state: 'error', + count: 2, + items: [ + { id: 'syslog-default', type: 'syslog', state: 'running', rawState: 'listening' }, + { id: 'kafka-default', type: 'kafka', state: 'error', rawState: 'failed' }, + ], + }, + }, + })], + loading: false, + error: null, + refetch: vi.fn(), + }); + + render(); + + expect(within(screen.getByLabelText('API：启用')).getByText('启用')).toHaveClass('text-green-600'); + expect(within(screen.getByLabelText('Syslog：启用')).getByText('启用')).toHaveClass('text-green-600'); + const kafkaStatus = screen.getByLabelText('Kafka：关闭（异常）'); + expect(within(kafkaStatus).getByText('关闭')).toHaveClass('text-red-600'); + expect(kafkaStatus).toHaveAttribute('title', 'Kafka：关闭（异常）'); + expect(screen.queryByText('草稿')).not.toBeInTheDocument(); + }); + + it('未配置发布能力时使用灰色状态', () => { + mockUseWorkflows.mockReturnValue({ + workflows: [makeWorkflow()], + loading: false, + error: null, + refetch: vi.fn(), + }); + + render(); + + const apiStatus = screen.getByLabelText('API：关闭（未配置）'); + const triggerStatus = screen.getByLabelText('Trigger：关闭（未配置）'); + expect(within(apiStatus).getByText('关闭')).toHaveClass('text-gray-400'); + expect(within(triggerStatus).getByText('关闭')).toHaveClass('text-gray-400'); + expect(apiStatus).toHaveAttribute('title', 'API：关闭（未配置）'); + }); + + it('Trigger 状态过多时保持单行并通过浮层展示其余状态', async () => { + const user = userEvent.setup(); + mockUseWorkflows.mockReturnValue({ + workflows: [makeWorkflow({ + integrationStatus: { + api: { configured: true, state: 'running' }, + trigger: { + configured: true, + state: 'error', + count: 4, + items: [ + { id: 'syslog-default', type: 'syslog', state: 'running' }, + { id: 'kafka-default', type: 'kafka', state: 'running' }, + { id: 'schedule-default', type: 'schedule', state: 'error' }, + { id: 'webhook-default', type: 'webhook', state: 'unconfigured' }, + ], + }, + }, + })], + loading: false, + error: null, + refetch: vi.fn(), + }); + + render(); + + const overflowStatus = screen.getByRole('button', { name: '更多状态：2' }); + expect(overflowStatus.parentElement).toHaveClass('h-5'); + expect(screen.queryByText('Schedule')).not.toBeInTheDocument(); + expect(screen.queryByText('Webhook')).not.toBeInTheDocument(); + + await user.hover(overflowStatus); + const hoverTooltip = screen.getByRole('tooltip'); + expect(within(hoverTooltip).getByText('Schedule')).toBeInTheDocument(); + expect(within(hoverTooltip).getByText('Webhook')).toBeInTheDocument(); + expect(within(hoverTooltip).queryByText('异常')).not.toBeInTheDocument(); + expect(within(hoverTooltip).queryByText('未配置')).not.toBeInTheDocument(); + expect(hoverTooltip).not.toHaveClass('-translate-y-full'); + expect(hoverTooltip).toHaveStyle({ maxHeight: '240px' }); + expect(hoverTooltip.querySelector('.overflow-y-auto')).toBeInTheDocument(); + + await user.unhover(overflowStatus); + await waitFor(() => expect(screen.queryByRole('tooltip')).not.toBeInTheDocument()); + + vi.spyOn(overflowStatus, 'getBoundingClientRect').mockReturnValue({ + x: 280, + y: 700, + width: 20, + height: 16, + top: 700, + right: 300, + bottom: 716, + left: 280, + toJSON: () => ({}), + }); + await user.click(overflowStatus); + expect(screen.getByRole('tooltip')).toHaveClass('-translate-y-full'); + expect(mockNavigate).not.toHaveBeenCalled(); + + await user.unhover(overflowStatus); + await waitFor(() => expect(screen.queryByRole('tooltip')).not.toBeInTheDocument()); + }); }); diff --git a/webui/src/pages/Workflow/index.tsx b/webui/src/pages/Workflow/index.tsx index ec2acd37b..ae175ad8a 100644 --- a/webui/src/pages/Workflow/index.tsx +++ b/webui/src/pages/Workflow/index.tsx @@ -1,4 +1,5 @@ -import { useState, useMemo, useEffect } from 'react'; +import { useState, useMemo, useEffect, useId, useCallback, useRef } from 'react'; +import { createPortal } from 'react-dom'; import { useNavigate } from 'react-router-dom'; import { Workflow as WorkflowIcon, @@ -14,43 +15,15 @@ import PageHeader from '@/components/common/PageHeader'; import LoadingSpinner from '@/components/common/LoadingSpinner'; import EmptyState from '@/components/common/EmptyState'; import { useWorkflows } from '@/hooks/useWorkflow'; -import { Workflow } from '@/api/workflow'; +import { + WorkflowCapabilityState, + WorkflowSummary, + WorkflowTriggerStatusSummary, + WorkflowTriggerType, +} from '@/api/workflow'; import { getWorkflowDisplayName } from '@/utils/workflowDisplay'; -// --------------------------------------------------------------------------- -// Color helpers (mirrors Agent page) -// --------------------------------------------------------------------------- - -const WORKFLOW_PALETTE = [ - '#3b82f6', // blue-500 - '#8b5cf6', // violet-500 - '#06b6d4', // cyan-500 - '#10b981', // emerald-500 - '#f59e0b', // amber-500 - '#ef4444', // red-500 - '#ec4899', // pink-500 - '#6366f1', // indigo-500 -]; - -function resolveWorkflowColor(workflow: Workflow): string { - let h = 0; - const seed = workflow.id || workflow.name; - for (let i = 0; i < seed.length; i++) { - h = seed.charCodeAt(i) + ((h << 5) - h); - } - return WORKFLOW_PALETTE[Math.abs(h) % WORKFLOW_PALETTE.length]; -} - -function hexAlpha(hex: string, alpha: number): string { - const h = hex.replace('#', ''); - const full = h.length === 3 ? h.split('').map(c => c + c).join('') : h; - const r = parseInt(full.slice(0, 2), 16); - const g = parseInt(full.slice(2, 4), 16); - const b = parseInt(full.slice(4, 6), 16); - return `rgba(${r},${g},${b},${alpha})`; -} - -function isBuiltin(workflow: Workflow): boolean { +function isBuiltin(workflow: WorkflowSummary): boolean { return workflow.source === 'project'; } @@ -60,6 +33,45 @@ function isBuiltin(workflow: Workflow): boolean { type SourceFilter = 'all' | 'builtin' | 'custom'; const PAGE_SIZE = 12; +const MAX_VISIBLE_TRIGGER_STATUSES = 2; +const STATUS_TOOLTIP_WIDTH = 224; +const STATUS_TOOLTIP_MAX_HEIGHT = 240; +const STATUS_TOOLTIP_GAP = 8; + +const CAPABILITY_STATUS_STYLES: Record = { + unconfigured: { text: 'text-gray-400', dot: 'bg-gray-300' }, + starting: { text: 'text-gray-500', dot: 'bg-gray-400' }, + running: { text: 'text-green-600', dot: 'bg-green-500' }, + stopped: { text: 'text-red-600', dot: 'bg-red-500' }, + error: { text: 'text-red-600', dot: 'bg-red-500' }, +}; + +const TRIGGER_TYPE_LABEL_KEYS = { + manual: 'integration.triggerType.manual', + schedule: 'integration.triggerType.schedule', + webhook: 'integration.triggerType.webhook', + syslog: 'integration.triggerType.syslog', + kafka: 'integration.triggerType.kafka', + internal_event: 'integration.triggerType.internal_event', + custom_webhook: 'integration.triggerType.custom_webhook', + custom_adapter: 'integration.triggerType.custom_adapter', + plugin: 'integration.triggerType.plugin', +} as const satisfies Record; + +const CAPABILITY_STATUS_LABEL_KEYS = { + unconfigured: 'integration.state.disabled', + starting: 'integration.state.disabled', + running: 'integration.state.enabled', + stopped: 'integration.state.disabled', + error: 'integration.state.disabled', +} as const satisfies Record; + +const CAPABILITY_DETAIL_LABEL_KEYS = { + unconfigured: 'integration.detailState.unconfigured', + starting: 'integration.detailState.starting', + stopped: 'integration.detailState.stopped', + error: 'integration.detailState.error', +} as const; // --------------------------------------------------------------------------- // WorkflowPage @@ -256,7 +268,7 @@ function WorkflowSection({ }: { title: string; icon: React.ReactNode; - workflows: Workflow[]; + workflows: WorkflowSummary[]; }) { const [page, setPage] = useState(1); const totalPages = Math.max(1, Math.ceil(workflows.length / PAGE_SIZE)); @@ -333,10 +345,9 @@ function WorkflowSection({ // WorkflowCard // --------------------------------------------------------------------------- -function WorkflowCard({ workflow }: { workflow: Workflow }) { +function WorkflowCard({ workflow }: { workflow: WorkflowSummary }) { const { t, i18n } = useTranslation('workflow'); const navigate = useNavigate(); - const color = resolveWorkflowColor(workflow); const builtin = isBuiltin(workflow); const displayName = getWorkflowDisplayName(workflow, i18n?.language); @@ -344,6 +355,16 @@ function WorkflowCard({ workflow }: { workflow: Workflow }) { workflow.stats.callCount > 0 ? ((workflow.stats.successCount / workflow.stats.callCount) * 100).toFixed(1) : '—'; + const apiStatus = workflow.integrationStatus?.api ?? { configured: false, state: 'unconfigured' as const }; + const triggerStatus = workflow.integrationStatus?.trigger ?? { + configured: false, + state: 'unconfigured' as const, + count: 0, + items: [], + }; + const triggerItems = triggerStatus.items ?? []; + const visibleTriggerItems = triggerItems.slice(0, MAX_VISIBLE_TRIGGER_STATUSES); + const hiddenTriggerItems = triggerItems.slice(MAX_VISIBLE_TRIGGER_STATUSES); return (
- {/* Top accent bar */} -
- {/* Card body */}
{/* Avatar + name row */}
-
- +
+
@@ -374,23 +389,18 @@ function WorkflowCard({ workflow }: { workflow: Workflow }) { {/* Source badge */} {builtin ? ( + bg-blue-50 text-blue-600"> {t('badge.builtin')} ) : ( + bg-teal-50 text-teal-600"> {t('badge.custom')} )} - {/* Status badge */} - - {t(`status.${workflow.status}` as any) ?? workflow.status} - {/* Node count */} - {workflow.workflowJson.nodes.length} {t('stats.nodes')} + {workflow.nodeCount} {t('stats.nodes')}
@@ -402,6 +412,36 @@ function WorkflowCard({ workflow }: { workflow: Workflow }) {

{workflow.description || t('noDescription')}

+ +
+
+ + {triggerItems.length > 0 ? visibleTriggerItems.map(trigger => ( + + )) : ( + + )} +
+ {hiddenTriggerItems.length > 0 && ( + + )} +
{/* Stats footer — kept from original, cleaned to white bg */} @@ -430,3 +470,162 @@ function WorkflowCard({ workflow }: { workflow: Workflow }) {
); } + +function CapabilityStatus({ + label, + state, + statusLabel, +}: { + label: string; + state: WorkflowCapabilityState; + statusLabel: string; +}) { + const { t } = useTranslation('workflow'); + const detailLabel = state === 'running' + ? null + : t(CAPABILITY_DETAIL_LABEL_KEYS[state]); + const accessibleLabel = detailLabel + ? `${label}：${statusLabel}（${detailLabel}）` + : `${label}：${statusLabel}`; + const styles = CAPABILITY_STATUS_STYLES[state]; + return ( + + + {label} + {statusLabel} + + ); +} + +function OverflowStatusList({ items }: { items: WorkflowTriggerStatusSummary[] }) { + const { t } = useTranslation('workflow'); + const tooltipId = useId(); + const hideTimeoutRef = useRef(null); + const [position, setPosition] = useState<{ + x: number; + y: number; + placement: 'top' | 'bottom'; + maxHeight: number; + } | null>(null); + + const cancelScheduledHide = useCallback(() => { + if (hideTimeoutRef.current !== null) { + window.clearTimeout(hideTimeoutRef.current); + hideTimeoutRef.current = null; + } + }, []); + + const showTooltip = useCallback((target: HTMLElement) => { + cancelScheduledHide(); + const rect = target.getBoundingClientRect(); + const viewportWidth = window.innerWidth; + const viewportHeight = window.innerHeight; + const estimatedHeight = Math.min(STATUS_TOOLTIP_MAX_HEIGHT, 36 + items.length * 24); + const availableAbove = Math.max(0, rect.top - STATUS_TOOLTIP_GAP * 2); + const availableBelow = Math.max(0, viewportHeight - rect.bottom - STATUS_TOOLTIP_GAP * 2); + const placement = availableAbove >= estimatedHeight || availableAbove >= availableBelow + ? 'top' + : 'bottom'; + const availableHeight = placement === 'top' ? availableAbove : availableBelow; + const x = Math.min( + Math.max(rect.right, STATUS_TOOLTIP_WIDTH + STATUS_TOOLTIP_GAP), + viewportWidth - STATUS_TOOLTIP_GAP, + ); + setPosition({ + x, + y: placement === 'top' + ? rect.top - STATUS_TOOLTIP_GAP + : rect.bottom + STATUS_TOOLTIP_GAP, + placement, + maxHeight: Math.max(48, Math.min(STATUS_TOOLTIP_MAX_HEIGHT, availableHeight)), + }); + }, [cancelScheduledHide, items.length]); + + const hideTooltip = useCallback(() => { + cancelScheduledHide(); + setPosition(null); + }, [cancelScheduledHide]); + + const scheduleHideTooltip = useCallback(() => { + cancelScheduledHide(); + hideTimeoutRef.current = window.setTimeout(hideTooltip, 100); + }, [cancelScheduledHide, hideTooltip]); + + useEffect(() => { + if (!position) return; + const handleKeyDown = (event: KeyboardEvent) => { + if (event.key === 'Escape') hideTooltip(); + }; + document.addEventListener('pointerdown', hideTooltip); + document.addEventListener('keydown', handleKeyDown); + return () => { + document.removeEventListener('pointerdown', hideTooltip); + document.removeEventListener('keydown', handleKeyDown); + }; + }, [hideTooltip, position]); + + useEffect(() => () => cancelScheduledHide(), [cancelScheduledHide]); + + return ( + <> + + {position && createPortal( + ); @@ -1181,7 +1219,7 @@ function PublishSection({
{configExpanded ? (
- {renderDriverSelector({ showApply: true })} + {renderPublishConfig({ showApply: true })}
Invoke URL
@@ -1252,7 +1290,7 @@ function PublishSection({ {configExpanded ? (
- {renderDriverSelector()} + {renderPublishConfig()}
) : null}
diff --git a/webui/src/types/index.ts b/webui/src/types/index.ts index e78dbfd28..7314c7725 100644 --- a/webui/src/types/index.ts +++ b/webui/src/types/index.ts @@ -7,6 +7,7 @@ export interface Session { id: string; slug: string; projectID: string; + effectiveProjectID?: string; directory: string; parentID?: string; summary?: SessionSummary; @@ -397,10 +398,8 @@ export interface MCPCatalogStats { export interface ProviderCredentials { secret_id?: string | null; /** - * Write-only compatibility field. Credential reads never return the raw key; - * they return ``api_key: null`` plus ``api_key_masked`` and - * ``has_credential``. UI code must never submit the masked value as a new - * secret. Leave ``api_key`` out of an update to preserve the stored key. + * Raw API key returned only by the admin-only reveal endpoint. Ordinary LLM + * and API-service credential reads return null and expose only masked values. */ api_key?: string | null; api_key_masked?: string | null; diff --git a/webui/src/utils/workflowDisplay.ts b/webui/src/utils/workflowDisplay.ts index 49d962beb..0d12d757c 100644 --- a/webui/src/utils/workflowDisplay.ts +++ b/webui/src/utils/workflowDisplay.ts @@ -1,6 +1,7 @@ -import type { Workflow } from '@/api/workflow'; +import type { Workflow, WorkflowSummary } from '@/api/workflow'; type LocalizedNameSource = Record | null | undefined; +type WorkflowDisplaySource = Workflow | WorkflowSummary; function isChineseLocale(language?: string): boolean { if (!language) return false; @@ -21,11 +22,11 @@ function localizedNamesFrom(value: LocalizedNameSource): Record return names; } -function collectLocalizedNames(workflow: Workflow | null | undefined): Record { +function collectLocalizedNames(workflow: WorkflowDisplaySource | null | undefined): Record { if (!workflow) return {}; - const raw = workflow as Workflow & Record; - const workflowJson = workflow.workflowJson as unknown as Record | undefined; - const metadata = workflow.workflowJson?.metadata as Record | undefined; + const raw = workflow as WorkflowDisplaySource & Record; + const workflowJson = raw.workflowJson as Record | undefined; + const metadata = workflowJson?.metadata as Record | undefined; const names: Record = {}; [ @@ -70,7 +71,7 @@ function pickLocalizedName(names: Record, language?: string): st } export function getWorkflowDisplayName( - workflow: Workflow | null | undefined, + workflow: WorkflowDisplaySource | null | undefined, language?: string, ): string { if (!workflow) return '';