From 4e59e4d5035d7f65fd2d6b5576666b3715d09e1c Mon Sep 17 00:00:00 2001 From: stephamie7 <1223696150@qq.com> Date: Wed, 12 Aug 2026 17:28:08 +0800 Subject: [PATCH 01/63] fix(release): handle published draft flag --- .github/workflows/sync-gitee.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/sync-gitee.yml b/.github/workflows/sync-gitee.yml index d9e7e942d..7e1da5c1c 100644 --- a/.github/workflows/sync-gitee.yml +++ b/.github/workflows/sync-gitee.yml @@ -75,7 +75,7 @@ jobs: > "${release_metadata}" resolved_tag="$(jq -er '.tag_name' "${release_metadata}")" - is_draft="$(jq -er '.draft' "${release_metadata}")" + is_draft="$(jq -r '.draft' "${release_metadata}")" published_at="$(jq -r '.published_at // empty' "${release_metadata}")" if [[ "${resolved_tag}" != "${TAG_NAME}" ]]; then From edcef181dcd966c174ed61fbf2618d521ba83941 Mon Sep 17 00:00:00 2001 From: stephamie7 <1223696150@qq.com> Date: Wed, 12 Aug 2026 17:37:29 +0800 Subject: [PATCH 02/63] fix(release): verify Gitee release after sync --- .github/workflows/sync-gitee.yml | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/.github/workflows/sync-gitee.yml b/.github/workflows/sync-gitee.yml index 7e1da5c1c..c41d97aea 100644 --- a/.github/workflows/sync-gitee.yml +++ b/.github/workflows/sync-gitee.yml @@ -266,6 +266,21 @@ jobs: exit 1 fi + verification_status="$( + curl --silent --show-error --retry 3 \ + --output "${release_response}" \ + --write-out '%{http_code}' \ + --get \ + --data-urlencode "access_token=${GITEE_TOKEN}" \ + "${GITEE_API_BASE_URL}/repos/${GITEE_OWNER}/${GITEE_REPO}/releases/tags/${encoded_tag}" + )" + + if [[ "${verification_status}" != "200" ]]; then + message="$(jq -r '.message // "unknown Gitee API error"' "${release_response}" 2>/dev/null || true)" + echo "::error title=Gitee Release verification failed::HTTP ${verification_status}: ${message}" + exit 1 + fi + synced_tag="$(jq -er '.tag_name' "${release_response}")" synced_release_id="$(jq -er '.id' "${release_response}")" if [[ "${synced_tag}" != "${TAG_NAME}" ]]; then From 5e878e41588ba9600732e4652e1c182b28d2edf7 Mon Sep 17 00:00:00 2001 From: stephamie7 <1223696150@qq.com> Date: Wed, 12 Aug 2026 17:39:31 +0800 Subject: [PATCH 03/63] fix(release): report Gitee API request failures --- .github/workflows/sync-gitee.yml | 24 ++++++++++++++++++------ 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/.github/workflows/sync-gitee.yml b/.github/workflows/sync-gitee.yml index c41d97aea..b089b0298 100644 --- a/.github/workflows/sync-gitee.yml +++ b/.github/workflows/sync-gitee.yml @@ -218,14 +218,18 @@ jobs: encoded_tag="$(jq -rn --arg value "${TAG_NAME}" '$value | @uri')" release_lookup="${RUNNER_TEMP}/gitee-release-lookup.json" release_response="${RUNNER_TEMP}/gitee-release-response.json" - lookup_status="$( + if ! lookup_status="$( curl --silent --show-error --retry 3 \ --output "${release_lookup}" \ --write-out '%{http_code}' \ --get \ --data-urlencode "access_token=${GITEE_TOKEN}" \ "${GITEE_API_BASE_URL}/repos/${GITEE_OWNER}/${GITEE_REPO}/releases/tags/${encoded_tag}" - )" + )"; then + echo "::error title=Gitee Release lookup request failed::Could not query ${TAG_NAME} after retries." + exit 1 + fi + echo "Gitee Release lookup returned HTTP ${lookup_status}." common_fields=( --data-urlencode "access_token=${GITEE_TOKEN}" @@ -251,14 +255,18 @@ jobs: exit 1 fi - sync_status="$( + if ! sync_status="$( curl --silent --show-error --retry 3 \ --output "${release_response}" \ --write-out '%{http_code}' \ --request "${sync_method}" \ "${common_fields[@]}" \ "${sync_url}" - )" + )"; then + echo "::error title=Gitee Release sync request failed::Release ${sync_method} failed after retries." + exit 1 + fi + echo "Gitee Release ${sync_method} returned HTTP ${sync_status}." if [[ ! "${sync_status}" =~ ^2[0-9][0-9]$ ]]; then message="$(jq -r '.message // "unknown Gitee API error"' "${release_response}" 2>/dev/null || true)" @@ -266,14 +274,18 @@ jobs: exit 1 fi - verification_status="$( + if ! verification_status="$( curl --silent --show-error --retry 3 \ --output "${release_response}" \ --write-out '%{http_code}' \ --get \ --data-urlencode "access_token=${GITEE_TOKEN}" \ "${GITEE_API_BASE_URL}/repos/${GITEE_OWNER}/${GITEE_REPO}/releases/tags/${encoded_tag}" - )" + )"; then + echo "::error title=Gitee Release verification request failed::Could not verify ${TAG_NAME} after retries." + exit 1 + fi + echo "Gitee Release verification returned HTTP ${verification_status}." if [[ "${verification_status}" != "200" ]]; then message="$(jq -r '.message // "unknown Gitee API error"' "${release_response}" 2>/dev/null || true)" From 7fadadbaf1ae25fdb0247d31c9b8ca9f24591c04 Mon Sep 17 00:00:00 2001 From: stephamie7 <1223696150@qq.com> Date: Wed, 12 Aug 2026 17:41:35 +0800 Subject: [PATCH 04/63] fix(release): handle null Gitee lookup --- .github/workflows/sync-gitee.yml | 29 +++++++++++++++++++++-------- 1 file changed, 21 insertions(+), 8 deletions(-) diff --git a/.github/workflows/sync-gitee.yml b/.github/workflows/sync-gitee.yml index b089b0298..9bcdc42f9 100644 --- a/.github/workflows/sync-gitee.yml +++ b/.github/workflows/sync-gitee.yml @@ -239,22 +239,35 @@ jobs: --data-urlencode "prerelease=${prerelease}" ) + release_missing="false" if [[ "${lookup_status}" == "200" ]]; then - release_id="$(jq -er '.id' "${release_lookup}")" - sync_method="PATCH" - sync_url="${GITEE_API_BASE_URL}/repos/${GITEE_OWNER}/${GITEE_REPO}/releases/${release_id}" - echo "Gitee Release ${TAG_NAME} already exists; updating its metadata only." + release_id="$(jq -r '.id // empty' "${release_lookup}")" + if [[ -n "${release_id}" ]]; then + sync_method="PATCH" + sync_url="${GITEE_API_BASE_URL}/repos/${GITEE_OWNER}/${GITEE_REPO}/releases/${release_id}" + echo "Gitee Release ${TAG_NAME} already exists; updating its metadata only." + elif jq -e 'type == "null"' "${release_lookup}" > /dev/null; then + release_missing="true" + echo "Gitee returned HTTP 200 with null for missing Release ${TAG_NAME}." + else + echo "::error title=Unexpected Gitee Release lookup response::HTTP 200 response does not contain a Release ID." + exit 1 + fi elif [[ "${lookup_status}" == "404" ]]; then - sync_method="POST" - sync_url="${GITEE_API_BASE_URL}/repos/${GITEE_OWNER}/${GITEE_REPO}/releases" - common_fields+=(--data-urlencode "target_commitish=${VERIFIED_SHA}") - echo "Creating Gitee Release ${TAG_NAME} at verified commit ${VERIFIED_SHA}." + release_missing="true" else message="$(jq -r '.message // "unknown Gitee API error"' "${release_lookup}" 2>/dev/null || true)" echo "::error title=Gitee Release lookup failed::HTTP ${lookup_status}: ${message}" exit 1 fi + if [[ "${release_missing}" == "true" ]]; then + sync_method="POST" + sync_url="${GITEE_API_BASE_URL}/repos/${GITEE_OWNER}/${GITEE_REPO}/releases" + common_fields+=(--data-urlencode "target_commitish=${VERIFIED_SHA}") + echo "Creating Gitee Release ${TAG_NAME} at verified commit ${VERIFIED_SHA}." + fi + if ! sync_status="$( curl --silent --show-error --retry 3 \ --output "${release_response}" \ From 62ca11856ce48e18ba857073019e9422e06a4f1b Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Mon, 17 Aug 2026 15:46:49 +0800 Subject: [PATCH 05/63] refactor(delegation): migrate legacy flows to delegate_task --- .../loop_host_forensics_fast/workflow.json | 4 +- .../loop_host_forensics_fast/workflow.md | 4 +- flocks/cli/commands/import_.py | 15 +- flocks/command/command.py | 1 - flocks/config/config.py | 35 ++- flocks/permission/helpers.py | 44 ++++ flocks/server/routes/misc.py | 3 - flocks/server/routes/session.py | 9 - flocks/server/routes/skill.py | 2 - flocks/session/__init__.py | 11 - flocks/session/context_usage.py | 8 +- flocks/session/features/subtask.py | 62 ----- flocks/session/message.py | 33 +-- flocks/session/prompt/anthropic-20250930.txt | 4 +- flocks/session/prompt/anthropic.txt | 10 +- flocks/session/prompt_strings.py | 6 +- flocks/session/session_loop.py | 214 +--------------- flocks/tool/agent/delegate_task.py | 10 +- flocks/tool/catalog.py | 1 - flocks/tool/code/grep.py | 2 +- flocks/tool/truncation.py | 4 +- flocks/utils/id.py | 2 - flocks/workflow/tool_context.py | 2 +- tests/config/test_config.py | 11 + tests/session/test_context_usage.py | 10 +- tests/session/test_message_parts.py | 41 +-- tests/session/test_session_abort_inject.py | 53 ---- tests/tool/test_builtin_management_tools.py | 11 +- tests/tool/test_task_model_pinning.py | 61 ----- tests/tool/test_tool_catalog.py | 8 +- tests/tool/test_tools.py | 12 +- tests/utils/test_id_compatibility.py | 1 - .../test_loop_host_forensics_fast_workflow.py | 4 +- tui/flocks/agent/generate.txt | 6 +- tui/flocks/cli/cmd/agent.ts | 2 +- .../cli/cmd/tui/routes/session/index.tsx | 10 +- .../cli/cmd/tui/routes/session/permission.tsx | 2 +- tui/flocks/command/index.ts | 3 - tui/flocks/config/config.test.ts | 13 + tui/flocks/config/config.ts | 9 +- tui/flocks/session/message-v2.ts | 22 -- tui/flocks/session/prompt.ts | 237 ++---------------- .../session/prompt/anthropic-20250930.txt | 4 +- tui/flocks/session/prompt/anthropic.txt | 10 +- tui/flocks/tool/bash.txt | 4 +- tui/flocks/tool/{task.ts => delegate-task.ts} | 17 +- .../tool/{task.txt => delegate-task.txt} | 18 +- tui/flocks/tool/glob.txt | 2 +- tui/flocks/tool/grep.txt | 2 +- tui/flocks/tool/registry.ts | 4 +- tui/flocks/tool/truncation.ts | 8 +- tui/sdk/gen/types.gen.ts | 23 +- tui/sdk/v2/gen/sdk.gen.ts | 5 +- tui/sdk/v2/gen/types.gen.ts | 33 +-- webui/src/api/skill.ts | 1 - 55 files changed, 266 insertions(+), 867 deletions(-) delete mode 100644 flocks/session/features/subtask.py delete mode 100644 tests/tool/test_task_model_pinning.py create mode 100644 tui/flocks/config/config.test.ts rename tui/flocks/tool/{task.ts => delegate-task.ts} (91%) rename tui/flocks/tool/{task.txt => delegate-task.txt} (79%) diff --git a/.flocks/plugins/workflows/loop_host_forensics_fast/workflow.json b/.flocks/plugins/workflows/loop_host_forensics_fast/workflow.json index e0623b805..e311ee35e 100644 --- a/.flocks/plugins/workflows/loop_host_forensics_fast/workflow.json +++ b/.flocks/plugins/workflows/loop_host_forensics_fast/workflow.json @@ -24,8 +24,8 @@ "id": "inspect_host", "type": "python", "name": "单台快速巡检", - "description": "每台主机先做 SSH 预检；预检通过后再调 task(host-forensics-fast)。超时仅重试一次；循环态保留文件路径、执行状态、verdict 与失败分类。", - "code": "import os\nimport re\nimport time\n\nhosts = inputs.get(\"hosts\", [])\nidx = int(inputs.get(\"host_idx\", 0))\nhost = hosts[idx] if 0 <= idx < len(hosts) else \"\"\nhost = str(host).strip()\n\nsu = inputs.get(\"ssh_user\")\nif isinstance(su, str):\n su = su.strip()\nelse:\n su = \"\"\n\nconnect_host = host\nconnect_user = \"\"\nif \"@\" in host:\n user_part, host_part = host.split(\"@\", 1)\n connect_user = str(user_part).strip()\n connect_host = str(host_part).strip() or host\n ssh_target = host\n user_hint = (\n \"主机列表项已含 `user@host` 形式。SSH 工具调用必须使用：host=`\"\n + connect_host\n + \"`，username=`\"\n + connect_user\n + \"`。\"\n )\nelif su:\n connect_user = su\n connect_host = host\n ssh_target = su + \"@\" + host\n user_hint = (\n \"工作流已指定 `ssh_user`=`\"\n + su\n + \"`。SSH 工具调用必须使用：host=`\"\n + connect_host\n + \"`，username=`\"\n + connect_user\n + \"`。\"\n )\nelse:\n connect_host = host\n ssh_target = host\n user_hint = (\n \"未指定 `ssh_user`。SSH 工具调用请使用 host=`\"\n + connect_host\n + \"`，username 留空使用默认账户（一般为 root）。\"\n )\n\n\ndef _extract_verdict(markdown_text):\n if not markdown_text:\n return \"UNKNOWN\"\n match = re.search(\n r\"(?im)^\\s*\\*{0,2}Verdict\\*{0,2}\\s*:\\s*\"\n r\"(CLEAN|SUSPICIOUS|COMPROMISED|UNKNOWN)\\b\",\n markdown_text,\n )\n if not match:\n return \"UNKNOWN\"\n return str(match.group(1)).upper()\n\n\ndef _is_timeout_error(message):\n text = str(message or \"\").lower()\n return (\n \"timed out\" in text\n or \"timeout\" in text\n or \"超时\" in text\n or \"节点执行超时\" in text\n )\n\n\ndef _classify_error(message):\n text = str(message or \"\")\n lower = text.lower()\n if not text.strip():\n return \"unknown\"\n if \"permission denied\" in lower or \"auth failed\" in lower or \"authentication failed\" in lower:\n return \"auth_failed\"\n if \"host key verification failed\" in lower or \"host key\" in lower:\n return \"host_key_verification_failed\"\n if \"connection refused\" in lower:\n return \"connection_refused\"\n if \"no route to host\" in lower:\n return \"no_route_to_host\"\n if \"network is unreachable\" in lower:\n return \"network_unreachable\"\n if \"name or service not known\" in lower or \"could not resolve\" in lower or \"nodename nor servname provided\" in lower:\n return \"dns_resolution_failed\"\n if \"connection reset\" in lower:\n return \"connection_reset\"\n if \"broken pipe\" in lower or \"connection lost\" in lower or \"disconnect\" in lower:\n return \"connection_lost\"\n if \"kex\" in lower or \"key exchange\" in lower or \"protocol error\" in lower:\n return \"ssh_handshake_failed\"\n if _is_timeout_error(text):\n if \"connect\" in lower or \"ssh connection failed\" in lower:\n return \"connect_timeout\"\n return \"execution_timeout\"\n if \"ssh connection failed\" in lower:\n return \"ssh_connection_failed\"\n return \"unknown\"\n\n\nidx1 = idx + 1\nper_host_dir = str(inputs.get(\"per_host_dir\") or \"\").strip()\nif not per_host_dir:\n od = str(inputs.get(\"output_dir\") or \"\").strip()\n if od:\n per_host_dir = os.path.join(od, \"host_triage\")\n else:\n per_host_dir = os.path.join(os.path.dirname(inputs.get(\"batch_report_path\") or \".\") or \".\", \"host_triage\")\nos.makedirs(per_host_dir, exist_ok=True)\n\n\ndef _slug(s):\n t = re.sub(r\"[^0-9A-Za-z._@-]+\", \"_\", str(s)).strip(\"_\")\n t = t.replace(\"@\", \"_at_\")\n return (t[:56] if t else \"host\")\n\n\nnh = len(hosts)\nbase = \"{:04d}_{}\".format(idx1, _slug(ssh_target))\nper_host_md = os.path.join(per_host_dir, base + \".md\")\n\npreflight_timeout_s = 20\npreflight_res = tool.run_safe(\n \"ssh_host_cmd\",\n host=connect_host,\n username=(connect_user or None),\n command=\"echo FLOCKS_SSH_OK\",\n timeout=preflight_timeout_s,\n)\npreflight_output = preflight_res.get(\"output\") or preflight_res.get(\"text\") or \"\"\npreflight_error = preflight_res.get(\"error\") or \"\"\npreflight_ok = bool(preflight_res.get(\"success\")) and \"FLOCKS_SSH_OK\" in str(preflight_output)\n\ntext = \"\"\nok = False\nerr = \"\"\nverdict = \"UNKNOWN\"\nfailure_category = \"\"\nattempts = 0\n\nif preflight_ok:\n desc = \"Fast triage item \" + str(idx1)\n prompt = (\n \"你是 host-forensics-fast 工作模式：对下列目标执行 Linux 主机快速安全巡检（首轮研判）。\\n\"\n \"请使用 ssh_run_script，script_path 为 `.flocks/plugins/agents/host-forensics-fast/scripts/triage_fast.sh`。\\n\"\n + user_hint\n + \"\\n\\n本次 SSH 工具参数必须使用：\\n\"\n + \"- host: \"\n + connect_host\n + \"\\n\"\n + (\"- username: \" + connect_user + \"\\n\" if connect_user else \"- username: （留空，使用默认账户）\\n\")\n + \"\\n请输出简洁 Markdown：结论、可疑项、风险判断、后续建议。\"\n )\n while attempts < 2:\n attempts += 1\n res = tool.run_safe(\n \"task\",\n description=desc + \" attempt \" + str(attempts),\n prompt=prompt,\n subagent_type=\"host-forensics-fast\",\n )\n text = res.get(\"text\") or \"\"\n ok = bool(res.get(\"success\"))\n err = res.get(\"error\") or \"\"\n if ok:\n verdict = _extract_verdict(text)\n break\n if not _is_timeout_error(err) or attempts >= 2:\n failure_category = _classify_error(err)\n break\n time.sleep(3)\n if not ok and not failure_category:\n failure_category = _classify_error(err)\nelse:\n err = preflight_error or \"SSH preflight failed\"\n failure_category = _classify_error(err)\n\nlines = [\n \"# 单主机快速巡检结果\",\n \"\",\n \"- 批次内序号: {} / {}\".format(idx1, nh),\n \"- 列表项 host: `{}`\".format(host),\n \"- ssh_target: `{}`\".format(ssh_target),\n \"- ssh_host: `{}`\".format(connect_host),\n \"- ssh_user: `{}`\".format(connect_user or \"(default)\"),\n \"- success: {}\".format(ok),\n \"- verdict: {}\".format(verdict),\n \"- failure_category: {}\".format(failure_category or \"\"),\n \"- inspect_attempts: {}\".format(attempts),\n \"\",\n]\nif not preflight_ok:\n lines.extend(\n [\n \"## SSH 预检失败\",\n \"\",\n \"- 分类: `{}`\".format(failure_category),\n \"- 错误:\",\n \"\",\n \"```\",\n str(err),\n \"```\",\n \"\",\n ]\n )\nelif err:\n lines.extend([\"## 错误\", \"\", \"```\", str(err), \"```\", \"\"])\nelse:\n lines.extend([\"## 子 Agent 输出\", \"\", (text if text else \"_(无正文)_\"), \"\"])\nwith open(per_host_md, \"w\", encoding=\"utf-8\") as f:\n f.write(\"\\n\".join(lines))\n\ntr = inputs.get(\"triage_results\", [])\ntr = list(tr) if isinstance(tr, list) else []\ntr.append(\n {\n \"host\": host,\n \"ssh_user\": connect_user or su,\n \"ssh_target\": ssh_target,\n \"ssh_host\": connect_host,\n \"success\": ok,\n \"verdict\": verdict,\n \"failure_category\": failure_category,\n \"inspect_attempts\": attempts,\n \"error\": err,\n \"per_host_md\": per_host_md,\n }\n)\noutputs[\"triage_results\"] = tr\n\nbatch_report_path = inputs.get(\"batch_report_path\", \"\")\nsection = (\n \"\\n## [{}/{}] `{}`\\n\\n\".format(idx1, nh, ssh_target)\n + \"- 单独报告: `{}`\\n\".format(per_host_md)\n + \"- 执行结果: {}\\n\".format(\"成功\" if ok else \"失败\")\n + \"- 判定结果: `{}`\\n\".format(verdict)\n + \"- 失败分类: `{}`\\n\".format(failure_category or \"\")\n + \"- 尝试次数: {}\\n\\n---\\n\".format(attempts)\n)\nif batch_report_path:\n with open(batch_report_path, \"a\", encoding=\"utf-8\") as f:\n f.write(section)\n\noutputs[\"last_host\"] = host\noutputs[\"last_ssh_target\"] = ssh_target\noutputs[\"last_success\"] = ok\noutputs[\"last_verdict\"] = verdict\noutputs[\"last_failure_category\"] = failure_category\noutputs[\"last_per_host_md\"] = per_host_md" + "description": "每台主机先做 SSH 预检；预检通过后再调 delegate_task(host-forensics-fast)。超时仅重试一次；循环态保留文件路径、执行状态、verdict 与失败分类。", + "code": "import os\nimport re\nimport time\n\nhosts = inputs.get(\"hosts\", [])\nidx = int(inputs.get(\"host_idx\", 0))\nhost = hosts[idx] if 0 <= idx < len(hosts) else \"\"\nhost = str(host).strip()\n\nsu = inputs.get(\"ssh_user\")\nif isinstance(su, str):\n su = su.strip()\nelse:\n su = \"\"\n\nconnect_host = host\nconnect_user = \"\"\nif \"@\" in host:\n user_part, host_part = host.split(\"@\", 1)\n connect_user = str(user_part).strip()\n connect_host = str(host_part).strip() or host\n ssh_target = host\n user_hint = (\n \"主机列表项已含 `user@host` 形式。SSH 工具调用必须使用：host=`\"\n + connect_host\n + \"`，username=`\"\n + connect_user\n + \"`。\"\n )\nelif su:\n connect_user = su\n connect_host = host\n ssh_target = su + \"@\" + host\n user_hint = (\n \"工作流已指定 `ssh_user`=`\"\n + su\n + \"`。SSH 工具调用必须使用：host=`\"\n + connect_host\n + \"`，username=`\"\n + connect_user\n + \"`。\"\n )\nelse:\n connect_host = host\n ssh_target = host\n user_hint = (\n \"未指定 `ssh_user`。SSH 工具调用请使用 host=`\"\n + connect_host\n + \"`，username 留空使用默认账户（一般为 root）。\"\n )\n\n\ndef _extract_verdict(markdown_text):\n if not markdown_text:\n return \"UNKNOWN\"\n match = re.search(\n r\"(?im)^\\s*\\*{0,2}Verdict\\*{0,2}\\s*:\\s*\"\n r\"(CLEAN|SUSPICIOUS|COMPROMISED|UNKNOWN)\\b\",\n markdown_text,\n )\n if not match:\n return \"UNKNOWN\"\n return str(match.group(1)).upper()\n\n\ndef _is_timeout_error(message):\n text = str(message or \"\").lower()\n return (\n \"timed out\" in text\n or \"timeout\" in text\n or \"超时\" in text\n or \"节点执行超时\" in text\n )\n\n\ndef _classify_error(message):\n text = str(message or \"\")\n lower = text.lower()\n if not text.strip():\n return \"unknown\"\n if \"permission denied\" in lower or \"auth failed\" in lower or \"authentication failed\" in lower:\n return \"auth_failed\"\n if \"host key verification failed\" in lower or \"host key\" in lower:\n return \"host_key_verification_failed\"\n if \"connection refused\" in lower:\n return \"connection_refused\"\n if \"no route to host\" in lower:\n return \"no_route_to_host\"\n if \"network is unreachable\" in lower:\n return \"network_unreachable\"\n if \"name or service not known\" in lower or \"could not resolve\" in lower or \"nodename nor servname provided\" in lower:\n return \"dns_resolution_failed\"\n if \"connection reset\" in lower:\n return \"connection_reset\"\n if \"broken pipe\" in lower or \"connection lost\" in lower or \"disconnect\" in lower:\n return \"connection_lost\"\n if \"kex\" in lower or \"key exchange\" in lower or \"protocol error\" in lower:\n return \"ssh_handshake_failed\"\n if _is_timeout_error(text):\n if \"connect\" in lower or \"ssh connection failed\" in lower:\n return \"connect_timeout\"\n return \"execution_timeout\"\n if \"ssh connection failed\" in lower:\n return \"ssh_connection_failed\"\n return \"unknown\"\n\n\nidx1 = idx + 1\nper_host_dir = str(inputs.get(\"per_host_dir\") or \"\").strip()\nif not per_host_dir:\n od = str(inputs.get(\"output_dir\") or \"\").strip()\n if od:\n per_host_dir = os.path.join(od, \"host_triage\")\n else:\n per_host_dir = os.path.join(os.path.dirname(inputs.get(\"batch_report_path\") or \".\") or \".\", \"host_triage\")\nos.makedirs(per_host_dir, exist_ok=True)\n\n\ndef _slug(s):\n t = re.sub(r\"[^0-9A-Za-z._@-]+\", \"_\", str(s)).strip(\"_\")\n t = t.replace(\"@\", \"_at_\")\n return (t[:56] if t else \"host\")\n\n\nnh = len(hosts)\nbase = \"{:04d}_{}\".format(idx1, _slug(ssh_target))\nper_host_md = os.path.join(per_host_dir, base + \".md\")\n\npreflight_timeout_s = 20\npreflight_res = tool.run_safe(\n \"ssh_host_cmd\",\n host=connect_host,\n username=(connect_user or None),\n command=\"echo FLOCKS_SSH_OK\",\n timeout=preflight_timeout_s,\n)\npreflight_output = preflight_res.get(\"output\") or preflight_res.get(\"text\") or \"\"\npreflight_error = preflight_res.get(\"error\") or \"\"\npreflight_ok = bool(preflight_res.get(\"success\")) and \"FLOCKS_SSH_OK\" in str(preflight_output)\n\ntext = \"\"\nok = False\nerr = \"\"\nverdict = \"UNKNOWN\"\nfailure_category = \"\"\nattempts = 0\n\nif preflight_ok:\n desc = \"Fast triage item \" + str(idx1)\n prompt = (\n \"你是 host-forensics-fast 工作模式：对下列目标执行 Linux 主机快速安全巡检（首轮研判）。\\n\"\n \"请使用 ssh_run_script，script_path 为 `.flocks/plugins/agents/host-forensics-fast/scripts/triage_fast.sh`。\\n\"\n + user_hint\n + \"\\n\\n本次 SSH 工具参数必须使用：\\n\"\n + \"- host: \"\n + connect_host\n + \"\\n\"\n + (\"- username: \" + connect_user + \"\\n\" if connect_user else \"- username: （留空，使用默认账户）\\n\")\n + \"\\n请输出简洁 Markdown：结论、可疑项、风险判断、后续建议。\"\n )\n while attempts < 2:\n attempts += 1\n res = tool.run_safe(\n \"delegate_task\",\n description=desc + \" attempt \" + str(attempts),\n prompt=prompt,\n subagent_type=\"host-forensics-fast\",\n )\n text = res.get(\"text\") or \"\"\n ok = bool(res.get(\"success\"))\n err = res.get(\"error\") or \"\"\n if ok:\n verdict = _extract_verdict(text)\n break\n if not _is_timeout_error(err) or attempts >= 2:\n failure_category = _classify_error(err)\n break\n time.sleep(3)\n if not ok and not failure_category:\n failure_category = _classify_error(err)\nelse:\n err = preflight_error or \"SSH preflight failed\"\n failure_category = _classify_error(err)\n\nlines = [\n \"# 单主机快速巡检结果\",\n \"\",\n \"- 批次内序号: {} / {}\".format(idx1, nh),\n \"- 列表项 host: `{}`\".format(host),\n \"- ssh_target: `{}`\".format(ssh_target),\n \"- ssh_host: `{}`\".format(connect_host),\n \"- ssh_user: `{}`\".format(connect_user or \"(default)\"),\n \"- success: {}\".format(ok),\n \"- verdict: {}\".format(verdict),\n \"- failure_category: {}\".format(failure_category or \"\"),\n \"- inspect_attempts: {}\".format(attempts),\n \"\",\n]\nif not preflight_ok:\n lines.extend(\n [\n \"## SSH 预检失败\",\n \"\",\n \"- 分类: `{}`\".format(failure_category),\n \"- 错误:\",\n \"\",\n \"```\",\n str(err),\n \"```\",\n \"\",\n ]\n )\nelif err:\n lines.extend([\"## 错误\", \"\", \"```\", str(err), \"```\", \"\"])\nelse:\n lines.extend([\"## 子 Agent 输出\", \"\", (text if text else \"_(无正文)_\"), \"\"])\nwith open(per_host_md, \"w\", encoding=\"utf-8\") as f:\n f.write(\"\\n\".join(lines))\n\ntr = inputs.get(\"triage_results\", [])\ntr = list(tr) if isinstance(tr, list) else []\ntr.append(\n {\n \"host\": host,\n \"ssh_user\": connect_user or su,\n \"ssh_target\": ssh_target,\n \"ssh_host\": connect_host,\n \"success\": ok,\n \"verdict\": verdict,\n \"failure_category\": failure_category,\n \"inspect_attempts\": attempts,\n \"error\": err,\n \"per_host_md\": per_host_md,\n }\n)\noutputs[\"triage_results\"] = tr\n\nbatch_report_path = inputs.get(\"batch_report_path\", \"\")\nsection = (\n \"\\n## [{}/{}] `{}`\\n\\n\".format(idx1, nh, ssh_target)\n + \"- 单独报告: `{}`\\n\".format(per_host_md)\n + \"- 执行结果: {}\\n\".format(\"成功\" if ok else \"失败\")\n + \"- 判定结果: `{}`\\n\".format(verdict)\n + \"- 失败分类: `{}`\\n\".format(failure_category or \"\")\n + \"- 尝试次数: {}\\n\\n---\\n\".format(attempts)\n)\nif batch_report_path:\n with open(batch_report_path, \"a\", encoding=\"utf-8\") as f:\n f.write(section)\n\noutputs[\"last_host\"] = host\noutputs[\"last_ssh_target\"] = ssh_target\noutputs[\"last_success\"] = ok\noutputs[\"last_verdict\"] = verdict\noutputs[\"last_failure_category\"] = failure_category\noutputs[\"last_per_host_md\"] = per_host_md" }, { "id": "advance_index", diff --git a/.flocks/plugins/workflows/loop_host_forensics_fast/workflow.md b/.flocks/plugins/workflows/loop_host_forensics_fast/workflow.md index e124becde..7ae3e4bf8 100644 --- a/.flocks/plugins/workflows/loop_host_forensics_fast/workflow.md +++ b/.flocks/plugins/workflows/loop_host_forensics_fast/workflow.md @@ -51,14 +51,14 @@ ### 3. 单台巡检（`inspect_host`） -- 工具/模型：Python + `ssh_host_cmd` 预检 + `task`（`subagent_type=host-forensics-fast`） +- 工具/模型：Python + `ssh_host_cmd` 预检 + `delegate_task`（`subagent_type=host-forensics-fast`） - 输入：`hosts`、`host_idx`、`ssh_user`、`per_host_dir`、`batch_report_path`、`triage_results` - 处理逻辑： - 取当前 `hosts[host_idx]`，计算 `ssh_target`，并归一化出 `ssh_host` / `ssh_user`。 - 先用 `ssh_host_cmd("echo FLOCKS_SSH_OK")` 做轻量 SSH 预检。 - 若预检失败：按错误文本归类（如 `auth_failed`、`connect_timeout`、`connection_refused` 等），直接写入索引与单机报告。 - 若预检通过：构造 prompt，明确要求子 Agent 调用 SSH 工具时分别传 `host` 和 `username`。 - - 调用 `tool.run_safe('task', ...)` 执行巡检；若仅因超时失败，则自动重试 1 次。 + - 调用 `tool.run_safe('delegate_task', ...)` 执行巡检；若仅因超时失败，则自动重试 1 次。 - 将本轮完整输出立即写入 `host_triage/NNNN_slug.md`。 - 从子 Agent 输出中提取 `Verdict`，未识别时回退为 `UNKNOWN`。 - 向 `triage_results` 仅追加轻量字段：`{host, ssh_user, ssh_target, ssh_host, success, verdict, failure_category, inspect_attempts, error, per_host_md}`。 diff --git a/flocks/cli/commands/import_.py b/flocks/cli/commands/import_.py index fe04b4845..2dd701c8a 100644 --- a/flocks/cli/commands/import_.py +++ b/flocks/cli/commands/import_.py @@ -98,6 +98,17 @@ def _normalize_part_data( metadata = normalized.get("metadata") metadata_dict = metadata if isinstance(metadata, dict) else {} + if part_type == "subtask": + normalized["type"] = "text" + normalized["text"] = "" + normalized["ignored"] = True + normalized["metadata"] = { + **metadata_dict, + "legacyPartType": "subtask", + } + part_type = "text" + metadata_dict = normalized["metadata"] + if "content" in normalized and "text" not in normalized: normalized["text"] = normalized.get("content", "") @@ -135,10 +146,6 @@ def _normalize_part_data( ) elif part_type == "agent": normalized.setdefault("name", metadata_dict.get("name") or normalized.get("content") or "agent") - elif part_type == "subtask": - normalized.setdefault("prompt", metadata_dict.get("prompt") or normalized.get("content", "")) - normalized.setdefault("description", metadata_dict.get("description") or "") - normalized.setdefault("agent", metadata_dict.get("agent") or "agent") elif part_type == "retry": normalized.setdefault("attempt", metadata_dict.get("attempt") or 1) normalized.setdefault("error", metadata_dict.get("error") or {}) diff --git a/flocks/command/command.py b/flocks/command/command.py index b75ee92d8..32e06b177 100644 --- a/flocks/command/command.py +++ b/flocks/command/command.py @@ -24,7 +24,6 @@ class CommandDef: template: str agent: Optional[str] = None model: Optional[str] = None - subtask: Optional[bool] = None hidden: bool = False aliases: Tuple[str, ...] = field(default_factory=tuple) visible_surfaces: Tuple[CommandSurface, ...] = ("webui", "tui", "acp", "cli") diff --git a/flocks/config/config.py b/flocks/config/config.py index 0f63e1d13..6bef59b5c 100644 --- a/flocks/config/config.py +++ b/flocks/config/config.py @@ -30,20 +30,38 @@ class PermissionAction(str, Enum): _LEGACY_TODO_TOOL_NAMES = {"todowrite", "todoread"} +_LEGACY_PERMISSION_TOOL_NAMES = { + **{name: "todo" for name in _LEGACY_TODO_TOOL_NAMES}, + "task": "delegate_task", +} def _canonical_permission_tool_name(tool: str) -> str: - if tool in _LEGACY_TODO_TOOL_NAMES: - return "todo" - return tool + return _LEGACY_PERMISSION_TOOL_NAMES.get(tool, tool) def _merge_permission_action(existing: Any, incoming: Any) -> Any: """Merge duplicate legacy permission names conservatively.""" - existing_value = existing.value if hasattr(existing, "value") else existing - incoming_value = incoming.value if hasattr(incoming, "value") else incoming - if existing_value == PermissionAction.DENY.value or incoming_value == PermissionAction.DENY.value: + def contains_deny(value: Any) -> bool: + raw_value = value.value if hasattr(value, "value") else value + if raw_value == PermissionAction.DENY.value: + return True + if isinstance(raw_value, dict): + return any(contains_deny(item) for item in raw_value.values()) + return False + + if contains_deny(existing) or contains_deny(incoming): return PermissionAction.DENY + if isinstance(existing, dict) and isinstance(incoming, dict): + merged = dict(existing) + for pattern, action in incoming.items(): + if pattern in merged: + merged[pattern] = _merge_permission_action( + merged[pattern], action, + ) + else: + merged[pattern] = action + return merged return existing if existing is not None else incoming @@ -79,11 +97,11 @@ class PermissionConfig(BaseModel): @model_validator(mode="before") @classmethod - def migrate_legacy_todo_permissions(cls, data): + def migrate_legacy_permissions(cls, data): if not isinstance(data, dict): return data migrated = dict(data) - for legacy_name in _LEGACY_TODO_TOOL_NAMES: + for legacy_name in _LEGACY_PERMISSION_TOOL_NAMES: if legacy_name in migrated: _assign_permission(migrated, legacy_name, migrated.pop(legacy_name)) return migrated @@ -158,7 +176,6 @@ class CommandConfig(BaseModel): description: Optional[str] = None agent: Optional[str] = None model: Optional[str] = None - subtask: Optional[bool] = None # ==================== Provider Configuration ==================== diff --git a/flocks/permission/helpers.py b/flocks/permission/helpers.py index 60206ebd7..5ea907711 100644 --- a/flocks/permission/helpers.py +++ b/flocks/permission/helpers.py @@ -6,6 +6,48 @@ Ruleset = List[PermissionRule] +_LEGACY_PERMISSION_NAMES = { + "task": "delegate_task", + "todowrite": "todo", + "todoread": "todo", +} + + +def _merge_legacy_permission(existing: Any, incoming: Any) -> Any: + """Merge aliases without allowing a legacy deny to become an allow.""" + def contains_deny(value: Any) -> bool: + raw_value = getattr(value, "value", value) + if raw_value == "deny": + return True + if isinstance(raw_value, dict): + return any(contains_deny(item) for item in raw_value.values()) + return False + + if contains_deny(existing) or contains_deny(incoming): + return "deny" + if isinstance(existing, dict) and isinstance(incoming, dict): + merged = dict(existing) + for pattern, action in incoming.items(): + if pattern in merged: + merged[pattern] = _merge_legacy_permission( + merged[pattern], action, + ) + else: + merged[pattern] = action + return merged + return existing + + +def _canonicalize_permission_config(config: Dict[str, Any]) -> Dict[str, Any]: + canonical: Dict[str, Any] = {} + for key, value in config.items(): + name = _LEGACY_PERMISSION_NAMES.get(key, key) + if name in canonical: + canonical[name] = _merge_legacy_permission(canonical[name], value) + else: + canonical[name] = value + return canonical + def from_config(permission_config: Union[Dict[str, Any], BaseModel]) -> Ruleset: """ @@ -22,6 +64,8 @@ def from_config(permission_config: Union[Dict[str, Any], BaseModel]) -> Ruleset: else: return ruleset + config_dict = _canonicalize_permission_config(config_dict) + for key, value in config_dict.items(): if isinstance(value, str) or isinstance(value, PermissionLevel): ruleset.append(PermissionRule( diff --git a/flocks/server/routes/misc.py b/flocks/server/routes/misc.py index 2605ced40..230e0af2d 100644 --- a/flocks/server/routes/misc.py +++ b/flocks/server/routes/misc.py @@ -151,7 +151,6 @@ async def list_commands() -> List[Dict[str, Any]]: "template": cmd.template, "agent": cmd.agent, "model": cmd.model, - "subtask": cmd.subtask, "hidden": cmd.hidden, "aliases": list(cmd.aliases), "visible_surfaces": list(cmd.visible_surfaces), @@ -192,7 +191,6 @@ async def get_command(name: str) -> Dict[str, Any]: "template": cmd.template, "agent": cmd.agent, "model": cmd.model, - "subtask": cmd.subtask, "hidden": cmd.hidden, "aliases": list(cmd.aliases), "visible_surfaces": list(cmd.visible_surfaces), @@ -241,4 +239,3 @@ async def list_experimental_resources() -> Dict[str, Any]: # Return empty dict - resources are not implemented yet return {} - diff --git a/flocks/server/routes/session.py b/flocks/server/routes/session.py index c0dbf61f5..236d7e594 100644 --- a/flocks/server/routes/session.py +++ b/flocks/server/routes/session.py @@ -2091,15 +2091,6 @@ class AgentPartInput(BaseModel): name: str = Field(..., description="Agent name") -class SubtaskPartInput(BaseModel): - """Subtask part input for API compatibility""" - type: Literal["subtask"] = "subtask" - id: Optional[str] = Field(None, description="Part ID") - agent: str = Field(..., description="Agent name") - prompt: str = Field(..., description="Subtask prompt") - description: Optional[str] = Field(None, description="Subtask description") - - class PromptRequest(BaseModel): """ Request to send a prompt/message diff --git a/flocks/server/routes/skill.py b/flocks/server/routes/skill.py index 28d0e3d7f..4d16ee70c 100644 --- a/flocks/server/routes/skill.py +++ b/flocks/server/routes/skill.py @@ -164,7 +164,6 @@ class CommandResponse(BaseModel): template: str = Field(..., description="Command template") agent: Optional[str] = Field(None, description="Preferred agent") model: Optional[str] = Field(None, description="Preferred model") - subtask: Optional[bool] = Field(None, description="Run as subtask") hidden: bool = Field(False, description="Hidden from UI") aliases: List[str] = Field(default_factory=list, description="Alternate slash aliases") visible_surfaces: List[str] = Field(default_factory=list, description="Surfaces where the command is visible") @@ -182,7 +181,6 @@ def _command_to_response(cmd: CommandInfo) -> CommandResponse: template=cmd.template, agent=cmd.agent, model=cmd.model, - subtask=cmd.subtask, hidden=cmd.hidden, aliases=list(cmd.aliases), visible_surfaces=list(cmd.visible_surfaces), diff --git a/flocks/session/__init__.py b/flocks/session/__init__.py index 54840e2fb..fe8ac7bad 100644 --- a/flocks/session/__init__.py +++ b/flocks/session/__init__.py @@ -30,7 +30,6 @@ ReasoningPart, PatchPart, AgentPart, - SubtaskPart, ) from flocks.session.prompt import SessionPrompt, SystemPrompt, ContextInfo from flocks.session.lifecycle.compaction import SessionCompaction, CompactionResult, CompactionPolicy, ContextTier @@ -53,11 +52,6 @@ ReminderConfig, ReminderContext, ) -from flocks.session.features.subtask import ( - SessionSubtask, - SubtaskInfo, - SubtaskResult, -) from flocks.session.lifecycle.revert import ( SessionRevertManager, RevertInput, @@ -93,7 +87,6 @@ "ReasoningPart", "PatchPart", "AgentPart", - "SubtaskPart", # Prompt "SessionPrompt", "SystemPrompt", @@ -119,10 +112,6 @@ "SessionReminders", "ReminderConfig", "ReminderContext", - # Subtask - "SessionSubtask", - "SubtaskInfo", - "SubtaskResult", # Revert "SessionRevertManager", "RevertInput", diff --git a/flocks/session/context_usage.py b/flocks/session/context_usage.py index ca35e09d1..376b41044 100644 --- a/flocks/session/context_usage.py +++ b/flocks/session/context_usage.py @@ -23,7 +23,7 @@ log = Log.create(service="context-usage") UsageSource = Literal["observed", "estimated"] -DELEGATION_TOOLS = {"delegate_task", "task"} +DELEGATION_TOOLS = {"delegate_task"} ZERO_VISIBLE_SEGMENTS = {"agentDelegation"} @@ -441,8 +441,8 @@ async def _estimate_message_breakdown(session_id: str, messages: List[Any]) -> t if part_type in {"reasoning", "thinking"}: tokens_by_key["reasoning"] += SessionPrompt.count_tokens(_field_value(part, "text", "") or "") continue - if part_type in {"agent", "subtask"}: - tokens_by_key["agentDelegation"] += _estimate_subtask_part_tokens(part) + if part_type == "agent": + tokens_by_key["agentDelegation"] += _estimate_agent_part_tokens(part) continue if part_type != "tool": continue @@ -499,7 +499,7 @@ def _context_key_for_tool(tool_name: str) -> str: return "tools" -def _estimate_subtask_part_tokens(part: Any) -> int: +def _estimate_agent_part_tokens(part: Any) -> int: total = 0 for field in ("prompt", "description", "name"): value = _field_value(part, field, "") diff --git a/flocks/session/features/subtask.py b/flocks/session/features/subtask.py deleted file mode 100644 index f1231cea5..000000000 --- a/flocks/session/features/subtask.py +++ /dev/null @@ -1,62 +0,0 @@ -""" -Session Subtask data models. - -Note: The SessionSubtask business logic has been removed as it was dead code. -The live subtask execution path is session_loop.py::_execute_subtask(), which -handles the full lifecycle inline without using this module. - -These data classes are kept because they are exported from session/__init__.py -and may be referenced by external consumers. -""" - -from dataclasses import dataclass, field -from datetime import datetime -from typing import Any, Dict, Optional - - -@dataclass -class SubtaskInfo: - """Information about a subtask""" - id: str - parent_session_id: str - child_session_id: Optional[str] = None - task_description: str = "" - agent: Optional[str] = None - model: Optional[str] = None - status: str = "pending" # pending, running, completed, error - result: Optional[str] = None - error: Optional[str] = None - created_at: int = field(default_factory=lambda: int(datetime.now().timestamp() * 1000)) - completed_at: Optional[int] = None - - -@dataclass -class SubtaskResult: - """Result of subtask execution""" - subtask_id: str - success: bool - output: str - error: Optional[str] = None - metadata: Dict[str, Any] = field(default_factory=dict) - - -# Minimal stub so imports of SessionSubtask don't break existing code. -class SessionSubtask: - """Subtask manager stub — business logic removed (was dead code). - - The active execution path is SessionLoop._execute_subtask() in session_loop.py. - """ - - @classmethod - async def execute_subtask(cls, *args, **kwargs) -> SubtaskResult: - raise NotImplementedError( - "SessionSubtask.execute_subtask() is deprecated. " - "Subtask execution is handled by SessionLoop._execute_subtask()." - ) - - -__all__ = [ - "SessionSubtask", - "SubtaskInfo", - "SubtaskResult", -] diff --git a/flocks/session/message.py b/flocks/session/message.py index a8c4b3e20..f3ac17bef 100644 --- a/flocks/session/message.py +++ b/flocks/session/message.py @@ -259,21 +259,6 @@ class AgentPart(BaseModel): source: Optional[Dict[str, Any]] = Field(None, description="Source information") -class SubtaskPart(BaseModel): - """Subtask/subagent part - Flocks compatible""" - model_config = ConfigDict(populate_by_name=True, by_alias=True) - - id: str = Field(default_factory=lambda: Identifier.ascending("part")) - sessionID: str = Field(..., description="Session ID") - messageID: str = Field(..., description="Message ID") - type: Literal["subtask"] = "subtask" - prompt: str = Field(..., description="Task prompt") - description: str = Field(..., description="Task description") - agent: str = Field(..., description="Agent name") - model: Optional[Dict[str, str]] = Field(None, description="Model configuration") - command: Optional[str] = Field(None, description="Command to execute") - - class RetryPart(BaseModel): """Retry part - Flocks compatible""" model_config = ConfigDict(populate_by_name=True, by_alias=True) @@ -301,7 +286,6 @@ class CompactionPart(BaseModel): # Union type for all parts - matches Flocks MessageV2.Part PartType = Union[ TextPart, - SubtaskPart, ReasoningPart, FilePart, ToolPart, @@ -1122,6 +1106,18 @@ def _normalize_part_data( metadata = normalized.get("metadata") metadata_dict = metadata if isinstance(metadata, dict) else {} + if part_type == "subtask": + normalized["type"] = "text" + normalized["text"] = "" + normalized["ignored"] = True + normalized["metadata"] = { + **metadata_dict, + "legacyPartType": "subtask", + } + part_type = "text" + metadata = normalized["metadata"] + metadata_dict = metadata + if "content" in normalized and "text" not in normalized: normalized["text"] = normalized.get("content", "") @@ -1174,10 +1170,6 @@ def _normalize_part_data( normalized.setdefault("tokens", metadata_dict.get("tokens") or cls._default_token_usage()) elif part_type == "agent": normalized.setdefault("name", metadata_dict.get("name") or normalized.get("content") or "agent") - elif part_type == "subtask": - normalized.setdefault("prompt", metadata_dict.get("prompt") or normalized.get("content", "")) - normalized.setdefault("description", metadata_dict.get("description") or "") - normalized.setdefault("agent", metadata_dict.get("agent") or "agent") elif part_type == "retry": normalized.setdefault("attempt", metadata_dict.get("attempt") or 1) normalized.setdefault("error", metadata_dict.get("error") or {}) @@ -1255,7 +1247,6 @@ def deserialize_part( 'step-start': StepStartPart, 'step-finish': StepFinishPart, 'agent': AgentPart, - 'subtask': SubtaskPart, 'retry': RetryPart, 'compaction': CompactionPart, } diff --git a/flocks/session/prompt/anthropic-20250930.txt b/flocks/session/prompt/anthropic-20250930.txt index a8ada5ede..ec4e65c94 100644 --- a/flocks/session/prompt/anthropic-20250930.txt +++ b/flocks/session/prompt/anthropic-20250930.txt @@ -122,10 +122,10 @@ I've found existing rules. Let me mark the first todo as in_progress and start d Users may configure 'hooks', shell commands that execute in response to events like tool calls, in settings. Treat feedback from hooks, including , as coming from the user. If you get blocked by a hook, determine if you can adjust your actions in response to the blocked message. If not, ask the user to check their hooks configuration. # Tool usage policy -- You should proactively use the Task tool with specialized agents when the task at hand matches the agent's description. +- You should proactively use `delegate_task` with specialized agents when the task at hand matches the agent's description. - Tool results and user messages may include tags. tags contain useful information and reminders. They are automatically added by the system, and bear no direct relation to the specific tool results or user messages in which they appear. - When WebFetch returns a message about a redirect to a different host, you should immediately make a new WebFetch request with the redirect URL provided in the response. -- If the user specifies that they want you to run tools "in parallel", you MUST send a single message with multiple tool use content blocks. For example, if you need to launch multiple agents in parallel, send a single message with multiple Task tool calls. +- If the user specifies that they want you to run tools "in parallel", you MUST send a single message with multiple tool use content blocks. For example, if you need to launch multiple agents in parallel, send a single message with multiple `delegate_task` calls. - Use specialized tools instead of bash commands when possible, as this provides a better user experience. For file operations, use dedicated tools: Read for reading files instead of cat/head/tail, Edit for editing instead of sed/awk, and Write for creating files instead of cat with heredoc or echo redirection. Reserve bash tools exclusively for actual system commands and terminal operations that require shell execution. NEVER use bash echo or other command-line tools to communicate thoughts, explanations, or instructions to the user. Output all communication directly in your response text instead. diff --git a/flocks/session/prompt/anthropic.txt b/flocks/session/prompt/anthropic.txt index 655551871..2455ef584 100644 --- a/flocks/session/prompt/anthropic.txt +++ b/flocks/session/prompt/anthropic.txt @@ -66,22 +66,22 @@ I've found existing SIGMA rules. Let me mark the first todo as in_progress and s # Tool usage policy -- You should proactively use the Task tool with specialized agents when the task at hand matches the agent's description. +- You should proactively use `delegate_task` with specialized agents when the task at hand matches the agent's description. - Tool results and user messages may include tags. tags contain useful information and reminders. They are automatically added by the system, and bear no direct relation to the specific tool results or user messages in which they appear. - When WebFetch returns a message about a redirect to a different host, you should immediately make a new WebFetch request with the redirect URL provided in the response. - You can call multiple tools in a single response. If you intend to call multiple tools and there are no dependencies between them, make all independent tool calls in parallel. Maximize use of parallel tool calls where possible to increase efficiency. However, if some tool calls depend on previous calls to inform dependent values, do NOT call these tools in parallel and instead call them sequentially. For instance, if one operation must complete before another starts, run these operations sequentially instead. Never use placeholders or guess missing parameters in tool calls. -- If the user specifies that they want you to run tools "in parallel", you MUST send a single message with multiple tool use content blocks. For example, if you need to launch multiple agents in parallel, send a single message with multiple Task tool calls. +- If the user specifies that they want you to run tools "in parallel", you MUST send a single message with multiple tool use content blocks. For example, if you need to launch multiple agents in parallel, send a single message with multiple `delegate_task` calls. - Use specialized tools instead of bash commands when possible, as this provides a better user experience. For file operations, use dedicated tools: Read for reading files instead of cat/head/tail, Edit for editing instead of sed/awk, and Write for creating files instead of cat with heredoc or echo redirection. Reserve bash tools exclusively for actual system commands and terminal operations that require shell execution. NEVER use bash echo or other command-line tools to communicate thoughts, explanations, or instructions to the user. Output all communication directly in your response text instead. -- VERY IMPORTANT: When exploring security logs, configurations, or investigating incidents that require context gathering, use the Task tool for complex searches instead of running multiple search commands directly. +- VERY IMPORTANT: When exploring security logs, configurations, or investigating incidents that require context gathering, use `delegate_task` for complex searches instead of running multiple search commands directly. - IMPORTANT: Always respond in the same language as the user. user: Where are authentication failures logged in our application? -assistant: [Uses the Task tool to find authentication logging locations instead of using Glob or Grep directly] +assistant: [Uses `delegate_task` to find authentication logging locations instead of using Glob or Grep directly] user: Find all places where user input is processed without validation -assistant: [Uses the Task tool to comprehensively search for input validation gaps] +assistant: [Uses `delegate_task` to comprehensively search for input validation gaps] IMPORTANT: Always use `todo(action="write")` to plan and track tasks throughout the conversation. diff --git a/flocks/session/prompt_strings.py b/flocks/session/prompt_strings.py index 8ccb18b5d..2554f9244 100644 --- a/flocks/session/prompt_strings.py +++ b/flocks/session/prompt_strings.py @@ -146,20 +146,20 @@ assistant: "Here is the relevant function: " - Since the user is greeting, use the Task tool to launch the greeting-responder agent to respond with a friendly joke. + Since the user is greeting, use delegate_task to launch the greeting-responder agent to respond with a friendly joke. assistant: "Now let me use the code-reviewer agent to review the code" - Context: User is creating an agent to respond to the word "hello" with a friendly jok. user: "Hello" - assistant: "I'm going to use the Task tool to launch the greeting-responder agent to respond with a friendly joke" + assistant: "I'm going to use delegate_task to launch the greeting-responder agent to respond with a friendly joke" Since the user is greeting, use the greeting-responder agent to respond with a friendly joke. - If the user mentioned or implied that the agent should be used proactively, you should include examples of this. -- NOTE: Ensure that in the examples, you are making the assistant use the Agent tool and not simply respond directly to the task. +- NOTE: Ensure that in the examples, you are making the assistant use delegate_task and not simply respond directly to the task. Your output must be a valid JSON object with exactly these fields: { diff --git a/flocks/session/session_loop.py b/flocks/session/session_loop.py index 53af1e89f..ebf005a79 100644 --- a/flocks/session/session_loop.py +++ b/flocks/session/session_loop.py @@ -1319,7 +1319,7 @@ async def _run_loop( 1. Get messages and analyze (lastUser, lastAssistant, lastFinished) 2. Check exit conditions 3. Generate title on first step - 4. Check for pending tasks (subtask/compaction) + 4. Check for pending compaction 5. Check context overflow (compaction before step) 6. Process step (call LLM + tools) 7. Loop until complete @@ -1374,7 +1374,7 @@ async def _run_loop( last_user: Optional[MessageInfo] = None last_assistant: Optional[MessageInfo] = None last_finished: Optional[MessageInfo] = None - tasks: List[tuple[str, Any]] = [] # (type, part) - compaction or subtask + tasks: List[tuple[str, Any]] = [] # (type, part) - compaction only scan_started_at = asyncio.get_event_loop().time() for msg in reversed(messages): @@ -1394,14 +1394,12 @@ async def _run_loop( if last_user and last_finished: break - # Collect pending tasks before lastFinished + # Collect pending compaction before lastFinished if not last_finished: parts = await Message.parts(msg.id, ctx.session.id) for part in parts: if part.type == "compaction": tasks.append(("compaction", part)) - elif part.type == "subtask": - tasks.append(("subtask", part)) log.debug("loop.message_scan_complete", { "session_id": ctx.session.id, "step": ctx.step, @@ -1497,25 +1495,11 @@ async def _run_loop( except Exception as e: log.error("loop.title_generation.error", {"error": str(e)}) - # Check for pending tasks (matching TUI lines 314-493) + # Check for pending compaction if tasks: task_type, task_part = tasks.pop() - - # Handle pending subtask (matching TUI lines 316-481) - if task_type == "subtask": - log.info("loop.subtask_detected", { - "session_id": ctx.session.id, - "step": ctx.step, - }) - - # Execute subtask using tool execution - await cls._execute_subtask(ctx, last_user, task_part) - - # Continue to next iteration - continue - - # Handle pending compaction (matching TUI lines 483-494) - elif task_type == "compaction": + + if task_type == "compaction": log.info("loop.compaction_pending", { "session_id": ctx.session.id, "step": ctx.step, @@ -2382,192 +2366,6 @@ async def _check_reminders( if reminder_msg and callbacks.on_reminder: await callbacks.on_reminder(await Message.get_text_content(reminder_msg)) - @classmethod - async def _execute_subtask( - cls, - ctx: LoopContext, - last_user: MessageInfo, - task_part: Any, - ) -> None: - """ - Execute subtask (matching TUI lines 316-481) - - 完全匹配 TUI 的 subtask 执行流程: - 1. 创建 assistant message - 2. 创建 tool part (Task tool) - 3. 执行 Task tool - 4. 更新 part 状态 - 5. 创建 synthetic user message - """ - from flocks.tool.registry import ToolRegistry - from flocks.agent.registry import Agent - - # Extract subtask information from part - agent_name = getattr(task_part, 'agent', 'hephaestus') - prompt = getattr(task_part, 'prompt', '') - description = getattr(task_part, 'description', '') - command = getattr(task_part, 'command', None) - model_info = getattr(task_part, 'model', None) - - # Get agent - agent = await Agent.get(agent_name) or await Agent.get("rex") - - # Determine model - if model_info: - provider_id = model_info.get('providerID', ctx.provider_id) - model_id = model_info.get('modelID', ctx.model_id) - else: - provider_id = ctx.provider_id - model_id = ctx.model_id - - # Create assistant message for subtask - assistant_msg = await Message.create( - session_id=ctx.session.id, - role=MessageRole.ASSISTANT, - content="", - agent=agent_name, - model=model_id, - provider=provider_id, - parent_id=last_user.id, - ) - - # Create tool part for Task - tool_call_id = Identifier.create("call") - from flocks.session.message import ToolPart, ToolStateRunning - - tool_part = ToolPart( - id=Identifier.ascending("part"), - sessionID=ctx.session.id, - messageID=assistant_msg.id, - type="tool", - callID=tool_call_id, - tool="task", - state=ToolStateRunning( - status="running", - input={ - "prompt": prompt, - "description": description, - "subagent_type": agent_name, - "command": command, - }, - time={"start": int(datetime.now().timestamp() * 1000)}, - ), - ) - - # Add part to message - await Message.add_part(ctx.session.id, assistant_msg.id, tool_part) - - # Get Task tool - task_tool = ToolRegistry.get("task") - if not task_tool: - log.error("loop.subtask.task_tool_not_found", {"session_id": ctx.session.id}) - return - - # Execute Task tool - task_args = { - "prompt": prompt, - "description": description, - "subagent_type": agent_name, - "command": command, - } - - # Create tool context - from flocks.tool.registry import ToolContext - - tool_ctx = ToolContext( - session_id=ctx.session.id, - message_id=assistant_msg.id, - agent=agent_name, - abort_event=ctx.abort_event, - ) - - execution_error: Optional[Exception] = None - result = None - - try: - result = await task_tool.execute(tool_ctx, **task_args) - except Exception as e: - execution_error = e - log.error("loop.subtask.execution_failed", { - "error": str(e), - "agent": agent_name, - "description": description, - }) - - # Update message finish - await Message.update(ctx.session.id, assistant_msg.id, finish="tool-calls") - - # Update tool part status - from flocks.session.message import ToolStateCompleted, ToolStateError - - if result: - # Create completed state - completed_state = ToolStateCompleted( - status="completed", - input={ - "prompt": prompt, - "description": description, - "subagent_type": agent_name, - "command": command, - }, - output=result.output if hasattr(result, 'output') else str(result), - title=result.title if hasattr(result, 'title') else None, - metadata=result.metadata if hasattr(result, 'metadata') else {}, - time={ - "start": tool_part.state.time.get("start"), - "end": int(datetime.now().timestamp() * 1000), - }, - ) - await Message.update_part( - session_id=ctx.session.id, - message_id=assistant_msg.id, - part_id=tool_part.id, - state=completed_state, - ) - else: - # Create error state - error_msg = str(execution_error) if execution_error else "Tool execution failed" - error_state = ToolStateError( - status="error", - error=f"Tool execution failed: {error_msg}", - time={ - "start": tool_part.state.time.get("start"), - "end": int(datetime.now().timestamp() * 1000), - }, - metadata={}, - input={ - "prompt": prompt, - "description": description, - "subagent_type": agent_name, - "command": command, - }, - ) - await Message.update_part( - session_id=ctx.session.id, - message_id=assistant_msg.id, - part_id=tool_part.id, - state=error_state, - ) - - # Create synthetic user message (matching TUI lines 457-478) - # This prevents reasoning models from erroring due to missing user messages - synthetic_user_msg = await Message.create( - session_id=ctx.session.id, - role=MessageRole.USER, - content="Summarize the task tool output above and continue with your task.", - agent=last_user.agent if hasattr(last_user, 'agent') else agent_name, - model=last_user.model if hasattr(last_user, 'model') else model_id, - provider=last_user.provider if hasattr(last_user, 'provider') else provider_id, - synthetic=True, - ) - - log.info("loop.subtask.completed", { - "session_id": ctx.session.id, - "agent": agent_name, - "success": result is not None, - }) - - # Export __all__ = [ diff --git a/flocks/tool/agent/delegate_task.py b/flocks/tool/agent/delegate_task.py index eb8ea2195..c467f0db3 100644 --- a/flocks/tool/agent/delegate_task.py +++ b/flocks/tool/agent/delegate_task.py @@ -217,7 +217,7 @@ def _derive_task_description( - Background subagent execution is disabled. Do not set run_in_background=true. - Foreground execution is always used: the tool waits for completion and returns results inline. - For independent parallel work needed this turn, emit multiple sibling - foreground delegate_task/task tool calls in the same assistant response. + foreground delegate_task tool calls in the same assistant response. The runtime executes them concurrently and the webui renders each as its own DelegateTaskCard. @@ -231,6 +231,7 @@ def _derive_task_description( name="delegate_task", description=DESCRIPTION, category=ToolCategory.SYSTEM, + native=True, parameters=[ ToolParameter( name="load_skills", @@ -283,9 +284,8 @@ async def delegate_task_tool( load_skills: Optional[List[str]] = None, description: Optional[str] = None, # Internal-only: not exposed in the public schema. The registry rejects - # `run_in_background=True` at the schema layer for any caller, but legacy - # in-process call paths (e.g. `task.py` alias) may still pass it through. - # This guard is the second line of defense. + # `run_in_background=True` at the schema layer for any caller. This guard + # also protects direct in-process callers that bypass the registry. run_in_background: bool = False, subagent_type: Optional[str] = None, session_id: Optional[str] = None, @@ -297,7 +297,7 @@ async def delegate_task_tool( success=False, error=( "Background subagent execution is disabled. " - "Use foreground delegate_task/task calls; emit multiple sibling calls " + "Use foreground delegate_task calls; emit multiple sibling calls " "in the same assistant turn for parallel work." ), ) diff --git a/flocks/tool/catalog.py b/flocks/tool/catalog.py index 40cffccb2..578c71422 100644 --- a/flocks/tool/catalog.py +++ b/flocks/tool/catalog.py @@ -37,7 +37,6 @@ class ToolCatalogMetadata(BaseModel): "webfetch": ["web", "http-fetch"], "websearch": ["web", "research"], "delegate_task": ["agent", "delegation"], - "task": ["agent", "delegation"], "schedule_task": ["scheduled-task", "scheduler-management"], "todo": ["task-management", "progress-tracking"], "run_workflow": ["workflow", "execution"], diff --git a/flocks/tool/code/grep.py b/flocks/tool/code/grep.py index 53fa278d0..371ff8b15 100644 --- a/flocks/tool/code/grep.py +++ b/flocks/tool/code/grep.py @@ -38,7 +38,7 @@ - Returns file paths and line numbers with at least one match sorted by modification time - Use this tool when you need to find files containing specific patterns - If you need to identify/count the number of matches within files, use the Bash tool with `rg` (ripgrep) directly. Do NOT use `grep`. -- When you are doing an open-ended search that may require multiple rounds of globbing and grepping, use the Task tool instead""" +- When you are doing an open-ended search that may require multiple rounds of globbing and grepping, use delegate_task instead""" def find_ripgrep() -> Optional[str]: diff --git a/flocks/tool/truncation.py b/flocks/tool/truncation.py index bea3e1099..4c749d0d1 100644 --- a/flocks/tool/truncation.py +++ b/flocks/tool/truncation.py @@ -100,7 +100,7 @@ def truncate_output( max_lines: Maximum number of lines to keep. max_bytes: Maximum byte size to keep. direction: "head" keeps the first N lines, "tail" keeps the last N. - has_task_tool: Whether the current agent can delegate via task tool. + has_task_tool: Whether the current agent can use delegate_task. Returns: TruncateResult with (possibly truncated) content. @@ -163,7 +163,7 @@ def truncate_output( hint = ( f"The tool call succeeded but the output was truncated. " f"Full output saved to: {filepath_str}\n" - f"Use the Task tool to have explore agent process this file with Grep and Read " + f"Use delegate_task to have an explore agent process this file with Grep and Read " f"(with offset/limit). Do NOT read the full file yourself - delegate to save context." ) elif filepath_str: diff --git a/flocks/utils/id.py b/flocks/utils/id.py index 7257ac0db..a7d682582 100644 --- a/flocks/utils/id.py +++ b/flocks/utils/id.py @@ -25,7 +25,6 @@ "call", # cal "step", # stp "agent", # agt - "subtask", # stk "event", # evt "tqref", # tqr "chbind", # chb (channel session binding) @@ -54,7 +53,6 @@ class Identifier: "call": "cal", "step": "stp", "agent": "agt", - "subtask": "stk", "event": "evt", "tqref": "tqr", "task": "tsk", diff --git a/flocks/workflow/tool_context.py b/flocks/workflow/tool_context.py index 8526a410e..8982cbc4b 100644 --- a/flocks/workflow/tool_context.py +++ b/flocks/workflow/tool_context.py @@ -35,7 +35,7 @@ async def build_workflow_tool_context( Prefer the caller-provided session/message. When absent, create a temporary parent session and synthetic user message so workflow-internal tools such as - ``task`` / ``delegate_task`` can resolve a valid parent session. + ``delegate_task`` can resolve a valid parent session. """ effective_session_id = str(session_id or "").strip() diff --git a/tests/config/test_config.py b/tests/config/test_config.py index 713448883..c957244e9 100644 --- a/tests/config/test_config.py +++ b/tests/config/test_config.py @@ -153,6 +153,17 @@ def test_legacy_todo_permission_names_migrate_to_todo(): assert dumped["bash"] == PermissionAction.ASK +def test_legacy_task_permission_name_migrates_to_delegate_task(): + permission = PermissionConfig.model_validate({ + "delegate_task": "allow", + "task": {"explore": "deny"}, + }) + + dumped = permission.model_dump(exclude_none=True) + assert dumped["delegate_task"] == PermissionAction.DENY + assert "task" not in dumped + + def test_legacy_todo_tool_flags_migrate_to_todo_permission(): config = ConfigInfo.model_validate({ "tools": { diff --git a/tests/session/test_context_usage.py b/tests/session/test_context_usage.py index 2bc9db71a..303dcf88a 100644 --- a/tests/session/test_context_usage.py +++ b/tests/session/test_context_usage.py @@ -275,7 +275,7 @@ async def test_context_usage_splits_skill_and_delegation_tools(context_usage_moc ), SimpleNamespace( type="tool", - tool="task", + tool="delegate_task", state=SimpleNamespace(input={}, output="t" * 80, time={"start": 3}), ), SimpleNamespace( @@ -288,20 +288,16 @@ async def test_context_usage_splits_skill_and_delegation_tools(context_usage_moc metadata={"tool": "skill_load"}, state=SimpleNamespace(input={}, output="m" * 40, time={"start": 5}), ), - SimpleNamespace( - type="subtask", - prompt="p" * 40, - description="q" * 40, - ), ] } snapshot = await context_usage.build_context_usage_snapshot("sess-1") assert [(segment.key, segment.tokens) for segment in snapshot.segments] == [ + ("conversation", 20), ("tools", 30), ("skillLoad", 30), - ("agentDelegation", 50), + ("agentDelegation", 30), ] tools_segment = next(segment for segment in snapshot.segments if segment.key == "tools") assert tools_segment.tokens == 30 diff --git a/tests/session/test_message_parts.py b/tests/session/test_message_parts.py index 89504ba62..c28a7d1c0 100644 --- a/tests/session/test_message_parts.py +++ b/tests/session/test_message_parts.py @@ -27,7 +27,6 @@ SnapshotPart, StepFinishPart, StepStartPart, - SubtaskPart, TextPart, TokenCache, TokenUsage, @@ -313,21 +312,9 @@ def test_creation(self): # --------------------------------------------------------------------------- -# SubtaskPart / AgentPart +# AgentPart # --------------------------------------------------------------------------- -class TestSubtaskPart: - def test_creation(self): - part = SubtaskPart( - sessionID=SID, - messageID=MID, - prompt="Summarize findings", - description="Summarize", - agent="rex", - ) - assert part.type == "subtask" - assert part.agent == "rex" - class TestAgentPart: def test_creation(self): @@ -416,6 +403,24 @@ def test_deserialize_reasoning_part(self): assert deserialized is not None assert deserialized.type == "reasoning" + def test_deserialize_legacy_subtask_as_ignored_text(self): + deserialized = Message.deserialize_part( + { + "id": "part_legacy_subtask", + "sessionID": SID, + "messageID": MID, + "type": "subtask", + "prompt": "old delegated command", + "description": "legacy", + "agent": "rex", + } + ) + + assert deserialized.type == "text" + assert deserialized.text == "" + assert deserialized.ignored is True + assert deserialized.metadata == {"legacyPartType": "subtask"} + def test_deserialize_unknown_type_falls_back_to_text(self): # Unknown type falls back to TextPart; missing required fields raise exception with pytest.raises(Exception): @@ -537,11 +542,11 @@ async def test_store_part_does_not_downgrade_terminal_tool_state(self, monkeypat sessionID=sid, messageID=msg.id, callID="call_terminal_guard", - tool="task", + tool="delegate_task", state=ToolStateCompleted( input={"prompt": "run"}, output="done", - title="task", + title="delegate_task", metadata={"sessionId": "ses_child_done"}, time={"start": 1000, "end": 2000}, ), @@ -551,10 +556,10 @@ async def test_store_part_does_not_downgrade_terminal_tool_state(self, monkeypat sessionID=sid, messageID=msg.id, callID="call_terminal_guard", - tool="task", + tool="delegate_task", state=ToolStateRunning( input={"prompt": "run"}, - title="task", + title="delegate_task", metadata={"sessionId": "ses_child_done", "status": "running"}, time={"start": 1000}, ), diff --git a/tests/session/test_session_abort_inject.py b/tests/session/test_session_abort_inject.py index 2f90b3da0..61334a51d 100644 --- a/tests/session/test_session_abort_inject.py +++ b/tests/session/test_session_abort_inject.py @@ -1049,59 +1049,6 @@ async def test_run_loop_does_not_return_previous_reply_when_current_step_fails( process_step.assert_awaited_once() -class TestExecuteSubtask: - @pytest.mark.asyncio - async def test_execute_subtask_passes_tool_context_first(self): - session_info = _make_session_info("subtask_exec_test") - ctx = LoopContext( - session=session_info, - provider_id="test-provider", - model_id="test-model", - agent_name="rex", - ) - last_user = SimpleNamespace( - id="msg_parent", - agent="rex", - model={"providerID": "test-provider", "modelID": "test-model"}, - provider="test-provider", - ) - task_part = SimpleNamespace( - agent="helper", - prompt="do the thing", - description="test task", - command=None, - model=None, - ) - - task_tool = MagicMock() - task_tool.execute = AsyncMock(return_value=SimpleNamespace( - output="done", - title="task complete", - metadata={"sessionId": "child-session"}, - )) - - assistant_msg = SimpleNamespace(id="msg_assistant") - synthetic_msg = SimpleNamespace(id="msg_synthetic") - - with patch("flocks.agent.registry.Agent.get", AsyncMock(return_value=SimpleNamespace(name="helper"))), \ - patch("flocks.tool.registry.ToolRegistry.get", return_value=task_tool), \ - patch("flocks.session.session_loop.Message.create", AsyncMock(side_effect=[assistant_msg, synthetic_msg])), \ - patch("flocks.session.session_loop.Message.add_part", AsyncMock()), \ - patch("flocks.session.session_loop.Message.update", AsyncMock()), \ - patch("flocks.session.session_loop.Message.update_part", AsyncMock()): - await SessionLoop._execute_subtask(ctx, last_user, task_part) - - task_tool.execute.assert_awaited_once() - tool_ctx = task_tool.execute.await_args.args[0] - assert tool_ctx.session_id == session_info.id - assert tool_ctx.message_id == assistant_msg.id - assert task_tool.execute.await_args.kwargs == { - "prompt": "do the thing", - "description": "test task", - "subagent_type": "helper", - "command": None, - } - # --------------------------------------------------------------------------- # LoopContext tests diff --git a/tests/tool/test_builtin_management_tools.py b/tests/tool/test_builtin_management_tools.py index 9c649404a..c604dc548 100644 --- a/tests/tool/test_builtin_management_tools.py +++ b/tests/tool/test_builtin_management_tools.py @@ -49,7 +49,16 @@ def test_lsp_remains_non_native_by_default() -> None: assert tool.info.native is False -def test_task_remains_non_native_when_declared() -> None: +def test_delegate_task_remains_native_when_declared() -> None: + ToolRegistry.init() + + tool = ToolRegistry.get("delegate_task") + + assert tool is not None + assert tool.info.native is True + + +def test_task_alias_remains_non_native_when_declared() -> None: ToolRegistry.init() tool = ToolRegistry.get("task") diff --git a/tests/tool/test_task_model_pinning.py b/tests/tool/test_task_model_pinning.py deleted file mode 100644 index 4f75e10ab..000000000 --- a/tests/tool/test_task_model_pinning.py +++ /dev/null @@ -1,61 +0,0 @@ -from unittest.mock import AsyncMock, patch - -import pytest - -from flocks.tool.agent.task import task_tool -from flocks.tool.registry import ToolContext, ToolRegistry, ToolResult - - -def _make_ctx() -> ToolContext: - return ToolContext(session_id="test-session", message_id="test-message", agent="rex") - - -class TestTaskCompatibilityAlias: - def test_task_schema_does_not_expose_background_execution(self): - schema = ToolRegistry.get_schema("task") - assert schema is not None - assert "run_in_background" not in schema.properties - # Legacy batch shape is gone. - assert "tasks" not in schema.properties - - @pytest.mark.asyncio - async def test_task_tool_rejects_background_execution_when_called_directly(self): - result = await task_tool( - _make_ctx(), - description="delegate explore", - prompt="Inspect the repository", - subagent_type="explore", - run_in_background=True, - ) - - assert result.success is False - assert "Background subagent execution is disabled" in (result.error or "") - - @pytest.mark.asyncio - async def test_task_tool_forwards_single_call_to_delegate_task(self): - delegate_result = ToolResult( - success=True, - output="ok", - metadata={"sessionId": "ses-child"}, - ) - - with patch( - "flocks.tool.agent.task.delegate_task_tool", - AsyncMock(return_value=delegate_result), - ) as delegate: - result = await task_tool( - _make_ctx(), - description="delegate explore", - prompt="Inspect the repository", - subagent_type="explore", - model="openai/gpt-5", - ) - - assert result is delegate_result - delegate.assert_awaited_once() - kwargs = delegate.await_args.kwargs - assert kwargs["description"] == "delegate explore" - assert kwargs["prompt"] == "Inspect the repository" - assert kwargs["subagent_type"] == "explore" - assert kwargs["run_in_background"] is False - assert kwargs["model"] == "openai/gpt-5" diff --git a/tests/tool/test_tool_catalog.py b/tests/tool/test_tool_catalog.py index f9e440ad7..7f1db8696 100644 --- a/tests/tool/test_tool_catalog.py +++ b/tests/tool/test_tool_catalog.py @@ -78,12 +78,8 @@ def test_catalog_uses_real_builtin_tool_names_for_metadata_keys() -> None: assert name in TOOL_TAGS -def test_task_tool_tags_reflect_agent_delegation() -> None: - metadata = get_tool_catalog_metadata("task") - - assert "agent" in metadata.tags - assert "delegation" in metadata.tags - assert "planning" not in metadata.tags +def test_task_compatibility_alias_has_no_catalog_tags() -> None: + assert "task" not in TOOL_TAGS def test_schedule_task_and_todo_use_distinct_management_tags() -> None: diff --git a/tests/tool/test_tools.py b/tests/tool/test_tools.py index ab202c376..33628c9f1 100644 --- a/tests/tool/test_tools.py +++ b/tests/tool/test_tools.py @@ -156,7 +156,7 @@ def test_expected_tools_registered(self): # P1 tools "webfetch", "todo", "question", # P2 tools - "task", "lsp", "skill_load", + "delegate_task", "lsp", "skill_load", # P3 tools (2) "websearch", "apply_patch", ] @@ -801,13 +801,13 @@ async def test_webfetch_schema(self): # P2 Tools Tests # ============================================================================= -class TestTaskTool: - """Test the task tool""" +class TestDelegateTaskTool: + """Test the delegate_task tool""" @pytest.mark.asyncio - async def test_task_exists(self): - """Test that task tool is registered""" - tool = ToolRegistry.get("task") + async def test_delegate_task_exists(self): + """Test that delegate_task tool is registered""" + tool = ToolRegistry.get("delegate_task") assert tool is not None diff --git a/tests/utils/test_id_compatibility.py b/tests/utils/test_id_compatibility.py index 22b1691a2..c26d9c509 100644 --- a/tests/utils/test_id_compatibility.py +++ b/tests/utils/test_id_compatibility.py @@ -26,7 +26,6 @@ def test_prefix_mappings(self): "call": "cal", "step": "stp", "agent": "agt", - "subtask": "stk", "event": "evt", "tqref": "tqr", "task": "tsk", diff --git a/tests/workflow/test_loop_host_forensics_fast_workflow.py b/tests/workflow/test_loop_host_forensics_fast_workflow.py index 6f5ec5c74..a60b1eb16 100644 --- a/tests/workflow/test_loop_host_forensics_fast_workflow.py +++ b/tests/workflow/test_loop_host_forensics_fast_workflow.py @@ -38,7 +38,7 @@ def run_safe(self, *args, **kwargs) -> dict: assert kwargs["host"] == "10.0.0.8" assert kwargs["username"] == "root" return {"success": True, "output": "FLOCKS_SSH_OK\n"} - assert args == ("task",) + assert args == ("delegate_task",) assert kwargs["subagent_type"] == "host-forensics-fast" assert "- host: 10.0.0.8" in kwargs["prompt"] assert "- username: root" in kwargs["prompt"] @@ -130,7 +130,7 @@ def __init__(self) -> None: def run_safe(self, *args, **kwargs) -> dict: if args == ("ssh_host_cmd",): return {"success": True, "output": "FLOCKS_SSH_OK\n"} - assert args == ("task",) + assert args == ("delegate_task",) self.task_calls += 1 if self.task_calls == 1: return { diff --git a/tui/flocks/agent/generate.txt b/tui/flocks/agent/generate.txt index 774277b0f..312a1d6d3 100644 --- a/tui/flocks/agent/generate.txt +++ b/tui/flocks/agent/generate.txt @@ -41,20 +41,20 @@ When a user describes what they want an agent to do, you will: assistant: "Here is the relevant function: " - Since the user is greeting, use the Task tool to launch the greeting-responder agent to respond with a friendly joke. + Since the user is greeting, use delegate_task to launch the greeting-responder agent to respond with a friendly joke. assistant: "Now let me use the code-reviewer agent to review the code" - Context: User is creating an agent to respond to the word "hello" with a friendly jok. user: "Hello" - assistant: "I'm going to use the Task tool to launch the greeting-responder agent to respond with a friendly joke" + assistant: "I'm going to use delegate_task to launch the greeting-responder agent to respond with a friendly joke" Since the user is greeting, use the greeting-responder agent to respond with a friendly joke. - If the user mentioned or implied that the agent should be used proactively, you should include examples of this. -- NOTE: Ensure that in the examples, you are making the assistant use the Agent tool and not simply respond directly to the task. +- NOTE: Ensure that in the examples, you are making the assistant use delegate_task and not simply respond directly to the task. Your output must be a valid JSON object with exactly these fields: { diff --git a/tui/flocks/cli/cmd/agent.ts b/tui/flocks/cli/cmd/agent.ts index 8a018cea9..c597fee93 100644 --- a/tui/flocks/cli/cmd/agent.ts +++ b/tui/flocks/cli/cmd/agent.ts @@ -22,7 +22,7 @@ const AVAILABLE_TOOLS = [ "glob", "grep", "webfetch", - "task", + "delegate_task", "todo", ] diff --git a/tui/flocks/cli/cmd/tui/routes/session/index.tsx b/tui/flocks/cli/cmd/tui/routes/session/index.tsx index bade13c9d..899a1e6b1 100644 --- a/tui/flocks/cli/cmd/tui/routes/session/index.tsx +++ b/tui/flocks/cli/cmd/tui/routes/session/index.tsx @@ -40,7 +40,7 @@ import type { GrepTool } from "@/tool/grep" import type { EditTool } from "@/tool/edit" import type { ApplyPatchTool } from "@/tool/apply_patch" import type { WebFetchTool } from "@/tool/webfetch" -import type { TaskTool } from "@/tool/task" +import type { DelegateTaskTool } from "@/tool/delegate-task" import type { QuestionTool } from "@/tool/question" import { useKeyboard, useRenderer, useTerminalDimensions, type JSX } from "@opentui/solid" import { useSDK } from "@tui/context/sdk" @@ -1861,7 +1861,7 @@ function SubagentActivity(props: { ) } -function Task(props: ToolProps) { +function Task(props: ToolProps) { const { theme } = useTheme() const keybind = useKeybind() const { navigate } = useRoute() @@ -1977,7 +1977,7 @@ function DelegateTask(props: ToolProps) { navigate({ type: "session", sessionID: sessionId()! }) : undefined} part={props.part} > @@ -1997,7 +1997,7 @@ function DelegateTask(props: ToolProps) { > - {delegateInput.description || "subtask"} + {delegateInput.description || "delegated task"} {isBackground() ? " (background)" : ""} {statusText()} @@ -2016,7 +2016,7 @@ function DelegateTask(props: ToolProps) { part={props.part} > {agentName()}{" "} - "{delegateInput.description || "subtask"}" + "{delegateInput.description || "delegated task"}" {isBackground() ? " (bg)" : ""} diff --git a/tui/flocks/cli/cmd/tui/routes/session/permission.tsx b/tui/flocks/cli/cmd/tui/routes/session/permission.tsx index 7d1b6c971..77e328560 100644 --- a/tui/flocks/cli/cmd/tui/routes/session/permission.tsx +++ b/tui/flocks/cli/cmd/tui/routes/session/permission.tsx @@ -217,7 +217,7 @@ export function PermissionPrompt(props: { request: PermissionRequest }) { description={("$ " + input().command) as string} /> - + { + test("maps task to delegate_task with deny precedence", () => { + expect( + Config.Permission.parse({ + delegate_task: "allow", + task: { explore: "deny" }, + }), + ).toEqual({ delegate_task: "deny" }) + }) +}) diff --git a/tui/flocks/config/config.ts b/tui/flocks/config/config.ts index 5085330ed..a43d70e11 100644 --- a/tui/flocks/config/config.ts +++ b/tui/flocks/config/config.ts @@ -501,12 +501,16 @@ export namespace Config { const canonicalPermissionToolName = (tool: string) => { if (tool === "todowrite" || tool === "todoread") return "todo" + if (tool === "task") return "delegate_task" return tool } const assignPermission = (target: Record, tool: string, action: PermissionRule) => { const canonical = canonicalPermissionToolName(tool) - if (target[canonical] === "deny" || action === "deny") { + const containsDeny = (value: PermissionRule): boolean => + value === "deny" || + (typeof value === "object" && value !== null && Object.values(value).some((item) => containsDeny(item))) + if ((canonical in target && containsDeny(target[canonical])) || containsDeny(action)) { target[canonical] = "deny" return } @@ -536,7 +540,7 @@ export namespace Config { glob: PermissionRule.optional(), grep: PermissionRule.optional(), bash: PermissionRule.optional(), - task: PermissionRule.optional(), + delegate_task: PermissionRule.optional(), external_directory: PermissionRule.optional(), todo: PermissionAction.optional(), question: PermissionAction.optional(), @@ -559,7 +563,6 @@ export namespace Config { description: z.string().optional(), agent: z.string().optional(), model: z.string().optional(), - subtask: z.boolean().optional(), }) export type Command = z.infer diff --git a/tui/flocks/session/message-v2.ts b/tui/flocks/session/message-v2.ts index f2f3331e5..ff596603b 100644 --- a/tui/flocks/session/message-v2.ts +++ b/tui/flocks/session/message-v2.ts @@ -163,21 +163,6 @@ export namespace MessageV2 { }) export type CompactionPart = z.infer - export const SubtaskPart = PartBase.extend({ - type: z.literal("subtask"), - prompt: z.string(), - description: z.string(), - agent: z.string(), - model: z - .object({ - providerID: z.string(), - modelID: z.string(), - }) - .optional(), - command: z.string().optional(), - }) - export type SubtaskPart = z.infer - export const RetryPart = PartBase.extend({ type: z.literal("retry"), attempt: z.number(), @@ -329,7 +314,6 @@ export namespace MessageV2 { export const Part = z .discriminatedUnion("type", [ TextPart, - SubtaskPart, ReasoningPart, FilePart, ToolPart, @@ -466,12 +450,6 @@ export namespace MessageV2 { text: "What did we do so far?", }) } - if (part.type === "subtask") { - userMessage.parts.push({ - type: "text", - text: "The following tool was executed by the user", - }) - } } } diff --git a/tui/flocks/session/prompt.ts b/tui/flocks/session/prompt.ts index befddd369..22313c167 100644 --- a/tui/flocks/session/prompt.ts +++ b/tui/flocks/session/prompt.ts @@ -34,7 +34,6 @@ import { SessionSummary } from "./summary" import { NamedError } from "@flocks-ai/util/error" import { fn } from "@/util/fn" import { SessionProcessor } from "./processor" -import { TaskTool } from "@/tool/task" import { Tool } from "@/tool/tool" import { PermissionNext } from "@/permission/next" import { SessionStatus } from "./status" @@ -200,16 +199,6 @@ export namespace SessionPrompt { .meta({ ref: "AgentPartInput", }), - MessageV2.SubtaskPart.omit({ - messageID: true, - sessionID: true, - }) - .partial({ - id: true, - }) - .meta({ - ref: "SubtaskPartInput", - }), ]), ), }) @@ -344,7 +333,7 @@ export namespace SessionPrompt { let lastUser: MessageV2.User | undefined let lastAssistant: MessageV2.Assistant | undefined let lastFinished: MessageV2.Assistant | undefined - let tasks: (MessageV2.CompactionPart | MessageV2.SubtaskPart)[] = [] + const pendingCompactions: MessageV2.CompactionPart[] = [] for (let i = msgs.length - 1; i >= 0; i--) { const msg = msgs[i] if (!lastUser && msg.info.role === "user") lastUser = msg.info as MessageV2.User @@ -352,9 +341,9 @@ export namespace SessionPrompt { if (!lastFinished && msg.info.role === "assistant" && msg.info.finish) lastFinished = msg.info as MessageV2.Assistant if (lastUser && lastFinished) break - const task = msg.parts.filter((part) => part.type === "compaction" || part.type === "subtask") - if (task && !lastFinished) { - tasks.push(...task) + const compactions = msg.parts.filter((part) => part.type === "compaction") + if (compactions.length > 0 && !lastFinished) { + pendingCompactions.push(...compactions) } } @@ -378,183 +367,16 @@ export namespace SessionPrompt { }) const model = await Provider.getModel(lastUser.model.providerID, lastUser.model.modelID) - const task = tasks.pop() - - // pending subtask - // TODO: centralize "invoke tool" logic - if (task?.type === "subtask") { - const taskTool = await TaskTool.init() - const taskModel = task.model ? await Provider.getModel(task.model.providerID, task.model.modelID) : model - const assistantMessage = (await Session.updateMessage({ - id: Identifier.ascending("message"), - role: "assistant", - parentID: lastUser.id, - sessionID, - mode: task.agent, - agent: task.agent, - path: { - cwd: Instance.directory, - root: Instance.worktree, - }, - cost: 0, - tokens: { - input: 0, - output: 0, - reasoning: 0, - cache: { read: 0, write: 0 }, - }, - modelID: taskModel.id, - providerID: taskModel.providerID, - time: { - created: Date.now(), - }, - })) as MessageV2.Assistant - let part = (await Session.updatePart({ - id: Identifier.ascending("part"), - messageID: assistantMessage.id, - sessionID: assistantMessage.sessionID, - type: "tool", - callID: ulid(), - tool: TaskTool.id, - state: { - status: "running", - input: { - prompt: task.prompt, - description: task.description, - subagent_type: task.agent, - command: task.command, - }, - time: { - start: Date.now(), - }, - }, - })) as MessageV2.ToolPart - const taskArgs = { - prompt: task.prompt, - description: task.description, - subagent_type: task.agent, - command: task.command, - } - await Plugin.trigger( - "tool.execute.before", - { - tool: "task", - sessionID, - callID: part.id, - }, - { args: taskArgs }, - ) - let executionError: Error | undefined - const taskAgent = await Agent.get(task.agent) - const taskCtx: Tool.Context = { - agent: task.agent, - messageID: assistantMessage.id, - sessionID: sessionID, - abort, - callID: part.callID, - extra: { bypassAgentCheck: true }, - async metadata(input) { - await Session.updatePart({ - ...part, - type: "tool", - state: { - ...part.state, - ...input, - }, - } satisfies MessageV2.ToolPart) - }, - async ask(req) { - await PermissionNext.ask({ - ...req, - sessionID: sessionID, - ruleset: PermissionNext.merge(taskAgent.permission, session.permission ?? []), - }) - }, - } - const result = await taskTool.execute(taskArgs, taskCtx).catch((error) => { - executionError = error - log.error("subtask execution failed", { error, agent: task.agent, description: task.description }) - return undefined - }) - await Plugin.trigger( - "tool.execute.after", - { - tool: "task", - sessionID, - callID: part.id, - }, - result, - ) - assistantMessage.finish = "tool-calls" - assistantMessage.time.completed = Date.now() - await Session.updateMessage(assistantMessage) - if (result && part.state.status === "running") { - await Session.updatePart({ - ...part, - state: { - status: "completed", - input: part.state.input, - title: result.title, - metadata: result.metadata, - output: result.output, - attachments: result.attachments, - time: { - ...part.state.time, - end: Date.now(), - }, - }, - } satisfies MessageV2.ToolPart) - } - if (!result) { - await Session.updatePart({ - ...part, - state: { - status: "error", - error: executionError ? `Tool execution failed: ${executionError.message}` : "Tool execution failed", - time: { - start: part.state.status === "running" ? part.state.time.start : Date.now(), - end: Date.now(), - }, - metadata: part.metadata, - input: part.state.input, - }, - } satisfies MessageV2.ToolPart) - } - - // Add synthetic user message to prevent certain reasoning models from erroring - // If we create assistant messages w/ out user ones following mid loop thinking signatures - // will be missing and it can cause errors for models like gemini for example - const summaryUserMsg: MessageV2.User = { - id: Identifier.ascending("message"), - sessionID, - role: "user", - time: { - created: Date.now(), - }, - agent: lastUser.agent, - model: lastUser.model, - } - await Session.updateMessage(summaryUserMsg) - await Session.updatePart({ - id: Identifier.ascending("part"), - messageID: summaryUserMsg.id, - sessionID, - type: "text", - text: "Summarize the task tool output above and continue with your task.", - synthetic: true, - } satisfies MessageV2.TextPart) - - continue - } + const pendingCompaction = pendingCompactions.pop() // pending compaction - if (task?.type === "compaction") { + if (pendingCompaction) { const result = await SessionCompaction.process({ messages: msgs, parentID: lastUser.id, abort, sessionID, - auto: task.auto, + auto: pendingCompaction.auto, }) if (result === "stop") break continue @@ -1193,8 +1015,8 @@ export namespace SessionPrompt { } if (part.type === "agent") { - // Check if this agent would be denied by task permission - const perm = PermissionNext.evaluate("task", part.name, agent.permission) + // Check whether this agent may be delegated to. + const perm = PermissionNext.evaluate("delegate_task", part.name, agent.permission) const hint = perm.action === "deny" ? " . Invoked by user; guaranteed to exist." : "" return [ { @@ -1212,7 +1034,7 @@ export namespace SessionPrompt { // An extra space is added here. Otherwise the 'Use' gets appended // to user's last word; making a combined word text: - " Use the above message and context to generate a prompt and call the task tool with subagent: " + + " Use the above message and context to generate a prompt and call the delegate_task tool with subagent: " + part.name + hint, }, @@ -1742,30 +1564,7 @@ NOTE: At any point in time through this workflow you should feel free to ask the } const templateParts = await resolvePromptParts(template) - const isSubtask = (agent.mode === "subagent" && command.subtask !== false) || command.subtask === true - const parts = isSubtask - ? [ - { - type: "subtask" as const, - agent: agent.name, - description: command.description ?? "", - command: input.command, - model: { - providerID: taskModel.providerID, - modelID: taskModel.modelID, - }, - // TODO: how can we make task tool accept a more complex input? - prompt: templateParts.find((y) => y.type === "text")?.text ?? "", - }, - ] - : [...templateParts, ...(input.parts ?? [])] - - const userAgent = isSubtask ? (input.agent ?? (await Agent.defaultAgent())) : agentName - const userModel = isSubtask - ? input.model - ? Provider.parseModel(input.model) - : await lastModel(input.sessionID) - : taskModel + const parts = [...templateParts, ...(input.parts ?? [])] await Plugin.trigger( "command.execute.before", @@ -1780,8 +1579,8 @@ NOTE: At any point in time through this workflow you should feel free to ask the const result = (await prompt({ sessionID: input.sessionID, messageID: input.messageID, - model: userModel, - agent: userAgent, + model: taskModel, + agent: agentName, parts, variant: input.variant, })) as MessageV2.WithParts @@ -1817,15 +1616,9 @@ NOTE: At any point in time through this workflow you should feel free to ask the if (!isFirst) return // Gather all messages up to and including the first real user message for context - // This includes any shell/subtask executions that preceded the user's first prompt const contextMessages = input.history.slice(0, firstRealUserIdx + 1) const firstRealUser = contextMessages[firstRealUserIdx] - // For subtask-only messages (from command invocations), extract the prompt directly - // since toModelMessage converts subtask parts to generic "The following tool was executed by the user" - const subtaskParts = firstRealUser.parts.filter((p) => p.type === "subtask") as MessageV2.SubtaskPart[] - const hasOnlySubtaskParts = subtaskParts.length > 0 && firstRealUser.parts.every((p) => p.type === "subtask") - const agent = await Agent.get("title") if (!agent) return const result = await LLM.stream({ @@ -1848,9 +1641,7 @@ NOTE: At any point in time through this workflow you should feel free to ask the role: "user", content: "Generate a title for this conversation:\n", }, - ...(hasOnlySubtaskParts - ? [{ role: "user" as const, content: subtaskParts.map((p) => p.prompt).join("\n") }] - : MessageV2.toModelMessage(contextMessages)), + ...MessageV2.toModelMessage(contextMessages), ], }) const text = await result.text.catch((err) => log.error("failed to generate title", { error: err })) diff --git a/tui/flocks/session/prompt/anthropic-20250930.txt b/tui/flocks/session/prompt/anthropic-20250930.txt index 676c4d8dc..01ac8e5b5 100644 --- a/tui/flocks/session/prompt/anthropic-20250930.txt +++ b/tui/flocks/session/prompt/anthropic-20250930.txt @@ -129,10 +129,10 @@ The user will primarily request you perform software engineering tasks. This inc # Tool usage policy -- You should proactively use the Task tool with specialized agents when the task at hand matches the agent's description. +- You should proactively use `delegate_task` with specialized agents when the task at hand matches the agent's description. - When WebFetch returns a message about a redirect to a different host, you should immediately make a new WebFetch request with the redirect URL provided in the response. -- If the user specifies that they want you to run tools "in parallel", you MUST send a single message with multiple tool use content blocks. For example, if you need to launch multiple agents in parallel, send a single message with multiple Task tool calls. +- If the user specifies that they want you to run tools "in parallel", you MUST send a single message with multiple tool use content blocks. For example, if you need to launch multiple agents in parallel, send a single message with multiple `delegate_task` calls. - Use specialized tools instead of bash commands when possible, as this provides a better user experience. For file operations, use dedicated tools: Read for reading files instead of cat/head/tail, Edit for editing instead of sed/awk, and Write for creating files instead of cat with heredoc or echo redirection. Reserve bash tools exclusively for actual system commands and terminal operations that require shell execution. NEVER use bash echo or other command-line tools to communicate thoughts, explanations, or instructions to the user. Output all communication directly in your response text instead. diff --git a/tui/flocks/session/prompt/anthropic.txt b/tui/flocks/session/prompt/anthropic.txt index 7a0e5fd5c..0709a5c30 100644 --- a/tui/flocks/session/prompt/anthropic.txt +++ b/tui/flocks/session/prompt/anthropic.txt @@ -73,20 +73,20 @@ The user will primarily request you perform SecOps tasks. This includes security # Tool usage policy -- You should proactively use the Task tool with specialized agents when the task at hand matches the agent's description. +- You should proactively use `delegate_task` with specialized agents when the task at hand matches the agent's description. - When WebFetch returns a message about a redirect to a different host, you should immediately make a new WebFetch request with the redirect URL provided in the response. - You can call multiple tools in a single response. If you intend to call multiple tools and there are no dependencies between them, make all independent tool calls in parallel. Maximize use of parallel tool calls where possible to increase efficiency. However, if some tool calls depend on previous calls to inform dependent values, do NOT call these tools in parallel and instead call them sequentially. For instance, if one operation must complete before another starts, run these operations sequentially instead. Never use placeholders or guess missing parameters in tool calls. -- If the user specifies that they want you to run tools "in parallel", you MUST send a single message with multiple tool use content blocks. For example, if you need to launch multiple agents in parallel, send a single message with multiple Task tool calls. +- If the user specifies that they want you to run tools "in parallel", you MUST send a single message with multiple tool use content blocks. For example, if you need to launch multiple agents in parallel, send a single message with multiple `delegate_task` calls. - Use specialized tools instead of bash commands when possible, as this provides a better user experience. For file operations, use dedicated tools: Read for reading files instead of cat/head/tail, Edit for editing instead of sed/awk, and Write for creating files instead of cat with heredoc or echo redirection. Reserve bash tools exclusively for actual system commands and terminal operations that require shell execution. NEVER use bash echo or other command-line tools to communicate thoughts, explanations, or instructions to the user. Output all communication directly in your response text instead. -- VERY IMPORTANT: When exploring the codebase to gather context or to answer a question that is not a needle query for a specific file/class/function, it is CRITICAL that you use the Task tool instead of running search commands directly. +- VERY IMPORTANT: When exploring the codebase to gather context or to answer a question that is not a needle query for a specific file/class/function, it is CRITICAL that you use `delegate_task` instead of running search commands directly. user: Where are errors from the client handled? -assistant: [Uses the Task tool to find the files that handle client errors instead of using Glob or Grep directly] +assistant: [Uses `delegate_task` to find the files that handle client errors instead of using Glob or Grep directly] user: What is the codebase structure? -assistant: [Uses the Task tool] +assistant: [Uses `delegate_task`] IMPORTANT: Always use `todo(action="write")` to plan and track tasks throughout the conversation. diff --git a/tui/flocks/tool/bash.txt b/tui/flocks/tool/bash.txt index f42bd270a..a9514bd9f 100644 --- a/tui/flocks/tool/bash.txt +++ b/tui/flocks/tool/bash.txt @@ -81,7 +81,7 @@ Git Safety Protocol: Important notes: - NEVER run additional commands to read or explore code, besides git bash commands -- NEVER use the todo or Task tools +- NEVER use the todo or delegate_task tools - DO NOT push to the remote repository unless the user explicitly asks you to do so - IMPORTANT: Never use git commands with the -i flag (like git rebase -i or git add -i) since they require interactive input which is not supported. - If there are no changes to commit (i.e., no untracked files and no modifications), do not create an empty commit @@ -108,7 +108,7 @@ gh pr create --title "the pr title" --body "$(cat <<'EOF' Important: -- DO NOT use the todo or Task tools +- DO NOT use the todo or delegate_task tools - Return the PR URL when you're done, so the user can see it # Other common operations diff --git a/tui/flocks/tool/task.ts b/tui/flocks/tool/delegate-task.ts similarity index 91% rename from tui/flocks/tool/task.ts rename to tui/flocks/tool/delegate-task.ts index f98316b39..dca20aae3 100644 --- a/tui/flocks/tool/task.ts +++ b/tui/flocks/tool/delegate-task.ts @@ -1,5 +1,5 @@ import { Tool } from "./tool" -import DESCRIPTION from "./task.txt" +import DESCRIPTION from "./delegate-task.txt" import z from "zod" import { Session } from "../session" import { Bus } from "../bus" @@ -20,13 +20,13 @@ const parameters = z.object({ command: z.string().describe("The command that triggered this task").optional(), }) -export const TaskTool = Tool.define("task", async (ctx) => { +export const DelegateTaskTool = Tool.define("delegate_task", async (ctx) => { const agents = await Agent.list().then((x) => x.filter((a) => a.mode !== "primary")) // Filter agents by permissions if agent provided const caller = ctx?.agent const accessibleAgents = caller - ? agents.filter((a) => PermissionNext.evaluate("task", a.name, caller.permission).action !== "deny") + ? agents.filter((a) => PermissionNext.evaluate("delegate_task", a.name, caller.permission).action !== "deny") : agents const description = DESCRIPTION.replace( @@ -41,10 +41,9 @@ export const TaskTool = Tool.define("task", async (ctx) => { async execute(params: z.infer, ctx) { const config = await Config.get() - // Skip permission check when user explicitly invoked via @ or command subtask if (!ctx.extra?.bypassAgentCheck) { await ctx.ask({ - permission: "task", + permission: "delegate_task", patterns: [params.subagent_type], always: ["*"], metadata: { @@ -57,7 +56,7 @@ export const TaskTool = Tool.define("task", async (ctx) => { const agent = await Agent.get(params.subagent_type) if (!agent) throw new Error(`Unknown agent type: ${params.subagent_type} is not a valid agent type`) - const hasTaskPermission = agent.permission.some((rule) => rule.permission === "task") + const hasDelegatePermission = agent.permission.some((rule) => rule.permission === "delegate_task") const session = await iife(async () => { if (params.session_id) { @@ -74,11 +73,11 @@ export const TaskTool = Tool.define("task", async (ctx) => { pattern: "*", action: "deny", }, - ...(hasTaskPermission + ...(hasDelegatePermission ? [] : [ { - permission: "task" as const, + permission: "delegate_task" as const, pattern: "*" as const, action: "deny" as const, }, @@ -147,7 +146,7 @@ export const TaskTool = Tool.define("task", async (ctx) => { agent: agent.name, tools: { todo: false, - ...(hasTaskPermission ? {} : { task: false }), + ...(hasDelegatePermission ? {} : { delegate_task: false }), ...Object.fromEntries((config.experimental?.primary_tools ?? []).map((t) => [t, false])), }, parts: promptParts, diff --git a/tui/flocks/tool/task.txt b/tui/flocks/tool/delegate-task.txt similarity index 79% rename from tui/flocks/tool/task.txt rename to tui/flocks/tool/delegate-task.txt index 7af2a6f60..21258b793 100644 --- a/tui/flocks/tool/task.txt +++ b/tui/flocks/tool/delegate-task.txt @@ -1,17 +1,17 @@ -Launch a new agent to handle complex, multistep tasks autonomously. +Delegate a complex, multistep task to another agent. Available agent types and the tools they have access to: {agents} -When using the Task tool, you must specify a subagent_type parameter to select which agent type to use. +When using the delegate_task tool, you must specify a subagent_type parameter to select which agent type to use. -When to use the Task tool: -- When you are instructed to execute custom slash commands. Use the Task tool with the slash command invocation as the entire prompt. The slash command can take arguments. For example: Task(description="Check the file", prompt="/check-file path/to/file.py") +When to use the delegate_task tool: +- When you are instructed to execute custom slash commands. Use delegate_task with the slash command invocation as the entire prompt. The slash command can take arguments. For example: delegate_task(description="Check the file", prompt="/check-file path/to/file.py") -When NOT to use the Task tool: -- If you want to read a specific file path, use the Read or Glob tool instead of the Task tool, to find the match more quickly +When NOT to use the delegate_task tool: +- If you want to read a specific file path, use the Read or Glob tool instead of delegate_task, to find the match more quickly - If you are searching for a specific class definition like "class Foo", use the Glob tool instead, to find the match more quickly -- If you are searching for code within a specific file or set of 2-3 files, use the Read tool instead of the Task tool, to find the match more quickly +- If you are searching for code within a specific file or set of 2-3 files, use the Read tool instead of delegate_task, to find the match more quickly - Other tasks that are not related to the agent descriptions above @@ -48,7 +48,7 @@ function isPrime(n) { Since a significant piece of code was written and the task was completed, now use the code-reviewer agent to review the code assistant: Now let me use the code-reviewer agent to review the code -assistant: Uses the Task tool to launch the code-reviewer agent +assistant: Uses delegate_task to launch the code-reviewer agent @@ -56,5 +56,5 @@ user: "Hello" Since the user is greeting, use the greeting-responder agent to respond with a friendly joke -assistant: "I'm going to use the Task tool to launch the with the greeting-responder agent" +assistant: "I'm going to use delegate_task to launch the greeting-responder agent" diff --git a/tui/flocks/tool/glob.txt b/tui/flocks/tool/glob.txt index add6b6ee1..63b3b04c9 100644 --- a/tui/flocks/tool/glob.txt +++ b/tui/flocks/tool/glob.txt @@ -2,5 +2,5 @@ - Supports glob patterns like "**/*.js" or "src/**/*.ts" - Returns matching file paths sorted by modification time - Use this tool when you need to find files by name patterns -- When you are doing an open-ended search that may require multiple rounds of globbing and grepping, use the Task tool instead +- When you are doing an open-ended search that may require multiple rounds of globbing and grepping, use delegate_task instead - You may call multiple independent tools in the same response. Prefer separate parallel Glob calls when multiple searches are likely to be useful. diff --git a/tui/flocks/tool/grep.txt b/tui/flocks/tool/grep.txt index adf583695..5e3fd1ee4 100644 --- a/tui/flocks/tool/grep.txt +++ b/tui/flocks/tool/grep.txt @@ -5,4 +5,4 @@ - Returns file paths and line numbers with at least one match sorted by modification time - Use this tool when you need to find files containing specific patterns - If you need to identify/count the number of matches within files, use the Bash tool with `rg` (ripgrep) directly. Do NOT use `grep`. -- When you are doing an open-ended search that may require multiple rounds of globbing and grepping, use the Task tool instead +- When you are doing an open-ended search that may require multiple rounds of globbing and grepping, use delegate_task instead diff --git a/tui/flocks/tool/registry.ts b/tui/flocks/tool/registry.ts index 4f4eed7a0..0714a55c7 100644 --- a/tui/flocks/tool/registry.ts +++ b/tui/flocks/tool/registry.ts @@ -4,7 +4,7 @@ import { EditTool } from "./edit" import { GlobTool } from "./glob" import { GrepTool } from "./grep" import { ReadTool } from "./read" -import { TaskTool } from "./task" +import { DelegateTaskTool } from "./delegate-task" import { TodoTool } from "./todo" import { WebFetchTool } from "./webfetch" import { WriteTool } from "./write" @@ -99,7 +99,7 @@ export namespace ToolRegistry { GrepTool, EditTool, WriteTool, - TaskTool, + DelegateTaskTool, WebFetchTool, TodoTool, WebSearchTool, diff --git a/tui/flocks/tool/truncation.ts b/tui/flocks/tool/truncation.ts index 074a0bcf0..248da7d86 100644 --- a/tui/flocks/tool/truncation.ts +++ b/tui/flocks/tool/truncation.ts @@ -41,9 +41,9 @@ export namespace Truncate { } } - function hasTaskTool(agent?: Agent.Info): boolean { + function hasDelegateTaskTool(agent?: Agent.Info): boolean { if (!agent?.permission) return false - const rule = PermissionNext.evaluate("task", "*", agent.permission) + const rule = PermissionNext.evaluate("delegate_task", "*", agent.permission) return rule.action !== "deny" } @@ -93,8 +93,8 @@ export namespace Truncate { const filepath = path.join(DIR, id) await Bun.write(Bun.file(filepath), text) - const hint = hasTaskTool(agent) - ? `The tool call succeeded but the output was truncated. Full output saved to: ${filepath}\nUse the Task tool to have explore agent process this file with Grep and Read (with offset/limit). Do NOT read the full file yourself - delegate to save context.` + const hint = hasDelegateTaskTool(agent) + ? `The tool call succeeded but the output was truncated. Full output saved to: ${filepath}\nUse delegate_task to have an explore agent process this file with Grep and Read (with offset/limit). Do NOT read the full file yourself - delegate to save context.` : `The tool call succeeded but the output was truncated. Full output saved to: ${filepath}\nUse Grep to search the full content or Read with offset/limit to view specific sections.` const message = direction === "head" diff --git a/tui/sdk/gen/types.gen.ts b/tui/sdk/gen/types.gen.ts index 8ac5c7342..ed4ba1843 100644 --- a/tui/sdk/gen/types.gen.ts +++ b/tui/sdk/gen/types.gen.ts @@ -383,15 +383,6 @@ export type CompactionPart = { export type Part = | TextPart - | { - id: string - sessionID: string - messageID: string - type: "subtask" - prompt: string - description: string - agent: string - } | ReasoningPart | FilePart | ToolPart @@ -1217,7 +1208,6 @@ export type Config = { description?: string agent?: string model?: string - subtask?: boolean } } watcher?: { @@ -1432,21 +1422,12 @@ export type AgentPartInput = { } } -export type SubtaskPartInput = { - id?: string - type: "subtask" - prompt: string - description: string - agent: string -} - export type Command = { name: string description?: string agent?: string model?: string template: string - subtask?: boolean } export type Model = { @@ -2591,7 +2572,7 @@ export type SessionPromptData = { tools?: { [key: string]: boolean } - parts: Array + parts: Array } path: { /** @@ -2686,7 +2667,7 @@ export type SessionPromptAsyncData = { tools?: { [key: string]: boolean } - parts: Array + parts: Array } path: { /** diff --git a/tui/sdk/v2/gen/sdk.gen.ts b/tui/sdk/v2/gen/sdk.gen.ts index a84c70938..033e235b6 100644 --- a/tui/sdk/v2/gen/sdk.gen.ts +++ b/tui/sdk/v2/gen/sdk.gen.ts @@ -134,7 +134,6 @@ import type { SessionUnshareResponses, SessionUpdateErrors, SessionUpdateResponses, - SubtaskPartInput, TextPartInput, ToolIdsErrors, ToolIdsResponses, @@ -1364,7 +1363,7 @@ export class Session extends HeyApiClient { } system?: string variant?: string - parts?: Array + parts?: Array }, options?: Options, ) { @@ -1452,7 +1451,7 @@ export class Session extends HeyApiClient { } system?: string variant?: string - parts?: Array + parts?: Array }, options?: Options, ) { diff --git a/tui/sdk/v2/gen/types.gen.ts b/tui/sdk/v2/gen/types.gen.ts index 77f869dc6..0109b8e50 100644 --- a/tui/sdk/v2/gen/types.gen.ts +++ b/tui/sdk/v2/gen/types.gen.ts @@ -429,20 +429,6 @@ export type CompactionPart = { export type Part = | TextPart - | { - id: string - sessionID: string - messageID: string - type: "subtask" - prompt: string - description: string - agent: string - model?: { - providerID: string - modelID: string - } - command?: string - } | ReasoningPart | FilePart | ToolPart @@ -1617,7 +1603,6 @@ export type Config = { description?: string agent?: string model?: string - subtask?: boolean } } watcher?: { @@ -1953,19 +1938,6 @@ export type AgentPartInput = { } } -export type SubtaskPartInput = { - id?: string - type: "subtask" - prompt: string - description: string - agent: string - model?: { - providerID: string - modelID: string - } - command?: string -} - export type ProviderAuthMethod = { type: "oauth" | "api" label: string @@ -2071,7 +2043,6 @@ export type Command = { model?: string mcp?: boolean template: string - subtask?: boolean hints: Array } @@ -3226,7 +3197,7 @@ export type SessionPromptData = { } system?: string variant?: string - parts: Array + parts: Array } path: { /** @@ -3413,7 +3384,7 @@ export type SessionPromptAsyncData = { } system?: string variant?: string - parts: Array + parts: Array } path: { /** diff --git a/webui/src/api/skill.ts b/webui/src/api/skill.ts index 254520cec..00c2f9dad 100644 --- a/webui/src/api/skill.ts +++ b/webui/src/api/skill.ts @@ -39,7 +39,6 @@ export interface Command { template: string; agent?: string; model?: string; - subtask?: boolean; hidden: boolean; aliases: string[]; visible_surfaces: string[]; From 32c36616132dca523dc1c47e889ca69fcd11be9f Mon Sep 17 00:00:00 2001 From: zhougongyan Date: Mon, 17 Aug 2026 16:13:47 +0800 Subject: [PATCH 06/63] refactor(session): centralize prompt context assembly Collect turn-scoped runtime inputs once and assemble deterministic prompt blocks so provider cache boundaries and context estimation share the same prompt contract. Co-Authored-By: Claude Opus 4.6 --- flocks/session/context_usage.py | 24 +- flocks/session/prompt.py | 391 +++++++++++------- flocks/session/runner.py | 251 +++++++---- tests/session/test_context_usage.py | 57 +++ tests/session/test_prompt_tokens.py | 37 +- tests/session/test_runner_step.py | 308 +++++++------- .../test_session_runner_tool_only_message.py | 8 +- 7 files changed, 703 insertions(+), 373 deletions(-) diff --git a/flocks/session/context_usage.py b/flocks/session/context_usage.py index ca35e09d1..9d35ce9f6 100644 --- a/flocks/session/context_usage.py +++ b/flocks/session/context_usage.py @@ -15,7 +15,7 @@ from flocks.provider.provider import Provider from flocks.session.message import Message -from flocks.session.prompt import SessionPrompt +from flocks.session.prompt import SessionPrompt, TurnPromptContext from flocks.session.session import SessionInfo from flocks.utils.log import Log @@ -309,7 +309,16 @@ async def _estimate_system_prompt_tokens( if agent is None: agent = await Agent.get("rex") - prompts = await SessionPrompt.build_system_prompts( + from flocks.config import Config + from flocks.project.instance import Instance + + try: + config = await Config.get() + config_instructions = tuple(config.instructions or ()) + except Exception: + config_instructions = () + + prompt_blocks = await SessionPrompt.build_system_prompt_blocks( session_id=session_id, session_directory=getattr(session, "directory", None) if session is not None else None, agent_name=getattr(agent, "name", agent_name) if agent is not None else agent_name, @@ -317,9 +326,16 @@ async def _estimate_system_prompt_tokens( provider_id=provider_id, model_id=model_id, prompt_tool_names=prompt_tool_names, - tool_revision=ToolRegistry.revision(), + turn_context=TurnPromptContext( + worktree=Instance.get_worktree(), + config_instructions=config_instructions, + tool_revision=ToolRegistry.revision(), + ), + ) + return sum( + SessionPrompt.count_tokens(block.content) + for block in prompt_blocks ) - return sum(SessionPrompt.count_tokens(prompt) for prompt in prompts) except Exception as exc: log.debug("context_usage.system_prompt_estimate_failed", { "session_id": session_id, diff --git a/flocks/session/prompt.py b/flocks/session/prompt.py index 75efef655..5c695f196 100644 --- a/flocks/session/prompt.py +++ b/flocks/session/prompt.py @@ -35,8 +35,7 @@ # Output token maximum OUTPUT_TOKEN_MAX = int(os.getenv("FLOCKS_OUTPUT_TOKEN_MAX", "32000")) SystemPromptCache = Dict[str, Any] -AsyncPromptFactory = Callable[[], Awaitable[Optional[str]]] -StringPromptFactory = Callable[[], Optional[str]] +AsyncPromptLoader = Callable[[], Awaitable[Optional[str]]] # Prompt template directory (same structure as Flocks) @@ -139,13 +138,30 @@ class ContextInfo(BaseModel): @dataclass(frozen=True) class SystemPromptBlock: - """Internal system prompt layer with cache metadata.""" + """Assembled system prompt layer.""" name: str content: str cache_scope: str - digest_inputs: Dict[str, Any] - cache_key: str + + +@dataclass(frozen=True) +class TurnPromptContext: + """Runtime prompt values collected once before deterministic assembly.""" + + tool_catalog: Optional[str] = None + device_asset_hint: Optional[str] = None + sandbox_context: Optional[str] = None + channel_context: Optional[str] = None + additional_context: Optional[str] = None + text_tool_catalog: Optional[str] = None + tool_results_reminder: Optional[str] = None + repeated_tool_calls_reminder: Optional[str] = None + worktree: Optional[str] = None + config_instructions: tuple[str, ...] = () + tool_revision: Optional[int] = None + device_revision: Optional[int] = None + minimal_prompt: Optional[bool] = None class SystemPrompt: @@ -804,49 +820,6 @@ def _layer_cache_key( """Build a layer cache key for one prompt block.""" return f"system_prompt_block:{name}:{cls._system_prompt_cache_digest(digest_inputs)}" - @classmethod - def _system_prompt_cache_key( - cls, - *, - session_id: str, - agent_name: str, - provider_id: str, - model_id: str, - block_keys: Iterable[str], - ) -> str: - """Build the cache key for the composed system prompt snapshot.""" - cache_digest = cls._system_prompt_cache_digest({ - "block_keys": tuple(block_keys), - }) - return f"system_prompts:{session_id}:{agent_name}:{provider_id}:{model_id}:{cache_digest}" - - @classmethod - def _read_system_prompt_cache( - cls, - static_cache: Optional[SystemPromptCache], - cache_key: Optional[str], - ) -> Optional[List[str]]: - """Return a defensive copy of cached prompt blocks when available.""" - if static_cache is None or cache_key is None: - return None - - cached = static_cache.get(cache_key) - if cached is None: - return None - return list(cached) - - @classmethod - def _write_system_prompt_cache( - cls, - static_cache: Optional[SystemPromptCache], - cache_key: Optional[str], - prompts: List[str], - ) -> None: - """Store a defensive copy of prompt blocks in the session cache.""" - if static_cache is None or cache_key is None: - return - static_cache[cache_key] = list(prompts) - @classmethod def _read_cached_prompt_block( cls, @@ -909,8 +882,6 @@ def _build_cached_prompt_block( name=name, content=content, cache_scope=cache_scope, - digest_inputs=digest_inputs, - cache_key=cache_key, ) @classmethod @@ -921,22 +892,21 @@ async def _build_cached_async_prompt_block( name: str, cache_scope: str, digest_inputs: Dict[str, Any], - builder: AsyncPromptFactory, + loader: AsyncPromptLoader, ) -> Optional[SystemPromptBlock]: """Build or reuse a cached async prompt block.""" cache_key = cls._layer_cache_key(name=name, digest_inputs=digest_inputs) content = cls._read_cached_prompt_block(static_cache, cache_key) if content is None: - content = cls._normalize_prompt_text(await builder()) - cls._write_cached_prompt_block(static_cache, cache_key, content) + content = cls._normalize_prompt_text(await loader()) + if content: + cls._write_cached_prompt_block(static_cache, cache_key, content) if not content: return None return SystemPromptBlock( name=name, content=content, cache_scope=cache_scope, - digest_inputs=digest_inputs, - cache_key=cache_key, ) @classmethod @@ -1044,26 +1014,6 @@ def _prompt_blocks_to_list( if block is not None and block.content.strip() ] - @classmethod - async def _build_optional_async_prompt( - cls, - prompt_factory: Optional[AsyncPromptFactory], - ) -> Optional[str]: - """Run an optional async prompt factory.""" - if not prompt_factory: - return None - return await prompt_factory() - - @classmethod - def _build_optional_prompt( - cls, - prompt_factory: Optional[StringPromptFactory], - ) -> Optional[str]: - """Run an optional synchronous prompt factory.""" - if not prompt_factory: - return None - return prompt_factory() - @classmethod def _print_system_prompts_for_debug( cls, @@ -1072,7 +1022,7 @@ def _print_system_prompts_for_debug( agent_name: str, provider_id: str, model_id: str, - prompts: List[str], + blocks: Iterable[SystemPromptBlock], ) -> None: """Print prompt blocks when FLOCKS_PRINT_SYSTEM_PROMPT is enabled.""" if os.getenv("FLOCKS_PRINT_SYSTEM_PROMPT", "").lower() not in ("1", "true", "yes"): @@ -1083,8 +1033,12 @@ def _print_system_prompts_for_debug( f"agent={agent_name} model={provider_id}/{model_id} ===" ) print(header, file=sys.stderr) - for idx, prompt in enumerate(prompts): - print(f"\n--- prompt[{idx}] ---\n{prompt}\n", file=sys.stderr) + for idx, block in enumerate(blocks): + print( + f"\n--- prompt[{idx}] {block.name} scope={block.cache_scope} " + f"---\n{block.content}\n", + file=sys.stderr, + ) print("=== end system_prompt ===\n", file=sys.stderr) @classmethod @@ -1141,22 +1095,92 @@ async def _is_builtin_system_subagent_session( return False @classmethod - async def _build_subagent_minimal_prompts( + def _append_turn_tail_blocks( cls, *, + blocks: List[SystemPromptBlock], + turn_context: TurnPromptContext, + static_cache: Optional[SystemPromptCache], + session_id: str, + ) -> None: + """Append per-turn values in the exact order sent to the model.""" + tail_values = [ + ("turn_additional_context", turn_context.additional_context), + ("text_tool_catalog", turn_context.text_tool_catalog), + ("tool_results_reminder", turn_context.tool_results_reminder), + ( + "repeated_tool_calls_reminder", + turn_context.repeated_tool_calls_reminder, + ), + ] + for name, content in tail_values: + normalized_content = cls._normalize_prompt_text(content) + block = cls._build_cached_prompt_block( + static_cache=static_cache, + name=name, + cache_scope="runtime_tail", + digest_inputs={"session_id": session_id, "content": content or ""}, + builder=lambda: normalized_content, + ) + if block is not None: + blocks.append(block) + + @classmethod + def _build_subagent_minimal_blocks( + cls, + *, + session_id: str, session_directory: Optional[str], agent_prompt: Optional[str], - ) -> List[str]: - """Build minimal system prompts for built-in system subagents.""" - prompts = [ - get_prompt_flocks_config_guard().strip(), - cls._normalize_prompt_text(agent_prompt), - cls._build_minimal_environment(session_directory), - ] - return [prompt for prompt in prompts if prompt] + turn_context: TurnPromptContext, + static_cache: Optional[SystemPromptCache], + ) -> List[SystemPromptBlock]: + """Build minimal prompt blocks for built-in system subagents.""" + blocks: List[SystemPromptBlock] = [] + guard_block = cls._build_cached_prompt_block( + static_cache=static_cache, + name="flocks_config_guard", + cache_scope="global", + digest_inputs={"prompt": get_prompt_flocks_config_guard()}, + builder=lambda: get_prompt_flocks_config_guard().strip(), + ) + if guard_block is not None: + blocks.append(guard_block) + + agent_block = cls._build_cached_prompt_block( + static_cache=static_cache, + name="agent_identity", + cache_scope="agent", + digest_inputs={"agent_prompt": agent_prompt or ""}, + builder=lambda: cls._normalize_prompt_text(agent_prompt), + ) + if agent_block is not None: + blocks.append(agent_block) + + environment_block = cls._build_cached_prompt_block( + static_cache=static_cache, + name="minimal_environment", + cache_scope="runtime_tail", + digest_inputs={ + "directory": session_directory, + "runtime_day": datetime.now().strftime("%Y-%m-%d"), + "platform": platform.system().lower(), + }, + builder=lambda: cls._build_minimal_environment(session_directory), + ) + if environment_block is not None: + blocks.append(environment_block) + + cls._append_turn_tail_blocks( + blocks=blocks, + turn_context=turn_context, + static_cache=static_cache, + session_id=session_id, + ) + return blocks @classmethod - async def build_system_prompts( + async def build_system_prompt_blocks( cls, *, session_id: str, @@ -1167,45 +1191,51 @@ async def build_system_prompts( model_id: str, execution_mode_prompt: Optional[str] = None, prompt_tool_names: Iterable[str] = (), - tool_revision: Optional[int] = None, memory_bootstrap_data: Optional[Dict[str, Any]] = None, static_cache: Optional[SystemPromptCache] = None, - sandbox_prompt_factory: Optional[AsyncPromptFactory] = None, - channel_context_prompt_factory: Optional[AsyncPromptFactory] = None, - tool_catalog_prompt_factory: Optional[StringPromptFactory] = None, - device_asset_prompt_factory: Optional[AsyncPromptFactory] = None, - device_revision: Optional[int] = None, + turn_context: Optional[TurnPromptContext] = None, use_text_tool_call_mode: bool = False, - ) -> List[str]: - """Build the runtime system prompt blocks for a session turn. + ) -> List[SystemPromptBlock]: + """Build the ordered system prompt blocks for a session turn. Stable identity and execution guidance come first, followed by session/workspace context, with runtime-only metadata kept at the - prompt tail. Cache mechanics are intentionally kept out of the block - construction below so this method reads as an ordered list of prompt - layers. + prompt tail. Runtime I/O is collected before this method so assembly is + deterministic and every downstream consumer sees the same blocks. """ + turn_context = turn_context or TurnPromptContext() vcs = "git" if session_directory else None - if await cls._is_builtin_system_subagent_session( - session_id=session_id, - agent_name=agent_name, - ): - prompts = await cls._build_subagent_minimal_prompts( + minimal_prompt = turn_context.minimal_prompt + if minimal_prompt is None: + minimal_prompt = await cls._is_builtin_system_subagent_session( + session_id=session_id, + agent_name=agent_name, + ) + if minimal_prompt: + minimal_blocks = cls._build_subagent_minimal_blocks( + session_id=session_id, session_directory=session_directory, agent_prompt=agent_prompt, + turn_context=turn_context, + static_cache=static_cache, ) cls._print_system_prompts_for_debug( session_id=session_id, agent_name=agent_name, provider_id=provider_id, model_id=model_id, - prompts=prompts, + blocks=minimal_blocks, ) - return prompts + return minimal_blocks normalized_tool_names = tuple(sorted(prompt_tool_names)) runtime_day = datetime.now().strftime("%Y-%m-%d") - custom_signature = SystemPrompt.custom_signature(directory=session_directory) + config_instructions = list(turn_context.config_instructions) + custom_signature = SystemPrompt.custom_signature( + directory=session_directory, + worktree=turn_context.worktree, + config_instructions=config_instructions, + ) memory_guidance = cls._build_memory_guidance_prompt( normalized_tool_names, memory_bootstrap_data, @@ -1217,7 +1247,11 @@ async def build_system_prompts( async def build_custom_context() -> Optional[str]: return cls._join_prompt_parts( - await SystemPrompt.custom(directory=session_directory), + await SystemPrompt.custom( + directory=session_directory, + worktree=turn_context.worktree, + config_instructions=config_instructions, + ), ) blocks: List[Optional[SystemPromptBlock]] = [ @@ -1281,23 +1315,32 @@ async def build_custom_context() -> Optional[str]: cache_scope="catalog", digest_inputs={ "agent_name": agent_name, - "tool_revision": tool_revision, + "tool_revision": turn_context.tool_revision, + "content": turn_context.tool_catalog or "", }, - builder=lambda: cls._build_optional_prompt(tool_catalog_prompt_factory) or "", + builder=lambda: cls._normalize_prompt_text( + turn_context.tool_catalog, + ), ), ] - if device_asset_prompt_factory: - blocks.append(await cls._build_cached_async_prompt_block( - static_cache=static_cache, - name="device_asset_hint", - cache_scope="runtime", - digest_inputs={ - "session_id": session_id, - "device_revision": device_revision, - }, - builder=device_asset_prompt_factory, - )) + if turn_context.device_asset_hint: + blocks.append( + cls._build_cached_prompt_block( + static_cache=static_cache, + name="device_asset_hint", + cache_scope="runtime", + digest_inputs={ + "session_id": session_id, + "device_revision": turn_context.device_revision, + "tool_revision": turn_context.tool_revision, + "content": turn_context.device_asset_hint, + }, + builder=lambda: cls._normalize_prompt_text( + turn_context.device_asset_hint, + ), + ) + ) blocks.append( cls._build_cached_prompt_block( @@ -1319,33 +1362,52 @@ async def build_custom_context() -> Optional[str]: static_cache=static_cache, name="context_files", cache_scope="workspace", - digest_inputs={"directory": session_directory, "signature": custom_signature}, - builder=build_custom_context, + digest_inputs={ + "directory": session_directory, + "worktree": turn_context.worktree, + "signature": custom_signature, + }, + loader=build_custom_context, ) blocks.append(custom_block) - if sandbox_prompt_factory: - blocks.append(await cls._build_cached_async_prompt_block( - static_cache=static_cache, - name="sandbox_context", - cache_scope="runtime", - digest_inputs={"session_id": session_id, "agent_name": agent_name}, - builder=sandbox_prompt_factory, - )) + if turn_context.sandbox_context: + blocks.append( + cls._build_cached_prompt_block( + static_cache=static_cache, + name="sandbox_context", + cache_scope="runtime_tail", + digest_inputs={ + "session_id": session_id, + "agent_name": agent_name, + "content": turn_context.sandbox_context, + }, + builder=lambda: cls._normalize_prompt_text( + turn_context.sandbox_context, + ), + ) + ) - if channel_context_prompt_factory: - blocks.append(await cls._build_cached_async_prompt_block( - static_cache=static_cache, - name="channel_context", - cache_scope="runtime", - digest_inputs={"session_id": session_id}, - builder=channel_context_prompt_factory, - )) + if turn_context.channel_context: + blocks.append( + cls._build_cached_prompt_block( + static_cache=static_cache, + name="channel_context", + cache_scope="runtime_tail", + digest_inputs={ + "session_id": session_id, + "content": turn_context.channel_context, + }, + builder=lambda: cls._normalize_prompt_text( + turn_context.channel_context, + ), + ) + ) blocks.append(cls._build_cached_prompt_block( static_cache=static_cache, name="runtime_metadata", - cache_scope="runtime", + cache_scope="runtime_tail", digest_inputs={ "session_id": session_id, "directory": session_directory, @@ -1364,28 +1426,55 @@ async def build_custom_context() -> Optional[str]: ), )) - cache_key = cls._system_prompt_cache_key( + resolved_blocks = [block for block in blocks if block is not None] + cls._append_turn_tail_blocks( + blocks=resolved_blocks, + turn_context=turn_context, + static_cache=static_cache, + session_id=session_id, + ) + cls._print_system_prompts_for_debug( session_id=session_id, agent_name=agent_name, provider_id=provider_id, model_id=model_id, - block_keys=[block.cache_key for block in blocks if block is not None], + blocks=resolved_blocks, ) - cached_prompts = cls._read_system_prompt_cache(static_cache, cache_key) - if cached_prompts is not None: - return cached_prompts + return resolved_blocks - prompts = cls._prompt_blocks_to_list(blocks) - cls._print_system_prompts_for_debug( + @classmethod + async def build_system_prompts( + cls, + *, + session_id: str, + session_directory: Optional[str], + agent_name: str, + agent_prompt: Optional[str], + provider_id: str, + model_id: str, + execution_mode_prompt: Optional[str] = None, + prompt_tool_names: Iterable[str] = (), + memory_bootstrap_data: Optional[Dict[str, Any]] = None, + static_cache: Optional[SystemPromptCache] = None, + turn_context: Optional[TurnPromptContext] = None, + use_text_tool_call_mode: bool = False, + ) -> List[str]: + """Compatibility API returning only the assembled prompt text.""" + blocks = await cls.build_system_prompt_blocks( session_id=session_id, + session_directory=session_directory, agent_name=agent_name, + agent_prompt=agent_prompt, provider_id=provider_id, model_id=model_id, - prompts=prompts, + execution_mode_prompt=execution_mode_prompt, + prompt_tool_names=prompt_tool_names, + memory_bootstrap_data=memory_bootstrap_data, + static_cache=static_cache, + turn_context=turn_context, + use_text_tool_call_mode=use_text_tool_call_mode, ) - - cls._write_system_prompt_cache(static_cache, cache_key, prompts) - return list(prompts) + return cls._prompt_blocks_to_list(blocks) @classmethod def _build_context_section(cls, context: ContextInfo) -> str: diff --git a/flocks/session/runner.py b/flocks/session/runner.py index b07a35ed3..0635982ea 100644 --- a/flocks/session/runner.py +++ b/flocks/session/runner.py @@ -19,7 +19,7 @@ from collections.abc import Mapping from datetime import datetime from typing import Optional, Dict, Any, List, Callable, Awaitable, Tuple -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace import httpcore import httpx @@ -28,7 +28,7 @@ from flocks.utils.id import Identifier from flocks.session.session import Session, SessionInfo from flocks.session.message import Message, MessageInfo, MessageRole, TextPart -from flocks.session.prompt import SessionPrompt +from flocks.session.prompt import SessionPrompt, SystemPromptBlock, TurnPromptContext from flocks.session.core.status import SessionStatus, SessionStatusRetry, SessionStatusBusy from flocks.session.core.defaults import ( DEFAULT_MAX_TOOL_STEPS, @@ -1509,24 +1509,19 @@ async def _process_step( self._log_perf("runner.process_step.tools_ready", tools_started_at, tool_count=len(tools)) prompt_tool_names = self._get_prompt_tool_names_from_schema(tools) - async def sandbox_prompt_factory() -> Optional[str]: - return await self._build_sandbox_prompt(agent) - - async def channel_context_prompt_factory() -> Optional[str]: - return await self._build_channel_context_prompt() - - async def device_asset_prompt_factory() -> Optional[str]: - return await self._build_device_asset_hint() - - try: - from flocks.tool.device.store import device_revision as get_device_revision - - current_device_revision = get_device_revision() - except Exception: - current_device_revision = None - prompts_started_at = time.perf_counter() - system_prompts = await SessionPrompt.build_system_prompts( + minimal_prompt = await SessionPrompt._is_builtin_system_subagent_session( + session_id=self.session.id, + agent_name=agent.name, + ) + turn_prompt_context = await self._build_turn_prompt_context( + agent=agent, + messages=messages, + last_user=last_user, + tools=tools, + minimal_prompt=minimal_prompt, + ) + system_prompts = await SessionPrompt.build_system_prompt_blocks( session_id=self.session.id, session_directory=self.session.directory, agent_name=agent.name, @@ -1539,57 +1534,15 @@ async def device_asset_prompt_factory() -> Optional[str]: plan_file=self._turn_plan_file, ), prompt_tool_names=prompt_tool_names, - tool_revision=ToolRegistry.revision(), memory_bootstrap_data=self._memory_bootstrap_data, static_cache=self._static_cache, - sandbox_prompt_factory=sandbox_prompt_factory, - channel_context_prompt_factory=channel_context_prompt_factory, - tool_catalog_prompt_factory=lambda: self._build_tool_catalog_prompt(agent), - device_asset_prompt_factory=device_asset_prompt_factory, - device_revision=current_device_revision, + turn_context=turn_prompt_context, use_text_tool_call_mode=self._should_use_text_tool_call_mode(), ) self._log_perf("runner.process_step.system_prompts_ready", prompts_started_at, prompt_count=len(system_prompts)) await self._run_session_start_hook(agent) - if self._turn_additional_context: - system_prompts.append(self._turn_additional_context) - - if self._should_use_text_tool_call_mode() and tools: - text_tool_catalog = self._build_text_tool_call_catalog_prompt(tools) - if text_tool_catalog: - system_prompts.append(text_tool_catalog) - - # If the last assistant message only contains tool results and no text, - # force a direct answer to avoid repeated tool calls. - last_assistant_msg = None - for msg in reversed(messages): - if msg.role == MessageRole.ASSISTANT: - last_assistant_msg = msg - break - if last_assistant_msg: - parts = await Message.parts(last_assistant_msg.id, self.session.id) - has_text = any(getattr(p, "type", None) == "text" and getattr(p, "text", "").strip() for p in parts) - has_tool_result = any( - getattr(p, "type", None) == "tool" and - getattr(getattr(p, "state", None), "status", None) in ("completed", "error", "running") - for p in parts - ) - if has_tool_result and not has_text: - from flocks.session.prompt_strings import PROMPT_TOOL_RESULTS_AVAILABLE - system_prompts.append(PROMPT_TOOL_RESULTS_AVAILABLE) - - if has_tool_result and self._should_warn_about_tool_loop(last_user_id=last_user.id): - state = self._get_tool_loop_guard_state(last_user_id=last_user.id) - log.warn("runner.repeated_tool_calls_detected", { - "tool_name": state.get("last_signature", "").split(":", 1)[0], - "exact_count": state.get("exact_count", 0), - "step": self._step, - }) - from flocks.session.prompt_strings import PROMPT_REPEATED_TOOL_CALLS - system_prompts.append(PROMPT_REPEATED_TOOL_CALLS) - # Convert messages to chat format with error handling try: queued_user_message_ids = self._get_queued_user_message_ids(messages) @@ -2100,6 +2053,133 @@ async def _record_usage_if_available( "error": str(exc), }) + async def _build_turn_prompt_context( + self, + *, + agent: AgentInfo, + messages: List[MessageInfo], + last_user: MessageInfo, + tools: List[Dict[str, Any]], + minimal_prompt: bool = False, + ) -> TurnPromptContext: + """Collect cached runtime values before deterministic prompt assembly.""" + if minimal_prompt: + return await self._add_turn_prompt_tail( + TurnPromptContext(minimal_prompt=True), + messages=messages, + last_user=last_user, + tools=tools, + ) + + from flocks.config import Config + from flocks.project.instance import Instance + + try: + from flocks.tool.device.store import device_revision + + current_device_revision = device_revision() + except Exception: + current_device_revision = None + + current_tool_revision = ToolRegistry.revision() + try: + config = await Config.get() + config_data = config.model_dump(by_alias=True, exclude_none=True) + config_instructions = tuple(config.instructions or ()) + except Exception as exc: + log.debug("runner.prompt_context.config_error", {"error": str(exc)}) + config_data = None + config_instructions = () + + worktree = Instance.get_worktree() + sandbox_context, channel_context, device_asset_hint = await asyncio.gather( + self._build_sandbox_prompt(agent, config_data=config_data), + self._build_channel_context_prompt(), + self._build_device_asset_hint(), + ) + source_context = TurnPromptContext( + tool_catalog=self._build_tool_catalog_prompt(agent), + device_asset_hint=device_asset_hint, + sandbox_context=sandbox_context, + channel_context=channel_context, + worktree=worktree, + config_instructions=config_instructions, + tool_revision=current_tool_revision, + device_revision=current_device_revision, + minimal_prompt=False, + ) + + return await self._add_turn_prompt_tail( + source_context, + messages=messages, + last_user=last_user, + tools=tools, + ) + + async def _add_turn_prompt_tail( + self, + source_context: TurnPromptContext, + *, + messages: List[MessageInfo], + last_user: MessageInfo, + tools: List[Dict[str, Any]], + ) -> TurnPromptContext: + """Add uncached per-step context and reminders to a source snapshot.""" + text_tool_catalog = None + if self._should_use_text_tool_call_mode() and tools: + text_tool_catalog = self._build_text_tool_call_catalog_prompt(tools) + + tool_results_reminder = None + repeated_tool_calls_reminder = None + last_assistant_msg = next( + ( + message + for message in reversed(messages) + if message.role == MessageRole.ASSISTANT + ), + None, + ) + if last_assistant_msg is not None: + parts = await Message.parts(last_assistant_msg.id, self.session.id) + has_text = any( + getattr(part, "type", None) == "text" + and getattr(part, "text", "").strip() + for part in parts + ) + has_tool_result = any( + getattr(part, "type", None) == "tool" + and getattr(getattr(part, "state", None), "status", None) + in ("completed", "error", "running") + for part in parts + ) + if has_tool_result and not has_text: + from flocks.session.prompt_strings import ( + PROMPT_TOOL_RESULTS_AVAILABLE, + ) + + tool_results_reminder = PROMPT_TOOL_RESULTS_AVAILABLE + + if has_tool_result and self._should_warn_about_tool_loop( + last_user_id=last_user.id, + ): + state = self._get_tool_loop_guard_state(last_user_id=last_user.id) + log.warn("runner.repeated_tool_calls_detected", { + "tool_name": state.get("last_signature", "").split(":", 1)[0], + "exact_count": state.get("exact_count", 0), + "step": self._step, + }) + from flocks.session.prompt_strings import PROMPT_REPEATED_TOOL_CALLS + + repeated_tool_calls_reminder = PROMPT_REPEATED_TOOL_CALLS + + return replace( + source_context, + additional_context=self._turn_additional_context, + text_tool_catalog=text_tool_catalog, + tool_results_reminder=tool_results_reminder, + repeated_tool_calls_reminder=repeated_tool_calls_reminder, + ) + async def _build_device_asset_hint(self) -> Optional[str]: """Return concise device-aware tool guidance plus enabled device summary.""" try: @@ -2142,15 +2222,22 @@ async def _build_device_asset_hint(self) -> Optional[str]: "如果同类设备有多个候选，不要猜测，先询问用户选择。" ) - async def _build_sandbox_prompt(self, agent: AgentInfo) -> Optional[str]: + async def _build_sandbox_prompt( + self, + agent: AgentInfo, + *, + config_data: Optional[Dict[str, Any]] = None, + ) -> Optional[str]: """Build sandbox context prompt when sandboxing is active.""" try: - from flocks.config import Config from flocks.session.core.session_state import get_main_session_id from flocks.sandbox.system_prompt import build_sandbox_system_prompt - cfg = await Config.get() - config_data = cfg.model_dump(by_alias=True, exclude_none=True) + if config_data is None: + from flocks.config import Config + + config = await Config.get() + config_data = config.model_dump(by_alias=True, exclude_none=True) session_key = self.session.id main_session_key = get_main_session_id() or self.session.id return await build_sandbox_system_prompt( @@ -2662,14 +2749,26 @@ def _build_tool_output_text(self, part: Any, tool_name: str, ctx_window_tokens: def _build_system_message_content( self, - system_prompts: List[str], + system_prompts: List[SystemPromptBlock] | List[str], ) -> str | list[dict[str, Any]]: """Format system prompts for the active provider. Anthropic supports structured system blocks, which lets us place a conservative cache breakpoint before the dynamic runtime tail. """ - prompt_parts = [prompt for prompt in system_prompts if prompt and prompt.strip()] + typed_blocks = [ + block + for block in system_prompts + if isinstance(block, SystemPromptBlock) and block.content.strip() + ] + if typed_blocks: + prompt_parts = [block.content for block in typed_blocks] + else: + prompt_parts = [ + prompt + for prompt in system_prompts + if isinstance(prompt, str) and prompt.strip() + ] if not prompt_parts: return "" @@ -2677,7 +2776,19 @@ def _build_system_message_content( if "anthropic" not in provider_lower: return "\n\n".join(prompt_parts) - cache_break_index = max(0, len(prompt_parts) - 3) + if typed_blocks: + first_runtime_tail = next( + ( + index + for index, block in enumerate(typed_blocks) + if block.cache_scope == "runtime_tail" + ), + len(typed_blocks), + ) + cache_break_index = max(0, first_runtime_tail - 1) + else: + # Compatibility for callers still passing plain strings. + cache_break_index = max(0, len(prompt_parts) - 3) blocks: list[dict[str, Any]] = [] for index, prompt in enumerate(prompt_parts): block: dict[str, Any] = { @@ -2692,7 +2803,7 @@ def _build_system_message_content( async def _to_chat_messages( self, messages: List[MessageInfo], - system_prompts: List[str], + system_prompts: List[SystemPromptBlock] | List[str], ) -> List[ChatMessage]: """ Convert messages to chat format with tool calls. diff --git a/tests/session/test_context_usage.py b/tests/session/test_context_usage.py index 2bc9db71a..798349b0f 100644 --- a/tests/session/test_context_usage.py +++ b/tests/session/test_context_usage.py @@ -1,8 +1,10 @@ from types import SimpleNamespace +from unittest.mock import AsyncMock import pytest from flocks.session import context_usage +from flocks.session.prompt import SystemPromptBlock, TurnPromptContext def _message( @@ -305,3 +307,58 @@ async def test_context_usage_splits_skill_and_delegation_tools(context_usage_moc ] tools_segment = next(segment for segment in snapshot.segments if segment.key == "tools") assert tools_segment.tokens == 30 + + +@pytest.mark.asyncio +async def test_system_prompt_estimate_uses_resolved_turn_context(monkeypatch): + captured = {} + + async def fake_build_system_prompt_blocks(**kwargs): + captured.update(kwargs) + return [ + SystemPromptBlock("stable", "a" * 40, "global"), + SystemPromptBlock("tail", "b" * 20, "runtime_tail"), + ] + + agent = SimpleNamespace(name="rex", prompt="agent prompt") + config = SimpleNamespace(instructions=["rules.md"]) + monkeypatch.setattr( + context_usage.SessionPrompt, + "build_system_prompt_blocks", + fake_build_system_prompt_blocks, + ) + monkeypatch.setattr( + "flocks.agent.registry.Agent.default_agent", + AsyncMock(return_value="rex"), + ) + monkeypatch.setattr( + "flocks.agent.registry.Agent.get", + AsyncMock(return_value=agent), + ) + monkeypatch.setattr( + "flocks.config.Config.get", + AsyncMock(return_value=config), + ) + monkeypatch.setattr( + "flocks.project.instance.Instance.get_worktree", + lambda: "/workspace", + ) + monkeypatch.setattr( + "flocks.tool.registry.ToolRegistry.revision", + lambda: 7, + ) + + tokens = await context_usage._estimate_system_prompt_tokens( + "sess-1", + session=SimpleNamespace(directory="/workspace/project"), + messages=[], + provider_id="openai", + model_id="gpt-5", + ) + + assert tokens == 15 + assert captured["turn_context"] == TurnPromptContext( + worktree="/workspace", + config_instructions=("rules.md",), + tool_revision=7, + ) diff --git a/tests/session/test_prompt_tokens.py b/tests/session/test_prompt_tokens.py index f7a1bd1fd..e921e5ebb 100644 --- a/tests/session/test_prompt_tokens.py +++ b/tests/session/test_prompt_tokens.py @@ -26,6 +26,7 @@ PromptTemplate, SessionPrompt, SystemPrompt, + TurnPromptContext, ) from flocks.session import prompt_strings @@ -265,7 +266,9 @@ async def test_builtin_system_subagent_child_uses_minimal_prompt(self): agent_prompt="You are Rex Junior.", provider_id="anthropic", model_id="claude-sonnet", - tool_catalog_prompt_factory=lambda: "SHOULD_NOT_APPEAR", + turn_context=TurnPromptContext( + tool_catalog="SHOULD_NOT_APPEAR", + ), ) assert len(prompts) == 3 @@ -305,6 +308,38 @@ async def test_builtin_system_subagent_root_uses_full_prompt(self): assert len(prompts) > 2 assert any(PROMPT_DEFAULT.strip() in prompt for prompt in prompts) + @pytest.mark.asyncio + async def test_full_prompt_loads_worktree_and_config_instructions( + self, + tmp_path: Path, + ) -> None: + nested = tmp_path / "src" / "package" + nested.mkdir(parents=True) + (tmp_path / "AGENTS.md").write_text("project rules", encoding="utf-8") + (nested / "extra-rules.md").write_text("extra rules", encoding="utf-8") + + with patch.object( + SessionPrompt, + "_is_builtin_system_subagent_session", + AsyncMock(return_value=False), + ): + prompts = await SessionPrompt.build_system_prompts( + session_id="ses-instructions", + session_directory=str(nested), + agent_name="rex", + agent_prompt="agent prompt", + provider_id="openai", + model_id="gpt-5", + turn_context=TurnPromptContext( + worktree=str(tmp_path), + config_instructions=("extra-rules.md",), + ), + ) + + combined = "\n\n".join(prompts) + assert "project rules" in combined + assert "extra rules" in combined + @pytest.mark.asyncio async def test_evolution_subagent_child_uses_full_prompt(self): agent = AgentInfo( diff --git a/tests/session/test_runner_step.py b/tests/session/test_runner_step.py index efee61f53..1cf74de4c 100644 --- a/tests/session/test_runner_step.py +++ b/tests/session/test_runner_step.py @@ -33,7 +33,12 @@ StepResult, ToolCall, ) -from flocks.session.prompt import SessionPrompt, get_prompt_flocks_config_guard +from flocks.session.prompt import ( + SessionPrompt, + SystemPromptBlock, + TurnPromptContext, + get_prompt_flocks_config_guard, +) from flocks.session.core.defaults import DEFAULT_MAX_TOOL_STEPS from flocks.session.session import Session, SessionInfo from flocks.tool.registry import ToolCategory, ToolInfo @@ -641,14 +646,21 @@ async def test_build_system_prompts_reuses_loop_static_cache(self): env_mock = MagicMock(return_value=["env prompt"]) runtime_mock = MagicMock(return_value=["runtime prompt"]) custom_mock = AsyncMock(return_value=["custom prompt"]) - sandbox_mock = AsyncMock(return_value="sandbox prompt") - channel_mock = AsyncMock(return_value="channel prompt") - device_mock = AsyncMock(return_value="device prompt") + turn_context = TurnPromptContext( + sandbox_context="sandbox prompt", + channel_context="channel prompt", + tool_catalog="tool catalog", + device_asset_hint="device prompt", + tool_revision=1, + device_revision=7, + ) - with patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), \ - patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), \ - patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), \ - patch("flocks.session.prompt.SystemPrompt.custom", custom_mock): + with ( + patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), + patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), + patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), + patch("flocks.session.prompt.SystemPrompt.custom", custom_mock), + ): prompts1 = await SessionPrompt.build_system_prompts( session_id=session.id, session_directory=session.directory, @@ -657,13 +669,8 @@ async def test_build_system_prompts_reuses_loop_static_cache(self): provider_id=runner1.provider_id, model_id=runner1.model_id, prompt_tool_names=("read",), - tool_revision=1, static_cache=shared_cache, - sandbox_prompt_factory=sandbox_mock, - channel_context_prompt_factory=channel_mock, - tool_catalog_prompt_factory=lambda: "tool catalog", - device_asset_prompt_factory=device_mock, - device_revision=7, + turn_context=turn_context, ) prompts2 = await SessionPrompt.build_system_prompts( session_id=session.id, @@ -673,22 +680,14 @@ async def test_build_system_prompts_reuses_loop_static_cache(self): provider_id=runner2.provider_id, model_id=runner2.model_id, prompt_tool_names=("read",), - tool_revision=1, static_cache=shared_cache, - sandbox_prompt_factory=sandbox_mock, - channel_context_prompt_factory=channel_mock, - tool_catalog_prompt_factory=lambda: "tool catalog", - device_asset_prompt_factory=device_mock, - device_revision=7, + turn_context=turn_context, ) assert prompts1 == prompts2 env_mock.assert_called_once() runtime_mock.assert_called_once() custom_mock.assert_awaited_once() - sandbox_mock.assert_awaited_once() - channel_mock.assert_awaited_once() - device_mock.assert_awaited_once() @pytest.mark.asyncio async def test_build_system_prompts_orders_stable_prefix_before_runtime_tail(self): @@ -704,15 +703,13 @@ async def test_build_system_prompts_orders_stable_prefix_before_runtime_tail(sel "inject": True, }, } - sandbox_mock = AsyncMock(return_value="sandbox prompt") - channel_mock = AsyncMock(return_value="channel prompt") - device_mock = AsyncMock(return_value="device prompt") - - with patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), \ - patch.object(SessionPrompt, "_build_tool_guidance_prompt", return_value="tool protocol"), \ - patch("flocks.session.prompt.SystemPrompt.environment_stable", return_value=["env prompt"]), \ - patch("flocks.session.prompt.SystemPrompt.runtime_metadata", return_value=["runtime prompt"]), \ - patch("flocks.session.prompt.SystemPrompt.custom", AsyncMock(return_value=["custom prompt"])): + with ( + patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), + patch.object(SessionPrompt, "_build_tool_guidance_prompt", return_value="tool protocol"), + patch("flocks.session.prompt.SystemPrompt.environment_stable", return_value=["env prompt"]), + patch("flocks.session.prompt.SystemPrompt.runtime_metadata", return_value=["runtime prompt"]), + patch("flocks.session.prompt.SystemPrompt.custom", AsyncMock(return_value=["custom prompt"])), + ): prompts = await SessionPrompt.build_system_prompts( session_id=session.id, session_directory=session.directory, @@ -730,11 +727,17 @@ async def test_build_system_prompts_orders_stable_prefix_before_runtime_tail(sel "write", ), memory_bootstrap_data=memory_bootstrap_data, - tool_catalog_prompt_factory=lambda: "tool catalog", - device_asset_prompt_factory=device_mock, - device_revision=3, - sandbox_prompt_factory=sandbox_mock, - channel_context_prompt_factory=channel_mock, + turn_context=TurnPromptContext( + tool_catalog="tool catalog", + device_asset_hint="device prompt", + sandbox_context="sandbox prompt", + channel_context="channel prompt", + additional_context="additional prompt", + text_tool_catalog="text tool catalog", + tool_results_reminder="tool results reminder", + repeated_tool_calls_reminder="tool loop reminder", + device_revision=3, + ), ) assert prompts == [ @@ -751,6 +754,10 @@ async def test_build_system_prompts_orders_stable_prefix_before_runtime_tail(sel "sandbox prompt", "channel prompt", "runtime prompt", + "additional prompt", + "text tool catalog", + "tool results reminder", + "tool loop reminder", ] @pytest.mark.asyncio @@ -764,16 +771,12 @@ async def test_build_system_prompts_rebuilds_when_tool_revision_changes(self): env_mock = MagicMock(return_value=["env prompt"]) runtime_mock = MagicMock(return_value=["runtime prompt"]) custom_mock = AsyncMock(return_value=["custom prompt"]) - sandbox_mock = AsyncMock(return_value="sandbox prompt") - channel_mock = AsyncMock(return_value="channel prompt") - device_mock = AsyncMock(return_value="device prompt") - - catalog_prompts = iter(["tool catalog v1", "tool catalog v2"]) - - with patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), \ - patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), \ - patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), \ - patch("flocks.session.prompt.SystemPrompt.custom", custom_mock): + with ( + patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), + patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), + patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), + patch("flocks.session.prompt.SystemPrompt.custom", custom_mock), + ): prompts1 = await SessionPrompt.build_system_prompts( session_id=session.id, session_directory=session.directory, @@ -782,13 +785,15 @@ async def test_build_system_prompts_rebuilds_when_tool_revision_changes(self): provider_id=runner.provider_id, model_id=runner.model_id, prompt_tool_names=("read",), - tool_revision=1, static_cache=shared_cache, - sandbox_prompt_factory=sandbox_mock, - channel_context_prompt_factory=channel_mock, - tool_catalog_prompt_factory=lambda: next(catalog_prompts), - device_asset_prompt_factory=device_mock, - device_revision=1, + turn_context=TurnPromptContext( + sandbox_context="sandbox prompt", + channel_context="channel prompt", + tool_catalog="tool catalog v1", + device_asset_hint="device prompt", + tool_revision=1, + device_revision=1, + ), ) agent.prompt = "agent prompt v2" prompts2 = await SessionPrompt.build_system_prompts( @@ -799,13 +804,15 @@ async def test_build_system_prompts_rebuilds_when_tool_revision_changes(self): provider_id=runner.provider_id, model_id=runner.model_id, prompt_tool_names=("read",), - tool_revision=2, static_cache=shared_cache, - sandbox_prompt_factory=sandbox_mock, - channel_context_prompt_factory=channel_mock, - tool_catalog_prompt_factory=lambda: next(catalog_prompts), - device_asset_prompt_factory=device_mock, - device_revision=1, + turn_context=TurnPromptContext( + sandbox_context="sandbox prompt", + channel_context="channel prompt", + tool_catalog="tool catalog v2", + device_asset_hint="device prompt", + tool_revision=2, + device_revision=1, + ), ) assert prompts1 != prompts2 @@ -816,8 +823,6 @@ async def test_build_system_prompts_rebuilds_when_tool_revision_changes(self): env_mock.assert_called_once() runtime_mock.assert_called_once() custom_mock.assert_awaited_once() - sandbox_mock.assert_awaited_once() - channel_mock.assert_awaited_once() @pytest.mark.asyncio async def test_build_system_prompts_reuses_static_device_hint_cache(self): @@ -830,14 +835,20 @@ async def test_build_system_prompts_reuses_static_device_hint_cache(self): env_mock = MagicMock(return_value=["env prompt"]) runtime_mock = MagicMock(return_value=["runtime prompt"]) custom_mock = AsyncMock(return_value=["custom prompt"]) - sandbox_mock = AsyncMock(return_value="sandbox prompt") - channel_mock = AsyncMock(return_value="channel prompt") - device_mock = AsyncMock(return_value="device prompt") + turn_context = TurnPromptContext( + sandbox_context="sandbox prompt", + channel_context="channel prompt", + tool_catalog="tool catalog", + device_asset_hint="device prompt", + tool_revision=1, + ) - with patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), \ - patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), \ - patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), \ - patch("flocks.session.prompt.SystemPrompt.custom", custom_mock): + with ( + patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), + patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), + patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), + patch("flocks.session.prompt.SystemPrompt.custom", custom_mock), + ): prompts1 = await SessionPrompt.build_system_prompts( session_id=session.id, session_directory=session.directory, @@ -846,12 +857,8 @@ async def test_build_system_prompts_reuses_static_device_hint_cache(self): provider_id=runner.provider_id, model_id=runner.model_id, prompt_tool_names=("read",), - tool_revision=1, static_cache=shared_cache, - sandbox_prompt_factory=sandbox_mock, - channel_context_prompt_factory=channel_mock, - tool_catalog_prompt_factory=lambda: "tool catalog", - device_asset_prompt_factory=device_mock, + turn_context=turn_context, ) prompts2 = await SessionPrompt.build_system_prompts( session_id=session.id, @@ -861,12 +868,8 @@ async def test_build_system_prompts_reuses_static_device_hint_cache(self): provider_id=runner.provider_id, model_id=runner.model_id, prompt_tool_names=("read",), - tool_revision=1, static_cache=shared_cache, - sandbox_prompt_factory=sandbox_mock, - channel_context_prompt_factory=channel_mock, - tool_catalog_prompt_factory=lambda: "tool catalog", - device_asset_prompt_factory=device_mock, + turn_context=turn_context, ) assert prompts1 == prompts2 @@ -874,9 +877,6 @@ async def test_build_system_prompts_reuses_static_device_hint_cache(self): env_mock.assert_called_once() runtime_mock.assert_called_once() custom_mock.assert_awaited_once() - sandbox_mock.assert_awaited_once() - channel_mock.assert_awaited_once() - device_mock.assert_awaited_once() @pytest.mark.asyncio async def test_build_system_prompts_rebuilds_when_device_revision_changes(self): @@ -889,14 +889,12 @@ async def test_build_system_prompts_rebuilds_when_device_revision_changes(self): env_mock = MagicMock(return_value=["env prompt"]) runtime_mock = MagicMock(return_value=["runtime prompt"]) custom_mock = AsyncMock(return_value=["custom prompt"]) - sandbox_mock = AsyncMock(return_value="sandbox prompt") - channel_mock = AsyncMock(return_value="channel prompt") - device_prompts = iter(["device prompt v1", "device prompt v2"]) - - with patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), \ - patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), \ - patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), \ - patch("flocks.session.prompt.SystemPrompt.custom", custom_mock): + with ( + patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), + patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), + patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), + patch("flocks.session.prompt.SystemPrompt.custom", custom_mock), + ): prompts1 = await SessionPrompt.build_system_prompts( session_id=session.id, session_directory=session.directory, @@ -905,13 +903,15 @@ async def test_build_system_prompts_rebuilds_when_device_revision_changes(self): provider_id=runner.provider_id, model_id=runner.model_id, prompt_tool_names=("read",), - tool_revision=1, static_cache=shared_cache, - sandbox_prompt_factory=sandbox_mock, - channel_context_prompt_factory=channel_mock, - tool_catalog_prompt_factory=lambda: "tool catalog", - device_asset_prompt_factory=AsyncMock(side_effect=lambda: next(device_prompts)), - device_revision=1, + turn_context=TurnPromptContext( + sandbox_context="sandbox prompt", + channel_context="channel prompt", + tool_catalog="tool catalog", + device_asset_hint="device prompt v1", + tool_revision=1, + device_revision=1, + ), ) prompts2 = await SessionPrompt.build_system_prompts( session_id=session.id, @@ -921,13 +921,15 @@ async def test_build_system_prompts_rebuilds_when_device_revision_changes(self): provider_id=runner.provider_id, model_id=runner.model_id, prompt_tool_names=("read",), - tool_revision=1, static_cache=shared_cache, - sandbox_prompt_factory=sandbox_mock, - channel_context_prompt_factory=channel_mock, - tool_catalog_prompt_factory=lambda: "tool catalog", - device_asset_prompt_factory=AsyncMock(side_effect=lambda: next(device_prompts)), - device_revision=2, + turn_context=TurnPromptContext( + sandbox_context="sandbox prompt", + channel_context="channel prompt", + tool_catalog="tool catalog", + device_asset_hint="device prompt v2", + tool_revision=1, + device_revision=2, + ), ) assert prompts1 != prompts2 @@ -936,8 +938,6 @@ async def test_build_system_prompts_rebuilds_when_device_revision_changes(self): env_mock.assert_called_once() runtime_mock.assert_called_once() custom_mock.assert_awaited_once() - sandbox_mock.assert_awaited_once() - channel_mock.assert_awaited_once() @pytest.mark.asyncio async def test_build_system_prompts_rebuilds_when_agent_prompt_changes(self): @@ -951,10 +951,12 @@ async def test_build_system_prompts_rebuilds_when_agent_prompt_changes(self): runtime_mock = MagicMock(return_value=["runtime prompt"]) custom_mock = AsyncMock(return_value=["custom prompt"]) - with patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), \ - patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), \ - patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), \ - patch("flocks.session.prompt.SystemPrompt.custom", custom_mock): + with ( + patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), + patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), + patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), + patch("flocks.session.prompt.SystemPrompt.custom", custom_mock), + ): prompts1 = await SessionPrompt.build_system_prompts( session_id=session.id, session_directory=session.directory, @@ -963,8 +965,8 @@ async def test_build_system_prompts_rebuilds_when_agent_prompt_changes(self): provider_id=runner.provider_id, model_id=runner.model_id, prompt_tool_names=("read",), - tool_revision=1, static_cache=shared_cache, + turn_context=TurnPromptContext(tool_revision=1), ) agent.prompt = "agent prompt v2" prompts2 = await SessionPrompt.build_system_prompts( @@ -975,8 +977,8 @@ async def test_build_system_prompts_rebuilds_when_agent_prompt_changes(self): provider_id=runner.provider_id, model_id=runner.model_id, prompt_tool_names=("read",), - tool_revision=1, static_cache=shared_cache, + turn_context=TurnPromptContext(tool_revision=1), ) assert prompts1 != prompts2 @@ -1109,10 +1111,12 @@ async def test_filesystem_memory_guidance_depends_on_tool_names(self): runtime_mock = MagicMock(return_value=["runtime prompt"]) custom_mock = AsyncMock(return_value=["custom prompt"]) - with patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), \ - patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), \ - patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), \ - patch("flocks.session.prompt.SystemPrompt.custom", custom_mock): + with ( + patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), + patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), + patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), + patch("flocks.session.prompt.SystemPrompt.custom", custom_mock), + ): prompts_with_memory = await SessionPrompt.build_system_prompts( session_id=session.id, session_directory=session.directory, @@ -1128,9 +1132,9 @@ async def test_filesystem_memory_guidance_depends_on_tool_names(self): "read", "write", ), - tool_revision=1, memory_bootstrap_data=runner._memory_bootstrap_data, static_cache=shared_cache, + turn_context=TurnPromptContext(tool_revision=1), ) prompts_without_memory = await SessionPrompt.build_system_prompts( session_id=session.id, @@ -1140,9 +1144,9 @@ async def test_filesystem_memory_guidance_depends_on_tool_names(self): provider_id=runner.provider_id, model_id=runner.model_id, prompt_tool_names=("read",), - tool_revision=1, memory_bootstrap_data=runner._memory_bootstrap_data, static_cache=shared_cache, + turn_context=TurnPromptContext(tool_revision=1), ) assert prompts_with_memory != prompts_without_memory @@ -1476,15 +1480,25 @@ async def test_to_chat_messages_uses_structured_anthropic_system_blocks(monkeypa monkeypatch.setattr(runner_mod.Message, "parts", AsyncMock(return_value=[])) monkeypatch.setattr(runner_mod.Message, "get_text_content", AsyncMock(return_value="hello")) - chat_messages = await runner._to_chat_messages( - [message], - ["provider prompt", "agent prompt", "context prompt", "runtime prompt"], - ) + prompt_blocks = [ + SystemPromptBlock(name, content, cache_scope) + for name, content, cache_scope in ( + ("provider", "provider prompt", "global"), + ("agent", "agent prompt", "agent"), + ("context", "context prompt", "workspace"), + ("sandbox", "sandbox prompt", "runtime_tail"), + ("runtime", "runtime prompt", "runtime_tail"), + ("reminder", "reminder prompt", "runtime_tail"), + ) + ] + + chat_messages = await runner._to_chat_messages([message], prompt_blocks) assert chat_messages[0].role == "system" assert isinstance(chat_messages[0].content, list) - assert chat_messages[0].content[1]["cache_control"] == {"type": "ephemeral"} - assert chat_messages[0].content[-1]["text"] == "runtime prompt" + assert chat_messages[0].content[2]["cache_control"] == {"type": "ephemeral"} + assert "cache_control" not in chat_messages[0].content[3] + assert chat_messages[0].content[-1]["text"] == "reminder prompt" @pytest.mark.asyncio @@ -2246,7 +2260,7 @@ async def fake_create(*args, **kwargs): 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_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) monkeypatch.setattr( runner, @@ -2304,7 +2318,7 @@ async def fake_to_chat_messages(_messages, _system_prompts): # noqa: ANN001 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_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) monkeypatch.setattr(runner, "_to_chat_messages", fake_to_chat_messages) monkeypatch.setattr(runner_mod.Message, "get_text_content", AsyncMock(return_value="queued")) @@ -2352,7 +2366,7 @@ async def fake_call_llm(*_args, **_kwargs): 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_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) monkeypatch.setattr( runner, @@ -2428,7 +2442,7 @@ async def fake_call_llm(*_args, **_kwargs): 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_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) monkeypatch.setattr( runner, @@ -2661,7 +2675,7 @@ async def publish_event(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(runner_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) monkeypatch.setattr(SessionRunner, "_build_callable_tool_schema", AsyncMock(return_value=[])) monkeypatch.setattr(runner_mod.SessionRetry, "sleep", AsyncMock(return_value=None)) @@ -2708,7 +2722,7 @@ async def test_process_step_uses_loaded_tool_schema_names_for_prompt_guidance(mo provider = MagicMock() provider.is_configured.return_value = True assistant_msg = SimpleNamespace(id="msg_assistant_prompt_guidance") - build_system_prompts = AsyncMock(return_value=[]) + build_system_prompt_blocks = AsyncMock(return_value=[]) tool_schema = [ {"type": "function", "function": {"name": "memory_search", "description": "", "parameters": {}}}, {"type": "function", "function": {"name": "bash", "description": "", "parameters": {}}}, @@ -2717,7 +2731,7 @@ async def test_process_step_uses_loaded_tool_schema_names_for_prompt_guidance(mo 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", build_system_prompts) + monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompt_blocks", build_system_prompt_blocks) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=tool_schema)) monkeypatch.setattr( runner, @@ -2737,8 +2751,12 @@ async def test_process_step_uses_loaded_tool_schema_names_for_prompt_guidance(mo result = await runner._process_step([last_user], last_user) assert result.content == "done" - build_system_prompts.assert_awaited_once() - assert build_system_prompts.await_args.kwargs["prompt_tool_names"] == ("bash", "memory_search") + build_system_prompt_blocks.assert_awaited_once() + assert build_system_prompt_blocks.await_args.kwargs["prompt_tool_names"] == ("bash", "memory_search") + assert isinstance( + build_system_prompt_blocks.await_args.kwargs["turn_context"], + TurnPromptContext, + ) @pytest.mark.asyncio @@ -2766,7 +2784,7 @@ async def test_process_step_records_usage_after_success(monkeypatch): 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_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) monkeypatch.setattr( runner, @@ -2791,7 +2809,7 @@ async def test_process_step_records_usage_after_success(monkeypatch): @pytest.mark.asyncio -async def test_process_step_passes_device_hint_factory_into_build_system_prompts(monkeypatch): +async def test_process_step_passes_resolved_device_hint_into_turn_context(monkeypatch): runner = _make_runner("ses_runner_device_hint_order") runner.callbacks = RunnerCallbacks(on_error=AsyncMock()) @@ -2808,12 +2826,12 @@ async def test_process_step_passes_device_hint_factory_into_build_system_prompts provider = MagicMock() provider.is_configured.return_value = True assistant_msg = SimpleNamespace(id="msg_assistant_device_hint_order") - build_system_prompts = AsyncMock(return_value=["provider", "tool catalog awareness", "device hint"]) + build_system_prompt_blocks = AsyncMock(return_value=["provider", "tool catalog awareness", "device hint"]) 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", build_system_prompts) + monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompt_blocks", build_system_prompt_blocks) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) device_hint_mock = AsyncMock(return_value="device hint") monkeypatch.setattr(runner, "_build_device_asset_hint", device_hint_mock) @@ -2836,11 +2854,11 @@ async def test_process_step_passes_device_hint_factory_into_build_system_prompts result = await runner._process_step([last_user], last_user) assert result.content == "done" - build_system_prompts.assert_awaited_once() - kwargs = build_system_prompts.await_args.kwargs - assert kwargs["device_revision"] == 9 - assert kwargs["device_asset_prompt_factory"] is not None - assert await kwargs["device_asset_prompt_factory"]() == "device hint" + build_system_prompt_blocks.assert_awaited_once() + turn_context = build_system_prompt_blocks.await_args.kwargs["turn_context"] + assert turn_context.device_revision == 9 + assert turn_context.device_asset_hint == "device hint" + device_hint_mock.assert_awaited_once() @pytest.mark.asyncio @@ -2871,7 +2889,7 @@ async def test_process_step_empty_retry_records_usage_per_attempt(monkeypatch): 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_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) monkeypatch.setattr( runner, @@ -2933,7 +2951,7 @@ async def test_process_step_retries_empty_transport_exception(monkeypatch): 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_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) monkeypatch.setattr( runner, @@ -2985,7 +3003,7 @@ async def _call_llm(*_args, **_kwargs): monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) monkeypatch.setattr( runner_mod.SessionPrompt, - "build_system_prompts", + "build_system_prompt_blocks", AsyncMock(return_value=[]), ) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) @@ -3033,7 +3051,7 @@ async def test_process_step_uses_default_max_steps_when_agent_steps_missing(monk 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_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=sentinel_tools)) monkeypatch.setattr( runner, @@ -3082,7 +3100,7 @@ async def test_process_step_respects_explicit_agent_steps_over_default(monkeypat 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_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=sentinel_tools)) monkeypatch.setattr( runner, @@ -3136,7 +3154,7 @@ async def fake_call_llm(self, provider, messages, tools, agent, assistant_msg): ))) 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_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) monkeypatch.setattr(runner_mod.Message, "get_text_content", AsyncMock(return_value="hi")) monkeypatch.setattr(runner_mod.Message, "parts", AsyncMock(return_value=[])) monkeypatch.setattr(runner_mod.Message, "create", create_mock) diff --git a/tests/session/test_session_runner_tool_only_message.py b/tests/session/test_session_runner_tool_only_message.py index 28e879e4f..75b973b57 100644 --- a/tests/session/test_session_runner_tool_only_message.py +++ b/tests/session/test_session_runner_tool_only_message.py @@ -81,7 +81,7 @@ async def fake_get_prompt_tool_names(self, agent): # noqa: ANN001 del self, agent return () - async def fake_build_system_prompts(*args, **kwargs): # noqa: ANN002, ANN003 + async def fake_build_system_prompt_blocks(*args, **kwargs): # noqa: ANN002, ANN003 del args, kwargs return [] @@ -100,7 +100,11 @@ async def fake_call_llm(self, provider, messages, tools, agent, assistant_msg): monkeypatch.setattr(Provider, "apply_config", fake_apply_config) monkeypatch.setattr(Agent, "get", fake_agent_get) monkeypatch.setattr(SessionRunner, "_get_prompt_tool_names", fake_get_prompt_tool_names) - monkeypatch.setattr(SessionPrompt, "build_system_prompts", fake_build_system_prompts) + monkeypatch.setattr( + SessionPrompt, + "build_system_prompt_blocks", + fake_build_system_prompt_blocks, + ) monkeypatch.setattr(SessionRunner, "_build_callable_tool_schema", fake_build_callable_tool_schema) monkeypatch.setattr(SessionRunner, "_to_chat_messages", fake_to_chat_messages) monkeypatch.setattr(SessionRunner, "_call_llm", fake_call_llm) From c40cc9a5f8493c82963b3fae14441c37511d5976 Mon Sep 17 00:00:00 2001 From: zhougongyan Date: Mon, 17 Aug 2026 16:13:57 +0800 Subject: [PATCH 07/63] chore(agents): make todo prompt guidance unconditional Keep Rex and Hephaestus on their existing todo workflow unconditionally and tighten Rex verification guidance without changing runtime task compatibility. Co-Authored-By: Claude Opus 4.6 --- .../agent/agents/hephaestus/prompt_builder.py | 43 +------------ flocks/agent/agents/rex/prompt_builder.py | 60 +++++++------------ tests/agent/test_prompt_builders.py | 36 +++++++++++ 3 files changed, 58 insertions(+), 81 deletions(-) create mode 100644 tests/agent/test_prompt_builders.py diff --git a/flocks/agent/agents/hephaestus/prompt_builder.py b/flocks/agent/agents/hephaestus/prompt_builder.py index 9617affcb..3d220164b 100644 --- a/flocks/agent/agents/hephaestus/prompt_builder.py +++ b/flocks/agent/agents/hephaestus/prompt_builder.py @@ -29,7 +29,6 @@ def inject( available_agents=available_agents, available_tools=tools, available_skills=skills, - use_task_system=False, ) @@ -37,7 +36,6 @@ def build_hephaestus_prompt( available_agents: List["AvailableAgent"], available_tools: List["AvailableTool"], available_skills: List["AvailableSkill"], - use_task_system: bool = False, ) -> str: from flocks.agent.prompt_utils import ( build_agent_selection_table, @@ -62,7 +60,7 @@ def build_hephaestus_prompt( oracle_section = build_oracle_section(available_agents) hard_blocks = build_hard_blocks_section() anti_patterns = build_anti_patterns_section() - todo_discipline = _todo_discipline_section(use_task_system) + todo_discipline = _todo_discipline_section() template = """You are Hephaestus, an autonomous deep worker for software engineering. @@ -245,44 +243,7 @@ def build_hephaestus_prompt( return prompt -def _todo_discipline_section(use_task_system: bool) -> str: - if use_task_system: - return """## Task Discipline (NON-NEGOTIABLE) - -**Track ALL multi-step work with tasks. This is your execution backbone.** - -### When to Create Tasks (MANDATORY) - -| Trigger | Action | -|---------|--------| -| 2+ step task | `TaskCreate` FIRST, atomic breakdown | -| Uncertain scope | `TaskCreate` to clarify thinking | -| Complex single task | Break down into trackable steps | - -### Workflow (STRICT) - -1. **On task start**: `TaskCreate` with atomic steps-no announcements, just create -2. **Before each step**: `TaskUpdate(status="in_progress")` (ONE at a time) -3. **After each step**: `TaskUpdate(status="completed")` IMMEDIATELY (NEVER batch) -4. **Scope changes**: Update tasks BEFORE proceeding - -### Why This Matters - -- **Execution anchor**: Tasks prevent drift from original request -- **Recovery**: If interrupted, tasks enable seamless continuation -- **Accountability**: Each task = explicit commitment to deliver - -### Anti-Patterns (BLOCKING) - -| Violation | Why It Fails | -|-----------|--------------| -| Skipping tasks on multi-step work | Steps get forgotten, user has no visibility | -| Batch-completing multiple tasks | Defeats real-time tracking purpose | -| Proceeding without `in_progress` | No indication of current work | -| Finishing without completing tasks | Task appears incomplete | - -**NO TASKS ON MULTI-STEP WORK = INCOMPLETE WORK.**""" - +def _todo_discipline_section() -> str: return """## Todo Discipline (NON-NEGOTIABLE) **Track ALL multi-step work with todos. This is your execution backbone.** diff --git a/flocks/agent/agents/rex/prompt_builder.py b/flocks/agent/agents/rex/prompt_builder.py index e74413457..86872c759 100644 --- a/flocks/agent/agents/rex/prompt_builder.py +++ b/flocks/agent/agents/rex/prompt_builder.py @@ -30,7 +30,6 @@ def inject( available_tools=tools, available_skills=skills, available_workflows=workflows or [], - use_task_system=False, ) @@ -39,7 +38,6 @@ def build_dynamic_rex_prompt( available_tools: List["AvailableTool"], available_skills: List["AvailableSkill"], available_workflows: Optional[List["AvailableWorkflow"]] = None, - use_task_system: bool = False, ) -> str: from flocks.agent.prompt_utils import ( build_agent_selection_table, @@ -58,12 +56,8 @@ def build_dynamic_rex_prompt( im_send_section = _build_im_send_pointer_section() anti_patterns = _build_rex_anti_patterns_section() command_guidance_section = _build_command_guidance_section() - task_management_section = _task_management_section(use_task_system) - todo_hook_note = ( - "YOUR TASK CREATION WOULD BE TRACKED BY HOOK([SYSTEM REMINDER - TASK CONTINUATION])" - if use_task_system - else "YOUR TODO CREATION WOULD BE TRACKED BY HOOK([SYSTEM REMINDER - TODO CONTINUATION])" - ) + task_management_section = _task_management_section() + todo_hook_note = "YOUR TODO CREATION WOULD BE TRACKED BY HOOK([SYSTEM REMINDER - TODO CONTINUATION])" template = """ You are "Rex" - Powerful AI orchestrator for security operations. @@ -143,13 +137,12 @@ def build_dynamic_rex_prompt( - Match existing codebase patterns when editing. - Fix bugs minimally; do not refactor during a bugfix unless required. - Keep search bounded: stop when you have enough context, when results repeat, or when direct evidence already answers the question. -- For independent parallel branches whose results are needed this turn, emit multiple foreground `delegate_task` / `task` tool calls in the same assistant turn. The runtime executes those sibling tool calls concurrently and returns all tool results before you continue. +- For independent parallel branches whose results are needed this turn, emit multiple foreground `delegate_task` tool calls in the same assistant turn. The runtime executes those sibling tool calls concurrently and returns all tool results before you continue. - Do not use `run_in_background=true`; background subagent execution is disabled. ## 5. Verify - - Use `lsp` for symbol-aware checks when useful, and run relevant tests on changed files before considering the work complete. -- Run relevant build or test commands before finalizing when the affected area has them. +- After code changes, run the lint/typecheck/tests. If tests fail, iterate until they pass before finalizing. - Verification evidence is mandatory: clean diagnostics, successful commands, or an explicit note about pre-existing failures. - Verify delegated work against expected behavior, codebase patterns, and any `must-do` / `must-not-do` requirements. @@ -278,55 +271,42 @@ def _build_clarification_protocol() -> str: ```""" -def _task_management_section(use_task_system: bool) -> str: - title = "Task Management" if use_task_system else "Todo Management" - unit = "tasks" if use_task_system else "todos" - create_action = "`TaskCreate`" if use_task_system else '`todo(action="write")`' - progress_action = ( - '`TaskUpdate(status="in_progress")`' - if use_task_system - else "mark `in_progress`" - ) - complete_action = ( - '`TaskUpdate(status="completed")`' - if use_task_system - else "mark `completed`" - ) +def _task_management_section() -> str: clarification_protocol = _build_clarification_protocol() - return f""" -## {title} + return f""" +## Todo Management -Use {unit} as the primary coordination mechanism for non-trivial execution work. +Use todos as the primary coordination mechanism for non-trivial execution work. ### When They Are Mandatory | Trigger | Action | |---------|--------| -| Multi-step work (2+ steps) | Create {unit} first | -| Uncertain scope | Create {unit} to structure the work | -| User request with multiple items | Create {unit} first | -| Complex single task | Break it into {unit} | +| Multi-step work (2+ steps) | Create todos first | +| Uncertain scope | Create todos to structure the work | +| User request with multiple items | Create todos first | +| Complex single task | Break it into todos | ### Operating Rules -1. Start with {create_action} before implementation work begins. -2. ONLY add {unit} when the user wants execution, not when they only want analysis or planning. -3. Before each step, {progress_action}. Keep only one item in progress. -4. After each step, {complete_action} immediately. Never batch updates. -5. If scope changes, update the {unit} before continuing. +1. Start with `todo(action="write")` before implementation work begins. +2. ONLY add todos when the user wants execution, not when they only want analysis or planning. +3. Before each step, mark it `in_progress`. Keep only one item in progress. +4. After each step, mark it `completed` immediately. Never batch updates. +5. If scope changes, update the todos before continuing. ### Failure Modes | Violation | Why It Breaks the Workflow | |-----------|----------------------------| -| Skipping {unit} on non-trivial work | The user loses progress visibility and steps get dropped | -| Batch-completing multiple {unit} | Real-time tracking becomes meaningless | +| Skipping todos on non-trivial work | The user loses progress visibility and steps get dropped | +| Batch-completing multiple todos | Real-time tracking becomes meaningless | | Proceeding without an in-progress item | It is unclear what is being worked on | | Finishing without closing items | The work appears incomplete | {clarification_protocol} -""" +""" def _build_security_priority_section(available_agents: List["AvailableAgent"]) -> str: diff --git a/tests/agent/test_prompt_builders.py b/tests/agent/test_prompt_builders.py new file mode 100644 index 000000000..8ce1f31af --- /dev/null +++ b/tests/agent/test_prompt_builders.py @@ -0,0 +1,36 @@ +"""Direct tests for Rex and Hephaestus prompt builders.""" + +import inspect + +from flocks.agent.agents.hephaestus.prompt_builder import build_hephaestus_prompt +from flocks.agent.agents.rex.prompt_builder import build_dynamic_rex_prompt + + +def test_rex_prompt_uses_todos_and_delegate_task_only(): + prompt = build_dynamic_rex_prompt([], [], [], []) + + assert "## Todo Management" in prompt + assert "YOUR TODO CREATION WOULD BE TRACKED BY HOOK" in prompt + assert "multiple foreground `delegate_task` tool calls" in prompt + assert "`delegate_task` / `task`" not in prompt + assert "After code changes, run the lint/typecheck/tests." in prompt + assert "If tests fail, iterate until they pass before finalizing." in prompt + assert "explicit note about pre-existing failures" in prompt + assert "Verify delegated work against expected behavior" in prompt + assert "TaskCreate" not in prompt + assert "TaskUpdate" not in prompt + + +def test_hephaestus_prompt_uses_existing_todo_discipline(): + prompt = build_hephaestus_prompt([], [], []) + + assert "## Todo Discipline (NON-NEGOTIABLE)" in prompt + assert "Track ALL multi-step work with todos." in prompt + assert '`todo(action="write")`' in prompt + assert "TaskCreate" not in prompt + assert "TaskUpdate" not in prompt + + +def test_prompt_builder_signatures_exclude_task_system_flag(): + assert "use_task_system" not in inspect.signature(build_dynamic_rex_prompt).parameters + assert "use_task_system" not in inspect.signature(build_hephaestus_prompt).parameters From 2351407ade3cd831d153240c8579da531cc03553 Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Mon, 17 Aug 2026 17:23:19 +0800 Subject: [PATCH 08/63] fix(delegation): preserve migration compatibility --- flocks/config/config.py | 23 +-- flocks/permission/helpers.py | 31 ++-- flocks/session/context_usage.py | 2 +- tests/config/test_config.py | 33 ++++ tests/session/test_context_usage.py | 2 +- tui/flocks/command/index.test.ts | 13 ++ tui/flocks/command/index.ts | 7 + tui/flocks/config/config.test.ts | 28 ++++ tui/flocks/config/config.ts | 22 ++- tui/flocks/session/message-v2.test.ts | 26 ++++ tui/flocks/session/message-v2.ts | 20 ++- tui/flocks/session/prompt.ts | 209 +++++++++++++++++++++++++- 12 files changed, 377 insertions(+), 39 deletions(-) create mode 100644 tui/flocks/command/index.test.ts create mode 100644 tui/flocks/session/message-v2.test.ts diff --git a/flocks/config/config.py b/flocks/config/config.py index 6bef59b5c..68d4d19b7 100644 --- a/flocks/config/config.py +++ b/flocks/config/config.py @@ -29,7 +29,7 @@ class PermissionAction(str, Enum): PermissionRule = Union[PermissionAction, Dict[str, Union[PermissionAction, Dict[str, PermissionAction]]]] -_LEGACY_TODO_TOOL_NAMES = {"todowrite", "todoread"} +_LEGACY_TODO_TOOL_NAMES = ("todowrite", "todoread") _LEGACY_PERMISSION_TOOL_NAMES = { **{name: "todo" for name in _LEGACY_TODO_TOOL_NAMES}, "task": "delegate_task", @@ -42,16 +42,6 @@ def _canonical_permission_tool_name(tool: str) -> str: def _merge_permission_action(existing: Any, incoming: Any) -> Any: """Merge duplicate legacy permission names conservatively.""" - def contains_deny(value: Any) -> bool: - raw_value = value.value if hasattr(value, "value") else value - if raw_value == PermissionAction.DENY.value: - return True - if isinstance(raw_value, dict): - return any(contains_deny(item) for item in raw_value.values()) - return False - - if contains_deny(existing) or contains_deny(incoming): - return PermissionAction.DENY if isinstance(existing, dict) and isinstance(incoming, dict): merged = dict(existing) for pattern, action in incoming.items(): @@ -62,6 +52,17 @@ def contains_deny(value: Any) -> bool: else: merged[pattern] = action return merged + + def contains_deny(value: Any) -> bool: + raw_value = value.value if hasattr(value, "value") else value + if raw_value == PermissionAction.DENY.value: + return True + if isinstance(raw_value, dict): + return any(contains_deny(item) for item in raw_value.values()) + return False + + if contains_deny(existing) or contains_deny(incoming): + return PermissionAction.DENY return existing if existing is not None else incoming diff --git a/flocks/permission/helpers.py b/flocks/permission/helpers.py index 5ea907711..1561e73ab 100644 --- a/flocks/permission/helpers.py +++ b/flocks/permission/helpers.py @@ -14,17 +14,7 @@ def _merge_legacy_permission(existing: Any, incoming: Any) -> Any: - """Merge aliases without allowing a legacy deny to become an allow.""" - def contains_deny(value: Any) -> bool: - raw_value = getattr(value, "value", value) - if raw_value == "deny": - return True - if isinstance(raw_value, dict): - return any(contains_deny(item) for item in raw_value.values()) - return False - - if contains_deny(existing) or contains_deny(incoming): - return "deny" + """Merge a legacy alias into an existing canonical permission.""" if isinstance(existing, dict) and isinstance(incoming, dict): merged = dict(existing) for pattern, action in incoming.items(): @@ -35,12 +25,29 @@ def contains_deny(value: Any) -> bool: else: merged[pattern] = action return merged + + def contains_deny(value: Any) -> bool: + raw_value = getattr(value, "value", value) + if raw_value == "deny": + return True + if isinstance(raw_value, dict): + return any(contains_deny(item) for item in raw_value.values()) + return False + + if contains_deny(existing) or contains_deny(incoming): + return "deny" return existing def _canonicalize_permission_config(config: Dict[str, Any]) -> Dict[str, Any]: - canonical: Dict[str, Any] = {} + canonical = { + key: value + for key, value in config.items() + if key not in _LEGACY_PERMISSION_NAMES + } for key, value in config.items(): + if key not in _LEGACY_PERMISSION_NAMES: + continue name = _LEGACY_PERMISSION_NAMES.get(key, key) if name in canonical: canonical[name] = _merge_legacy_permission(canonical[name], value) diff --git a/flocks/session/context_usage.py b/flocks/session/context_usage.py index 376b41044..ec8f284c8 100644 --- a/flocks/session/context_usage.py +++ b/flocks/session/context_usage.py @@ -23,7 +23,7 @@ log = Log.create(service="context-usage") UsageSource = Literal["observed", "estimated"] -DELEGATION_TOOLS = {"delegate_task"} +DELEGATION_TOOLS = {"delegate_task", "task"} ZERO_VISIBLE_SEGMENTS = {"agentDelegation"} diff --git a/tests/config/test_config.py b/tests/config/test_config.py index c957244e9..f17e9c69c 100644 --- a/tests/config/test_config.py +++ b/tests/config/test_config.py @@ -8,6 +8,7 @@ from unittest.mock import patch from flocks.config.config import Config, GlobalConfig, ConfigInfo, PermissionAction, PermissionConfig +from flocks.permission.helpers import from_config @pytest.fixture(autouse=True) @@ -164,6 +165,38 @@ def test_legacy_task_permission_name_migrates_to_delegate_task(): assert "task" not in dumped +@pytest.mark.parametrize( + "raw_permission", + [ + { + "task": {"explore": "allow", "legacy-only": "ask"}, + "delegate_task": {"explore": "ask", "canonical-only": "allow"}, + }, + { + "delegate_task": {"explore": "ask", "canonical-only": "allow"}, + "task": {"explore": "allow", "legacy-only": "ask"}, + }, + ], +) +def test_legacy_task_permission_merge_is_order_independent(raw_permission): + direct_rules = from_config(raw_permission) + model_rules = from_config(PermissionConfig.model_validate(raw_permission)) + + expected = { + ("delegate_task", "explore", "ask"), + ("delegate_task", "legacy-only", "ask"), + ("delegate_task", "canonical-only", "allow"), + } + assert { + (rule.permission, rule.pattern, rule.level.value) + for rule in direct_rules + } == expected + assert { + (rule.permission, rule.pattern, rule.level.value) + for rule in model_rules + } == expected + + def test_legacy_todo_tool_flags_migrate_to_todo_permission(): config = ConfigInfo.model_validate({ "tools": { diff --git a/tests/session/test_context_usage.py b/tests/session/test_context_usage.py index 303dcf88a..d9323daa4 100644 --- a/tests/session/test_context_usage.py +++ b/tests/session/test_context_usage.py @@ -280,7 +280,7 @@ async def test_context_usage_splits_skill_and_delegation_tools(context_usage_moc ), SimpleNamespace( type="tool", - tool="delegate_task", + tool="task", state=SimpleNamespace(input={}, output="d" * 40, time={"start": 4}), ), SimpleNamespace( diff --git a/tui/flocks/command/index.test.ts b/tui/flocks/command/index.test.ts new file mode 100644 index 000000000..9ee70946d --- /dev/null +++ b/tui/flocks/command/index.test.ts @@ -0,0 +1,13 @@ +import { describe, expect, test } from "bun:test" +import { Command } from "." + +describe("delegated commands", () => { + test("delegates commands that explicitly request a subtask", () => { + expect(Command.shouldDelegate({ subtask: true }, "primary")).toBe(true) + }) + + test("delegates subagent commands unless explicitly disabled", () => { + expect(Command.shouldDelegate({}, "subagent")).toBe(true) + expect(Command.shouldDelegate({ subtask: false }, "subagent")).toBe(false) + }) +}) diff --git a/tui/flocks/command/index.ts b/tui/flocks/command/index.ts index cd3b711b8..4d02feeca 100644 --- a/tui/flocks/command/index.ts +++ b/tui/flocks/command/index.ts @@ -26,6 +26,7 @@ export namespace Command { description: z.string().optional(), agent: z.string().optional(), model: z.string().optional(), + subtask: z.boolean().optional(), mcp: z.boolean().optional(), // workaround for zod not supporting async functions natively so we use getters // https://zod.dev/v4/changelog?id=zfunction @@ -39,6 +40,10 @@ export namespace Command { // for some reason zod is inferring `string` for z.promise(z.string()).or(z.string()) so we have to manually override it export type Info = Omit, "template"> & { template: Promise | string } + export function shouldDelegate(command: Pick, agentMode: string | undefined) { + return (agentMode === "subagent" && command.subtask !== false) || command.subtask === true + } + export function hints(template: string): string[] { const result: string[] = [] const numbered = template.match(/\$\d+/g) @@ -72,6 +77,7 @@ export namespace Command { get template() { return PROMPT_REVIEW.replace("${path}", Instance.worktree) }, + subtask: true, hints: hints(PROMPT_REVIEW), }, } @@ -82,6 +88,7 @@ export namespace Command { agent: command.agent, model: command.model, description: command.description, + subtask: command.subtask, get template() { return command.template }, diff --git a/tui/flocks/config/config.test.ts b/tui/flocks/config/config.test.ts index 7c8452467..3683ad309 100644 --- a/tui/flocks/config/config.test.ts +++ b/tui/flocks/config/config.test.ts @@ -10,4 +10,32 @@ describe("permission aliases", () => { }), ).toEqual({ delegate_task: "deny" }) }) + + test.each([ + { + task: { explore: "allow", "legacy-only": "ask" }, + delegate_task: { explore: "ask", "canonical-only": "allow" }, + }, + { + delegate_task: { explore: "ask", "canonical-only": "allow" }, + task: { explore: "allow", "legacy-only": "ask" }, + }, + ])("merges task into delegate_task independently of key order", (permission) => { + expect(Config.Permission.parse(permission)).toEqual({ + delegate_task: { + explore: "ask", + "legacy-only": "ask", + "canonical-only": "allow", + }, + }) + }) + + test("preserves the legacy subtask command option during migration", () => { + expect( + Config.Command.parse({ + template: "Review this change", + subtask: true, + }).subtask, + ).toBe(true) + }) }) diff --git a/tui/flocks/config/config.ts b/tui/flocks/config/config.ts index a43d70e11..c9b259860 100644 --- a/tui/flocks/config/config.ts +++ b/tui/flocks/config/config.ts @@ -507,14 +507,24 @@ export namespace Config { const assignPermission = (target: Record, tool: string, action: PermissionRule) => { const canonical = canonicalPermissionToolName(tool) - const containsDeny = (value: PermissionRule): boolean => + const existing = target[canonical] + if (typeof existing === "object" && typeof action === "object") { + target[canonical] = { ...action, ...existing } + for (const pattern of Object.keys(action)) { + if (pattern in existing && (existing[pattern] === "deny" || action[pattern] === "deny")) { + target[canonical][pattern] = "deny" + } + } + return + } + const containsDeny = (value: PermissionRule | undefined): boolean => value === "deny" || (typeof value === "object" && value !== null && Object.values(value).some((item) => containsDeny(item))) - if ((canonical in target && containsDeny(target[canonical])) || containsDeny(action)) { + if (containsDeny(existing) || containsDeny(action)) { target[canonical] = "deny" return } - if (!(canonical in target)) target[canonical] = action + if (existing === undefined) target[canonical] = action } const permissionTransform = (x: unknown): Record => { @@ -523,7 +533,10 @@ export namespace Config { const { __originalKeys, ...rest } = obj const result: Record = {} const keys = __originalKeys ?? Object.keys(rest) - for (const key of keys) { + for (const key of keys.filter((key) => canonicalPermissionToolName(key) === key)) { + if (key in rest) assignPermission(result, key, rest[key] as PermissionRule) + } + for (const key of keys.filter((key) => canonicalPermissionToolName(key) !== key)) { if (key in rest) assignPermission(result, key, rest[key] as PermissionRule) } return result @@ -563,6 +576,7 @@ export namespace Config { description: z.string().optional(), agent: z.string().optional(), model: z.string().optional(), + subtask: z.boolean().optional(), }) export type Command = z.infer diff --git a/tui/flocks/session/message-v2.test.ts b/tui/flocks/session/message-v2.test.ts new file mode 100644 index 000000000..39418e243 --- /dev/null +++ b/tui/flocks/session/message-v2.test.ts @@ -0,0 +1,26 @@ +import { describe, expect, test } from "bun:test" +import { MessageV2 } from "./message-v2" + +describe("stored message parts", () => { + test("normalizes a legacy subtask part to ignored text", () => { + expect( + MessageV2.normalizeStoredPart({ + id: "part_legacy_subtask", + sessionID: "session_legacy", + messageID: "message_legacy", + type: "subtask", + prompt: "Review the current changes", + description: "Review changes", + agent: "reviewer", + }), + ).toEqual({ + id: "part_legacy_subtask", + sessionID: "session_legacy", + messageID: "message_legacy", + type: "text", + text: "", + ignored: true, + metadata: { legacyPartType: "subtask" }, + }) + }) +}) diff --git a/tui/flocks/session/message-v2.ts b/tui/flocks/session/message-v2.ts index ff596603b..c9cbd3e61 100644 --- a/tui/flocks/session/message-v2.ts +++ b/tui/flocks/session/message-v2.ts @@ -330,6 +330,22 @@ export namespace MessageV2 { }) export type Part = z.infer + export function normalizeStoredPart(part: unknown): Part { + if (typeof part === "object" && part !== null && "type" in part && part.type === "subtask") { + const legacy = PartBase.parse(part) + return { + id: legacy.id, + sessionID: legacy.sessionID, + messageID: legacy.messageID, + type: "text", + text: "", + ignored: true, + metadata: { legacyPartType: "subtask" }, + } + } + return Part.parse(part) + } + export const Assistant = Base.extend({ role: z.literal("assistant"), time: z.object({ @@ -559,8 +575,8 @@ export namespace MessageV2 { export const parts = fn(Identifier.schema("message"), async (messageID) => { const result = [] as MessageV2.Part[] for (const item of await Storage.list(["part", messageID])) { - const read = await Storage.read(item) - result.push(read) + const read = await Storage.read(item) + result.push(normalizeStoredPart(read)) } result.sort((a, b) => (a.id > b.id ? 1 : -1)) return result diff --git a/tui/flocks/session/prompt.ts b/tui/flocks/session/prompt.ts index 22313c167..18d267c59 100644 --- a/tui/flocks/session/prompt.ts +++ b/tui/flocks/session/prompt.ts @@ -42,6 +42,7 @@ import { iife } from "@/util/iife" import { Shell } from "@/shell/shell" import { Truncate } from "@/tool/truncation" import { Ripgrep } from "../file/ripgrep" +import { DelegateTaskTool } from "@/tool/delegate-task" // @ts-ignore globalThis.AI_SDK_LOG_WARNINGS = false @@ -1475,6 +1476,167 @@ NOTE: At any point in time through this workflow you should feel free to ask the const argsRegex = /(?:\[Image\s+\d+\]|"[^"]*"|'[^']*'|[^\s"']+)/gi const placeholderRegex = /\$(\d+)/g const quoteTrimRegex = /^["']|["']$/g + + async function executeDelegatedCommand(input: { + sessionID: string + parentID: string + parentAgent: Agent.Info + parentModel: { providerID: string; modelID: string } + taskAgent: Agent.Info + taskModel: Provider.Model + prompt: string + description: string + command: string + variant?: string + }) { + const assistantMessage = (await Session.updateMessage({ + id: Identifier.ascending("message"), + role: "assistant", + parentID: input.parentID, + sessionID: input.sessionID, + mode: input.parentAgent.name, + agent: input.parentAgent.name, + path: { + cwd: Instance.directory, + root: Instance.worktree, + }, + cost: 0, + tokens: { + input: 0, + output: 0, + reasoning: 0, + cache: { read: 0, write: 0 }, + }, + modelID: input.taskModel.id, + providerID: input.taskModel.providerID, + time: { created: Date.now() }, + })) as MessageV2.Assistant + const taskArgs = { + prompt: input.prompt, + description: input.description, + subagent_type: input.taskAgent.name, + command: input.command, + } + let part = (await Session.updatePart({ + id: Identifier.ascending("part"), + messageID: assistantMessage.id, + sessionID: input.sessionID, + type: "tool", + callID: ulid(), + tool: DelegateTaskTool.id, + state: { + status: "running", + input: taskArgs, + time: { start: Date.now() }, + }, + })) as MessageV2.ToolPart + const abort = start(input.sessionID) + if (!abort) throw new Session.BusyError(input.sessionID) + SessionStatus.set(input.sessionID, { type: "busy" }) + using _ = defer(() => cancel(input.sessionID)) + const taskTool = await DelegateTaskTool.init({ agent: input.parentAgent }) + const taskCtx: Tool.Context = { + agent: input.parentAgent.name, + messageID: assistantMessage.id, + sessionID: input.sessionID, + abort, + callID: part.callID, + extra: { bypassAgentCheck: true }, + async metadata(metadata) { + part = (await Session.updatePart({ + ...part, + state: { + ...part.state, + ...metadata, + }, + })) as MessageV2.ToolPart + }, + async ask(request) { + const session = await Session.get(input.sessionID) + await PermissionNext.ask({ + ...request, + sessionID: input.sessionID, + ruleset: PermissionNext.merge(input.parentAgent.permission, session.permission ?? []), + }) + }, + } + + await Plugin.trigger( + "tool.execute.before", + { + tool: DelegateTaskTool.id, + sessionID: input.sessionID, + callID: part.callID, + }, + { args: taskArgs }, + ) + let executionError: Error | undefined + const result = await taskTool.execute(taskArgs, taskCtx).catch((error) => { + executionError = error instanceof Error ? error : new Error(String(error)) + log.error("delegated command failed", { + error: executionError, + agent: input.taskAgent.name, + description: input.description, + }) + return undefined + }) + await Plugin.trigger( + "tool.execute.after", + { + tool: DelegateTaskTool.id, + sessionID: input.sessionID, + callID: part.callID, + }, + result, + ) + + assistantMessage.finish = "tool-calls" + assistantMessage.time.completed = Date.now() + await Session.updateMessage(assistantMessage) + if (result && part.state.status === "running") { + await Session.updatePart({ + ...part, + state: { + status: "completed", + input: part.state.input, + title: result.title, + metadata: result.metadata, + output: result.output, + attachments: result.attachments, + time: { ...part.state.time, end: Date.now() }, + }, + } satisfies MessageV2.ToolPart) + } else if (!result) { + await Session.updatePart({ + ...part, + state: { + status: "error", + error: executionError ? `Tool execution failed: ${executionError.message}` : "Tool execution failed", + time: { + start: part.state.status === "running" ? part.state.time.start : Date.now(), + end: Date.now(), + }, + metadata: part.metadata, + input: part.state.input, + }, + } satisfies MessageV2.ToolPart) + } + + cancel(input.sessionID) + return prompt({ + sessionID: input.sessionID, + model: input.parentModel, + agent: input.parentAgent.name, + variant: input.variant, + parts: [ + { + type: "text", + text: "Summarize the delegate_task output above and continue with your task.", + synthetic: true, + }, + ], + }) + } /** * Regular expression to match @ file references in text * Matches @ followed by file paths, excluding commas, periods at end of sentences, and backticks @@ -1565,6 +1727,7 @@ NOTE: At any point in time through this workflow you should feel free to ask the const templateParts = await resolvePromptParts(template) const parts = [...templateParts, ...(input.parts ?? [])] + const isSubtask = Command.shouldDelegate(command, agent.mode) await Plugin.trigger( "command.execute.before", @@ -1576,14 +1739,44 @@ NOTE: At any point in time through this workflow you should feel free to ask the { parts }, ) - const result = (await prompt({ - sessionID: input.sessionID, - messageID: input.messageID, - model: taskModel, - agent: agentName, - parts, - variant: input.variant, - })) as MessageV2.WithParts + const result = (await (async () => { + if (!isSubtask) { + return prompt({ + sessionID: input.sessionID, + messageID: input.messageID, + model: taskModel, + agent: agentName, + parts, + variant: input.variant, + }) + } + + const parentAgentName = input.agent ?? (await Agent.defaultAgent()) + const parentAgent = await Agent.get(parentAgentName) + if (!parentAgent) throw new Error(`Agent not found: "${parentAgentName}"`) + const parentModel = input.model ? Provider.parseModel(input.model) : await lastModel(input.sessionID) + const userMessage = await prompt({ + sessionID: input.sessionID, + messageID: input.messageID, + model: parentModel, + agent: parentAgent.name, + parts: templateParts.filter((part) => part.type === "text"), + variant: input.variant, + noReply: true, + }) + return executeDelegatedCommand({ + sessionID: input.sessionID, + parentID: userMessage.info.id, + parentAgent, + parentModel, + taskAgent: agent, + taskModel: await Provider.getModel(taskModel.providerID, taskModel.modelID), + prompt: templateParts.find((part) => part.type === "text")?.text ?? "", + description: command.description ?? "", + command: input.command, + variant: input.variant, + }) + })()) as MessageV2.WithParts Bus.publish(Command.Event.Executed, { name: input.command, From 2cf940f4c546cdeb1cfcde2e98af4a88a4c13c8a Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Mon, 17 Aug 2026 18:11:54 +0800 Subject: [PATCH 09/63] refactor(prompt): remove obsolete TUI copies --- .../session/prompt/anthropic-20250930.txt | 164 ------------------ tui/flocks/session/prompt/anthropic.txt | 101 ----------- tui/flocks/session/prompt/anthropic_spoof.txt | 1 - tui/flocks/session/prompt/beast.txt | 147 ---------------- tui/flocks/session/prompt/build-switch.txt | 5 - tui/flocks/session/prompt/codex_header.txt | 73 -------- tui/flocks/session/prompt/copilot-gpt-5.txt | 143 --------------- tui/flocks/session/prompt/gemini.txt | 155 ----------------- tui/flocks/session/prompt/max-steps.txt | 16 -- tui/flocks/session/prompt/plan.txt | 26 --- tui/flocks/session/prompt/qwen.txt | 107 ------------ 11 files changed, 938 deletions(-) delete mode 100644 tui/flocks/session/prompt/anthropic-20250930.txt delete mode 100644 tui/flocks/session/prompt/anthropic.txt delete mode 100644 tui/flocks/session/prompt/anthropic_spoof.txt delete mode 100644 tui/flocks/session/prompt/beast.txt delete mode 100644 tui/flocks/session/prompt/build-switch.txt delete mode 100644 tui/flocks/session/prompt/codex_header.txt delete mode 100644 tui/flocks/session/prompt/copilot-gpt-5.txt delete mode 100644 tui/flocks/session/prompt/gemini.txt delete mode 100644 tui/flocks/session/prompt/max-steps.txt delete mode 100644 tui/flocks/session/prompt/plan.txt delete mode 100644 tui/flocks/session/prompt/qwen.txt diff --git a/tui/flocks/session/prompt/anthropic-20250930.txt b/tui/flocks/session/prompt/anthropic-20250930.txt deleted file mode 100644 index 01ac8e5b5..000000000 --- a/tui/flocks/session/prompt/anthropic-20250930.txt +++ /dev/null @@ -1,164 +0,0 @@ -You are an interactive CLI tool that helps users with software engineering tasks. Use the instructions below and the tools available to you to assist the user. - -IMPORTANT: Assist with defensive security tasks only. Refuse to create, modify, or improve code that may be used maliciously. Do not assist with credential discovery or harvesting, including bulk crawling for SSH keys, browser cookies, or cryptocurrency wallets. Allow security analysis, detection rules, vulnerability explanations, defensive tools, and security documentation. -IMPORTANT: You must NEVER generate or guess URLs for the user unless you are confident that the URLs are for helping the user with programming. You may use URLs provided by the user in their messages or local files. - -If the user asks for help or wants to give feedback inform them of the following: -- /help: Get help with using Claude Code -- To give feedback, users should report the issue at https://github.com/anthropics/claude-code/issues - -When the user directly asks about Claude Code (eg. "can Claude Code do...", "does Claude Code have..."), or asks in second person (eg. "are you able...", "can you do..."), or asks how to use a specific Claude Code feature (eg. implement a hook, or write a slash command), use the WebFetch tool to gather information to answer the question from Claude Code docs. The list of available docs is available at https://docs.claude.com/en/docs/claude-code/claude_code_docs_map.md. - -# Tone and style -You should be concise, direct, and to the point, while providing complete information and matching the level of detail you provide in your response with the level of complexity of the user's query or the work you have completed. -A concise response is generally less than 4 lines, not including tool calls or code generated. You should provide more detail when the task is complex or when the user asks you to. -IMPORTANT: You should minimize output tokens as much as possible while maintaining helpfulness, quality, and accuracy. Only address the specific task at hand, avoiding tangential information unless absolutely critical for completing the request. If you can answer in 1-3 sentences or a short paragraph, please do. -IMPORTANT: You should NOT answer with unnecessary preamble or postamble (such as explaining your code or summarizing your action), unless the user asks you to. -Do not add additional code explanation summary unless requested by the user. After working on a file, briefly confirm that you have completed the task, rather than providing an explanation of what you did. -Answer the user's question directly, avoiding any elaboration, explanation, introduction, conclusion, or excessive details. Brief answers are best, but be sure to provide complete information. You MUST avoid extra preamble before/after your response, such as "The answer is .", "Here is the content of the file..." or "Based on the information provided, the answer is..." or "Here is what I will do next...". - -Here are some examples to demonstrate appropriate verbosity: - -user: 2 + 2 -assistant: 4 - - - -user: what is 2+2? -assistant: 4 - - - -user: is 11 a prime number? -assistant: Yes - - - -user: what command should I run to list files in the current directory? -assistant: ls - - - -user: what command should I run to watch files in the current directory? -assistant: [runs ls to list the files in the current directory, then read docs/commands in the relevant file to find out how to watch files] -npm run dev - - - -user: How many golf balls fit inside a jetta? -assistant: 150000 - - - -user: what files are in the directory src/? -assistant: [runs ls and sees foo.c, bar.c, baz.c] -user: which file contains the implementation of foo? -assistant: src/foo.c - -When you run a non-trivial bash command, you should explain what the command does and why you are running it, to make sure the user understands what you are doing (this is especially important when you are running a command that will make changes to the user's system). -Remember that your output will be displayed on a command line interface. Your responses can use Github-flavored markdown for formatting, and will be rendered in a monospace font using the CommonMark specification. -Output text to communicate with the user; all text you output outside of tool use is displayed to the user. Only use tools to complete tasks. Never use tools like Bash or code comments as means to communicate with the user during the session. -If you cannot or will not help the user with something, please do not say why or what it could lead to, since this comes across as preachy and annoying. Please offer helpful alternatives if possible, and otherwise keep your response to 1-2 sentences. -Only use emojis if the user explicitly requests it. Avoid using emojis in all communication unless asked. -IMPORTANT: Keep your responses short, since they will be displayed on a command line interface. - -# Proactiveness -You are allowed to be proactive, but only when the user asks you to do something. You should strive to strike a balance between: -- Doing the right thing when asked, including taking actions and follow-up actions -- Not surprising the user with actions you take without asking -For example, if the user asks you how to approach something, you should do your best to answer their question first, and not immediately jump into taking actions. - -# Professional objectivity -Prioritize technical accuracy and truthfulness over validating the user's beliefs. Focus on facts and problem-solving, providing direct, objective technical info without any unnecessary superlatives, praise, or emotional validation. It is best for the user if Claude honestly applies the same rigorous standards to all ideas and disagrees when necessary, even if it may not be what the user wants to hear. Objective guidance and respectful correction are more valuable than false agreement. Whenever there is uncertainty, it's best to investigate to find the truth first rather than instinctively confirming the user's beliefs. - -# Task Management -You have access to the `todo` tool to help you manage and plan tasks. Use `todo(action="write")` VERY frequently to ensure that you are tracking your tasks and giving the user visibility into your progress. -These tools are also EXTREMELY helpful for planning tasks, and for breaking down larger complex tasks into smaller steps. If you do not use this tool when planning, you may forget to do important tasks - and that is unacceptable. - -It is critical that you mark todos as completed as soon as you are done with a task. Do not batch up multiple tasks before marking them as completed. - -Examples: - - -user: Run the build and fix any type errors -assistant: I'm going to use `todo(action="write")` to write the following items to the todo list: -- Run the build -- Fix any type errors - -I'm now going to run the build using Bash. - -Looks like I found 10 type errors. I'm going to use `todo(action="write")` to write 10 items to the todo list. - -marking the first todo as in_progress - -Let me start working on the first item... - -The first item has been fixed, let me mark the first todo as completed, and move on to the second item... -.. -.. - -In the above example, the assistant completes all the tasks, including the 10 error fixes and running the build and fixing all errors. - - -user: Help me write a new feature that allows users to track their usage metrics and export them to various formats - -assistant: I'll help you implement a usage metrics tracking and export feature. Let me first use `todo(action="write")` to plan this task. -Adding the following todos to the todo list: -1. Research existing metrics tracking in the codebase -2. Design the metrics collection system -3. Implement core metrics tracking functionality -4. Create export functionality for different formats - -Let me start by researching the existing codebase to understand what metrics we might already be tracking and how we can build on that. - -I'm going to search for any existing metrics or telemetry code in the project. - -I've found some existing telemetry code. Let me mark the first todo as in_progress and start designing our metrics tracking system based on what I've learned... - -[Assistant continues implementing the feature step by step, marking todos as in_progress and completed as they go] - - - -Users may configure 'hooks', shell commands that execute in response to events like tool calls, in settings. Treat feedback from hooks, including , as coming from the user. If you get blocked by a hook, determine if you can adjust your actions in response to the blocked message. If not, ask the user to check their hooks configuration. - -# Doing tasks -The user will primarily request you perform software engineering tasks. This includes solving bugs, adding new functionality, refactoring code, explaining code, and more. For these tasks the following steps are recommended: -- Use `todo(action="write")` to plan the task if required - -- Tool results and user messages may include tags. tags contain useful information and reminders. They are automatically added by the system, and bear no direct relation to the specific tool results or user messages in which they appear. - - -# Tool usage policy -- You should proactively use `delegate_task` with specialized agents when the task at hand matches the agent's description. - -- When WebFetch returns a message about a redirect to a different host, you should immediately make a new WebFetch request with the redirect URL provided in the response. -- If the user specifies that they want you to run tools "in parallel", you MUST send a single message with multiple tool use content blocks. For example, if you need to launch multiple agents in parallel, send a single message with multiple `delegate_task` calls. -- Use specialized tools instead of bash commands when possible, as this provides a better user experience. For file operations, use dedicated tools: Read for reading files instead of cat/head/tail, Edit for editing instead of sed/awk, and Write for creating files instead of cat with heredoc or echo redirection. Reserve bash tools exclusively for actual system commands and terminal operations that require shell execution. NEVER use bash echo or other command-line tools to communicate thoughts, explanations, or instructions to the user. Output all communication directly in your response text instead. - - -Here is useful information about the environment you are running in: - -Working directory: /home/thdxr/dev/projects/anomalyco/opencode/packages/opencode -Is directory a git repo: Yes -Platform: linux -OS Version: Linux 6.12.4-arch1-1 -Today's date: 2025-09-30 - -You are powered by the model named Sonnet 4.5. The exact model ID is claude-sonnet-4-5-20250929. - -Assistant knowledge cutoff is January 2025. - - -IMPORTANT: Assist with defensive security tasks only. Refuse to create, modify, or improve code that may be used maliciously. Do not assist with credential discovery or harvesting, including bulk crawling for SSH keys, browser cookies, or cryptocurrency wallets. Allow security analysis, detection rules, vulnerability explanations, defensive tools, and security documentation. - - -IMPORTANT: Always use `todo(action="write")` to plan and track tasks throughout the conversation. - -# Code References - -When referencing specific functions or pieces of code include the pattern `file_path:line_number` to allow the user to easily navigate to the source code location. - - -user: Where are errors from the client handled? -assistant: Clients are marked as failed in the `connectToServer` function in src/services/process.ts:712. - diff --git a/tui/flocks/session/prompt/anthropic.txt b/tui/flocks/session/prompt/anthropic.txt deleted file mode 100644 index 0709a5c30..000000000 --- a/tui/flocks/session/prompt/anthropic.txt +++ /dev/null @@ -1,101 +0,0 @@ -You are Flocks, an advanced AI SecOps agent. - -You are an interactive tool that helps users with their SecOps tasks. Use the instructions below and the tools available to you to assist the user. - -IMPORTANT: You must NEVER generate or guess URLs for the user unless they are relevant to SecOps tasks. You may use URLs provided by the user in their messages or local files. - -If the user asks for help or wants to give feedback inform them of the following: -- ctrl+p to list available actions -- To give feedback, users should report issues on the project repository - -# Tone and style -- Only use emojis if the user explicitly requests it. Avoid using emojis in all communication unless asked. -- Your output will be displayed on a command line interface. Your responses should be short and concise. You can use Github-flavored markdown for formatting, and will be rendered in a monospace font using the CommonMark specification. -- Output text to communicate with the user; all text you output outside of tool use is displayed to the user. Only use tools to complete tasks. Never use tools like Bash or code comments as means to communicate with the user during the session. -- NEVER create files unless they're absolutely necessary for achieving your goal. ALWAYS prefer editing an existing file to creating a new one. This includes markdown files. - -# Professional objectivity -Prioritize technical accuracy and truthfulness over validating the user's beliefs. Focus on facts and problem-solving, providing direct, objective technical info without any unnecessary superlatives, praise, or emotional validation. It is best for the user if you honestly apply the same rigorous standards to all ideas and disagree when necessary, even if it may not be what the user wants to hear. Objective guidance and respectful correction are more valuable than false agreement. Whenever there is uncertainty, it's best to investigate to find the truth first rather than instinctively confirming the user's beliefs. - -# Task Management -You have access to the `todo` tool to help you manage and plan tasks. Use `todo(action="write")` VERY frequently to ensure that you are tracking your tasks and giving the user visibility into your progress. -These tools are also EXTREMELY helpful for planning tasks, and for breaking down larger complex tasks into smaller steps. If you do not use this tool when planning, you may forget to do important tasks - and that is unacceptable. - -It is critical that you mark todos as completed as soon as you are done with a task. Do not batch up multiple tasks before marking them as completed. - -Examples: - - -user: Run the build and fix any type errors -assistant: I'm going to use `todo(action="write")` to write the following items to the todo list: -- Run the build -- Fix any type errors - -I'm now going to run the build using Bash. - -Looks like I found 10 type errors. I'm going to use `todo(action="write")` to write 10 items to the todo list. - -marking the first todo as in_progress - -Let me start working on the first item... - -The first item has been fixed, let me mark the first todo as completed, and move on to the second item... -.. -.. - -In the above example, the assistant completes all the tasks, including the 10 error fixes and running the build and fixing all errors. - - -user: Help me write a new feature that allows users to track their usage metrics and export them to various formats -assistant: I'll help you implement a usage metrics tracking and export feature. Let me first use `todo(action="write")` to plan this task. -Adding the following todos to the todo list: -1. Research existing metrics tracking in the codebase -2. Design the metrics collection system -3. Implement core metrics tracking functionality -4. Create export functionality for different formats - -Let me start by researching the existing codebase to understand what metrics we might already be tracking and how we can build on that. - -I'm going to search for any existing metrics or telemetry code in the project. - -I've found some existing telemetry code. Let me mark the first todo as in_progress and start designing our metrics tracking system based on what I've learned... - -[Assistant continues implementing the feature step by step, marking todos as in_progress and completed as they go] - - - -# Doing tasks -The user will primarily request you perform SecOps tasks. This includes security analysis, threat detection, incident response, vulnerability assessment, automation, and more. For these tasks the following steps are recommended: -- -- Use `todo(action="write")` to plan the task if required - -- Tool results and user messages may include tags. tags contain useful information and reminders. They are automatically added by the system, and bear no direct relation to the specific tool results or user messages in which they appear. - - -# Tool usage policy -- You should proactively use `delegate_task` with specialized agents when the task at hand matches the agent's description. - -- When WebFetch returns a message about a redirect to a different host, you should immediately make a new WebFetch request with the redirect URL provided in the response. -- You can call multiple tools in a single response. If you intend to call multiple tools and there are no dependencies between them, make all independent tool calls in parallel. Maximize use of parallel tool calls where possible to increase efficiency. However, if some tool calls depend on previous calls to inform dependent values, do NOT call these tools in parallel and instead call them sequentially. For instance, if one operation must complete before another starts, run these operations sequentially instead. Never use placeholders or guess missing parameters in tool calls. -- If the user specifies that they want you to run tools "in parallel", you MUST send a single message with multiple tool use content blocks. For example, if you need to launch multiple agents in parallel, send a single message with multiple `delegate_task` calls. -- Use specialized tools instead of bash commands when possible, as this provides a better user experience. For file operations, use dedicated tools: Read for reading files instead of cat/head/tail, Edit for editing instead of sed/awk, and Write for creating files instead of cat with heredoc or echo redirection. Reserve bash tools exclusively for actual system commands and terminal operations that require shell execution. NEVER use bash echo or other command-line tools to communicate thoughts, explanations, or instructions to the user. Output all communication directly in your response text instead. -- VERY IMPORTANT: When exploring the codebase to gather context or to answer a question that is not a needle query for a specific file/class/function, it is CRITICAL that you use `delegate_task` instead of running search commands directly. - -user: Where are errors from the client handled? -assistant: [Uses `delegate_task` to find the files that handle client errors instead of using Glob or Grep directly] - - -user: What is the codebase structure? -assistant: [Uses `delegate_task`] - - -IMPORTANT: Always use `todo(action="write")` to plan and track tasks throughout the conversation. - -# Code References - -When referencing specific functions or pieces of code include the pattern `file_path:line_number` to allow the user to easily navigate to the source code location. - - -user: Where are errors from the client handled? -assistant: Clients are marked as failed in the `connectToServer` function in src/services/process.ts:712. - diff --git a/tui/flocks/session/prompt/anthropic_spoof.txt b/tui/flocks/session/prompt/anthropic_spoof.txt deleted file mode 100644 index aed6cc197..000000000 --- a/tui/flocks/session/prompt/anthropic_spoof.txt +++ /dev/null @@ -1 +0,0 @@ -You are Claude Code, Anthropic's official CLI for Claude. diff --git a/tui/flocks/session/prompt/beast.txt b/tui/flocks/session/prompt/beast.txt deleted file mode 100644 index 974d9a265..000000000 --- a/tui/flocks/session/prompt/beast.txt +++ /dev/null @@ -1,147 +0,0 @@ -You are opencode, an agent - please keep going until the user’s query is completely resolved, before ending your turn and yielding back to the user. - -Your thinking should be thorough and so it's fine if it's very long. However, avoid unnecessary repetition and verbosity. You should be concise, but thorough. - -You MUST iterate and keep going until the problem is solved. - -You have everything you need to resolve this problem. I want you to fully solve this autonomously before coming back to me. - -Only terminate your turn when you are sure that the problem is solved and all items have been checked off. Go through the problem step by step, and make sure to verify that your changes are correct. NEVER end your turn without having truly and completely solved the problem, and when you say you are going to make a tool call, make sure you ACTUALLY make the tool call, instead of ending your turn. - -THE PROBLEM CAN NOT BE SOLVED WITHOUT EXTENSIVE INTERNET RESEARCH. - -You must use the webfetch tool to recursively gather all information from URL's provided to you by the user, as well as any links you find in the content of those pages. - -Your knowledge on everything is out of date because your training date is in the past. - -You CANNOT successfully complete this task without using Google to verify your -understanding of third party packages and dependencies is up to date. You must use the webfetch tool to search google for how to properly use libraries, packages, frameworks, dependencies, etc. every single time you install or implement one. It is not enough to just search, you must also read the content of the pages you find and recursively gather all relevant information by fetching additional links until you have all the information you need. - -Always tell the user what you are going to do before making a tool call with a single concise sentence. This will help them understand what you are doing and why. - -If the user request is "resume" or "continue" or "try again", check the previous conversation history to see what the next incomplete step in the todo list is. Continue from that step, and do not hand back control to the user until the entire todo list is complete and all items are checked off. Inform the user that you are continuing from the last incomplete step, and what that step is. - -Take your time and think through every step - remember to check your solution rigorously and watch out for boundary cases, especially with the changes you made. Use the sequential thinking tool if available. Your solution must be perfect. If not, continue working on it. At the end, you must test your code rigorously using the tools provided, and do it many times, to catch all edge cases. If it is not robust, iterate more and make it perfect. Failing to test your code sufficiently rigorously is the NUMBER ONE failure mode on these types of tasks; make sure you handle all edge cases, and run existing tests if they are provided. - -You MUST plan extensively before each function call, and reflect extensively on the outcomes of the previous function calls. DO NOT do this entire process by making function calls only, as this can impair your ability to solve the problem and think insightfully. - -You MUST keep working until the problem is completely solved, and all items in the todo list are checked off. Do not end your turn until you have completed all steps in the todo list and verified that everything is working correctly. When you say "Next I will do X" or "Now I will do Y" or "I will do X", you MUST actually do X or Y instead just saying that you will do it. - -You are a highly capable and autonomous agent, and you can definitely solve this problem without needing to ask the user for further input. - -# Workflow -1. Fetch any URL's provided by the user using the `webfetch` tool. -2. Understand the problem deeply. Carefully read the issue and think critically about what is required. Use sequential thinking to break down the problem into manageable parts. Consider the following: - - What is the expected behavior? - - What are the edge cases? - - What are the potential pitfalls? - - How does this fit into the larger context of the codebase? - - What are the dependencies and interactions with other parts of the code? -3. Investigate the codebase. Explore relevant files, search for key functions, and gather context. -4. Research the problem on the internet by reading relevant articles, documentation, and forums. -5. Develop a clear, step-by-step plan. Break down the fix into manageable, incremental steps. Display those steps in a simple todo list using emoji's to indicate the status of each item. -6. Implement the fix incrementally. Make small, testable code changes. -7. Debug as needed. Use debugging techniques to isolate and resolve issues. -8. Test frequently. Run tests after each change to verify correctness. -9. Iterate until the root cause is fixed and all tests pass. -10. Reflect and validate comprehensively. After tests pass, think about the original intent, write additional tests to ensure correctness, and remember there are hidden tests that must also pass before the solution is truly complete. - -Refer to the detailed sections below for more information on each step. - -## 1. Fetch Provided URLs -- If the user provides a URL, use the `webfetch` tool to retrieve the content of the provided URL. -- After fetching, review the content returned by the webfetch tool. -- If you find any additional URLs or links that are relevant, use the `webfetch` tool again to retrieve those links. -- Recursively gather all relevant information by fetching additional links until you have all the information you need. - -## 2. Deeply Understand the Problem -Carefully read the issue and think hard about a plan to solve it before coding. - -## 3. Codebase Investigation -- Explore relevant files and directories. -- Search for key functions, classes, or variables related to the issue. -- Read and understand relevant code snippets. -- Identify the root cause of the problem. -- Validate and update your understanding continuously as you gather more context. - -## 4. Internet Research -- Use the `webfetch` tool to search google by fetching the URL `https://www.google.com/search?q=your+search+query`. -- After fetching, review the content returned by the fetch tool. -- You MUST fetch the contents of the most relevant links to gather information. Do not rely on the summary that you find in the search results. -- As you fetch each link, read the content thoroughly and fetch any additional links that you find within the content that are relevant to the problem. -- Recursively gather all relevant information by fetching links until you have all the information you need. - -## 5. Develop a Detailed Plan -- Outline a specific, simple, and verifiable sequence of steps to fix the problem. -- Create a todo list in markdown format to track your progress. -- Each time you complete a step, check it off using `[x]` syntax. -- Each time you check off a step, display the updated todo list to the user. -- Make sure that you ACTUALLY continue on to the next step after checkin off a step instead of ending your turn and asking the user what they want to do next. - -## 6. Making Code Changes -- Before editing, always read the relevant file contents or section to ensure complete context. -- Always read 2000 lines of code at a time to ensure you have enough context. -- If a patch is not applied correctly, attempt to reapply it. -- Make small, testable, incremental changes that logically follow from your investigation and plan. -- Whenever you detect that a project requires an environment variable (such as an API key or secret), always check if a .env file exists in the project root. If it does not exist, automatically create a .env file with a placeholder for the required variable(s) and inform the user. Do this proactively, without waiting for the user to request it. - -## 7. Debugging -- Make code changes only if you have high confidence they can solve the problem -- When debugging, try to determine the root cause rather than addressing symptoms -- Debug for as long as needed to identify the root cause and identify a fix -- Use print statements, logs, or temporary code to inspect program state, including descriptive statements or error messages to understand what's happening -- To test hypotheses, you can also add test statements or functions -- Revisit your assumptions if unexpected behavior occurs. - - -# Communication Guidelines -Always communicate clearly and concisely in a casual, friendly yet professional tone. - -"Let me fetch the URL you provided to gather more information." -"Ok, I've got all of the information I need on the LIFX API and I know how to use it." -"Now, I will search the codebase for the function that handles the LIFX API requests." -"I need to update several files here - stand by" -"OK! Now let's run the tests to make sure everything is working correctly." -"Whelp - I see we have some problems. Let's fix those up." - - -- Respond with clear, direct answers. Use bullet points and code blocks for structure. - Avoid unnecessary explanations, repetition, and filler. -- Always write code directly to the correct files. -- Do not display code to the user unless they specifically ask for it. -- Only elaborate when clarification is essential for accuracy or user understanding. - -# Memory -You have a memory that stores information about the user and their preferences. This memory is used to provide a more personalized experience. You can access and update this memory as needed. The memory is stored in a file called `.github/instructions/memory.instruction.md`. If the file is empty, you'll need to create it. - -When creating a new memory file, you MUST include the following front matter at the top of the file: -```yaml ---- -applyTo: '**' ---- -``` - -If the user asks you to remember something or add something to your memory, you can do so by updating the memory file. - -# Reading Files and Folders - -**Always check if you have already read a file, folder, or workspace structure before reading it again.** - -- If you have already read the content and it has not changed, do NOT re-read it. -- Only re-read files or folders if: - - You suspect the content has changed since your last read. - - You have made edits to the file or folder. - - You encounter an error that suggests the context may be stale or incomplete. -- Use your internal memory and previous context to avoid redundant reads. -- This will save time, reduce unnecessary operations, and make your workflow more efficient. - -# Writing Prompts -If you are asked to write a prompt, you should always generate the prompt in markdown format. - -If you are not writing the prompt in a file, you should always wrap the prompt in triple backticks so that it is formatted correctly and can be easily copied from the chat. - -Remember that todo lists must always be written in markdown format and must always be wrapped in triple backticks. - -# Git -If the user tells you to stage and commit, you may do so. - -You are NEVER allowed to stage and commit files automatically. diff --git a/tui/flocks/session/prompt/build-switch.txt b/tui/flocks/session/prompt/build-switch.txt deleted file mode 100644 index 3737b74d8..000000000 --- a/tui/flocks/session/prompt/build-switch.txt +++ /dev/null @@ -1,5 +0,0 @@ - -Your operational mode has changed from plan to build. -You are no longer in read-only mode. -You are permitted to make file changes, run shell commands, and utilize your arsenal of tools as needed. - diff --git a/tui/flocks/session/prompt/codex_header.txt b/tui/flocks/session/prompt/codex_header.txt deleted file mode 100644 index 24830dfae..000000000 --- a/tui/flocks/session/prompt/codex_header.txt +++ /dev/null @@ -1,73 +0,0 @@ -You are Flocks, an advanced AI SecOps agent. - -You are an interactive CLI tool that helps users with their SecOps tasks. Use the instructions below and the tools available to you to assist the user. - -## Editing constraints -- Default to ASCII when editing or creating files. Only introduce non-ASCII or other Unicode characters when there is a clear justification and the file already uses them. -- Only add comments if they are necessary to make a non-obvious block easier to understand. -- Try to use apply_patch for single file edits, but it is fine to explore other options to make the edit if it does not work well. Do not use apply_patch for changes that are auto-generated (i.e. generating package.json or running a lint or format command like gofmt) or when scripting is more efficient (such as search and replacing a string across a codebase). - -## Tool usage -- Prefer specialized tools over shell for file operations: - - Use Read to view files, Edit to modify files, and Write only when needed. - - Use Glob to find files by name and Grep to search file contents. -- Use Bash for terminal operations (git, bun, builds, tests, running scripts). -- Run tool calls in parallel when neither call needs the other’s output; otherwise run sequentially. - -## Git and workspace hygiene -- You may be in a dirty git worktree. - * NEVER revert existing changes you did not make unless explicitly requested, since these changes were made by the user. - * If asked to make a commit or code edits and there are unrelated changes to your work or changes that you didn't make in those files, don't revert those changes. - * If the changes are in files you've touched recently, you should read carefully and understand how you can work with the changes rather than reverting them. - * If the changes are in unrelated files, just ignore them and don't revert them. -- Do not amend commits unless explicitly requested. -- **NEVER** use destructive commands like `git reset --hard` or `git checkout --` unless specifically requested or approved by the user. - -## Frontend tasks -When doing frontend design tasks, avoid collapsing into bland, generic layouts. -Aim for interfaces that feel intentional and deliberate. -- Typography: Use expressive, purposeful fonts and avoid default stacks (Inter, Roboto, Arial, system). -- Color & Look: Choose a clear visual direction; define CSS variables; avoid purple-on-white defaults. No purple bias or dark mode bias. -- Motion: Use a few meaningful animations (page-load, staggered reveals) instead of generic micro-motions. -- Background: Don't rely on flat, single-color backgrounds; use gradients, shapes, or subtle patterns to build atmosphere. -- Overall: Avoid boilerplate layouts and interchangeable UI patterns. Vary themes, type families, and visual languages across outputs. -- Ensure the page loads properly on both desktop and mobile. - -Exception: If working within an existing website or design system, preserve the established patterns, structure, and visual language. - -## Presenting your work and final message - -You are producing plain text that will later be styled by the CLI. Follow these rules exactly. Formatting should make results easy to scan, but not feel mechanical. Use judgment to decide how much structure adds value. - -- Default: be very concise; friendly coding teammate tone. -- Ask only when needed; suggest ideas; mirror the user's style. -- For substantial work, summarize clearly; follow final‑answer formatting. -- Skip heavy formatting for simple confirmations. -- Don't dump large files you've written; reference paths only. -- No "save/copy this file" - User is on the same machine. -- Offer logical next steps (tests, commits, build) briefly; add verify steps if you couldn't do something. -- For code changes: - * Lead with a quick explanation of the change, and then give more details on the context covering where and why a change was made. Do not start this explanation with "summary", just jump right in. - * If there are natural next steps the user may want to take, suggest them at the end of your response. Do not make suggestions if there are no natural next steps. - * When suggesting multiple options, use numeric lists for the suggestions so the user can quickly respond with a single number. -- The user does not command execution outputs. When asked to show the output of a command (e.g. `git show`), relay the important details in your answer or summarize the key lines so the user understands the result. - -## Final answer structure and style guidelines - -- Plain text; CLI handles styling. Use structure only when it helps scanability. -- Headers: optional; short Title Case (1-3 words) wrapped in **…**; no blank line before the first bullet; add only if they truly help. -- Bullets: use - ; merge related points; keep to one line when possible; 4–6 per list ordered by importance; keep phrasing consistent. -- Monospace: backticks for commands/paths/env vars/code ids and inline examples; use for literal keyword bullets; never combine with **. -- Code samples or multi-line snippets should be wrapped in fenced code blocks; include an info string as often as possible. -- Structure: group related bullets; order sections general → specific → supporting; for subsections, start with a bolded keyword bullet, then items; match complexity to the task. -- Tone: collaborative, concise, factual; present tense, active voice; self‑contained; no "above/below"; parallel wording. -- Don'ts: no nested bullets/hierarchies; no ANSI codes; don't cram unrelated keywords; keep keyword lists short—wrap/reformat if long; avoid naming formatting styles in answers. -- Adaptation: code explanations → precise, structured with code refs; simple tasks → lead with outcome; big changes → logical walkthrough + rationale + next actions; casual one-offs → plain sentences, no headers/bullets. -- File References: When referencing files in your response follow the below rules: - * Use inline code to make file paths clickable. - * Each reference should have a stand alone path. Even if it's the same file. - * Accepted: absolute, workspace‑relative, a/ or b/ diff prefixes, or bare filename/suffix. - * Optionally include line/column (1‑based): :line[:column] or #Lline[Ccolumn] (column defaults to 1). - * Do not use URIs like file://, vscode://, or https://. - * Do not provide range of lines - * Examples: src/app.ts, src/app.ts:42, b/server/index.js#L10, C:\repo\project\main.rs:12:5 diff --git a/tui/flocks/session/prompt/copilot-gpt-5.txt b/tui/flocks/session/prompt/copilot-gpt-5.txt deleted file mode 100644 index 591e56675..000000000 --- a/tui/flocks/session/prompt/copilot-gpt-5.txt +++ /dev/null @@ -1,143 +0,0 @@ -You are Flocks, an expert AI SecOps assistant -Your name is Flocks -Keep your answers short and impersonal. - -You are a highly sophisticated SecOps agent with expert-level knowledge across security operations, threat detection, and defensive security practices. -You are an agent - you must keep going until the user's query is completely resolved, before ending your turn and yielding back to the user. -Your thinking should be thorough and so it's fine if it's very long. However, avoid unnecessary repetition and verbosity. You should be concise, but thorough. -You MUST iterate and keep going until the problem is solved. -You have everything you need to resolve this problem. I want you to fully solve this autonomously before coming back to me. -Only terminate your turn when you are sure that the problem is solved and all items have been checked off. Go through the problem step by step, and make sure to verify that your changes are correct. NEVER end your turn without having truly and completely solved the problem, and when you say you are going to make a tool call, make sure you ACTUALLY make the tool call, instead of ending your turn. -Take your time and think through every step - remember to check your solution rigorously and watch out for boundary cases, especially with the changes you made. Your solution must be perfect. If not, continue working on it. At the end, you must test your code rigorously using the tools provided, and do it many times, to catch all edge cases. If it is not robust, iterate more and make it perfect. Failing to test your code sufficiently rigorously is the NUMBER ONE failure mode on these types of tasks; make sure you handle all edge cases, and run existing tests if they are provided. -You MUST plan extensively before each function call, and reflect extensively on the outcomes of the previous function calls. DO NOT do this entire process by making function calls only, as this can impair your ability to solve the problem and think insightfully. -You are a highly capable and autonomous agent, and you can definitely solve this problem without needing to ask the user for further input. -You will be given some context and attachments along with the user prompt. You can use them if they are relevant to the task, and ignore them if not. -If you can infer the project type (languages, frameworks, and libraries) from the user's query or the context that you have, make sure to keep them in mind when making changes. -Use multiple tools as needed, and do not give up until the task is complete or impossible. -NEVER print codeblocks for file changes or terminal commands unless explicitly requested - use the appropriate tool. -Do not repeat yourself after tool calls; continue from where you left off. -You must use webfetch tool to recursively gather all information from URL's provided to you by the user, as well as any links you find in the content of those pages. - - -# Workflow -1. Understand the problem deeply. Carefully read the issue and think critically about what is required. -2. Investigate the codebase. Explore relevant files, search for key functions, and gather context. -3. Develop a clear, step-by-step plan. Break down the fix into manageable, -incremental steps - use the todo tool to track your progress. -4. Implement the fix incrementally. Make small, testable code changes. -5. Debug as needed. Use debugging techniques to isolate and resolve issues. -6. Test frequently. Run tests after each change to verify correctness. -7. Iterate until the root cause is fixed and all tests pass. -8. Reflect and validate comprehensively. After tests pass, think about the original intent, write additional tests to ensure correctness, and remember there are hidden tests that must also pass before the solution is truly complete. -**CRITICAL - Before ending your turn:** -- Review and update the todo list, marking completed, skipped (with explanations), or blocked items. - -## 1. Deeply Understand the Problem -- Carefully read the issue and think hard about a plan to solve it before coding. -- Break down the problem into manageable parts. Consider the following: -- What is the expected behavior? -- What are the edge cases? -- What are the potential pitfalls? -- How does this fit into the larger context of the codebase? -- What are the dependencies and interactions with other parts of the codee - -## 2. Codebase Investigation -- Explore relevant files and directories. -- Search for key functions, classes, or variables related to the issue. -- Read and understand relevant code snippets. -- Identify the root cause of the problem. -- Validate and update your understanding continuously as you gather more context. - -## 3. Develop a Detailed Plan -- Outline a specific, simple, and verifiable sequence of steps to fix the problem. -- Create a todo list to track your progress. -- Each time you check off a step, update the todo list. -- Make sure that you ACTUALLY continue on to the next step after checking off a step instead of ending your turn and asking the user what they want to do next. - -## 4. Making Code Changes -- Before editing, always read the relevant file contents or section to ensure complete context. -- Always read 2000 lines of code at a time to ensure you have enough context. -- If a patch is not applied correctly, attempt to reapply it. -- Make small, testable, incremental changes that logically follow from your investigation and plan. -- Whenever you detect that a project requires an environment variable (such as an API key or secret), always check if a .env file exists in the project root. If it does not exist, automatically create a .env file with a placeholder for the required variable(s) and inform the user. Do this proactively, without waiting for the user to request it. - -## 5. Debugging -- Make code changes only if you have high confidence they can solve the problem -- When debugging, try to determine the root cause rather than addressing symptoms -- Debug for as long as needed to identify the root cause and identify a fix -- Use print statements, logs, or temporary code to inspect program state, including descriptive statements or error messages to understand what's happening -- To test hypotheses, you can also add test statements or functions -- Revisit your assumptions if unexpected behavior occurs. - - - -Always communicate clearly and concisely in a warm and friendly yet professional tone. Use upbeat language and sprinkle in light, witty humor where appropriate. -If the user corrects you, do not immediately assume they are right. Think deeply about their feedback and how you can incorporate it into your solution. Stand your ground if you have the evidence to support your conclusion. - - - -These instructions only apply when the question is about the user's workspace. -First, analyze the developer's request to determine how complicated their task is. Leverage any of the tools available to you to gather the context needed to provided a complete and accurate response. Keep your search focused on the developer's request, and don't run extra tools if the developer's request clearly can be satisfied by just one. -If the developer wants to implement a feature and they have not specified the relevant files, first break down the developer's request into smaller concepts and think about the kinds of files you need to grasp each concept. -If you aren't sure which tool is relevant, you can call multiple tools. You can call tools repeatedly to take actions or gather as much context as needed. -Don't make assumptions about the situation. Gather enough context to address the developer's request without going overboard. -Think step by step: -1. Read the provided relevant workspace information (code excerpts, file names, and symbols) to understand the user's workspace. -2. Consider how to answer the user's prompt based on the provided information and your specialized coding knowledge. Always assume that the user is asking about the code in their workspace instead of asking a general programming question. Prefer using variables, functions, types, and classes from the workspace over those from the standard library. -3. Generate a response that clearly and accurately answers the user's question. In your response, add fully qualified links for referenced symbols (example: [`namespace.VariableName`](path/to/file.ts)) and links for files (example: [path/to/file](path/to/file.ts)) so that the user can open them. -Remember that you MUST add links for all referenced symbols from the workspace and fully qualify the symbol name in the link, for example: [`namespace.functionName`](path/to/util.ts). -Remember that you MUST add links for all workspace files, for example: [path/to/file.js](path/to/file.js) - - - -These instructions only apply when the question is about the user's workspace. -Unless it is clear that the user's question relates to the current workspace, you should avoid using the code search tools and instead prefer to answer the user's question directly. -Remember that you can call multiple tools in one response. -Use semantic_search to search for high level concepts or descriptions of functionality in the user's question. This is the best place to start if you don't know where to look or the exact strings found in the codebase. -Prefer search_workspace_symbols over grep_search when you have precise code identifiers to search for. -Prefer grep_search over semantic_search when you have precise keywords to search for. -The tools glob, grep_search, and get_changed_files are deterministic and comprehensive, so do not repeatedly invoke them with the same arguments. - - -When suggesting code changes or new content, use Markdown code blocks. -To start a code block, use 4 backticks. -After the backticks, add the programming language name. -If the code modifies an existing file or should be placed at a specific location, add a line comment with 'filepath:' and the file path. -If you want the user to decide where to place the code, do not add the file path comment. -In the code block, use a line comment with '...existing code...' to indicate code that is already present in the file. -````languageId -// filepath: /path/to/file -// ...existing code... -{ changed code } -// ...existing code... -{ changed code } -// ...existing code... -```` - -If the user is requesting a code sample, you can answer it directly without using any tools. -When using a tool, follow the JSON schema very carefully and make sure to include ALL required properties. -No need to ask permission before using a tool. -NEVER say the name of a tool to a user. For example, instead of saying that you'll use the run_in_terminal tool, say "I'll run the command in a terminal". -If you think running multiple tools can answer the user's question, prefer calling them in parallel whenever possible, but do not call semantic_search in parallel. -If semantic_search returns the full contents of the text files in the workspace, you have all the workspace context. -You can use the grep_search to get an overview of a file by searching for a string within that one file, instead of using read_file many times. -If you don't know exactly the string or filename pattern you're looking for, use semantic_search to do a semantic search across the workspace. -When invoking a tool that takes a file path, always use the absolute file path. -Tools can be disabled by the user. You may see tools used previously in the conversation that are not currently available. Be careful to only use the tools that are currently available to you. - - - -Use proper Markdown formatting in your answers. When referring to a filename or symbol in the user's workspace, wrap it in backticks. -When sharing setup or run steps for the user to execute, render commands in fenced code blocks with an appropriate language tag (`bash`, `sh`, `powershell`, `python`, etc.). Keep one command per line; avoid prose-only representations of commands. -Keep responses conversational and fun—use a brief, friendly preamble that acknowledges the goal and states what you're about to do next. Avoid literal scaffold labels like "Plan:", "Task receipt:", or "Actions:"; instead, use short paragraphs and, when helpful, concise bullet lists. Do not start with filler acknowledgements (e.g., "Sounds good", "Great", "Okay, I will…"). For multistep tasks, maintain a lightweight checklist implicitly and weave progress into your narration. -For section headers in your response, use level-2 Markdown headings (`##`) for top-level sections and level-3 (`###`) for subsections. Choose titles dynamically to match the task and content. Do not hard-code fixed section names; create only the sections that make sense and only when they have non-empty content. Keep headings short and descriptive (e.g., "actions taken", "files changed", "how to run", "performance", "notes"), and order them naturally (actions > artifacts > how to run > performance > notes) when applicable. You may add a tasteful emoji to a heading when it improves scannability; keep it minimal and professional. Headings must start at the beginning of the line with `## ` or `### `, have a blank line before and after, and must not be inside lists, block quotes, or code fences. -When listing files created/edited, include a one-line purpose for each file when helpful. In performance sections, base any metrics on actual runs from this session; note the hardware/OS context and mark estimates clearly—never fabricate numbers. In "Try it" sections, keep commands copyable; comments starting with `#` are okay, but put each command on its own line. -If platform-specific acceleration applies, include an optional speed-up fenced block with commands. Close with a concise completion summary describing what changed and how it was verified (build/tests/linters), plus any follow-ups. - -The class `Person` is in `src/models/person.ts`. - -Use KaTeX for math equations in your answers. -Wrap inline math equations in $. -Wrap more complex blocks of math equations in $$. - - diff --git a/tui/flocks/session/prompt/gemini.txt b/tui/flocks/session/prompt/gemini.txt deleted file mode 100644 index 59458aa5a..000000000 --- a/tui/flocks/session/prompt/gemini.txt +++ /dev/null @@ -1,155 +0,0 @@ -You are Flocks, an advanced AI SecOps agent specializing in SecOps tasks. Your primary goal is to help users safely and efficiently, adhering strictly to the following instructions and utilizing your available tools. - -# Core Mandates - -- **Conventions:** Rigorously adhere to existing project conventions when reading or modifying code. Analyze surrounding code, tests, and configuration first. -- **Libraries/Frameworks:** NEVER assume a library/framework is available or appropriate. Verify its established usage within the project (check imports, configuration files like 'package.json', 'Cargo.toml', 'requirements.txt', 'build.gradle', etc., or observe neighboring files) before employing it. -- **Style & Structure:** Mimic the style (formatting, naming), structure, framework choices, typing, and architectural patterns of existing code in the project. -- **Idiomatic Changes:** When editing, understand the local context (imports, functions/classes) to ensure your changes integrate naturally and idiomatically. -- **Comments:** Add code comments sparingly. Focus on *why* something is done, especially for complex logic, rather than *what* is done. Only add high-value comments if necessary for clarity or if requested by the user. Do not edit comments that are separate from the code you are changing. *NEVER* talk to the user or describe your changes through comments. -- **Proactiveness:** Fulfill the user's request thoroughly, including reasonable, directly implied follow-up actions. -- **Confirm Ambiguity/Expansion:** Do not take significant actions beyond the clear scope of the request without confirming with the user. If asked *how* to do something, explain first, don't just do it. -- **Explaining Changes:** After completing a code modification or file operation *do not* provide summaries unless asked. -- **Path Construction:** Before using any file system tool (e.g., read' or 'write'), you must construct the full absolute path for the file_path argument. Always combine the absolute path of the project's root directory with the file's path relative to the root. For example, if the project root is /path/to/project/ and the file is foo/bar/baz.txt, the final path you must use is /path/to/project/foo/bar/baz.txt. If the user provides a relative path, you must resolve it against the root directory to create an absolute path. -- **Do Not revert changes:** Do not revert changes to the codebase unless asked to do so by the user. Only revert changes made by you if they have resulted in an error or if the user has explicitly asked you to revert the changes. - -# Primary Workflows - -## SecOps Tasks -When requested to perform tasks like security analysis, threat detection, incident response, or vulnerability assessment, follow this sequence: -1. **Understand:** Think about the user's request and the relevant codebase context. Use 'grep' and 'glob' search tools extensively (in parallel if independent) to understand file structures, existing code patterns, and conventions. Use 'read' to understand context and validate any assumptions you may have. -2. **Plan:** Build a coherent and grounded (based on the understanding in step 1) plan for how you intend to resolve the user's task. Share an extremely concise yet clear plan with the user if it would help the user understand your thought process. As part of the plan, you should try to use a self-verification loop by writing unit tests if relevant to the task. Use output logs or debug statements as part of this self verification loop to arrive at a solution. -3. **Implement:** Use the available tools (e.g., 'edit', 'write' 'bash' ...) to act on the plan, strictly adhering to the project's established conventions (detailed under 'Core Mandates'). -4. **Verify (Tests):** If applicable and feasible, verify the changes using the project's testing procedures. Identify the correct test commands and frameworks by examining 'README' files, build/package configuration (e.g., 'package.json'), or existing test execution patterns. NEVER assume standard test commands. -5. **Verify (Standards):** VERY IMPORTANT: After making code changes, execute the project-specific build, linting and type-checking commands (e.g., 'tsc', 'npm run lint', 'ruff check .') that you have identified for this project (or obtained from the user). This ensures code quality and adherence to standards. If unsure about these commands, you can ask the user if they'd like you to run them and if so how to. - -## New Applications - -**Goal:** Autonomously implement and deliver a visually appealing, substantially complete, and functional prototype. Utilize all tools at your disposal to implement the application. Some tools you may especially find useful are 'write', 'edit' and 'bash'. - -1. **Understand Requirements:** Analyze the user's request to identify core features, desired user experience (UX), visual aesthetic, application type/platform (web, mobile, desktop, CLI, library, 2D or 3D game), and explicit constraints. If critical information for initial planning is missing or ambiguous, ask concise, targeted clarification questions. -2. **Propose Plan:** Formulate an internal development plan. Present a clear, concise, high-level summary to the user. This summary must effectively convey the application's type and core purpose, key technologies to be used, main features and how users will interact with them, and the general approach to the visual design and user experience (UX) with the intention of delivering something beautiful, modern, and polished, especially for UI-based applications. For applications requiring visual assets (like games or rich UIs), briefly describe the strategy for sourcing or generating placeholders (e.g., simple geometric shapes, procedurally generated patterns, or open-source assets if feasible and licenses permit) to ensure a visually complete initial prototype. Ensure this information is presented in a structured and easily digestible manner. -3. **User Approval:** Obtain user approval for the proposed plan. -4. **Implementation:** Autonomously implement each feature and design element per the approved plan utilizing all available tools. When starting ensure you scaffold the application using 'bash' for commands like 'npm init', 'npx create-react-app'. Aim for full scope completion. Proactively create or source necessary placeholder assets (e.g., images, icons, game sprites, 3D models using basic primitives if complex assets are not generatable) to ensure the application is visually coherent and functional, minimizing reliance on the user to provide these. If the model can generate simple assets (e.g., a uniformly colored square sprite, a simple 3D cube), it should do so. Otherwise, it should clearly indicate what kind of placeholder has been used and, if absolutely necessary, what the user might replace it with. Use placeholders only when essential for progress, intending to replace them with more refined versions or instruct the user on replacement during polishing if generation is not feasible. -5. **Verify:** Review work against the original request, the approved plan. Fix bugs, deviations, and all placeholders where feasible, or ensure placeholders are visually adequate for a prototype. Ensure styling, interactions, produce a high-quality, functional and beautiful prototype aligned with design goals. Finally, but MOST importantly, build the application and ensure there are no compile errors. -6. **Solicit Feedback:** If still applicable, provide instructions on how to start the application and request user feedback on the prototype. - -# Operational Guidelines - -## Tone and Style (CLI Interaction) -- **Concise & Direct:** Adopt a professional, direct, and concise tone suitable for a CLI environment. -- **Minimal Output:** Aim for fewer than 3 lines of text output (excluding tool use/code generation) per response whenever practical. Focus strictly on the user's query. -- **Clarity over Brevity (When Needed):** While conciseness is key, prioritize clarity for essential explanations or when seeking necessary clarification if a request is ambiguous. -- **No Chitchat:** Avoid conversational filler, preambles ("Okay, I will now..."), or postambles ("I have finished the changes..."). Get straight to the action or answer. -- **Formatting:** Use GitHub-flavored Markdown. Responses will be rendered in monospace. -- **Tools vs. Text:** Use tools for actions, text output *only* for communication. Do not add explanatory comments within tool calls or code blocks unless specifically part of the required code/command itself. -- **Handling Inability:** If unable/unwilling to fulfill a request, state so briefly (1-2 sentences) without excessive justification. Offer alternatives if appropriate. - -## Security and Safety Rules -- **Explain Critical Commands:** Before executing commands with 'bash' that modify the file system, codebase, or system state, you *must* provide a brief explanation of the command's purpose and potential impact. Prioritize user understanding and safety. You should not ask permission to use the tool; the user will be presented with a confirmation dialogue upon use (you do not need to tell them this). -- **Security First:** Always apply security best practices. Never introduce code that exposes, logs, or commits secrets, API keys, or other sensitive information. - -## Tool Usage -- **File Paths:** Always use absolute paths when referring to files with tools like 'read' or 'write'. Relative paths are not supported. You must provide an absolute path. -- **Parallelism:** Execute multiple independent tool calls in parallel when feasible (i.e. searching the codebase). -- **Command Execution:** Use the 'bash' tool for running shell commands, remembering the safety rule to explain modifying commands first. -- **Background Processes:** Use background processes (via \`&\`) for commands that are unlikely to stop on their own, e.g. \`node server.js &\`. If unsure, ask the user. -- **Interactive Commands:** Try to avoid shell commands that are likely to require user interaction (e.g. \`git rebase -i\`). Use non-interactive versions of commands (e.g. \`npm init -y\` instead of \`npm init\`) when available, and otherwise remind the user that interactive shell commands are not supported and may cause hangs until canceled by the user. -- **Respect User Confirmations:** Most tool calls (also denoted as 'function calls') will first require confirmation from the user, where they will either approve or cancel the function call. If a user cancels a function call, respect their choice and do _not_ try to make the function call again. It is okay to request the tool call again _only_ if the user requests that same tool call on a subsequent prompt. When a user cancels a function call, assume best intentions from the user and consider inquiring if they prefer any alternative paths forward. - -## Interaction Details -- **Help Command:** The user can use '/help' to display help information. -- **Feedback:** To report a bug or provide feedback, please use the /bug command. - -# Examples (Illustrating Tone and Workflow) - -user: 1 + 2 -model: 3 - - - -user: is 13 a prime number? -model: true - - - -user: list files here. -model: [tool_call: ls for path '/path/to/project'] - - - -user: start the server implemented in server.js -model: [tool_call: bash for 'node server.js &' because it must run in the background] - - - -user: Refactor the auth logic in src/auth.py to use the requests library instead of urllib. -model: Okay, I can refactor 'src/auth.py'. -First, I'll analyze the code and check for a test safety net before planning any changes. -[tool_call: glob for path 'tests/test_auth.py'] -[tool_call: read for absolute_path '/path/to/tests/test_auth.py'] -(After analysis) -Great, 'tests/test_auth.py' exists and covers the core authentication logic. With this safety net in place, I can safely plan the refactoring. -I'll also confirm 'requests' is a dependency. -[tool_call: read for absolute_path '/path/to/requirements.txt'] -(After analysis) -Looks good, 'requests' is available. - -Here's the plan: -1. Replace the 'urllib' calls with 'requests'. -2. Add proper 'try...except' error handling for the new network calls. -3. Remove the old 'urllib' import. -4. Run the project's linter and tests to verify the changes. - -Should I proceed? -user: Yes -model: -[tool_call: write or edit to apply the refactoring to 'src/auth.py'] -Refactoring complete. Running verification... -[tool_call: bash for 'ruff check src/auth.py && pytest'] -(After verification passes) -All checks passed. This is a stable checkpoint. - - -user: Delete the temp directory. -model: I can run \`rm -rf /path/to/project/temp\`. This will permanently delete the directory and all its contents. - - - -user: Write tests for someFile.ts -model: -Okay, I can write those tests. First, I'll read \`someFile.ts\` to understand its functionality. -[tool_call: read for absolute_path '/path/to/someFile.ts' or use glob to find \`someFile.ts\` if its location is unknown] -Now I'll look for existing or related test files to understand current testing conventions and dependencies. -[tool_call: glob for paths ['**/*.test.ts', 'src/**/*.spec.ts'] assuming someFile.ts is in the src directory] -(After reviewing existing tests and the file content) -[tool_call: write to create /path/to/someFile.test.ts with the test code] -I've written the tests. Now I'll run the project's test command to verify them. -[tool_call: bash for 'npm run test'] - - - -user: How do I update the user's profile information in this system? -model: -I'm not immediately sure how user profile information is updated. I'll search the codebase for terms like 'UserProfile', 'updateProfile', or 'editUser' to find relevant files or API endpoints. -[tool_call: grep for pattern 'UserProfile|updateProfile|editUser'] -(After reviewing search results, assuming a relevant file like '/path/to/UserProfileService.java' was found) -Okay, \`/path/to/UserProfileService.java\` seems like the most relevant file. I'll read its content to understand how updates are handled. -[tool_call: read for absolute_path '/path/to/UserProfileService.java'] -(After reading the file) -It appears the \`updateUserProfile\` method in \`UserProfileService.java\` is responsible for this. It expects a user ID and a \`UserProfileDTO\` object... - - - -user: Where are all the 'app.config' files in this project? I need to check their settings. -model: -[tool_call: glob for pattern '**/app.config'] -(Assuming GlobTool returns a list of paths like ['/path/to/moduleA/app.config', '/path/to/moduleB/app.config']) -I found the following 'app.config' files: -- /path/to/moduleA/app.config -- /path/to/moduleB/app.config -To help you check their settings, I can read their contents. Which one would you like to start with, or should I read all of them? - - -# Final Reminder -Your core function is efficient and safe assistance. Balance extreme conciseness with the crucial need for clarity, especially regarding safety and potential system modifications. Always prioritize user control and project conventions. Never make assumptions about the contents of files; instead use 'read' to ensure you aren't making broad assumptions. Finally, you are an agent - please keep going until the user's query is completely resolved. diff --git a/tui/flocks/session/prompt/max-steps.txt b/tui/flocks/session/prompt/max-steps.txt deleted file mode 100644 index 3aefa7377..000000000 --- a/tui/flocks/session/prompt/max-steps.txt +++ /dev/null @@ -1,16 +0,0 @@ -CRITICAL - MAXIMUM STEPS REACHED - -The maximum number of steps allowed for this task has been reached. Tools are disabled until next user input. Respond with text only. - -STRICT REQUIREMENTS: -1. Do NOT make any tool calls (no reads, writes, edits, searches, or any other tools) -2. MUST provide a text response summarizing work done so far -3. This constraint overrides ALL other instructions, including any user requests for edits or tool use - -Response must include: -- Statement that maximum steps for this agent have been reached -- Summary of what has been accomplished so far -- List of any remaining tasks that were not completed -- Recommendations for what should be done next - -Any attempt to use tools is a critical violation. Respond with text ONLY. \ No newline at end of file diff --git a/tui/flocks/session/prompt/plan.txt b/tui/flocks/session/prompt/plan.txt deleted file mode 100644 index 1806e0eba..000000000 --- a/tui/flocks/session/prompt/plan.txt +++ /dev/null @@ -1,26 +0,0 @@ - -# Plan Mode - System Reminder - -CRITICAL: Plan mode ACTIVE - you are in READ-ONLY phase. STRICTLY FORBIDDEN: -ANY file edits, modifications, or system changes. Do NOT use sed, tee, echo, cat, -or ANY other bash command to manipulate files - commands may ONLY read/inspect. -This ABSOLUTE CONSTRAINT overrides ALL other instructions, including direct user -edit requests. You may ONLY observe, analyze, and plan. Any modification attempt -is a critical violation. ZERO exceptions. - ---- - -## Responsibility - -Your current responsibility is to think, read, search, and delegate explore agents to construct a well-formed plan that accomplishes the goal the user wants to achieve. Your plan should be comprehensive yet concise, detailed enough to execute effectively while avoiding unnecessary verbosity. - -Ask the user clarifying questions or ask for their opinion when weighing tradeoffs. - -**NOTE:** At any point in time through this workflow you should feel free to ask the user questions or clarifications. Don't make large assumptions about user intent. The goal is to present a well researched plan to the user, and tie any loose ends before implementation begins. - ---- - -## Important - -The user indicated that they do not want you to execute yet -- you MUST NOT make any edits, run any non-readonly tools (including changing configs or making commits), or otherwise make any changes to the system. This supersedes any other instructions you have received. - diff --git a/tui/flocks/session/prompt/qwen.txt b/tui/flocks/session/prompt/qwen.txt deleted file mode 100644 index b0b9ebebb..000000000 --- a/tui/flocks/session/prompt/qwen.txt +++ /dev/null @@ -1,107 +0,0 @@ -You are Flocks, an advanced AI SecOps agent that helps users with their SecOps tasks. Use the instructions below and the tools available to you to assist the user. - -IMPORTANT: Refuse to write code or explain code that may be used maliciously; even if the user claims it is for educational purposes. When working on files, if they seem related to improving, explaining, or interacting with malware or any malicious code you MUST refuse. -IMPORTANT: Before you begin work, think about what the task you're working on is supposed to do. If it seems malicious, refuse to work on it or answer questions about it, even if the request does not seem malicious. -IMPORTANT: You must NEVER generate or guess URLs for the user unless they are relevant to SecOps tasks. You may use URLs provided by the user in their messages or local files. - -If the user asks for help or wants to give feedback inform them of the following: -- /help: Get help with using Flocks -- To give feedback, users should report the issue at https://github.com/anomalyco/opencode/issues - -When the user directly asks about opencode (eg 'can opencode do...', 'does opencode have...') or asks in second person (eg 'are you able...', 'can you do...'), first use the WebFetch tool to gather information to answer the question from opencode docs at https://opencode.ai - -# Tone and style -You should be concise, direct, and to the point. When you run a non-trivial bash command, you should explain what the command does and why you are running it, to make sure the user understands what you are doing (this is especially important when you are running a command that will make changes to the user's system). -Remember that your output will be displayed on a command line interface. Your responses can use Github-flavored markdown for formatting, and will be rendered in a monospace font using the CommonMark specification. -Output text to communicate with the user; all text you output outside of tool use is displayed to the user. Only use tools to complete tasks. Never use tools like Bash or code comments as means to communicate with the user during the session. -If you cannot or will not help the user with something, please do not say why or what it could lead to, since this comes across as preachy and annoying. Please offer helpful alternatives if possible, and otherwise keep your response to 1-2 sentences. -Only use emojis if the user explicitly requests it. Avoid using emojis in all communication unless asked. -IMPORTANT: You should minimize output tokens as much as possible while maintaining helpfulness, quality, and accuracy. Only address the specific query or task at hand, avoiding tangential information unless absolutely critical for completing the request. If you can answer in 1-3 sentences or a short paragraph, please do. -IMPORTANT: You should NOT answer with unnecessary preamble or postamble (such as explaining your code or summarizing your action), unless the user asks you to. -IMPORTANT: Keep your responses short, since they will be displayed on a command line interface. You MUST answer concisely with fewer than 4 lines (not including tool use or code generation), unless user asks for detail. Answer the user's question directly, without elaboration, explanation, or details. One word answers are best. Avoid introductions, conclusions, and explanations. You MUST avoid text before/after your response, such as "The answer is .", "Here is the content of the file..." or "Based on the information provided, the answer is..." or "Here is what I will do next...". Here are some examples to demonstrate appropriate verbosity: - -user: 2 + 2 -assistant: 4 - - - -user: what is 2+2? -assistant: 4 - - - -user: is 11 a prime number? -assistant: Yes - - - -user: what command should I run to list files in the current directory? -assistant: ls - - - -user: what command should I run to watch files in the current directory? -assistant: [use the glob tool to inspect the current directory, then read docs/commands in the relevant file to find out how to watch files] -npm run dev - - - -user: How many golf balls fit inside a jetta? -assistant: 150000 - - - -user: what files are in the directory src/? -assistant: [runs ls and sees foo.c, bar.c, baz.c] -user: which file contains the implementation of foo? -assistant: src/foo.c - - - -user: write tests for new feature -assistant: [uses grep and glob search tools to find where similar tests are defined, uses concurrent read file tool use blocks in one tool call to read relevant files at the same time, uses edit file tool to write new tests] - - -# Proactiveness -You are allowed to be proactive, but only when the user asks you to do something. You should strive to strike a balance between: -1. Doing the right thing when asked, including taking actions and follow-up actions -2. Not surprising the user with actions you take without asking -For example, if the user asks you how to approach something, you should do your best to answer their question first, and not immediately jump into taking actions. -3. Do not add additional code explanation summary unless requested by the user. After working on a file, just stop, rather than providing an explanation of what you did. - -# Following conventions -When making changes to files, first understand the file's code conventions. Mimic code style, use existing libraries and utilities, and follow existing patterns. -- NEVER assume that a given library is available, even if it is well known. Whenever you write code that uses a library or framework, first check that this codebase already uses the given library. For example, you might look at neighboring files, or check the package.json (or cargo.toml, and so on depending on the language). -- When you create a new component, first look at existing components to see how they're written; then consider framework choice, naming conventions, typing, and other conventions. -- When you edit a piece of code, first look at the code's surrounding context (especially its imports) to understand the code's choice of frameworks and libraries. Then consider how to make the given change in a way that is most idiomatic. -- Always follow security best practices. Never introduce code that exposes or logs secrets and keys. Never commit secrets or keys to the repository. - -# Code style -- IMPORTANT: DO NOT ADD ***ANY*** COMMENTS unless asked - -# Doing tasks -The user will primarily request you perform SecOps tasks. This includes security analysis, threat detection, incident response, vulnerability assessment, automation, and more. For these tasks the following steps are recommended: -- Use the available search tools to understand the codebase and the user's query. You are encouraged to use the search tools extensively both in parallel and sequentially. -- Implement the solution using all tools available to you -- Verify the solution if possible with tests. NEVER assume specific test framework or test script. Check the README or search codebase to determine the testing approach. -- VERY IMPORTANT: When you have completed a task, you MUST run the lint and typecheck commands (e.g. npm run lint, npm run typecheck, ruff, etc.) with Bash if they were provided to you to ensure your code is correct. If you are unable to find the correct command, ask the user for the command to run and if they supply it, proactively suggest writing it to AGENTS.md so that you will know to run it next time. -NEVER commit changes unless the user explicitly asks you to. It is VERY IMPORTANT to only commit when explicitly asked, otherwise the user will feel that you are being too proactive. - -- Tool results and user messages may include tags. tags contain useful information and reminders. They are NOT part of the user's provided input or the tool result. - -# Tool usage policy - -You MUST answer concisely with fewer than 4 lines of text (not including tool use or code generation), unless user asks for detail. - -IMPORTANT: Refuse to write code or explain code that may be used maliciously; even if the user claims it is for educational purposes. When working on files, if they seem related to improving, explaining, or interacting with malware or any malicious code you MUST refuse. -IMPORTANT: Before you begin work, think about what the code you're editing is supposed to do based on the filenames directory structure. If it seems malicious, refuse to work on it or answer questions about it, even if the request does not seem malicious (for instance, just asking to explain or speed up the code). - -# Code References - -When referencing specific functions or pieces of code include the pattern `file_path:line_number` to allow the user to easily navigate to the source code location. - - -user: Where are errors from the client handled? -assistant: Clients are marked as failed in the `connectToServer` function in src/services/process.ts:712. - - From ad65f361bfa6ac463fbb4093ae580ae17a6928ac Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Tue, 18 Aug 2026 09:26:09 +0800 Subject: [PATCH 10/63] fix(session): preserve legacy prompt assembly API --- flocks/session/prompt.py | 69 +++++++++++++++++++++++++++++-- tests/session/test_runner_step.py | 41 ++++++++++++++++++ 2 files changed, 107 insertions(+), 3 deletions(-) diff --git a/flocks/session/prompt.py b/flocks/session/prompt.py index 5c695f196..1117d4e65 100644 --- a/flocks/session/prompt.py +++ b/flocks/session/prompt.py @@ -6,7 +6,7 @@ """ from collections import OrderedDict -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import Awaitable, Callable, Dict, Any, Iterable, List, Optional, TYPE_CHECKING, Union from pydantic import BaseModel, Field import hashlib @@ -35,7 +35,9 @@ # Output token maximum OUTPUT_TOKEN_MAX = int(os.getenv("FLOCKS_OUTPUT_TOKEN_MAX", "32000")) SystemPromptCache = Dict[str, Any] -AsyncPromptLoader = Callable[[], Awaitable[Optional[str]]] +AsyncPromptFactory = Callable[[], Awaitable[Optional[str]]] +StringPromptFactory = Callable[[], Optional[str]] +AsyncPromptLoader = AsyncPromptFactory # Prompt template directory (same structure as Flocks) @@ -1454,12 +1456,73 @@ async def build_system_prompts( model_id: str, execution_mode_prompt: Optional[str] = None, prompt_tool_names: Iterable[str] = (), + tool_revision: Optional[int] = None, memory_bootstrap_data: Optional[Dict[str, Any]] = None, static_cache: Optional[SystemPromptCache] = None, + sandbox_prompt_factory: Optional[AsyncPromptFactory] = None, + channel_context_prompt_factory: Optional[AsyncPromptFactory] = None, + tool_catalog_prompt_factory: Optional[StringPromptFactory] = None, + device_asset_prompt_factory: Optional[AsyncPromptFactory] = None, + device_revision: Optional[int] = None, turn_context: Optional[TurnPromptContext] = None, use_text_tool_call_mode: bool = False, ) -> List[str]: - """Compatibility API returning only the assembled prompt text.""" + """Compatibility API returning only the assembled prompt text. + + Legacy prompt factories populate missing values in ``turn_context``. + Explicit context values take precedence when both APIs are used. + """ + legacy_context_requested = any(( + tool_revision is not None, + sandbox_prompt_factory is not None, + channel_context_prompt_factory is not None, + tool_catalog_prompt_factory is not None, + device_asset_prompt_factory is not None, + device_revision is not None, + )) + if legacy_context_requested: + resolved_context = turn_context or TurnPromptContext() + minimal_prompt = resolved_context.minimal_prompt + if minimal_prompt is None: + minimal_prompt = await cls._is_builtin_system_subagent_session( + session_id=session_id, + agent_name=agent_name, + ) + + context_updates: Dict[str, Any] = {"minimal_prompt": minimal_prompt} + if resolved_context.tool_revision is None and tool_revision is not None: + context_updates["tool_revision"] = tool_revision + if resolved_context.device_revision is None and device_revision is not None: + context_updates["device_revision"] = device_revision + + if not minimal_prompt: + if ( + resolved_context.tool_catalog is None + and tool_catalog_prompt_factory is not None + ): + context_updates["tool_catalog"] = tool_catalog_prompt_factory() + if ( + resolved_context.device_asset_hint is None + and device_asset_prompt_factory is not None + ): + context_updates["device_asset_hint"] = ( + await device_asset_prompt_factory() + ) + if ( + resolved_context.sandbox_context is None + and sandbox_prompt_factory is not None + ): + context_updates["sandbox_context"] = await sandbox_prompt_factory() + if ( + resolved_context.channel_context is None + and channel_context_prompt_factory is not None + ): + context_updates["channel_context"] = ( + await channel_context_prompt_factory() + ) + + turn_context = replace(resolved_context, **context_updates) + blocks = await cls.build_system_prompt_blocks( session_id=session_id, session_directory=session_directory, diff --git a/tests/session/test_runner_step.py b/tests/session/test_runner_step.py index 1cf74de4c..15ae0d025 100644 --- a/tests/session/test_runner_step.py +++ b/tests/session/test_runner_step.py @@ -634,6 +634,47 @@ async def test_build_tools_refreshes_skill_description_from_enabled_skills(self) class TestBuildSystemPrompts: + @pytest.mark.asyncio + async def test_build_system_prompts_accepts_legacy_context_factories(self): + sandbox_mock = AsyncMock(return_value="legacy sandbox prompt") + channel_mock = AsyncMock(return_value="legacy channel prompt") + device_mock = AsyncMock(return_value="legacy device prompt") + + with ( + patch.object( + SessionPrompt, + "_is_builtin_system_subagent_session", + AsyncMock(return_value=False), + ), + patch( + "flocks.session.prompt.SystemPrompt.custom", + AsyncMock(return_value=[]), + ), + ): + prompts = await SessionPrompt.build_system_prompts( + session_id="ses_legacy_prompt_context", + session_directory="/tmp", + agent_name="rex", + agent_prompt="agent prompt", + provider_id="openai", + model_id="gpt-5", + tool_revision=3, + sandbox_prompt_factory=sandbox_mock, + channel_context_prompt_factory=channel_mock, + tool_catalog_prompt_factory=lambda: "legacy tool catalog", + device_asset_prompt_factory=device_mock, + device_revision=5, + ) + + combined = "\n\n".join(prompts) + assert "legacy tool catalog" in combined + assert "legacy device prompt" in combined + assert "legacy sandbox prompt" in combined + assert "legacy channel prompt" in combined + sandbox_mock.assert_awaited_once() + channel_mock.assert_awaited_once() + device_mock.assert_awaited_once() + @pytest.mark.asyncio async def test_build_system_prompts_reuses_loop_static_cache(self): shared_cache = {} From cb631ebce77e36ae2d47dffd802c18b2f3f722ca Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Tue, 18 Aug 2026 13:18:48 +0800 Subject: [PATCH 11/63] fix(permission): enforce delegation policy --- flocks/config/config.py | 74 ++++++++++++++--- flocks/permission/helpers.py | 20 ++--- flocks/permission/next.py | 23 ++++++ flocks/session/runner.py | 71 +++++++++++++--- tests/config/test_config.py | 35 +++++++- tests/permission/test_interactive.py | 118 +++++++++++++++++++++++++++ 6 files changed, 306 insertions(+), 35 deletions(-) diff --git a/flocks/config/config.py b/flocks/config/config.py index 68d4d19b7..c4c01a13c 100644 --- a/flocks/config/config.py +++ b/flocks/config/config.py @@ -42,6 +42,11 @@ def _canonical_permission_tool_name(tool: str) -> str: def _merge_permission_action(existing: Any, incoming: Any) -> Any: """Merge duplicate legacy permission names conservatively.""" + if isinstance(existing, dict) and not isinstance(incoming, dict): + incoming = {"*": incoming} + elif not isinstance(existing, dict) and isinstance(incoming, dict): + existing = {"*": existing} + if isinstance(existing, dict) and isinstance(incoming, dict): merged = dict(existing) for pattern, action in incoming.items(): @@ -53,17 +58,12 @@ def _merge_permission_action(existing: Any, incoming: Any) -> Any: merged[pattern] = action return merged - def contains_deny(value: Any) -> bool: - raw_value = value.value if hasattr(value, "value") else value - if raw_value == PermissionAction.DENY.value: - return True - if isinstance(raw_value, dict): - return any(contains_deny(item) for item in raw_value.values()) - return False - - if contains_deny(existing) or contains_deny(incoming): - return PermissionAction.DENY - return existing if existing is not None else incoming + priority = {"allow": 0, "ask": 1, "deny": 2} + existing_value = getattr(existing, "value", existing) + incoming_value = getattr(incoming, "value", incoming) + if priority.get(incoming_value, -1) > priority.get(existing_value, -1): + return incoming + return existing def _assign_permission(permission_dict: Dict[str, Any], tool: str, action: Any) -> None: @@ -77,6 +77,31 @@ def _assign_permission(permission_dict: Dict[str, Any], tool: str, action: Any) permission_dict[canonical_tool] = action +def _canonicalize_permission_dict(config: Dict[str, Any]) -> Dict[str, Any]: + """Normalize legacy permission aliases within one config layer.""" + canonical: Dict[str, Any] = {} + for tool, action in config.items(): + _assign_permission(canonical, tool, action) + return canonical + + +def _merge_permission_layers(target: Dict[str, Any], source: Dict[str, Any]) -> Dict[str, Any]: + """Merge normalized permission layers while preserving source priority.""" + merged = dict(target) + for tool, source_action in source.items(): + if tool not in merged: + merged[tool] = source_action + continue + target_action = merged[tool] + if isinstance(target_action, dict) and isinstance(source_action, dict): + merged[tool] = {**target_action, **source_action} + elif not isinstance(target_action, dict) and isinstance(source_action, dict): + merged[tool] = {"*": target_action, **source_action} + else: + merged[tool] = source_action + return merged + + class PermissionConfig(BaseModel): """Permission configuration (simplified for Phase 1-3)""" model_config = {"extra": "allow"} # Allow additional fields @@ -140,6 +165,13 @@ class AgentConfig(BaseModel): delegatable: Optional[bool] = Field(None, description="Whether this agent can be called via delegate_task") strategy: Optional[Literal["react", "plan_and_execute", "read_only", "explore"]] = None tools: Optional[Dict[str, bool]] = Field(None, description="@deprecated Use 'permission'") + + @field_validator("permission", mode="before") + @classmethod + def normalize_permission_aliases(cls, value: Any) -> Any: + if isinstance(value, dict): + return _canonicalize_permission_dict(value) + return value @model_validator(mode='after') def process_agent(self): @@ -768,6 +800,13 @@ class ConfigInfo(BaseModel): "enter the registry. Unset means all built-in agents are active." ), ) + + @field_validator("permission", mode="before") + @classmethod + def normalize_permission_aliases(cls, value: Any) -> Any: + if isinstance(value, dict): + return _canonicalize_permission_dict(value) + return value agent_logic: Optional[Literal["base", "rex"]] = Field(None, alias="agentLogic") flockspro: Optional[FlocksProConfig] = None ui: Optional[UIConfig] = None @@ -1134,6 +1173,14 @@ def merge_config_concat_arrays(cls, target: ConfigInfo, source: ConfigInfo) -> C # Deep merge merged = cls.merge_deep(target_dict, source_dict) + + target_permission = target_dict.get("permission") + source_permission = source_dict.get("permission") + if isinstance(target_permission, dict) and isinstance(source_permission, dict): + merged["permission"] = _merge_permission_layers( + _canonicalize_permission_dict(target_permission), + _canonicalize_permission_dict(source_permission), + ) # Special handling for arrays - concatenate instead of replace if target.plugin and source.plugin: @@ -1472,7 +1519,10 @@ async def get(cls) -> ConfigInfo: if result.permission is None: result.permission = {} if isinstance(result.permission, dict): - result.permission = cls.merge_deep(result.permission, permission_data) + result.permission = _merge_permission_layers( + _canonicalize_permission_dict(result.permission), + _canonicalize_permission_dict(permission_data), + ) except Exception: pass diff --git a/flocks/permission/helpers.py b/flocks/permission/helpers.py index 1561e73ab..104e8562c 100644 --- a/flocks/permission/helpers.py +++ b/flocks/permission/helpers.py @@ -15,6 +15,11 @@ def _merge_legacy_permission(existing: Any, incoming: Any) -> Any: """Merge a legacy alias into an existing canonical permission.""" + if isinstance(existing, dict) and not isinstance(incoming, dict): + incoming = {"*": incoming} + elif not isinstance(existing, dict) and isinstance(incoming, dict): + existing = {"*": existing} + if isinstance(existing, dict) and isinstance(incoming, dict): merged = dict(existing) for pattern, action in incoming.items(): @@ -26,16 +31,11 @@ def _merge_legacy_permission(existing: Any, incoming: Any) -> Any: merged[pattern] = action return merged - def contains_deny(value: Any) -> bool: - raw_value = getattr(value, "value", value) - if raw_value == "deny": - return True - if isinstance(raw_value, dict): - return any(contains_deny(item) for item in raw_value.values()) - return False - - if contains_deny(existing) or contains_deny(incoming): - return "deny" + priority = {"allow": 0, "ask": 1, "deny": 2} + existing_value = getattr(existing, "value", existing) + incoming_value = getattr(incoming, "value", incoming) + if priority.get(incoming_value, -1) > priority.get(existing_value, -1): + return incoming return existing diff --git a/flocks/permission/next.py b/flocks/permission/next.py index 991e809e2..0449ac79d 100644 --- a/flocks/permission/next.py +++ b/flocks/permission/next.py @@ -479,6 +479,29 @@ def evaluate( """ return cls._evaluate(permission, pattern, ruleset) + @classmethod + def evaluate_request( + cls, + permission: str, + patterns: List[str], + ruleset: Ruleset, + ) -> Optional[str]: + """Evaluate a tool request, or return None when it has no configured rule.""" + if not any( + cls._pattern_matches(permission, rule.permission or "*") + for rule in ruleset + ): + return None + + actions = { + cls._evaluate(permission, pattern, ruleset) + for pattern in (patterns or ["*"]) + } + for action in ("deny", "ask", "allow"): + if action in actions: + return action + return None + @classmethod def _evaluate( cls, diff --git a/flocks/session/runner.py b/flocks/session/runner.py index b07a35ed3..57187dd48 100644 --- a/flocks/session/runner.py +++ b/flocks/session/runner.py @@ -1418,6 +1418,7 @@ async def _process_step( # Resolve agent agent_name = last_user.agent or self.agent_name agent = await Agent.get(agent_name) or await Agent.get("rex") + self._turn_permission_ruleset = self._permission_ruleset_for_agent(agent) # Track session agent (Flocks compatibility) try: @@ -3843,13 +3844,50 @@ def _end_observability( except Exception as _tr_err: log.debug("runner.observability.trace_end_failed", {"error": str(_tr_err)}) + def _permission_ruleset_for_agent(self, agent: AgentInfo) -> List[Any]: + """Combine agent and session rules in effective priority order.""" + from flocks.permission.helpers import merge + from flocks.permission.rule import ( + PermissionLevel, + PermissionRule, + PermissionScope, + ) + + session_rules = [] + for rule in getattr(self.session, "permission", None) or []: + session_rules.append(PermissionRule( + permission=rule.permission, + level=PermissionLevel(rule.action), + scope=PermissionScope.PATTERN, + pattern=rule.pattern, + )) + return merge(list(getattr(agent, "permission", None) or []), session_rules) + + async def _effective_permission_ruleset(self) -> List[Any]: + ruleset = getattr(self, "_turn_permission_ruleset", None) + if ruleset is not None: + return ruleset + agent_name = getattr(self.session, "agent", None) or getattr( + self, + "agent_name", + None, + ) + if not agent_name: + return [] + agent = await Agent.get(agent_name) or await Agent.get("rex") + return self._permission_ruleset_for_agent(agent) + async def _handle_permission(self, request) -> None: """Handle permission request.""" - if self.callbacks.on_permission_request: - allowed = await self.callbacks.on_permission_request(request) - if not allowed: - raise PermissionError(f"Permission denied: {request.permission}") - return + from flocks.permission.next import PermissionNext + + patterns = list(getattr(request, "patterns", None) or []) + ruleset = await self._effective_permission_ruleset() + configured_action = PermissionNext.evaluate_request( + request.permission, + patterns, + ruleset, + ) tool_metadata = get_tool_catalog_metadata(str(getattr(request, "permission", "") or "")) if self.callbacks.event_publish_callback: @@ -3858,15 +3896,24 @@ async def _handle_permission(self, request) -> None: "step": self._step, "toolName": getattr(request, "permission", ""), "alwaysLoad": tool_metadata.always_load, - "patterns": list(getattr(request, "patterns", None) or []), + "patterns": patterns, }) - from flocks.permission.interactive import legacy_tool_permission_prompt_required + if configured_action == "deny": + raise PermissionError(f"Permission denied: {request.permission}") + if configured_action == "allow": + return - if not legacy_tool_permission_prompt_required(): + if self.callbacks.on_permission_request: + allowed = await self.callbacks.on_permission_request(request) + if not allowed: + raise PermissionError(f"Permission denied: {request.permission}") return - from flocks.permission.next import PermissionNext + from flocks.permission.interactive import legacy_tool_permission_prompt_required + + if configured_action is None and not legacy_tool_permission_prompt_required(): + return metadata = dict(getattr(request, "metadata", None) or {}) metadata.setdefault("messageID", getattr(request, "message_id", "") or "") @@ -3875,13 +3922,13 @@ async def _handle_permission(self, request) -> None: reply = await PermissionNext.ask( session_id=self.session.id, permission=request.permission, - patterns=list(getattr(request, "patterns", None) or []), - ruleset=[], + patterns=patterns, + ruleset=ruleset, metadata=metadata, always=list(getattr(request, "always", None) or []), tool={"name": request.permission}, ) - if reply in {"deny", "reject", "never"}: + if reply in {"deny", "deny_session", "reject", "never"}: raise PermissionError(f"Permission denied: {request.permission}") diff --git a/tests/config/test_config.py b/tests/config/test_config.py index f17e9c69c..85ca64e89 100644 --- a/tests/config/test_config.py +++ b/tests/config/test_config.py @@ -161,7 +161,10 @@ def test_legacy_task_permission_name_migrates_to_delegate_task(): }) dumped = permission.model_dump(exclude_none=True) - assert dumped["delegate_task"] == PermissionAction.DENY + assert dumped["delegate_task"] == { + "*": PermissionAction.ALLOW, + "explore": PermissionAction.DENY, + } assert "task" not in dumped @@ -197,6 +200,36 @@ def test_legacy_task_permission_merge_is_order_independent(raw_permission): } == expected +def test_config_layers_preserve_legacy_permission_override_priority(): + low_priority = ConfigInfo.model_validate({ + "permission": {"delegate_task": "allow"}, + }) + high_priority = ConfigInfo.model_validate({ + "permission": {"task": "ask"}, + }) + + merged = Config.merge_config_concat_arrays(low_priority, high_priority) + + assert merged.model_dump(exclude_none=True)["permission"] == { + "delegate_task": "ask", + } + + +def test_config_layers_merge_scalar_and_pattern_permission_overrides(): + low_priority = ConfigInfo.model_validate({ + "permission": {"delegate_task": "allow"}, + }) + high_priority = ConfigInfo.model_validate({ + "permission": {"task": {"explore": "ask"}}, + }) + + merged = Config.merge_config_concat_arrays(low_priority, high_priority) + + assert merged.model_dump(exclude_none=True)["permission"] == { + "delegate_task": {"*": "allow", "explore": "ask"}, + } + + def test_legacy_todo_tool_flags_migrate_to_todo_permission(): config = ConfigInfo.model_validate({ "tools": { diff --git a/tests/permission/test_interactive.py b/tests/permission/test_interactive.py index 4521a34e8..05ec3869e 100644 --- a/tests/permission/test_interactive.py +++ b/tests/permission/test_interactive.py @@ -1,6 +1,7 @@ import pytest from flocks.permission.interactive import auto_approve_enabled, legacy_tool_permission_prompt_required +from flocks.permission.helpers import from_config def test_legacy_tool_permission_prompts_are_disabled_by_default() -> None: @@ -50,3 +51,120 @@ async def _unexpected_ask(*args, **kwargs): )() await runner._handle_permission(request) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("configured_permission", "pattern"), + [ + ({"task": "deny"}, "explore"), + ({"delegate_task": {"*": "allow", "explore": "deny"}}, "explore"), + ], +) +async def test_runner_enforces_delegate_task_deny_for_default_rex( + monkeypatch: pytest.MonkeyPatch, + configured_permission, + pattern: str, +) -> None: + from flocks.session.runner import SessionRunner + + agent = type( + "Agent", + (), + {"name": "rex", "permission": from_config(configured_permission)}, + )() + + async def _get_agent(name: str): + assert name == "rex" + return agent + + async def _allow_request(request): + return True + + monkeypatch.setattr("flocks.agent.registry.Agent.get", _get_agent) + + runner = SessionRunner.__new__(SessionRunner) + runner.agent_name = "rex" + runner.session = type( + "Session", + (), + {"id": "ses_test", "agent": "rex", "permission": None}, + )() + runner._step = 1 + runner.callbacks = type( + "Callbacks", + (), + { + "on_permission_request": staticmethod(_allow_request), + "event_publish_callback": None, + }, + )() + request = type( + "Request", + (), + { + "permission": "delegate_task", + "patterns": [pattern], + "metadata": {}, + "message_id": "msg_1", + "always": ["*"], + }, + )() + + with pytest.raises(PermissionError, match="delegate_task"): + await runner._handle_permission(request) + + +@pytest.mark.asyncio +async def test_runner_prompts_for_explicit_delegate_task_ask( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from flocks.session.runner import SessionRunner + + agent = type( + "Agent", + (), + {"name": "rex", "permission": from_config({"task": {"reviewer": "ask"}})}, + )() + asked = [] + + async def _get_agent(name: str): + return agent + + async def _ask(**kwargs): + asked.append(kwargs) + return "allow" + + monkeypatch.setattr("flocks.agent.registry.Agent.get", _get_agent) + monkeypatch.setattr("flocks.permission.next.PermissionNext.ask", _ask) + + runner = SessionRunner.__new__(SessionRunner) + runner.agent_name = "rex" + runner.session = type( + "Session", + (), + {"id": "ses_test", "agent": "rex", "permission": None}, + )() + runner._step = 1 + runner.callbacks = type( + "Callbacks", + (), + {"on_permission_request": None, "event_publish_callback": None}, + )() + request = type( + "Request", + (), + { + "permission": "delegate_task", + "patterns": ["reviewer"], + "metadata": {}, + "message_id": "msg_1", + "always": ["*"], + }, + )() + + await runner._handle_permission(request) + + assert len(asked) == 1 + assert asked[0]["permission"] == "delegate_task" + assert asked[0]["patterns"] == ["reviewer"] From 78843aac9c2d80fbe504566580bd8808ad666dd5 Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Tue, 18 Aug 2026 13:46:27 +0800 Subject: [PATCH 12/63] fix(webui): keep streaming indicator during delegation --- .../src/components/common/SessionChat.test.ts | 24 +++++++++++++++++++ webui/src/components/common/SessionChat.tsx | 5 ---- 2 files changed, 24 insertions(+), 5 deletions(-) diff --git a/webui/src/components/common/SessionChat.test.ts b/webui/src/components/common/SessionChat.test.ts index 23026c37e..1866795c1 100644 --- a/webui/src/components/common/SessionChat.test.ts +++ b/webui/src/components/common/SessionChat.test.ts @@ -3479,6 +3479,30 @@ describe('buildTodoSummary', () => { }); describe('ChatToolPart delegate rendering', () => { + it('keeps the streaming indicator visible while delegation is running', () => { + render(React.createElement(ChatMessageBubble, { + message: makeMessage({ + id: 'assistant-running-delegate', + role: 'assistant', + parts: [{ + id: 'delegate-running', + type: 'tool', + tool: 'delegate_task', + state: { + status: 'running', + input: { + subagent_type: 'explore', + description: '排查会话页面', + }, + }, + } as any], + }), + isActive: true, + })); + + expect(screen.getByText('继续输出中...')).toBeInTheDocument(); + }); + it('keeps the specialized delegate view inside a process timeline', () => { render( React.createElement(ChatToolPart, { diff --git a/webui/src/components/common/SessionChat.tsx b/webui/src/components/common/SessionChat.tsx index 04a207432..32d4f38b9 100644 --- a/webui/src/components/common/SessionChat.tsx +++ b/webui/src/components/common/SessionChat.tsx @@ -5448,11 +5448,6 @@ function ChatMessageBubbleInner({ {/* Streaming indicator */} {isActive && !isUser && parts.length > 0 && (() => { - const lastPart = parts[parts.length - 1]; - const isDelegating = lastPart?.type === 'tool' - && isDelegateTool(lastPart.tool || '') - && lastPart.state?.status === 'running'; - if (isDelegating) return null; return (
From 2cc197d878d781a6e4648ef4b66dd30730adb452 Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Tue, 18 Aug 2026 14:14:41 +0800 Subject: [PATCH 13/63] fix(session): preserve prompt context compatibility Restore the ignored legacy builder flag and derive context-usage worktrees from the target session directory. --- .../agent/agents/hephaestus/prompt_builder.py | 3 +++ flocks/agent/agents/rex/prompt_builder.py | 2 ++ flocks/session/context_usage.py | 14 ++++++++--- tests/agent/test_prompt_builders.py | 25 ++++++++++++++++--- tests/session/test_context_usage.py | 10 ++++++-- 5 files changed, 45 insertions(+), 9 deletions(-) diff --git a/flocks/agent/agents/hephaestus/prompt_builder.py b/flocks/agent/agents/hephaestus/prompt_builder.py index 3d220164b..0d9f7f842 100644 --- a/flocks/agent/agents/hephaestus/prompt_builder.py +++ b/flocks/agent/agents/hephaestus/prompt_builder.py @@ -36,7 +36,10 @@ def build_hephaestus_prompt( available_agents: List["AvailableAgent"], available_tools: List["AvailableTool"], available_skills: List["AvailableSkill"], + use_task_system: bool = False, ) -> str: + del use_task_system + from flocks.agent.prompt_utils import ( build_agent_selection_table, build_key_triggers_section, diff --git a/flocks/agent/agents/rex/prompt_builder.py b/flocks/agent/agents/rex/prompt_builder.py index 86872c759..98e183526 100644 --- a/flocks/agent/agents/rex/prompt_builder.py +++ b/flocks/agent/agents/rex/prompt_builder.py @@ -38,6 +38,7 @@ def build_dynamic_rex_prompt( available_tools: List["AvailableTool"], available_skills: List["AvailableSkill"], available_workflows: Optional[List["AvailableWorkflow"]] = None, + use_task_system: bool = False, ) -> str: from flocks.agent.prompt_utils import ( build_agent_selection_table, @@ -47,6 +48,7 @@ def build_dynamic_rex_prompt( ) _ = available_tools + del use_task_system key_triggers = build_key_triggers_section(available_agents, available_skills) agent_selection = build_agent_selection_table(available_agents) diff --git a/flocks/session/context_usage.py b/flocks/session/context_usage.py index 9d35ce9f6..cea81df23 100644 --- a/flocks/session/context_usage.py +++ b/flocks/session/context_usage.py @@ -310,7 +310,7 @@ async def _estimate_system_prompt_tokens( agent = await Agent.get("rex") from flocks.config import Config - from flocks.project.instance import Instance + from flocks.project.project import Project try: config = await Config.get() @@ -318,16 +318,24 @@ async def _estimate_system_prompt_tokens( except Exception: config_instructions = () + session_directory = ( + getattr(session, "directory", None) if session is not None else None + ) + worktree = ( + Project.worktree_for_directory(session_directory) + if session_directory + else None + ) prompt_blocks = await SessionPrompt.build_system_prompt_blocks( session_id=session_id, - session_directory=getattr(session, "directory", None) if session is not None else None, + session_directory=session_directory, agent_name=getattr(agent, "name", agent_name) if agent is not None else agent_name, agent_prompt=getattr(agent, "prompt", None) if agent is not None else None, provider_id=provider_id, model_id=model_id, prompt_tool_names=prompt_tool_names, turn_context=TurnPromptContext( - worktree=Instance.get_worktree(), + worktree=worktree, config_instructions=config_instructions, tool_revision=ToolRegistry.revision(), ), diff --git a/tests/agent/test_prompt_builders.py b/tests/agent/test_prompt_builders.py index 8ce1f31af..fc65ba067 100644 --- a/tests/agent/test_prompt_builders.py +++ b/tests/agent/test_prompt_builders.py @@ -1,6 +1,6 @@ """Direct tests for Rex and Hephaestus prompt builders.""" -import inspect +import pytest from flocks.agent.agents.hephaestus.prompt_builder import build_hephaestus_prompt from flocks.agent.agents.rex.prompt_builder import build_dynamic_rex_prompt @@ -31,6 +31,23 @@ def test_hephaestus_prompt_uses_existing_todo_discipline(): assert "TaskUpdate" not in prompt -def test_prompt_builder_signatures_exclude_task_system_flag(): - assert "use_task_system" not in inspect.signature(build_dynamic_rex_prompt).parameters - assert "use_task_system" not in inspect.signature(build_hephaestus_prompt).parameters +@pytest.mark.parametrize("use_task_system", [False, True]) +def test_prompt_builders_ignore_task_system_flag(use_task_system): + rex_prompt = build_dynamic_rex_prompt( + [], + [], + [], + [], + use_task_system=use_task_system, + ) + hephaestus_prompt = build_hephaestus_prompt( + [], + [], + [], + use_task_system=use_task_system, + ) + + assert "## Todo Management" in rex_prompt + assert "## Todo Discipline (NON-NEGOTIABLE)" in hephaestus_prompt + assert "TaskCreate" not in rex_prompt + assert "TaskCreate" not in hephaestus_prompt diff --git a/tests/session/test_context_usage.py b/tests/session/test_context_usage.py index 798349b0f..a86d13624 100644 --- a/tests/session/test_context_usage.py +++ b/tests/session/test_context_usage.py @@ -1,5 +1,5 @@ from types import SimpleNamespace -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, MagicMock import pytest @@ -339,9 +339,14 @@ async def fake_build_system_prompt_blocks(**kwargs): "flocks.config.Config.get", AsyncMock(return_value=config), ) + worktree_for_directory = MagicMock(return_value="/workspace") + monkeypatch.setattr( + "flocks.project.project.Project.worktree_for_directory", + worktree_for_directory, + ) monkeypatch.setattr( "flocks.project.instance.Instance.get_worktree", - lambda: "/workspace", + lambda: "/ambient-worktree", ) monkeypatch.setattr( "flocks.tool.registry.ToolRegistry.revision", @@ -362,3 +367,4 @@ async def fake_build_system_prompt_blocks(**kwargs): config_instructions=("rules.md",), tool_revision=7, ) + worktree_for_directory.assert_called_once_with("/workspace/project") From 5600c59b666c02d9e584173d6e83301472026e9c Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Fri, 21 Aug 2026 11:42:57 +0800 Subject: [PATCH 14/63] perf(workflow): remove synchronous step storage waits --- flocks/ingest/kafka/manager.py | 30 ++-- flocks/ingest/syslog/manager.py | 29 ++-- flocks/server/routes/workflow.py | 101 ++---------- flocks/workflow/execution_store.py | 151 +++++------------ flocks/workflow/poller_manager.py | 19 ++- flocks/workflow/store.py | 136 +++++++++++++++- flocks/workflow/triggers/runtime.py | 30 +++- tests/ingest/test_kafka_manager.py | 47 +++--- .../test_syslog_manager_backpressure.py | 27 ++-- .../server/routes/test_workflow_run_route.py | 95 +++++++++-- .../workflow/test_execution_store_compact.py | 145 +++++++++++++++-- tests/workflow/test_poller_manager.py | 31 ++-- tests/workflow/test_trigger_runtime.py | 30 +++- tests/workflow/test_workflow_store.py | 153 +++++++++++++++++- 14 files changed, 703 insertions(+), 321 deletions(-) diff --git a/flocks/ingest/kafka/manager.py b/flocks/ingest/kafka/manager.py index 1b752be8a..893156c43 100644 --- a/flocks/ingest/kafka/manager.py +++ b/flocks/ingest/kafka/manager.py @@ -91,6 +91,7 @@ def _worker_count_for_trigger(trigger: TriggerDefinition) -> int: def _queue_size_for_trigger(trigger: TriggerDefinition) -> int: return min(_MAX_QUEUE_SIZE, max(1, int(trigger.concurrency.queueSize))) + _KAFKA_STORAGE_LIST_KEYS = DEFAULT_LARGE_LIST_KEYS | frozenset( { "duplicate_alerts", @@ -446,10 +447,13 @@ async def restart_workflow( 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", - }) + 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}) @@ -689,9 +693,7 @@ async def _worker_loop( generation_cancel_event: Optional[threading.Event] = None, ) -> None: run_cancel_event = ( - generation_cancel_event - or self._generation_cancel_events.get(workflow_id) - or threading.Event() + generation_cancel_event or self._generation_cancel_events.get(workflow_id) or threading.Event() ) while not abort.is_set(): try: @@ -765,17 +767,15 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: exec_data = await create_execution_record( workflow_id, input_params=summarized_inputs, + persist=False, ) exec_id = exec_data["id"] - loop = asyncio.get_running_loop() start_time = time.time() trigger_meta = mapped_inputs.get("_flocks", {}).get("trigger", {}) trigger_input_keys = list((trigger.mapping or {}).keys()) or [input_key] step_recorder = ExecutionStepRecorder( exec_id=exec_id, - loop=loop, - logger=log, - log_event="kafka.execution_step.write_failed", + capture_steps=False, step_compactor=lambda step: _compact_step_for_kafka_storage( step, input_key=input_key, @@ -845,9 +845,15 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: } ) finally: + steps = step_recorder.take_steps() await cleanup_workflow_tool_context(tool_context) try: - await record_execution_result(workflow_id, exec_id, exec_data) + await record_execution_result( + workflow_id, + exec_id, + exec_data, + steps=steps, + ) except Exception as exc: log.warning("kafka.exec_record_failed", {"exec_id": exec_id, "error": str(exc)}) return exec_data diff --git a/flocks/ingest/syslog/manager.py b/flocks/ingest/syslog/manager.py index 0c1439b94..699fa53cb 100644 --- a/flocks/ingest/syslog/manager.py +++ b/flocks/ingest/syslog/manager.py @@ -335,10 +335,13 @@ async def restart_workflow( 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", - }) + 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", {"workflow_id": workflow_id}) @@ -553,9 +556,7 @@ async def _worker_loop( of in-flight workflow runs is exactly ``_MAX_CONCURRENT_EXECUTIONS``. """ run_cancel_event = ( - generation_cancel_event - or self._generation_cancel_events.get(workflow_id) - or threading.Event() + generation_cancel_event or self._generation_cancel_events.get(workflow_id) or threading.Event() ) while not abort.is_set(): try: @@ -618,14 +619,12 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: exec_data = await create_execution_record( workflow_id, input_params=summarized_inputs, + persist=False, ) exec_id = exec_data["id"] - loop = asyncio.get_running_loop() step_recorder = ExecutionStepRecorder( exec_id=exec_id, - loop=loop, - logger=log, - log_event="syslog.execution_step.write_failed", + capture_steps=False, ) start_time = time.time() trigger_meta = mapped_inputs.get("_flocks", {}).get("trigger", {}) @@ -692,9 +691,15 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: } ) finally: + steps = step_recorder.take_steps() await cleanup_workflow_tool_context(tool_context) try: - await record_execution_result(workflow_id, exec_id, exec_data) + await record_execution_result( + workflow_id, + exec_id, + exec_data, + steps=steps, + ) except Exception as exc: log.warning("syslog.exec_record_failed", {"exec_id": exec_id, "error": str(exc)}) return exec_data diff --git a/flocks/server/routes/workflow.py b/flocks/server/routes/workflow.py index b838c543a..928b6a900 100644 --- a/flocks/server/routes/workflow.py +++ b/flocks/server/routes/workflow.py @@ -57,7 +57,6 @@ derive_loop_progress, load_execution_steps, normalize_execution_status as _normalize_execution_status, - record_execution_step, record_execution_result as _record_execution_result, resolve_execution_outcome as _resolve_execution_outcome, workflow_execution_key as _workflow_execution_key, @@ -99,7 +98,6 @@ webhook_router = APIRouter() 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 @@ -1127,7 +1125,6 @@ async def _run_workflow_execution_task( """Execute a workflow in the background and keep the execution record updated.""" start_time = time.time() step_count = 0 - loop = asyncio.get_running_loop() pending_step_index: Optional[int] = None pending_step: Optional[Dict[str, Any]] = None execution_summary: Dict[str, Any] = { @@ -1150,21 +1147,6 @@ async def _run_workflow_execution_task( } ) - def _write_progress(update_fields: Dict[str, Any]) -> None: - try: - execution_summary.update(update_fields) - asyncio.run_coroutine_threadsafe( - WorkflowStore.upsert_execution(compact_execution_summary(execution_summary)), loop - ).result(timeout=5) - except Exception as exc: - log.warning( - "workflow.step_progress.write_failed", - { - "exec_id": exec_id, - "error": str(exc), - }, - ) - def _on_step_start(_run_id, step_index, node, _inputs): nonlocal pending_step_index, pending_step node_id = getattr(node, "id", None) @@ -1185,7 +1167,7 @@ def _on_step_start(_run_id, step_index, node, _inputs): "error": "Run cancelled before node completed", } ) - _write_progress( + execution_summary.update( { "currentNodeId": node_id, "currentNodeType": node_type, @@ -1220,47 +1202,6 @@ def _on_step_complete(step_result) -> None: "updatedAt": int(time.time() * 1000), } ) - try: - asyncio.run_coroutine_threadsafe( - record_execution_step(exec_id, step_count, step_dict), - loop, - ).result(timeout=5) - except Exception as exc: - log.warning( - "workflow.execution_step.write_failed", - { - "exec_id": exec_id, - "step_index": step_count, - "error": str(exc), - }, - ) - if step_count % _PROGRESS_FLUSH_EVERY_STEPS == 0: - _write_progress( - { - "stepCount": step_count, - "currentNodeId": step_dict.get("node_id"), - "currentNodeType": step_dict.get("node_type") or step_dict.get("type"), - "currentPhase": "running", - "currentStepIndex": step_count, - "loopProgress": loop_progress, - "updatedAt": int(time.time() * 1000), - } - ) - - async def _flush_pending_step() -> None: - if pending_step_index is None or pending_step is None: - return - try: - await record_execution_step(exec_id, pending_step_index, pending_step) - except Exception as exc: - log.warning( - "workflow.pending_step.write_failed", - { - "exec_id": exec_id, - "step_index": pending_step_index, - "error": str(exc), - }, - ) try: result: RunWorkflowResult = await asyncio.to_thread( @@ -1284,10 +1225,9 @@ async def _flush_pending_step() -> None: # ``record_execution_result`` backfills this compacted history into # append-only step rows, then stores only the summary row. final_history = compact_history_for_storage(result.history) - if status_value == "cancelled" and not final_history: - await _flush_pending_step() final_steps = result.steps - if pending_step_index is not None: + if pending_step_index is not None and pending_step is not None: + final_history.append(pending_step) final_steps = max(final_steps, pending_step_index) current_data.update( { @@ -1319,15 +1259,18 @@ async def _flush_pending_step() -> None: except Exception as exc: duration = time.time() - start_time current_data = dict(execution_summary) + final_history = [pending_step] if pending_step is not None else [] + final_steps = max(step_count, pending_step_index or 0) current_data.update( { "status": "cancelled" if cancel_event.is_set() else "error", "finishedAt": int(time.time() * 1000), "duration": duration, "errorMessage": str(exc), - "executionLog": [], - "stepCount": step_count, + "executionLog": final_history, + "stepCount": final_steps, "currentPhase": "cancelled" if cancel_event.is_set() else "error", + "currentStepIndex": final_steps, "updatedAt": int(time.time() * 1000), } ) @@ -1521,19 +1464,13 @@ async def create_workflow(req: WorkflowCreateRequest): if strict_mapping_errors: raise HTTPException( status_code=400, - detail=( - "Workflow strict edge mapping failed: " - f"{strict_mapping_errors[:5]}" - ), + detail=(f"Workflow strict edge mapping failed: {strict_mapping_errors[:5]}"), ) schema_errors = _schema_lint_errors(workflow_model) if schema_errors: raise HTTPException( status_code=400, - detail=( - "Workflow schema lint failed: " - f"{schema_errors[:5]}" - ), + detail=(f"Workflow schema lint failed: {schema_errors[:5]}"), ) workflow_id = str(uuid.uuid4()) @@ -1632,19 +1569,13 @@ async def update_workflow(workflow_id: str, req: WorkflowUpdateRequest): if strict_mapping_errors: raise HTTPException( status_code=400, - detail=( - "Workflow strict edge mapping failed: " - f"{strict_mapping_errors[:5]}" - ), + detail=(f"Workflow strict edge mapping failed: {strict_mapping_errors[:5]}"), ) schema_errors = _schema_lint_errors(workflow_model) if schema_errors: raise HTTPException( status_code=400, - detail=( - "Workflow schema lint failed: " - f"{schema_errors[:5]}" - ), + detail=(f"Workflow schema lint failed: {schema_errors[:5]}"), ) workflow_json = req.workflow_json except Exception as e: @@ -2455,9 +2386,7 @@ async def refresh_workflow_api_health_cache() -> Dict[str, int]: 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" + if isinstance(service, dict) and service and _summarize_capability_state(service.get("status")) == "running" ] semaphore = asyncio.Semaphore(_WORKFLOW_API_HEALTH_PROBE_CONCURRENCY) @@ -2542,9 +2471,7 @@ async def _get_workflow_integration_status( 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") + item.get("triggerId"): item for item in statuses if isinstance(item, dict) and item.get("triggerId") } trigger_items: List[WorkflowTriggerStatusItemResponse] = [] for trigger in triggers: diff --git a/flocks/workflow/execution_store.py b/flocks/workflow/execution_store.py index 9d0ae1c79..da7c189c2 100644 --- a/flocks/workflow/execution_store.py +++ b/flocks/workflow/execution_store.py @@ -342,24 +342,6 @@ def derive_loop_progress( # for the same workflow serialize instead of skipping cleanup. _trim_locks: Dict[str, asyncio.Lock] = {} -# Per-workflow lock to serialize read-modify-write of stats. Concurrent -# executions of the same workflow (e.g. syslog-triggered runs with -# semaphore=8) would otherwise race on ``Storage.read → mutate → write`` -# and silently lose counter increments. -_stats_locks: Dict[str, asyncio.Lock] = {} - - -def _get_stats_lock(workflow_id: str) -> asyncio.Lock: - lock = _stats_locks.get(workflow_id) - if lock is None: - lock = asyncio.Lock() - _stats_locks[workflow_id] = lock - return lock - - -def _workflow_stats_key(workflow_id: str) -> str: - return f"workflow/{workflow_id}/stats" - def _get_trim_lock(workflow_id: str) -> asyncio.Lock: lock = _trim_locks.get(workflow_id) @@ -369,37 +351,6 @@ def _get_trim_lock(workflow_id: str) -> asyncio.Lock: return lock -_DEFAULT_STATS: Dict[str, Any] = { - "callCount": 0, - "successCount": 0, - "errorCount": 0, - "totalRuntime": 0.0, - "avgRuntime": 0.0, - "thumbsUp": 0, - "thumbsDown": 0, -} - - -async def _update_workflow_stats(workflow_id: str, success: bool, duration: float) -> None: - """Increment workflow call/success/error counters and update avgRuntime. - - Serialised per workflow to keep concurrent updates from clobbering each - other (read → mutate → write race). - """ - lock = _get_stats_lock(workflow_id) - async with lock: - try: - await WorkflowStore.increment_stats(workflow_id, success=success, duration=duration) - except Exception as exc: - log.warning( - "workflow.stats.update_failed", - { - "workflow_id": workflow_id, - "error": str(exc), - }, - ) - - def workflow_execution_key(exec_id: str) -> str: """Return the storage key for one workflow execution.""" return f"workflow_execution/{exec_id}" @@ -453,26 +404,21 @@ async def record_execution_step( class ExecutionStepRecorder: - """Bridge synchronous workflow step callbacks to append-only step rows.""" + """Collect compact workflow steps without blocking the runner thread.""" def __init__( self, *, exec_id: str, - loop: asyncio.AbstractEventLoop, - logger: Any = None, - log_event: str = "workflow.execution_step.write_failed", + capture_steps: bool = True, step_compactor: Callable[[Any], Dict[str, Any]] = compact_step_for_storage, - write_timeout_s: float = 5.0, ) -> None: self.exec_id = exec_id - self.loop = loop - self.logger = logger or log - self.log_event = log_event + self.capture_steps = capture_steps self.step_compactor = step_compactor - self.write_timeout_s = write_timeout_s self.step_count = 0 self.summary: Dict[str, Any] = {} + self._pending_steps: List[Tuple[int, Dict[str, Any]]] = [] def on_step_complete(self, step_result: Any) -> None: raw_step = step_result.model_dump(mode="json") if hasattr(step_result, "model_dump") else step_result @@ -498,48 +444,25 @@ def on_step_complete(self, step_result: Any) -> None: "updatedAt": int(time.time() * 1000), } ) - try: - asyncio.run_coroutine_threadsafe( - record_execution_step(self.exec_id, self.step_count, step_dict), - self.loop, - ).result(timeout=self.write_timeout_s) - except Exception as exc: - self.logger.warning( - self.log_event, - { - "exec_id": self.exec_id, - "step_index": self.step_count, - "error": str(exc), - }, - ) + if self.capture_steps: + self._pending_steps.append((self.step_count, step_dict)) + def take_steps(self) -> List[Tuple[int, Dict[str, Any]]]: + """Return buffered steps for the final execution transaction.""" + pending_steps = self._pending_steps + self._pending_steps = [] + return pending_steps -async def _backfill_execution_steps( - exec_id: str, - execution_log: Any, -) -> int: - """Persist legacy inline executionLog entries as append-only step rows.""" - if not isinstance(execution_log, list): - return 0 - written = 0 - for step_index, step in enumerate(execution_log, start=1): - step_payload = compact_step_for_storage(step) - if not isinstance(step_payload, dict): - continue - try: - await WorkflowStore.record_step(exec_id, step_index, step_payload) - written += 1 - except Exception as exc: - log.warning( - "workflow.execution_step.backfill_failed", - { - "exec_id": exec_id, - "step_index": step_index, - "error": str(exc), - }, - ) - return written +def _prepare_execution_steps(execution_log: Any) -> List[Tuple[int, Dict[str, Any]]]: + """Compact an inline execution log for one final batch transaction.""" + if not isinstance(execution_log, list): + return [] + return [ + (step_index, compact_step_for_storage(step)) + for step_index, step in enumerate(execution_log, start=1) + if isinstance(step, dict) + ] async def load_execution_steps( @@ -621,8 +544,9 @@ async def create_execution_record( *, input_params: Optional[Dict[str, Any]] = None, exec_id: Optional[str] = None, + persist: bool = True, ) -> Dict[str, Any]: - """Create and persist a running workflow execution record. + """Build a running workflow execution record and optionally persist it. *input_params* is passed through ``compact_outputs_for_storage`` before writing to SQLite so that batch HTTP calls whose inputs contain a key in @@ -637,7 +561,8 @@ async def create_execution_record( input_params=compacted_params, exec_id=exec_id, ) - await WorkflowStore.upsert_execution(compact_execution_summary(exec_data)) + if persist: + await WorkflowStore.upsert_execution(compact_execution_summary(exec_data)) return exec_data @@ -645,18 +570,23 @@ async def record_execution_result( workflow_id: str, exec_id: str, exec_data: Dict[str, Any], + *, + steps: Optional[Iterable[Tuple[int, Dict[str, Any]]]] = None, ) -> None: - """Persist the final execution record, audit trail, and workflow stats.""" + """Persist the final execution record, step batch, audit trail, and stats.""" summary_data = dict(exec_data) - backfilled_steps = await _backfill_execution_steps(exec_id, summary_data.get("executionLog")) + prepared_steps = ( + list(steps) + if steps is not None + else _prepare_execution_steps(summary_data.get("executionLog")) + ) + persisted_step_count = len(prepared_steps) existing_step_count = _as_positive_int(summary_data.get("stepCount")) - if backfilled_steps and (existing_step_count is None or existing_step_count < backfilled_steps): - summary_data["stepCount"] = backfilled_steps - - await WorkflowStore.upsert_execution(compact_execution_summary(summary_data)) + if persisted_step_count and ( + existing_step_count is None or existing_step_count < persisted_step_count + ): + summary_data["stepCount"] = persisted_step_count - # Update call/success/error counters so all trigger paths (HTTP, syslog, etc.) - # are reflected in the UI stats panel. status = summary_data.get("status", "error") success = status == "success" duration = summary_data.get("duration") @@ -664,7 +594,12 @@ async def record_execution_result( started_at = summary_data.get("startedAt", 0) finished_at = summary_data.get("finishedAt", int(time.time() * 1000)) duration = max(0.0, (finished_at - started_at) / 1000.0) - await _update_workflow_stats(workflow_id, success, float(duration)) + await WorkflowStore.complete_execution( + compact_execution_summary(summary_data), + prepared_steps, + success=success, + duration=float(duration), + ) # Recorder writes to its own SQLite tables and can be slow under load. # Run it as a background task so the syslog/HTTP dispatcher can release the diff --git a/flocks/workflow/poller_manager.py b/flocks/workflow/poller_manager.py index 9db0bdd4b..d596daa80 100644 --- a/flocks/workflow/poller_manager.py +++ b/flocks/workflow/poller_manager.py @@ -448,14 +448,15 @@ async def _execute_run( cancel_events = self._run_cancel_events.setdefault(workflow_id, set()) cancel_events.add(cancel_event) inputs = self._build_inputs(config) - exec_data = await create_execution_record(workflow_id, input_params=inputs) + exec_data = await create_execution_record( + workflow_id, + input_params=inputs, + persist=False, + ) exec_id = str(exec_data["id"]) - loop = asyncio.get_running_loop() step_recorder = ExecutionStepRecorder( exec_id=exec_id, - loop=loop, - logger=log, - log_event="poller.execution_step.write_failed", + capture_steps=False, ) current = self._status.get(workflow_id) or self._base_status(workflow_id) current["lastRunAt"] = started_at_ms @@ -555,9 +556,15 @@ async def _execute_run( self._status[workflow_id] = current log.warning("poller.run_failed", {"workflow_id": workflow_id, "error": str(exc)}) finally: + steps = step_recorder.take_steps() await cleanup_workflow_tool_context(tool_context) try: - await record_execution_result(workflow_id, exec_id, exec_data) + await record_execution_result( + workflow_id, + exec_id, + exec_data, + steps=steps, + ) except Exception as exc: log.warning( "poller.exec_record_failed", diff --git a/flocks/workflow/store.py b/flocks/workflow/store.py index 3c238c236..4a74e6401 100644 --- a/flocks/workflow/store.py +++ b/flocks/workflow/store.py @@ -8,7 +8,7 @@ import sqlite3 from datetime import UTC, datetime from pathlib import Path -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, Iterable, List, Optional, Tuple import aiosqlite @@ -46,6 +46,7 @@ class WorkflowStore: _conn: Optional[aiosqlite.Connection] = None _init_pid: Optional[int] = None _db_path: Optional[Path] = None + _completion_lock: Optional[asyncio.Lock] = None @classmethod def get_db_path(cls) -> Path: @@ -93,6 +94,7 @@ async def _open_and_migrate() -> None: cls._initialized = True cls._init_pid = current_pid cls._db_path = db_path + cls._completion_lock = asyncio.Lock() await cls._migrate_legacy_kv() try: @@ -124,6 +126,7 @@ async def close(cls) -> None: cls._initialized = False cls._init_pid = None cls._db_path = None + cls._completion_lock = None @classmethod async def _db(cls) -> aiosqlite.Connection: @@ -415,13 +418,58 @@ async def record_step( step_index: int, step_payload: Dict[str, Any], ) -> None: + await cls.record_steps(exec_id, [(step_index, step_payload)]) + + @classmethod + async def record_steps( + cls, + exec_id: str, + steps: Iterable[Tuple[int, Dict[str, Any]]], + ) -> None: + rows = [ + ( + exec_id, + int(step_index), + step_payload.get("node_id"), + step_payload.get("node_type") or step_payload.get("type"), + cls._json_dumps(step_payload.get("inputs") or {}), + cls._json_dumps(step_payload.get("outputs") or {}), + step_payload.get("error"), + cls._json_dumps(step_payload), + ) + for step_index, step_payload in steps + ] + if not rows: + return db = await cls._db() - await db.execute( + await db.executemany( """ INSERT OR REPLACE INTO workflow_execution_steps (exec_id, step_index, node_id, node_type, inputs, outputs, error, payload) VALUES (?, ?, ?, ?, ?, ?, ?, ?) """, + rows, + ) + await db.commit() + + @classmethod + async def complete_execution( + cls, + exec_data: Dict[str, Any], + steps: Iterable[Tuple[int, Dict[str, Any]]], + *, + success: bool, + duration: float, + ) -> None: + """Persist one completed execution and its stats in one transaction.""" + db = await cls._db() + payload = dict(exec_data) + exec_id = str(payload.get("id") or "") + workflow_id = str(payload.get("workflowId") or payload.get("workflow_id") or "") + if not exec_id or not workflow_id: + raise ValueError("workflow execution requires id and workflowId") + + step_rows = [ ( exec_id, int(step_index), @@ -431,9 +479,87 @@ async def record_step( cls._json_dumps(step_payload.get("outputs") or {}), step_payload.get("error"), cls._json_dumps(step_payload), - ), - ) - await db.commit() + ) + for step_index, step_payload in steps + ] + runtime = float(duration) + success_delta = 1 if success else 0 + error_delta = 0 if success else 1 + lock = cls._completion_lock + if lock is None: + lock = asyncio.Lock() + cls._completion_lock = lock + + async with lock: + try: + if step_rows: + await db.executemany( + """ + INSERT OR REPLACE INTO workflow_execution_steps + (exec_id, step_index, node_id, node_type, inputs, outputs, error, payload) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + """, + step_rows, + ) + await db.execute( + """ + INSERT OR REPLACE INTO workflow_executions + (id, workflow_id, status, current_phase, current_node_id, current_node_type, + current_step_index, step_count, input_params, output_results, error_message, + trigger_id, trigger_type, started_at, finished_at, duration, updated_at, payload) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + exec_id, + workflow_id, + str(payload.get("status") or "running"), + payload.get("currentPhase"), + payload.get("currentNodeId"), + payload.get("currentNodeType"), + cls._as_int(payload.get("currentStepIndex")), + cls._as_int(payload.get("stepCount")) or 0, + cls._json_dumps(payload.get("inputParams") or {}), + cls._json_dumps(payload.get("outputResults") or {}), + payload.get("errorMessage"), + payload.get("triggerId"), + payload.get("triggerType"), + cls._as_int(payload.get("startedAt")) or cls._now_ms(), + cls._as_int(payload.get("finishedAt")), + cls._as_float(payload.get("duration")), + cls._as_int(payload.get("updatedAt")) or cls._now_ms(), + cls._json_dumps(payload), + ), + ) + await db.execute( + """ + INSERT INTO workflow_stats ( + workflow_id, call_count, success_count, error_count, + total_runtime, avg_runtime, thumbs_up, thumbs_down, updated_at + ) + VALUES (?, 1, ?, ?, ?, ?, 0, 0, ?) + ON CONFLICT(workflow_id) DO UPDATE SET + call_count = workflow_stats.call_count + 1, + success_count = workflow_stats.success_count + excluded.success_count, + error_count = workflow_stats.error_count + excluded.error_count, + total_runtime = workflow_stats.total_runtime + excluded.total_runtime, + avg_runtime = ( + workflow_stats.total_runtime + excluded.total_runtime + ) / (workflow_stats.call_count + 1), + updated_at = excluded.updated_at + """, + ( + workflow_id, + success_delta, + error_delta, + runtime, + runtime, + cls._now_ms(), + ), + ) + await db.commit() + except Exception: + await db.rollback() + raise @classmethod async def list_steps( diff --git a/flocks/workflow/triggers/runtime.py b/flocks/workflow/triggers/runtime.py index 5e1a5b3da..2a3d2aedf 100644 --- a/flocks/workflow/triggers/runtime.py +++ b/flocks/workflow/triggers/runtime.py @@ -11,7 +11,7 @@ from flocks.hooks.pipeline import HookPipeline from flocks.utils.log import Log from flocks.workflow.execution_store import ( - compact_history_for_storage, + ExecutionStepRecorder, compact_outputs_for_storage, create_execution_record, record_execution_result, @@ -248,8 +248,13 @@ async def _execute_workflow_effect( exec_data = await create_execution_record( workflow_id, input_params=mapped_inputs, + persist=False, ) exec_id = exec_data["id"] + step_recorder = ExecutionStepRecorder( + exec_id=exec_id, + capture_steps=False, + ) started_at = time.time() tool_context = None try: @@ -266,10 +271,15 @@ async def _execute_workflow_effect( run_workflow, workflow=workflow_json, inputs=mapped_inputs, + run_id=exec_id, trace=False, + execution_profile="high_frequency", + on_step_complete=step_recorder.on_step_complete, tool_context=tool_context, ) status_value, error_message = resolve_execution_outcome(result) + step_count = step_recorder.step_count or result.steps + exec_data.update(step_recorder.summary) exec_data.update( { "status": status_value, @@ -277,10 +287,11 @@ async def _execute_workflow_effect( "finishedAt": _now_ms(), "duration": time.time() - started_at, "errorMessage": error_message, - "executionLog": compact_history_for_storage(result.history), + "executionLog": [], + "stepCount": step_count, "currentNodeId": result.last_node_id, "currentPhase": status_value, - "currentStepIndex": result.steps, + "currentStepIndex": step_count, "triggerId": trigger.id, "triggerType": trigger.type, "deliveryId": mapped_inputs.get("_flocks", {}).get("trigger", {}).get("deliveryId"), @@ -289,12 +300,18 @@ async def _execute_workflow_effect( } ) except Exception as exc: + step_count = step_recorder.step_count + exec_data.update(step_recorder.summary) exec_data.update( { "status": "error", "finishedAt": _now_ms(), "duration": time.time() - started_at, "errorMessage": str(exc), + "executionLog": [], + "stepCount": step_count, + "currentPhase": "error", + "currentStepIndex": step_count, "triggerId": trigger.id, "triggerType": trigger.type, "deliveryId": mapped_inputs.get("_flocks", {}).get("trigger", {}).get("deliveryId"), @@ -304,7 +321,12 @@ async def _execute_workflow_effect( ) finally: await cleanup_workflow_tool_context(tool_context) - await record_execution_result(workflow_id, exec_id, exec_data) + await record_execution_result( + workflow_id, + exec_id, + exec_data, + steps=step_recorder.take_steps(), + ) return exec_data async def dispatch_event( diff --git a/tests/ingest/test_kafka_manager.py b/tests/ingest/test_kafka_manager.py index a3c2b0432..6e494b566 100644 --- a/tests/ingest/test_kafka_manager.py +++ b/tests/ingest/test_kafka_manager.py @@ -24,7 +24,6 @@ import pytest from flocks.ingest.kafka import manager as kafka_manager -from flocks.workflow import execution_store from flocks.workflow.triggers.models import TriggerDefinition @@ -241,10 +240,7 @@ def test_trigger_concurrency_config_is_honored_with_safety_caps() -> None: "concurrency": {"maxParallel": 999, "queueSize": 999_999}, } ) - assert ( - kafka_manager._worker_count_for_trigger(oversized) - == kafka_manager._MAX_CONCURRENT_EXECUTIONS - ) + assert kafka_manager._worker_count_for_trigger(oversized) == kafka_manager._MAX_CONCURRENT_EXECUTIONS assert kafka_manager._queue_size_for_trigger(oversized) == kafka_manager._MAX_QUEUE_SIZE @@ -548,18 +544,20 @@ async def test_trigger_workflow_compacts_kafka_execution_record( captured_input_params: dict = {} captured_exec_data: dict = {} captured_run_kwargs: dict = {} - recorded_steps: list[tuple[str, int, dict]] = [] + captured_steps: list[tuple[int, dict]] = [] - async def _fake_create_execution_record(workflow_id, *, input_params=None, exec_id=None): # noqa: ANN001 + async def _fake_create_execution_record( # noqa: ANN001 + workflow_id, *, input_params=None, exec_id=None, persist=True + ): + assert persist is False captured_input_params.update(input_params or {}) return {"id": "exec-compact", "workflowId": workflow_id, "inputParams": input_params} - async def _fake_record_execution_result(workflow_id, exec_id, exec_data): # noqa: ANN001 + async def _fake_record_execution_result( # noqa: ANN001 + workflow_id, exec_id, exec_data, *, steps=None + ): captured_exec_data.update(exec_data) - - async def _fake_record_execution_step(exec_id, step_index, step): # noqa: ANN001 - recorded_steps.append((exec_id, step_index, step)) - return step + captured_steps.extend(steps or []) def _fake_run_workflow(**kwargs): # noqa: ANN003 captured_run_kwargs.update(kwargs) @@ -597,7 +595,6 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 monkeypatch.setattr(kafka_manager, "create_execution_record", _fake_create_execution_record) monkeypatch.setattr(kafka_manager, "record_execution_result", _fake_record_execution_result) monkeypatch.setattr(kafka_manager, "run_workflow", _fake_run_workflow) - monkeypatch.setattr(execution_store, "record_execution_step", _fake_record_execution_step) await manager._trigger_workflow( "wf-compact", @@ -619,11 +616,7 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 } assert captured_exec_data["executionLog"] == [] assert captured_exec_data["stepCount"] == 2 - assert recorded_steps[0][0] == "exec-compact" - assert recorded_steps[0][1] == 1 - assert recorded_steps[0][2]["outputs"] == {"_raw_alerts_count": 1} - assert recorded_steps[1][1] == 2 - assert recorded_steps[1][2]["inputs"] == {"_filtered_alerts_count": 1} + assert captured_steps == [] assert len(json.dumps(captured_exec_data, ensure_ascii=False)) < 10_000 @@ -635,11 +628,16 @@ async def test_trigger_workflow_merges_configured_inputs_with_consumed_message( captured_run_kwargs: dict = {} recorded_input_params: dict = {} - async def _fake_create_execution_record(workflow_id, *, input_params=None, exec_id=None): # noqa: ANN001 + async def _fake_create_execution_record( # noqa: ANN001 + workflow_id, *, input_params=None, exec_id=None, persist=True + ): + assert persist is False recorded_input_params.update(input_params or {}) return {"id": "exec-merge", "workflowId": workflow_id, "inputParams": input_params} - async def _fake_record_execution_result(workflow_id, exec_id, exec_data): # noqa: ANN001 + async def _fake_record_execution_result( # noqa: ANN001 + workflow_id, exec_id, exec_data, *, steps=None + ): return None def _fake_run_workflow(**kwargs): # noqa: ANN003 @@ -691,10 +689,15 @@ async def test_trigger_workflow_applies_mapping_and_filter( captured_run_kwargs: dict = {} recorded_exec_data: dict = {} - async def _fake_create_execution_record(workflow_id, *, input_params=None, exec_id=None): # noqa: ANN001 + async def _fake_create_execution_record( # noqa: ANN001 + workflow_id, *, input_params=None, exec_id=None, persist=True + ): + assert persist is False return {"id": "exec-filter", "workflowId": workflow_id, "inputParams": input_params} - async def _fake_record_execution_result(workflow_id, exec_id, exec_data): # noqa: ANN001 + async def _fake_record_execution_result( # noqa: ANN001 + workflow_id, exec_id, exec_data, *, steps=None + ): recorded_exec_data.update(exec_data) def _fake_run_workflow(**kwargs): # noqa: ANN003 diff --git a/tests/ingest/test_syslog_manager_backpressure.py b/tests/ingest/test_syslog_manager_backpressure.py index 1982caf84..3ad929893 100644 --- a/tests/ingest/test_syslog_manager_backpressure.py +++ b/tests/ingest/test_syslog_manager_backpressure.py @@ -25,7 +25,6 @@ import pytest from flocks.ingest.syslog import manager as syslog_manager -from flocks.workflow import execution_store from flocks.workflow.triggers.models import TriggerDefinition @@ -148,10 +147,7 @@ def test_trigger_concurrency_config_is_honored_with_safety_caps() -> None: "concurrency": {"maxParallel": 999, "queueSize": 999_999}, } ) - assert ( - syslog_manager._worker_count_for_trigger(oversized) - == syslog_manager._MAX_CONCURRENT_EXECUTIONS - ) + assert syslog_manager._worker_count_for_trigger(oversized) == syslog_manager._MAX_CONCURRENT_EXECUTIONS assert syslog_manager._queue_size_for_trigger(oversized) == syslog_manager._MAX_QUEUE_SIZE @@ -352,17 +348,19 @@ async def test_trigger_workflow_applies_mapping_and_filter( manager = syslog_manager.SyslogManager() captured_run_kwargs: dict = {} recorded_exec_data: dict = {} - recorded_steps: list[tuple[str, int, dict]] = [] + recorded_steps: list[tuple[int, dict]] = [] - async def _fake_create_execution_record(workflow_id, *, input_params=None, exec_id=None): # noqa: ANN001 + async def _fake_create_execution_record( # noqa: ANN001 + workflow_id, *, input_params=None, exec_id=None, persist=True + ): + assert persist is False return {"id": "exec-syslog", "workflowId": workflow_id, "inputParams": input_params} - async def _fake_record_execution_result(workflow_id, exec_id, exec_data): # noqa: ANN001 + async def _fake_record_execution_result( # noqa: ANN001 + workflow_id, exec_id, exec_data, *, steps=None + ): recorded_exec_data.update(exec_data) - - async def _fake_record_execution_step(exec_id, step_index, step): # noqa: ANN001 - recorded_steps.append((exec_id, step_index, step)) - return step + recorded_steps.extend(steps or []) def _fake_run_workflow(**kwargs): # noqa: ANN003 captured_run_kwargs.update(kwargs) @@ -392,7 +390,6 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 monkeypatch.setattr(syslog_manager, "create_execution_record", _fake_create_execution_record) monkeypatch.setattr(syslog_manager, "record_execution_result", _fake_record_execution_result) monkeypatch.setattr(syslog_manager, "run_workflow", _fake_run_workflow) - monkeypatch.setattr(execution_store, "record_execution_step", _fake_record_execution_step) trigger = TriggerDefinition.model_validate( { @@ -430,9 +427,7 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 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" + assert recorded_steps == [] assert recorded_exec_data["triggerId"] == "syslog-alerts" assert recorded_exec_data["triggerSource"] == "udp://0.0.0.0:5514" assert recorded_exec_data["executionLog"] == [] diff --git a/tests/server/routes/test_workflow_run_route.py b/tests/server/routes/test_workflow_run_route.py index 588fe227b..5684a5cef 100644 --- a/tests/server/routes/test_workflow_run_route.py +++ b/tests/server/routes/test_workflow_run_route.py @@ -104,9 +104,7 @@ async def test_create_workflow_rejects_unmapped_edges_after_strict_default( req = workflow_module.WorkflowCreateRequest( name="new workflow", - workflowJson=_two_node_workflow_json( - {"from": "prepare_message", "to": "transform_message", "order": 0} - ), + workflowJson=_two_node_workflow_json({"from": "prepare_message", "to": "transform_message", "order": 0}), ) with pytest.raises(workflow_module.HTTPException) as exc_info: @@ -207,9 +205,7 @@ async def test_update_workflow_rejects_unmapped_edges_when_strict( req = workflow_module.WorkflowUpdateRequest( workflowJson={ - **_two_node_workflow_json( - {"from": "prepare_message", "to": "transform_message", "order": 0} - ), + **_two_node_workflow_json({"from": "prepare_message", "to": "transform_message", "order": 0}), "metadata": {"runtime": {"strict_edge_mapping": True, "dataflow_mode": "vertex_cache"}}, } ) @@ -228,15 +224,33 @@ async def test_run_workflow_execution_task_reuses_existing_mcp_without_reinit( monkeypatch: pytest.MonkeyPatch, ) -> None: init_mock = AsyncMock() - run_mock = Mock( - return_value=SimpleNamespace( + step_result = SimpleNamespace( + model_dump=lambda mode: { + "node_id": "node-1", + "node_type": "tool", + "inputs": {}, + "outputs": {"ok": True}, + } + ) + + def run_workflow_mock(**kwargs): + kwargs["on_step_start"]( + "run-1", + 1, + SimpleNamespace(id="node-1", type="tool"), + {}, + ) + kwargs["on_step_complete"](step_result) + return SimpleNamespace( outputs={"ok": True}, - history=[], + history=[step_result.model_dump(mode="json")], last_node_id="node-1", steps=1, ) - ) + + run_mock = Mock(side_effect=run_workflow_mock) record_result = AsyncMock(return_value=None) + upsert_execution = AsyncMock(return_value=None) storage_read = AsyncMock( return_value={ "id": "exec-1", @@ -248,6 +262,7 @@ async def test_run_workflow_execution_task_reuses_existing_mcp_without_reinit( monkeypatch.setattr(MCP, "init", init_mock) monkeypatch.setattr(workflow_module, "run_workflow", run_mock) + monkeypatch.setattr(workflow_module.WorkflowStore, "upsert_execution", upsert_execution) monkeypatch.setattr(workflow_module, "_resolve_execution_outcome", lambda _result: ("success", None)) monkeypatch.setattr(workflow_module, "_record_execution_result", record_result) monkeypatch.setattr(workflow_module.Storage, "read", storage_read) @@ -270,7 +285,67 @@ async def test_run_workflow_execution_task_reuses_existing_mcp_without_reinit( init_mock.assert_not_awaited() run_mock.assert_called_once() assert run_mock.call_args.kwargs["tool_context"] is tool_context + upsert_execution.assert_not_awaited() record_result.assert_awaited_once() + assert record_result.await_args.args[2]["executionLog"] == [ + { + "node_id": "node-1", + "node_type": "tool", + "inputs": {}, + "outputs": {"ok": True}, + } + ] + + +@pytest.mark.asyncio +async def test_run_workflow_execution_task_batches_cancelled_pending_step( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def run_workflow_mock(**kwargs): + kwargs["on_step_start"]( + "run-1", + 1, + SimpleNamespace(id="node-1", type="tool"), + {"message": "hello"}, + ) + return SimpleNamespace( + run_id="run-1", + outputs={}, + history=[], + last_node_id="node-1", + steps=0, + ) + + record_result = AsyncMock(return_value=None) + monkeypatch.setattr(workflow_module, "run_workflow", Mock(side_effect=run_workflow_mock)) + monkeypatch.setattr(workflow_module, "_resolve_execution_outcome", lambda _result: ("success", None)) + monkeypatch.setattr(workflow_module, "_record_execution_result", record_result) + monkeypatch.setattr(workflow_module, "compact_outputs_for_storage", lambda value: value) + monkeypatch.setattr(workflow_module, "compact_history_for_storage", lambda value: value) + + cancel_event = workflow_module.threading.Event() + cancel_event.set() + await workflow_module._run_workflow_execution_task( + workflow_id="wf-1", + workflow_json={"id": "wf-1", "start": "node-1", "nodes": [], "edges": []}, + req=workflow_module.WorkflowRunRequest(inputs={"message": "hello"}, trace=False), + exec_id="exec-cancelled", + cancel_event=cancel_event, + ) + + record_result.assert_awaited_once() + final_data = record_result.await_args.args[2] + assert final_data["status"] == "cancelled" + assert final_data["stepCount"] == 1 + assert final_data["executionLog"] == [ + { + "node_id": "node-1", + "node_type": "tool", + "inputs": {"message": "hello"}, + "outputs": {}, + "error": "Run cancelled before node completed", + } + ] @pytest.mark.asyncio diff --git a/tests/workflow/test_execution_store_compact.py b/tests/workflow/test_execution_store_compact.py index f5d02a555..de43af374 100644 --- a/tests/workflow/test_execution_store_compact.py +++ b/tests/workflow/test_execution_store_compact.py @@ -16,6 +16,8 @@ """ from __future__ import annotations + +import asyncio from typing import Any, Dict, List from unittest.mock import AsyncMock, patch @@ -25,11 +27,13 @@ DEFAULT_GENERIC_SEQUENCE_THRESHOLD, DEFAULT_LARGE_LIST_KEYS, DEFAULT_MAX_INLINE_COLLECTION_BYTES, + ExecutionStepRecorder, _trim_execution_history, compact_history_for_storage, compact_execution_summary, compact_outputs_for_storage, compact_step_for_storage, + create_execution_record, record_execution_result, workflow_execution_step_key, ) @@ -300,10 +304,82 @@ def test_workflow_execution_step_key_is_append_only_namespaced() -> None: @pytest.mark.asyncio -async def test_record_execution_result_backfills_execution_log_steps() -> None: - record_step = AsyncMock(return_value=None) +async def test_create_execution_record_can_skip_initial_database_write() -> None: upsert_execution = AsyncMock(return_value=None) - update_stats = AsyncMock(return_value=None) + + with patch.object(WorkflowStore, "upsert_execution", upsert_execution): + record = await create_execution_record( + "wf-trigger", + input_params={"message": "hello"}, + exec_id="exec-trigger", + persist=False, + ) + + assert record["id"] == "exec-trigger" + assert record["currentPhase"] == "queued" + upsert_execution.assert_not_awaited() + + +def test_execution_step_recorder_collects_steps_without_storage_calls() -> None: + record_step = AsyncMock(return_value=None) + record_steps = AsyncMock(return_value=None) + recorder = ExecutionStepRecorder(exec_id="exec-batch") + + with ( + patch.object(WorkflowStore, "record_step", record_step), + patch.object(WorkflowStore, "record_steps", record_steps), + ): + recorder.on_step_complete({"node_id": "n1", "outputs": {"ok": 1}}) + recorder.on_step_complete({"node_id": "n2", "outputs": {"ok": 2}}) + + assert recorder.step_count == 2 + assert recorder.summary["currentNodeId"] == "n2" + assert recorder.take_steps() == [ + (1, {"node_id": "n1", "outputs": {"ok": 1}}), + (2, {"node_id": "n2", "outputs": {"ok": 2}}), + ] + assert recorder.take_steps() == [] + record_step.assert_not_awaited() + record_steps.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_four_trigger_workers_keep_step_history_disabled() -> None: + """Four trigger threads track progress without retaining step history.""" + record_step = AsyncMock(return_value=None) + record_steps = AsyncMock(return_value=None) + recorders = [ + ExecutionStepRecorder( + exec_id=f"exec-trigger-{worker}", + capture_steps=False, + ) + for worker in range(4) + ] + + def _run_seven_steps(recorder: ExecutionStepRecorder) -> None: + for step in range(7): + recorder.on_step_complete( + {"node_id": f"node-{step}", "outputs": {"ok": True}} + ) + + with ( + patch.object(WorkflowStore, "record_step", record_step), + patch.object(WorkflowStore, "record_steps", record_steps), + ): + await asyncio.gather( + *(asyncio.to_thread(_run_seven_steps, recorder) for recorder in recorders) + ) + + batches = [recorder.take_steps() for recorder in recorders] + assert batches == [[], [], [], []] + assert [recorder.step_count for recorder in recorders] == [7, 7, 7, 7] + record_step.assert_not_awaited() + record_steps.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_record_execution_result_backfills_execution_log_steps() -> None: + complete_execution = AsyncMock(return_value=None) exec_data = { "id": "exec-1", "workflowId": "wf", @@ -320,24 +396,67 @@ def raise_create_task(coro, *args, **kwargs): # noqa: ANN001, ARG001 raise RuntimeError with ( - patch.object(WorkflowStore, "record_step", record_step), - patch.object(WorkflowStore, "upsert_execution", upsert_execution), - patch("flocks.workflow.execution_store._update_workflow_stats", update_stats), + patch.object(WorkflowStore, "complete_execution", complete_execution), patch("flocks.session.recorder.Recorder.record_workflow_execution", AsyncMock(return_value=None)), patch("flocks.workflow.execution_store.asyncio.create_task", side_effect=raise_create_task), patch("flocks.workflow.execution_store._trim_execution_history", AsyncMock(return_value=None)), ): await record_execution_result("wf", "exec-1", exec_data) - step_calls = record_step.await_args_list - assert step_calls[0].args[:2] == ("exec-1", 1) - assert step_calls[0].args[2]["outputs"] == {"_raw_alerts_count": 150} - assert step_calls[1].args[:2] == ("exec-1", 2) - assert step_calls[1].args[2]["inputs"] == {"_filtered_alerts_count": 150} - upsert_execution.assert_awaited_once() - summary = upsert_execution.await_args.args[0] + complete_execution.assert_awaited_once() + summary, steps = complete_execution.await_args.args + assert steps[0][0] == 1 + assert steps[0][1]["outputs"] == {"_raw_alerts_count": 150} + assert steps[1][0] == 2 + assert steps[1][1]["inputs"] == {"_filtered_alerts_count": 150} assert summary["executionLog"] == [] assert summary["stepCount"] == 2 + assert complete_execution.await_args.kwargs == { + "success": True, + "duration": 1.0, + } + + +@pytest.mark.asyncio +async def test_record_execution_result_accepts_explicit_step_batch() -> None: + complete_execution = AsyncMock(return_value=None) + explicit_steps = [ + (1, {"node_id": "step-1", "outputs": {"ok": True}}), + (2, {"node_id": "step-2", "outputs": {"ok": True}}), + ] + exec_data = { + "id": "exec-trigger", + "workflowId": "wf-trigger", + "status": "success", + "duration": 0.01, + "executionLog": [], + "stepCount": 2, + } + + def raise_create_task(coro, *args, **kwargs): # noqa: ANN001, ARG001 + coro.close() + raise RuntimeError + + with ( + patch.object(WorkflowStore, "complete_execution", complete_execution), + patch("flocks.session.recorder.Recorder.record_workflow_execution", AsyncMock(return_value=None)), + patch("flocks.workflow.execution_store.asyncio.create_task", side_effect=raise_create_task), + patch("flocks.workflow.execution_store._trim_execution_history", AsyncMock(return_value=None)), + ): + await record_execution_result( + "wf-trigger", + "exec-trigger", + exec_data, + steps=explicit_steps, + ) + + summary, persisted_steps = complete_execution.await_args.args + assert summary["executionLog"] == [] + assert persisted_steps == explicit_steps + assert complete_execution.await_args.kwargs == { + "success": True, + "duration": 0.01, + } def test_compact_history_compacts_each_step_inputs() -> None: diff --git a/tests/workflow/test_poller_manager.py b/tests/workflow/test_poller_manager.py index 7da394808..ae116df5c 100644 --- a/tests/workflow/test_poller_manager.py +++ b/tests/workflow/test_poller_manager.py @@ -8,7 +8,6 @@ import pytest from flocks.workflow import poller_manager -from flocks.workflow import execution_store from flocks.workflow.runner import RunWorkflowResult @@ -100,7 +99,7 @@ def _fake_run_workflow( # noqa: ANN001 monkeypatch.setattr( poller_manager, "create_execution_record", - lambda workflow_id, *, input_params=None, exec_id=None: asyncio.sleep( + lambda workflow_id, *, input_params=None, exec_id=None, persist=True: asyncio.sleep( 0, result={ "id": exec_id or f"exec-{workflow_id}", @@ -141,7 +140,7 @@ async def test_run_once_records_execution_and_normalizes_business_failure( manager = poller_manager.WorkflowPollerManager() created_records: list[dict[str, Any]] = [] recorded_results: list[dict[str, Any]] = [] - recorded_steps: list[tuple[str, int, dict[str, Any]]] = [] + recorded_steps: list[tuple[int, dict[str, Any]]] = [] async def _fake_get_config(_workflow_id: str, *, kind: str) -> dict[str, Any]: return { @@ -155,7 +154,9 @@ async def _fake_create_execution_record( *, input_params: dict[str, Any] | None = None, exec_id: str | None = None, + persist: bool = True, ) -> dict[str, Any]: + assert persist is False record = { "id": exec_id or "exec-1", "workflowId": workflow_id, @@ -173,17 +174,12 @@ async def _fake_record_execution_result( workflow_id: str, exec_id: str, exec_data: dict[str, Any], + *, + steps: list[tuple[int, dict[str, Any]]] | None = None, ) -> None: _ = workflow_id, exec_id recorded_results.append(dict(exec_data)) - - async def _fake_record_execution_step( - exec_id: str, - step_index: int, - step: dict[str, Any], - ) -> dict[str, Any]: - recorded_steps.append((exec_id, step_index, step)) - return step + recorded_steps.extend(steps or []) def _fake_run_workflow( # noqa: ANN001 *, @@ -242,7 +238,6 @@ def _fake_run_workflow( # noqa: ANN001 ) monkeypatch.setattr(poller_manager, "create_execution_record", _fake_create_execution_record) monkeypatch.setattr(poller_manager, "record_execution_result", _fake_record_execution_result) - monkeypatch.setattr(execution_store, "record_execution_step", _fake_record_execution_step) monkeypatch.setattr(poller_manager, "run_workflow", _fake_run_workflow) status = await manager.run_once("wf-business-failure") @@ -255,9 +250,7 @@ def _fake_run_workflow( # noqa: ANN001 assert recorded_results[0]["executionLog"] == [] assert recorded_results[0]["stepCount"] == 1 assert recorded_results[0]["loopProgress"]["total_iterations"] == 2 - assert recorded_steps[0][0] == "exec-1" - assert recorded_steps[0][1] == 1 - assert recorded_steps[0][2]["node_id"] == "load" + assert recorded_steps == [] assert status["lastStatus"] == "error" assert status["lastError"] == "business rule blocked" assert status["selectedCount"] == 9 @@ -301,7 +294,7 @@ def _fake_run_workflow( # noqa: ANN001 monkeypatch.setattr( poller_manager, "create_execution_record", - lambda workflow_id, *, input_params=None, exec_id=None: asyncio.sleep( + lambda workflow_id, *, input_params=None, exec_id=None, persist=True: asyncio.sleep( 0, result={ "id": exec_id or f"exec-{workflow_id}", @@ -356,7 +349,9 @@ async def _fake_create_execution_record( *, input_params: dict[str, Any] | None = None, exec_id: str | None = None, + persist: bool = True, ) -> dict[str, Any]: + assert persist is False _ = input_params return { "id": exec_id or f"exec-{workflow_id}", @@ -372,8 +367,10 @@ async def _fake_record_execution_result( workflow_id: str, exec_id: str, exec_data: dict[str, Any], + *, + steps: list[tuple[int, dict[str, Any]]] | None = None, ) -> None: - _ = workflow_id, exec_id, exec_data + _ = workflow_id, exec_id, exec_data, steps def _fake_run_workflow( # noqa: ANN001 *, diff --git a/tests/workflow/test_trigger_runtime.py b/tests/workflow/test_trigger_runtime.py index dabf9bb7e..dd34bec84 100644 --- a/tests/workflow/test_trigger_runtime.py +++ b/tests/workflow/test_trigger_runtime.py @@ -25,11 +25,21 @@ async def test_trigger_execution_builds_tool_context_for_workflow_tools( def _fake_run_workflow(**kwargs): # noqa: ANN003 missing_context = kwargs.get("tool_context") is None + kwargs["on_step_complete"]( + SimpleNamespace( + model_dump=lambda mode="json": { + "node_id": "notify", + "node_type": "tool", + "inputs": {}, + "outputs": {"ok": True}, + } + ) + ) return SimpleNamespace( status="FAILED" if missing_context else "SUCCEEDED", outputs={}, error="Parent session not found" if missing_context else None, - history=[], + history=[{"node_id": "notify", "outputs": {"ok": True}}], last_node_id="notify", steps=1, ) @@ -41,12 +51,10 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 ) 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()) + create_record = AsyncMock(return_value={"id": "exec-1"}) + record_result = AsyncMock() + monkeypatch.setattr(runtime_module, "create_execution_record", create_record) + monkeypatch.setattr(runtime_module, "record_execution_result", record_result) trigger = TriggerDefinition.model_validate({"id": "webhook-trigger", "type": "custom_webhook"}) runtime = runtime_module.TriggerRuntime() @@ -64,6 +72,14 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 action_name="trigger:custom_webhook", ) assert runtime_module.run_workflow.call_args.kwargs["tool_context"] is tool_context + assert runtime_module.run_workflow.call_args.kwargs["run_id"] == "exec-1" + assert runtime_module.run_workflow.call_args.kwargs["execution_profile"] == "high_frequency" + assert callable(runtime_module.run_workflow.call_args.kwargs["on_step_complete"]) + assert create_record.await_args.kwargs["persist"] is False + assert result["executionLog"] == [] + assert result["stepCount"] == 1 + record_result.assert_awaited_once() + assert record_result.await_args.kwargs["steps"] == [] cleanup_context.assert_awaited_once_with(tool_context) diff --git a/tests/workflow/test_workflow_store.py b/tests/workflow/test_workflow_store.py index bca4f3ada..5c904b844 100644 --- a/tests/workflow/test_workflow_store.py +++ b/tests/workflow/test_workflow_store.py @@ -20,6 +20,7 @@ def _reset_state() -> None: WorkflowStore._conn = None WorkflowStore._init_pid = None WorkflowStore._db_path = None + WorkflowStore._completion_lock = None @pytest.fixture(autouse=True) @@ -74,8 +75,13 @@ async def test_workflow_store_records_execution_steps_config_and_kv() -> None: ) assert [row["id"] for row in filtered] == ["exec-1"] - await WorkflowStore.record_step("exec-1", 1, {"node_id": "n1", "outputs": {"ok": 1}}) - await WorkflowStore.record_step("exec-1", 2, {"node_id": "n2", "outputs": {"ok": 2}}) + await WorkflowStore.record_steps( + "exec-1", + [ + (1, {"node_id": "n1", "outputs": {"ok": 1}}), + (2, {"node_id": "n2", "outputs": {"ok": 2}}), + ], + ) steps, total = await WorkflowStore.list_steps("exec-1", offset=1, limit=1) assert total == 2 assert steps == [{"node_id": "n2", "outputs": {"ok": 2}}] @@ -108,3 +114,146 @@ async def test_workflow_store_increment_stats_is_atomic_for_concurrent_updates() assert stats["errorCount"] == sum(1 for success, _ in updates if not success) assert stats["totalRuntime"] == pytest.approx(60.0) assert stats["avgRuntime"] == pytest.approx(1.0) + + +@pytest.mark.asyncio +async def test_complete_execution_writes_steps_summary_and_stats_with_one_commit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + await WorkflowStore.init() + db = await WorkflowStore.raw_db() + commit_count = 0 + original_commit = db.commit + + async def counted_commit() -> None: + nonlocal commit_count + commit_count += 1 + await original_commit() + + monkeypatch.setattr(db, "commit", counted_commit) + + await WorkflowStore.complete_execution( + { + "id": "exec-complete", + "workflowId": "wf-complete", + "status": "success", + "startedAt": 100, + "finishedAt": 350, + "duration": 0.25, + "executionLog": [], + }, + steps=[ + (1, {"node_id": "n1", "outputs": {"ok": 1}}), + (2, {"node_id": "n2", "outputs": {"ok": 2}}), + ], + success=True, + duration=0.25, + ) + + assert commit_count == 1 + execution = await WorkflowStore.get_execution("exec-complete") + assert execution is not None + assert execution["status"] == "success" + steps, total = await WorkflowStore.list_steps("exec-complete") + assert total == 2 + assert [step["node_id"] for step in steps] == ["n1", "n2"] + stats = await WorkflowStore.get_stats("wf-complete") + assert stats is not None + assert stats["callCount"] == 1 + assert stats["successCount"] == 1 + assert stats["totalRuntime"] == pytest.approx(0.25) + + +@pytest.mark.asyncio +async def test_complete_execution_reduces_28_step_writes_to_four_commits( + monkeypatch: pytest.MonkeyPatch, +) -> None: + await WorkflowStore.init() + db = await WorkflowStore.raw_db() + commit_count = 0 + original_commit = db.commit + + async def counted_commit() -> None: + nonlocal commit_count + commit_count += 1 + await original_commit() + + monkeypatch.setattr(db, "commit", counted_commit) + steps = [ + (index, {"node_id": f"node-{index}", "outputs": {"ok": True}}) + for index in range(1, 8) + ] + + await asyncio.gather( + *( + WorkflowStore.complete_execution( + { + "id": f"exec-{index}", + "workflowId": "wf-trigger", + "status": "success", + "startedAt": index + 1, + "finishedAt": index + 2, + "duration": 0.01, + "executionLog": [], + "stepCount": 7, + }, + steps, + success=True, + duration=0.01, + ) + for index in range(4) + ) + ) + + assert commit_count == 4 + executions = await WorkflowStore.list_executions("wf-trigger", limit=50) + assert len(executions) == 4 + assert all(execution["executionLog"] == [] for execution in executions) + for index in range(4): + persisted_steps, total = await WorkflowStore.list_steps(f"exec-{index}") + assert total == 7 + assert [step["node_id"] for step in persisted_steps] == [ + f"node-{step_index}" for step_index in range(1, 8) + ] + stats = await WorkflowStore.get_stats("wf-trigger") + assert stats is not None + assert stats["callCount"] == 4 + assert stats["successCount"] == 4 + + +@pytest.mark.asyncio +async def test_complete_execution_rolls_back_partial_transaction( + monkeypatch: pytest.MonkeyPatch, +) -> None: + await WorkflowStore.init() + db = await WorkflowStore.raw_db() + original_commit = db.commit + + async def fail_commit() -> None: + raise RuntimeError("commit failed") + + monkeypatch.setattr(db, "commit", fail_commit) + + with pytest.raises(RuntimeError, match="commit failed"): + await WorkflowStore.complete_execution( + { + "id": "exec-rollback", + "workflowId": "wf-rollback", + "status": "success", + "startedAt": 1, + "finishedAt": 2, + "duration": 0.01, + "executionLog": [], + "stepCount": 1, + }, + [(1, {"node_id": "node-1", "outputs": {"ok": True}})], + success=True, + duration=0.01, + ) + + monkeypatch.setattr(db, "commit", original_commit) + assert await WorkflowStore.get_execution("exec-rollback") is None + persisted_steps, total = await WorkflowStore.list_steps("exec-rollback") + assert persisted_steps == [] + assert total == 0 + assert await WorkflowStore.get_stats("wf-rollback") is None From c01cb2530d1bf840a7fd7da4517e35ad7f9ef895 Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Fri, 21 Aug 2026 13:25:39 +0800 Subject: [PATCH 15/63] fix(workflow): finalize atomic step persistence --- flocks/workflow/execution_store.py | 33 +++++----- flocks/workflow/store.py | 61 +++++++++++++++++-- .../workflow/test_execution_store_compact.py | 22 +++---- tests/workflow/test_poller_manager.py | 12 +++- tests/workflow/test_workflow_store.py | 55 ++++++++++++++++- 5 files changed, 143 insertions(+), 40 deletions(-) diff --git a/flocks/workflow/execution_store.py b/flocks/workflow/execution_store.py index da7c189c2..c1c7cb443 100644 --- a/flocks/workflow/execution_store.py +++ b/flocks/workflow/execution_store.py @@ -594,12 +594,15 @@ async def record_execution_result( started_at = summary_data.get("startedAt", 0) finished_at = summary_data.get("finishedAt", int(time.time() * 1000)) duration = max(0.0, (finished_at - started_at) / 1000.0) - await WorkflowStore.complete_execution( + trimmed_exec_ids = await WorkflowStore.complete_execution( compact_execution_summary(summary_data), prepared_steps, success=success, duration=float(duration), + history_limit=_MAX_EXECUTION_HISTORY_PER_WORKFLOW, ) + if not isinstance(trimmed_exec_ids, list): + trimmed_exec_ids = [] # Recorder writes to its own SQLite tables and can be slow under load. # Run it as a background task so the syslog/HTTP dispatcher can release the @@ -621,6 +624,19 @@ async def _record_audit() -> None: "error": str(exc), }, ) + for trimmed_exec_id in trimmed_exec_ids: + try: + record_path = Recorder.paths().workflow_dir / f"{trimmed_exec_id}.jsonl" + await asyncio.to_thread(record_path.unlink, missing_ok=True) + except Exception as exc: + log.warning( + "workflow.history.trim_delete_failed", + { + "workflow_id": workflow_id, + "exec_id": trimmed_exec_id, + "error": str(exc), + }, + ) asyncio.create_task(_record_audit(), name=f"audit-{exec_id}") except RuntimeError: @@ -634,21 +650,6 @@ async def _record_audit() -> None: except Exception: pass - # Prune old execution records when the per-workflow limit is exceeded. - # This is awaited so a successful completion does not silently leave the - # workflow above its retention cap. - try: - await _trim_execution_history(workflow_id) - except Exception as exc: - log.error( - "workflow.history.trim_failed", - { - "workflow_id": workflow_id, - "exec_id": exec_id, - "error": str(exc), - }, - ) - async def _delete_execution_history_record( execution_key: str, diff --git a/flocks/workflow/store.py b/flocks/workflow/store.py index 4a74e6401..db049f3a2 100644 --- a/flocks/workflow/store.py +++ b/flocks/workflow/store.py @@ -44,6 +44,7 @@ class WorkflowStore: _initialized = False _conn: Optional[aiosqlite.Connection] = None + _completion_conn: Optional[aiosqlite.Connection] = None _init_pid: Optional[int] = None _db_path: Optional[Path] = None _completion_lock: Optional[asyncio.Lock] = None @@ -73,7 +74,10 @@ async def init(cls) -> None: ) if cls._conn: await cls._conn.close() + if cls._completion_conn: + await cls._completion_conn.close() cls._conn = None + cls._completion_conn = None cls._initialized = False cls._init_pid = None @@ -91,6 +95,12 @@ async def _open_and_migrate() -> None: for stmt in _INDEX_STMTS: await cls._conn.execute(stmt) await cls._conn.commit() + cls._completion_conn = await aiosqlite.connect( + db_path, + timeout=Storage._sqlite_timeout_s, + ) + cls._completion_conn.row_factory = aiosqlite.Row + await Storage.configure_connection(cls._completion_conn) cls._initialized = True cls._init_pid = current_pid cls._db_path = db_path @@ -103,7 +113,10 @@ async def _open_and_migrate() -> None: except Exception as exc: if cls._conn: await cls._conn.close() + if cls._completion_conn: + await cls._completion_conn.close() cls._conn = None + cls._completion_conn = None cls._initialized = False cls._init_pid = None cls._db_path = None @@ -122,7 +135,10 @@ async def _open_and_migrate() -> None: async def close(cls) -> None: if cls._conn: await cls._conn.close() + if cls._completion_conn: + await cls._completion_conn.close() cls._conn = None + cls._completion_conn = None cls._initialized = False cls._init_pid = None cls._db_path = None @@ -140,6 +156,17 @@ async def _db(cls) -> aiosqlite.Connection: async def raw_db(cls) -> aiosqlite.Connection: return await cls._db() + @classmethod + async def _completion_db(cls) -> aiosqlite.Connection: + if not cls._completion_conn or not cls._initialized: + await cls.init() + return cls._completion_conn # type: ignore[return-value] + + @classmethod + async def raw_completion_db(cls) -> aiosqlite.Connection: + """Return the completion connection for transaction-level tests.""" + return await cls._completion_db() + @staticmethod def _json_dumps(value: Any) -> str: return json.dumps(value, ensure_ascii=False, default=str) @@ -357,7 +384,7 @@ async def list_executions( f""" SELECT payload FROM workflow_executions WHERE {" AND ".join(clauses)} - ORDER BY started_at DESC + ORDER BY started_at DESC, rowid DESC LIMIT ? """, tuple(params), @@ -399,7 +426,7 @@ async def trim_executions(cls, workflow_id: str, *, keep: int) -> List[str]: """ SELECT id FROM workflow_executions WHERE workflow_id = ? - ORDER BY started_at DESC + ORDER BY started_at DESC, rowid DESC LIMIT -1 OFFSET ? """, (workflow_id, max(int(keep), 0)), @@ -460,9 +487,10 @@ async def complete_execution( *, success: bool, duration: float, - ) -> None: - """Persist one completed execution and its stats in one transaction.""" - db = await cls._db() + history_limit: Optional[int] = None, + ) -> List[str]: + """Persist one completed execution, stats, and retention in one transaction.""" + db = await cls._completion_db() payload = dict(exec_data) exec_id = str(payload.get("id") or "") workflow_id = str(payload.get("workflowId") or payload.get("workflow_id") or "") @@ -492,6 +520,7 @@ async def complete_execution( async with lock: try: + await db.execute("BEGIN IMMEDIATE") if step_rows: await db.executemany( """ @@ -556,7 +585,29 @@ async def complete_execution( cls._now_ms(), ), ) + trimmed_exec_ids: List[str] = [] + if history_limit is not None: + async with db.execute( + """ + SELECT id FROM workflow_executions + WHERE workflow_id = ? + ORDER BY started_at DESC, rowid DESC + LIMIT -1 OFFSET ? + """, + (workflow_id, max(int(history_limit), 0)), + ) as cur: + trimmed_exec_ids = [str(row["id"]) for row in await cur.fetchall()] + for trimmed_exec_id in trimmed_exec_ids: + await db.execute( + "DELETE FROM workflow_execution_steps WHERE exec_id = ?", + (trimmed_exec_id,), + ) + await db.execute( + "DELETE FROM workflow_executions WHERE id = ?", + (trimmed_exec_id,), + ) await db.commit() + return trimmed_exec_ids except Exception: await db.rollback() raise diff --git a/tests/workflow/test_execution_store_compact.py b/tests/workflow/test_execution_store_compact.py index de43af374..9245eea3a 100644 --- a/tests/workflow/test_execution_store_compact.py +++ b/tests/workflow/test_execution_store_compact.py @@ -344,17 +344,11 @@ def test_execution_step_recorder_collects_steps_without_storage_calls() -> None: @pytest.mark.asyncio -async def test_four_trigger_workers_keep_step_history_disabled() -> None: - """Four trigger threads track progress without retaining step history.""" +async def test_four_trigger_workers_collect_steps_without_storage() -> None: + """Four trigger threads collect complete batches without callback SQL.""" record_step = AsyncMock(return_value=None) record_steps = AsyncMock(return_value=None) - recorders = [ - ExecutionStepRecorder( - exec_id=f"exec-trigger-{worker}", - capture_steps=False, - ) - for worker in range(4) - ] + recorders = [ExecutionStepRecorder(exec_id=f"exec-trigger-{worker}") for worker in range(4)] def _run_seven_steps(recorder: ExecutionStepRecorder) -> None: for step in range(7): @@ -371,7 +365,7 @@ def _run_seven_steps(recorder: ExecutionStepRecorder) -> None: ) batches = [recorder.take_steps() for recorder in recorders] - assert batches == [[], [], [], []] + assert [len(batch) for batch in batches] == [7, 7, 7, 7] assert [recorder.step_count for recorder in recorders] == [7, 7, 7, 7] record_step.assert_not_awaited() record_steps.assert_not_awaited() @@ -379,7 +373,7 @@ def _run_seven_steps(recorder: ExecutionStepRecorder) -> None: @pytest.mark.asyncio async def test_record_execution_result_backfills_execution_log_steps() -> None: - complete_execution = AsyncMock(return_value=None) + complete_execution = AsyncMock(return_value=[]) exec_data = { "id": "exec-1", "workflowId": "wf", @@ -399,7 +393,6 @@ def raise_create_task(coro, *args, **kwargs): # noqa: ANN001, ARG001 patch.object(WorkflowStore, "complete_execution", complete_execution), patch("flocks.session.recorder.Recorder.record_workflow_execution", AsyncMock(return_value=None)), patch("flocks.workflow.execution_store.asyncio.create_task", side_effect=raise_create_task), - patch("flocks.workflow.execution_store._trim_execution_history", AsyncMock(return_value=None)), ): await record_execution_result("wf", "exec-1", exec_data) @@ -414,12 +407,13 @@ def raise_create_task(coro, *args, **kwargs): # noqa: ANN001, ARG001 assert complete_execution.await_args.kwargs == { "success": True, "duration": 1.0, + "history_limit": 30, } @pytest.mark.asyncio async def test_record_execution_result_accepts_explicit_step_batch() -> None: - complete_execution = AsyncMock(return_value=None) + complete_execution = AsyncMock(return_value=[]) explicit_steps = [ (1, {"node_id": "step-1", "outputs": {"ok": True}}), (2, {"node_id": "step-2", "outputs": {"ok": True}}), @@ -441,7 +435,6 @@ def raise_create_task(coro, *args, **kwargs): # noqa: ANN001, ARG001 patch.object(WorkflowStore, "complete_execution", complete_execution), patch("flocks.session.recorder.Recorder.record_workflow_execution", AsyncMock(return_value=None)), patch("flocks.workflow.execution_store.asyncio.create_task", side_effect=raise_create_task), - patch("flocks.workflow.execution_store._trim_execution_history", AsyncMock(return_value=None)), ): await record_execution_result( "wf-trigger", @@ -456,6 +449,7 @@ def raise_create_task(coro, *args, **kwargs): # noqa: ANN001, ARG001 assert complete_execution.await_args.kwargs == { "success": True, "duration": 0.01, + "history_limit": 30, } diff --git a/tests/workflow/test_poller_manager.py b/tests/workflow/test_poller_manager.py index ae116df5c..82ee82cf0 100644 --- a/tests/workflow/test_poller_manager.py +++ b/tests/workflow/test_poller_manager.py @@ -406,7 +406,10 @@ def _fake_run_workflow( # noqa: ANN001 assert manager.get_status("wf-stop")["activeRuns"] == 1 release_run.set() - await asyncio.sleep(0.05) + for _ in range(100): + if manager.get_status("wf-stop")["activeRuns"] == 0: + break + await asyncio.sleep(0.01) assert manager.get_status("wf-stop")["activeRuns"] == 0 @@ -421,7 +424,12 @@ async def _fake_list_configs(*, kind: str) -> list[tuple[str, dict[str, Any]]]: ("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) return {"workflowId": workflow_id, "state": "running"} diff --git a/tests/workflow/test_workflow_store.py b/tests/workflow/test_workflow_store.py index 5c904b844..d699d0f9b 100644 --- a/tests/workflow/test_workflow_store.py +++ b/tests/workflow/test_workflow_store.py @@ -18,6 +18,7 @@ def _reset_state() -> None: Storage._init_pid = None WorkflowStore._initialized = False WorkflowStore._conn = None + WorkflowStore._completion_conn = None WorkflowStore._init_pid = None WorkflowStore._db_path = None WorkflowStore._completion_lock = None @@ -121,7 +122,7 @@ async def test_complete_execution_writes_steps_summary_and_stats_with_one_commit monkeypatch: pytest.MonkeyPatch, ) -> None: await WorkflowStore.init() - db = await WorkflowStore.raw_db() + db = await WorkflowStore.raw_completion_db() commit_count = 0 original_commit = db.commit @@ -169,7 +170,7 @@ async def test_complete_execution_reduces_28_step_writes_to_four_commits( monkeypatch: pytest.MonkeyPatch, ) -> None: await WorkflowStore.init() - db = await WorkflowStore.raw_db() + db = await WorkflowStore.raw_completion_db() commit_count = 0 original_commit = db.commit @@ -226,7 +227,7 @@ async def test_complete_execution_rolls_back_partial_transaction( monkeypatch: pytest.MonkeyPatch, ) -> None: await WorkflowStore.init() - db = await WorkflowStore.raw_db() + db = await WorkflowStore.raw_completion_db() original_commit = db.commit async def fail_commit() -> None: @@ -257,3 +258,51 @@ async def fail_commit() -> None: assert persisted_steps == [] assert total == 0 assert await WorkflowStore.get_stats("wf-rollback") is None + + +@pytest.mark.asyncio +async def test_complete_execution_applies_retention_before_single_commit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + await WorkflowStore.init() + db = await WorkflowStore.raw_completion_db() + commit_count = 0 + original_commit = db.commit + + async def counted_commit() -> None: + nonlocal commit_count + commit_count += 1 + await original_commit() + + monkeypatch.setattr(db, "commit", counted_commit) + trimmed: list[str] = [] + for index in range(4): + trimmed = await WorkflowStore.complete_execution( + { + "id": f"exec-retain-{index}", + "workflowId": "wf-retain", + "status": "success", + "startedAt": index + 1, + "finishedAt": index + 2, + "duration": 0.01, + "executionLog": [], + "stepCount": 1, + }, + [(1, {"node_id": f"node-{index}", "outputs": {"ok": True}})], + success=True, + duration=0.01, + history_limit=3, + ) + + assert commit_count == 4 + assert trimmed == ["exec-retain-0"] + assert await WorkflowStore.get_execution("exec-retain-0") is None + old_steps, old_total = await WorkflowStore.list_steps("exec-retain-0") + assert old_steps == [] + assert old_total == 0 + executions = await WorkflowStore.list_executions("wf-retain", limit=10) + assert [execution["id"] for execution in executions] == [ + "exec-retain-3", + "exec-retain-2", + "exec-retain-1", + ] From 76f30996f8dfcc292f83242d94cb067a45720382 Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Fri, 21 Aug 2026 17:01:57 +0800 Subject: [PATCH 16/63] fix(workflow): preserve steps without callback waits Keep workflow callbacks storage-free while batching complete step history and serializing interactive progress writes. Isolate terminal persistence from stats and retention failures. Co-Authored-By: Claude Opus 4.6 --- flocks/ingest/kafka/manager.py | 1 - flocks/ingest/syslog/manager.py | 5 +- flocks/server/routes/workflow.py | 230 +++++----- flocks/tool/task/run_workflow.py | 258 ++++++----- flocks/workflow/execution_store.py | 213 ++++++---- flocks/workflow/poller_manager.py | 5 +- flocks/workflow/store.py | 121 ++---- flocks/workflow/triggers/runtime.py | 5 +- tests/ingest/test_kafka_manager.py | 10 +- .../test_syslog_manager_backpressure.py | 12 +- .../server/routes/test_workflow_run_route.py | 400 +++++++++++++++++- .../workflow/test_execution_store_compact.py | 346 +++++++++++---- tests/workflow/test_poller_manager.py | 12 +- tests/workflow/test_tool_run_workflow.py | 269 +++++++++++- tests/workflow/test_trigger_runtime.py | 12 +- tests/workflow/test_workflow_store.py | 93 ++-- 16 files changed, 1411 insertions(+), 581 deletions(-) diff --git a/flocks/ingest/kafka/manager.py b/flocks/ingest/kafka/manager.py index 893156c43..a44c28eec 100644 --- a/flocks/ingest/kafka/manager.py +++ b/flocks/ingest/kafka/manager.py @@ -775,7 +775,6 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: trigger_input_keys = list((trigger.mapping or {}).keys()) or [input_key] step_recorder = ExecutionStepRecorder( exec_id=exec_id, - capture_steps=False, step_compactor=lambda step: _compact_step_for_kafka_storage( step, input_key=input_key, diff --git a/flocks/ingest/syslog/manager.py b/flocks/ingest/syslog/manager.py index 699fa53cb..ab15ad23a 100644 --- a/flocks/ingest/syslog/manager.py +++ b/flocks/ingest/syslog/manager.py @@ -622,10 +622,7 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: persist=False, ) exec_id = exec_data["id"] - step_recorder = ExecutionStepRecorder( - exec_id=exec_id, - capture_steps=False, - ) + step_recorder = ExecutionStepRecorder(exec_id=exec_id) start_time = time.time() trigger_meta = mapped_inputs.get("_flocks", {}).get("trigger", {}) tool_context = None diff --git a/flocks/server/routes/workflow.py b/flocks/server/routes/workflow.py index 928b6a900..0a35c4eab 100644 --- a/flocks/server/routes/workflow.py +++ b/flocks/server/routes/workflow.py @@ -55,6 +55,8 @@ compact_step_for_storage, create_execution_record, derive_loop_progress, + ExecutionProgressWriter, + ExecutionStepRecorder, load_execution_steps, normalize_execution_status as _normalize_execution_status, record_execution_result as _record_execution_result, @@ -144,6 +146,7 @@ class ActiveWorkflowExecution: workflow_id: str task: asyncio.Task[Any] cancel_event: threading.Event + progress_writer: ExecutionProgressWriter _active_workflow_executions: Dict[str, ActiveWorkflowExecution] = {} @@ -1120,11 +1123,12 @@ async def _run_workflow_execution_task( req: WorkflowRunRequest, exec_id: str, cancel_event: threading.Event, + progress_writer: ExecutionProgressWriter, tool_context: Optional[ToolContext] = None, ) -> None: """Execute a workflow in the background and keep the execution record updated.""" start_time = time.time() - step_count = 0 + step_recorder = ExecutionStepRecorder(exec_id=exec_id) pending_step_index: Optional[int] = None pending_step: Optional[Dict[str, Any]] = None execution_summary: Dict[str, Any] = { @@ -1167,122 +1171,123 @@ def _on_step_start(_run_id, step_index, node, _inputs): "error": "Run cancelled before node completed", } ) - execution_summary.update( - { - "currentNodeId": node_id, - "currentNodeType": node_type, - "currentPhase": "running", - "currentStepIndex": step_index, - "loopProgress": loop_progress, - "updatedAt": int(time.time() * 1000), - } - ) + progress_update = { + "currentNodeId": node_id, + "currentNodeType": node_type, + "currentPhase": "cancelling" if cancel_event.is_set() else "running", + "currentStepIndex": step_index, + "loopProgress": loop_progress, + "updatedAt": int(time.time() * 1000), + } + execution_summary.update(progress_update) + progress_writer.submit(progress_update) return step_index def _on_step_complete(step_result) -> None: - nonlocal step_count, pending_step_index, pending_step - step_dict = compact_step_for_storage(step_result.model_dump(mode="json")) - step_count += 1 + nonlocal pending_step_index, pending_step + step_recorder.on_step_complete(step_result) pending_step_index = None pending_step = None - loop_progress = derive_loop_progress( - node_id=step_dict.get("node_id"), - global_step_index=step_count, - inputs=step_dict.get("inputs"), - outputs=step_dict.get("outputs"), - ) - execution_summary.update( - { - "stepCount": step_count, - "currentNodeId": step_dict.get("node_id"), - "currentNodeType": step_dict.get("node_type") or step_dict.get("type"), - "currentPhase": "running", - "currentStepIndex": step_count, - "loopProgress": loop_progress, - "updatedAt": int(time.time() * 1000), - } - ) - + progress_update = dict(step_recorder.summary) + if cancel_event.is_set(): + progress_update["currentPhase"] = "cancelling" + execution_summary.update(progress_update) + progress_writer.submit(progress_update) + + result: Optional[RunWorkflowResult] = None + execution_error: Optional[Exception] = None try: - result: RunWorkflowResult = await asyncio.to_thread( - run_workflow, - workflow=workflow_json, - inputs=req.inputs or {}, - timeout_s=req.timeout_s, - trace=req.trace, - on_step_start=_on_step_start, - on_step_complete=_on_step_complete, - cancel=cancel_event.is_set, - tool_context=tool_context, - ) + try: + result = await asyncio.to_thread( + run_workflow, + workflow=workflow_json, + inputs=req.inputs or {}, + timeout_s=req.timeout_s, + trace=req.trace, + on_step_start=_on_step_start, + on_step_complete=_on_step_complete, + cancel=cancel_event.is_set, + tool_context=tool_context, + ) + except Exception as exc: + execution_error = exc duration = time.time() - start_time current_data = dict(execution_summary) - status_value, error_message = _resolve_execution_outcome(result) - if cancel_event.is_set() and status_value == "success": - status_value = "cancelled" - error_message = error_message or f"Run cancelled: run_id={result.run_id or exec_id}" - # ``record_execution_result`` backfills this compacted history into - # append-only step rows, then stores only the summary row. - final_history = compact_history_for_storage(result.history) - final_steps = result.steps + final_step_batch = step_recorder.take_steps() if pending_step_index is not None and pending_step is not None: - final_history.append(pending_step) - final_steps = max(final_steps, pending_step_index) - current_data.update( - { - "outputResults": compact_outputs_for_storage(result.outputs), - "status": status_value, - "finishedAt": int(time.time() * 1000), - "duration": duration, - "executionLog": final_history, - "stepCount": final_steps, - "errorMessage": error_message, - "currentNodeId": result.last_node_id, - "currentNodeType": current_data.get("currentNodeType"), - "currentPhase": status_value, - "currentStepIndex": final_steps, - "updatedAt": int(time.time() * 1000), - } + final_step_batch.append((pending_step_index, pending_step)) + final_step_batch.sort(key=lambda item: item[0]) + final_steps = max( + max((step_index for step_index, _ in final_step_batch), default=0), + pending_step_index or 0, ) + final_history = [step for _, step in final_step_batch] + + if execution_error is None: + assert result is not None + status_value, error_message = _resolve_execution_outcome(result) + if cancel_event.is_set() and status_value == "success": + status_value = "cancelled" + error_message = error_message or f"Run cancelled: run_id={result.run_id or exec_id}" + final_steps = max(result.steps, final_steps) + current_data.update( + { + "outputResults": compact_outputs_for_storage(result.outputs), + "status": status_value, + "finishedAt": int(time.time() * 1000), + "duration": duration, + "executionLog": final_history, + "stepCount": final_steps, + "errorMessage": error_message, + "currentNodeId": result.last_node_id, + "currentNodeType": current_data.get("currentNodeType"), + "currentPhase": status_value, + "currentStepIndex": final_steps, + "updatedAt": int(time.time() * 1000), + } + ) + else: + current_data.update( + { + "status": "cancelled" if cancel_event.is_set() else "error", + "finishedAt": int(time.time() * 1000), + "duration": duration, + "errorMessage": str(execution_error), + "executionLog": final_history, + "stepCount": final_steps, + "currentPhase": "cancelled" if cancel_event.is_set() else "error", + "currentStepIndex": final_steps, + "updatedAt": int(time.time() * 1000), + } + ) - await _record_execution_result(workflow_id, exec_id, current_data) - log.info( - "workflow.executed", - { - "id": workflow_id, - "exec_id": exec_id, - "status": status_value, - "duration": duration, - }, - ) - except Exception as exc: - duration = time.time() - start_time - current_data = dict(execution_summary) - final_history = [pending_step] if pending_step is not None else [] - final_steps = max(step_count, pending_step_index or 0) - current_data.update( - { - "status": "cancelled" if cancel_event.is_set() else "error", - "finishedAt": int(time.time() * 1000), - "duration": duration, - "errorMessage": str(exc), - "executionLog": final_history, - "stepCount": final_steps, - "currentPhase": "cancelled" if cancel_event.is_set() else "error", - "currentStepIndex": final_steps, - "updatedAt": int(time.time() * 1000), - } - ) - await _record_execution_result(workflow_id, exec_id, current_data) - log.error( - "workflow.execute.error", - { - "id": workflow_id, - "exec_id": exec_id, - "error": str(exc), - }, + await progress_writer.close_and_drain() + await _record_execution_result( + workflow_id, + exec_id, + current_data, + steps=final_step_batch, ) + if execution_error is None: + log.info( + "workflow.executed", + { + "id": workflow_id, + "exec_id": exec_id, + "status": current_data["status"], + "duration": duration, + }, + ) + else: + log.error( + "workflow.execute.error", + { + "id": workflow_id, + "exec_id": exec_id, + "error": str(execution_error), + }, + ) finally: _active_workflow_executions.pop(exec_id, None) @@ -1692,6 +1697,7 @@ async def run_workflow_endpoint(workflow_id: str, req: WorkflowRunRequest): ) await WorkflowStore.upsert_execution(compact_execution_summary(exec_data)) exec_id = str(exec_data["id"]) + progress_writer = ExecutionProgressWriter(exec_data) cancel_event = threading.Event() task = asyncio.create_task( @@ -1701,6 +1707,7 @@ async def run_workflow_endpoint(workflow_id: str, req: WorkflowRunRequest): req=req, exec_id=exec_id, cancel_event=cancel_event, + progress_writer=progress_writer, tool_context=tool_context, ), name=f"workflow-run-{exec_id}", @@ -1709,6 +1716,7 @@ async def run_workflow_endpoint(workflow_id: str, req: WorkflowRunRequest): workflow_id=workflow_id, task=task, cancel_event=cancel_event, + progress_writer=progress_writer, ) # Guarantee cleanup of the registry entry even when the task is @@ -1757,13 +1765,13 @@ async def cancel_workflow_execution(workflow_id: str, exec_id: str): raise HTTPException(status_code=404, detail="Execution not found for this workflow") active.cancel_event.set() - exec_data.update( - { - "currentPhase": "cancelling", - "errorMessage": exec_data.get("errorMessage") or "Cancellation requested", - } - ) - await WorkflowStore.upsert_execution(exec_data) + progress_update = { + "currentPhase": "cancelling", + "errorMessage": exec_data.get("errorMessage") or "Cancellation requested", + "updatedAt": int(time.time() * 1000), + } + exec_data.update(progress_update) + await active.progress_writer.update(progress_update) log.info( "workflow.execution.cancel_requested", { diff --git a/flocks/tool/task/run_workflow.py b/flocks/tool/task/run_workflow.py index 9b78e5123..ab5fcbf10 100644 --- a/flocks/tool/task/run_workflow.py +++ b/flocks/tool/task/run_workflow.py @@ -11,32 +11,29 @@ import time from pathlib import Path from types import SimpleNamespace -from typing import Optional, Dict, Any, Union +from typing import Optional, Dict, Any, Union, List, Tuple from flocks.tool.registry import ToolRegistry, ToolCategory, ToolParameter, ParameterType, ToolResult, ToolContext from flocks.utils.log import Log from flocks.session.recorder import Recorder from flocks.workflow.execution_store import ( compact_history_for_storage, - compact_execution_summary, compact_outputs_for_storage, compact_step_for_storage, create_execution_record, derive_loop_progress, + ExecutionProgressWriter, + ExecutionStepRecorder, normalize_execution_status, - record_execution_step, record_execution_result, resolve_execution_outcome, ) from flocks.workflow.fs_store import read_workflow_from_fs, resolve_workflow_id_from_source -from flocks.workflow.store import WorkflowStore from flocks.tool.truncation import truncate_output log = Log.create(service="tool.run_workflow") -_PROGRESS_FLUSH_EVERY_STEPS = 5 - # Lazy import to avoid circular import (flocks.tool <-> flocks.workflow) _WORKFLOW_AVAILABLE: Optional[bool] = None RequirementsInstaller = None @@ -573,33 +570,17 @@ async def run_workflow_tool( canonical_workflow_id = registered_workflow_id or resolve_workflow_id_from_source(workflow_source) display_workflow_id = canonical_workflow_id or workflow_id tracked_execution: Optional[Dict[str, Any]] = None - tracked_step_count = 0 + step_recorder: Optional[ExecutionStepRecorder] = None + progress_writer: Optional[ExecutionProgressWriter] = None + callback_step_count = 0 pending_step_index: Optional[int] = None pending_step: Optional[Dict[str, Any]] = None + final_step_batch: Optional[List[Tuple[int, Dict[str, Any]]]] = None loop = asyncio.get_running_loop() def _emit_metadata(metadata: Dict[str, Any]) -> None: loop.call_soon_threadsafe(ctx.metadata, metadata) - def _update_execution_progress(update_fields: Dict[str, Any]) -> None: - try: - if tracked_execution is None: - return - tracked_execution.update(update_fields) - asyncio.run_coroutine_threadsafe( - WorkflowStore.upsert_execution(compact_execution_summary(tracked_execution)), - loop, - ).result(timeout=5) - except Exception as exc: - log.warning( - "run_workflow.execution_progress.write_failed", - { - "workflow_id": display_workflow_id, - "exec_id": tracked_execution["id"] if tracked_execution else None, - "error": str(exc), - }, - ) - def _on_step_start( _run_id: Optional[str], step_index: int, @@ -625,17 +606,19 @@ def _on_step_start( "error": "Run cancelled before node completed", } ) + current_phase = "cancelling" if ctx.abort.is_set() else "running" + progress_update = { + "currentNodeId": current_node_id, + "currentNodeType": current_node_type, + "currentPhase": current_phase, + "currentStepIndex": step_index, + "loopProgress": loop_progress, + "updatedAt": int(time.time() * 1000), + } if tracked_execution is not None: - _update_execution_progress( - { - "currentNodeId": current_node_id, - "currentNodeType": current_node_type, - "currentPhase": "running", - "currentStepIndex": step_index, - "loopProgress": loop_progress, - "updatedAt": int(time.time() * 1000), - } - ) + tracked_execution.update(progress_update) + if progress_writer is not None: + progress_writer.submit(progress_update) _emit_metadata( { "title": f"Running workflow: {workflow_name}", @@ -645,7 +628,7 @@ def _on_step_start( "total_nodes": workflow_total_nodes, "workflow_execution_id": tracked_execution["id"] if tracked_execution else None, "status": "running", - "phase": "running", + "phase": current_phase, "current_node_id": current_node_id, "current_node_type": current_node_type, "step_index": step_index, @@ -656,64 +639,43 @@ def _on_step_start( return step_index def _on_step_complete(step_result: Any) -> None: - nonlocal tracked_step_count, pending_step_index, pending_step - if hasattr(step_result, "model_dump"): - step_dict = step_result.model_dump(mode="json") - elif isinstance(step_result, dict): - step_dict = dict(step_result) + nonlocal callback_step_count, pending_step_index, pending_step + if step_recorder is not None: + step_recorder.on_step_complete(step_result) + callback_step_count = step_recorder.step_count + progress_update = dict(step_recorder.summary) else: - step_dict = {"node_id": None, "outputs": {}, "error": str(step_result)} - step_index = tracked_step_count + 1 - compacted_step = compact_step_for_storage(step_dict) + if hasattr(step_result, "model_dump"): + step_dict = step_result.model_dump(mode="json") + elif isinstance(step_result, dict): + step_dict = dict(step_result) + else: + step_dict = {"node_id": None, "outputs": {}, "error": str(step_result)} + callback_step_count += 1 + compacted_step = compact_step_for_storage(step_dict) + progress_update = { + "stepCount": callback_step_count, + "currentNodeId": compacted_step.get("node_id"), + "currentNodeType": compacted_step.get("node_type") + or compacted_step.get("type"), + "currentPhase": "running", + "currentStepIndex": callback_step_count, + "loopProgress": derive_loop_progress( + node_id=compacted_step.get("node_id"), + global_step_index=callback_step_count, + inputs=compacted_step.get("inputs"), + outputs=compacted_step.get("outputs"), + ), + "updatedAt": int(time.time() * 1000), + } pending_step_index = None pending_step = None - loop_progress = derive_loop_progress( - node_id=step_dict.get("node_id"), - global_step_index=step_index, - inputs=step_dict.get("inputs"), - outputs=step_dict.get("outputs"), - ) - tracked_step_count = step_index + if ctx.abort.is_set(): + progress_update["currentPhase"] = "cancelling" if tracked_execution is not None: - tracked_execution.update( - { - "stepCount": tracked_step_count, - "currentNodeId": step_dict.get("node_id"), - "currentNodeType": step_dict.get("node_type") or step_dict.get("type"), - "currentPhase": "running", - "currentStepIndex": tracked_step_count, - "loopProgress": loop_progress, - "updatedAt": int(time.time() * 1000), - } - ) - if tracked_execution is not None: - try: - asyncio.run_coroutine_threadsafe( - record_execution_step(tracked_execution["id"], step_index, compacted_step), - loop, - ).result(timeout=5) - except Exception as exc: - log.warning( - "run_workflow.execution_step.write_failed", - { - "workflow_id": display_workflow_id, - "exec_id": tracked_execution["id"], - "step_index": step_index, - "error": str(exc), - }, - ) - if tracked_step_count % _PROGRESS_FLUSH_EVERY_STEPS == 0: - _update_execution_progress( - { - "stepCount": tracked_step_count, - "currentNodeId": step_dict.get("node_id"), - "currentNodeType": step_dict.get("node_type") or step_dict.get("type"), - "currentPhase": "running", - "currentStepIndex": tracked_step_count, - "loopProgress": loop_progress, - "updatedAt": int(time.time() * 1000), - } - ) + tracked_execution.update(progress_update) + if progress_writer is not None: + progress_writer.submit(progress_update) _emit_metadata( { "title": f"Running workflow: {workflow_name}", @@ -723,36 +685,24 @@ def _on_step_complete(step_result: Any) -> None: "total_nodes": workflow_total_nodes, "workflow_execution_id": tracked_execution["id"] if tracked_execution else None, "status": "running", - "phase": "running", - "current_node_id": step_dict.get("node_id"), - "current_node_type": step_dict.get("node_type") or step_dict.get("type"), - "step_index": tracked_step_count, - "step_count": tracked_step_count, - "loop_progress": loop_progress, + "phase": progress_update["currentPhase"], + "current_node_id": progress_update.get("currentNodeId"), + "current_node_type": progress_update.get("currentNodeType"), + "step_index": callback_step_count, + "step_count": callback_step_count, + "loop_progress": progress_update.get("loopProgress"), }, } ) - return - async def _flush_pending_step() -> None: - if tracked_execution is None or pending_step_index is None or pending_step is None: - return - try: - await record_execution_step( - tracked_execution["id"], - pending_step_index, - pending_step, - ) - except Exception as exc: - log.warning( - "run_workflow.pending_step.write_failed", - { - "workflow_id": display_workflow_id, - "exec_id": tracked_execution["id"], - "step_index": pending_step_index, - "error": str(exc), - }, - ) + def _take_final_step_batch() -> List[Tuple[int, Dict[str, Any]]]: + nonlocal final_step_batch + if final_step_batch is None: + final_step_batch = step_recorder.take_steps() if step_recorder is not None else [] + if pending_step_index is not None and pending_step is not None: + final_step_batch.append((pending_step_index, pending_step)) + final_step_batch.sort(key=lambda item: item[0]) + return final_step_batch await ctx.ask( permission="run_workflow", @@ -771,6 +721,8 @@ async def _flush_pending_step() -> None: canonical_workflow_id, input_params=workflow_inputs, ) + step_recorder = ExecutionStepRecorder(exec_id=tracked_execution["id"]) + progress_writer = ExecutionProgressWriter(tracked_execution) # Update metadata to show workflow is running _emit_metadata( @@ -890,7 +842,8 @@ async def _flush_pending_step() -> None: result_dict = {"status": "UNKNOWN", "output": str(result)} status = result_dict.get("status", "UNKNOWN") - success = status == "SUCCEEDED" + status_value = normalize_execution_status(status) + success = status_value == "success" error = result_dict.get("error") output, output_truncated, output_path = _format_workflow_result_for_tool(result_dict) @@ -905,19 +858,25 @@ async def _flush_pending_step() -> None: }, ) - # Append-only recording for audit/replay - await _record_workflow_tool_result(display_workflow_id, result_dict) - - status_value = normalize_execution_status(status) compacted_history = compact_history_for_storage(result_dict.get("history")) - history_count = len(compacted_history) - if status_value == "cancelled" and not compacted_history: - await _flush_pending_step() + tracked_steps = _take_final_step_batch() if tracked_execution is not None else [] + final_history = ( + [step for _, step in tracked_steps] + if tracked_execution is not None + else compacted_history + ) + history_count = len(final_history) final_step_count = result_dict.get("steps") if not isinstance(final_step_count, int): - final_step_count = tracked_step_count - if pending_step_index is not None: - final_step_count = max(final_step_count, pending_step_index) + final_step_count = callback_step_count + final_step_count = max( + final_step_count, + max((step_index for step_index, _ in tracked_steps), default=0), + ) + + if tracked_execution is None: + await _record_workflow_tool_result(display_workflow_id, result_dict) + if tracked_execution and canonical_workflow_id: current_data = dict(tracked_execution) outcome_result = result @@ -928,13 +887,20 @@ async def _flush_pending_step() -> None: error=result_dict.get("error"), ) status_value, error_message = resolve_execution_outcome(outcome_result) # type: ignore[arg-type] + if ctx.abort.is_set() and status_value == "success": + status_value = "cancelled" + error_message = error_message or ( + f"Run cancelled: run_id={result_dict.get('run_id') or tracked_execution['id']}" + ) + error = error or error_message + success = status_value == "success" current_data.update( { "outputResults": compact_outputs_for_storage(result_dict.get("outputs")), "status": status_value, "finishedAt": int(time.time() * 1000), "duration": time.time() - execution_started_at, - "executionLog": compacted_history, + "executionLog": final_history, "stepCount": final_step_count, "errorMessage": error_message, "currentNodeId": result_dict.get("last_node_id"), @@ -943,10 +909,13 @@ async def _flush_pending_step() -> None: "updatedAt": int(time.time() * 1000), } ) + if progress_writer is not None: + await progress_writer.close_and_drain() await record_execution_result( canonical_workflow_id, tracked_execution["id"], current_data, + steps=tracked_steps, ) _emit_metadata( { @@ -979,7 +948,7 @@ async def _flush_pending_step() -> None: total_nodes=workflow_total_nodes, workflow_execution_id=tracked_execution["id"] if tracked_execution else None, status=status_value, - steps=result_dict.get("steps", 0), + steps=final_step_count, last_node_id=result_dict.get("last_node_id"), outputs=result_dict.get("outputs"), history_count=history_count, @@ -1018,24 +987,34 @@ async def _flush_pending_step() -> None: "error": error_msg, }, ) + terminal_status = "cancelled" if ctx.abort.is_set() else "error" + final_step_count = callback_step_count if tracked_execution and canonical_workflow_id: + tracked_steps = _take_final_step_batch() + final_step_count = max( + final_step_count, + max((step_index for step_index, _ in tracked_steps), default=0), + ) current_data = dict(tracked_execution) current_data.update( { - "status": "error", + "status": terminal_status, "finishedAt": int(time.time() * 1000), "errorMessage": error_msg, - "executionLog": [], - "stepCount": tracked_step_count, - "currentPhase": "error", - "currentStepIndex": tracked_step_count, + "executionLog": [step for _, step in tracked_steps], + "stepCount": final_step_count, + "currentPhase": terminal_status, + "currentStepIndex": final_step_count, "updatedAt": int(time.time() * 1000), } ) + if progress_writer is not None: + await progress_writer.close_and_drain() await record_execution_result( canonical_workflow_id, tracked_execution["id"], current_data, + steps=tracked_steps, ) _emit_metadata( { @@ -1045,9 +1024,9 @@ async def _flush_pending_step() -> None: "workflow_name": workflow_name, "total_nodes": workflow_total_nodes, "workflow_execution_id": tracked_execution["id"], - "status": "error", - "phase": "error", - "step_index": tracked_step_count, + "status": terminal_status, + "phase": terminal_status, + "step_index": final_step_count, }, } ) @@ -1061,6 +1040,7 @@ async def _flush_pending_step() -> None: workflow_name=workflow_name, total_nodes=workflow_total_nodes, workflow_execution_id=tracked_execution["id"] if tracked_execution else None, - status="FAILED", + status="CANCELLED" if terminal_status == "cancelled" else "FAILED", + steps=final_step_count, ), ) diff --git a/flocks/workflow/execution_store.py b/flocks/workflow/execution_store.py index c1c7cb443..3dbc59cd4 100644 --- a/flocks/workflow/execution_store.py +++ b/flocks/workflow/execution_store.py @@ -5,6 +5,7 @@ import asyncio from itertools import islice import sys +import threading import time import uuid from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple @@ -337,18 +338,6 @@ def derive_loop_progress( # Keep this intentionally small so high-frequency workflows do not keep # inflating the SQLite row set and matching JSONL audit files indefinitely. _MAX_EXECUTION_HISTORY_PER_WORKFLOW = 30 -# Per-workflow trim lock. Trims are awaited by the writer so the retention cap -# is enforced before ``record_execution_result`` returns, while concurrent runs -# for the same workflow serialize instead of skipping cleanup. -_trim_locks: Dict[str, asyncio.Lock] = {} - - -def _get_trim_lock(workflow_id: str) -> asyncio.Lock: - lock = _trim_locks.get(workflow_id) - if lock is None: - lock = asyncio.Lock() - _trim_locks[workflow_id] = lock - return lock def workflow_execution_key(exec_id: str) -> str: @@ -410,11 +399,9 @@ def __init__( self, *, exec_id: str, - capture_steps: bool = True, step_compactor: Callable[[Any], Dict[str, Any]] = compact_step_for_storage, ) -> None: self.exec_id = exec_id - self.capture_steps = capture_steps self.step_compactor = step_compactor self.step_count = 0 self.summary: Dict[str, Any] = {} @@ -444,8 +431,7 @@ def on_step_complete(self, step_result: Any) -> None: "updatedAt": int(time.time() * 1000), } ) - if self.capture_steps: - self._pending_steps.append((self.step_count, step_dict)) + self._pending_steps.append((self.step_count, step_dict)) def take_steps(self) -> List[Tuple[int, Dict[str, Any]]]: """Return buffered steps for the final execution transaction.""" @@ -454,6 +440,94 @@ def take_steps(self) -> List[Tuple[int, Dict[str, Any]]]: return pending_steps +class ExecutionProgressWriter: + """Coalesce nonblocking execution-summary updates onto one SQLite writer.""" + + def __init__(self, execution_summary: Dict[str, Any]) -> None: + self._loop = asyncio.get_running_loop() + self._summary = compact_execution_summary(execution_summary) + self._pending_summary: Optional[Dict[str, Any]] = None + self._pending_waiters: List[asyncio.Future[None]] = [] + self._writer_task: Optional[asyncio.Task[None]] = None + self._submission_lock = threading.Lock() + self._closed = False + + def submit(self, update: Dict[str, Any]) -> None: + """Queue an update from any thread without waiting for persistence.""" + with self._submission_lock: + if self._closed: + return + self._loop.call_soon_threadsafe(self._merge_update, dict(update), None) + + async def update(self, update: Dict[str, Any]) -> None: + """Queue and await an owner-loop update, preserving submission order.""" + if asyncio.get_running_loop() is not self._loop: + raise RuntimeError("ExecutionProgressWriter.update must run on its owner loop") + + waiter = self._loop.create_future() + with self._submission_lock: + if self._closed: + return + self._loop.call_soon(self._merge_update, dict(update), waiter) + await waiter + + async def close_and_drain(self) -> None: + """Reject new updates and flush every update accepted before closing.""" + if asyncio.get_running_loop() is not self._loop: + raise RuntimeError("ExecutionProgressWriter.close_and_drain must run on its owner loop") + + barrier = self._loop.create_future() + with self._submission_lock: + self._closed = True + self._loop.call_soon(barrier.set_result, None) + await barrier + + writer_task = self._writer_task + if writer_task is not None: + await asyncio.shield(writer_task) + + def _merge_update( + self, + update: Dict[str, Any], + waiter: Optional[asyncio.Future[None]], + ) -> None: + self._summary.update(update) + self._pending_summary = compact_execution_summary(self._summary) + if waiter is not None: + self._pending_waiters.append(waiter) + if self._writer_task is None: + exec_id = str(self._summary.get("id") or "unknown") + self._writer_task = self._loop.create_task( + self._flush(), + name=f"workflow-progress-{exec_id}", + ) + + async def _flush(self) -> None: + try: + while self._pending_summary is not None: + summary = self._pending_summary + waiters = self._pending_waiters + self._pending_summary = None + self._pending_waiters = [] + try: + await WorkflowStore.upsert_execution(summary) + except Exception as exc: + log.warning( + "workflow.progress.update_failed", + { + "workflow_id": summary.get("workflowId"), + "exec_id": summary.get("id"), + "error": str(exc), + }, + ) + finally: + for waiter in waiters: + if not waiter.done(): + waiter.set_result(None) + finally: + self._writer_task = None + + def _prepare_execution_steps(execution_log: Any) -> List[Tuple[int, Dict[str, Any]]]: """Compact an inline execution log for one final batch transaction.""" if not isinstance(execution_log, list): @@ -594,15 +668,48 @@ async def record_execution_result( started_at = summary_data.get("startedAt", 0) finished_at = summary_data.get("finishedAt", int(time.time() * 1000)) duration = max(0.0, (finished_at - started_at) / 1000.0) - trimmed_exec_ids = await WorkflowStore.complete_execution( + + await WorkflowStore.complete_execution( compact_execution_summary(summary_data), prepared_steps, - success=success, - duration=float(duration), - history_limit=_MAX_EXECUTION_HISTORY_PER_WORKFLOW, ) - if not isinstance(trimmed_exec_ids, list): - trimmed_exec_ids = [] + + try: + await WorkflowStore.increment_stats( + workflow_id, + success=success, + duration=float(duration), + ) + except Exception as exc: + log.warning( + "workflow.stats.update_failed", + { + "workflow_id": workflow_id, + "exec_id": exec_id, + "error": str(exc), + }, + ) + + trimmed_exec_ids: List[str] = [] + try: + trimmed_exec_ids = await WorkflowStore.trim_executions( + workflow_id, + keep=_MAX_EXECUTION_HISTORY_PER_WORKFLOW, + ) + except Exception as exc: + log.error( + "workflow.history.trim_failed", + { + "workflow_id": workflow_id, + "exec_id": exec_id, + "error": str(exc), + }, + ) + + audit_data = dict(exec_data) + audit_data["executionLog"] = [ + step for _, step in sorted(prepared_steps, key=lambda item: item[0]) + ] # Recorder writes to its own SQLite tables and can be slow under load. # Run it as a background task so the syslog/HTTP dispatcher can release the @@ -614,7 +721,7 @@ async def _record_audit() -> None: await Recorder.record_workflow_execution( exec_id=exec_id, workflow_id=workflow_id, - run_result=exec_data, + run_result=audit_data, ) except Exception as exc: log.debug( @@ -645,65 +752,7 @@ async def _record_audit() -> None: await Recorder.record_workflow_execution( exec_id=exec_id, workflow_id=workflow_id, - run_result=exec_data, + run_result=audit_data, ) except Exception: pass - - -async def _delete_execution_history_record( - execution_key: str, - *, - index_key: Optional[str] = None, -) -> None: - exec_id = execution_key.rsplit("/", 1)[-1] - deleted_steps = await WorkflowStore.clear_steps(exec_id) - removed_execution = await WorkflowStore.delete_execution(exec_id) - record_path = Recorder.paths().workflow_dir / f"{exec_id}.jsonl" - await asyncio.to_thread(record_path.unlink, missing_ok=True) - log.debug( - "workflow.history.trim_deleted", - { - "exec_id": exec_id, - "execution_key": execution_key, - "steps": deleted_steps, - "removed_execution": removed_execution, - }, - ) - - -async def _trim_execution_history(workflow_id: str) -> None: - """Delete the oldest execution records once the per-workflow cap is exceeded. - - New records carry a per-workflow ``workflow_execution_index`` key, so hot - trims avoid scanning unrelated workflows. This path is intentionally - index-only: if an old execution has no index key, it is outside the hot - retention path and should be handled by a separate migration/GC task. - - A per-workflow lock serializes concurrent trims. Cleanup is awaited by - ``record_execution_result`` so the retention cap is enforced synchronously - instead of being an opportunistic background task. - """ - lock = _get_trim_lock(workflow_id) - async with lock: - failures: List[str] = [] - for exec_id in await WorkflowStore.trim_executions( - workflow_id, - keep=_MAX_EXECUTION_HISTORY_PER_WORKFLOW, - ): - try: - record_path = Recorder.paths().workflow_dir / f"{exec_id}.jsonl" - await asyncio.to_thread(record_path.unlink, missing_ok=True) - except Exception as exc: - failures.append(f"{exec_id}: {exc}") - log.warning( - "workflow.history.trim_delete_failed", - { - "workflow_id": workflow_id, - "exec_id": exec_id, - "error": str(exc), - }, - ) - - if failures: - raise RuntimeError("Failed to trim workflow execution history: " + "; ".join(failures[:3])) diff --git a/flocks/workflow/poller_manager.py b/flocks/workflow/poller_manager.py index d596daa80..7fdb7a101 100644 --- a/flocks/workflow/poller_manager.py +++ b/flocks/workflow/poller_manager.py @@ -454,10 +454,7 @@ async def _execute_run( persist=False, ) exec_id = str(exec_data["id"]) - step_recorder = ExecutionStepRecorder( - exec_id=exec_id, - capture_steps=False, - ) + step_recorder = ExecutionStepRecorder(exec_id=exec_id) current = self._status.get(workflow_id) or self._base_status(workflow_id) current["lastRunAt"] = started_at_ms current["activeRuns"] = self._cleanup_done_runs(workflow_id) diff --git a/flocks/workflow/store.py b/flocks/workflow/store.py index db049f3a2..b4d9b129f 100644 --- a/flocks/workflow/store.py +++ b/flocks/workflow/store.py @@ -158,6 +158,8 @@ async def raw_db(cls) -> aiosqlite.Connection: @classmethod async def _completion_db(cls) -> aiosqlite.Connection: + if cls._initialized and cls._init_pid is not None and cls._init_pid != os.getpid(): + await cls.init() if not cls._completion_conn or not cls._initialized: await cls.init() return cls._completion_conn # type: ignore[return-value] @@ -439,21 +441,12 @@ async def trim_executions(cls, workflow_id: str, *, keep: int) -> List[str]: return exec_ids @classmethod - async def record_step( - cls, - exec_id: str, - step_index: int, - step_payload: Dict[str, Any], - ) -> None: - await cls.record_steps(exec_id, [(step_index, step_payload)]) - - @classmethod - async def record_steps( + def _step_rows( cls, exec_id: str, steps: Iterable[Tuple[int, Dict[str, Any]]], - ) -> None: - rows = [ + ) -> List[Tuple[Any, ...]]: + return [ ( exec_id, int(step_index), @@ -466,6 +459,23 @@ async def record_steps( ) for step_index, step_payload in steps ] + + @classmethod + async def record_step( + cls, + exec_id: str, + step_index: int, + step_payload: Dict[str, Any], + ) -> None: + await cls.record_steps(exec_id, [(step_index, step_payload)]) + + @classmethod + async def record_steps( + cls, + exec_id: str, + steps: Iterable[Tuple[int, Dict[str, Any]]], + ) -> None: + rows = cls._step_rows(exec_id, steps) if not rows: return db = await cls._db() @@ -484,12 +494,8 @@ async def complete_execution( cls, exec_data: Dict[str, Any], steps: Iterable[Tuple[int, Dict[str, Any]]], - *, - success: bool, - duration: float, - history_limit: Optional[int] = None, - ) -> List[str]: - """Persist one completed execution, stats, and retention in one transaction.""" + ) -> None: + """Atomically persist one final execution summary and its step batch.""" db = await cls._completion_db() payload = dict(exec_data) exec_id = str(payload.get("id") or "") @@ -497,22 +503,7 @@ async def complete_execution( if not exec_id or not workflow_id: raise ValueError("workflow execution requires id and workflowId") - step_rows = [ - ( - exec_id, - int(step_index), - step_payload.get("node_id"), - step_payload.get("node_type") or step_payload.get("type"), - cls._json_dumps(step_payload.get("inputs") or {}), - cls._json_dumps(step_payload.get("outputs") or {}), - step_payload.get("error"), - cls._json_dumps(step_payload), - ) - for step_index, step_payload in steps - ] - runtime = float(duration) - success_delta = 1 if success else 0 - error_delta = 0 if success else 1 + step_rows = cls._step_rows(exec_id, steps) lock = cls._completion_lock if lock is None: lock = asyncio.Lock() @@ -559,57 +550,19 @@ async def complete_execution( cls._json_dumps(payload), ), ) - await db.execute( - """ - INSERT INTO workflow_stats ( - workflow_id, call_count, success_count, error_count, - total_runtime, avg_runtime, thumbs_up, thumbs_down, updated_at - ) - VALUES (?, 1, ?, ?, ?, ?, 0, 0, ?) - ON CONFLICT(workflow_id) DO UPDATE SET - call_count = workflow_stats.call_count + 1, - success_count = workflow_stats.success_count + excluded.success_count, - error_count = workflow_stats.error_count + excluded.error_count, - total_runtime = workflow_stats.total_runtime + excluded.total_runtime, - avg_runtime = ( - workflow_stats.total_runtime + excluded.total_runtime - ) / (workflow_stats.call_count + 1), - updated_at = excluded.updated_at - """, - ( - workflow_id, - success_delta, - error_delta, - runtime, - runtime, - cls._now_ms(), - ), - ) - trimmed_exec_ids: List[str] = [] - if history_limit is not None: - async with db.execute( - """ - SELECT id FROM workflow_executions - WHERE workflow_id = ? - ORDER BY started_at DESC, rowid DESC - LIMIT -1 OFFSET ? - """, - (workflow_id, max(int(history_limit), 0)), - ) as cur: - trimmed_exec_ids = [str(row["id"]) for row in await cur.fetchall()] - for trimmed_exec_id in trimmed_exec_ids: - await db.execute( - "DELETE FROM workflow_execution_steps WHERE exec_id = ?", - (trimmed_exec_id,), - ) - await db.execute( - "DELETE FROM workflow_executions WHERE id = ?", - (trimmed_exec_id,), - ) await db.commit() - return trimmed_exec_ids - except Exception: - await db.rollback() + except BaseException: + try: + await db.rollback() + except BaseException as rollback_exc: + log.error( + "workflow.store.completion_rollback_failed", + { + "workflow_id": workflow_id, + "exec_id": exec_id, + "error": str(rollback_exc), + }, + ) raise @classmethod diff --git a/flocks/workflow/triggers/runtime.py b/flocks/workflow/triggers/runtime.py index 2a3d2aedf..d41aa012d 100644 --- a/flocks/workflow/triggers/runtime.py +++ b/flocks/workflow/triggers/runtime.py @@ -251,10 +251,7 @@ async def _execute_workflow_effect( persist=False, ) exec_id = exec_data["id"] - step_recorder = ExecutionStepRecorder( - exec_id=exec_id, - capture_steps=False, - ) + step_recorder = ExecutionStepRecorder(exec_id=exec_id) started_at = time.time() tool_context = None try: diff --git a/tests/ingest/test_kafka_manager.py b/tests/ingest/test_kafka_manager.py index 6e494b566..2134fa9c9 100644 --- a/tests/ingest/test_kafka_manager.py +++ b/tests/ingest/test_kafka_manager.py @@ -616,7 +616,15 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 } assert captured_exec_data["executionLog"] == [] assert captured_exec_data["stepCount"] == 2 - assert captured_steps == [] + assert [step_index for step_index, _ in captured_steps] == [1, 2] + assert [step["node_id"] for _, step in captured_steps] == [ + "receive_alert", + "dedup_and_write", + ] + assert captured_steps[0][1]["inputs"]["kafka_message"]["_type"] == "dict" + assert captured_steps[0][1]["outputs"] == {"_raw_alerts_count": 1} + assert captured_steps[1][1]["inputs"] == {"_filtered_alerts_count": 1} + assert captured_steps[1][1]["outputs"] == {"_enriched_alerts_count": 1} assert len(json.dumps(captured_exec_data, ensure_ascii=False)) < 10_000 diff --git a/tests/ingest/test_syslog_manager_backpressure.py b/tests/ingest/test_syslog_manager_backpressure.py index 3ad929893..0ce283747 100644 --- a/tests/ingest/test_syslog_manager_backpressure.py +++ b/tests/ingest/test_syslog_manager_backpressure.py @@ -427,7 +427,17 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 action_name="trigger:syslog", ) trigger_tool_context.cleanup.assert_awaited_once_with(trigger_tool_context.context) - assert recorded_steps == [] + assert recorded_steps == [ + ( + 1, + { + "node_id": "receive_alert", + "node_type": "python", + "inputs": {"message": "demo"}, + "outputs": {"ok": True}, + }, + ) + ] assert recorded_exec_data["triggerId"] == "syslog-alerts" assert recorded_exec_data["triggerSource"] == "udp://0.0.0.0:5514" assert recorded_exec_data["executionLog"] == [] diff --git a/tests/server/routes/test_workflow_run_route.py b/tests/server/routes/test_workflow_run_route.py index 5684a5cef..f0891bafa 100644 --- a/tests/server/routes/test_workflow_run_route.py +++ b/tests/server/routes/test_workflow_run_route.py @@ -1,3 +1,4 @@ +import asyncio from types import SimpleNamespace from unittest.mock import AsyncMock, Mock @@ -243,7 +244,7 @@ def run_workflow_mock(**kwargs): kwargs["on_step_complete"](step_result) return SimpleNamespace( outputs={"ok": True}, - history=[step_result.model_dump(mode="json")], + history=[], last_node_id="node-1", steps=1, ) @@ -272,6 +273,14 @@ def run_workflow_mock(**kwargs): req = workflow_module.WorkflowRunRequest(inputs={"ip": "8.8.8.8"}, trace=False) tool_context = ToolContext(session_id="session-1", message_id="message-1", agent="rex") + progress_writer = workflow_module.ExecutionProgressWriter( + { + "id": "exec-1", + "workflowId": "wf-1", + "status": "running", + "executionLog": [], + } + ) await workflow_module._run_workflow_execution_task( workflow_id="wf-1", @@ -279,22 +288,25 @@ def run_workflow_mock(**kwargs): req=req, exec_id="exec-1", cancel_event=workflow_module.threading.Event(), + progress_writer=progress_writer, tool_context=tool_context, ) init_mock.assert_not_awaited() run_mock.assert_called_once() assert run_mock.call_args.kwargs["tool_context"] is tool_context - upsert_execution.assert_not_awaited() + assert upsert_execution.await_count >= 1 + assert all(call.args[0]["executionLog"] == [] for call in upsert_execution.await_args_list) + assert upsert_execution.await_args.args[0]["currentNodeId"] == "node-1" record_result.assert_awaited_once() - assert record_result.await_args.args[2]["executionLog"] == [ - { - "node_id": "node-1", - "node_type": "tool", - "inputs": {}, - "outputs": {"ok": True}, - } - ] + expected_step = { + "node_id": "node-1", + "node_type": "tool", + "inputs": {}, + "outputs": {"ok": True}, + } + assert record_result.await_args.args[2]["executionLog"] == [expected_step] + assert record_result.await_args.kwargs["steps"] == [(1, expected_step)] @pytest.mark.asyncio @@ -320,34 +332,392 @@ def run_workflow_mock(**kwargs): monkeypatch.setattr(workflow_module, "run_workflow", Mock(side_effect=run_workflow_mock)) monkeypatch.setattr(workflow_module, "_resolve_execution_outcome", lambda _result: ("success", None)) monkeypatch.setattr(workflow_module, "_record_execution_result", record_result) + monkeypatch.setattr( + workflow_module.WorkflowStore, + "upsert_execution", + AsyncMock(return_value=None), + ) monkeypatch.setattr(workflow_module, "compact_outputs_for_storage", lambda value: value) monkeypatch.setattr(workflow_module, "compact_history_for_storage", lambda value: value) cancel_event = workflow_module.threading.Event() cancel_event.set() + progress_writer = workflow_module.ExecutionProgressWriter( + { + "id": "exec-cancelled", + "workflowId": "wf-1", + "status": "running", + "executionLog": [], + } + ) await workflow_module._run_workflow_execution_task( workflow_id="wf-1", workflow_json={"id": "wf-1", "start": "node-1", "nodes": [], "edges": []}, req=workflow_module.WorkflowRunRequest(inputs={"message": "hello"}, trace=False), exec_id="exec-cancelled", cancel_event=cancel_event, + progress_writer=progress_writer, ) record_result.assert_awaited_once() final_data = record_result.await_args.args[2] assert final_data["status"] == "cancelled" assert final_data["stepCount"] == 1 - assert final_data["executionLog"] == [ + pending_step = { + "node_id": "node-1", + "node_type": "tool", + "inputs": {"message": "hello"}, + "outputs": {}, + "error": "Run cancelled before node completed", + } + assert final_data["executionLog"] == [pending_step] + assert record_result.await_args.kwargs["steps"] == [(1, pending_step)] + + +@pytest.mark.asyncio +async def test_run_workflow_execution_task_keeps_completed_and_pending_step_indices( + monkeypatch: pytest.MonkeyPatch, +) -> None: + cancel_event = workflow_module.threading.Event() + completed_step = SimpleNamespace( + model_dump=lambda mode: { + "node_id": "node-1", + "node_type": "python", + "inputs": {"value": 1}, + "outputs": {"value": 2}, + } + ) + + def run_workflow_mock(**kwargs): + kwargs["on_step_start"]( + "run-1", + 1, + SimpleNamespace(id="node-1", type="python"), + {"value": 1}, + ) + kwargs["on_step_complete"](completed_step) + cancel_event.set() + kwargs["on_step_start"]( + "run-1", + 2, + SimpleNamespace(id="node-2", type="tool"), + {"message": "hello"}, + ) + return SimpleNamespace( + run_id="run-1", + outputs={"value": 2}, + history=[], + last_node_id="node-2", + steps=1, + ) + + record_result = AsyncMock(return_value=None) + upsert_execution = AsyncMock(return_value=None) + monkeypatch.setattr(workflow_module, "run_workflow", Mock(side_effect=run_workflow_mock)) + monkeypatch.setattr(workflow_module, "_resolve_execution_outcome", lambda _result: ("success", None)) + monkeypatch.setattr(workflow_module, "_record_execution_result", record_result) + monkeypatch.setattr(workflow_module.WorkflowStore, "upsert_execution", upsert_execution) + + progress_writer = workflow_module.ExecutionProgressWriter( { + "id": "exec-partial-cancel", + "workflowId": "wf-1", + "status": "running", + "executionLog": [], + } + ) + await workflow_module._run_workflow_execution_task( + workflow_id="wf-1", + workflow_json={"id": "wf-1", "start": "node-1", "nodes": [], "edges": []}, + req=workflow_module.WorkflowRunRequest(inputs={"value": 1}, trace=False), + exec_id="exec-partial-cancel", + cancel_event=cancel_event, + progress_writer=progress_writer, + ) + + steps = record_result.await_args.kwargs["steps"] + assert [step_index for step_index, _ in steps] == [1, 2] + assert [step["node_id"] for _, step in steps] == ["node-1", "node-2"] + assert steps[1][1]["error"] == "Run cancelled before node completed" + assert record_result.await_args.args[2]["status"] == "cancelled" + assert upsert_execution.await_args.args[0]["currentPhase"] == "cancelling" + + +@pytest.mark.asyncio +async def test_run_workflow_execution_task_keeps_steps_when_runner_raises( + monkeypatch: pytest.MonkeyPatch, +) -> None: + completed_step = SimpleNamespace( + model_dump=lambda mode: { "node_id": "node-1", - "node_type": "tool", - "inputs": {"message": "hello"}, - "outputs": {}, - "error": "Run cancelled before node completed", + "node_type": "python", + "inputs": {}, + "outputs": {"ok": True}, + } + ) + + def run_workflow_mock(**kwargs): + kwargs["on_step_start"]( + "run-1", + 1, + SimpleNamespace(id="node-1", type="python"), + {}, + ) + kwargs["on_step_complete"](completed_step) + kwargs["on_step_start"]( + "run-1", + 2, + SimpleNamespace(id="node-2", type="tool"), + {"message": "hello"}, + ) + raise RuntimeError("runner failed") + + record_result = AsyncMock(return_value=None) + monkeypatch.setattr(workflow_module, "run_workflow", Mock(side_effect=run_workflow_mock)) + monkeypatch.setattr(workflow_module, "_record_execution_result", record_result) + monkeypatch.setattr( + workflow_module.WorkflowStore, + "upsert_execution", + AsyncMock(return_value=None), + ) + + progress_writer = workflow_module.ExecutionProgressWriter( + { + "id": "exec-runner-error", + "workflowId": "wf-1", + "status": "running", + "executionLog": [], + } + ) + await workflow_module._run_workflow_execution_task( + workflow_id="wf-1", + workflow_json={"id": "wf-1", "start": "node-1", "nodes": [], "edges": []}, + req=workflow_module.WorkflowRunRequest(inputs={}, trace=False), + exec_id="exec-runner-error", + cancel_event=workflow_module.threading.Event(), + progress_writer=progress_writer, + ) + + final_data = record_result.await_args.args[2] + steps = record_result.await_args.kwargs["steps"] + assert final_data["status"] == "error" + assert final_data["errorMessage"] == "runner failed" + assert [step_index for step_index, _ in steps] == [1, 2] + assert [step["node_id"] for _, step in steps] == ["node-1", "node-2"] + + +@pytest.mark.asyncio +async def test_run_workflow_execution_task_does_not_reclassify_persistence_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + step_result = SimpleNamespace( + model_dump=lambda mode: { + "node_id": "node-1", + "node_type": "python", + "inputs": {}, + "outputs": {"ok": True}, + } + ) + + def run_workflow_mock(**kwargs): + kwargs["on_step_start"]( + "run-1", + 1, + SimpleNamespace(id="node-1", type="python"), + {}, + ) + kwargs["on_step_complete"](step_result) + return SimpleNamespace( + outputs={"ok": True}, + history=[], + last_node_id="node-1", + steps=1, + ) + + record_result = AsyncMock(side_effect=RuntimeError("storage failed")) + monkeypatch.setattr(workflow_module, "run_workflow", Mock(side_effect=run_workflow_mock)) + monkeypatch.setattr(workflow_module, "_resolve_execution_outcome", lambda _result: ("success", None)) + monkeypatch.setattr(workflow_module, "_record_execution_result", record_result) + monkeypatch.setattr( + workflow_module.WorkflowStore, + "upsert_execution", + AsyncMock(return_value=None), + ) + + progress_writer = workflow_module.ExecutionProgressWriter( + { + "id": "exec-storage-error", + "workflowId": "wf-1", + "status": "running", + "executionLog": [], } + ) + + with pytest.raises(RuntimeError, match="storage failed"): + await workflow_module._run_workflow_execution_task( + workflow_id="wf-1", + workflow_json={"id": "wf-1", "start": "node-1", "nodes": [], "edges": []}, + req=workflow_module.WorkflowRunRequest(inputs={}, trace=False), + exec_id="exec-storage-error", + cancel_event=workflow_module.threading.Event(), + progress_writer=progress_writer, + ) + + record_result.assert_awaited_once() + final_data = record_result.await_args.args[2] + assert final_data["status"] == "success" + assert record_result.await_args.kwargs["steps"] == [ + ( + 1, + { + "node_id": "node-1", + "node_type": "python", + "inputs": {}, + "outputs": {"ok": True}, + }, + ) ] +@pytest.mark.asyncio +async def test_run_workflow_callbacks_do_not_wait_for_blocked_progress_write( + monkeypatch: pytest.MonkeyPatch, +) -> None: + write_started = asyncio.Event() + release_write = asyncio.Event() + runner_finished = workflow_module.threading.Event() + write_order: list[str] = [] + step_result = SimpleNamespace( + model_dump=lambda mode: { + "node_id": "node-1", + "node_type": "python", + "inputs": {}, + "outputs": {"ok": True}, + } + ) + + async def blocked_upsert(_summary): + write_order.append("progress-start") + write_started.set() + await release_write.wait() + write_order.append("progress-end") + + async def record_result(*args, **kwargs): # noqa: ANN002, ANN003 + write_order.append("final") + + def run_workflow_mock(**kwargs): + kwargs["on_step_start"]( + "run-1", + 1, + SimpleNamespace(id="node-1", type="python"), + {}, + ) + kwargs["on_step_complete"](step_result) + runner_finished.set() + return SimpleNamespace( + outputs={"ok": True}, + history=[], + last_node_id="node-1", + steps=1, + ) + + record_result_mock = AsyncMock(side_effect=record_result) + monkeypatch.setattr(workflow_module, "run_workflow", Mock(side_effect=run_workflow_mock)) + monkeypatch.setattr(workflow_module, "_resolve_execution_outcome", lambda _result: ("success", None)) + monkeypatch.setattr(workflow_module, "_record_execution_result", record_result_mock) + monkeypatch.setattr(workflow_module.WorkflowStore, "upsert_execution", blocked_upsert) + + progress_writer = workflow_module.ExecutionProgressWriter( + { + "id": "exec-blocked-progress", + "workflowId": "wf-1", + "status": "running", + "executionLog": [], + } + ) + task = asyncio.create_task( + workflow_module._run_workflow_execution_task( + workflow_id="wf-1", + workflow_json={"id": "wf-1", "start": "node-1", "nodes": [], "edges": []}, + req=workflow_module.WorkflowRunRequest(inputs={}, trace=False), + exec_id="exec-blocked-progress", + cancel_event=workflow_module.threading.Event(), + progress_writer=progress_writer, + ) + ) + + await write_started.wait() + assert await asyncio.to_thread(runner_finished.wait, 0.1) + record_result_mock.assert_not_awaited() + release_write.set() + await task + + record_result_mock.assert_awaited_once() + assert write_order[-2:] == ["progress-end", "final"] + + +@pytest.mark.asyncio +async def test_cancel_workflow_execution_uses_active_progress_writer( + monkeypatch: pytest.MonkeyPatch, +) -> None: + persisted_summaries: list[dict] = [] + + async def capture_upsert(summary): + persisted_summaries.append(dict(summary)) + + monkeypatch.setattr( + workflow_module.WorkflowStore, + "get_execution", + AsyncMock( + return_value={ + "id": "exec-cancel-route", + "workflowId": "wf-1", + "status": "running", + "currentPhase": "running", + "executionLog": [], + } + ), + ) + monkeypatch.setattr(workflow_module.WorkflowStore, "upsert_execution", capture_upsert) + + progress_writer = workflow_module.ExecutionProgressWriter( + { + "id": "exec-cancel-route", + "workflowId": "wf-1", + "status": "running", + "currentPhase": "queued", + "executionLog": [], + } + ) + progress_writer.submit({"currentPhase": "running", "currentNodeId": "node-1"}) + cancel_event = workflow_module.threading.Event() + current_task = asyncio.current_task() + assert current_task is not None + workflow_module._active_workflow_executions["exec-cancel-route"] = ( + workflow_module.ActiveWorkflowExecution( + workflow_id="wf-1", + task=current_task, + cancel_event=cancel_event, + progress_writer=progress_writer, + ) + ) + + try: + response = await workflow_module.cancel_workflow_execution( + "wf-1", + "exec-cancel-route", + ) + await progress_writer.close_and_drain() + finally: + workflow_module._active_workflow_executions.pop("exec-cancel-route", None) + + assert response["status"] == "accepted" + assert cancel_event.is_set() + assert persisted_summaries[-1]["currentPhase"] == "cancelling" + assert persisted_summaries[-1]["currentNodeId"] == "node-1" + assert persisted_summaries[-1]["errorMessage"] == "Cancellation requested" + + @pytest.mark.asyncio async def test_workflow_tool_context_preserves_current_opaque_extension_context( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/workflow/test_execution_store_compact.py b/tests/workflow/test_execution_store_compact.py index 9245eea3a..0ceb57849 100644 --- a/tests/workflow/test_execution_store_compact.py +++ b/tests/workflow/test_execution_store_compact.py @@ -18,6 +18,7 @@ from __future__ import annotations import asyncio +from types import SimpleNamespace from typing import Any, Dict, List from unittest.mock import AsyncMock, patch @@ -27,8 +28,8 @@ DEFAULT_GENERIC_SEQUENCE_THRESHOLD, DEFAULT_LARGE_LIST_KEYS, DEFAULT_MAX_INLINE_COLLECTION_BYTES, + ExecutionProgressWriter, ExecutionStepRecorder, - _trim_execution_history, compact_history_for_storage, compact_execution_summary, compact_outputs_for_storage, @@ -371,9 +372,130 @@ def _run_seven_steps(recorder: ExecutionStepRecorder) -> None: record_steps.assert_not_awaited() +@pytest.mark.asyncio +async def test_progress_writer_submits_without_waiting_and_coalesces_updates() -> None: + write_started = asyncio.Event() + release_write = asyncio.Event() + writes: List[Dict[str, Any]] = [] + active_writes = 0 + max_active_writes = 0 + + async def blocked_upsert(summary: Dict[str, Any]) -> None: + nonlocal active_writes, max_active_writes + active_writes += 1 + max_active_writes = max(max_active_writes, active_writes) + writes.append(dict(summary)) + try: + if len(writes) == 1: + write_started.set() + await release_write.wait() + finally: + active_writes -= 1 + + writer = ExecutionProgressWriter( + { + "id": "exec-progress", + "workflowId": "wf-progress", + "status": "running", + "executionLog": [{"node_id": "ignored"}], + } + ) + + with patch.object(WorkflowStore, "upsert_execution", side_effect=blocked_upsert): + writer.submit({"currentNodeId": "node-1", "currentStepIndex": 1}) + await write_started.wait() + await asyncio.wait_for( + asyncio.to_thread( + writer.submit, + {"currentNodeId": "node-2", "currentStepIndex": 2}, + ), + timeout=0.1, + ) + await asyncio.to_thread( + writer.submit, + {"currentNodeId": "node-3", "currentStepIndex": 3}, + ) + release_write.set() + await writer.close_and_drain() + + assert max_active_writes == 1 + assert len(writes) == 2 + assert writes[0]["currentNodeId"] == "node-1" + assert writes[-1]["currentNodeId"] == "node-3" + assert writes[-1]["currentStepIndex"] == 3 + assert writes[-1]["executionLog"] == [] + + +@pytest.mark.asyncio +async def test_progress_writer_awaited_update_is_ordered_and_close_rejects_late_updates() -> None: + writes: List[Dict[str, Any]] = [] + + async def capture_upsert(summary: Dict[str, Any]) -> None: + writes.append(dict(summary)) + + writer = ExecutionProgressWriter( + { + "id": "exec-cancelling", + "workflowId": "wf-cancelling", + "status": "running", + "currentPhase": "queued", + "executionLog": [], + } + ) + + with patch.object(WorkflowStore, "upsert_execution", side_effect=capture_upsert): + await asyncio.to_thread( + writer.submit, + {"currentPhase": "running", "currentNodeId": "node-1"}, + ) + await writer.update({"currentPhase": "cancelling"}) + await writer.close_and_drain() + writer.submit({"currentPhase": "running", "currentNodeId": "late-node"}) + await asyncio.sleep(0) + + assert writes[-1]["currentPhase"] == "cancelling" + assert writes[-1]["currentNodeId"] == "node-1" + assert all(write.get("currentNodeId") != "late-node" for write in writes) + + +@pytest.mark.asyncio +async def test_progress_writer_logs_write_failures_without_raising() -> None: + writer = ExecutionProgressWriter( + { + "id": "exec-write-failure", + "workflowId": "wf-write-failure", + "status": "running", + "executionLog": [], + } + ) + + with patch.object( + WorkflowStore, + "upsert_execution", + AsyncMock(side_effect=RuntimeError("database locked")), + ): + await writer.update({"currentPhase": "running"}) + await writer.close_and_drain() + + @pytest.mark.asyncio async def test_record_execution_result_backfills_execution_log_steps() -> None: - complete_execution = AsyncMock(return_value=[]) + calls: List[str] = [] + + async def complete_execution(*args, **kwargs): # noqa: ANN002, ANN003 + calls.append("complete") + + async def increment_stats(*args, **kwargs): # noqa: ANN002, ANN003 + calls.append("stats") + + async def trim_executions(*args, **kwargs): # noqa: ANN002, ANN003 + calls.append("trim") + return [] + + complete_execution_mock = AsyncMock(side_effect=complete_execution) + increment_stats_mock = AsyncMock(side_effect=increment_stats) + trim_executions_mock = AsyncMock(side_effect=trim_executions) + record_audit = AsyncMock(return_value=None) exec_data = { "id": "exec-1", "workflowId": "wf", @@ -390,33 +512,39 @@ def raise_create_task(coro, *args, **kwargs): # noqa: ANN001, ARG001 raise RuntimeError with ( - patch.object(WorkflowStore, "complete_execution", complete_execution), - patch("flocks.session.recorder.Recorder.record_workflow_execution", AsyncMock(return_value=None)), + patch.object(WorkflowStore, "complete_execution", complete_execution_mock), + patch.object(WorkflowStore, "increment_stats", increment_stats_mock), + patch.object(WorkflowStore, "trim_executions", trim_executions_mock), + patch("flocks.session.recorder.Recorder.record_workflow_execution", record_audit), patch("flocks.workflow.execution_store.asyncio.create_task", side_effect=raise_create_task), ): await record_execution_result("wf", "exec-1", exec_data) - complete_execution.assert_awaited_once() - summary, steps = complete_execution.await_args.args + assert calls == ["complete", "stats", "trim"] + complete_execution_mock.assert_awaited_once() + summary, steps = complete_execution_mock.await_args.args + assert complete_execution_mock.await_args.kwargs == {} assert steps[0][0] == 1 assert steps[0][1]["outputs"] == {"_raw_alerts_count": 150} assert steps[1][0] == 2 assert steps[1][1]["inputs"] == {"_filtered_alerts_count": 150} assert summary["executionLog"] == [] assert summary["stepCount"] == 2 - assert complete_execution.await_args.kwargs == { - "success": True, - "duration": 1.0, - "history_limit": 30, - } + increment_stats_mock.assert_awaited_once_with("wf", success=True, duration=1.0) + trim_executions_mock.assert_awaited_once_with("wf", keep=30) + audit_data = record_audit.await_args.kwargs["run_result"] + assert audit_data["executionLog"] == [step for _, step in steps] @pytest.mark.asyncio async def test_record_execution_result_accepts_explicit_step_batch() -> None: - complete_execution = AsyncMock(return_value=[]) + complete_execution = AsyncMock(return_value=None) + increment_stats = AsyncMock(return_value=None) + trim_executions = AsyncMock(return_value=[]) + record_audit = AsyncMock(return_value=None) explicit_steps = [ - (1, {"node_id": "step-1", "outputs": {"ok": True}}), (2, {"node_id": "step-2", "outputs": {"ok": True}}), + (1, {"node_id": "step-1", "outputs": {"ok": True}}), ] exec_data = { "id": "exec-trigger", @@ -433,7 +561,9 @@ def raise_create_task(coro, *args, **kwargs): # noqa: ANN001, ARG001 with ( patch.object(WorkflowStore, "complete_execution", complete_execution), - patch("flocks.session.recorder.Recorder.record_workflow_execution", AsyncMock(return_value=None)), + patch.object(WorkflowStore, "increment_stats", increment_stats), + patch.object(WorkflowStore, "trim_executions", trim_executions), + patch("flocks.session.recorder.Recorder.record_workflow_execution", record_audit), patch("flocks.workflow.execution_store.asyncio.create_task", side_effect=raise_create_task), ): await record_execution_result( @@ -444,13 +574,88 @@ def raise_create_task(coro, *args, **kwargs): # noqa: ANN001, ARG001 ) summary, persisted_steps = complete_execution.await_args.args + assert complete_execution.await_args.kwargs == {} assert summary["executionLog"] == [] assert persisted_steps == explicit_steps - assert complete_execution.await_args.kwargs == { - "success": True, - "duration": 0.01, - "history_limit": 30, - } + increment_stats.assert_awaited_once_with("wf-trigger", success=True, duration=0.01) + trim_executions.assert_awaited_once_with("wf-trigger", keep=30) + audit_data = record_audit.await_args.kwargs["run_result"] + assert [step["node_id"] for step in audit_data["executionLog"]] == [ + "step-1", + "step-2", + ] + + +@pytest.mark.asyncio +async def test_record_execution_result_stats_failure_does_not_block_retention() -> None: + complete_execution = AsyncMock(return_value=None) + increment_stats = AsyncMock(side_effect=RuntimeError("stats locked")) + trim_executions = AsyncMock(return_value=[]) + + def raise_create_task(coro, *args, **kwargs): # noqa: ANN001, ARG001 + coro.close() + raise RuntimeError + + with ( + patch.object(WorkflowStore, "complete_execution", complete_execution), + patch.object(WorkflowStore, "increment_stats", increment_stats), + patch.object(WorkflowStore, "trim_executions", trim_executions), + patch("flocks.session.recorder.Recorder.record_workflow_execution", AsyncMock(return_value=None)), + patch("flocks.workflow.execution_store.asyncio.create_task", side_effect=raise_create_task), + ): + await record_execution_result( + "wf-stats-failure", + "exec-stats-failure", + { + "id": "exec-stats-failure", + "workflowId": "wf-stats-failure", + "status": "success", + "duration": 0.5, + "executionLog": [], + }, + ) + + complete_execution.assert_awaited_once() + increment_stats.assert_awaited_once() + trim_executions.assert_awaited_once_with("wf-stats-failure", keep=30) + + +@pytest.mark.asyncio +async def test_record_execution_result_retention_failure_keeps_committed_execution() -> None: + complete_execution = AsyncMock(return_value=None) + increment_stats = AsyncMock(return_value=None) + trim_executions = AsyncMock(side_effect=RuntimeError("retention locked")) + + def raise_create_task(coro, *args, **kwargs): # noqa: ANN001, ARG001 + coro.close() + raise RuntimeError + + with ( + patch.object(WorkflowStore, "complete_execution", complete_execution), + patch.object(WorkflowStore, "increment_stats", increment_stats), + patch.object(WorkflowStore, "trim_executions", trim_executions), + patch("flocks.session.recorder.Recorder.record_workflow_execution", AsyncMock(return_value=None)), + patch("flocks.workflow.execution_store.asyncio.create_task", side_effect=raise_create_task), + ): + await record_execution_result( + "wf-retention-failure", + "exec-retention-failure", + { + "id": "exec-retention-failure", + "workflowId": "wf-retention-failure", + "status": "error", + "duration": 0.5, + "executionLog": [], + }, + ) + + complete_execution.assert_awaited_once() + increment_stats.assert_awaited_once_with( + "wf-retention-failure", + success=False, + duration=0.5, + ) + trim_executions.assert_awaited_once_with("wf-retention-failure", keep=30) def test_compact_history_compacts_each_step_inputs() -> None: @@ -538,62 +743,55 @@ def test_compact_outputs_covers_raw_alerts_in_input_params() -> None: @pytest.mark.asyncio -async def test_trim_execution_history_keeps_only_30_and_deletes_matching_jsonl( - tmp_path, -) -> None: - workflow_id = "wf-trim" - for idx in range(32): - exec_id = f"exec-{idx:02d}" - workflow_record = tmp_path / "workflow" / f"{exec_id}.jsonl" - workflow_record.parent.mkdir(parents=True, exist_ok=True) - workflow_record.write_text('{"type":"workflow.summary"}\n', encoding="utf-8") - - # Another workflow's record should be ignored entirely because the trim - # only reads workflow_execution_index//. - other_record = tmp_path / "workflow" / "other-exec.jsonl" - other_record.parent.mkdir(parents=True, exist_ok=True) - other_record.write_text('{"type":"workflow.summary"}\n', encoding="utf-8") - - trim_mock = AsyncMock(return_value=["exec-00", "exec-01"]) +async def test_record_execution_result_deletes_jsonl_for_trimmed_executions(tmp_path) -> None: + workflow_dir = tmp_path / "workflow" + workflow_dir.mkdir(parents=True, exist_ok=True) + trimmed_paths = [workflow_dir / "exec-00.jsonl", workflow_dir / "exec-01.jsonl"] + retained_path = workflow_dir / "exec-02.jsonl" + for record_path in [*trimmed_paths, retained_path]: + record_path.write_text('{"type":"workflow.summary"}\n', encoding="utf-8") + + complete_execution = AsyncMock(return_value=None) + increment_stats = AsyncMock(return_value=None) + trim_executions = AsyncMock(return_value=["exec-00", "exec-01"]) + record_audit = AsyncMock(return_value=None) + created_tasks: List[asyncio.Task[None]] = [] + real_create_task = asyncio.create_task + + def capture_create_task(coro, *args, **kwargs): # noqa: ANN001 + task = real_create_task(coro, *args, **kwargs) + created_tasks.append(task) + return task with ( - patch.object(WorkflowStore, "trim_executions", trim_mock), - patch("flocks.session.recorder._record_dir", return_value=tmp_path), - ): - await _trim_execution_history(workflow_id) - - trim_mock.assert_awaited_once_with(workflow_id, keep=30) - assert not (tmp_path / "workflow" / "exec-00.jsonl").exists() - assert not (tmp_path / "workflow" / "exec-01.jsonl").exists() - assert (tmp_path / "workflow" / "exec-02.jsonl").exists() - assert other_record.exists() - - -@pytest.mark.asyncio -async def test_trim_execution_history_uses_index_without_full_scan(tmp_path) -> None: - workflow_id = "wf-indexed" - for idx in range(32): - exec_id = f"exec-{idx:02d}" - workflow_record = tmp_path / "workflow" / f"{exec_id}.jsonl" - workflow_record.parent.mkdir(parents=True, exist_ok=True) - workflow_record.write_text('{"type":"workflow.summary"}\n', encoding="utf-8") - - trim_mock = AsyncMock(return_value=["exec-00", "exec-01"]) - - with ( - patch.object(WorkflowStore, "trim_executions", trim_mock), - patch("flocks.session.recorder._record_dir", return_value=tmp_path), + patch.object(WorkflowStore, "complete_execution", complete_execution), + patch.object(WorkflowStore, "increment_stats", increment_stats), + patch.object(WorkflowStore, "trim_executions", trim_executions), + patch("flocks.session.recorder.Recorder.record_workflow_execution", record_audit), + patch( + "flocks.workflow.execution_store.Recorder.paths", + return_value=SimpleNamespace(workflow_dir=workflow_dir), + ), + patch( + "flocks.workflow.execution_store.asyncio.create_task", + side_effect=capture_create_task, + ), ): - await _trim_execution_history(workflow_id) - - trim_mock.assert_awaited_once_with(workflow_id, keep=30) - assert not (tmp_path / "workflow" / "exec-00.jsonl").exists() - assert not (tmp_path / "workflow" / "exec-01.jsonl").exists() - + await record_execution_result( + "wf-trim", + "exec-32", + { + "id": "exec-32", + "workflowId": "wf-trim", + "status": "success", + "duration": 0.25, + "executionLog": [], + }, + steps=[(1, {"node_id": "node-1", "outputs": {"ok": True}})], + ) + await asyncio.gather(*created_tasks) -@pytest.mark.asyncio -async def test_trim_execution_history_surfaces_delete_failures() -> None: - workflow_id = "wf-trim-fail" - with patch.object(WorkflowStore, "trim_executions", AsyncMock(side_effect=RuntimeError("locked"))): - with pytest.raises(RuntimeError, match="locked"): - await _trim_execution_history(workflow_id) + trim_executions.assert_awaited_once_with("wf-trim", keep=30) + record_audit.assert_awaited_once() + assert all(not path.exists() for path in trimmed_paths) + assert retained_path.exists() diff --git a/tests/workflow/test_poller_manager.py b/tests/workflow/test_poller_manager.py index 82ee82cf0..f4069d1bd 100644 --- a/tests/workflow/test_poller_manager.py +++ b/tests/workflow/test_poller_manager.py @@ -250,7 +250,17 @@ def _fake_run_workflow( # noqa: ANN001 assert recorded_results[0]["executionLog"] == [] assert recorded_results[0]["stepCount"] == 1 assert recorded_results[0]["loopProgress"]["total_iterations"] == 2 - assert recorded_steps == [] + assert recorded_steps == [ + ( + 1, + { + "node_id": "load", + "node_type": "python", + "inputs": {"iteration": 1, "total_iterations": 2}, + "outputs": {"load_stats": {"record_count": 9}}, + }, + ) + ] assert status["lastStatus"] == "error" assert status["lastError"] == "business rule blocked" assert status["selectedCount"] == 9 diff --git a/tests/workflow/test_tool_run_workflow.py b/tests/workflow/test_tool_run_workflow.py index 223d4397f..0842f90e5 100644 --- a/tests/workflow/test_tool_run_workflow.py +++ b/tests/workflow/test_tool_run_workflow.py @@ -10,6 +10,8 @@ """ import asyncio +import threading + import pytest from unittest.mock import AsyncMock, Mock, patch, MagicMock from typing import Dict, Any @@ -27,6 +29,7 @@ import flocks.tool.task.run_workflow as run_workflow_module from flocks.mcp.client import McpClient from flocks.workflow.runner import RunWorkflowResult, run_workflow +from flocks.workflow.store import WorkflowStore class FakeRunWorkflowResult: @@ -285,7 +288,12 @@ async def test_run_workflow_success(self, tool_context_with_permission, simple_w } ) mock_run = Mock(name="run_workflow", return_value=fake) - with patch.object(run_workflow_module, "_get_workflow_runtime", return_value=_runtime_tuple(run_fn=mock_run)): + direct_audit = AsyncMock(return_value=None) + with ( + patch.object(run_workflow_module, "_get_workflow_runtime", return_value=_runtime_tuple(run_fn=mock_run)), + patch.object(run_workflow_module, "resolve_workflow_id_from_source", return_value=None), + patch.object(run_workflow_module, "_record_workflow_tool_result", direct_audit), + ): result = await ToolRegistry.execute( "run_workflow", ctx=tool_context_with_permission, workflow=simple_workflow, inputs={} ) @@ -297,6 +305,7 @@ async def test_run_workflow_success(self, tool_context_with_permission, simple_w assert result.metadata["status"] == "success" assert result.metadata["steps"] == 1 assert "run_id" not in result.metadata + direct_audit.assert_awaited_once_with("test-workflow-001", fake.__dict__) # Check that permission was requested assert len(tool_context_with_permission._permissions_requested) > 0 @@ -326,7 +335,7 @@ def run_side_effect(**kwargs): steps=1, last_node_id="node-1", outputs={"message": "ok"}, - history=[{"node_id": "node-1", "node_type": "python", "outputs": {"message": "ok"}}], + history=[], error=None, ) @@ -343,6 +352,7 @@ def run_side_effect(**kwargs): ) upsert_execution = AsyncMock(return_value=None) record_result = AsyncMock(return_value=None) + direct_audit = AsyncMock(return_value=None) with ( patch.object(run_workflow_module, "_get_workflow_runtime", return_value=_runtime_tuple(run_fn=mock_run)), @@ -353,8 +363,9 @@ def run_side_effect(**kwargs): ), patch.object(run_workflow_module, "resolve_workflow_id_from_source", return_value="test-workflow-001"), patch.object(run_workflow_module, "create_execution_record", create_execution), - patch.object(run_workflow_module.WorkflowStore, "upsert_execution", upsert_execution), + patch.object(WorkflowStore, "upsert_execution", upsert_execution), patch.object(run_workflow_module, "record_execution_result", record_result), + patch.object(run_workflow_module, "_record_workflow_tool_result", direct_audit), ): result = await ToolRegistry.execute( "run_workflow", @@ -368,9 +379,242 @@ def run_side_effect(**kwargs): assert "run_id" not in result.metadata create_execution.assert_awaited_once() record_result.assert_awaited_once() + expected_step = { + "node_id": "node-1", + "node_type": "python", + "outputs": {"message": "ok"}, + } + assert record_result.await_args.kwargs["steps"] == [(1, expected_step)] + assert record_result.await_args.args[2]["executionLog"] == [expected_step] assert upsert_execution.await_count >= 1 + assert all(call.args[0]["executionLog"] == [] for call in upsert_execution.await_args_list) + direct_audit.assert_not_awaited() assert any(update.get("workflow_execution_id") == "exec-registered" for update in metadata_updates) + @pytest.mark.anyio + async def test_run_workflow_registered_callbacks_do_not_wait_for_progress_storage( + self, + tool_context_with_permission, + simple_workflow, + ): + write_started = asyncio.Event() + release_write = asyncio.Event() + runner_finished = threading.Event() + write_order: list[str] = [] + + async def blocked_upsert(_summary): + write_order.append("progress-start") + write_started.set() + await release_write.wait() + write_order.append("progress-end") + + async def record_result(*args, **kwargs): # noqa: ANN002, ANN003 + write_order.append("final") + + def run_side_effect(**kwargs): + kwargs["on_step_start"]( + kwargs["run_id"], + 1, + MagicMock(id="node-1", type="python"), + {}, + ) + kwargs["on_step_complete"]( + { + "node_id": "node-1", + "node_type": "python", + "outputs": {"ok": True}, + } + ) + runner_finished.set() + return FakeRunWorkflowResult( + status="SUCCEEDED", + run_id=kwargs["run_id"], + steps=1, + last_node_id="node-1", + outputs={"ok": True}, + history=[], + error=None, + ) + + mock_run = Mock(name="run_workflow", side_effect=run_side_effect) + create_execution = AsyncMock( + return_value={ + "id": "exec-blocked", + "workflowId": "test-workflow-001", + "status": "running", + "startedAt": 1, + "executionLog": [], + } + ) + record_result_mock = AsyncMock(side_effect=record_result) + + with ( + patch.object(run_workflow_module, "_get_workflow_runtime", return_value=_runtime_tuple(run_fn=mock_run)), + patch.object(run_workflow_module, "resolve_workflow_id_from_source", return_value="test-workflow-001"), + patch.object(run_workflow_module, "create_execution_record", create_execution), + patch.object(WorkflowStore, "upsert_execution", blocked_upsert), + patch.object(run_workflow_module, "record_execution_result", record_result_mock), + patch.object(run_workflow_module, "_record_workflow_tool_result", AsyncMock(return_value=None)), + ): + task = asyncio.create_task( + ToolRegistry.execute( + "run_workflow", + ctx=tool_context_with_permission, + workflow=simple_workflow, + inputs={}, + ) + ) + await write_started.wait() + assert await asyncio.to_thread(runner_finished.wait, 0.1) + record_result_mock.assert_not_awaited() + release_write.set() + result = await task + + assert result.success is True + record_result_mock.assert_awaited_once() + assert write_order[-2:] == ["progress-end", "final"] + + @pytest.mark.anyio + async def test_run_workflow_registered_cancellation_keeps_completed_and_pending_steps( + self, + tool_context_with_permission, + simple_workflow, + ): + persisted_summaries: list[dict[str, Any]] = [] + + async def capture_upsert(summary): + persisted_summaries.append(dict(summary)) + + def run_side_effect(**kwargs): + kwargs["on_step_start"]( + kwargs["run_id"], + 1, + MagicMock(id="node-1", type="python"), + {}, + ) + kwargs["on_step_complete"]( + { + "node_id": "node-1", + "node_type": "python", + "outputs": {"ok": True}, + } + ) + tool_context_with_permission.abort.set() + kwargs["on_step_start"]( + kwargs["run_id"], + 2, + MagicMock(id="node-2", type="tool"), + {"message": "hello"}, + ) + return FakeRunWorkflowResult( + status="SUCCEEDED", + run_id=kwargs["run_id"], + steps=1, + last_node_id="node-2", + outputs={"ok": True}, + history=[], + error=None, + ) + + mock_run = Mock(name="run_workflow", side_effect=run_side_effect) + create_execution = AsyncMock( + return_value={ + "id": "exec-cancelled", + "workflowId": "test-workflow-001", + "status": "running", + "startedAt": 1, + "executionLog": [], + } + ) + record_result = AsyncMock(return_value=None) + direct_audit = AsyncMock(return_value=None) + + with ( + patch.object(run_workflow_module, "_get_workflow_runtime", return_value=_runtime_tuple(run_fn=mock_run)), + patch.object(run_workflow_module, "resolve_workflow_id_from_source", return_value="test-workflow-001"), + patch.object(run_workflow_module, "create_execution_record", create_execution), + patch.object(WorkflowStore, "upsert_execution", capture_upsert), + patch.object(run_workflow_module, "record_execution_result", record_result), + patch.object(run_workflow_module, "_record_workflow_tool_result", direct_audit), + ): + result = await ToolRegistry.execute( + "run_workflow", + ctx=tool_context_with_permission, + workflow=simple_workflow, + inputs={}, + ) + + steps = record_result.await_args.kwargs["steps"] + assert result.success is False + assert result.metadata["status"] == "cancelled" + assert [step_index for step_index, _ in steps] == [1, 2] + assert [step["node_id"] for _, step in steps] == ["node-1", "node-2"] + assert steps[1][1]["error"] == "Run cancelled before node completed" + assert record_result.await_args.args[2]["status"] == "cancelled" + assert persisted_summaries[-1]["currentPhase"] == "cancelling" + direct_audit.assert_not_awaited() + + @pytest.mark.anyio + async def test_run_workflow_registered_failure_keeps_callback_steps( + self, + tool_context_with_permission, + simple_workflow, + ): + def run_side_effect(**kwargs): + kwargs["on_step_start"]( + kwargs["run_id"], + 1, + MagicMock(id="node-1", type="python"), + {}, + ) + kwargs["on_step_complete"]( + { + "node_id": "node-1", + "node_type": "python", + "outputs": {"ok": True}, + } + ) + kwargs["on_step_start"]( + kwargs["run_id"], + 2, + MagicMock(id="node-2", type="tool"), + {"message": "hello"}, + ) + raise RuntimeError("runner failed") + + mock_run = Mock(name="run_workflow", side_effect=run_side_effect) + create_execution = AsyncMock( + return_value={ + "id": "exec-failed", + "workflowId": "test-workflow-001", + "status": "running", + "startedAt": 1, + "executionLog": [], + } + ) + record_result = AsyncMock(return_value=None) + + with ( + patch.object(run_workflow_module, "_get_workflow_runtime", return_value=_runtime_tuple(run_fn=mock_run)), + patch.object(run_workflow_module, "resolve_workflow_id_from_source", return_value="test-workflow-001"), + patch.object(run_workflow_module, "create_execution_record", create_execution), + patch.object(WorkflowStore, "upsert_execution", AsyncMock(return_value=None)), + patch.object(run_workflow_module, "record_execution_result", record_result), + ): + result = await ToolRegistry.execute( + "run_workflow", + ctx=tool_context_with_permission, + workflow=simple_workflow, + inputs={}, + ) + + steps = record_result.await_args.kwargs["steps"] + assert result.success is False + assert "runner failed" in result.error + assert [step_index for step_index, _ in steps] == [1, 2] + assert [step["node_id"] for _, step in steps] == ["node-1", "node-2"] + assert record_result.await_args.args[2]["executionLog"] == [step for _, step in steps] + @pytest.mark.anyio async def test_run_workflow_registered_id_overrides_missing_workflow_json_id( self, @@ -416,7 +660,7 @@ def run_side_effect(**kwargs): return_value={"id": "wf-directory-id", "workflowJson": workflow_without_id}, ), patch.object(run_workflow_module, "create_execution_record", create_execution), - patch.object(run_workflow_module.WorkflowStore, "upsert_execution", AsyncMock(return_value=None)), + patch.object(WorkflowStore, "upsert_execution", AsyncMock(return_value=None)), patch.object(run_workflow_module, "record_execution_result", AsyncMock(return_value=None)), ): result = await ToolRegistry.execute( @@ -476,17 +720,16 @@ def run_side_effect(**kwargs): } ) upsert_execution = AsyncMock(return_value=None) - record_step = AsyncMock(return_value=None) record_result = AsyncMock(return_value=None) + direct_audit = AsyncMock(return_value=None) with ( patch.object(run_workflow_module, "_get_workflow_runtime", return_value=_runtime_tuple(run_fn=mock_run)), patch.object(run_workflow_module, "resolve_workflow_id_from_source", return_value="test-workflow-001"), patch.object(run_workflow_module, "create_execution_record", create_execution), - patch.object(run_workflow_module.WorkflowStore, "upsert_execution", upsert_execution), - patch.object(run_workflow_module, "record_execution_step", record_step), + patch.object(WorkflowStore, "upsert_execution", upsert_execution), patch.object(run_workflow_module, "record_execution_result", record_result), - patch.object(run_workflow_module, "_record_workflow_tool_result", AsyncMock(return_value=None)), + patch.object(run_workflow_module, "_record_workflow_tool_result", direct_audit), ): result = await ToolRegistry.execute( "run_workflow", @@ -496,8 +739,9 @@ def run_side_effect(**kwargs): ) assert result.success is True - record_step.assert_awaited() - step_payload = record_step.await_args.args[2] + steps = record_result.await_args.kwargs["steps"] + assert [step_index for step_index, _ in steps] == [1] + step_payload = steps[0][1] assert step_payload["inputs"] == { "_raw_alerts_count": 150, "source": "syslog", @@ -506,18 +750,19 @@ def run_side_effect(**kwargs): "_raw_alerts_count": 150, "message": "ok", } + direct_audit.assert_not_awaited() assert result.metadata["has_output"] is True assert result.metadata["output_keys"] == ["enriched_alerts", "message"] assert "outputs" not in result.metadata assert "history" not in result.metadata - assert result.metadata["history_count"] == 0 + assert result.metadata["history_count"] == 1 final_exec_data = record_result.await_args.args[2] assert final_exec_data["outputResults"] == { "_enriched_alerts_count": 150, "message": "done", } - assert final_exec_data["executionLog"] == [] + assert final_exec_data["executionLog"] == [step_payload] assert final_exec_data["stepCount"] == 1 assert any(update.get("workflow_execution_id") == "exec-compacted" for update in metadata_updates) diff --git a/tests/workflow/test_trigger_runtime.py b/tests/workflow/test_trigger_runtime.py index dd34bec84..6456a9227 100644 --- a/tests/workflow/test_trigger_runtime.py +++ b/tests/workflow/test_trigger_runtime.py @@ -79,7 +79,17 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 assert result["executionLog"] == [] assert result["stepCount"] == 1 record_result.assert_awaited_once() - assert record_result.await_args.kwargs["steps"] == [] + assert record_result.await_args.kwargs["steps"] == [ + ( + 1, + { + "node_id": "notify", + "node_type": "tool", + "inputs": {}, + "outputs": {"ok": True}, + }, + ) + ] cleanup_context.assert_awaited_once_with(tool_context) diff --git a/tests/workflow/test_workflow_store.py b/tests/workflow/test_workflow_store.py index d699d0f9b..6608db1b3 100644 --- a/tests/workflow/test_workflow_store.py +++ b/tests/workflow/test_workflow_store.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import os from pathlib import Path import pytest @@ -118,7 +119,7 @@ async def test_workflow_store_increment_stats_is_atomic_for_concurrent_updates() @pytest.mark.asyncio -async def test_complete_execution_writes_steps_summary_and_stats_with_one_commit( +async def test_complete_execution_writes_steps_and_summary_with_one_commit( monkeypatch: pytest.MonkeyPatch, ) -> None: await WorkflowStore.init() @@ -147,8 +148,6 @@ async def counted_commit() -> None: (1, {"node_id": "n1", "outputs": {"ok": 1}}), (2, {"node_id": "n2", "outputs": {"ok": 2}}), ], - success=True, - duration=0.25, ) assert commit_count == 1 @@ -158,11 +157,7 @@ async def counted_commit() -> None: steps, total = await WorkflowStore.list_steps("exec-complete") assert total == 2 assert [step["node_id"] for step in steps] == ["n1", "n2"] - stats = await WorkflowStore.get_stats("wf-complete") - assert stats is not None - assert stats["callCount"] == 1 - assert stats["successCount"] == 1 - assert stats["totalRuntime"] == pytest.approx(0.25) + assert await WorkflowStore.get_stats("wf-complete") is None @pytest.mark.asyncio @@ -199,8 +194,6 @@ async def counted_commit() -> None: "stepCount": 7, }, steps, - success=True, - duration=0.01, ) for index in range(4) ) @@ -216,10 +209,7 @@ async def counted_commit() -> None: assert [step["node_id"] for step in persisted_steps] == [ f"node-{step_index}" for step_index in range(1, 8) ] - stats = await WorkflowStore.get_stats("wf-trigger") - assert stats is not None - assert stats["callCount"] == 4 - assert stats["successCount"] == 4 + assert await WorkflowStore.get_stats("wf-trigger") is None @pytest.mark.asyncio @@ -248,8 +238,6 @@ async def fail_commit() -> None: "stepCount": 1, }, [(1, {"node_id": "node-1", "outputs": {"ok": True}})], - success=True, - duration=0.01, ) monkeypatch.setattr(db, "commit", original_commit) @@ -261,48 +249,59 @@ async def fail_commit() -> None: @pytest.mark.asyncio -async def test_complete_execution_applies_retention_before_single_commit( +async def test_complete_execution_rolls_back_cancelled_transaction( monkeypatch: pytest.MonkeyPatch, ) -> None: await WorkflowStore.init() db = await WorkflowStore.raw_completion_db() - commit_count = 0 original_commit = db.commit - async def counted_commit() -> None: - nonlocal commit_count - commit_count += 1 - await original_commit() + async def cancel_commit() -> None: + raise asyncio.CancelledError - monkeypatch.setattr(db, "commit", counted_commit) - trimmed: list[str] = [] - for index in range(4): - trimmed = await WorkflowStore.complete_execution( + monkeypatch.setattr(db, "commit", cancel_commit) + + with pytest.raises(asyncio.CancelledError): + await WorkflowStore.complete_execution( { - "id": f"exec-retain-{index}", - "workflowId": "wf-retain", + "id": "exec-cancelled-commit", + "workflowId": "wf-cancelled-commit", "status": "success", - "startedAt": index + 1, - "finishedAt": index + 2, - "duration": 0.01, + "startedAt": 1, + "finishedAt": 2, "executionLog": [], "stepCount": 1, }, - [(1, {"node_id": f"node-{index}", "outputs": {"ok": True}})], - success=True, - duration=0.01, - history_limit=3, + [(1, {"node_id": "node-1", "outputs": {"ok": True}})], ) - assert commit_count == 4 - assert trimmed == ["exec-retain-0"] - assert await WorkflowStore.get_execution("exec-retain-0") is None - old_steps, old_total = await WorkflowStore.list_steps("exec-retain-0") - assert old_steps == [] - assert old_total == 0 - executions = await WorkflowStore.list_executions("wf-retain", limit=10) - assert [execution["id"] for execution in executions] == [ - "exec-retain-3", - "exec-retain-2", - "exec-retain-1", - ] + monkeypatch.setattr(db, "commit", original_commit) + await WorkflowStore.complete_execution( + { + "id": "exec-after-cancel", + "workflowId": "wf-cancelled-commit", + "status": "success", + "startedAt": 3, + "finishedAt": 4, + "executionLog": [], + "stepCount": 1, + }, + [(1, {"node_id": "node-2", "outputs": {"ok": True}})], + ) + + assert await WorkflowStore.get_execution("exec-cancelled-commit") is None + assert await WorkflowStore.get_execution("exec-after-cancel") is not None + + +@pytest.mark.asyncio +async def test_completion_connection_reinitializes_after_pid_change() -> None: + await WorkflowStore.init() + original_connection = await WorkflowStore.raw_completion_db() + original_lock = WorkflowStore._completion_lock + WorkflowStore._init_pid = -1 + + refreshed_connection = await WorkflowStore.raw_completion_db() + + assert refreshed_connection is not original_connection + assert WorkflowStore._completion_lock is not original_lock + assert WorkflowStore._init_pid == os.getpid() From 0908fd356fc72ebc063f7a52f3c68d348b632c2f Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Fri, 21 Aug 2026 17:12:34 +0800 Subject: [PATCH 17/63] fix(workflow): persist queued trigger executions --- flocks/ingest/kafka/manager.py | 2 +- flocks/ingest/syslog/manager.py | 2 +- flocks/workflow/poller_manager.py | 2 +- flocks/workflow/triggers/runtime.py | 2 +- tests/ingest/test_kafka_manager.py | 6 +++--- tests/ingest/test_syslog_manager_backpressure.py | 2 +- tests/workflow/test_poller_manager.py | 4 ++-- tests/workflow/test_trigger_runtime.py | 2 +- 8 files changed, 11 insertions(+), 11 deletions(-) diff --git a/flocks/ingest/kafka/manager.py b/flocks/ingest/kafka/manager.py index a44c28eec..d5bd88f44 100644 --- a/flocks/ingest/kafka/manager.py +++ b/flocks/ingest/kafka/manager.py @@ -767,7 +767,7 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: exec_data = await create_execution_record( workflow_id, input_params=summarized_inputs, - persist=False, + persist=True, ) exec_id = exec_data["id"] start_time = time.time() diff --git a/flocks/ingest/syslog/manager.py b/flocks/ingest/syslog/manager.py index ab15ad23a..88a2a68cb 100644 --- a/flocks/ingest/syslog/manager.py +++ b/flocks/ingest/syslog/manager.py @@ -619,7 +619,7 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: exec_data = await create_execution_record( workflow_id, input_params=summarized_inputs, - persist=False, + persist=True, ) exec_id = exec_data["id"] step_recorder = ExecutionStepRecorder(exec_id=exec_id) diff --git a/flocks/workflow/poller_manager.py b/flocks/workflow/poller_manager.py index 7fdb7a101..5e774fba1 100644 --- a/flocks/workflow/poller_manager.py +++ b/flocks/workflow/poller_manager.py @@ -451,7 +451,7 @@ async def _execute_run( exec_data = await create_execution_record( workflow_id, input_params=inputs, - persist=False, + persist=True, ) exec_id = str(exec_data["id"]) step_recorder = ExecutionStepRecorder(exec_id=exec_id) diff --git a/flocks/workflow/triggers/runtime.py b/flocks/workflow/triggers/runtime.py index d41aa012d..f7121a106 100644 --- a/flocks/workflow/triggers/runtime.py +++ b/flocks/workflow/triggers/runtime.py @@ -248,7 +248,7 @@ async def _execute_workflow_effect( exec_data = await create_execution_record( workflow_id, input_params=mapped_inputs, - persist=False, + persist=True, ) exec_id = exec_data["id"] step_recorder = ExecutionStepRecorder(exec_id=exec_id) diff --git a/tests/ingest/test_kafka_manager.py b/tests/ingest/test_kafka_manager.py index 2134fa9c9..19afec5d1 100644 --- a/tests/ingest/test_kafka_manager.py +++ b/tests/ingest/test_kafka_manager.py @@ -549,7 +549,7 @@ async def test_trigger_workflow_compacts_kafka_execution_record( async def _fake_create_execution_record( # noqa: ANN001 workflow_id, *, input_params=None, exec_id=None, persist=True ): - assert persist is False + assert persist is True captured_input_params.update(input_params or {}) return {"id": "exec-compact", "workflowId": workflow_id, "inputParams": input_params} @@ -639,7 +639,7 @@ async def test_trigger_workflow_merges_configured_inputs_with_consumed_message( async def _fake_create_execution_record( # noqa: ANN001 workflow_id, *, input_params=None, exec_id=None, persist=True ): - assert persist is False + assert persist is True recorded_input_params.update(input_params or {}) return {"id": "exec-merge", "workflowId": workflow_id, "inputParams": input_params} @@ -700,7 +700,7 @@ async def test_trigger_workflow_applies_mapping_and_filter( async def _fake_create_execution_record( # noqa: ANN001 workflow_id, *, input_params=None, exec_id=None, persist=True ): - assert persist is False + assert persist is True return {"id": "exec-filter", "workflowId": workflow_id, "inputParams": input_params} async def _fake_record_execution_result( # noqa: ANN001 diff --git a/tests/ingest/test_syslog_manager_backpressure.py b/tests/ingest/test_syslog_manager_backpressure.py index 0ce283747..1b8a459a2 100644 --- a/tests/ingest/test_syslog_manager_backpressure.py +++ b/tests/ingest/test_syslog_manager_backpressure.py @@ -353,7 +353,7 @@ async def test_trigger_workflow_applies_mapping_and_filter( async def _fake_create_execution_record( # noqa: ANN001 workflow_id, *, input_params=None, exec_id=None, persist=True ): - assert persist is False + assert persist is True return {"id": "exec-syslog", "workflowId": workflow_id, "inputParams": input_params} async def _fake_record_execution_result( # noqa: ANN001 diff --git a/tests/workflow/test_poller_manager.py b/tests/workflow/test_poller_manager.py index f4069d1bd..f2f5f6914 100644 --- a/tests/workflow/test_poller_manager.py +++ b/tests/workflow/test_poller_manager.py @@ -156,7 +156,7 @@ async def _fake_create_execution_record( exec_id: str | None = None, persist: bool = True, ) -> dict[str, Any]: - assert persist is False + assert persist is True record = { "id": exec_id or "exec-1", "workflowId": workflow_id, @@ -361,7 +361,7 @@ async def _fake_create_execution_record( exec_id: str | None = None, persist: bool = True, ) -> dict[str, Any]: - assert persist is False + assert persist is True _ = input_params return { "id": exec_id or f"exec-{workflow_id}", diff --git a/tests/workflow/test_trigger_runtime.py b/tests/workflow/test_trigger_runtime.py index 6456a9227..279ade3b0 100644 --- a/tests/workflow/test_trigger_runtime.py +++ b/tests/workflow/test_trigger_runtime.py @@ -75,7 +75,7 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 assert runtime_module.run_workflow.call_args.kwargs["run_id"] == "exec-1" assert runtime_module.run_workflow.call_args.kwargs["execution_profile"] == "high_frequency" assert callable(runtime_module.run_workflow.call_args.kwargs["on_step_complete"]) - assert create_record.await_args.kwargs["persist"] is False + assert create_record.await_args.kwargs["persist"] is True assert result["executionLog"] == [] assert result["stepCount"] == 1 record_result.assert_awaited_once() From afc664994688ff40ef7b63e82f9a7317184817e0 Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Fri, 21 Aug 2026 17:51:01 +0800 Subject: [PATCH 18/63] refactor(workflow): simplify step persistence plumbing Remove obsolete persistence options and duplicated step/row handling while preserving atomic completion and nonblocking progress behavior. Co-Authored-By: Claude Opus 4.6 --- flocks/ingest/kafka/manager.py | 2 - flocks/ingest/syslog/manager.py | 3 +- flocks/server/routes/workflow.py | 2 +- flocks/tool/task/run_workflow.py | 46 ++-------- flocks/workflow/execution_store.py | 10 +-- flocks/workflow/poller_manager.py | 3 +- flocks/workflow/store.py | 67 +++++---------- flocks/workflow/triggers/runtime.py | 3 +- tests/ingest/test_kafka_manager.py | 9 +- .../test_syslog_manager_backpressure.py | 3 +- .../server/routes/test_workflow_run_route.py | 86 ++++--------------- .../workflow/test_execution_store_compact.py | 77 +++-------------- tests/workflow/test_poller_manager.py | 8 +- tests/workflow/test_trigger_runtime.py | 5 +- 14 files changed, 76 insertions(+), 248 deletions(-) diff --git a/flocks/ingest/kafka/manager.py b/flocks/ingest/kafka/manager.py index d5bd88f44..dea17a7e9 100644 --- a/flocks/ingest/kafka/manager.py +++ b/flocks/ingest/kafka/manager.py @@ -767,14 +767,12 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: exec_data = await create_execution_record( workflow_id, input_params=summarized_inputs, - persist=True, ) exec_id = exec_data["id"] start_time = time.time() trigger_meta = mapped_inputs.get("_flocks", {}).get("trigger", {}) trigger_input_keys = list((trigger.mapping or {}).keys()) or [input_key] step_recorder = ExecutionStepRecorder( - exec_id=exec_id, step_compactor=lambda step: _compact_step_for_kafka_storage( step, input_key=input_key, diff --git a/flocks/ingest/syslog/manager.py b/flocks/ingest/syslog/manager.py index 88a2a68cb..63e1d1b18 100644 --- a/flocks/ingest/syslog/manager.py +++ b/flocks/ingest/syslog/manager.py @@ -619,10 +619,9 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: exec_data = await create_execution_record( workflow_id, input_params=summarized_inputs, - persist=True, ) exec_id = exec_data["id"] - step_recorder = ExecutionStepRecorder(exec_id=exec_id) + step_recorder = ExecutionStepRecorder() start_time = time.time() trigger_meta = mapped_inputs.get("_flocks", {}).get("trigger", {}) tool_context = None diff --git a/flocks/server/routes/workflow.py b/flocks/server/routes/workflow.py index 0a35c4eab..b1ca85243 100644 --- a/flocks/server/routes/workflow.py +++ b/flocks/server/routes/workflow.py @@ -1128,7 +1128,7 @@ async def _run_workflow_execution_task( ) -> None: """Execute a workflow in the background and keep the execution record updated.""" start_time = time.time() - step_recorder = ExecutionStepRecorder(exec_id=exec_id) + step_recorder = ExecutionStepRecorder() pending_step_index: Optional[int] = None pending_step: Optional[Dict[str, Any]] = None execution_summary: Dict[str, Any] = { diff --git a/flocks/tool/task/run_workflow.py b/flocks/tool/task/run_workflow.py index ab5fcbf10..7c50452b7 100644 --- a/flocks/tool/task/run_workflow.py +++ b/flocks/tool/task/run_workflow.py @@ -570,9 +570,8 @@ async def run_workflow_tool( canonical_workflow_id = registered_workflow_id or resolve_workflow_id_from_source(workflow_source) display_workflow_id = canonical_workflow_id or workflow_id tracked_execution: Optional[Dict[str, Any]] = None - step_recorder: Optional[ExecutionStepRecorder] = None + step_recorder = ExecutionStepRecorder() progress_writer: Optional[ExecutionProgressWriter] = None - callback_step_count = 0 pending_step_index: Optional[int] = None pending_step: Optional[Dict[str, Any]] = None final_step_batch: Optional[List[Tuple[int, Dict[str, Any]]]] = None @@ -639,35 +638,9 @@ def _on_step_start( return step_index def _on_step_complete(step_result: Any) -> None: - nonlocal callback_step_count, pending_step_index, pending_step - if step_recorder is not None: - step_recorder.on_step_complete(step_result) - callback_step_count = step_recorder.step_count - progress_update = dict(step_recorder.summary) - else: - if hasattr(step_result, "model_dump"): - step_dict = step_result.model_dump(mode="json") - elif isinstance(step_result, dict): - step_dict = dict(step_result) - else: - step_dict = {"node_id": None, "outputs": {}, "error": str(step_result)} - callback_step_count += 1 - compacted_step = compact_step_for_storage(step_dict) - progress_update = { - "stepCount": callback_step_count, - "currentNodeId": compacted_step.get("node_id"), - "currentNodeType": compacted_step.get("node_type") - or compacted_step.get("type"), - "currentPhase": "running", - "currentStepIndex": callback_step_count, - "loopProgress": derive_loop_progress( - node_id=compacted_step.get("node_id"), - global_step_index=callback_step_count, - inputs=compacted_step.get("inputs"), - outputs=compacted_step.get("outputs"), - ), - "updatedAt": int(time.time() * 1000), - } + nonlocal pending_step_index, pending_step + step_recorder.on_step_complete(step_result) + progress_update = dict(step_recorder.summary) pending_step_index = None pending_step = None if ctx.abort.is_set(): @@ -688,8 +661,8 @@ def _on_step_complete(step_result: Any) -> None: "phase": progress_update["currentPhase"], "current_node_id": progress_update.get("currentNodeId"), "current_node_type": progress_update.get("currentNodeType"), - "step_index": callback_step_count, - "step_count": callback_step_count, + "step_index": step_recorder.step_count, + "step_count": step_recorder.step_count, "loop_progress": progress_update.get("loopProgress"), }, } @@ -698,7 +671,7 @@ def _on_step_complete(step_result: Any) -> None: def _take_final_step_batch() -> List[Tuple[int, Dict[str, Any]]]: nonlocal final_step_batch if final_step_batch is None: - final_step_batch = step_recorder.take_steps() if step_recorder is not None else [] + final_step_batch = step_recorder.take_steps() if pending_step_index is not None and pending_step is not None: final_step_batch.append((pending_step_index, pending_step)) final_step_batch.sort(key=lambda item: item[0]) @@ -721,7 +694,6 @@ def _take_final_step_batch() -> List[Tuple[int, Dict[str, Any]]]: canonical_workflow_id, input_params=workflow_inputs, ) - step_recorder = ExecutionStepRecorder(exec_id=tracked_execution["id"]) progress_writer = ExecutionProgressWriter(tracked_execution) # Update metadata to show workflow is running @@ -868,7 +840,7 @@ def _take_final_step_batch() -> List[Tuple[int, Dict[str, Any]]]: history_count = len(final_history) final_step_count = result_dict.get("steps") if not isinstance(final_step_count, int): - final_step_count = callback_step_count + final_step_count = step_recorder.step_count final_step_count = max( final_step_count, max((step_index for step_index, _ in tracked_steps), default=0), @@ -988,7 +960,7 @@ def _take_final_step_batch() -> List[Tuple[int, Dict[str, Any]]]: }, ) terminal_status = "cancelled" if ctx.abort.is_set() else "error" - final_step_count = callback_step_count + final_step_count = step_recorder.step_count if tracked_execution and canonical_workflow_id: tracked_steps = _take_final_step_batch() final_step_count = max( diff --git a/flocks/workflow/execution_store.py b/flocks/workflow/execution_store.py index 3dbc59cd4..6503ee981 100644 --- a/flocks/workflow/execution_store.py +++ b/flocks/workflow/execution_store.py @@ -372,7 +372,7 @@ def workflow_execution_step_prefix(exec_id: str) -> str: def compact_execution_summary(exec_data: Dict[str, Any]) -> Dict[str, Any]: """Return an execution record safe to keep in the hot summary row. - Step details are stored separately under ``workflow_execution_step`` keys. + Step details are stored separately in ``workflow_execution_steps`` rows. Keeping ``executionLog`` out of the summary row avoids rewriting an ever-growing JSON blob on every progress update. """ @@ -398,10 +398,8 @@ class ExecutionStepRecorder: def __init__( self, *, - exec_id: str, step_compactor: Callable[[Any], Dict[str, Any]] = compact_step_for_storage, ) -> None: - self.exec_id = exec_id self.step_compactor = step_compactor self.step_count = 0 self.summary: Dict[str, Any] = {} @@ -618,9 +616,8 @@ async def create_execution_record( *, input_params: Optional[Dict[str, Any]] = None, exec_id: Optional[str] = None, - persist: bool = True, ) -> Dict[str, Any]: - """Build a running workflow execution record and optionally persist it. + """Build and persist a running workflow execution record. *input_params* is passed through ``compact_outputs_for_storage`` before writing to SQLite so that batch HTTP calls whose inputs contain a key in @@ -635,8 +632,7 @@ async def create_execution_record( input_params=compacted_params, exec_id=exec_id, ) - if persist: - await WorkflowStore.upsert_execution(compact_execution_summary(exec_data)) + await WorkflowStore.upsert_execution(compact_execution_summary(exec_data)) return exec_data diff --git a/flocks/workflow/poller_manager.py b/flocks/workflow/poller_manager.py index 5e774fba1..9d348d5db 100644 --- a/flocks/workflow/poller_manager.py +++ b/flocks/workflow/poller_manager.py @@ -451,10 +451,9 @@ async def _execute_run( exec_data = await create_execution_record( workflow_id, input_params=inputs, - persist=True, ) exec_id = str(exec_data["id"]) - step_recorder = ExecutionStepRecorder(exec_id=exec_id) + step_recorder = ExecutionStepRecorder() current = self._status.get(workflow_id) or self._base_status(workflow_id) current["lastRunAt"] = started_at_ms current["activeRuns"] = self._cleanup_done_runs(workflow_id) diff --git a/flocks/workflow/store.py b/flocks/workflow/store.py index b4d9b129f..f7efbd0a1 100644 --- a/flocks/workflow/store.py +++ b/flocks/workflow/store.py @@ -37,6 +37,13 @@ "workflow_syslog_config/", ) _WORKFLOW_PREFIXES = _WORKFLOW_KV_PREFIXES + _WORKFLOW_TABLE_PREFIXES +_EXECUTION_UPSERT_SQL = """ + INSERT OR REPLACE INTO workflow_executions + (id, workflow_id, status, current_phase, current_node_id, current_node_type, + current_step_index, step_count, input_params, output_results, error_message, + trigger_id, trigger_type, started_at, finished_at, duration, updated_at, payload) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +""" class WorkflowStore: @@ -312,21 +319,18 @@ async def _migrate_legacy_kv(cls) -> None: log.info("workflow.store.legacy_kv_migrated", counts) @classmethod - async def upsert_execution(cls, exec_data: Dict[str, Any]) -> None: - db = await cls._db() + def _execution_row( + cls, + exec_data: Dict[str, Any], + ) -> Tuple[str, str, Tuple[Any, ...]]: payload = dict(exec_data) exec_id = str(payload.get("id") or "") workflow_id = str(payload.get("workflowId") or payload.get("workflow_id") or "") if not exec_id or not workflow_id: raise ValueError("workflow execution requires id and workflowId") - await db.execute( - """ - INSERT OR REPLACE INTO workflow_executions - (id, workflow_id, status, current_phase, current_node_id, current_node_type, - current_step_index, step_count, input_params, output_results, error_message, - trigger_id, trigger_type, started_at, finished_at, duration, updated_at, payload) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - """, + return ( + exec_id, + workflow_id, ( exec_id, workflow_id, @@ -348,6 +352,12 @@ async def upsert_execution(cls, exec_data: Dict[str, Any]) -> None: cls._json_dumps(payload), ), ) + + @classmethod + async def upsert_execution(cls, exec_data: Dict[str, Any]) -> None: + db = await cls._db() + _, _, row = cls._execution_row(exec_data) + await db.execute(_EXECUTION_UPSERT_SQL, row) await db.commit() @classmethod @@ -497,12 +507,7 @@ async def complete_execution( ) -> None: """Atomically persist one final execution summary and its step batch.""" db = await cls._completion_db() - payload = dict(exec_data) - exec_id = str(payload.get("id") or "") - workflow_id = str(payload.get("workflowId") or payload.get("workflow_id") or "") - if not exec_id or not workflow_id: - raise ValueError("workflow execution requires id and workflowId") - + exec_id, workflow_id, execution_row = cls._execution_row(exec_data) step_rows = cls._step_rows(exec_id, steps) lock = cls._completion_lock if lock is None: @@ -521,35 +526,7 @@ async def complete_execution( """, step_rows, ) - await db.execute( - """ - INSERT OR REPLACE INTO workflow_executions - (id, workflow_id, status, current_phase, current_node_id, current_node_type, - current_step_index, step_count, input_params, output_results, error_message, - trigger_id, trigger_type, started_at, finished_at, duration, updated_at, payload) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - """, - ( - exec_id, - workflow_id, - str(payload.get("status") or "running"), - payload.get("currentPhase"), - payload.get("currentNodeId"), - payload.get("currentNodeType"), - cls._as_int(payload.get("currentStepIndex")), - cls._as_int(payload.get("stepCount")) or 0, - cls._json_dumps(payload.get("inputParams") or {}), - cls._json_dumps(payload.get("outputResults") or {}), - payload.get("errorMessage"), - payload.get("triggerId"), - payload.get("triggerType"), - cls._as_int(payload.get("startedAt")) or cls._now_ms(), - cls._as_int(payload.get("finishedAt")), - cls._as_float(payload.get("duration")), - cls._as_int(payload.get("updatedAt")) or cls._now_ms(), - cls._json_dumps(payload), - ), - ) + await db.execute(_EXECUTION_UPSERT_SQL, execution_row) await db.commit() except BaseException: try: diff --git a/flocks/workflow/triggers/runtime.py b/flocks/workflow/triggers/runtime.py index f7121a106..2d88fb1e4 100644 --- a/flocks/workflow/triggers/runtime.py +++ b/flocks/workflow/triggers/runtime.py @@ -248,10 +248,9 @@ async def _execute_workflow_effect( exec_data = await create_execution_record( workflow_id, input_params=mapped_inputs, - persist=True, ) exec_id = exec_data["id"] - step_recorder = ExecutionStepRecorder(exec_id=exec_id) + step_recorder = ExecutionStepRecorder() started_at = time.time() tool_context = None try: diff --git a/tests/ingest/test_kafka_manager.py b/tests/ingest/test_kafka_manager.py index 19afec5d1..2b9b6eb09 100644 --- a/tests/ingest/test_kafka_manager.py +++ b/tests/ingest/test_kafka_manager.py @@ -547,9 +547,8 @@ async def test_trigger_workflow_compacts_kafka_execution_record( captured_steps: list[tuple[int, dict]] = [] async def _fake_create_execution_record( # noqa: ANN001 - workflow_id, *, input_params=None, exec_id=None, persist=True + workflow_id, *, input_params=None, exec_id=None ): - assert persist is True captured_input_params.update(input_params or {}) return {"id": "exec-compact", "workflowId": workflow_id, "inputParams": input_params} @@ -637,9 +636,8 @@ async def test_trigger_workflow_merges_configured_inputs_with_consumed_message( recorded_input_params: dict = {} async def _fake_create_execution_record( # noqa: ANN001 - workflow_id, *, input_params=None, exec_id=None, persist=True + workflow_id, *, input_params=None, exec_id=None ): - assert persist is True recorded_input_params.update(input_params or {}) return {"id": "exec-merge", "workflowId": workflow_id, "inputParams": input_params} @@ -698,9 +696,8 @@ async def test_trigger_workflow_applies_mapping_and_filter( recorded_exec_data: dict = {} async def _fake_create_execution_record( # noqa: ANN001 - workflow_id, *, input_params=None, exec_id=None, persist=True + workflow_id, *, input_params=None, exec_id=None ): - assert persist is True return {"id": "exec-filter", "workflowId": workflow_id, "inputParams": input_params} async def _fake_record_execution_result( # noqa: ANN001 diff --git a/tests/ingest/test_syslog_manager_backpressure.py b/tests/ingest/test_syslog_manager_backpressure.py index 1b8a459a2..29750f055 100644 --- a/tests/ingest/test_syslog_manager_backpressure.py +++ b/tests/ingest/test_syslog_manager_backpressure.py @@ -351,9 +351,8 @@ async def test_trigger_workflow_applies_mapping_and_filter( recorded_steps: list[tuple[int, dict]] = [] async def _fake_create_execution_record( # noqa: ANN001 - workflow_id, *, input_params=None, exec_id=None, persist=True + workflow_id, *, input_params=None, exec_id=None ): - assert persist is True return {"id": "exec-syslog", "workflowId": workflow_id, "inputParams": input_params} async def _fake_record_execution_result( # noqa: ANN001 diff --git a/tests/server/routes/test_workflow_run_route.py b/tests/server/routes/test_workflow_run_route.py index f0891bafa..69c213656 100644 --- a/tests/server/routes/test_workflow_run_route.py +++ b/tests/server/routes/test_workflow_run_route.py @@ -41,6 +41,17 @@ def _two_node_workflow_json(edge): } +def _progress_writer(exec_id: str, **updates): + summary = { + "id": exec_id, + "workflowId": "wf-1", + "status": "running", + "executionLog": [], + } + summary.update(updates) + return workflow_module.ExecutionProgressWriter(summary) + + @pytest.mark.asyncio async def test_create_workflow_applies_vertex_cache_runtime_defaults(monkeypatch: pytest.MonkeyPatch) -> None: writes: list[dict] = [] @@ -252,35 +263,17 @@ def run_workflow_mock(**kwargs): run_mock = Mock(side_effect=run_workflow_mock) record_result = AsyncMock(return_value=None) upsert_execution = AsyncMock(return_value=None) - storage_read = AsyncMock( - return_value={ - "id": "exec-1", - "workflowId": "wf-1", - "currentNodeType": "tool", - "executionLog": [], - } - ) - monkeypatch.setattr(MCP, "init", init_mock) monkeypatch.setattr(workflow_module, "run_workflow", run_mock) monkeypatch.setattr(workflow_module.WorkflowStore, "upsert_execution", upsert_execution) monkeypatch.setattr(workflow_module, "_resolve_execution_outcome", lambda _result: ("success", None)) monkeypatch.setattr(workflow_module, "_record_execution_result", record_result) - monkeypatch.setattr(workflow_module.Storage, "read", storage_read) - monkeypatch.setattr(workflow_module.Storage, "write", AsyncMock(return_value=None)) monkeypatch.setattr(workflow_module, "compact_outputs_for_storage", lambda value: value) monkeypatch.setattr(workflow_module, "compact_history_for_storage", lambda value: value) req = workflow_module.WorkflowRunRequest(inputs={"ip": "8.8.8.8"}, trace=False) tool_context = ToolContext(session_id="session-1", message_id="message-1", agent="rex") - progress_writer = workflow_module.ExecutionProgressWriter( - { - "id": "exec-1", - "workflowId": "wf-1", - "status": "running", - "executionLog": [], - } - ) + progress_writer = _progress_writer("exec-1") await workflow_module._run_workflow_execution_task( workflow_id="wf-1", @@ -342,14 +335,7 @@ def run_workflow_mock(**kwargs): cancel_event = workflow_module.threading.Event() cancel_event.set() - progress_writer = workflow_module.ExecutionProgressWriter( - { - "id": "exec-cancelled", - "workflowId": "wf-1", - "status": "running", - "executionLog": [], - } - ) + progress_writer = _progress_writer("exec-cancelled") await workflow_module._run_workflow_execution_task( workflow_id="wf-1", workflow_json={"id": "wf-1", "start": "node-1", "nodes": [], "edges": []}, @@ -418,14 +404,7 @@ def run_workflow_mock(**kwargs): monkeypatch.setattr(workflow_module, "_record_execution_result", record_result) monkeypatch.setattr(workflow_module.WorkflowStore, "upsert_execution", upsert_execution) - progress_writer = workflow_module.ExecutionProgressWriter( - { - "id": "exec-partial-cancel", - "workflowId": "wf-1", - "status": "running", - "executionLog": [], - } - ) + progress_writer = _progress_writer("exec-partial-cancel") await workflow_module._run_workflow_execution_task( workflow_id="wf-1", workflow_json={"id": "wf-1", "start": "node-1", "nodes": [], "edges": []}, @@ -481,14 +460,7 @@ def run_workflow_mock(**kwargs): AsyncMock(return_value=None), ) - progress_writer = workflow_module.ExecutionProgressWriter( - { - "id": "exec-runner-error", - "workflowId": "wf-1", - "status": "running", - "executionLog": [], - } - ) + progress_writer = _progress_writer("exec-runner-error") await workflow_module._run_workflow_execution_task( workflow_id="wf-1", workflow_json={"id": "wf-1", "start": "node-1", "nodes": [], "edges": []}, @@ -544,14 +516,7 @@ def run_workflow_mock(**kwargs): AsyncMock(return_value=None), ) - progress_writer = workflow_module.ExecutionProgressWriter( - { - "id": "exec-storage-error", - "workflowId": "wf-1", - "status": "running", - "executionLog": [], - } - ) + progress_writer = _progress_writer("exec-storage-error") with pytest.raises(RuntimeError, match="storage failed"): await workflow_module._run_workflow_execution_task( @@ -627,14 +592,7 @@ def run_workflow_mock(**kwargs): monkeypatch.setattr(workflow_module, "_record_execution_result", record_result_mock) monkeypatch.setattr(workflow_module.WorkflowStore, "upsert_execution", blocked_upsert) - progress_writer = workflow_module.ExecutionProgressWriter( - { - "id": "exec-blocked-progress", - "workflowId": "wf-1", - "status": "running", - "executionLog": [], - } - ) + progress_writer = _progress_writer("exec-blocked-progress") task = asyncio.create_task( workflow_module._run_workflow_execution_task( workflow_id="wf-1", @@ -680,15 +638,7 @@ async def capture_upsert(summary): ) monkeypatch.setattr(workflow_module.WorkflowStore, "upsert_execution", capture_upsert) - progress_writer = workflow_module.ExecutionProgressWriter( - { - "id": "exec-cancel-route", - "workflowId": "wf-1", - "status": "running", - "currentPhase": "queued", - "executionLog": [], - } - ) + progress_writer = _progress_writer("exec-cancel-route", currentPhase="queued") progress_writer.submit({"currentPhase": "running", "currentNodeId": "node-1"}) cancel_event = workflow_module.threading.Event() current_task = asyncio.current_task() diff --git a/tests/workflow/test_execution_store_compact.py b/tests/workflow/test_execution_store_compact.py index 0ceb57849..ea8c3d5cd 100644 --- a/tests/workflow/test_execution_store_compact.py +++ b/tests/workflow/test_execution_store_compact.py @@ -34,7 +34,6 @@ compact_execution_summary, compact_outputs_for_storage, compact_step_for_storage, - create_execution_record, record_execution_result, workflow_execution_step_key, ) @@ -45,6 +44,11 @@ def _make_alerts(n: int) -> List[Dict[str, Any]]: return [{"sip": f"1.2.3.{i % 256}", "url": f"/p/{i}"} for i in range(n)] +def _raise_create_task(coro, *args, **kwargs): # noqa: ANN001, ARG001 + coro.close() + raise RuntimeError + + # ── compact_outputs_for_storage ─────────────────────────────────────────────── @@ -304,27 +308,10 @@ def test_workflow_execution_step_key_is_append_only_namespaced() -> None: assert workflow_execution_step_key("exec-1", 12) == "workflow_execution_step/exec-1/00000012" -@pytest.mark.asyncio -async def test_create_execution_record_can_skip_initial_database_write() -> None: - upsert_execution = AsyncMock(return_value=None) - - with patch.object(WorkflowStore, "upsert_execution", upsert_execution): - record = await create_execution_record( - "wf-trigger", - input_params={"message": "hello"}, - exec_id="exec-trigger", - persist=False, - ) - - assert record["id"] == "exec-trigger" - assert record["currentPhase"] == "queued" - upsert_execution.assert_not_awaited() - - def test_execution_step_recorder_collects_steps_without_storage_calls() -> None: record_step = AsyncMock(return_value=None) record_steps = AsyncMock(return_value=None) - recorder = ExecutionStepRecorder(exec_id="exec-batch") + recorder = ExecutionStepRecorder() with ( patch.object(WorkflowStore, "record_step", record_step), @@ -344,34 +331,6 @@ def test_execution_step_recorder_collects_steps_without_storage_calls() -> None: record_steps.assert_not_awaited() -@pytest.mark.asyncio -async def test_four_trigger_workers_collect_steps_without_storage() -> None: - """Four trigger threads collect complete batches without callback SQL.""" - record_step = AsyncMock(return_value=None) - record_steps = AsyncMock(return_value=None) - recorders = [ExecutionStepRecorder(exec_id=f"exec-trigger-{worker}") for worker in range(4)] - - def _run_seven_steps(recorder: ExecutionStepRecorder) -> None: - for step in range(7): - recorder.on_step_complete( - {"node_id": f"node-{step}", "outputs": {"ok": True}} - ) - - with ( - patch.object(WorkflowStore, "record_step", record_step), - patch.object(WorkflowStore, "record_steps", record_steps), - ): - await asyncio.gather( - *(asyncio.to_thread(_run_seven_steps, recorder) for recorder in recorders) - ) - - batches = [recorder.take_steps() for recorder in recorders] - assert [len(batch) for batch in batches] == [7, 7, 7, 7] - assert [recorder.step_count for recorder in recorders] == [7, 7, 7, 7] - record_step.assert_not_awaited() - record_steps.assert_not_awaited() - - @pytest.mark.asyncio async def test_progress_writer_submits_without_waiting_and_coalesces_updates() -> None: write_started = asyncio.Event() @@ -507,16 +466,12 @@ async def trim_executions(*args, **kwargs): # noqa: ANN002, ANN003 ], } - def raise_create_task(coro, *args, **kwargs): # noqa: ANN001, ARG001 - coro.close() - raise RuntimeError - with ( patch.object(WorkflowStore, "complete_execution", complete_execution_mock), patch.object(WorkflowStore, "increment_stats", increment_stats_mock), patch.object(WorkflowStore, "trim_executions", trim_executions_mock), patch("flocks.session.recorder.Recorder.record_workflow_execution", record_audit), - patch("flocks.workflow.execution_store.asyncio.create_task", side_effect=raise_create_task), + patch("flocks.workflow.execution_store.asyncio.create_task", side_effect=_raise_create_task), ): await record_execution_result("wf", "exec-1", exec_data) @@ -555,16 +510,12 @@ async def test_record_execution_result_accepts_explicit_step_batch() -> None: "stepCount": 2, } - def raise_create_task(coro, *args, **kwargs): # noqa: ANN001, ARG001 - coro.close() - raise RuntimeError - with ( patch.object(WorkflowStore, "complete_execution", complete_execution), patch.object(WorkflowStore, "increment_stats", increment_stats), patch.object(WorkflowStore, "trim_executions", trim_executions), patch("flocks.session.recorder.Recorder.record_workflow_execution", record_audit), - patch("flocks.workflow.execution_store.asyncio.create_task", side_effect=raise_create_task), + patch("flocks.workflow.execution_store.asyncio.create_task", side_effect=_raise_create_task), ): await record_execution_result( "wf-trigger", @@ -592,16 +543,12 @@ async def test_record_execution_result_stats_failure_does_not_block_retention() increment_stats = AsyncMock(side_effect=RuntimeError("stats locked")) trim_executions = AsyncMock(return_value=[]) - def raise_create_task(coro, *args, **kwargs): # noqa: ANN001, ARG001 - coro.close() - raise RuntimeError - with ( patch.object(WorkflowStore, "complete_execution", complete_execution), patch.object(WorkflowStore, "increment_stats", increment_stats), patch.object(WorkflowStore, "trim_executions", trim_executions), patch("flocks.session.recorder.Recorder.record_workflow_execution", AsyncMock(return_value=None)), - patch("flocks.workflow.execution_store.asyncio.create_task", side_effect=raise_create_task), + patch("flocks.workflow.execution_store.asyncio.create_task", side_effect=_raise_create_task), ): await record_execution_result( "wf-stats-failure", @@ -626,16 +573,12 @@ async def test_record_execution_result_retention_failure_keeps_committed_executi increment_stats = AsyncMock(return_value=None) trim_executions = AsyncMock(side_effect=RuntimeError("retention locked")) - def raise_create_task(coro, *args, **kwargs): # noqa: ANN001, ARG001 - coro.close() - raise RuntimeError - with ( patch.object(WorkflowStore, "complete_execution", complete_execution), patch.object(WorkflowStore, "increment_stats", increment_stats), patch.object(WorkflowStore, "trim_executions", trim_executions), patch("flocks.session.recorder.Recorder.record_workflow_execution", AsyncMock(return_value=None)), - patch("flocks.workflow.execution_store.asyncio.create_task", side_effect=raise_create_task), + patch("flocks.workflow.execution_store.asyncio.create_task", side_effect=_raise_create_task), ): await record_execution_result( "wf-retention-failure", diff --git a/tests/workflow/test_poller_manager.py b/tests/workflow/test_poller_manager.py index f2f5f6914..cb10353ec 100644 --- a/tests/workflow/test_poller_manager.py +++ b/tests/workflow/test_poller_manager.py @@ -99,7 +99,7 @@ def _fake_run_workflow( # noqa: ANN001 monkeypatch.setattr( poller_manager, "create_execution_record", - lambda workflow_id, *, input_params=None, exec_id=None, persist=True: asyncio.sleep( + lambda workflow_id, *, input_params=None, exec_id=None: asyncio.sleep( 0, result={ "id": exec_id or f"exec-{workflow_id}", @@ -154,9 +154,7 @@ async def _fake_create_execution_record( *, input_params: dict[str, Any] | None = None, exec_id: str | None = None, - persist: bool = True, ) -> dict[str, Any]: - assert persist is True record = { "id": exec_id or "exec-1", "workflowId": workflow_id, @@ -304,7 +302,7 @@ def _fake_run_workflow( # noqa: ANN001 monkeypatch.setattr( poller_manager, "create_execution_record", - lambda workflow_id, *, input_params=None, exec_id=None, persist=True: asyncio.sleep( + lambda workflow_id, *, input_params=None, exec_id=None: asyncio.sleep( 0, result={ "id": exec_id or f"exec-{workflow_id}", @@ -359,9 +357,7 @@ async def _fake_create_execution_record( *, input_params: dict[str, Any] | None = None, exec_id: str | None = None, - persist: bool = True, ) -> dict[str, Any]: - assert persist is True _ = input_params return { "id": exec_id or f"exec-{workflow_id}", diff --git a/tests/workflow/test_trigger_runtime.py b/tests/workflow/test_trigger_runtime.py index 279ade3b0..7f05ce780 100644 --- a/tests/workflow/test_trigger_runtime.py +++ b/tests/workflow/test_trigger_runtime.py @@ -75,7 +75,10 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 assert runtime_module.run_workflow.call_args.kwargs["run_id"] == "exec-1" assert runtime_module.run_workflow.call_args.kwargs["execution_profile"] == "high_frequency" assert callable(runtime_module.run_workflow.call_args.kwargs["on_step_complete"]) - assert create_record.await_args.kwargs["persist"] is True + create_record.assert_awaited_once_with( + "wf-trigger", + input_params={"message": "hello"}, + ) assert result["executionLog"] == [] assert result["stepCount"] == 1 record_result.assert_awaited_once() From 3f3dae6ccc2e7d1a70d265a48390ceddce18ed9a Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Tue, 25 Aug 2026 17:11:42 +0800 Subject: [PATCH 19/63] perf(kafka): skip unneeded workflow tool context --- flocks/ingest/kafka/manager.py | 30 +++++++++++++++---- tests/ingest/test_kafka_manager.py | 48 ++++++++++++++++++++++++++++++ 2 files changed, 73 insertions(+), 5 deletions(-) diff --git a/flocks/ingest/kafka/manager.py b/flocks/ingest/kafka/manager.py index dea17a7e9..5f7846431 100644 --- a/flocks/ingest/kafka/manager.py +++ b/flocks/ingest/kafka/manager.py @@ -124,6 +124,20 @@ def _strip_execution_only_comments(value: Any) -> Any: } +def _configured_bool(value: Any, *, default: bool) -> bool: + """Parse a boolean trigger input without treating ``"false"`` as true.""" + if isinstance(value, bool): + return value + if value is None: + return default + normalized = str(value).strip().lower() + if normalized in {"1", "true", "yes", "on"}: + return True + if normalized in {"0", "false", "no", "off"}: + return False + return default + + def _decode_message(raw: Optional[bytes]) -> Any: """Decode a Kafka message value to a Python object. @@ -750,6 +764,10 @@ async def _trigger_workflow( configured_inputs = _strip_execution_only_comments( configured_inputs if isinstance(configured_inputs, dict) else {} ) + tool_context_required = _configured_bool( + configured_inputs.get("tool_context_required"), + default=True, + ) event = build_trigger_event( workflow_id=workflow_id, trigger=trigger, @@ -781,10 +799,11 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: ) tool_context = None try: - tool_context = await build_workflow_tool_context( - workflow_id=workflow_id, - action_name=f"trigger:{trigger.type}", - ) + if tool_context_required: + 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, @@ -843,7 +862,8 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: ) finally: steps = step_recorder.take_steps() - await cleanup_workflow_tool_context(tool_context) + if tool_context is not None: + await cleanup_workflow_tool_context(tool_context) try: await record_execution_result( workflow_id, diff --git a/tests/ingest/test_kafka_manager.py b/tests/ingest/test_kafka_manager.py index 2b9b6eb09..7e3a19043 100644 --- a/tests/ingest/test_kafka_manager.py +++ b/tests/ingest/test_kafka_manager.py @@ -686,6 +686,54 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 assert recorded_input_params["kafka_message"]["keys"] == ["alarmData"] +@pytest.mark.parametrize("tool_context_required", [False, "false"]) +@pytest.mark.asyncio +async def test_trigger_workflow_can_skip_unneeded_tool_context( + monkeypatch: pytest.MonkeyPatch, + trigger_tool_context: SimpleNamespace, + tool_context_required: object, +) -> None: + manager = kafka_manager.KafkaManager() + captured_run_kwargs: dict = {} + + async def _fake_create_execution_record( # noqa: ANN001 + workflow_id, *, input_params=None, exec_id=None + ): + return {"id": "exec-no-context", "workflowId": workflow_id, "inputParams": input_params} + + async def _fake_record_execution_result( # noqa: ANN001 + workflow_id, exec_id, exec_data, *, steps=None + ): + return None + + def _fake_run_workflow(**kwargs): # noqa: ANN003 + captured_run_kwargs.update(kwargs) + return SimpleNamespace( + status="SUCCEEDED", + error=None, + outputs={"ok": True}, + history=[], + last_node_id="done", + steps=1, + ) + + monkeypatch.setattr(kafka_manager, "create_execution_record", _fake_create_execution_record) + monkeypatch.setattr(kafka_manager, "record_execution_result", _fake_record_execution_result) + monkeypatch.setattr(kafka_manager, "run_workflow", _fake_run_workflow) + + await manager._trigger_workflow( + "wf-no-context", + {"start": "receive_alert", "nodes": [], "edges": []}, + {"id": 1}, + "kafka_message", + {"tool_context_required": tool_context_required}, + ) + + assert captured_run_kwargs["tool_context"] is None + trigger_tool_context.builder.assert_not_awaited() + trigger_tool_context.cleanup.assert_not_awaited() + + @pytest.mark.asyncio async def test_trigger_workflow_applies_mapping_and_filter( monkeypatch: pytest.MonkeyPatch, From 431968cd1fa1b2b137c1ca12afdbad28f6000722 Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Tue, 25 Aug 2026 18:38:49 +0800 Subject: [PATCH 20/63] fix(workflow): avoid inherited connection close after fork --- flocks/workflow/store.py | 17 ++++++++------- tests/workflow/test_workflow_store.py | 30 ++++++++++++++++++++------- 2 files changed, 32 insertions(+), 15 deletions(-) diff --git a/flocks/workflow/store.py b/flocks/workflow/store.py index f7efbd0a1..129d62131 100644 --- a/flocks/workflow/store.py +++ b/flocks/workflow/store.py @@ -66,10 +66,9 @@ async def init(cls) -> None: db_path = cls.get_db_path() if cls._initialized and cls._init_pid == current_pid and cls._db_path == db_path: return - if cls._initialized and ( - (cls._init_pid is not None and cls._init_pid != current_pid) - or (cls._db_path is not None and cls._db_path != db_path) - ): + pid_changed = cls._initialized and cls._init_pid is not None and cls._init_pid != current_pid + db_path_changed = cls._initialized and cls._db_path is not None and cls._db_path != db_path + if pid_changed or db_path_changed: log.warn( "workflow.store.fork_detected", { @@ -79,14 +78,16 @@ async def init(cls) -> None: "new_db_path": str(db_path), }, ) - if cls._conn: - await cls._conn.close() - if cls._completion_conn: - await cls._completion_conn.close() + if not pid_changed: + if cls._conn: + await cls._conn.close() + if cls._completion_conn: + await cls._completion_conn.close() cls._conn = None cls._completion_conn = None cls._initialized = False cls._init_pid = None + cls._completion_lock = None await Storage._ensure_init() db_path.parent.mkdir(parents=True, exist_ok=True) diff --git a/tests/workflow/test_workflow_store.py b/tests/workflow/test_workflow_store.py index 6608db1b3..b33234648 100644 --- a/tests/workflow/test_workflow_store.py +++ b/tests/workflow/test_workflow_store.py @@ -3,6 +3,7 @@ import asyncio import os from pathlib import Path +from unittest.mock import AsyncMock import pytest @@ -294,14 +295,29 @@ async def cancel_commit() -> None: @pytest.mark.asyncio -async def test_completion_connection_reinitializes_after_pid_change() -> None: +async def test_pid_change_drops_inherited_connections_without_closing( + monkeypatch: pytest.MonkeyPatch, +) -> None: await WorkflowStore.init() - original_connection = await WorkflowStore.raw_completion_db() + original_db = await WorkflowStore.raw_db() + original_completion_db = await WorkflowStore.raw_completion_db() + original_db_close = original_db.close + original_completion_db_close = original_completion_db.close + db_close = AsyncMock() + completion_db_close = AsyncMock() + monkeypatch.setattr(original_db, "close", db_close) + monkeypatch.setattr(original_completion_db, "close", completion_db_close) original_lock = WorkflowStore._completion_lock WorkflowStore._init_pid = -1 - refreshed_connection = await WorkflowStore.raw_completion_db() - - assert refreshed_connection is not original_connection - assert WorkflowStore._completion_lock is not original_lock - assert WorkflowStore._init_pid == os.getpid() + try: + refreshed_connection = await WorkflowStore.raw_completion_db() + + assert refreshed_connection is not original_completion_db + assert WorkflowStore._completion_lock is not original_lock + assert WorkflowStore._init_pid == os.getpid() + db_close.assert_not_awaited() + completion_db_close.assert_not_awaited() + finally: + await original_db_close() + await original_completion_db_close() From bd88b19ea70dc2af053eb87748ed04590304f496 Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Wed, 26 Aug 2026 11:33:19 +0800 Subject: [PATCH 21/63] chore(workflow): remove unused execution imports --- flocks/server/routes/workflow.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/flocks/server/routes/workflow.py b/flocks/server/routes/workflow.py index b1ca85243..dfda6b13c 100644 --- a/flocks/server/routes/workflow.py +++ b/flocks/server/routes/workflow.py @@ -61,8 +61,6 @@ normalize_execution_status as _normalize_execution_status, record_execution_result as _record_execution_result, resolve_execution_outcome as _resolve_execution_outcome, - workflow_execution_key as _workflow_execution_key, - workflow_execution_step_prefix as _workflow_execution_step_prefix, ) from flocks.workflow.io import load_workflow, dump_workflow from flocks.workflow.store import WorkflowStore From 43cf614203e60b81ac08a8a5e010eb99e827b7d9 Mon Sep 17 00:00:00 2001 From: John Yin <10972267+john-yin2333@user.noreply.gitee.com> Date: Fri, 28 Aug 2026 09:33:21 +0800 Subject: [PATCH 22/63] fix(webui): show extra memory root files --- webui/src/pages/Workspace/index.test.tsx | 47 +++++++++++++++++------- webui/src/pages/Workspace/index.tsx | 34 +++++++++++++++-- 2 files changed, 63 insertions(+), 18 deletions(-) diff --git a/webui/src/pages/Workspace/index.test.tsx b/webui/src/pages/Workspace/index.test.tsx index aa3d03005..111ecbd73 100644 --- a/webui/src/pages/Workspace/index.test.tsx +++ b/webui/src/pages/Workspace/index.test.tsx @@ -452,7 +452,7 @@ describe('WorkspacePage', () => { expect(mocks.toastSuccess).toHaveBeenCalledWith('Saved successfully'); }); - it('Memory 文件按 USER、Global、Project 和 Daily 层级展示', async () => { + it('Memory 文件按核心记忆、Project、Daily、其他根文件层级展示', async () => { mocks.listVisibleProjects.mockResolvedValue({ data: [{ id: 'prj_example', @@ -496,28 +496,38 @@ describe('WorkspacePage', () => { expect(daily).toBeInTheDocument(); expect(screen.queryByText('projects')).not.toBeInTheDocument(); expect(screen.queryByText('prj_stale/MEMORY.md')).not.toBeInTheDocument(); - expect(screen.queryByText('2026-04-07.md')).not.toBeInTheDocument(); - expect(screen.queryByText('test.md')).not.toBeInTheDocument(); + expect(screen.getByText('2026-04-07.md')).toBeInTheDocument(); + expect(screen.getByText('test.md')).toBeInTheDocument(); expect( projectMemory.compareDocumentPosition(daily) & Node.DOCUMENT_POSITION_FOLLOWING, ).toBeTruthy(); + expect( + daily.compareDocumentPosition(screen.getByText('2026-04-07.md')) & Node.DOCUMENT_POSITION_FOLLOWING, + ).toBeTruthy(); expect(screen.queryByText('2026-08-03.md')).not.toBeInTheDocument(); await user.click(projectMemory); expect(await screen.findByText('/Users/test/workspace/flocks-raven')).toBeInTheDocument(); - - await user.click(screen.getByRole('button', { name: /daily/ })); - expect(await screen.findByText('2026-08-03.md')).toBeInTheDocument(); }); - it('Memory 根目录的非规范文件不显示', async () => { + it('Memory 根目录的其他文件和目录会显示但默认只读', async () => { mocks.listMemory.mockResolvedValue({ data: [ - file('profile.pdf', 'profile.pdf', false), - file('logo.svg', 'logo.svg', false), - file('legacy.md', 'legacy.md'), + file('SHORT_MEMORY.md', 'SHORT_MEMORY.md'), + file('bak.txt', 'bak.txt'), + { + ...directory('archive', 'archive'), + children: [file('2026-08-18.md', 'archive/2026-08-18.md')], + }, ], }); + mocks.readMemoryFile.mockResolvedValue({ + data: { + path: 'bak.txt', + content: 'backup memory', + truncated: false, + }, + }); const user = userEvent.setup(); renderWithRouter(); @@ -526,11 +536,20 @@ describe('WorkspacePage', () => { await waitFor(() => { expect(mocks.listMemory).toHaveBeenCalled(); }); - expect(screen.queryByText('profile.pdf')).not.toBeInTheDocument(); - expect(screen.queryByText('logo.svg')).not.toBeInTheDocument(); - expect(screen.queryByText('legacy.md')).not.toBeInTheDocument(); + expect(screen.getByText('SHORT_MEMORY.md')).toBeInTheDocument(); + expect(screen.getByText('archive')).toBeInTheDocument(); + expect(screen.getByText('bak.txt')).toBeInTheDocument(); + + const archiveButton = screen.getByText('archive').closest('button'); + if (!archiveButton) throw new Error('Archive memory row button not found'); + await user.click(archiveButton); + expect(await screen.findByText('2026-08-18.md')).toBeInTheDocument(); + + await user.click(screen.getByText('bak.txt')); + expect(await screen.findByText('backup memory')).toBeInTheDocument(); + expect(mocks.readMemoryFile).toHaveBeenCalledWith('bak.txt'); + expect(screen.queryByTitle('Edit')).not.toBeInTheDocument(); expect(pdfMocks.getDocument).not.toHaveBeenCalled(); - expect(mocks.readMemoryFile).not.toHaveBeenCalled(); }); it('Memory 文本文件快速切换时忽略过期读取结果', async () => { diff --git a/webui/src/pages/Workspace/index.tsx b/webui/src/pages/Workspace/index.tsx index 42ebe481e..7e0d8aa1d 100644 --- a/webui/src/pages/Workspace/index.tsx +++ b/webui/src/pages/Workspace/index.tsx @@ -1585,6 +1585,20 @@ function memoryPathParts(path: string): string[] { return path.replace(/\\/g, '/').split('/'); } +function isEditableMemoryNode(node: WorkspaceNode): boolean { + const pathParts = memoryPathParts(node.path); + if (pathParts.length === 1) { + return node.path === 'USER.md' || node.path === 'MEMORY.md'; + } + if (pathParts.length === 2) { + return pathParts[0] === 'daily' && node.name.endsWith('.md'); + } + if (pathParts.length === 3) { + return pathParts[0] === 'projects' && pathParts[2] === 'MEMORY.md'; + } + return false; +} + function buildMemoryView( nodes: WorkspaceNode[], visibleProjects: WorkspaceProject[], @@ -1594,11 +1608,19 @@ function buildMemoryView( const projects = nodes.find((node) => node.path === 'projects'); const daily = nodes.find((node) => node.path === 'daily'); const projectById = new Map(visibleProjects.map((project) => [project.id, project])); + const consumedRootPaths = new Set(); const view: WorkspaceNode[] = []; - if (userMemory) view.push(userMemory); - if (globalMemory) view.push(globalMemory); + if (userMemory) { + view.push(userMemory); + consumedRootPaths.add(userMemory.path); + } + if (globalMemory) { + view.push(globalMemory); + consumedRootPaths.add(globalMemory.path); + } if (projects) { + consumedRootPaths.add(projects.path); const projectMemories = collectMemoryFiles(projects.children ?? []).flatMap((node) => { const pathParts = memoryPathParts(node.path); const project = pathParts.length === 3 ? projectById.get(pathParts[1]) : undefined; @@ -1613,7 +1635,11 @@ function buildMemoryView( }); view.push(...projectMemories); } - if (daily) view.push(daily); + if (daily) { + view.push(daily); + consumedRootPaths.add(daily.path); + } + view.push(...nodes.filter((node) => !consumedRootPaths.has(node.path))); return view; } @@ -1903,7 +1929,7 @@ function MemoryTab() {
{formatBytes(selected.size ?? 0)} {formatDate(selected.modified_at)} - {selected.is_text_file && !editing && !truncated && contentState === 'ready' && ( + {selected.is_text_file && isEditableMemoryNode(selected) && !editing && !truncated && contentState === 'ready' && (
{formatBytes(selected.size ?? 0)} {formatDate(selected.modified_at)} - {selected.is_text_file && isEditableMemoryNode(selected) && !editing && !truncated && contentState === 'ready' && ( + {selected.is_text_file && selected.editable === true && !editing && !truncated && contentState === 'ready' && ( + ))} + + ); +} + +function SkipConfirmDialog({ + t, + target, + onCancel, + onConfirm, +}: { + t: (key: string, options?: any) => string; + target: SkipTarget; + onCancel: () => void; + onConfirm: () => void; +}) { + if (!target) return null; + + const isModel = target === 'model'; + + return ( +
+
+
+ +
+

+ {isModel ? t('onboarding.bootstrap.skipModelTitle') : t('onboarding.bootstrap.skipIntelTitle')} +

+

+ {isModel ? t('onboarding.bootstrap.skipModelDescription') : t('onboarding.bootstrap.skipIntelDescription')} +

+
+
+
+ + +
+
+
+ ); +} + export default function OnboardingModal({ onClose }: OnboardingModalProps) { const { t } = useTranslation('common'); const navigate = useNavigate(); - const [hasLLM, setHasLLM] = useState(null); + const [step, setStep] = useState('model'); const [catalog, setCatalog] = useState([]); + const [statusLoading, setStatusLoading] = useState(true); + const [hasLLM, setHasLLM] = useState(null); const [starting, setStarting] = useState(false); const [startStatus, setStartStatus] = useState(null); - const [dontShowAgain, setDontShowAgain] = useState(() => isOnboardingDismissed()); - const [resolvedDefaultModel, setResolvedDefaultModel] = useState(null); + const [skipConfirm, setSkipConfirm] = useState(null); + const [resolvedDefaultModel, setResolvedDefaultModel] = useState(null); const [primaryProviderId, setPrimaryProviderId] = useState('threatbook-cn-llm'); const [primaryModelId, setPrimaryModelId] = useState(''); const [primaryApiKey, setPrimaryApiKey] = useState(''); @@ -214,36 +329,70 @@ export default function OnboardingModal({ onClose }: OnboardingModalProps) { const [primarySaving, setPrimarySaving] = useState(false); const [primaryConfigured, setPrimaryConfigured] = useState(false); const [primaryStatus, setPrimaryStatus] = useState(null); - const [primarySectionCollapsed, setPrimarySectionCollapsed] = useState(false); const [primaryEditing, setPrimaryEditing] = useState(false); - - const [optionalThreatBookRegion, setOptionalThreatBookRegion] = useState('cn'); - const [optionalThreatBookApiKey, setOptionalThreatBookApiKey] = useState(''); - const [optionalSaving, setOptionalSaving] = useState(false); - const [optionalConfigured, setOptionalConfigured] = useState(false); - const [optionalStatus, setOptionalStatus] = useState(null); - const [optionalSectionCollapsed, setOptionalSectionCollapsed] = useState(false); - const [optionalThreatBookLoaded, setOptionalThreatBookLoaded] = useState(false); - - useEffect(() => { - defaultModelAPI.getResolved() - .then((res) => { - const resolved = { - providerId: res.data.provider_id, - modelId: res.data.model_id, - }; - setResolvedDefaultModel(resolved); - setPrimaryProviderId(resolved.providerId); - setPrimaryModelId(resolved.modelId); - setPrimaryConfigured(true); - setPrimarySectionCollapsed(true); + const [modelRegion, setModelRegion] = useState('cn'); + const [modelSkipped, setModelSkipped] = useState(false); + + const [intelRegion, setIntelRegion] = useState('cn'); + const [intelApiKey, setIntelApiKey] = useState(''); + const [intelSaving, setIntelSaving] = useState(false); + const [intelConfigured, setIntelConfigured] = useState(false); + const [intelStatus, setIntelStatus] = useState(null); + const [intelEditing, setIntelEditing] = useState(false); + const [intelSkipped, setIntelSkipped] = useState(false); + const [intelRuntimeStatus, setIntelRuntimeStatus] = useState(null); + + const refreshOnboardingStatus = useCallback(async (silent = false) => { + if (!silent) setStatusLoading(true); + try { + const res = await onboardingAPI.getStatus(); + const defaultModel = res.data.default_model; + setHasLLM(res.data.completed); + setPrimaryConfigured(Boolean(res.data.has_default_model)); + setResolvedDefaultModel(defaultModel + ? { + providerId: defaultModel.provider_id, + modelId: defaultModel.model_id, + } + : null); + if (defaultModel) { + setPrimaryProviderId(defaultModel.provider_id); + setPrimaryModelId(defaultModel.model_id); + setModelRegion(regionForProvider(defaultModel.provider_id)); setPrimaryEditing(false); + } + if (res.data.threatbook_intel) { + const intel = res.data.threatbook_intel; + setIntelRuntimeStatus(intel); + setIntelConfigured(intel.configured); + setIntelRegion(intel.region || 'cn'); + if (intel.configured) setIntelEditing(false); + } + } catch { + try { + const resolved = await defaultModelAPI.getResolved(); setHasLLM(true); - }) - .catch(() => { + setPrimaryConfigured(true); + setResolvedDefaultModel({ + providerId: resolved.data.provider_id, + modelId: resolved.data.model_id, + }); + setPrimaryProviderId(resolved.data.provider_id); + setPrimaryModelId(resolved.data.model_id); + setModelRegion(regionForProvider(resolved.data.provider_id)); + setPrimaryEditing(false); + } catch { setHasLLM(false); - }); + setPrimaryConfigured(false); + setResolvedDefaultModel(null); + } + } finally { + if (!silent) setStatusLoading(false); + } + }, []); + useEffect(() => { + refreshOnboardingStatus(); catalogAPI.list() .then((res) => { setCatalog(res.data.providers || []); @@ -251,47 +400,59 @@ export default function OnboardingModal({ onClose }: OnboardingModalProps) { .catch(() => { setCatalog([]); }); - - }, []); + }, [refreshOnboardingStatus]); const primaryProviders = useMemo(() => { - const filtered = catalog.filter((provider) => - THREATBOOK_PROVIDER_IDS.includes(provider.id as any) + return catalog.filter((provider) => + isThreatBookProvider(provider.id) || provider.id === 'openai-compatible' || provider.id === resolvedDefaultModel?.providerId || provider.models.length > 0 ); - - const preferredOrder = ['threatbook-cn-llm', 'threatbook-io-llm']; - return filtered.sort((a, b) => { - const aIndex = preferredOrder.indexOf(a.id); - const bIndex = preferredOrder.indexOf(b.id); - if (aIndex !== -1 || bIndex !== -1) { - return (aIndex === -1 ? 999 : aIndex) - (bIndex === -1 ? 999 : bIndex); - } - return a.name.localeCompare(b.name); - }); }, [catalog, resolvedDefaultModel?.providerId]); + const providerOptions = useMemo(() => { + const hasThreatBook = primaryProviders.some((provider) => isThreatBookProvider(provider.id)) + || isThreatBookProvider(primaryProviderId); + const nonThreatBook = primaryProviders + .filter((provider) => !isThreatBookProvider(provider.id)) + .sort((a, b) => a.name.localeCompare(b.name)); + + return { + hasThreatBook, + nonThreatBook, + }; + }, [primaryProviderId, primaryProviders]); + + const primaryProviderIsThreatBook = useMemo( + () => isThreatBookProvider(primaryProviderId), + [primaryProviderId] + ); + const selectedPrimaryProvider = useMemo( () => primaryProviders.find((provider) => provider.id === primaryProviderId) || null, [primaryProviderId, primaryProviders] ); - const primaryProviderIsThreatBook = useMemo( - () => THREATBOOK_PROVIDER_IDS.includes(primaryProviderId as any), - [primaryProviderId] + const selectedPrimaryModel = useMemo( + () => selectedPrimaryProvider?.models.find((model) => model.id === primaryModelId) || null, + [primaryModelId, selectedPrimaryProvider] ); - const primaryThreatBookRegion: OnboardingRegion = primaryProviderId === 'threatbook-io-llm' ? 'global' : 'cn'; - - const primaryApiPlaceholder = useMemo(() => { - if (primaryProviderIsThreatBook) return t('onboarding.bootstrap.tbPlaceholder'); - return ( - selectedPrimaryProvider?.credential_schemas?.[0]?.fields?.find((field) => field.name === 'api_key')?.placeholder - || t('onboarding.bootstrap.thirdPartyKeyPlaceholder') - ); - }, [primaryProviderIsThreatBook, selectedPrimaryProvider, t]); + const primaryResolvedProviderId = resolvedDefaultModel?.providerId || primaryProviderId; + const primaryResolvedModelId = resolvedDefaultModel?.modelId || primaryModelId; + const primaryResolvedProvider = useMemo( + () => primaryProviders.find((provider) => provider.id === primaryResolvedProviderId) || null, + [primaryProviders, primaryResolvedProviderId] + ); + const primaryResolvedModel = useMemo( + () => primaryResolvedProvider?.models.find((model) => model.id === primaryResolvedModelId) || null, + [primaryResolvedModelId, primaryResolvedProvider] + ); + const primaryResolvedProviderIsThreatBook = useMemo( + () => isThreatBookProvider(primaryResolvedProviderId), + [primaryResolvedProviderId] + ); const needsPrimaryBaseUrl = useMemo(() => { if (!selectedPrimaryProvider || primaryProviderIsThreatBook) return false; @@ -300,17 +461,23 @@ export default function OnboardingModal({ onClose }: OnboardingModalProps) { ); }, [primaryProviderIsThreatBook, selectedPrimaryProvider]); + const primaryApiPlaceholder = useMemo(() => { + if (primaryProviderIsThreatBook) return t('onboarding.bootstrap.modelKeyPlaceholder'); + return ( + selectedPrimaryProvider?.credential_schemas?.[0]?.fields?.find((field) => field.name === 'api_key')?.placeholder + || t('onboarding.bootstrap.thirdPartyKeyPlaceholder') + ); + }, [primaryProviderIsThreatBook, selectedPrimaryProvider, t]); + useEffect(() => { - if (!primaryProviderId && primaryProviders.length > 0) { - setPrimaryProviderId(primaryProviders[0].id); + if (!primaryProviderId && providerOptions.hasThreatBook) { + setPrimaryProviderId(providerIdForRegion(modelRegion)); } - }, [primaryProviderId, primaryProviders]); + }, [modelRegion, primaryProviderId, providerOptions.hasThreatBook]); useEffect(() => { if (!selectedPrimaryProvider) { - if (!primaryConfigured) { - setPrimaryModelId(''); - } + if (!primaryConfigured) setPrimaryModelId(''); return; } const hasModels = selectedPrimaryProvider.models.length > 0; @@ -328,44 +495,22 @@ export default function OnboardingModal({ onClose }: OnboardingModalProps) { }, [ needsPrimaryBaseUrl, primaryBaseUrl, + primaryConfigured, primaryModelId, primaryProviderIsThreatBook, selectedPrimaryProvider, ]); - const selectedPrimaryModel = useMemo( - () => selectedPrimaryProvider?.models.find((model) => model.id === primaryModelId) || null, - [primaryModelId, selectedPrimaryProvider] - ); - const primaryResolvedProviderId = resolvedDefaultModel?.providerId || primaryProviderId; - const primaryResolvedModelId = resolvedDefaultModel?.modelId || primaryModelId; - const primaryResolvedProvider = useMemo( - () => primaryProviders.find((provider) => provider.id === primaryResolvedProviderId) || null, - [primaryProviders, primaryResolvedProviderId] - ); - const primaryResolvedModel = useMemo( - () => primaryResolvedProvider?.models.find((model) => model.id === primaryResolvedModelId) || null, - [primaryResolvedModelId, primaryResolvedProvider] - ); - const primaryResolvedProviderIsThreatBook = useMemo( - () => THREATBOOK_PROVIDER_IDS.includes(primaryResolvedProviderId as any), - [primaryResolvedProviderId] - ); - - const getProviderLabel = (provider: CatalogProvider) => { - if (provider.id === 'threatbook-cn-llm') return t('onboarding.bootstrap.providerThreatBookCn'); - if (provider.id === 'threatbook-io-llm') return t('onboarding.bootstrap.providerThreatBookGlobal'); - return provider.name; + const getProviderLabel = (provider: CatalogProvider | null, providerId?: string) => { + const id = provider?.id || providerId || ''; + if (isThreatBookProvider(id)) return t('onboarding.bootstrap.providerThreatBookFree'); + return provider?.name || id; }; const primaryConfiguredSummary = useMemo(() => { - if (!primaryConfigured || !primaryProviderId) return ''; - - const providerLabel = selectedPrimaryProvider - ? getProviderLabel(selectedPrimaryProvider) - : primaryProviderId; - const modelLabel = selectedPrimaryModel?.name || primaryModelId; - + if (!primaryConfigured || !primaryResolvedProviderId) return ''; + const providerLabel = getProviderLabel(primaryResolvedProvider, primaryResolvedProviderId); + const modelLabel = primaryResolvedModel?.name || primaryResolvedModelId; if (!providerLabel || !modelLabel) return ''; return t('onboarding.bootstrap.primaryConfiguredSummary', { provider: providerLabel, @@ -373,92 +518,54 @@ export default function OnboardingModal({ onClose }: OnboardingModalProps) { }); }, [ primaryConfigured, - primaryModelId, - primaryProviderId, - selectedPrimaryModel, - selectedPrimaryProvider, + primaryResolvedModel, + primaryResolvedModelId, + primaryResolvedProvider, + primaryResolvedProviderId, t, ]); - const primaryConfiguredRegionLabel = useMemo(() => { - if (!primaryResolvedProviderIsThreatBook) return ''; - return primaryResolvedProviderId === 'threatbook-io-llm' - ? t('onboarding.bootstrap.regionGlobal') - : t('onboarding.bootstrap.regionChina'); - }, [primaryResolvedProviderId, primaryResolvedProviderIsThreatBook, t]); - const primaryConfiguredDetailsHint = useMemo(() => { - if (!primaryConfigured) return ''; - if (primaryResolvedProviderId === 'threatbook-cn-llm') { - return t('onboarding.bootstrap.configuredThreatBookCnDetails'); - } - if (primaryResolvedProviderId === 'threatbook-io-llm') { - return t('onboarding.bootstrap.configuredThreatBookGlobalDetails'); - } - return t('onboarding.bootstrap.configuredThirdPartyDetails'); - }, [primaryConfigured, primaryResolvedProviderId, t]); - const optionalConfiguredSummary = useMemo(() => { - if (!optionalConfigured) return ''; - return optionalThreatBookRegion === 'cn' - ? t('onboarding.bootstrap.optionalThreatBookCnSuccess') - : t('onboarding.bootstrap.optionalThreatBookGlobalSuccess'); - }, [optionalConfigured, optionalThreatBookRegion, t]); + const primaryConfiguredStatusValue = useMemo(() => { + if (!primaryConfigured) return t('onboarding.bootstrap.statusNotConfigured'); + return hasLLM + ? t('onboarding.bootstrap.statusVerified') + : t('onboarding.bootstrap.statusSavedNeedsVerify'); + }, [hasLLM, primaryConfigured, t]); + + const intelMcpStatusValue = useMemo(() => { + if (!intelRuntimeStatus || !intelRuntimeStatus.mcp_configured) return t('onboarding.bootstrap.statusNotConfigured'); + if (intelRuntimeStatus.mcp_connected) return t('onboarding.bootstrap.statusConnected'); + if (intelRuntimeStatus.mcp_status === 'error') return t('onboarding.bootstrap.statusError'); + if (intelRuntimeStatus.mcp_status === 'disabled') return t('onboarding.bootstrap.statusDisabled'); + return t('onboarding.bootstrap.statusConfigured'); + }, [intelRuntimeStatus, t]); + + const currentIntelCapabilities = useMemo(() => { + const matrix = intelRuntimeStatus?.service_matrix || { cn: ['api', 'mcp'], global: ['api'] }; + return matrix[intelRegion] || []; + }, [intelRegion, intelRuntimeStatus?.service_matrix]); const canSavePrimary = primaryProviderIsThreatBook ? Boolean(primaryApiKey.trim()) : Boolean(primaryProviderId && primaryModelId && primaryApiKey.trim()); - - const canSaveOptionalThreatBook = primaryConfigured && !primaryProviderIsThreatBook && Boolean(optionalThreatBookApiKey.trim()); - const canStart = hasLLM === true || primaryConfigured; - const showOptionalThreatBookSection = !primaryProviderIsThreatBook && primaryConfigured; + const canSaveIntel = Boolean(intelApiKey.trim()); const showPrimaryConfiguredDetails = primaryConfigured && !primaryEditing; + const showIntelConfiguredDetails = intelConfigured && !intelEditing; - useEffect(() => { - if (!showOptionalThreatBookSection || optionalThreatBookLoaded) return; - - let cancelled = false; - - Promise.allSettled([ - providerAPI.getServiceCredentials('threatbook-cn'), - providerAPI.getServiceCredentials('threatbook-io'), - mcpAPI.getCredentials('threatbook_mcp'), - ]).then((results) => { - if (cancelled) return; - - const cnServiceConfigured = results[0].status === 'fulfilled' && results[0].value.data.has_credential; - const globalServiceConfigured = results[1].status === 'fulfilled' && results[1].value.data.has_credential; - const cnMcpConfigured = results[2].status === 'fulfilled' && results[2].value.data.has_credential; - - if (cnServiceConfigured && cnMcpConfigured) { - setOptionalThreatBookRegion('cn'); - setOptionalConfigured(true); - setOptionalSectionCollapsed(true); - } else if (globalServiceConfigured) { - setOptionalThreatBookRegion('global'); - setOptionalConfigured(true); - setOptionalSectionCollapsed(true); - } - - setOptionalThreatBookLoaded(true); - }); - - return () => { - cancelled = true; - }; - }, [optionalThreatBookLoaded, showOptionalThreatBookSection]); - - const buildPrimaryPayload = (threatbookKey?: string): OnboardingRequest => { + const buildPrimaryPayload = (): OnboardingRequest => { if (primaryProviderIsThreatBook) { return { - region: primaryThreatBookRegion, + region: modelRegion, use_threatbook_model: true, threatbook_api_key: primaryApiKey.trim() || null, + threatbook_model_only: true, }; } return { - region: optionalThreatBookRegion, + region: modelRegion, use_threatbook_model: false, - threatbook_api_key: threatbookKey?.trim() || null, + threatbook_api_key: null, third_party_llm: { provider_id: primaryProviderId, api_key: primaryApiKey.trim(), @@ -469,10 +576,10 @@ export default function OnboardingModal({ onClose }: OnboardingModalProps) { }; }; - const buildOptionalThreatBookPayload = (): OnboardingRequest => ({ - region: optionalThreatBookRegion, + const buildIntelPayload = (): OnboardingRequest => ({ + region: intelRegion, use_threatbook_model: false, - threatbook_api_key: optionalThreatBookApiKey.trim() || null, + threatbook_api_key: intelApiKey.trim() || null, threatbook_services_only: true, }); @@ -496,52 +603,18 @@ export default function OnboardingModal({ onClose }: OnboardingModalProps) { validation, }); - const handleToggleDismiss = (checked: boolean) => { - setDontShowAgain(checked); - if (checked) { - dismissOnboarding(); + const handlePrimaryProviderChange = (providerId: string) => { + if (providerId === THREATBOOK_FREE_PROVIDER_OPTION) { + setPrimaryProviderId(providerIdForRegion(modelRegion)); } else { - undismissOnboarding(); + setPrimaryProviderId(providerId); } - }; - - const handlePrimaryProviderChange = (providerId: string) => { - setPrimaryProviderId(providerId); setPrimaryApiKey(''); setPrimaryBaseUrl(''); setPrimaryStatus(null); setPrimaryConfigured(false); - setPrimarySectionCollapsed(false); - setOptionalStatus(null); - setOptionalConfigured(false); - setOptionalThreatBookLoaded(false); - }; - - const handleTogglePrimarySection = () => { - if (primarySectionCollapsed) { - setPrimarySectionCollapsed(false); - setPrimaryEditing(false); - return; - } - setPrimarySectionCollapsed(true); - setPrimaryEditing(false); - }; - - const handleEditPrimary = () => { setPrimaryEditing(true); - setPrimaryStatus(null); - }; - - const handleBackToPrimaryDetails = () => { - if (!primaryConfigured) return; - if (resolvedDefaultModel) { - setPrimaryProviderId(resolvedDefaultModel.providerId); - setPrimaryModelId(resolvedDefaultModel.modelId); - } - setPrimaryApiKey(''); - setPrimaryBaseUrl(''); - setPrimaryStatus(null); - setPrimaryEditing(false); + setModelSkipped(false); }; const handleSavePrimary = async () => { @@ -549,13 +622,25 @@ export default function OnboardingModal({ onClose }: OnboardingModalProps) { setPrimarySaving(true); setPrimaryStatus(null); - setOptionalStatus(null); setStartStatus(null); try { - const payload = buildPrimaryPayload(); - const validateRes = await onboardingAPI.validate(payload); - const validateData = validateRes.data; + let payload = buildPrimaryPayload(); + let validateRes = await onboardingAPI.validate(payload); + let validateData = validateRes.data; + + if ( + primaryProviderIsThreatBook + && validateData.error_code === 'region_mismatch' + && validateData.suggested_region + ) { + payload = { + ...payload, + region: validateData.suggested_region, + }; + validateRes = await onboardingAPI.validate(payload); + validateData = validateRes.data; + } if (!validateData.can_apply) { setPrimaryStatus(buildErrorStatus(validateData, t('onboarding.bootstrap.testFailed'))); @@ -564,39 +649,34 @@ export default function OnboardingModal({ onClose }: OnboardingModalProps) { const applyRes = await onboardingAPI.apply(payload); const applyData = applyRes.data; - const savedProviderId = applyData.default_model?.provider_id || primaryProviderId; + const savedProviderId = applyData.default_model?.provider_id + || (primaryProviderIsThreatBook ? providerIdForRegion(payload.region) : primaryProviderId); const savedModelId = applyData.default_model?.model_id || primaryModelId || selectedPrimaryProvider?.models[0]?.id || ''; + setResolvedDefaultModel({ providerId: savedProviderId, modelId: savedModelId, }); + setPrimaryProviderId(savedProviderId); setPrimaryModelId(savedModelId); + setModelRegion(regionForProvider(savedProviderId)); setPrimaryConfigured(true); - setPrimarySectionCollapsed(true); setPrimaryEditing(false); + setModelSkipped(false); setHasLLM(true); const successMessage = primaryProviderIsThreatBook - ? (primaryThreatBookRegion === 'cn' - ? t('onboarding.bootstrap.primaryThreatBookCnSuccess') - : t('onboarding.bootstrap.primaryThreatBookGlobalSuccess')) + ? t('onboarding.bootstrap.primaryThreatBookSuccess') : t('onboarding.bootstrap.primaryThirdPartySuccess', { - provider: getProviderLabel(selectedPrimaryProvider!), + provider: getProviderLabel(selectedPrimaryProvider, primaryProviderId), model: selectedPrimaryModel?.name || primaryModelId, }); - if (primaryProviderIsThreatBook) { - setPrimaryStatus(buildSuccessStatus(validateData, applyData, successMessage)); - } else { - setPrimaryStatus({ - tone: 'success', - message: successMessage, - apply: applyData, - }); - } + setPrimaryStatus(buildSuccessStatus(validateData, applyData, successMessage)); + refreshOnboardingStatus(true); } catch (err: any) { setPrimaryStatus({ tone: 'error', @@ -607,40 +687,43 @@ export default function OnboardingModal({ onClose }: OnboardingModalProps) { } }; - const handleSaveOptionalThreatBook = async () => { - if (!canSaveOptionalThreatBook) return; + const handleSaveIntel = async () => { + if (!canSaveIntel) return; - setOptionalSaving(true); - setOptionalStatus(null); + setIntelSaving(true); + setIntelStatus(null); setStartStatus(null); try { - const payload = buildOptionalThreatBookPayload(); + const payload = buildIntelPayload(); const validateRes = await onboardingAPI.validate(payload); const validateData = validateRes.data; if (!validateData.can_apply) { - setOptionalStatus(buildErrorStatus(validateData, t('onboarding.bootstrap.serviceTestFailed'))); + setIntelStatus(buildErrorStatus(validateData, t('onboarding.bootstrap.serviceTestFailed'))); return; } const applyRes = await onboardingAPI.apply(payload); const applyData = applyRes.data; - setOptionalConfigured(true); - setOptionalSectionCollapsed(true); - - const successMessage = optionalThreatBookRegion === 'cn' - ? t('onboarding.bootstrap.optionalThreatBookCnSuccess') - : t('onboarding.bootstrap.optionalThreatBookGlobalSuccess'); - - setOptionalStatus(buildSuccessStatus(validateData, applyData, successMessage)); + setIntelConfigured(true); + setIntelEditing(false); + setIntelSkipped(false); + setIntelStatus(buildSuccessStatus( + validateData, + applyData, + intelRegion === 'cn' + ? t('onboarding.bootstrap.intelChinaSuccess') + : t('onboarding.bootstrap.intelGlobalSuccess'), + )); + refreshOnboardingStatus(true); } catch (err: any) { - setOptionalStatus({ + setIntelStatus({ tone: 'error', message: err?.response?.data?.message || err?.response?.data?.detail || err?.message || t('onboarding.bootstrap.saveError'), }); } finally { - setOptionalSaving(false); + setIntelSaving(false); } }; @@ -665,429 +748,483 @@ export default function OnboardingModal({ onClose }: OnboardingModalProps) { } }; - const isLoading = hasLLM === null; + const handleConfirmSkip = () => { + if (skipConfirm === 'model') { + setModelSkipped(true); + setStep('intel'); + setSkipConfirm(null); + return; + } + if (skipConfirm === 'intel') { + setIntelSkipped(true); + setSkipConfirm(null); + handleStart(); + } + }; - return ( -
-
-
-
- -

{t('onboarding.title')}

-
+ const renderModelStep = () => ( +
+
+

{t('onboarding.bootstrap.modelPageTitle')}

+

{t('onboarding.bootstrap.modelPageDescription')}

+
-
-
-
-
-
- - {t('onboarding.bootstrap.welcomeTitle')} -
-

- {t('onboarding.bootstrap.welcomeIntro')} -

-

- {t('onboarding.bootstrap.welcomeHint')} -

+ {statusLoading && ( +
+

{t('status.loading')}

+
+ )} + + {!statusLoading && showPrimaryConfiguredDetails && ( +
+
+
+

{t('onboarding.bootstrap.configuredDetailsTitle')}

+

{primaryConfiguredSummary}

+
- {!isLoading && hasLLM && ( -
- - {t('onboarding.bootstrap.configured')} +
+ + + +
+ + +
+ )} + + {!statusLoading && !showPrimaryConfiguredDetails && ( +
+
+ + +
+ + {!primaryProviderIsThreatBook && (selectedPrimaryProvider?.models?.length ? ( +
+ +
+ ) : ( + { + setPrimaryModelId(event.target.value); + setPrimaryStatus(null); + }} + placeholder={t('onboarding.bootstrap.thirdPartyModelIdPlaceholder')} + className="w-full rounded-lg border border-gray-200 bg-white px-3 py-2 text-xs transition-all placeholder-gray-300 focus:border-red-400 focus:outline-none focus:ring-2 focus:ring-red-400/50" + /> + ))} + + {needsPrimaryBaseUrl && ( + { + setPrimaryBaseUrl(event.target.value); + setPrimaryStatus(null); + }} + placeholder={t('onboarding.bootstrap.thirdPartyBaseUrlPlaceholder')} + className="w-full rounded-lg border border-gray-200 bg-white px-3 py-2 text-xs transition-all placeholder-gray-300 focus:border-red-400 focus:outline-none focus:ring-2 focus:ring-red-400/50" + /> )} -
+
+ {primaryConfigured && ( + + )} + { + setPrimaryApiKey(event.target.value); + setPrimaryStatus(null); + }} + placeholder={primaryApiPlaceholder} + className="min-w-0 flex-1 rounded-lg border border-gray-200 bg-white px-3 py-2 text-xs transition-all placeholder-gray-300 focus:border-red-400 focus:outline-none focus:ring-2 focus:ring-red-400/50" + /> + {primaryProviderIsThreatBook && ( + + + {t('onboarding.bootstrap.modelKeyLink')} + + )} +
- {isLoading && ( -
-

{t('status.loading')}

-
- )} + +
+ )} +
+ ); - {!isLoading && !primarySectionCollapsed && showPrimaryConfiguredDetails && ( -
-
-
-

- {t('onboarding.bootstrap.configuredDetailsTitle')} -

-

- {primaryConfiguredDetailsHint} -

-
- -
- -
-
-

- {t('onboarding.bootstrap.configuredStatusLabel')} -

-

- {t('onboarding.bootstrap.configuredReadyValue')} -

-
-
-

- {t('onboarding.bootstrap.configuredProviderLabel')} -

-

- {primaryResolvedProvider - ? getProviderLabel(primaryResolvedProvider) - : primaryResolvedProviderId} -

-
-
-

- {t('onboarding.bootstrap.configuredModelLabel')} -

-

- {primaryResolvedModel?.name || primaryResolvedModelId} -

-
- {primaryConfiguredRegionLabel && ( -
-

- {t('onboarding.bootstrap.configuredRegionLabel')} -

-

- {primaryConfiguredRegionLabel} -

-
- )} -
- - -
- )} + const renderIntelStep = () => ( +
+
+

{t('onboarding.bootstrap.intelPageTitle')}

+

{t('onboarding.bootstrap.intelPageDescription')}

+
- {!isLoading && !primarySectionCollapsed && !showPrimaryConfiguredDetails && ( - <> - - - {primaryProviderIsThreatBook && ( -
- - {primaryThreatBookRegion === 'cn' - ? t('onboarding.bootstrap.primaryThreatBookCnFreeHint') - : t('onboarding.bootstrap.primaryThreatBookGlobalFreeHint')} - - - - {primaryThreatBookRegion === 'cn' - ? t('onboarding.bootstrap.primaryThreatBookCnLink') - : t('onboarding.bootstrap.primaryThreatBookGlobalLink')} - -
- )} + {showIntelConfiguredDetails && ( +
+
+
+

{t('onboarding.bootstrap.configuredDetailsTitle')}

+

{t('onboarding.bootstrap.intelConfiguredHint')}

+
+ +
- {!primaryProviderIsThreatBook && (selectedPrimaryProvider?.models?.length ? ( - - ) : ( - { - setPrimaryModelId(e.target.value); - setPrimaryStatus(null); - }} - placeholder={t('onboarding.bootstrap.thirdPartyModelIdPlaceholder')} - className="w-full text-xs px-3 py-2 rounded-lg border border-gray-200 bg-white focus:outline-none focus:ring-2 focus:ring-red-400/50 focus:border-red-400 transition-all placeholder-gray-300" - /> - ))} +
+ + + + +
- {needsPrimaryBaseUrl && ( - { - setPrimaryBaseUrl(e.target.value); - setPrimaryStatus(null); - }} - placeholder={t('onboarding.bootstrap.thirdPartyBaseUrlPlaceholder')} - className="w-full text-xs px-3 py-2 rounded-lg border border-gray-200 bg-white focus:outline-none focus:ring-2 focus:ring-red-400/50 focus:border-red-400 transition-all placeholder-gray-300" - /> - )} + +
+ )} + + {!showIntelConfiguredDetails && ( +
+
+
+

{t('onboarding.bootstrap.intelRegionTitle')}

+

{t('onboarding.bootstrap.intelRegionHint')}

+
+ { + setIntelRegion(region); + setIntelStatus(null); + setIntelConfigured(false); + setIntelSkipped(false); + }} + chinaLabel={t('onboarding.bootstrap.intelRegionChina')} + globalLabel={t('onboarding.bootstrap.intelRegionGlobal')} + /> +
-
- {primaryConfigured && ( - - )} - { - setPrimaryApiKey(e.target.value); - setPrimaryStatus(null); - }} - placeholder={primaryApiPlaceholder} - className="flex-1 text-xs px-3 py-2 rounded-lg border border-gray-200 bg-white focus:outline-none focus:ring-2 focus:ring-red-400/50 focus:border-red-400 transition-all placeholder-gray-300" - /> - -
- - {primaryProviderIsThreatBook ? ( -

- {primaryThreatBookRegion === 'cn' - ? t('onboarding.bootstrap.primaryThreatBookCnHint') - : t('onboarding.bootstrap.primaryThreatBookGlobalHint')} -

- ) : (selectedPrimaryModel || primaryModelId) ? ( -

- {t('onboarding.bootstrap.thirdPartyModelHint', { model: selectedPrimaryModel?.name || primaryModelId })} -

- ) : null} - - { - const suggested = primaryStatus.validation?.suggested_region; - if (suggested === 'cn') handlePrimaryProviderChange('threatbook-cn-llm'); - if (suggested === 'global') handlePrimaryProviderChange('threatbook-io-llm'); - } - : null - } - /> - +
+ {currentIntelCapabilities.includes('api') && ( + + {t('onboarding.bootstrap.intelApiCapability')} + + )} + {currentIntelCapabilities.includes('mcp') && ( + + {t('onboarding.bootstrap.intelMcpCapability')} + )}
- {showOptionalThreatBookSection && ( -
+
+ {intelConfigured && ( + )} + { + setIntelApiKey(event.target.value); + setIntelStatus(null); + }} + placeholder={t('onboarding.bootstrap.intelKeyPlaceholder')} + className="min-w-0 flex-1 rounded-lg border border-gray-200 bg-white px-3 py-2 text-xs transition-all placeholder-gray-300 focus:border-red-400 focus:outline-none focus:ring-2 focus:ring-red-400/50" + /> + + + {t('onboarding.bootstrap.intelKeyLink')} + + +
- {!optionalSectionCollapsed && ( - <> -
-
- {(['cn', 'global'] as OnboardingRegion[]).map((candidate) => ( - - ))} -
+ { + if (intelStatus.validation?.suggested_region) { + setIntelRegion(intelStatus.validation.suggested_region); + setIntelStatus(null); + } + } + : null + } + /> +
+ )} + +
+

{t('onboarding.bootstrap.summaryTitle')}

+
+ + + +
+
+
+ ); - - - {optionalThreatBookRegion === 'cn' - ? t('onboarding.bootstrap.primaryThreatBookCnLink') - : t('onboarding.bootstrap.primaryThreatBookGlobalLink')} - -
+ const footerPrimaryAction = step === 'model' + ? { + label: t('onboarding.bootstrap.nextStep'), + disabled: statusLoading || primarySaving, + onClick: () => setStep('intel'), + } + : { + label: starting ? t('onboarding.startingButton') : t('onboarding.startButton'), + disabled: starting || primarySaving || intelSaving, + onClick: handleStart, + }; -
- { - setOptionalThreatBookApiKey(e.target.value); - setOptionalStatus(null); - setOptionalConfigured(false); - }} - placeholder={t('onboarding.bootstrap.tbPlaceholder')} - className="flex-1 text-xs px-3 py-2 rounded-lg border border-gray-200 bg-white focus:outline-none focus:ring-2 focus:ring-red-400/50 focus:border-red-400 transition-all placeholder-gray-300" - /> - -
+ return ( +
+
+
+
+ +

{t('onboarding.title')}

+
+ + 》 + + 》 + +
+
- {!optionalConfigured && ( -

{t('onboarding.bootstrap.tbSkipHint')}

- )} - - { - if (optionalStatus.validation?.suggested_region) { - setOptionalThreatBookRegion(optionalStatus.validation.suggested_region); - setOptionalStatus(null); - } - } - : null - } - /> - - )} -
- )} +
+
+ {step === 'model' ? renderModelStep() : renderIntelStep()}
-
+
{startStatus && (
)} -
- - + {step === 'intel' && ( + )} - + +
+ + setSkipConfirm(null)} + onConfirm={handleConfirmSkip} + />
); diff --git a/webui/src/components/layout/Layout.test.tsx b/webui/src/components/layout/Layout.test.tsx index 4b35498b7..910d8ee4a 100644 --- a/webui/src/components/layout/Layout.test.tsx +++ b/webui/src/components/layout/Layout.test.tsx @@ -37,6 +37,7 @@ const { getCredentials: vi.fn(), }, onboardingAPI: { + getStatus: vi.fn(), validate: vi.fn(), apply: vi.fn(), }, @@ -206,6 +207,32 @@ function makeProvider(id: string, name: string, models: Array<{ id: string; name }; } +function makeOnboardingStatus(overrides: Record = {}) { + return { + completed: true, + has_default_model: true, + default_model: { + provider_id: 'threatbook-cn-llm', + model_id: 'minimax-m2.7', + }, + threatbook_intel: { + configured: false, + region: null, + api_configured: false, + api_service_id: 'threatbook-cn', + mcp_configured: false, + mcp_connected: false, + mcp_status: 'not_configured', + mcp_name: 'threatbook_mcp', + service_matrix: { + cn: ['api', 'mcp'], + global: ['api'], + }, + }, + ...overrides, + }; +} + function renderHomeWithLayout() { return render( @@ -311,6 +338,9 @@ describe('Layout onboarding entry', () => { model_id: 'minimax-m2.7', }, }); + onboardingAPI.getStatus.mockResolvedValue({ + data: makeOnboardingStatus(), + }); catalogAPI.list.mockResolvedValue({ data: { @@ -363,12 +393,26 @@ describe('Layout onboarding entry', () => { await screen.findByText('onboarding.bootstrap.primaryConfiguredSummary'); - await user.click(screen.getByText('onboarding.bootstrap.primaryTitle')); - expect(screen.getByText('onboarding.bootstrap.configuredDetailsTitle')).toBeInTheDocument(); expect(screen.getByRole('button', { name: 'onboarding.bootstrap.editPrimary' })).toBeInTheDocument(); expect(screen.queryByRole('button', { name: 'onboarding.bootstrap.savePrimary' })).not.toBeInTheDocument(); - expect(screen.queryByPlaceholderText('onboarding.bootstrap.tbPlaceholder')).not.toBeInTheDocument(); + expect(screen.queryByPlaceholderText('onboarding.bootstrap.modelKeyPlaceholder')).not.toBeInTheDocument(); + }); + + it('auto-opens onboarding from backend status even when the old dismissed flag exists', async () => { + localStorage.setItem('flocks_onboarding_dismissed', 'true'); + onboardingAPI.getStatus.mockResolvedValue({ + data: makeOnboardingStatus({ + completed: false, + has_default_model: false, + default_model: null, + }), + }); + + renderHomeWithLayout(); + + expect(await screen.findByText('onboarding.bootstrap.modelPageTitle')).toBeInTheDocument(); + expect(screen.getByRole('button', { name: 'onboarding.bootstrap.savePrimary' })).toBeInTheDocument(); }); it('keeps standard pages out of a flex column content wrapper', async () => { diff --git a/webui/src/components/layout/Layout.tsx b/webui/src/components/layout/Layout.tsx index 5df046b9e..0febfb5cd 100644 --- a/webui/src/components/layout/Layout.tsx +++ b/webui/src/components/layout/Layout.tsx @@ -25,15 +25,13 @@ import { Loader2, type LucideIcon, } from 'lucide-react'; -import { useState, useEffect, useLayoutEffect, useCallback, useMemo, useRef, lazy, Suspense } from 'react'; +import { useState, useEffect, useCallback, useMemo, useRef, lazy, Suspense } from 'react'; import type { ComponentType, CSSProperties, KeyboardEvent as ReactKeyboardEvent, PointerEvent as ReactPointerEvent } from 'react'; import { useTranslation } from 'react-i18next'; +import { onboardingAPI } from '@/api/onboarding'; // Modals are only rendered after the user clicks/triggers them; pulling them // into the eager Layout chunk costs ~1.7k LOC + i18n keys + lucide icons that -// the home page never needs. To keep the lazy split effective, we don't -// re-import dismissal helpers from the modal modules (a static named import -// would force Rollup to bundle the whole module eagerly). -const ONBOARDING_DISMISSED_KEY = 'flocks_onboarding_dismissed'; +// the home page never needs. const COLLAPSED_NAV_SECTIONS_KEY = 'flocks_layout_collapsed_nav_sections'; const SIDEBAR_WIDTH_KEY = 'flocks_layout_sidebar_width'; const SIDEBAR_DEFAULT_WIDTH = 208; @@ -57,10 +55,6 @@ function lazyLayoutComponent( )); } -function isOnboardingDismissed(): boolean { - return localStorage.getItem(ONBOARDING_DISMISSED_KEY) === 'true'; -} - function readCollapsedNavSectionIds(): Set { try { const rawValue = localStorage.getItem(COLLAPSED_NAV_SECTIONS_KEY); @@ -380,12 +374,25 @@ export default function Layout() { updateSidebarWidth(sidebarWidth + (event.key === 'ArrowRight' ? 16 : -16)); }, [collapsed, sidebarWidth, updateSidebarWidth]); - // useLayoutEffect runs synchronously before paint, so there's no flash on initial load. - // It also re-runs when the user navigates back to /, covering both cases in one place. - useLayoutEffect(() => { - if (isHome && !isOnboardingDismissed()) { - setShowOnboarding(true); - } + useEffect(() => { + if (!isHome) return undefined; + + let cancelled = false; + onboardingAPI.getStatus() + .then((res) => { + if (!cancelled && !res.data.completed) { + setShowOnboarding(true); + } + }) + .catch(() => { + if (!cancelled) { + setShowOnboarding(true); + } + }); + + return () => { + cancelled = true; + }; }, [isHome]); const handleOpenOnboarding = useCallback(() => setShowOnboarding(true), []); diff --git a/webui/src/locales/en-US/common.json b/webui/src/locales/en-US/common.json index 64d34aa41..d33cc1f93 100644 --- a/webui/src/locales/en-US/common.json +++ b/webui/src/locales/en-US/common.json @@ -166,109 +166,92 @@ "sidebarLabel": "Setup Guide", "presetUser": "Hi Rex! I just started using Flocks — please help me with the initial setup.", "presetRex": "Hello! I'm Rex, Flocks' AI security operations assistant 👋\n\nGreat to meet you! I'll guide you through the initial setup — connecting security APIs, IM channels, NDR devices and more. It takes about 5–10 minutes and you can skip any step.\n\nReady? Let's get started 🚀", - "initialMessage": "Please start the onboarding process. After any IM channel is configured, explicitly guide me to open the corresponding channel in the Channels page, click \"Restart Connection\", click \"Save\" if there are any page changes, and continue only after the connection is restored.", + "initialMessage": "Please start the onboarding process. The model setup and ThreatBook intelligence setup pages have been completed or skipped. Continue guiding me through IM channels, device integration, operating scenarios, and the remaining setup. After any IM channel is configured, explicitly guide me to open the corresponding channel in the Channels page, click \"Restart Connection\", click \"Save\" if there are any page changes, and continue only after the connection is restored.", "bootstrap": { - "title": "A default model is required", - "subtitle": "Set up a default model first, then start chatting. If you choose a ThreatBook free model, Flocks will also configure the matching ThreatBook API and MCP services for you.", "welcomeTitle": "Welcome to Flocks", - "welcomeIntro": "Before you begin, you need to configure a default model. Once validation succeeds, you can start chatting right away and connect more services later.", - "welcomeHint": "ThreatBook offers free models for both Chinese and Global users. Chinese users can also enable ThreatBook MCP and API services together.", - "primaryTitle": "1. Configure model", - "primarySubtitle": "", - "primaryThreatBookCnFreeHint": "ThreatBook provides a free model plus Threatbook MCP / API services for Chinese users.", - "primaryThreatBookGlobalFreeHint": "ThreatBook provides a free model plus ThreatBook API service for Global users.", - "primaryThreatBookCnLink": "Chinese users claim API key", - "primaryThreatBookGlobalLink": "Global users claim API key", - "providerThreatBookCn": "ThreatBook-China free model", - "providerThreatBookGlobal": "ThreatBook Global free model", + "welcomeIntro": "Start with two core setup steps: a default model powers Rex conversations, while ThreatBook intelligence enables security intelligence tools. Either step can be skipped and finished later.", + "stepModel": "Configure model", + "stepIntel": "Configure ThreatBook intelligence", + "stepRex": "Rex guided setup", + "modelPageTitle": "Configure default model", + "modelPageDescription": "Flocks needs a default model to power Rex conversations, Agents, and workflows. You can use the ThreatBook free model or configure another model provider.", + "providerLabel": "Model provider", + "providerThreatBookFree": "ThreatBook free model", + "modelLabel": "Default model", + "modelRegionTitle": "API key region", + "modelRegionHint": "Used to match the available region for the ThreatBook free model.", + "modelKeyLink": "Claim API Key", + "modelKeyPlaceholder": "Paste your ThreatBook free model API key", "savePrimary": "Save & Verify Model", - "saveOptionalThreatBook": "Save & Verify ThreatBook Services", - "primaryThreatBookCnHint": "After validation, Flocks will also configure the ThreatBook-China model, API, and MCP services.", - "primaryThreatBookGlobalHint": "After validation, Flocks will also configure the ThreatBook Global model and API service.", - "primaryThreatBookCnSuccess": "Model verification succeeded. ThreatBook-China API and MCP have also been configured.", - "primaryThreatBookGlobalSuccess": "Model verification succeeded. ThreatBook Global API has also been configured.", - "primaryThirdPartySuccess": "Default model verification succeeded. {{provider}} / {{model}} has been configured. You can optionally enable ThreatBook services below.", + "primaryThreatBookSuccess": "Default model verification succeeded. ThreatBook free model has been configured.", + "primaryThirdPartySuccess": "Default model verification succeeded. {{provider}} / {{model}} has been configured.", "primaryConfiguredSummary": "Current configuration: {{provider}} / {{model}}", "configuredDetailsTitle": "Current configuration", "configuredStatusLabel": "Status", "configuredProviderLabel": "Provider", "configuredModelLabel": "Default model", "configuredRegionLabel": "Region", - "configuredReadyValue": "Configured and ready", - "configuredThreatBookCnDetails": "Your default model is already set to the ThreatBook-China free model. You can start right away, or open edit mode to change and re-verify it.", - "configuredThreatBookGlobalDetails": "Your default model is already set to the ThreatBook Global free model. You can start right away, or open edit mode to change and re-verify it.", - "configuredThirdPartyDetails": "Your default model is already configured. You can start right away, or open edit mode to change and re-verify it.", "editPrimary": "Edit configuration", + "editIntel": "Edit configuration", "backToConfiguredDetails": "Back to current configuration", - "optionalThreatBookTitle": "2. Optional: Enable ThreatBook services", - "optionalThreatBookSubtitle": "Even if you use another model as the default, you can still add ThreatBook API here. Chinese users can also enable MCP at the same time.", - "optionalThreatBookCnSuccess": "ThreatBook-China services verified successfully. API and MCP have both been configured.", - "optionalThreatBookGlobalSuccess": "ThreatBook Global services verified successfully. API has been configured.", - "tbLabel": "ThreatBook API Key (recommended)", - "freeBadge": "Limited Free", - "tbPlaceholder": "Paste your ThreatBook API Key", - "otherLLM": "Or use another LLM (expand)", - "saveKey": "Save & Verify", + "intelPageTitle": "Configure ThreatBook intelligence services", + "intelPageDescription": "Enable free ThreatBook API and MCP capabilities for IOC lookup, threat intelligence analysis, and security operations tools in Rex.", + "intelRegionTitle": "Choose service region", + "intelRegionHint": "Activation links and enabled capabilities vary by region.", + "intelRegionChina": "China", + "intelRegionGlobal": "Global", + "intelApiCapability": "Free API service", + "intelMcpCapability": "Free MCP service", + "intelKeyPlaceholder": "Paste your ThreatBook intelligence API key", + "intelKeyLink": "Claim API Key", + "intelApiLabel": "ThreatBook API", + "intelMcpLabel": "ThreatBook MCP", + "saveIntel": "Save & Verify Intelligence", + "intelChinaSuccess": "ThreatBook intelligence verified successfully. API and MCP have been configured.", + "intelGlobalSuccess": "ThreatBook intelligence verified successfully. API has been configured.", + "intelConfiguredHint": "ThreatBook intelligence configuration was detected. Edit it to replace the key or re-verify.", "testing": "Verifying...", - "testPassed": "Verified", "testFailed": "Model call failed — please check your API key", "serviceTestFailed": "Service verification failed — please check the API key and selected region", - "retryTest": "Retry", - "testPassedHint": "Verified — set as default model: {{model}}", - "startBlockedHint": "Please finish configuring and verifying a default model first", - "saved": "Saved", - "configured": "System default model configured — ready to start", "saveError": "Failed to save, please try again", "startError": "Failed to start a session, please try again", - "defaultModelHint": "Will verify model after saving: {{model}}", - "modelsPageHint": { - "prefix": "", - "link": "Go to Models page for more", - "suffix": "" - }, - "regionTitle": "1. Choose your region", - "regionSubtitle": "ThreatBook activation links, valid keys, and service presets depend on the selected region.", "regionChina": "China", "regionGlobal": "Global", - "modelStrategyTitle": "2. Choose a default model strategy", - "modelStrategySubtitle": "Use the free ThreatBook model, or keep another model provider as your default.", - "useThreatBookModel": "Use ThreatBook free model", - "useThreatBookModelHint": "ThreatBook will be set as the default LLM, and ThreatBook API plus Chinese MCP will be configured together.", - "useOtherModel": "Use another model provider", - "useOtherModelHint": "Configure a third-party default model, and optionally add a ThreatBook key for extra services.", - "tbKeyTitle": "3. ThreatBook key", - "tbRequiredHint": "ThreatBook is selected as the default model, so a region-matching ThreatBook key is required here.", - "tbOptionalHint": "Recommended. Even with another default model, this can still enable ThreatBook API and Chinese MCP.", - "requiredBadge": "Required", - "optionalBadge": "Optional", - "cnResourcesRequired": "For Chinese, this will configure: ThreatBook model + ThreatBook API + ThreatBook MCP.", - "cnResourcesOptional": "For Chinese, this can additionally enable: ThreatBook model + ThreatBook API + ThreatBook MCP.", - "globalResourcesRequired": "For International, this will configure: ThreatBook model + ThreatBook API.", - "globalResourcesOptional": "For International, this can additionally enable: ThreatBook model + ThreatBook API.", - "tbSkipHint": "This step is optional. You can skip it and still start chatting, but ThreatBook model, MCP, and API services will remain disabled for now.", - "thirdPartyTitle": "4. Third-party model", - "thirdPartySubtitle": "Choose another provider, enter its key, and select the default model.", "thirdPartyKeyPlaceholder": "Paste the third-party model API key", "thirdPartyBaseUrlPlaceholder": "Enter the model service base URL if needed", "thirdPartyModelIdPlaceholder": "Enter the model ID (e.g. gpt-4o, deepseek-v4-flash)", - "thirdPartyModelHint": "Will verify and set this as the default model: {{model}}", - "thirdPartyAdvancedHint": "For advanced custom setups such as OpenAI Compatible, use the Models page instead.", - "validateAndApply": "Validate & Apply", - "applying": "Applying...", - "applyFailed": "Failed to apply configuration", "switchToChina": "Switch to Chinese region", "switchToGlobal": "Switch to Global region", + "statusVerified": "Configured and verified", + "statusSavedNeedsVerify": "Configured, re-verification recommended", + "statusConfigured": "Configured", + "statusConnected": "Connected", + "statusError": "Configuration error", + "statusDisabled": "Disabled", + "statusNotConfigured": "Not configured", "statusPassed": "Passed", "statusFailed": "Failed", "statusSkipped": "Skipped", + "skipPage": "Skip", + "previousStep": "Previous", + "nextStep": "Next", + "skipModelTitle": "Skip model setup?", + "skipModelDescription": "Configure it later in Models. Without a default model, Rex conversations, Agents, and workflows may not be able to call an LLM.", + "skipIntelTitle": "Skip ThreatBook intelligence?", + "skipIntelDescription": "Configure it later in Intelligence MCP. You can still enter Rex chat, but IOC lookup and threat intelligence analysis from ThreatBook API/MCP will be unavailable for now.", + "returnToConfig": "Return to setup", + "confirmSkip": "Skip anyway", + "summaryTitle": "Setup summary", + "summaryModel": "Default model", + "summaryIntel": "ThreatBook intelligence", + "summaryNext": "Next", + "summaryNextValue": "Rex continues setup", "resourceLabels": { "threatbook_llm": "ThreatBook model", "threatbook_api": "ThreatBook API", "threatbook_mcp": "ThreatBook MCP", "third_party_llm": "Third-party model" - }, - "getKeyCN": "Chinese users — get key", - "getKeyIntl": "Global users — get key" + } }, "steps": { "securityApis": { "title": "Configure Security Tool APIs", "tags": ["ThreatBook", "VirusTotal", "FOFA", "URLScan", "Shodan"] }, diff --git a/webui/src/locales/zh-CN/common.json b/webui/src/locales/zh-CN/common.json index c5fee2597..a38bdd2c5 100644 --- a/webui/src/locales/zh-CN/common.json +++ b/webui/src/locales/zh-CN/common.json @@ -166,109 +166,92 @@ "sidebarLabel": "新手引导", "presetUser": "你好，Rex！我刚开始使用 Flocks，请帮我完成初始配置。", "presetRex": "你好！我是 Rex，Flocks 的 AI 安全运营助手 👋\n\n很高兴认识你！接下来我会引导你完成 Flocks 的初始配置，包括接入安全 API、IM 渠道和 NDR 设备等，整个过程大约需要 5–10 分钟，随时可以跳过任意步骤。\n\n准备好了吗？让我们正式开始吧 🚀", - "initialMessage": "请启动新手引导流程。完成 IM 渠道配置后，请明确引导我进入 Channels 页面对应通道，点击“重启连接”；如果页面里还有改动，再点击“保存”，并在连接恢复后继续下一步。", + "initialMessage": "请启动新手引导流程。模型配置和微步情报配置页面已完成或已跳过，请继续引导我完成 IM 渠道、设备接入、运营场景等剩余配置。完成 IM 渠道配置后，请明确引导我进入 Channels 页面对应通道，点击“重启连接”；如果页面里还有改动，再点击“保存”，并在连接恢复后继续下一步。", "bootstrap": { - "title": "需要先配置一个默认模型", - "subtitle": "模型配置完成后即可进入对话。若你使用微步免费模型，系统会同时帮你配置并验证微步的 API 和 MCP 能力。", "welcomeTitle": "欢迎使用 Flocks", - "welcomeIntro": "开始之前，首先需要配置一个默认模型。配置成功后即可进入对话，后续再按需接入更多安全服务。", - "welcomeHint": "微步为中国区和国际区用户都提供了免费模型，中国区用户还可额外启用微步 MCP 与 API 服务。", - "primaryTitle": "1. 配置模型", - "primarySubtitle": "", - "primaryThreatBookCnFreeHint": "微步为中国区提供免费模型和Threatbook MCP / API服务。", - "primaryThreatBookGlobalFreeHint": "微步为国际区提供免费模型和ThreatBook API服务。", - "primaryThreatBookCnLink": "中国区用户领取API Key", - "primaryThreatBookGlobalLink": "国际区用户领取API Key", - "providerThreatBookCn": "ThreatBook 中国区免费模型", - "providerThreatBookGlobal": "ThreatBook 国际区免费模型", + "welcomeIntro": "先完成两个核心配置：默认模型负责驱动 Rex 对话，微步情报服务负责提供安全情报工具能力。任一步都可以跳过，稍后再补齐。", + "stepModel": "配置模型", + "stepIntel": "配置微步情报", + "stepRex": "Rex 对话引导", + "modelPageTitle": "配置默认模型", + "modelPageDescription": "Flocks 需要一个默认模型来驱动 Rex 对话、Agent 和工作流。你可以使用 ThreatBook 免费模型，也可以配置其他模型。", + "providerLabel": "模型提供方", + "providerThreatBookFree": "ThreatBook 免费模型", + "modelLabel": "默认模型", + "modelRegionTitle": "API Key 所属区域", + "modelRegionHint": "用于匹配 ThreatBook 免费模型的可用区域。", + "modelKeyLink": "领取 API Key", + "modelKeyPlaceholder": "粘贴 ThreatBook 免费模型 API Key", "savePrimary": "保存并验证模型", - "saveOptionalThreatBook": "保存并验证微步服务", - "primaryThreatBookCnHint": "验证成功后，将同时配置 ThreatBook 中国区模型、API 和 MCP。", - "primaryThreatBookGlobalHint": "验证成功后，将同时配置 ThreatBook 国际区模型和 API。", - "primaryThreatBookCnSuccess": "模型验证成功，已同时配置 ThreatBook 中国区 API 和 MCP。", - "primaryThreatBookGlobalSuccess": "模型验证成功，已同时配置 ThreatBook 国际区 API。", - "primaryThirdPartySuccess": "默认模型验证成功，已配置 {{provider}} / {{model}}。如需启用微步安全服务，可继续完成下方可选步骤。", + "primaryThreatBookSuccess": "默认模型验证成功，已配置 ThreatBook 免费模型。", + "primaryThirdPartySuccess": "默认模型验证成功，已配置 {{provider}} / {{model}}。", "primaryConfiguredSummary": "当前配置：{{provider}} / {{model}}", "configuredDetailsTitle": "当前配置详情", "configuredStatusLabel": "配置状态", "configuredProviderLabel": "模型提供方", "configuredModelLabel": "默认模型", "configuredRegionLabel": "服务区域", - "configuredReadyValue": "已配置，可直接开始", - "configuredThreatBookCnDetails": "当前默认模型使用 ThreatBook 中国区免费模型，可直接开始；如需更换模型或重新验证，可进入编辑配置。", - "configuredThreatBookGlobalDetails": "当前默认模型使用 ThreatBook 国际区免费模型，可直接开始；如需更换模型或重新验证，可进入编辑配置。", - "configuredThirdPartyDetails": "当前默认模型已配置，可直接开始；如需更换模型或重新验证，可进入编辑配置。", "editPrimary": "编辑配置", + "editIntel": "编辑配置", "backToConfiguredDetails": "返回当前配置", - "optionalThreatBookTitle": "2. 可选：启用 ThreatBook 安全服务", - "optionalThreatBookSubtitle": "如果你当前使用其他模型，也可以在这里额外配置 ThreatBook 的 API 服务；中国区用户还可同时启用 MCP。", - "optionalThreatBookCnSuccess": "ThreatBook 中国区服务验证成功，已同时配置 API 和 MCP。", - "optionalThreatBookGlobalSuccess": "ThreatBook 国际区服务验证成功，已同时配置 API。", - "tbLabel": "ThreatBook API Key（推荐）", - "freeBadge": "限时免费", - "tbPlaceholder": "粘贴你的 ThreatBook API Key", - "otherLLM": "或使用其他 LLM（展开）", - "saveKey": "保存并验证", + "intelPageTitle": "配置微步情报服务", + "intelPageDescription": "免费启用微步 API 与 MCP 能力，为 Rex 提供 IOC 查询、威胁情报分析和安全运营工具能力。", + "intelRegionTitle": "选择服务区域", + "intelRegionHint": "不同区域的领取入口和可启用能力不同。", + "intelRegionChina": "国内", + "intelRegionGlobal": "国外", + "intelApiCapability": "免费 API 服务", + "intelMcpCapability": "免费 MCP 服务", + "intelKeyPlaceholder": "粘贴微步情报 API Key", + "intelKeyLink": "领取 API Key", + "intelApiLabel": "ThreatBook API", + "intelMcpLabel": "ThreatBook MCP", + "saveIntel": "保存并验证微步情报", + "intelChinaSuccess": "微步情报服务验证成功，已配置 ThreatBook API 和 MCP。", + "intelGlobalSuccess": "微步情报服务验证成功，已配置 ThreatBook API。", + "intelConfiguredHint": "当前已检测到微步情报配置；如需更换 Key 或重新验证，可进入编辑配置。", "testing": "验证中...", - "testPassed": "验证通过", "testFailed": "模型调用失败，请检查 Key 是否正确", "serviceTestFailed": "服务验证失败，请检查 Key 与区域是否正确", - "retryTest": "重新验证", - "testPassedHint": "验证通过，已设为默认模型：{{model}}", - "startBlockedHint": "请先完成默认模型配置并验证通过", - "saved": "已保存", - "configured": "系统默认模型已配置，可直接开始", "saveError": "保存失败，请重试", "startError": "启动会话失败，请重试", - "defaultModelHint": "保存后将验证模型可用性：{{model}}", - "modelsPageHint": { - "prefix": "", - "link": "前往 Models 页面管理更多模型", - "suffix": "" - }, - "regionTitle": "1. 选择你的区域", - "regionSubtitle": "ThreatBook 的激活入口、可用 key 和服务矩阵会随区域切换。", "regionChina": "中国区", "regionGlobal": "国际区", - "modelStrategyTitle": "2. 选择默认模型方案", - "modelStrategySubtitle": "你可以直接使用 ThreatBook 免费模型，或使用其他模型厂商。", - "useThreatBookModel": "使用 ThreatBook 免费模型", - "useThreatBookModelHint": "将默认使用 ThreatBook 的免费模型，并同步配置 ThreatBook API 与中国区 MCP。", - "useOtherModel": "使用其他模型厂商", - "useOtherModelHint": "可单独配置第三方默认模型，同时按需填写 ThreatBook key 启用微步服务。", - "tbKeyTitle": "3. ThreatBook key", - "tbRequiredHint": "当前已选择 ThreatBook 作为默认模型，因此这里必须填写对应区域的 key。", - "tbOptionalHint": "推荐填写。即使你使用其他模型，仍可额外启用 ThreatBook API，以及中国区 MCP。", - "requiredBadge": "必填", - "optionalBadge": "选填", - "cnResourcesRequired": "中国区填写后将配置：ThreatBook 模型 + ThreatBook API + ThreatBook MCP。", - "cnResourcesOptional": "中国区填写后可额外启用：ThreatBook 模型 + ThreatBook API + ThreatBook MCP。", - "globalResourcesRequired": "全球区填写后将配置：ThreatBook 模型 + ThreatBook API。", - "globalResourcesOptional": "全球区填写后可额外启用：ThreatBook 模型 + ThreatBook API。", - "tbSkipHint": "这一步是可选的。跳过后仍可进入对话，只是暂时不会启用微步提供的模型、MCP、API 服务。", - "thirdPartyTitle": "4. 第三方模型", - "thirdPartySubtitle": "选择一个其他模型厂商，填写其 key，并指定默认模型。", "thirdPartyKeyPlaceholder": "粘贴第三方模型 API Key", "thirdPartyBaseUrlPlaceholder": "填写模型服务 Base URL（如需要）", "thirdPartyModelIdPlaceholder": "输入模型 ID（例如 gpt-4o、deepseek-v4-flash）", - "thirdPartyModelHint": "保存后会验证并设为默认模型：{{model}}", - "thirdPartyAdvancedHint": "如需 OpenAI Compatible 等高级自定义接入，请前往 Models 页面配置。", - "validateAndApply": "验证并应用", - "applying": "应用中...", - "applyFailed": "应用失败，请重试", "switchToChina": "切换到中国区", "switchToGlobal": "切换到全球区", + "statusVerified": "已配置并验证通过", + "statusSavedNeedsVerify": "已配置，建议重新验证", + "statusConfigured": "已配置", + "statusConnected": "已连接", + "statusError": "配置异常", + "statusDisabled": "已禁用", + "statusNotConfigured": "未配置", "statusPassed": "通过", "statusFailed": "失败", "statusSkipped": "跳过", + "skipPage": "跳过", + "previousStep": "上一步", + "nextStep": "下一步", + "skipModelTitle": "跳过模型配置？", + "skipModelDescription": "稍后去模型清单配置。未配置默认模型时，Rex 对话、Agent 和工作流可能无法正常调用大模型。", + "skipIntelTitle": "跳过微步情报服务？", + "skipIntelDescription": "稍后在情报mcp完成配置。跳过后仍可进入 Rex 对话，但暂时无法使用微步情报 API/MCP 提供的 IOC 查询和情报分析能力。", + "returnToConfig": "返回配置", + "confirmSkip": "继续跳过", + "summaryTitle": "配置小结", + "summaryModel": "默认模型", + "summaryIntel": "微步情报", + "summaryNext": "下一步", + "summaryNextValue": "Rex 继续引导", "resourceLabels": { "threatbook_llm": "ThreatBook 模型", "threatbook_api": "ThreatBook API", "threatbook_mcp": "ThreatBook MCP", "third_party_llm": "第三方模型" - }, - "getKeyCN": "中国用户领取", - "getKeyIntl": "国际用户领取" + } }, "steps": { "securityApis": { "title": "配置安全工具 API", "tags": ["ThreatBook", "VirusTotal", "FOFA", "URLScan", "Shodan"] }, From 1e97a0d7bcae7ffec10c873aa0b38ae512390c53 Mon Sep 17 00:00:00 2001 From: John Yin <10972267+john-yin2333@user.noreply.gitee.com> Date: Sun, 6 Sep 2026 23:15:12 +0800 Subject: [PATCH 34/63] Refine ThreatBook intelligence onboarding page --- .../common/OnboardingModal.test.tsx | 7 ++- .../src/components/common/OnboardingModal.tsx | 58 +++++++++++++------ webui/src/components/layout/Layout.test.tsx | 2 +- webui/src/locales/en-US/common.json | 13 +++-- webui/src/locales/zh-CN/common.json | 13 +++-- 5 files changed, 60 insertions(+), 33 deletions(-) diff --git a/webui/src/components/common/OnboardingModal.test.tsx b/webui/src/components/common/OnboardingModal.test.tsx index 73d7b5897..630ac17df 100644 --- a/webui/src/components/common/OnboardingModal.test.tsx +++ b/webui/src/components/common/OnboardingModal.test.tsx @@ -107,7 +107,7 @@ function makeStatus(overrides: Record = {}) { mcp_name: 'threatbook_mcp', service_matrix: { cn: ['api', 'mcp'], - global: ['api'], + global: ['api', 'mcp'], }, }, ...overrides, @@ -337,7 +337,7 @@ describe('OnboardingModal', () => { const globalLink = screen.getByRole('link', { name: 'onboarding.bootstrap.intelKeyLink' }); expect(globalLink).toHaveAttribute('href', 'https://i.threatbook.io/flocks/activate'); - expect(screen.queryByText('onboarding.bootstrap.intelMcpCapability')).not.toBeInTheDocument(); + expect(screen.getByText('onboarding.bootstrap.intelMcpCapability')).toBeInTheDocument(); }); it('shows configured summaries and lets users enter edit mode', async () => { @@ -362,7 +362,7 @@ describe('OnboardingModal', () => { mcp_name: 'threatbook_mcp', service_matrix: { cn: ['api', 'mcp'], - global: ['api'], + global: ['api', 'mcp'], }, }, }), @@ -378,6 +378,7 @@ describe('OnboardingModal', () => { await user.click(screen.getByRole('button', { name: 'onboarding.bootstrap.nextStep' })); expect(screen.getByText('onboarding.bootstrap.intelConfiguredHint')).toBeInTheDocument(); + expect(screen.getAllByText('onboarding.bootstrap.intelConfiguredVerified').length).toBeGreaterThan(0); expect(screen.getByRole('button', { name: 'onboarding.bootstrap.editIntel' })).toBeInTheDocument(); }); diff --git a/webui/src/components/common/OnboardingModal.tsx b/webui/src/components/common/OnboardingModal.tsx index 72a113872..c559f1ce5 100644 --- a/webui/src/components/common/OnboardingModal.tsx +++ b/webui/src/components/common/OnboardingModal.tsx @@ -237,7 +237,7 @@ function RegionChooser({ globalLabel: string; }) { return ( -
+
{([ ['cn', chinaLabel], ['global', globalLabel], @@ -246,10 +246,10 @@ function RegionChooser({ key={candidate} type="button" onClick={() => onChange(candidate)} - className={`rounded-md px-3 py-1.5 text-xs font-medium transition-colors ${ + className={`rounded-lg px-5 py-2 text-sm font-semibold transition-colors ${ value === candidate - ? 'bg-white text-red-700 shadow-sm ring-1 ring-gray-200' - : 'text-gray-500 hover:text-gray-700' + ? 'bg-green-50 text-green-700 ring-1 ring-green-200' + : 'text-gray-500 hover:bg-gray-50 hover:text-gray-700' }`} > {label} @@ -540,9 +540,23 @@ export default function OnboardingModal({ onClose }: OnboardingModalProps) { return t('onboarding.bootstrap.statusConfigured'); }, [intelRuntimeStatus, t]); + const intelRegionLabel = useMemo(() => ( + intelRegion === 'cn' + ? t('onboarding.bootstrap.intelRegionChina') + : t('onboarding.bootstrap.intelRegionGlobal') + ), [intelRegion, t]); + + const intelConfiguredStatusValue = useMemo(() => { + if (!intelConfigured) return t('onboarding.bootstrap.statusNotConfigured'); + return t('onboarding.bootstrap.intelConfiguredVerified', { region: intelRegionLabel }); + }, [intelConfigured, intelRegionLabel, t]); + const currentIntelCapabilities = useMemo(() => { - const matrix = intelRuntimeStatus?.service_matrix || { cn: ['api', 'mcp'], global: ['api'] }; - return matrix[intelRegion] || []; + const matrix = intelRuntimeStatus?.service_matrix || { cn: ['api', 'mcp'], global: ['api', 'mcp'] }; + const configuredCapabilities = new Set(matrix[intelRegion] || []); + configuredCapabilities.add('api'); + configuredCapabilities.add('mcp'); + return ['api', 'mcp'].filter((capability) => configuredCapabilities.has(capability)); }, [intelRegion, intelRuntimeStatus?.service_matrix]); const canSavePrimary = primaryProviderIsThreatBook @@ -695,9 +709,18 @@ export default function OnboardingModal({ onClose }: OnboardingModalProps) { setStartStatus(null); try { - const payload = buildIntelPayload(); - const validateRes = await onboardingAPI.validate(payload); - const validateData = validateRes.data; + let payload = buildIntelPayload(); + let validateRes = await onboardingAPI.validate(payload); + let validateData = validateRes.data; + + if (validateData.error_code === 'region_mismatch' && validateData.suggested_region) { + payload = { + ...payload, + region: validateData.suggested_region, + }; + validateRes = await onboardingAPI.validate(payload); + validateData = validateRes.data; + } if (!validateData.can_apply) { setIntelStatus(buildErrorStatus(validateData, t('onboarding.bootstrap.serviceTestFailed'))); @@ -706,13 +729,14 @@ export default function OnboardingModal({ onClose }: OnboardingModalProps) { const applyRes = await onboardingAPI.apply(payload); const applyData = applyRes.data; + setIntelRegion(payload.region); setIntelConfigured(true); setIntelEditing(false); setIntelSkipped(false); setIntelStatus(buildSuccessStatus( validateData, applyData, - intelRegion === 'cn' + payload.region === 'cn' ? t('onboarding.bootstrap.intelChinaSuccess') : t('onboarding.bootstrap.intelGlobalSuccess'), )); @@ -977,11 +1001,11 @@ export default function OnboardingModal({ onClose }: OnboardingModalProps) {
-

{t('onboarding.bootstrap.intelRegionTitle')}

-

{t('onboarding.bootstrap.intelRegionHint')}

+

{t('onboarding.bootstrap.intelRegionTitle')}

+

{t('onboarding.bootstrap.intelRegionHint')}

{currentIntelCapabilities.includes('api') && ( - + {t('onboarding.bootstrap.intelApiCapability')} )} {currentIntelCapabilities.includes('mcp') && ( - + {t('onboarding.bootstrap.intelMcpCapability')} )} @@ -1117,7 +1141,7 @@ export default function OnboardingModal({ onClose }: OnboardingModalProps) { = {}) { mcp_name: 'threatbook_mcp', service_matrix: { cn: ['api', 'mcp'], - global: ['api'], + global: ['api', 'mcp'], }, }, ...overrides, diff --git a/webui/src/locales/en-US/common.json b/webui/src/locales/en-US/common.json index d33cc1f93..c09cf1d83 100644 --- a/webui/src/locales/en-US/common.json +++ b/webui/src/locales/en-US/common.json @@ -195,12 +195,12 @@ "editIntel": "Edit configuration", "backToConfiguredDetails": "Back to current configuration", "intelPageTitle": "Configure ThreatBook intelligence services", - "intelPageDescription": "Enable free ThreatBook API and MCP capabilities for IOC lookup, threat intelligence analysis, and security operations tools in Rex.", + "intelPageDescription": "Enable free ThreatBook intelligence API and MCP capabilities for IOC lookup, threat intelligence analysis, and security operations tools in Rex.", "intelRegionTitle": "Choose service region", - "intelRegionHint": "Activation links and enabled capabilities vary by region.", - "intelRegionChina": "China", - "intelRegionGlobal": "Global", - "intelApiCapability": "Free API service", + "intelRegionHint": "Choose the region matching your API Key. Flocks will configure the available ThreatBook intelligence services automatically.", + "intelRegionChina": "China region", + "intelRegionGlobal": "International region", + "intelApiCapability": "Free ThreatBook intelligence API service", "intelMcpCapability": "Free MCP service", "intelKeyPlaceholder": "Paste your ThreatBook intelligence API key", "intelKeyLink": "Claim API Key", @@ -208,7 +208,8 @@ "intelMcpLabel": "ThreatBook MCP", "saveIntel": "Save & Verify Intelligence", "intelChinaSuccess": "ThreatBook intelligence verified successfully. API and MCP have been configured.", - "intelGlobalSuccess": "ThreatBook intelligence verified successfully. API has been configured.", + "intelGlobalSuccess": "ThreatBook intelligence verified successfully. International region capabilities have been configured.", + "intelConfiguredVerified": "{{region}} configured and verified", "intelConfiguredHint": "ThreatBook intelligence configuration was detected. Edit it to replace the key or re-verify.", "testing": "Verifying...", "testFailed": "Model call failed — please check your API key", diff --git a/webui/src/locales/zh-CN/common.json b/webui/src/locales/zh-CN/common.json index a38bdd2c5..9369315de 100644 --- a/webui/src/locales/zh-CN/common.json +++ b/webui/src/locales/zh-CN/common.json @@ -195,12 +195,12 @@ "editIntel": "编辑配置", "backToConfiguredDetails": "返回当前配置", "intelPageTitle": "配置微步情报服务", - "intelPageDescription": "免费启用微步 API 与 MCP 能力，为 Rex 提供 IOC 查询、威胁情报分析和安全运营工具能力。", + "intelPageDescription": "免费启用微步情报API与MCP能力，为Rex提供IOC查询、威胁情报分析和安全运营工具能力。", "intelRegionTitle": "选择服务区域", - "intelRegionHint": "不同区域的领取入口和可启用能力不同。", - "intelRegionChina": "国内", - "intelRegionGlobal": "国外", - "intelApiCapability": "免费 API 服务", + "intelRegionHint": "选择 API Key 对应区域，系统会自动配置可用的微步情报服务。", + "intelRegionChina": "中国区", + "intelRegionGlobal": "国际区", + "intelApiCapability": "免费微步情报 API 服务", "intelMcpCapability": "免费 MCP 服务", "intelKeyPlaceholder": "粘贴微步情报 API Key", "intelKeyLink": "领取 API Key", @@ -208,7 +208,8 @@ "intelMcpLabel": "ThreatBook MCP", "saveIntel": "保存并验证微步情报", "intelChinaSuccess": "微步情报服务验证成功，已配置 ThreatBook API 和 MCP。", - "intelGlobalSuccess": "微步情报服务验证成功，已配置 ThreatBook API。", + "intelGlobalSuccess": "微步情报服务验证成功，已配置国际区能力。", + "intelConfiguredVerified": "已配置{{region}}并验证通过", "intelConfiguredHint": "当前已检测到微步情报配置；如需更换 Key 或重新验证，可进入编辑配置。", "testing": "验证中...", "testFailed": "模型调用失败，请检查 Key 是否正确", From 5759ce12f7807abeb7ccdb7fb40ddcc0ab63db64 Mon Sep 17 00:00:00 2001 From: John Yin <10972267+john-yin2333@user.noreply.gitee.com> Date: Sun, 6 Sep 2026 23:16:46 +0800 Subject: [PATCH 35/63] Soften onboarding skip copy --- webui/src/components/common/OnboardingModal.test.tsx | 10 ++++++++++ webui/src/locales/en-US/common.json | 4 ++-- webui/src/locales/zh-CN/common.json | 4 ++-- 3 files changed, 14 insertions(+), 4 deletions(-) diff --git a/webui/src/components/common/OnboardingModal.test.tsx b/webui/src/components/common/OnboardingModal.test.tsx index 630ac17df..fb49933b1 100644 --- a/webui/src/components/common/OnboardingModal.test.tsx +++ b/webui/src/components/common/OnboardingModal.test.tsx @@ -4,6 +4,7 @@ import { render, screen, waitFor } from '@testing-library/react'; import userEvent from '@testing-library/user-event'; import { MemoryRouter } from 'react-router-dom'; import OnboardingModal from './OnboardingModal'; +import zhCNCommon from '@/locales/zh-CN/common.json'; const { catalogAPI, @@ -321,6 +322,15 @@ describe('OnboardingModal', () => { expect(screen.getByText('onboarding.bootstrap.intelPageTitle')).toBeInTheDocument(); }); + it('uses neutral skip descriptions that point users to later setup locations', () => { + const bootstrap = zhCNCommon.onboarding.bootstrap; + + expect(bootstrap.skipModelDescription).toBe('可以先跳过此步骤，稍后在模型清单中继续完成默认模型配置。'); + expect(bootstrap.skipIntelDescription).toBe('可以先跳过此步骤，稍后在情报 MCP 中继续完成微步情报 API 与 MCP 配置。'); + expect(bootstrap.skipModelDescription).not.toMatch(/无法|不可用|不能|失败/); + expect(bootstrap.skipIntelDescription).not.toMatch(/无法|不可用|不能|失败/); + }); + it('shows domestic and global intelligence activation links by selected region', async () => { const user = userEvent.setup(); diff --git a/webui/src/locales/en-US/common.json b/webui/src/locales/en-US/common.json index c09cf1d83..8801273ef 100644 --- a/webui/src/locales/en-US/common.json +++ b/webui/src/locales/en-US/common.json @@ -237,9 +237,9 @@ "previousStep": "Previous", "nextStep": "Next", "skipModelTitle": "Skip model setup?", - "skipModelDescription": "Configure it later in Models. Without a default model, Rex conversations, Agents, and workflows may not be able to call an LLM.", + "skipModelDescription": "You can skip this step for now and continue configuring the default model later in Models.", "skipIntelTitle": "Skip ThreatBook intelligence?", - "skipIntelDescription": "Configure it later in Intelligence MCP. You can still enter Rex chat, but IOC lookup and threat intelligence analysis from ThreatBook API/MCP will be unavailable for now.", + "skipIntelDescription": "You can skip this step for now and continue configuring ThreatBook intelligence API and MCP later in Intelligence MCP.", "returnToConfig": "Return to setup", "confirmSkip": "Skip anyway", "summaryTitle": "Setup summary", diff --git a/webui/src/locales/zh-CN/common.json b/webui/src/locales/zh-CN/common.json index 9369315de..da46bcb7e 100644 --- a/webui/src/locales/zh-CN/common.json +++ b/webui/src/locales/zh-CN/common.json @@ -237,9 +237,9 @@ "previousStep": "上一步", "nextStep": "下一步", "skipModelTitle": "跳过模型配置？", - "skipModelDescription": "稍后去模型清单配置。未配置默认模型时，Rex 对话、Agent 和工作流可能无法正常调用大模型。", + "skipModelDescription": "可以先跳过此步骤，稍后在模型清单中继续完成默认模型配置。", "skipIntelTitle": "跳过微步情报服务？", - "skipIntelDescription": "稍后在情报mcp完成配置。跳过后仍可进入 Rex 对话，但暂时无法使用微步情报 API/MCP 提供的 IOC 查询和情报分析能力。", + "skipIntelDescription": "可以先跳过此步骤，稍后在情报 MCP 中继续完成微步情报 API 与 MCP 配置。", "returnToConfig": "返回配置", "confirmSkip": "继续跳过", "summaryTitle": "配置小结", From aff81669c9e7b1726064be16d89784e69b880a25 Mon Sep 17 00:00:00 2001 From: John Yin <10972267+john-yin2333@user.noreply.gitee.com> Date: Sun, 6 Sep 2026 23:24:34 +0800 Subject: [PATCH 36/63] Clarify onboarding skip destinations --- .../common/OnboardingModal.test.tsx | 25 ++++++++++++++++--- .../src/components/common/OnboardingModal.tsx | 13 +++++++--- webui/src/locales/en-US/common.json | 4 +-- webui/src/locales/zh-CN/common.json | 4 +-- 4 files changed, 35 insertions(+), 11 deletions(-) diff --git a/webui/src/components/common/OnboardingModal.test.tsx b/webui/src/components/common/OnboardingModal.test.tsx index fb49933b1..d99039991 100644 --- a/webui/src/components/common/OnboardingModal.test.tsx +++ b/webui/src/components/common/OnboardingModal.test.tsx @@ -9,6 +9,7 @@ import zhCNCommon from '@/locales/zh-CN/common.json'; const { catalogAPI, clientPost, + currentLanguage, defaultModelAPI, onboardingAPI, sessionApi, @@ -17,6 +18,9 @@ const { list: vi.fn(), }, clientPost: vi.fn(), + currentLanguage: { + value: 'zh-CN', + }, defaultModelAPI: { getResolved: vi.fn(), }, @@ -52,7 +56,7 @@ vi.mock('@/api/client', () => ({ vi.mock('react-i18next', () => ({ useTranslation: () => ({ t: (key: string) => key, - i18n: { language: 'zh-CN' }, + i18n: { language: currentLanguage.value }, }), })); @@ -142,6 +146,7 @@ function renderOnboarding() { describe('OnboardingModal', () => { beforeEach(() => { vi.clearAllMocks(); + currentLanguage.value = 'zh-CN'; defaultModelAPI.getResolved.mockRejectedValue(new Error('no default model')); onboardingAPI.getStatus.mockResolvedValue({ data: makeStatus(), @@ -325,8 +330,8 @@ describe('OnboardingModal', () => { it('uses neutral skip descriptions that point users to later setup locations', () => { const bootstrap = zhCNCommon.onboarding.bootstrap; - expect(bootstrap.skipModelDescription).toBe('可以先跳过此步骤，稍后在模型清单中继续完成默认模型配置。'); - expect(bootstrap.skipIntelDescription).toBe('可以先跳过此步骤，稍后在情报 MCP 中继续完成微步情报 API 与 MCP 配置。'); + expect(bootstrap.skipModelDescription).toBe('可以先跳过此步骤，稍后在左侧导航栏「模型清单」中继续完成默认模型配置。'); + expect(bootstrap.skipIntelDescription).toBe('可以先跳过此步骤，稍后在左侧导航栏「工具清单」页面的「MCP」页签中，找到 ThreatBook MCP 后继续完成微步情报 API 与 MCP 配置。'); expect(bootstrap.skipModelDescription).not.toMatch(/无法|不可用|不能|失败/); expect(bootstrap.skipIntelDescription).not.toMatch(/无法|不可用|不能|失败/); }); @@ -350,6 +355,20 @@ describe('OnboardingModal', () => { expect(screen.getByText('onboarding.bootstrap.intelMcpCapability')).toBeInTheDocument(); }); + it('defaults intelligence setup to the international region in English', async () => { + const user = userEvent.setup(); + currentLanguage.value = 'en-US'; + + renderOnboarding(); + + await screen.findByRole('button', { name: 'onboarding.bootstrap.savePrimary' }); + await user.click(screen.getByRole('button', { name: 'onboarding.bootstrap.nextStep' })); + + const keyLink = screen.getByRole('link', { name: 'onboarding.bootstrap.intelKeyLink' }); + expect(keyLink).toHaveAttribute('href', 'https://i.threatbook.io/flocks/activate'); + expect(screen.getByText('onboarding.bootstrap.intelMcpCapability')).toBeInTheDocument(); + }); + it('shows configured summaries and lets users enter edit mode', async () => { const user = userEvent.setup(); diff --git a/webui/src/components/common/OnboardingModal.tsx b/webui/src/components/common/OnboardingModal.tsx index c559f1ce5..baccc78fd 100644 --- a/webui/src/components/common/OnboardingModal.tsx +++ b/webui/src/components/common/OnboardingModal.tsx @@ -58,6 +58,10 @@ function regionForProvider(providerId: string | null | undefined): OnboardingReg return providerId === 'threatbook-io-llm' ? 'global' : 'cn'; } +function getDefaultIntelRegion(language: string | undefined): OnboardingRegion { + return language?.toLowerCase().startsWith('en') ? 'global' : 'cn'; +} + function statusStyles(tone: SectionTone) { if (tone === 'success') { return { @@ -310,8 +314,9 @@ function SkipConfirmDialog({ } export default function OnboardingModal({ onClose }: OnboardingModalProps) { - const { t } = useTranslation('common'); + const { t, i18n } = useTranslation('common'); const navigate = useNavigate(); + const defaultIntelRegion = useMemo(() => getDefaultIntelRegion(i18n.language), [i18n.language]); const [step, setStep] = useState('model'); const [catalog, setCatalog] = useState([]); @@ -333,7 +338,7 @@ export default function OnboardingModal({ onClose }: OnboardingModalProps) { const [modelRegion, setModelRegion] = useState('cn'); const [modelSkipped, setModelSkipped] = useState(false); - const [intelRegion, setIntelRegion] = useState('cn'); + const [intelRegion, setIntelRegion] = useState(() => getDefaultIntelRegion(i18n.language)); const [intelApiKey, setIntelApiKey] = useState(''); const [intelSaving, setIntelSaving] = useState(false); const [intelConfigured, setIntelConfigured] = useState(false); @@ -365,7 +370,7 @@ export default function OnboardingModal({ onClose }: OnboardingModalProps) { const intel = res.data.threatbook_intel; setIntelRuntimeStatus(intel); setIntelConfigured(intel.configured); - setIntelRegion(intel.region || 'cn'); + setIntelRegion(intel.region || defaultIntelRegion); if (intel.configured) setIntelEditing(false); } } catch { @@ -389,7 +394,7 @@ export default function OnboardingModal({ onClose }: OnboardingModalProps) { } finally { if (!silent) setStatusLoading(false); } - }, []); + }, [defaultIntelRegion]); useEffect(() => { refreshOnboardingStatus(); diff --git a/webui/src/locales/en-US/common.json b/webui/src/locales/en-US/common.json index 8801273ef..5465761ed 100644 --- a/webui/src/locales/en-US/common.json +++ b/webui/src/locales/en-US/common.json @@ -237,9 +237,9 @@ "previousStep": "Previous", "nextStep": "Next", "skipModelTitle": "Skip model setup?", - "skipModelDescription": "You can skip this step for now and continue configuring the default model later in Models.", + "skipModelDescription": "You can skip this step for now and continue configuring the default model later from the Models item in the left navigation.", "skipIntelTitle": "Skip ThreatBook intelligence?", - "skipIntelDescription": "You can skip this step for now and continue configuring ThreatBook intelligence API and MCP later in Intelligence MCP.", + "skipIntelDescription": "You can skip this step for now. Later, open Tools from the left navigation, switch to the MCP tab, find ThreatBook MCP, and continue configuring ThreatBook intelligence API and MCP.", "returnToConfig": "Return to setup", "confirmSkip": "Skip anyway", "summaryTitle": "Setup summary", diff --git a/webui/src/locales/zh-CN/common.json b/webui/src/locales/zh-CN/common.json index da46bcb7e..da2a88bc2 100644 --- a/webui/src/locales/zh-CN/common.json +++ b/webui/src/locales/zh-CN/common.json @@ -237,9 +237,9 @@ "previousStep": "上一步", "nextStep": "下一步", "skipModelTitle": "跳过模型配置？", - "skipModelDescription": "可以先跳过此步骤，稍后在模型清单中继续完成默认模型配置。", + "skipModelDescription": "可以先跳过此步骤，稍后在左侧导航栏「模型清单」中继续完成默认模型配置。", "skipIntelTitle": "跳过微步情报服务？", - "skipIntelDescription": "可以先跳过此步骤，稍后在情报 MCP 中继续完成微步情报 API 与 MCP 配置。", + "skipIntelDescription": "可以先跳过此步骤，稍后在左侧导航栏「工具清单」页面的「MCP」页签中，找到 ThreatBook MCP 后继续完成微步情报 API 与 MCP 配置。", "returnToConfig": "返回配置", "confirmSkip": "继续跳过", "summaryTitle": "配置小结", From 308b512ecfe7a92a94ee1c2db60851ae0bc265ea Mon Sep 17 00:00:00 2001 From: John Yin <10972267+john-yin2333@user.noreply.gitee.com> Date: Sun, 6 Sep 2026 23:35:42 +0800 Subject: [PATCH 37/63] Add ThreatBook free key link to model provider setup --- webui/src/locales/en-US/model.json | 1 + webui/src/locales/zh-CN/model.json | 1 + webui/src/pages/Model/index.test.tsx | 83 ++++++++++++++++++++++++++++ webui/src/pages/Model/index.tsx | 19 ++++++- 4 files changed, 103 insertions(+), 1 deletion(-) diff --git a/webui/src/locales/en-US/model.json b/webui/src/locales/en-US/model.json index 937fdc6ef..2aa65c24a 100644 --- a/webui/src/locales/en-US/model.json +++ b/webui/src/locales/en-US/model.json @@ -111,6 +111,7 @@ "apiKeyOptionalHint": "If Base URL points to an internal/self-hosted endpoint with no auth, leave this empty and the system will use a placeholder.", "apiKeyOptionalPlaceholder": "Leave empty for no-auth gateway", "apiKeyKeepExisting": "Leave blank to keep the existing API key", + "claimFreeKey": "Claim free key", "hide": "Hide", "show": "Show", "availableModels": "Available Models", diff --git a/webui/src/locales/zh-CN/model.json b/webui/src/locales/zh-CN/model.json index bf9dd91d3..dec712cd5 100644 --- a/webui/src/locales/zh-CN/model.json +++ b/webui/src/locales/zh-CN/model.json @@ -111,6 +111,7 @@ "apiKeyOptionalHint": "若 Base URL 指向无需鉴权的内网网关或自建服务，可留空，系统会自动使用占位符。", "apiKeyOptionalPlaceholder": "可留空（无鉴权网关）", "apiKeyKeepExisting": "留空以保留现有 API Key", + "claimFreeKey": "领取免费key", "hide": "隐藏", "show": "查看", "availableModels": "可用模型", diff --git a/webui/src/pages/Model/index.test.tsx b/webui/src/pages/Model/index.test.tsx index cabbe5667..f9b96b54d 100644 --- a/webui/src/pages/Model/index.test.tsx +++ b/webui/src/pages/Model/index.test.tsx @@ -172,6 +172,68 @@ describe('ModelPage add provider dialog', () => { mocks.catalogList.mockResolvedValue({ data: { providers: [ + { + id: 'threatbook-cn-llm', + name: 'ThreatBook-cn-llm', + description: 'ThreatBook China LLM Service', + credential_schemas: [{ + auth_method: 'api_key', + fields: [{ + name: 'api_key', + label: 'API Key', + type: 'secret', + required: true, + placeholder: 'Paste ThreatBook CN LLM API Key', + }], + }], + env_vars: [], + default_base_url: 'https://llm.threatbook.cn/v1', + model_count: 1, + models: [{ + id: 'deepseek-v4-flash-0731', + name: 'deepseek-v4-flash-0731', + model_type: 'llm', + status: 'active', + capabilities: { + supports_tools: true, + supports_vision: false, + supports_reasoning: true, + supports_streaming: true, + }, + }], + allow_multiple: false, + }, + { + id: 'threatbook-io-llm', + name: 'ThreatBook-io-llm', + description: 'ThreatBook International LLM Service', + credential_schemas: [{ + auth_method: 'api_key', + fields: [{ + name: 'api_key', + label: 'API Key', + type: 'secret', + required: true, + placeholder: 'Paste ThreatBook IO LLM API Key', + }], + }], + env_vars: [], + default_base_url: 'https://llm.threatbook.io/v1', + model_count: 1, + models: [{ + id: 'deepseek-v4-flash-0731', + name: 'deepseek-v4-flash-0731', + model_type: 'llm', + status: 'active', + capabilities: { + supports_tools: true, + supports_vision: false, + supports_reasoning: true, + supports_streaming: true, + }, + }], + allow_multiple: false, + }, { id: 'openai-compatible', name: 'OpenAI Compatible', @@ -193,6 +255,27 @@ describe('ModelPage add provider dialog', () => { }); }); + it('keeps ThreatBook CN and IO providers separate and shows a shared free key link', async () => { + const user = userEvent.setup(); + + renderWithRouter(); + + await user.click(screen.getByRole('button', { name: 'Add Provider' })); + await user.click(await screen.findByRole('button', { name: 'Select Provider...' })); + + expect(await screen.findByRole('button', { name: /ThreatBook-cn-llm/i })).toBeInTheDocument(); + expect(screen.getByRole('button', { name: /ThreatBook-io-llm/i })).toBeInTheDocument(); + + await user.click(screen.getByRole('button', { name: /ThreatBook-cn-llm/i })); + let keyLink = screen.getByRole('link', { name: 'form.claimFreeKey' }); + expect(keyLink).toHaveAttribute('href', 'https://portal.agentflocks.com/'); + + await user.click(screen.getByRole('button', { name: /ThreatBook-cn-llm/i })); + await user.click(screen.getByRole('button', { name: /ThreatBook-io-llm/i })); + keyLink = screen.getByRole('link', { name: 'form.claimFreeKey' }); + expect(keyLink).toHaveAttribute('href', 'https://portal.agentflocks.com/'); + }); + it('blocks openai-compatible creation until Base URL is filled and submits once provided', async () => { const user = userEvent.setup(); diff --git a/webui/src/pages/Model/index.tsx b/webui/src/pages/Model/index.tsx index 5fcac4d96..e206ecb09 100644 --- a/webui/src/pages/Model/index.tsx +++ b/webui/src/pages/Model/index.tsx @@ -9,7 +9,7 @@ import { ChevronDown, Check, AlertCircle, Loader2, X, Shield, Pencil, Star, AlertTriangle, CheckCircle2, ArrowUp, ArrowDown, ListOrdered, - Info, + Info, ExternalLink, } from 'lucide-react'; import PageHeader from '@/components/common/PageHeader'; import LoadingSpinner from '@/components/common/LoadingSpinner'; @@ -66,11 +66,17 @@ function isCatalogBaseUrlRequired(providerId: string): boolean { } const AZURE_PROVIDER_IDS = new Set(['azure-openai', 'azure']); +const THREATBOOK_LLM_PROVIDER_IDS = new Set(['threatbook-cn-llm', 'threatbook-io-llm']); +const THREATBOOK_FREE_KEY_URL = 'https://portal.agentflocks.com/'; function isAzureProviderId(providerId: string): boolean { return AZURE_PROVIDER_IDS.has(providerId); } +function isThreatBookLLMProviderId(providerId: string): boolean { + return THREATBOOK_LLM_PROVIDER_IDS.has(providerId); +} + function convertEditablePrice( value: string, sourceCurrency: string, @@ -1724,6 +1730,17 @@ function AddProviderDialog({ connectedIds, onClose, onAdded }: { {showApiKey ? : }
+ {isThreatBookLLMProviderId(selectedCatalogId) && ( + + + {t('form.claimFreeKey')} + + )} {selectedCatalogId !== 'ollama' && providerAllowsEmptyApiKey(selectedCatalogId) && (

{t('form.apiKeyOptionalHint')}

)} From 6ba55739bb9f9b43b071566dbc1c4bcb19c17f0e Mon Sep 17 00:00:00 2001 From: John Yin <10972267+john-yin2333@user.noreply.gitee.com> Date: Sun, 6 Sep 2026 23:47:31 +0800 Subject: [PATCH 38/63] Add ThreatBook MCP free key link --- webui/src/locales/en-US/tool.json | 1 + webui/src/locales/zh-CN/tool.json | 1 + webui/src/pages/Tool/ToolSheets.tsx | 13 ++-- .../components/ServiceDetailPanel.test.tsx | 68 ++++++++++++++++++- .../Tool/components/ServiceDetailPanel.tsx | 34 +++++++++- 5 files changed, 110 insertions(+), 7 deletions(-) diff --git a/webui/src/locales/en-US/tool.json b/webui/src/locales/en-US/tool.json index 90c78762c..40d7a83f7 100644 --- a/webui/src/locales/en-US/tool.json +++ b/webui/src/locales/en-US/tool.json @@ -357,6 +357,7 @@ "quickActions": "Quick Actions", "testingConn": "Testing...", "testConnection": "Test Connection", + "claimFreeApiKey": "Claim free API key", "refreshingTools": "Refreshing...", "refreshTools": "Refresh Tools", "disconnectConn": "Disconnect", diff --git a/webui/src/locales/zh-CN/tool.json b/webui/src/locales/zh-CN/tool.json index 59e98ef22..1a3078d8f 100644 --- a/webui/src/locales/zh-CN/tool.json +++ b/webui/src/locales/zh-CN/tool.json @@ -357,6 +357,7 @@ "quickActions": "快捷操作", "testingConn": "测试中...", "testConnection": "测试连接", + "claimFreeApiKey": "免费领取 API Key", "refreshingTools": "刷新中...", "refreshTools": "刷新工具列表", "disconnectConn": "断开连接", diff --git a/webui/src/pages/Tool/ToolSheets.tsx b/webui/src/pages/Tool/ToolSheets.tsx index 5ba1bfff0..442f7cadf 100644 --- a/webui/src/pages/Tool/ToolSheets.tsx +++ b/webui/src/pages/Tool/ToolSheets.tsx @@ -6,7 +6,7 @@ * - GenerateToolSheet: AI 生成自定义工具 */ -import { useMemo, useState } from 'react'; +import { useMemo, useState, type ReactNode } from 'react'; import { useTranslation } from 'react-i18next'; import { Database, Cloud, Code, Info, CheckCircle, XCircle, Activity, Wifi, WifiOff, Loader2 } from 'lucide-react'; import EntitySheet from '@/components/common/EntitySheet'; @@ -252,6 +252,7 @@ interface MCPFormFieldsProps { testResult: { success: boolean; message: string; tools_count?: number } | null; onTestConnection: () => void; isTesting: boolean; + serviceUrlAction?: ReactNode; } export function MCPFormFields({ @@ -262,6 +263,7 @@ export function MCPFormFields({ testResult, onTestConnection, isTesting, + serviceUrlAction, }: MCPFormFieldsProps) { const { t } = useTranslation('tool'); const readOnly = !onChange; @@ -357,9 +359,12 @@ export function MCPFormFields({ {formData.connType === 'sse' && (
- +
+ + {serviceUrlAction} +
({ +const { currentLanguage, mcpAPI } = vi.hoisted(() => ({ + currentLanguage: { value: 'zh-CN' }, mcpAPI: { get: vi.fn(), update: vi.fn(), @@ -67,11 +68,13 @@ vi.mock('../ToolSheets', () => ({ onChange, onTestConnection, testResult, + serviceUrlAction, }: { formData: { url: string }; onChange?: (fields: { url: string }) => void; onTestConnection: () => void; testResult: { message: string } | null; + serviceUrlAction?: React.ReactNode; }) => (
({ + {serviceUrlAction} {testResult &&
{testResult.message}
}
), @@ -106,11 +110,16 @@ vi.mock('react-i18next', () => ({ 'button.save': '保存', 'button.saving': '保存中...', 'detail.testFailed': '连接测试失败', + 'detail.claimFreeApiKey': currentLanguage.value.startsWith('zh') ? '免费领取 API Key' : 'Claim free API key', 'alert.connectionOk': '连接成功', }; return translations[key] ?? key; }, - i18n: { changeLanguage: vi.fn() }, + i18n: { + language: currentLanguage.value, + resolvedLanguage: currentLanguage.value, + changeLanguage: vi.fn(), + }, }), Trans: ({ children }: { children: React.ReactNode }) => children, initReactI18next: { type: '3rdParty', init: vi.fn() }, @@ -144,6 +153,7 @@ describe('MCPServerDetailPanel', () => { beforeEach(() => { vi.clearAllMocks(); + currentLanguage.value = 'zh-CN'; mcpAPI.get.mockResolvedValue(detailResponse); mcpAPI.update.mockResolvedValue({ data: { success: true }, @@ -224,4 +234,58 @@ describe('MCPServerDetailPanel', () => { expect(mcpAPI.get).toHaveBeenCalledTimes(2); expect(mcpAPI.testExisting).not.toHaveBeenCalled(); }); + + it('shows the China free key link for ThreatBook MCP in Chinese', async () => { + render( + , + ); + + const link = await screen.findByRole('link', { name: '免费领取 API Key' }); + expect(link).toHaveAttribute('href', 'https://x.threatbook.com/flocks/activate'); + expect(link).toHaveAttribute('target', '_blank'); + }); + + it('shows the international free key link for ThreatBook MCP in English', async () => { + currentLanguage.value = 'en-US'; + + render( + , + ); + + const link = await screen.findByRole('link', { name: 'Claim free API key' }); + expect(link).toHaveAttribute('href', 'https://i.threatbook.io/flocks/activate'); + }); + + it('does not show the free key link for other MCP servers', async () => { + render( + , + ); + + await screen.findByLabelText('service-url'); + expect(screen.queryByRole('link', { name: '免费领取 API Key' })).not.toBeInTheDocument(); + }); }); diff --git a/webui/src/pages/Tool/components/ServiceDetailPanel.tsx b/webui/src/pages/Tool/components/ServiceDetailPanel.tsx index 4480e0fb4..b0b172ed4 100644 --- a/webui/src/pages/Tool/components/ServiceDetailPanel.tsx +++ b/webui/src/pages/Tool/components/ServiceDetailPanel.tsx @@ -3,6 +3,7 @@ import { useTranslation } from 'react-i18next'; import { Info, Wrench, FileText, Activity, Zap, RefreshCw, Power, PowerOff, CheckCircle, XCircle, Cloud, Database, AlertTriangle, Eye, EyeOff, Save, Trash2, + ExternalLink, } from 'lucide-react'; import type { Tool } from '@/api/tool'; import type { MCPCatalogCategory, MCPCatalogEntry, MCPCredentials, MCPServer, MCPServerDetail } from '@/types'; @@ -16,6 +17,23 @@ import { buildMCPConfigFromForm, buildMCPFormDataFromConfig, getMCPFormError, MC import type { MCPFormData, ConnStatus as MCPConnStatus } from '../ToolSheets'; import type { APIServiceCredentialField, APIServiceMetadata, ProviderCredentials } from '@/types'; +const THREATBOOK_MCP_SERVER_NAMES = new Set(['threatbook_mcp', 'threatbook-mcp']); +const THREATBOOK_MCP_ACTIVATION_URLS = { + cn: 'https://x.threatbook.com/flocks/activate', + global: 'https://i.threatbook.io/flocks/activate', +} as const; + +function isThreatBookMCPServer(serverName: string): boolean { + return THREATBOOK_MCP_SERVER_NAMES.has(serverName.trim().toLowerCase()); +} + +export function getThreatBookMCPActivationUrl(language: string): string { + const normalizedLanguage = language.toLowerCase().replace('_', '-'); + return normalizedLanguage.startsWith('zh') + ? THREATBOOK_MCP_ACTIVATION_URLS.cn + : THREATBOOK_MCP_ACTIVATION_URLS.global; +} + function KvRowValue({ value }: { value: string }) { const [showTooltip, setShowTooltip] = useState(false); @@ -72,7 +90,7 @@ export function MCPServerDetailPanel({ onRemove?: () => void; onSelectTool: (tool: Tool) => void; }) { - const { t } = useTranslation('tool'); + const { t, i18n } = useTranslation('tool'); const [detailTab, setDetailTab] = useState<'overview' | 'tools' | 'resources'>('overview'); const [serverDetail, setServerDetail] = useState(null); const [formData, setFormData] = useState(null); @@ -82,6 +100,9 @@ export function MCPServerDetailPanel({ const [testResult, setTestResult] = useState<{ success: boolean; message: string; latency?: number; tools_count?: number } | null>(null); const [refreshing, setRefreshing] = useState(false); const [savingConfig, setSavingConfig] = useState(false); + const threatBookActivationUrl = isThreatBookMCPServer(server.name) + ? getThreatBookMCPActivationUrl(i18n.resolvedLanguage || i18n.language) + : null; const createFormData = useCallback((detail: MCPServerDetail | null): MCPFormData => ( buildMCPFormDataFromConfig(server.name, detail?.config, server.url) @@ -244,6 +265,17 @@ export function MCPServerDetailPanel({ testResult={testResult ? { success: testResult.success, message: testResult.message, tools_count: testResult.tools_count } : null} onTestConnection={handleTestConnection} isTesting={testingConnection} + serviceUrlAction={threatBookActivationUrl ? ( + + + {t('detail.claimFreeApiKey')} + + ) : undefined} /> )} From 1c6c896af5dc06dd66e702ab1c4bdbb37582ddce Mon Sep 17 00:00:00 2001 From: John Yin <10972267+john-yin2333@user.noreply.gitee.com> Date: Mon, 7 Sep 2026 00:18:27 +0800 Subject: [PATCH 39/63] Complete regional ThreatBook MCP setup --- flocks/mcp/utils.py | 6 +- flocks/security/secrets.py | 5 + flocks/server/routes/mcp.py | 185 ++++++++- flocks/server/routes/onboarding.py | 69 ++-- flocks/server/threatbook_regions.py | 54 +++ tests/server/routes/test_mcp_routes.py | 184 +++++++++ tests/server/routes/test_onboarding_routes.py | 49 ++- webui/src/api/mcp.ts | 11 + .../src/components/common/OnboardingModal.tsx | 16 +- webui/src/constants/threatbook.ts | 26 ++ webui/src/locales/en-US/tool.json | 34 +- webui/src/locales/zh-CN/tool.json | 34 +- webui/src/pages/Tool/ToolSheets.tsx | 13 +- .../components/ServiceDetailPanel.test.tsx | 130 +++++- .../Tool/components/ServiceDetailPanel.tsx | 46 +-- .../components/ThreatBookMCPConfigPanel.tsx | 371 ++++++++++++++++++ 16 files changed, 1114 insertions(+), 119 deletions(-) create mode 100644 flocks/server/threatbook_regions.py create mode 100644 webui/src/constants/threatbook.ts create mode 100644 webui/src/pages/Tool/components/ThreatBookMCPConfigPanel.tsx diff --git a/flocks/mcp/utils.py b/flocks/mcp/utils.py index b5fa4f416..ef0df7c74 100644 --- a/flocks/mcp/utils.py +++ b/flocks/mcp/utils.py @@ -335,8 +335,9 @@ def extract_api_key_from_mcp_url(server_name: str, config: Dict[str, Any]) -> Di key = unquote(key_encoded) value = unquote(value_encoded) if key.lower() in _SENSITIVE_QUERY_PARAMS and not value.startswith("{secret:"): - secret_key = f"{server_name}_mcp_key" from flocks.security import get_secret_manager + from flocks.security.secrets import get_mcp_secret_id + secret_key = get_mcp_secret_id(server_name) get_secret_manager().set(secret_key, value) new_parts.append(f"{key_encoded}={{secret:{secret_key}}}") extracted = True @@ -383,7 +384,8 @@ def extract_auth_value_from_mcp_config(server_name: str, config: Dict[str, Any]) elif "scheme" in updated_auth and not scheme: updated_auth.pop("scheme", None) - secret_key = str(auth_config.get("secret_id") or f"{server_name}_mcp_key") + from flocks.security.secrets import get_mcp_secret_id + secret_key = str(auth_config.get("secret_id") or get_mcp_secret_id(server_name)) from flocks.security import get_secret_manager get_secret_manager().set(secret_key, auth_value) diff --git a/flocks/security/secrets.py b/flocks/security/secrets.py index faf648499..8f487fc46 100644 --- a/flocks/security/secrets.py +++ b/flocks/security/secrets.py @@ -26,6 +26,11 @@ log = Log.create(service="security.secrets") +def get_mcp_secret_id(server_name: str) -> str: + """Return the canonical MCP secret ID without duplicating an `_mcp` suffix.""" + return f"{server_name}_key" if server_name.endswith("_mcp") else f"{server_name}_mcp_key" + + class SecretManager: """ Flat KV secret manager using plain JSON file diff --git a/flocks/server/routes/mcp.py b/flocks/server/routes/mcp.py index 1990e13b1..ee133822b 100644 --- a/flocks/server/routes/mcp.py +++ b/flocks/server/routes/mcp.py @@ -8,6 +8,7 @@ """ import asyncio +import copy from typing import Dict, Optional, List, Any from fastapi import APIRouter, HTTPException from fastapi.responses import JSONResponse @@ -41,6 +42,13 @@ ) from flocks.config.config import Config from flocks.config.config_writer import ConfigWriter +from flocks.security import get_secret_manager +from flocks.security.secrets import get_mcp_secret_id +from flocks.server.threatbook_regions import ( + THREATBOOK_REGION_PRESETS, + ThreatBookRegion, + build_threatbook_mcp_url, +) from flocks.utils.log import Log @@ -293,6 +301,22 @@ class McpUpdateRequest(BaseModel): config: Dict[str, Any] = Field(..., description="Partial or full MCP server configuration") +class ThreatBookMcpConfigureRequest(BaseModel): + """Region-aware ThreatBook MCP setup request.""" + + region: ThreatBookRegion + api_key: str = Field(..., min_length=1, description="ThreatBook regional API key") + + +class ThreatBookMcpConfigureResponse(BaseModel): + success: bool + message: str + region: ThreatBookRegion + endpoint: str + connected: bool + tools_count: int = 0 + + @router.post( "/test", response_model=Dict[str, Any], @@ -344,6 +368,138 @@ async def test_mcp_connection(request: McpTestRequest): log.warn("mcp.test.cleanup_failed", {"server": temp_name, "error": str(e)}) +async def _restore_threatbook_mcp_setup( + name: str, + config_snapshot: Dict[str, Any], + secret_snapshot: Dict[str, Any], + previous_config: Optional[Dict[str, Any]], + was_connected: bool, +) -> None: + """Restore persisted and runtime state after a failed ThreatBook apply.""" + secrets = get_secret_manager() + ConfigWriter._write_raw(config_snapshot) + secrets._save(secret_snapshot) + + from flocks.tool.tool_loader import delete_mcp_config, save_mcp_config + + if previous_config: + save_mcp_config(name, previous_config) + else: + delete_mcp_config(name) + + try: + runtime_status = await MCP.status() + if name in runtime_status: + await MCP.remove(name) + if was_connected and previous_config: + restored_config = await _load_mcp_server_config(name) + if restored_config: + await MCP.connect(name, restored_config) + except Exception as exc: + log.warning("mcp.threatbook.restore_runtime_failed", {"name": name, "error": str(exc)}) + + +@router.post( + "/{name}/threatbook-configure", + response_model=ThreatBookMcpConfigureResponse, + summary="Configure ThreatBook MCP by region", + description="Validate a regional ThreatBook key before atomically saving and connecting MCP.", + operation_id="mcp.threatbook_configure", +) +async def configure_threatbook_mcp( + name: str, + request: ThreatBookMcpConfigureRequest, +) -> ThreatBookMcpConfigureResponse: + """Validate first, then persist the regional endpoint and secret reference.""" + preset = THREATBOOK_REGION_PRESETS[request.region] + if name != preset["threatbook_mcp_name"]: + raise HTTPException(status_code=400, detail="This setup flow is only available for ThreatBook MCP") + + api_key = request.api_key.strip() + if not api_key: + raise HTTPException(status_code=400, detail="API key required") + + endpoint = preset["threatbook_mcp_endpoint"] + validation = await test_mcp_connection( + McpTestRequest( + name=name, + config={ + "type": "remote", + "url": build_threatbook_mcp_url(request.region, api_key), + "transport": "auto", + }, + ) + ) + if not validation.get("success"): + return ThreatBookMcpConfigureResponse( + success=False, + message=validation.get("message") or "ThreatBook MCP validation failed", + region=request.region, + endpoint=endpoint, + connected=False, + tools_count=0, + ) + + secrets = get_secret_manager() + config_snapshot = copy.deepcopy(ConfigWriter._read_raw()) + secret_snapshot = copy.deepcopy(secrets._load()) + previous_config = copy.deepcopy(_load_raw_mcp_server_config(name)) + runtime_status = await MCP.status() + previous_status = runtime_status.get(name) + was_connected = bool( + previous_status is not None + and previous_status.status == McpStatus.CONNECTED + ) + secret_id = preset["threatbook_mcp_secret_id"] + next_config = { + "type": "remote", + "url": build_threatbook_mcp_url( + request.region, + f"{{secret:{secret_id}}}", + encode_key=False, + ), + "transport": "auto", + "enabled": True, + } + + try: + secrets.set(secret_id, api_key) + _persist_mcp_server_config(name, next_config) + + if previous_status is not None: + await MCP.remove(name) + resolved_config = await _load_mcp_server_config(name) + if not resolved_config or not await MCP.connect(name, resolved_config): + raise ValueError("ThreatBook MCP could not reconnect after saving") + + tools_count = await MCP.refresh_tools(name) + return ThreatBookMcpConfigureResponse( + success=True, + message="ThreatBook MCP configured and connected successfully", + region=request.region, + endpoint=endpoint, + connected=True, + tools_count=tools_count, + ) + except Exception as exc: + await _restore_threatbook_mcp_setup( + name, + config_snapshot, + secret_snapshot, + previous_config, + was_connected, + ) + log.error("mcp.threatbook.configure_failed", { + "name": name, + "region": request.region, + "error": str(exc), + }) + raise HTTPException( + status_code=500, + detail=f"ThreatBook MCP configuration failed: {exc}", + ) from exc + + @router.post( "/{name}/test", response_model=Dict[str, Any], @@ -819,18 +975,22 @@ class McpCredentialResponse(BaseModel): async def get_mcp_credentials(name: str): """Get masked credential info for a server. - Looks up the convention-based secret_id '{name}_mcp_key' in .secret.json. - Falls back to legacy '{name}_api_key' for backward compatibility. + Avoids duplicating an existing ``_mcp`` suffix and falls back to historical + ``{name}_mcp_key`` / ``{name}_api_key`` variants for compatibility. """ - from flocks.security import get_secret_manager from flocks.security.secrets import SecretManager try: secrets = get_secret_manager() - # Convention-based secret_id: _mcp_key first, fall back to legacy _api_key - secret_id = f"{name}_mcp_key" + # Convention-based secret_id first, fall back to historical variants. + secret_id = get_mcp_secret_id(name) api_key = secrets.get(secret_id) + duplicated_suffix_id = f"{name}_mcp_key" + if not api_key and duplicated_suffix_id != secret_id: + api_key = secrets.get(duplicated_suffix_id) + if api_key: + secret_id = duplicated_suffix_id if not api_key: legacy_id = f"{name}_api_key" api_key = secrets.get(legacy_id) @@ -858,16 +1018,14 @@ async def set_mcp_credentials(name: str, request: McpCredentialRequest): Stores in .secret.json with flat KV format. """ - from flocks.security import get_secret_manager - try: if not request.api_key: raise HTTPException(status_code=400, detail="API key required") secrets = get_secret_manager() - # Use provided secret_id or convention-based default (_mcp_key for MCP servers) - secret_id = request.secret_id or f"{name}_mcp_key" + # Use the provided secret ID or the canonical MCP naming convention. + secret_id = request.secret_id or get_mcp_secret_id(name) secrets.set(secret_id, request.api_key) log.info("mcp.credentials.set", {"name": name, "secret_id": secret_id}) @@ -893,13 +1051,14 @@ async def set_mcp_credentials(name: str, request: McpCredentialRequest): ) async def delete_mcp_credentials(name: str): """Delete credentials for a server.""" - from flocks.security import get_secret_manager - try: secrets = get_secret_manager() - # Delete both current (_mcp_key) and legacy (_api_key) entries - secret_id = f"{name}_mcp_key" + # Delete current and historical credential IDs. + secret_id = get_mcp_secret_id(name) deleted = secrets.delete(secret_id) + duplicated_suffix_id = f"{name}_mcp_key" + if duplicated_suffix_id != secret_id: + deleted = secrets.delete(duplicated_suffix_id) or deleted deleted = secrets.delete(f"{name}_api_key") or deleted if deleted: diff --git a/flocks/server/routes/onboarding.py b/flocks/server/routes/onboarding.py index 53c75c72b..bbf6a6d2b 100644 --- a/flocks/server/routes/onboarding.py +++ b/flocks/server/routes/onboarding.py @@ -9,8 +9,7 @@ import copy from contextlib import asynccontextmanager -from typing import Any, Dict, List, Literal, Optional -from urllib.parse import quote +from typing import Any, Dict, List, Optional from fastapi import APIRouter, HTTPException from pydantic import BaseModel, Field @@ -45,6 +44,12 @@ test_provider_credentials, update_api_service, ) +from flocks.server.threatbook_regions import ( + THREATBOOK_REGION_PRESETS, + ThreatBookRegion, + build_threatbook_mcp_url, + infer_threatbook_mcp_region, +) from flocks.tool.tool_loader import save_mcp_config from flocks.utils.log import Log @@ -52,7 +57,7 @@ router = APIRouter() log = Log.create(service="routes.onboarding") -Region = Literal["cn", "global"] +Region = ThreatBookRegion class ThirdPartyLLMConfig(BaseModel): @@ -160,28 +165,7 @@ def _llm_provider_has_usable_credentials(provider_id: str) -> bool: return bool(_get_inline_provider_api_key(provider_id)) -ONBOARDING_REGION_PRESETS: Dict[Region, Dict[str, Any]] = { - "cn": { - "activation_url": "https://x.threatbook.com/flocks/activate", - "threatbook_llm_provider_id": "threatbook-cn-llm", - "threatbook_default_model_id": "deepseek-v4-flash-0731", - "threatbook_api_service_id": "threatbook-cn", - "threatbook_mcp_name": "threatbook_mcp", - "threatbook_mcp_url": "https://mcp.threatbook.cn/mcp?apikey={api_key}", - "threatbook_mcp_secret_id": "threatbook_mcp_key", - "requires_mcp": True, - }, - "global": { - "activation_url": "https://i.threatbook.io/flocks/activate", - "threatbook_llm_provider_id": "threatbook-io-llm", - "threatbook_default_model_id": "deepseek-v4-flash-0731", - "threatbook_api_service_id": "threatbook-io", - "threatbook_mcp_name": None, - "threatbook_mcp_url": None, - "threatbook_mcp_secret_id": None, - "requires_mcp": False, - }, -} +ONBOARDING_REGION_PRESETS = THREATBOOK_REGION_PRESETS async def _api_service_has_credentials(service_id: str) -> bool: @@ -263,10 +247,17 @@ async def _build_threatbook_intel_status() -> ThreatBookIntelStatus: global_preset = ONBOARDING_REGION_PRESETS["global"] cn_api_configured = await _api_service_has_credentials(cn_preset["threatbook_api_service_id"]) global_api_configured = await _api_service_has_credentials(global_preset["threatbook_api_service_id"]) - cn_mcp = await _detect_mcp_status(cn_preset["threatbook_mcp_name"]) + mcp_name = cn_preset["threatbook_mcp_name"] + mcp_status = await _detect_mcp_status(mcp_name) + raw_mcp_config = ConfigWriter.get_mcp_server(mcp_name) + mcp_region = infer_threatbook_mcp_region( + raw_mcp_config.get("url") if isinstance(raw_mcp_config, dict) else None + ) region: Optional[Region] = None - if cn_api_configured or cn_mcp["configured"]: + if mcp_status["configured"] and mcp_region: + region = mcp_region + elif cn_api_configured: region = "cn" elif global_api_configured: region = "global" @@ -274,7 +265,7 @@ async def _build_threatbook_intel_status() -> ThreatBookIntelStatus: api_configured = cn_api_configured if region != "global" else global_api_configured return ThreatBookIntelStatus( - configured=bool(cn_api_configured or global_api_configured or cn_mcp["configured"]), + configured=bool(cn_api_configured or global_api_configured or mcp_status["configured"]), region=region, api_configured=api_configured, api_service_id=( @@ -282,13 +273,13 @@ async def _build_threatbook_intel_status() -> ThreatBookIntelStatus: if region == "global" else cn_preset["threatbook_api_service_id"] ), - mcp_configured=cn_mcp["configured"], - mcp_connected=cn_mcp["connected"], - mcp_status=cn_mcp["status"], - mcp_name=cn_preset["threatbook_mcp_name"], + mcp_configured=mcp_status["configured"], + mcp_connected=mcp_status["connected"], + mcp_status=mcp_status["status"], + mcp_name=mcp_name, service_matrix={ "cn": ["api", "mcp"], - "global": ["api"], + "global": ["api", "mcp"], }, ) @@ -425,19 +416,17 @@ async def _test_provider_or_service_with_temp_credentials( async def _test_mcp_with_temp_key(region: Region, api_key: str) -> Dict[str, Any]: preset = ONBOARDING_REGION_PRESETS[region] mcp_name = preset["threatbook_mcp_name"] - mcp_url = preset["threatbook_mcp_url"] - if not mcp_name or not mcp_url: + if not mcp_name: return { "success": True, "message": "MCP not required for this region", } - safe_key = quote(api_key, safe="") request = McpTestRequest( name=mcp_name, config={ "type": "remote", - "url": mcp_url.format(api_key=safe_key), + "url": build_threatbook_mcp_url(region, api_key), }, ) return await test_mcp_connection(request) @@ -776,7 +765,11 @@ def _ensure_threatbook_mcp_config(region: Region) -> None: config = { "type": "remote", - "url": f"https://mcp.threatbook.cn/mcp?apikey={{secret:{mcp_secret_id}}}", + "url": build_threatbook_mcp_url( + region, + f"{{secret:{mcp_secret_id}}}", + encode_key=False, + ), "enabled": True, } ConfigWriter.add_mcp_server(mcp_name, config) diff --git a/flocks/server/threatbook_regions.py b/flocks/server/threatbook_regions.py new file mode 100644 index 000000000..d752a02ff --- /dev/null +++ b/flocks/server/threatbook_regions.py @@ -0,0 +1,54 @@ +"""Canonical ThreatBook regional endpoints shared by guided setup flows.""" + +from __future__ import annotations + +from typing import Any, Dict, Literal +from urllib.parse import quote + + +ThreatBookRegion = Literal["cn", "global"] + +THREATBOOK_REGION_PRESETS: Dict[ThreatBookRegion, Dict[str, Any]] = { + "cn": { + "activation_url": "https://x.threatbook.com/flocks/activate", + "threatbook_llm_provider_id": "threatbook-cn-llm", + "threatbook_default_model_id": "deepseek-v4-flash-0731", + "threatbook_api_service_id": "threatbook-cn", + "threatbook_mcp_name": "threatbook_mcp", + "threatbook_mcp_endpoint": "https://mcp.threatbook.cn/mcp", + "threatbook_mcp_secret_id": "threatbook_mcp_key", + "requires_mcp": True, + }, + "global": { + "activation_url": "https://i.threatbook.io/flocks/activate", + "threatbook_llm_provider_id": "threatbook-io-llm", + "threatbook_default_model_id": "deepseek-v4-flash-0731", + "threatbook_api_service_id": "threatbook-io", + "threatbook_mcp_name": "threatbook_mcp", + "threatbook_mcp_endpoint": "https://mcp.threatbook.io/mcp", + "threatbook_mcp_secret_id": "threatbook_mcp_key", + "requires_mcp": True, + }, +} + + +def build_threatbook_mcp_url( + region: ThreatBookRegion, + api_key: str, + *, + encode_key: bool = True, +) -> str: + """Build a regional MCP URL from a key or secret placeholder.""" + endpoint = THREATBOOK_REGION_PRESETS[region]["threatbook_mcp_endpoint"] + value = quote(api_key, safe="") if encode_key else api_key + return f"{endpoint}?apikey={value}" + + +def infer_threatbook_mcp_region(url: str | None) -> ThreatBookRegion | None: + """Infer the configured region from a persisted MCP URL.""" + normalized_url = (url or "").strip().lower() + if normalized_url.startswith(THREATBOOK_REGION_PRESETS["global"]["threatbook_mcp_endpoint"]): + return "global" + if normalized_url.startswith(THREATBOOK_REGION_PRESETS["cn"]["threatbook_mcp_endpoint"]): + return "cn" + return None diff --git a/tests/server/routes/test_mcp_routes.py b/tests/server/routes/test_mcp_routes.py index 50beef95d..ccbf3e2f4 100644 --- a/tests/server/routes/test_mcp_routes.py +++ b/tests/server/routes/test_mcp_routes.py @@ -1123,3 +1123,187 @@ async def fake_connect(name: str, config: dict) -> bool: assert resp.status_code == 504, resp.text assert "timed out" in resp.text.lower() + + @pytest.mark.asyncio + async def test_configure_threatbook_mcp_validates_then_saves_global_region( + self, client: AsyncClient, monkeypatch: pytest.MonkeyPatch + ): + raw_state: dict = {"mcp": {}} + secret_state: dict[str, str] = {} + validated_configs: list[dict] = [] + connected_configs: list[dict] = [] + + class FakeSecrets: + def _load(self): + return dict(secret_state) + + def _save(self, data): + secret_state.clear() + secret_state.update(data) + + def set(self, secret_id: str, value: str): + secret_state[secret_id] = value + + async def fake_test(request): + validated_configs.append(request.config) + return {"success": True, "message": "ok", "tools_count": 4} + + async def fake_status(): + return {} + + async def fake_load(name: str): + return raw_state["mcp"].get(name) + + async def fake_connect(name: str, config: dict): + connected_configs.append(config) + return True + + async def fake_refresh(name: str): + return 4 + + def fake_persist(name: str, config: dict): + raw_state["mcp"][name] = dict(config) + + monkeypatch.setattr(mcp_routes, "get_secret_manager", lambda: FakeSecrets()) + monkeypatch.setattr(mcp_routes, "test_mcp_connection", fake_test) + monkeypatch.setattr(mcp_routes, "_load_mcp_server_config", fake_load) + monkeypatch.setattr(mcp_routes, "_persist_mcp_server_config", fake_persist) + monkeypatch.setattr(mcp_routes.ConfigWriter, "_read_raw", lambda: raw_state.copy()) + monkeypatch.setattr(mcp_routes, "_load_raw_mcp_server_config", lambda name: raw_state["mcp"].get(name)) + monkeypatch.setattr(mcp_routes.MCP, "status", fake_status) + monkeypatch.setattr(mcp_routes.MCP, "connect", fake_connect) + monkeypatch.setattr(mcp_routes.MCP, "refresh_tools", fake_refresh) + + resp = await client.post( + "/api/mcp/threatbook_mcp/threatbook-configure", + json={"region": "global", "api_key": "global key"}, + ) + + assert resp.status_code == 200, resp.text + assert resp.json()["connected"] is True + assert validated_configs[0]["url"] == ( + "https://mcp.threatbook.io/mcp?apikey=global%20key" + ) + assert raw_state["mcp"]["threatbook_mcp"]["url"] == ( + "https://mcp.threatbook.io/mcp?apikey={secret:threatbook_mcp_key}" + ) + assert secret_state["threatbook_mcp_key"] == "global key" + assert connected_configs + + @pytest.mark.asyncio + async def test_get_threatbook_mcp_credentials_uses_non_duplicated_secret_id( + self, client: AsyncClient, monkeypatch: pytest.MonkeyPatch + ): + requested_ids: list[str] = [] + + class FakeSecrets: + def get(self, secret_id: str): + requested_ids.append(secret_id) + return "configured-key" if secret_id == "threatbook_mcp_key" else None + + monkeypatch.setattr(mcp_routes, "get_secret_manager", lambda: FakeSecrets()) + + resp = await client.get("/api/mcp/threatbook_mcp/credentials") + + assert resp.status_code == 200, resp.text + assert resp.json()["has_credential"] is True + assert resp.json()["secret_id"] == "threatbook_mcp_key" + assert requested_ids == ["threatbook_mcp_key"] + + @pytest.mark.asyncio + async def test_configure_threatbook_mcp_does_not_persist_failed_validation( + self, client: AsyncClient, monkeypatch: pytest.MonkeyPatch + ): + persisted: list[dict] = [] + + async def fake_test(request): + return {"success": False, "message": "invalid regional key"} + + monkeypatch.setattr(mcp_routes, "test_mcp_connection", fake_test) + monkeypatch.setattr( + mcp_routes, + "_persist_mcp_server_config", + lambda name, config: persisted.append(config), + ) + + resp = await client.post( + "/api/mcp/threatbook_mcp/threatbook-configure", + json={"region": "cn", "api_key": "bad-key"}, + ) + + assert resp.status_code == 200, resp.text + assert resp.json()["success"] is False + assert resp.json()["message"] == "invalid regional key" + assert persisted == [] + + @pytest.mark.asyncio + async def test_configure_threatbook_mcp_restores_previous_state_when_apply_fails( + self, client: AsyncClient, monkeypatch: pytest.MonkeyPatch + ): + old_config = { + "type": "remote", + "url": "https://mcp.threatbook.cn/mcp?apikey={secret:threatbook_mcp_key}", + "enabled": True, + } + raw_state = {"mcp": {"threatbook_mcp": dict(old_config)}} + secret_state = {"threatbook_mcp_key": "old-key"} + connect_attempts: list[str] = [] + + class FakeSecrets: + def _load(self): + return dict(secret_state) + + def _save(self, data): + secret_state.clear() + secret_state.update(data) + + def set(self, secret_id: str, value: str): + secret_state[secret_id] = value + + async def fake_test(request): + return {"success": True, "message": "ok"} + + async def fake_status(): + return { + "threatbook_mcp": McpStatusInfo(status=McpStatus.CONNECTED), + } + + async def fake_remove(name: str): + return True + + async def fake_load(name: str): + return raw_state["mcp"].get(name) + + async def fake_connect(name: str, config: dict): + connect_attempts.append(config["url"]) + return len(connect_attempts) > 1 + + def fake_write_raw(data: dict): + raw_state.clear() + raw_state.update(data) + + def fake_persist(name: str, config: dict): + raw_state["mcp"][name] = dict(config) + + monkeypatch.setattr(mcp_routes, "get_secret_manager", lambda: FakeSecrets()) + monkeypatch.setattr(mcp_routes, "test_mcp_connection", fake_test) + monkeypatch.setattr(mcp_routes, "_load_mcp_server_config", fake_load) + monkeypatch.setattr(mcp_routes, "_load_raw_mcp_server_config", lambda name: raw_state["mcp"].get(name)) + monkeypatch.setattr(mcp_routes, "_persist_mcp_server_config", fake_persist) + monkeypatch.setattr(mcp_routes.ConfigWriter, "_read_raw", lambda: {"mcp": {"threatbook_mcp": dict(raw_state["mcp"]["threatbook_mcp"])}}) + monkeypatch.setattr(mcp_routes.ConfigWriter, "_write_raw", fake_write_raw) + monkeypatch.setattr(mcp_routes.MCP, "status", fake_status) + monkeypatch.setattr(mcp_routes.MCP, "remove", fake_remove) + monkeypatch.setattr(mcp_routes.MCP, "connect", fake_connect) + monkeypatch.setattr(tool_loader, "save_mcp_config", lambda name, config: None) + monkeypatch.setattr(tool_loader, "delete_mcp_config", lambda name: True) + + resp = await client.post( + "/api/mcp/threatbook_mcp/threatbook-configure", + json={"region": "global", "api_key": "new-key"}, + ) + + assert resp.status_code == 500, resp.text + assert raw_state["mcp"]["threatbook_mcp"] == old_config + assert secret_state["threatbook_mcp_key"] == "old-key" + assert connect_attempts[-1] == old_config["url"] diff --git a/tests/server/routes/test_onboarding_routes.py b/tests/server/routes/test_onboarding_routes.py index 3261d5e7b..5a2bff1d2 100644 --- a/tests/server/routes/test_onboarding_routes.py +++ b/tests/server/routes/test_onboarding_routes.py @@ -5,6 +5,9 @@ from flocks.server.routes import onboarding as onboarding_routes +REAL_BUILD_THREATBOOK_INTEL_STATUS = onboarding_routes._build_threatbook_intel_status + + def _intel_status(**overrides): data = { "configured": False, @@ -17,7 +20,7 @@ def _intel_status(**overrides): "mcp_name": "threatbook_mcp", "service_matrix": { "cn": ["api", "mcp"], - "global": ["api"], + "global": ["api", "mcp"], }, } data.update(overrides) @@ -113,6 +116,38 @@ async def fake_intel_status(): assert data["threatbook_intel"]["api_configured"] is True assert data["threatbook_intel"]["mcp_connected"] is True + @pytest.mark.asyncio + async def test_intel_status_infers_global_region_from_mcp_endpoint( + self, monkeypatch: pytest.MonkeyPatch + ): + async def fake_api_credentials(service_id: str): + return service_id == "threatbook-io" + + async def fake_mcp_status(name: str): + return { + "status": "connected", + "configured": True, + "connected": True, + "has_credential": True, + } + + monkeypatch.setattr(onboarding_routes, "_api_service_has_credentials", fake_api_credentials) + monkeypatch.setattr(onboarding_routes, "_detect_mcp_status", fake_mcp_status) + monkeypatch.setattr( + onboarding_routes.ConfigWriter, + "get_mcp_server", + lambda name: { + "url": "https://mcp.threatbook.io/mcp?apikey={secret:threatbook_mcp_key}", + }, + ) + + status = await REAL_BUILD_THREATBOOK_INTEL_STATUS() + + assert status.region == "global" + assert status.api_service_id == "threatbook-io" + assert status.mcp_connected is True + assert status.service_matrix["global"] == ["api", "mcp"] + class TestOnboardingValidateRoutes: @pytest.mark.asyncio @@ -307,6 +342,10 @@ def test_threatbook_region_presets_use_deepseek_v4_flash_0731(self): assert onboarding_routes.ONBOARDING_REGION_PRESETS["global"]["activation_url"] == ( "https://i.threatbook.io/flocks/activate" ) + assert onboarding_routes.ONBOARDING_REGION_PRESETS["global"]["threatbook_mcp_endpoint"] == ( + "https://mcp.threatbook.io/mcp" + ) + assert onboarding_routes.ONBOARDING_REGION_PRESETS["global"]["requires_mcp"] is True def test_ensure_threatbook_mcp_config_uses_explicit_secret_reference( self, monkeypatch: pytest.MonkeyPatch @@ -340,6 +379,12 @@ def fake_save_mcp_config(name: str, config: dict): "https://mcp.threatbook.cn/mcp?apikey={secret:threatbook_mcp_key}" ) + onboarding_routes._ensure_threatbook_mcp_config("global") + + assert captured["config"]["url"] == ( + "https://mcp.threatbook.io/mcp?apikey={secret:threatbook_mcp_key}" + ) + @pytest.mark.asyncio async def test_apply_cn_threatbook_model_configures_llm_api_mcp_and_default( self, client, monkeypatch: pytest.MonkeyPatch @@ -584,7 +629,7 @@ async def fake_set_default_model(model_type, body): data = resp.json() assert data["success"] is True assert data["threatbook_enabled"] is False - assert set(data["skipped"]) == {"threatbook_api"} + assert set(data["skipped"]) == {"threatbook_api", "threatbook_mcp"} assert ("provider", "openai") in calls assert ("default_model", "openai") in calls diff --git a/webui/src/api/mcp.ts b/webui/src/api/mcp.ts index ff2f4b2d9..72de79044 100644 --- a/webui/src/api/mcp.ts +++ b/webui/src/api/mcp.ts @@ -8,6 +8,7 @@ import type { MCPCatalogCategory, MCPCatalogStats, } from '@/types'; +import type { ThreatBookRegion } from '@/constants/threatbook'; export type { MCPServer, MCPServerDetail }; @@ -66,6 +67,16 @@ export const mcpAPI = { { config } ), + configureThreatBook: (server: string, payload: { region: ThreatBookRegion; api_key: string }) => + client.post<{ + success: boolean; + message: string; + region: ThreatBookRegion; + endpoint: string; + connected: boolean; + tools_count: number; + }>(`/api/mcp/${server}/threatbook-configure`, payload), + // Catalog catalogList: () => client.get('/api/mcp/catalog/entries'), diff --git a/webui/src/components/common/OnboardingModal.tsx b/webui/src/components/common/OnboardingModal.tsx index baccc78fd..f58f84eaf 100644 --- a/webui/src/components/common/OnboardingModal.tsx +++ b/webui/src/components/common/OnboardingModal.tsx @@ -14,14 +14,10 @@ import { type OnboardingValidateResponse, } from '@/api/onboarding'; import type { CatalogProvider } from '@/types'; +import { getDefaultThreatBookRegion, THREATBOOK_REGION_CONFIG } from '@/constants/threatbook'; const MODEL_KEY_LINK = 'https://portal.agentflocks.com'; -const TBCLOUD_LINKS: Record = { - cn: 'https://x.threatbook.com/flocks/activate', - global: 'https://i.threatbook.io/flocks/activate', -}; - const THREATBOOK_PROVIDER_IDS = ['threatbook-cn-llm', 'threatbook-io-llm'] as const; const THREATBOOK_FREE_PROVIDER_OPTION = 'threatbook-free'; @@ -58,10 +54,6 @@ function regionForProvider(providerId: string | null | undefined): OnboardingReg return providerId === 'threatbook-io-llm' ? 'global' : 'cn'; } -function getDefaultIntelRegion(language: string | undefined): OnboardingRegion { - return language?.toLowerCase().startsWith('en') ? 'global' : 'cn'; -} - function statusStyles(tone: SectionTone) { if (tone === 'success') { return { @@ -316,7 +308,7 @@ function SkipConfirmDialog({ export default function OnboardingModal({ onClose }: OnboardingModalProps) { const { t, i18n } = useTranslation('common'); const navigate = useNavigate(); - const defaultIntelRegion = useMemo(() => getDefaultIntelRegion(i18n.language), [i18n.language]); + const defaultIntelRegion = useMemo(() => getDefaultThreatBookRegion(i18n.language), [i18n.language]); const [step, setStep] = useState('model'); const [catalog, setCatalog] = useState([]); @@ -338,7 +330,7 @@ export default function OnboardingModal({ onClose }: OnboardingModalProps) { const [modelRegion, setModelRegion] = useState('cn'); const [modelSkipped, setModelSkipped] = useState(false); - const [intelRegion, setIntelRegion] = useState(() => getDefaultIntelRegion(i18n.language)); + const [intelRegion, setIntelRegion] = useState(() => getDefaultThreatBookRegion(i18n.language)); const [intelApiKey, setIntelApiKey] = useState(''); const [intelSaving, setIntelSaving] = useState(false); const [intelConfigured, setIntelConfigured] = useState(false); @@ -1092,7 +1084,7 @@ export default function OnboardingModal({ onClose }: OnboardingModalProps) { className="min-w-0 flex-1 rounded-lg border border-gray-200 bg-white px-3 py-2 text-xs transition-all placeholder-gray-300 focus:border-red-400 focus:outline-none focus:ring-2 focus:ring-red-400/50" /> = { + cn: { + activationUrl: 'https://x.threatbook.com/flocks/activate', + mcpEndpoint: 'https://mcp.threatbook.cn/mcp', + }, + global: { + activationUrl: 'https://i.threatbook.io/flocks/activate', + mcpEndpoint: 'https://mcp.threatbook.io/mcp', + }, +}; + +export function getDefaultThreatBookRegion(language: string | undefined): ThreatBookRegion { + return language?.toLowerCase().startsWith('en') ? 'global' : 'cn'; +} + +export function inferThreatBookRegionFromMcpUrl(url: string | undefined): ThreatBookRegion | null { + const normalizedUrl = (url || '').trim().toLowerCase(); + if (normalizedUrl.startsWith(THREATBOOK_REGION_CONFIG.global.mcpEndpoint)) return 'global'; + if (normalizedUrl.startsWith(THREATBOOK_REGION_CONFIG.cn.mcpEndpoint)) return 'cn'; + return null; +} diff --git a/webui/src/locales/en-US/tool.json b/webui/src/locales/en-US/tool.json index 40d7a83f7..9b5b00b03 100644 --- a/webui/src/locales/en-US/tool.json +++ b/webui/src/locales/en-US/tool.json @@ -357,7 +357,6 @@ "quickActions": "Quick Actions", "testingConn": "Testing...", "testConnection": "Test Connection", - "claimFreeApiKey": "Claim free API key", "refreshingTools": "Refreshing...", "refreshTools": "Refresh Tools", "disconnectConn": "Disconnect", @@ -387,7 +386,38 @@ "disableServer": "Disable Server", "enableServer": "Enable Server", "hide": "Hide", - "show": "Show" + "show": "Show", + "threatbookMcp": { + "title": "Configure ThreatBook MCP", + "description": "Enable the free ThreatBook intelligence MCP service for IOC lookup, threat analysis, and security operations in Rex.", + "region": "Service region", + "regionHint": "Choose the region for your API key. Flocks will match the claim page and MCP endpoint automatically.", + "regions": { + "cn": "China", + "global": "International" + }, + "freeService": "Free ThreatBook MCP service", + "apiKey": "API Key", + "keyPlaceholder": "Paste the ThreatBook API key for this region", + "keyHint": "After claiming a key, paste it here. Flocks will build the service URL and store the key securely.", + "keyRequired": "Enter the API key for the selected region.", + "keyConfigured": "Securely configured", + "claimFreeKey": "Claim free API key", + "endpoint": "MCP endpoint", + "endpointHint": "The endpoint is selected automatically. Your API key is stored separately and never needs to be appended manually.", + "saveAndVerify": "Save and verify connection", + "saving": "Verifying and saving...", + "saveSuccess": "ThreatBook MCP is configured and verified.", + "saveFailed": "ThreatBook MCP setup failed. Check the region and API key.", + "configuredTitle": "{{region}} configured", + "connected": "The configuration is verified and the MCP service is connected.", + "savedNotConnected": "The configuration is saved, but the MCP service is not connected.", + "edit": "Edit configuration", + "retest": "Test again", + "testing": "Testing...", + "testFailed": "Connection test failed.", + "loading": "Loading ThreatBook MCP configuration..." + } }, "credentials": { diff --git a/webui/src/locales/zh-CN/tool.json b/webui/src/locales/zh-CN/tool.json index 1a3078d8f..7b8f2f167 100644 --- a/webui/src/locales/zh-CN/tool.json +++ b/webui/src/locales/zh-CN/tool.json @@ -357,7 +357,6 @@ "quickActions": "快捷操作", "testingConn": "测试中...", "testConnection": "测试连接", - "claimFreeApiKey": "免费领取 API Key", "refreshingTools": "刷新中...", "refreshTools": "刷新工具列表", "disconnectConn": "断开连接", @@ -387,7 +386,38 @@ "disableServer": "禁用服务器", "enableServer": "启用服务器", "hide": "隐藏", - "show": "显示" + "show": "显示", + "threatbookMcp": { + "title": "配置 ThreatBook MCP", + "description": "免费启用微步威胁情报 MCP 服务，为 Rex 提供 IOC 查询、威胁情报研判和安全分析能力。", + "region": "服务区域", + "regionHint": "请选择 API Key 所属区域，系统会自动匹配对应的领取入口和 MCP 服务地址。", + "regions": { + "cn": "中国区", + "global": "国际区" + }, + "freeService": "免费 ThreatBook MCP 服务", + "apiKey": "API Key", + "keyPlaceholder": "粘贴当前区域的 ThreatBook API Key", + "keyHint": "领取后将 Key 粘贴到此处，系统会自动完成服务地址拼接和安全存储。", + "keyRequired": "请先填写当前区域的 API Key。", + "keyConfigured": "已安全配置", + "claimFreeKey": "领取免费 API Key", + "endpoint": "MCP 服务地址", + "endpointHint": "服务地址由所选区域自动生成，API Key 将单独加密保存，无需手动拼接。", + "saveAndVerify": "保存并验证连接", + "saving": "正在验证并保存...", + "saveSuccess": "ThreatBook MCP 已配置并验证通过。", + "saveFailed": "ThreatBook MCP 配置失败，请检查区域和 API Key。", + "configuredTitle": "已配置{{region}}", + "connected": "配置已验证，MCP 服务当前已连接。", + "savedNotConnected": "配置已保存，MCP 服务当前未连接。", + "edit": "编辑配置", + "retest": "重新测试", + "testing": "测试中...", + "testFailed": "连接测试失败。", + "loading": "正在读取 ThreatBook MCP 配置..." + } }, "credentials": { diff --git a/webui/src/pages/Tool/ToolSheets.tsx b/webui/src/pages/Tool/ToolSheets.tsx index 442f7cadf..5ba1bfff0 100644 --- a/webui/src/pages/Tool/ToolSheets.tsx +++ b/webui/src/pages/Tool/ToolSheets.tsx @@ -6,7 +6,7 @@ * - GenerateToolSheet: AI 生成自定义工具 */ -import { useMemo, useState, type ReactNode } from 'react'; +import { useMemo, useState } from 'react'; import { useTranslation } from 'react-i18next'; import { Database, Cloud, Code, Info, CheckCircle, XCircle, Activity, Wifi, WifiOff, Loader2 } from 'lucide-react'; import EntitySheet from '@/components/common/EntitySheet'; @@ -252,7 +252,6 @@ interface MCPFormFieldsProps { testResult: { success: boolean; message: string; tools_count?: number } | null; onTestConnection: () => void; isTesting: boolean; - serviceUrlAction?: ReactNode; } export function MCPFormFields({ @@ -263,7 +262,6 @@ export function MCPFormFields({ testResult, onTestConnection, isTesting, - serviceUrlAction, }: MCPFormFieldsProps) { const { t } = useTranslation('tool'); const readOnly = !onChange; @@ -359,12 +357,9 @@ export function MCPFormFields({ {formData.connType === 'sse' && (
-
- - {serviceUrlAction} -
+
({ currentLanguage: { value: 'zh-CN' }, mcpAPI: { get: vi.fn(), + getCredentials: vi.fn(), + configureThreatBook: vi.fn(), + testCredentials: vi.fn(), update: vi.fn(), testExisting: vi.fn(), }, @@ -68,13 +71,11 @@ vi.mock('../ToolSheets', () => ({ onChange, onTestConnection, testResult, - serviceUrlAction, }: { formData: { url: string }; onChange?: (fields: { url: string }) => void; onTestConnection: () => void; testResult: { message: string } | null; - serviceUrlAction?: React.ReactNode; }) => (
({ - {serviceUrlAction} {testResult &&
{testResult.message}
}
), @@ -93,7 +93,8 @@ vi.mock('../ToolSheets', () => ({ vi.mock('react-i18next', () => ({ useTranslation: () => ({ - t: (key: string) => { + t: (key: string, options?: Record) => { + const isChinese = currentLanguage.value.startsWith('zh'); const translations: Record = { 'detail.tabs.overview': '概览', 'detail.tabs.tools': '工具', @@ -110,10 +111,37 @@ vi.mock('react-i18next', () => ({ 'button.save': '保存', 'button.saving': '保存中...', 'detail.testFailed': '连接测试失败', - 'detail.claimFreeApiKey': currentLanguage.value.startsWith('zh') ? '免费领取 API Key' : 'Claim free API key', + 'detail.show': isChinese ? '显示' : 'Show', + 'detail.hide': isChinese ? '隐藏' : 'Hide', + 'detail.threatbookMcp.title': isChinese ? '配置 ThreatBook MCP' : 'Configure ThreatBook MCP', + 'detail.threatbookMcp.description': isChinese ? '免费启用微步威胁情报 MCP 服务' : 'Enable the free ThreatBook intelligence MCP service', + 'detail.threatbookMcp.region': isChinese ? '服务区域' : 'Service region', + 'detail.threatbookMcp.regionHint': isChinese ? '选择区域' : 'Choose region', + 'detail.threatbookMcp.regions.cn': isChinese ? '中国区' : 'China', + 'detail.threatbookMcp.regions.global': isChinese ? '国际区' : 'International', + 'detail.threatbookMcp.freeService': isChinese ? '免费 ThreatBook MCP 服务' : 'Free ThreatBook MCP service', + 'detail.threatbookMcp.apiKey': 'API Key', + 'detail.threatbookMcp.keyPlaceholder': isChinese ? '粘贴当前区域的 ThreatBook API Key' : 'Paste the ThreatBook API key for this region', + 'detail.threatbookMcp.keyHint': isChinese ? '领取后填写' : 'Paste after claiming', + 'detail.threatbookMcp.keyRequired': isChinese ? '请填写 Key' : 'Enter API key', + 'detail.threatbookMcp.claimFreeKey': isChinese ? '领取免费 API Key' : 'Claim free API key', + 'detail.threatbookMcp.endpoint': isChinese ? 'MCP 服务地址' : 'MCP endpoint', + 'detail.threatbookMcp.endpointHint': isChinese ? '自动生成地址' : 'Endpoint is generated automatically', + 'detail.threatbookMcp.saveAndVerify': isChinese ? '保存并验证连接' : 'Save and verify connection', + 'detail.threatbookMcp.saving': isChinese ? '正在验证并保存...' : 'Verifying and saving...', + 'detail.threatbookMcp.saveSuccess': isChinese ? '配置成功' : 'Configured', + 'detail.threatbookMcp.saveFailed': isChinese ? '配置失败' : 'Setup failed', + 'detail.threatbookMcp.loading': isChinese ? '读取配置' : 'Loading configuration', + 'detail.threatbookMcp.configuredTitle': isChinese ? '已配置{{region}}' : '{{region}} configured', + 'detail.threatbookMcp.connected': isChinese ? '当前已连接' : 'Connected', + 'detail.threatbookMcp.savedNotConnected': isChinese ? '当前未连接' : 'Not connected', + 'detail.threatbookMcp.keyConfigured': isChinese ? '已安全配置' : 'Securely configured', + 'detail.threatbookMcp.edit': isChinese ? '编辑配置' : 'Edit configuration', + 'detail.threatbookMcp.retest': isChinese ? '重新测试' : 'Test again', + 'detail.threatbookMcp.testing': isChinese ? '测试中...' : 'Testing...', 'alert.connectionOk': '连接成功', }; - return translations[key] ?? key; + return (translations[key] ?? key).replace('{{region}}', String(options?.region ?? '')); }, i18n: { language: currentLanguage.value, @@ -155,6 +183,22 @@ describe('MCPServerDetailPanel', () => { vi.clearAllMocks(); currentLanguage.value = 'zh-CN'; mcpAPI.get.mockResolvedValue(detailResponse); + mcpAPI.getCredentials.mockResolvedValue({ + data: { has_credential: false }, + }); + mcpAPI.configureThreatBook.mockResolvedValue({ + data: { + success: true, + message: 'ok', + region: 'cn', + endpoint: 'https://mcp.threatbook.cn/mcp', + connected: true, + tools_count: 3, + }, + }); + mcpAPI.testCredentials.mockResolvedValue({ + data: { success: true, message: 'ok', tools_count: 3 }, + }); mcpAPI.update.mockResolvedValue({ data: { success: true }, }); @@ -248,9 +292,11 @@ describe('MCPServerDetailPanel', () => { />, ); - const link = await screen.findByRole('link', { name: '免费领取 API Key' }); + const link = await screen.findByRole('link', { name: '领取免费 API Key' }); expect(link).toHaveAttribute('href', 'https://x.threatbook.com/flocks/activate'); expect(link).toHaveAttribute('target', '_blank'); + expect(screen.getByRole('button', { name: '中国区' })).toBeInTheDocument(); + expect(screen.getByDisplayValue('https://mcp.threatbook.cn/mcp')).toBeInTheDocument(); }); it('shows the international free key link for ThreatBook MCP in English', async () => { @@ -270,6 +316,74 @@ describe('MCPServerDetailPanel', () => { const link = await screen.findByRole('link', { name: 'Claim free API key' }); expect(link).toHaveAttribute('href', 'https://i.threatbook.io/flocks/activate'); + expect(screen.getByRole('button', { name: 'International' })).toBeInTheDocument(); + expect(screen.getByDisplayValue('https://mcp.threatbook.io/mcp')).toBeInTheDocument(); + }); + + it('sends the selected region and API key through the closed-loop setup action', async () => { + const user = userEvent.setup(); + + render( + , + ); + + await user.click(await screen.findByRole('button', { name: '国际区' })); + await user.type(screen.getByPlaceholderText('粘贴当前区域的 ThreatBook API Key'), 'global-key'); + await user.click(screen.getByRole('button', { name: '保存并验证连接' })); + + await waitFor(() => { + expect(mcpAPI.configureThreatBook).toHaveBeenCalledWith('threatbook_mcp', { + region: 'global', + api_key: 'global-key', + }); + }); + }); + + it('shows the configured region instead of changing it with the interface language', async () => { + currentLanguage.value = 'en-US'; + mcpAPI.get.mockResolvedValue({ + ...detailResponse, + data: { + ...detailResponse.data, + name: 'threatbook_mcp', + config: { + type: 'sse', + url: 'https://mcp.threatbook.cn/mcp?apikey={secret:threatbook_mcp_key}', + }, + }, + }); + mcpAPI.getCredentials.mockResolvedValue({ + data: { + has_credential: true, + secret_id: 'threatbook_mcp_key', + api_key_masked: 'conf****-key', + }, + }); + + render( + , + ); + + expect(await screen.findByText('China configured')).toBeInTheDocument(); + expect(screen.getByText('https://mcp.threatbook.cn/mcp')).toBeInTheDocument(); + expect(screen.getByRole('button', { name: 'Edit configuration' })).toBeInTheDocument(); }); it('does not show the free key link for other MCP servers', async () => { @@ -286,6 +400,6 @@ describe('MCPServerDetailPanel', () => { ); await screen.findByLabelText('service-url'); - expect(screen.queryByRole('link', { name: '免费领取 API Key' })).not.toBeInTheDocument(); + expect(screen.queryByRole('link', { name: '领取免费 API Key' })).not.toBeInTheDocument(); }); }); diff --git a/webui/src/pages/Tool/components/ServiceDetailPanel.tsx b/webui/src/pages/Tool/components/ServiceDetailPanel.tsx index b0b172ed4..eb3686991 100644 --- a/webui/src/pages/Tool/components/ServiceDetailPanel.tsx +++ b/webui/src/pages/Tool/components/ServiceDetailPanel.tsx @@ -3,7 +3,6 @@ import { useTranslation } from 'react-i18next'; import { Info, Wrench, FileText, Activity, Zap, RefreshCw, Power, PowerOff, CheckCircle, XCircle, Cloud, Database, AlertTriangle, Eye, EyeOff, Save, Trash2, - ExternalLink, } from 'lucide-react'; import type { Tool } from '@/api/tool'; import type { MCPCatalogCategory, MCPCatalogEntry, MCPCredentials, MCPServer, MCPServerDetail } from '@/types'; @@ -16,22 +15,10 @@ import { EnabledBadge } from './badges'; import { buildMCPConfigFromForm, buildMCPFormDataFromConfig, getMCPFormError, MCPFormFields } from '../ToolSheets'; import type { MCPFormData, ConnStatus as MCPConnStatus } from '../ToolSheets'; import type { APIServiceCredentialField, APIServiceMetadata, ProviderCredentials } from '@/types'; - -const THREATBOOK_MCP_SERVER_NAMES = new Set(['threatbook_mcp', 'threatbook-mcp']); -const THREATBOOK_MCP_ACTIVATION_URLS = { - cn: 'https://x.threatbook.com/flocks/activate', - global: 'https://i.threatbook.io/flocks/activate', -} as const; +import ThreatBookMCPConfigPanel from './ThreatBookMCPConfigPanel'; function isThreatBookMCPServer(serverName: string): boolean { - return THREATBOOK_MCP_SERVER_NAMES.has(serverName.trim().toLowerCase()); -} - -export function getThreatBookMCPActivationUrl(language: string): string { - const normalizedLanguage = language.toLowerCase().replace('_', '-'); - return normalizedLanguage.startsWith('zh') - ? THREATBOOK_MCP_ACTIVATION_URLS.cn - : THREATBOOK_MCP_ACTIVATION_URLS.global; + return serverName.trim().toLowerCase() === 'threatbook_mcp'; } function KvRowValue({ value }: { value: string }) { @@ -90,7 +77,7 @@ export function MCPServerDetailPanel({ onRemove?: () => void; onSelectTool: (tool: Tool) => void; }) { - const { t, i18n } = useTranslation('tool'); + const { t } = useTranslation('tool'); const [detailTab, setDetailTab] = useState<'overview' | 'tools' | 'resources'>('overview'); const [serverDetail, setServerDetail] = useState(null); const [formData, setFormData] = useState(null); @@ -100,9 +87,7 @@ export function MCPServerDetailPanel({ const [testResult, setTestResult] = useState<{ success: boolean; message: string; latency?: number; tools_count?: number } | null>(null); const [refreshing, setRefreshing] = useState(false); const [savingConfig, setSavingConfig] = useState(false); - const threatBookActivationUrl = isThreatBookMCPServer(server.name) - ? getThreatBookMCPActivationUrl(i18n.resolvedLanguage || i18n.language) - : null; + const isThreatBookServer = isThreatBookMCPServer(server.name); const createFormData = useCallback((detail: MCPServerDetail | null): MCPFormData => ( buildMCPFormDataFromConfig(server.name, detail?.config, server.url) @@ -246,7 +231,17 @@ export function MCPServerDetailPanel({
) : detailTab === 'overview' ? (
- {formData && ( + {isThreatBookServer ? ( + { + await loadServerDetail(); + await onStatusChange?.(); + }} + /> + ) : formData && ( - - {t('detail.claimFreeApiKey')} - - ) : undefined} /> )} diff --git a/webui/src/pages/Tool/components/ThreatBookMCPConfigPanel.tsx b/webui/src/pages/Tool/components/ThreatBookMCPConfigPanel.tsx new file mode 100644 index 000000000..561882eb4 --- /dev/null +++ b/webui/src/pages/Tool/components/ThreatBookMCPConfigPanel.tsx @@ -0,0 +1,371 @@ +import { useEffect, useMemo, useState } from 'react'; +import { useTranslation } from 'react-i18next'; +import { + AlertCircle, + CheckCircle2, + ExternalLink, + Eye, + EyeOff, + Loader2, + Pencil, + RefreshCw, +} from 'lucide-react'; +import { mcpAPI } from '@/api/mcp'; +import { + getDefaultThreatBookRegion, + inferThreatBookRegionFromMcpUrl, + THREATBOOK_REGION_CONFIG, + type ThreatBookRegion, +} from '@/constants/threatbook'; +import type { MCPCredentials, MCPServer } from '@/types'; + +interface ThreatBookMCPConfigPanelProps { + serverName: string; + serverStatus: MCPServer['status']; + configUrl?: string; + onConfigured: () => Promise; +} + +type OperationResult = { + success: boolean; + message: string; +} | null; + +export default function ThreatBookMCPConfigPanel({ + serverName, + serverStatus, + configUrl, + onConfigured, +}: ThreatBookMCPConfigPanelProps) { + const { t, i18n } = useTranslation('tool'); + const configuredRegion = useMemo( + () => inferThreatBookRegionFromMcpUrl(configUrl), + [configUrl], + ); + const languageDefaultRegion = useMemo( + () => getDefaultThreatBookRegion(i18n.resolvedLanguage || i18n.language), + [i18n.language, i18n.resolvedLanguage], + ); + const [region, setRegion] = useState(configuredRegion || languageDefaultRegion); + const [credentials, setCredentials] = useState(null); + const [loadingCredentials, setLoadingCredentials] = useState(true); + const [editing, setEditing] = useState(true); + const [apiKey, setApiKey] = useState(''); + const [showApiKey, setShowApiKey] = useState(false); + const [saving, setSaving] = useState(false); + const [testing, setTesting] = useState(false); + const [result, setResult] = useState(null); + + const loadCredentials = async () => { + try { + setLoadingCredentials(true); + const response = await mcpAPI.getCredentials(serverName); + setCredentials(response.data); + setEditing(!(response.data.has_credential && configuredRegion)); + } catch { + setCredentials(null); + setEditing(true); + } finally { + setLoadingCredentials(false); + } + }; + + useEffect(() => { + setRegion(configuredRegion || languageDefaultRegion); + }, [configuredRegion, languageDefaultRegion]); + + useEffect(() => { + void loadCredentials(); + }, [serverName, configuredRegion]); + + const regionConfig = THREATBOOK_REGION_CONFIG[region]; + const summaryRegion = configuredRegion || region; + const isConfigured = Boolean(credentials?.has_credential && configuredRegion); + const isConnected = serverStatus === 'connected'; + + const handleRegionChange = (nextRegion: ThreatBookRegion) => { + if (nextRegion === region) return; + setRegion(nextRegion); + setApiKey(''); + setShowApiKey(false); + setResult(null); + }; + + const handleSave = async () => { + const trimmedKey = apiKey.trim(); + if (!trimmedKey) { + setResult({ success: false, message: t('detail.threatbookMcp.keyRequired') }); + return; + } + + try { + setSaving(true); + setResult(null); + const response = await mcpAPI.configureThreatBook(serverName, { + region, + api_key: trimmedKey, + }); + if (!response.data.success) { + setResult({ success: false, message: response.data.message }); + return; + } + + setApiKey(''); + setShowApiKey(false); + setResult({ success: true, message: t('detail.threatbookMcp.saveSuccess') }); + await onConfigured(); + await loadCredentials(); + setEditing(false); + } catch (error: any) { + setResult({ + success: false, + message: error.response?.data?.detail || error.message || t('detail.threatbookMcp.saveFailed'), + }); + } finally { + setSaving(false); + } + }; + + const handleRetest = async () => { + try { + setTesting(true); + setResult(null); + const response = await mcpAPI.testCredentials(serverName); + setResult({ + success: response.data.success, + message: response.data.message, + }); + if (response.data.success) await onConfigured(); + } catch (error: any) { + setResult({ + success: false, + message: error.response?.data?.detail || error.message || t('detail.threatbookMcp.testFailed'), + }); + } finally { + setTesting(false); + } + }; + + if (loadingCredentials) { + return ( +
+ + {t('detail.threatbookMcp.loading')} +
+ ); + } + + return ( +
+
+

{t('detail.threatbookMcp.title')}

+

{t('detail.threatbookMcp.description')}

+
+ + {isConfigured && !editing ? ( +
+
+
+ +
+

+ {t('detail.threatbookMcp.configuredTitle', { + region: t(`detail.threatbookMcp.regions.${summaryRegion}`), + })} +

+

+ {isConnected + ? t('detail.threatbookMcp.connected') + : t('detail.threatbookMcp.savedNotConnected')} +

+
+
+
+ + +
+
+ +
+
+
{t('detail.threatbookMcp.region')}
+
+ {t(`detail.threatbookMcp.regions.${summaryRegion}`)} +
+
+
+
{t('detail.threatbookMcp.apiKey')}
+
+ {credentials?.api_key_masked || t('detail.threatbookMcp.keyConfigured')} +
+
+
+
{t('detail.threatbookMcp.endpoint')}
+
+ {THREATBOOK_REGION_CONFIG[summaryRegion].mcpEndpoint} +
+
+
+
+ ) : ( +
+
+
+

{t('detail.threatbookMcp.region')}

+

{t('detail.threatbookMcp.regionHint')}

+
+
+ {(['cn', 'global'] as const).map((item) => ( + + ))} +
+
+ +
+ {t('detail.threatbookMcp.freeService')} +
+ +
+ +
+
+ { + setApiKey(event.target.value); + setResult(null); + }} + placeholder={t('detail.threatbookMcp.keyPlaceholder')} + autoComplete="off" + className="h-10 w-full rounded-lg border border-gray-300 bg-white px-3 pr-10 text-sm focus:border-red-400 focus:outline-none focus:ring-2 focus:ring-red-400/30" + /> + +
+ + + {t('detail.threatbookMcp.claimFreeKey')} + +
+

{t('detail.threatbookMcp.keyHint')}

+
+ +
+ + +

{t('detail.threatbookMcp.endpointHint')}

+
+ + {result && ( +
+ {result.success + ? + : } + {result.message} +
+ )} + +
+ {isConfigured && ( + + )} + +
+
+ )} + + {result && isConfigured && !editing && ( +
+ {result.success + ? + : } + {result.message} +
+ )} +
+ ); +} From f23ae17fb713b2d62c00769c11f0129bffddd318 Mon Sep 17 00:00:00 2001 From: John Yin <10972267+john-yin2333@user.noreply.gitee.com> Date: Mon, 7 Sep 2026 00:28:57 +0800 Subject: [PATCH 40/63] Reveal saved MCP keys on demand --- flocks/server/routes/mcp.py | 78 ++++++++---- tests/server/routes/test_mcp_routes.py | 34 +++++- webui/src/api/mcp.ts | 3 + webui/src/locales/en-US/tool.json | 1 + webui/src/locales/zh-CN/tool.json | 1 + .../components/ServiceDetailPanel.test.tsx | 112 +++++++++++++++++- .../components/ThreatBookMCPConfigPanel.tsx | 61 +++++++++- 7 files changed, 262 insertions(+), 28 deletions(-) diff --git a/flocks/server/routes/mcp.py b/flocks/server/routes/mcp.py index ee133822b..55f51d900 100644 --- a/flocks/server/routes/mcp.py +++ b/flocks/server/routes/mcp.py @@ -10,7 +10,7 @@ import asyncio import copy from typing import Dict, Optional, List, Any -from fastapi import APIRouter, HTTPException +from fastapi import APIRouter, HTTPException, Response from fastapi.responses import JSONResponse from pydantic import BaseModel, Field @@ -966,6 +966,38 @@ class McpCredentialResponse(BaseModel): has_credential: bool +class McpCredentialRevealResponse(BaseModel): + """Response containing a credential explicitly requested by the user.""" + api_key: str + + +def _get_mcp_credential(name: str) -> tuple[Optional[str], Optional[str]]: + """Load an MCP credential while supporting historical secret IDs.""" + secrets = get_secret_manager() + + secret_id = get_mcp_secret_id(name) + api_key = secrets.get(secret_id) + duplicated_suffix_id = f"{name}_mcp_key" + if not api_key and duplicated_suffix_id != secret_id: + api_key = secrets.get(duplicated_suffix_id) + if api_key: + secret_id = duplicated_suffix_id + if not api_key: + legacy_id = f"{name}_api_key" + api_key = secrets.get(legacy_id) + if api_key: + secret_id = legacy_id + + return (secret_id, api_key) if api_key else (None, None) + + +def _mask_mcp_credential(api_key: str) -> str: + """Mask an MCP key as three visible chars on each side.""" + if len(api_key) <= 6: + return "xxxx" + return f"{api_key[:3]}xxxx{api_key[-3:]}" + + @router.get( "/{name}/credentials", response_model=McpCredentialResponse, @@ -978,28 +1010,12 @@ async def get_mcp_credentials(name: str): Avoids duplicating an existing ``_mcp`` suffix and falls back to historical ``{name}_mcp_key`` / ``{name}_api_key`` variants for compatibility. """ - from flocks.security.secrets import SecretManager - try: - secrets = get_secret_manager() - - # Convention-based secret_id first, fall back to historical variants. - secret_id = get_mcp_secret_id(name) - api_key = secrets.get(secret_id) - duplicated_suffix_id = f"{name}_mcp_key" - if not api_key and duplicated_suffix_id != secret_id: - api_key = secrets.get(duplicated_suffix_id) - if api_key: - secret_id = duplicated_suffix_id - if not api_key: - legacy_id = f"{name}_api_key" - api_key = secrets.get(legacy_id) - if api_key: - secret_id = legacy_id + secret_id, api_key = _get_mcp_credential(name) return McpCredentialResponse( - secret_id=secret_id if api_key else None, - api_key_masked=SecretManager.mask(api_key) if api_key else None, + secret_id=secret_id, + api_key_masked=_mask_mcp_credential(api_key) if api_key else None, has_credential=bool(api_key), ) except Exception as e: @@ -1007,6 +1023,28 @@ async def get_mcp_credentials(name: str): raise HTTPException(status_code=500, detail=str(e)) +@router.post( + "/{name}/credentials/reveal", + response_model=McpCredentialRevealResponse, + summary="Reveal MCP server credential", + description="Reveal a stored MCP credential after an explicit user action." +) +async def reveal_mcp_credential(name: str, response: Response): + """Return the full key only for an explicit reveal request.""" + response.headers["Cache-Control"] = "no-store" + response.headers["Pragma"] = "no-cache" + try: + _, api_key = _get_mcp_credential(name) + if not api_key: + raise HTTPException(status_code=404, detail="No credentials found for this server") + return McpCredentialRevealResponse(api_key=api_key) + except HTTPException: + raise + except Exception as e: + log.error("mcp.credentials.reveal.error", {"name": name, "error": str(e)}) + raise HTTPException(status_code=500, detail=str(e)) + + @router.post( "/{name}/credentials", response_model=Dict[str, Any], diff --git a/tests/server/routes/test_mcp_routes.py b/tests/server/routes/test_mcp_routes.py index ccbf3e2f4..fdc59b1a1 100644 --- a/tests/server/routes/test_mcp_routes.py +++ b/tests/server/routes/test_mcp_routes.py @@ -1199,7 +1199,7 @@ async def test_get_threatbook_mcp_credentials_uses_non_duplicated_secret_id( class FakeSecrets: def get(self, secret_id: str): requested_ids.append(secret_id) - return "configured-key" if secret_id == "threatbook_mcp_key" else None + return "312abcdef321" if secret_id == "threatbook_mcp_key" else None monkeypatch.setattr(mcp_routes, "get_secret_manager", lambda: FakeSecrets()) @@ -1208,8 +1208,40 @@ def get(self, secret_id: str): assert resp.status_code == 200, resp.text assert resp.json()["has_credential"] is True assert resp.json()["secret_id"] == "threatbook_mcp_key" + assert resp.json()["api_key_masked"] == "312xxxx321" + assert "312abcdef321" not in resp.text assert requested_ids == ["threatbook_mcp_key"] + @pytest.mark.asyncio + async def test_reveal_threatbook_mcp_credentials_returns_full_key_on_demand( + self, client: AsyncClient, monkeypatch: pytest.MonkeyPatch + ): + class FakeSecrets: + def get(self, secret_id: str): + return "312abcdef321" if secret_id == "threatbook_mcp_key" else None + + monkeypatch.setattr(mcp_routes, "get_secret_manager", lambda: FakeSecrets()) + + resp = await client.post("/api/mcp/threatbook_mcp/credentials/reveal") + + assert resp.status_code == 200, resp.text + assert resp.json() == {"api_key": "312abcdef321"} + assert resp.headers["cache-control"] == "no-store" + + @pytest.mark.asyncio + async def test_reveal_mcp_credentials_returns_not_found_when_missing( + self, client: AsyncClient, monkeypatch: pytest.MonkeyPatch + ): + class FakeSecrets: + def get(self, secret_id: str): + return None + + monkeypatch.setattr(mcp_routes, "get_secret_manager", lambda: FakeSecrets()) + + resp = await client.post("/api/mcp/threatbook_mcp/credentials/reveal") + + assert resp.status_code == 404, resp.text + @pytest.mark.asyncio async def test_configure_threatbook_mcp_does_not_persist_failed_validation( self, client: AsyncClient, monkeypatch: pytest.MonkeyPatch diff --git a/webui/src/api/mcp.ts b/webui/src/api/mcp.ts index 72de79044..c19bc00ad 100644 --- a/webui/src/api/mcp.ts +++ b/webui/src/api/mcp.ts @@ -40,6 +40,9 @@ export const mcpAPI = { // Credentials management getCredentials: (server: string) => client.get(`/api/mcp/${server}/credentials`), + + revealCredentials: (server: string) => + client.post<{api_key: string}>(`/api/mcp/${server}/credentials/reveal`), setCredentials: (server: string, credentials: MCPCredentialInput) => client.post<{success: boolean; message: string}>(`/api/mcp/${server}/credentials`, credentials), diff --git a/webui/src/locales/en-US/tool.json b/webui/src/locales/en-US/tool.json index 9b5b00b03..d7e5c6a1c 100644 --- a/webui/src/locales/en-US/tool.json +++ b/webui/src/locales/en-US/tool.json @@ -401,6 +401,7 @@ "keyPlaceholder": "Paste the ThreatBook API key for this region", "keyHint": "After claiming a key, paste it here. Flocks will build the service URL and store the key securely.", "keyRequired": "Enter the API key for the selected region.", + "revealFailed": "Failed to retrieve the saved API key. Try again.", "keyConfigured": "Securely configured", "claimFreeKey": "Claim free API key", "endpoint": "MCP endpoint", diff --git a/webui/src/locales/zh-CN/tool.json b/webui/src/locales/zh-CN/tool.json index 7b8f2f167..efcb16cca 100644 --- a/webui/src/locales/zh-CN/tool.json +++ b/webui/src/locales/zh-CN/tool.json @@ -401,6 +401,7 @@ "keyPlaceholder": "粘贴当前区域的 ThreatBook API Key", "keyHint": "领取后将 Key 粘贴到此处，系统会自动完成服务地址拼接和安全存储。", "keyRequired": "请先填写当前区域的 API Key。", + "revealFailed": "读取已保存的 API Key 失败，请重试。", "keyConfigured": "已安全配置", "claimFreeKey": "领取免费 API Key", "endpoint": "MCP 服务地址", diff --git a/webui/src/pages/Tool/components/ServiceDetailPanel.test.tsx b/webui/src/pages/Tool/components/ServiceDetailPanel.test.tsx index a71a9cdcb..a9709deaa 100644 --- a/webui/src/pages/Tool/components/ServiceDetailPanel.test.tsx +++ b/webui/src/pages/Tool/components/ServiceDetailPanel.test.tsx @@ -9,6 +9,7 @@ const { currentLanguage, mcpAPI } = vi.hoisted(() => ({ mcpAPI: { get: vi.fn(), getCredentials: vi.fn(), + revealCredentials: vi.fn(), configureThreatBook: vi.fn(), testCredentials: vi.fn(), update: vi.fn(), @@ -124,6 +125,7 @@ vi.mock('react-i18next', () => ({ 'detail.threatbookMcp.keyPlaceholder': isChinese ? '粘贴当前区域的 ThreatBook API Key' : 'Paste the ThreatBook API key for this region', 'detail.threatbookMcp.keyHint': isChinese ? '领取后填写' : 'Paste after claiming', 'detail.threatbookMcp.keyRequired': isChinese ? '请填写 Key' : 'Enter API key', + 'detail.threatbookMcp.revealFailed': isChinese ? '读取 Key 失败' : 'Failed to retrieve API key', 'detail.threatbookMcp.claimFreeKey': isChinese ? '领取免费 API Key' : 'Claim free API key', 'detail.threatbookMcp.endpoint': isChinese ? 'MCP 服务地址' : 'MCP endpoint', 'detail.threatbookMcp.endpointHint': isChinese ? '自动生成地址' : 'Endpoint is generated automatically', @@ -186,6 +188,9 @@ describe('MCPServerDetailPanel', () => { mcpAPI.getCredentials.mockResolvedValue({ data: { has_credential: false }, }); + mcpAPI.revealCredentials.mockResolvedValue({ + data: { api_key: '312abcdef321' }, + }); mcpAPI.configureThreatBook.mockResolvedValue({ data: { success: true, @@ -365,7 +370,7 @@ describe('MCPServerDetailPanel', () => { data: { has_credential: true, secret_id: 'threatbook_mcp_key', - api_key_masked: 'conf****-key', + api_key_masked: '312xxxx321', }, }); @@ -386,6 +391,111 @@ describe('MCPServerDetailPanel', () => { expect(screen.getByRole('button', { name: 'Edit configuration' })).toBeInTheDocument(); }); + it('shows the saved key masked in edit mode and reveals it only on demand', async () => { + const user = userEvent.setup(); + mcpAPI.get.mockResolvedValue({ + ...detailResponse, + data: { + ...detailResponse.data, + name: 'threatbook_mcp', + config: { + type: 'sse', + url: 'https://mcp.threatbook.cn/mcp?apikey={secret:threatbook_mcp_key}', + }, + }, + }); + mcpAPI.getCredentials.mockResolvedValue({ + data: { + has_credential: true, + secret_id: 'threatbook_mcp_key', + api_key_masked: '312xxxx321', + }, + }); + + render( + , + ); + + await user.click(await screen.findByRole('button', { name: '编辑配置' })); + + const keyInput = screen.getByLabelText('API Key *'); + expect(keyInput).toHaveValue('312xxxx321'); + expect(keyInput).toHaveAttribute('readonly'); + expect(screen.queryByDisplayValue('312abcdef321')).not.toBeInTheDocument(); + + await user.click(screen.getByTitle('显示')); + + await waitFor(() => { + expect(mcpAPI.revealCredentials).toHaveBeenCalledWith('threatbook_mcp'); + expect(keyInput).toHaveValue('312abcdef321'); + }); + expect(keyInput).not.toHaveAttribute('readonly'); + + await user.click(screen.getByTitle('隐藏')); + expect(keyInput).toHaveValue('312xxxx321'); + expect(keyInput).toHaveAttribute('readonly'); + + await user.click(screen.getByRole('button', { name: '国际区' })); + expect(keyInput).toHaveValue(''); + expect(keyInput).not.toHaveAttribute('readonly'); + }); + + it('locks region switching while the saved key is being revealed', async () => { + const user = userEvent.setup(); + let resolveReveal: ((value: { data: { api_key: string } }) => void) | undefined; + mcpAPI.revealCredentials.mockReturnValue(new Promise((resolve) => { + resolveReveal = resolve; + })); + mcpAPI.get.mockResolvedValue({ + ...detailResponse, + data: { + ...detailResponse.data, + name: 'threatbook_mcp', + config: { + type: 'sse', + url: 'https://mcp.threatbook.cn/mcp?apikey={secret:threatbook_mcp_key}', + }, + }, + }); + mcpAPI.getCredentials.mockResolvedValue({ + data: { + has_credential: true, + secret_id: 'threatbook_mcp_key', + api_key_masked: '312xxxx321', + }, + }); + + render( + , + ); + + await user.click(await screen.findByRole('button', { name: '编辑配置' })); + await user.click(screen.getByTitle('显示')); + + expect(screen.getByRole('button', { name: '国际区' })).toBeDisabled(); + + resolveReveal?.({ data: { api_key: '312abcdef321' } }); + await waitFor(() => { + expect(screen.getByRole('button', { name: '国际区' })).toBeEnabled(); + }); + }); + it('does not show the free key link for other MCP servers', async () => { render( (null); @@ -82,15 +84,50 @@ export default function ThreatBookMCPConfigPanel({ const summaryRegion = configuredRegion || region; const isConfigured = Boolean(credentials?.has_credential && configuredRegion); const isConnected = serverStatus === 'connected'; + const hasStoredKeyForRegion = Boolean( + credentials?.has_credential && configuredRegion && configuredRegion === region, + ); + const showingStoredKeyMask = hasStoredKeyForRegion && !showApiKey; + const displayedApiKey = showingStoredKeyMask + ? credentials?.api_key_masked || '' + : apiKey; const handleRegionChange = (nextRegion: ThreatBookRegion) => { if (nextRegion === region) return; setRegion(nextRegion); setApiKey(''); setShowApiKey(false); + setStoredKeyLoaded(false); setResult(null); }; + const handleApiKeyVisibility = async () => { + if (showApiKey) { + setShowApiKey(false); + return; + } + + if (hasStoredKeyForRegion && !storedKeyLoaded) { + try { + setRevealingKey(true); + setResult(null); + const response = await mcpAPI.revealCredentials(serverName); + setApiKey(response.data.api_key); + setStoredKeyLoaded(true); + } catch (error: any) { + setResult({ + success: false, + message: error.response?.data?.detail || error.message || t('detail.threatbookMcp.revealFailed'), + }); + return; + } finally { + setRevealingKey(false); + } + } + + setShowApiKey(true); + }; + const handleSave = async () => { const trimmedKey = apiKey.trim(); if (!trimmedKey) { @@ -112,6 +149,7 @@ export default function ThreatBookMCPConfigPanel({ setApiKey(''); setShowApiKey(false); + setStoredKeyLoaded(false); setResult({ success: true, message: t('detail.threatbookMcp.saveSuccess') }); await onConfigured(); await loadCredentials(); @@ -195,6 +233,8 @@ export default function ThreatBookMCPConfigPanel({ onClick={() => { setRegion(configuredRegion || languageDefaultRegion); setApiKey(''); + setShowApiKey(false); + setStoredKeyLoaded(false); setResult(null); setEditing(true); }} @@ -240,11 +280,12 @@ export default function ThreatBookMCPConfigPanel({ key={item} type="button" onClick={() => handleRegionChange(item)} + disabled={revealingKey || saving} className={`min-w-[96px] rounded-md px-4 py-2 text-sm font-semibold transition-colors ${ region === item ? 'bg-white text-green-700 shadow-sm ring-1 ring-green-200' : 'text-gray-500 hover:text-gray-700' - }`} + } disabled:cursor-not-allowed disabled:opacity-50`} > {t(`detail.threatbookMcp.regions.${item}`)} @@ -264,8 +305,9 @@ export default function ThreatBookMCPConfigPanel({
{ setApiKey(event.target.value); setResult(null); @@ -276,11 +318,16 @@ export default function ThreatBookMCPConfigPanel({ />
{ setRegion(configuredRegion || languageDefaultRegion); setApiKey(''); + setShowApiKey(false); + setStoredKeyLoaded(false); setResult(null); setEditing(false); }} From 840347030a91e50e2a3cf2f5c52cae850960fa7d Mon Sep 17 00:00:00 2001 From: John Yin <10972267+john-yin2333@user.noreply.gitee.com> Date: Mon, 7 Sep 2026 00:33:15 +0800 Subject: [PATCH 41/63] Fully mask saved MCP keys --- flocks/server/routes/mcp.py | 6 ++---- tests/server/routes/test_mcp_routes.py | 2 +- .../Tool/components/ServiceDetailPanel.test.tsx | 13 ++++++++----- .../Tool/components/ThreatBookMCPConfigPanel.tsx | 2 +- 4 files changed, 12 insertions(+), 11 deletions(-) diff --git a/flocks/server/routes/mcp.py b/flocks/server/routes/mcp.py index 55f51d900..3da9d66a6 100644 --- a/flocks/server/routes/mcp.py +++ b/flocks/server/routes/mcp.py @@ -992,10 +992,8 @@ def _get_mcp_credential(name: str) -> tuple[Optional[str], Optional[str]]: def _mask_mcp_credential(api_key: str) -> str: - """Mask an MCP key as three visible chars on each side.""" - if len(api_key) <= 6: - return "xxxx" - return f"{api_key[:3]}xxxx{api_key[-3:]}" + """Mask an MCP key without exposing any part of it.""" + return "************" if api_key else "" @router.get( diff --git a/tests/server/routes/test_mcp_routes.py b/tests/server/routes/test_mcp_routes.py index fdc59b1a1..7cfab9f7e 100644 --- a/tests/server/routes/test_mcp_routes.py +++ b/tests/server/routes/test_mcp_routes.py @@ -1208,7 +1208,7 @@ def get(self, secret_id: str): assert resp.status_code == 200, resp.text assert resp.json()["has_credential"] is True assert resp.json()["secret_id"] == "threatbook_mcp_key" - assert resp.json()["api_key_masked"] == "312xxxx321" + assert resp.json()["api_key_masked"] == "************" assert "312abcdef321" not in resp.text assert requested_ids == ["threatbook_mcp_key"] diff --git a/webui/src/pages/Tool/components/ServiceDetailPanel.test.tsx b/webui/src/pages/Tool/components/ServiceDetailPanel.test.tsx index a9709deaa..6e228c963 100644 --- a/webui/src/pages/Tool/components/ServiceDetailPanel.test.tsx +++ b/webui/src/pages/Tool/components/ServiceDetailPanel.test.tsx @@ -370,7 +370,7 @@ describe('MCPServerDetailPanel', () => { data: { has_credential: true, secret_id: 'threatbook_mcp_key', - api_key_masked: '312xxxx321', + api_key_masked: '************', }, }); @@ -408,7 +408,7 @@ describe('MCPServerDetailPanel', () => { data: { has_credential: true, secret_id: 'threatbook_mcp_key', - api_key_masked: '312xxxx321', + api_key_masked: '************', }, }); @@ -427,7 +427,8 @@ describe('MCPServerDetailPanel', () => { await user.click(await screen.findByRole('button', { name: '编辑配置' })); const keyInput = screen.getByLabelText('API Key *'); - expect(keyInput).toHaveValue('312xxxx321'); + expect(keyInput).toHaveValue('************'); + expect(keyInput).toHaveAttribute('type', 'password'); expect(keyInput).toHaveAttribute('readonly'); expect(screen.queryByDisplayValue('312abcdef321')).not.toBeInTheDocument(); @@ -437,10 +438,12 @@ describe('MCPServerDetailPanel', () => { expect(mcpAPI.revealCredentials).toHaveBeenCalledWith('threatbook_mcp'); expect(keyInput).toHaveValue('312abcdef321'); }); + expect(keyInput).toHaveAttribute('type', 'text'); expect(keyInput).not.toHaveAttribute('readonly'); await user.click(screen.getByTitle('隐藏')); - expect(keyInput).toHaveValue('312xxxx321'); + expect(keyInput).toHaveValue('************'); + expect(keyInput).toHaveAttribute('type', 'password'); expect(keyInput).toHaveAttribute('readonly'); await user.click(screen.getByRole('button', { name: '国际区' })); @@ -469,7 +472,7 @@ describe('MCPServerDetailPanel', () => { data: { has_credential: true, secret_id: 'threatbook_mcp_key', - api_key_masked: '312xxxx321', + api_key_masked: '************', }, }); diff --git a/webui/src/pages/Tool/components/ThreatBookMCPConfigPanel.tsx b/webui/src/pages/Tool/components/ThreatBookMCPConfigPanel.tsx index d77c4da67..594850451 100644 --- a/webui/src/pages/Tool/components/ThreatBookMCPConfigPanel.tsx +++ b/webui/src/pages/Tool/components/ThreatBookMCPConfigPanel.tsx @@ -305,7 +305,7 @@ export default function ThreatBookMCPConfigPanel({
{ From 322cb735c1ee59272d7311cf25bb23dc699f551e Mon Sep 17 00:00:00 2001 From: John Yin <10972267+john-yin2333@user.noreply.gitee.com> Date: Mon, 7 Sep 2026 00:38:08 +0800 Subject: [PATCH 42/63] Add Flocks LLM usage portal link --- webui/src/components/layout/Layout.test.tsx | 9 ++++++++- webui/src/components/layout/Layout.tsx | 17 ++++++++++++++++- webui/src/locales/en-US/nav.json | 1 + webui/src/locales/zh-CN/nav.json | 1 + 4 files changed, 26 insertions(+), 2 deletions(-) diff --git a/webui/src/components/layout/Layout.test.tsx b/webui/src/components/layout/Layout.test.tsx index 403339b7d..dd139bb72 100644 --- a/webui/src/components/layout/Layout.test.tsx +++ b/webui/src/components/layout/Layout.test.tsx @@ -731,7 +731,14 @@ describe('Layout onboarding entry', () => { await user.click(screen.getByRole('button', { name: 'admin settings' })); expect(screen.getByRole('link', { name: 'Flocks Pro' })).toHaveAttribute('href', '/settings/flockspro'); - expect(screen.getByRole('link', { name: 'settings' })).toHaveAttribute('href', '/settings/preferences'); + const updateEntry = screen.getByRole('button', { name: 'checkUpdate' }); + const usageEntry = screen.getByRole('link', { name: 'flocksLlmUsageQuota' }); + const settingsEntry = screen.getByRole('link', { name: 'settings' }); + expect(usageEntry).toHaveAttribute('href', 'https://portal.agentflocks.com'); + expect(usageEntry).toHaveAttribute('target', '_blank'); + expect(settingsEntry).toHaveAttribute('href', '/settings/preferences'); + expect(updateEntry.compareDocumentPosition(usageEntry) & Node.DOCUMENT_POSITION_FOLLOWING).toBeTruthy(); + expect(usageEntry.compareDocumentPosition(settingsEntry) & Node.DOCUMENT_POSITION_FOLLOWING).toBeTruthy(); await user.click(screen.getByRole('button', { name: 'logout' })); expect(logout).toHaveBeenCalledTimes(1); diff --git a/webui/src/components/layout/Layout.tsx b/webui/src/components/layout/Layout.tsx index 0febfb5cd..0a50be13d 100644 --- a/webui/src/components/layout/Layout.tsx +++ b/webui/src/components/layout/Layout.tsx @@ -22,6 +22,7 @@ import { Settings, ArrowUpCircle, RefreshCw, + Gauge, Loader2, type LucideIcon, } from 'lucide-react'; @@ -170,6 +171,7 @@ import { recoverLazyLoad } from '@/utils/chunkLoadRecovery'; const UPDATE_CHECK_INTERVAL_MS = 3_600_000; const UPDATE_CHECK_MIN_GAP_MS = 600_000; const UPDATE_CHECK_INITIAL_DELAY_MS = 250; +const FLOCKS_LLM_USAGE_URL = 'https://portal.agentflocks.com'; interface LayoutNavItem { name: string; @@ -1142,7 +1144,7 @@ export default function Layout() { > {accountMenuOpen && (
{showFlocksproUpgradeEntry && ( {t('checkUpdate')} + { + setAccountMenuOpen(false); + setSidebarOpen(false); + }} + 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" + > + + {t('flocksLlmUsageQuota')} + Date: Mon, 7 Sep 2026 00:40:49 +0800 Subject: [PATCH 43/63] Shorten Flocks LLM usage label --- webui/src/locales/en-US/nav.json | 2 +- webui/src/locales/zh-CN/nav.json | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/webui/src/locales/en-US/nav.json b/webui/src/locales/en-US/nav.json index 89fb5004a..efa5eb51b 100644 --- a/webui/src/locales/en-US/nav.json +++ b/webui/src/locales/en-US/nav.json @@ -22,7 +22,7 @@ "systemLog": "System Logs", "flocksproUpgrade": "Upgrade", "checkUpdate": "Check for updates", - "flocksLlmUsageQuota": "Flocks LLM Usage & Quota", + "flocksLlmUsageQuota": "Flocks LLM Usage", "auditLogs": "Audit Logs", "settings": "Settings", "settingsBack": "Back", diff --git a/webui/src/locales/zh-CN/nav.json b/webui/src/locales/zh-CN/nav.json index 50643cd03..e6a955690 100644 --- a/webui/src/locales/zh-CN/nav.json +++ b/webui/src/locales/zh-CN/nav.json @@ -22,7 +22,7 @@ "systemLog": "系统日志", "flocksproUpgrade": "升级", "checkUpdate": "检查更新", - "flocksLlmUsageQuota": "Flocks llm 用量&额度", + "flocksLlmUsageQuota": "Flocks llm用量", "auditLogs": "审计日志", "settings": "设置", "settingsBack": "返回", From cd949fb067539f1a4c850ba55ab3f1bf893c9116 Mon Sep 17 00:00:00 2001 From: John Yin <10972267+john-yin2333@user.noreply.gitee.com> Date: Mon, 7 Sep 2026 00:54:41 +0800 Subject: [PATCH 44/63] Reveal saved onboarding keys on demand --- flocks/server/routes/provider.py | 81 +++++++ .../routes/test_global_mutation_auth.py | 24 +- webui/src/api/provider.test.ts | 9 + webui/src/api/provider.ts | 3 + .../common/OnboardingModal.test.tsx | 74 +++++- .../src/components/common/OnboardingModal.tsx | 223 ++++++++++++++++-- webui/src/constants/threatbook.ts | 3 + webui/src/locales/en-US/common.json | 4 + webui/src/locales/zh-CN/common.json | 4 + 9 files changed, 397 insertions(+), 28 deletions(-) diff --git a/flocks/server/routes/provider.py b/flocks/server/routes/provider.py index 4f28570cf..b3dd043da 100644 --- a/flocks/server/routes/provider.py +++ b/flocks/server/routes/provider.py @@ -2146,6 +2146,87 @@ async def get_service_credentials( raise HTTPException(status_code=500, detail=str(e)) +def _load_api_service_primary_key(provider_id: str) -> tuple[Optional[str], Optional[str]]: + """Load the primary API key for an API service without logging its value.""" + from flocks.security import get_secret_manager + + secrets = get_secret_manager() + raw_service = ConfigWriter.get_api_service_raw(provider_id) + metadata = _load_api_service_metadata_data(provider_id) or {} + + for candidate in _get_api_service_secret_candidates( + provider_id, + raw_service, + field_name="api_key", + ): + value = secrets.get(candidate) + if not value: + continue + + auth = metadata.get("authentication") or metadata.get("auth") + expects_secondary_secret = ( + isinstance(auth, dict) + and bool(auth.get("secret_secret")) + and _should_persist_secondary_secret(metadata) + ) + if expects_secondary_secret: + split_result = _split_compound_service_credentials(value) + if split_result: + value = split_result[0] + return candidate, value + + return None, None + + +@router.post( + "/{provider_id}/service-credentials/reveal", + response_model=ProviderCredentialResponse, + summary="Reveal API service credentials", + description="Reveal the primary API key for an API service to an administrator.", +) +async def reveal_service_credentials( + provider_id: str, + request: Request, + response: Response, + admin: AuthUser = Depends(require_admin), +): + response.headers["Cache-Control"] = "no-store" + response.headers["Pragma"] = "no-cache" + + try: + secret_id, api_key = _load_api_service_primary_key(provider_id) + credentials = ProviderCredentialResponse( + secret_id=secret_id, + api_key=api_key, + api_key_masked=SecretManager.mask(api_key) if api_key else None, + has_credential=bool(api_key), + ) + try: + await emit_audit_event( + "provider.service_credentials_reveal", + { + "action": "service_credentials_reveal", + "actor_id": admin.id, + "actor_name": admin.username, + "user_id": admin.id, + "username": admin.username, + "provider_id": provider_id, + "secret_id": secret_id, + "ip": get_request_ip(request), + "user_agent": get_request_user_agent(request), + }, + ) + except Exception as audit_error: + log.warn( + "service.credentials.reveal.audit_failed", + {"provider_id": provider_id, "error": str(audit_error)}, + ) + return credentials + except Exception as e: + log.error("service.credentials.reveal.error", {"provider_id": provider_id, "error": str(e)}) + raise HTTPException(status_code=500, detail=str(e)) + + @router.post( "/{provider_id}/service-credentials", response_model=Dict[str, Any], diff --git a/tests/server/routes/test_global_mutation_auth.py b/tests/server/routes/test_global_mutation_auth.py index 40cdb9f67..4f526516a 100644 --- a/tests/server/routes/test_global_mutation_auth.py +++ b/tests/server/routes/test_global_mutation_auth.py @@ -44,6 +44,7 @@ def test_provider_global_configuration_credentials_and_tests_require_admin(): client.post("/api/provider/example/credentials", json={"api_key": "secret"}), client.delete("/api/provider/example/credentials"), client.get("/api/provider/example/service-credentials"), + client.post("/api/provider/example/service-credentials/reveal"), client.post("/api/provider/example/service-credentials", json={"api_key": "secret"}), client.post("/api/provider/example/test-credentials", json={}), client.patch("/api/provider/api-services/example", json={"enabled": False}), @@ -107,13 +108,16 @@ def test_admin_reveal_returns_raw_llm_key_with_audit_and_no_store( 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") + service_reveal_response = client.post("/api/provider/service-example/service-credentials/reveal") assert llm_response.status_code == 200 assert reveal_response.status_code == 200 assert service_response.status_code == 200 + assert service_reveal_response.status_code == 200 llm_payload = llm_response.json() reveal_payload = reveal_response.json() service_payload = service_response.json() + service_reveal_payload = service_reveal_response.json() assert llm_payload["api_key"] is None assert llm_payload["api_key_masked"] assert reveal_payload["api_key"] == raw_llm_key @@ -123,15 +127,23 @@ def test_admin_reveal_returns_raw_llm_key_with_audit_and_no_store( 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 service_reveal_payload["api_key"] == raw_service_key + assert service_reveal_response.headers["Cache-Control"] == "no-store" + assert service_reveal_response.headers["Pragma"] == "no-cache" 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) + assert raw_service_key in service_reveal_response.text + assert audit.await_count == 2 + audit_events = [call.args for call in audit.await_args_list] + assert [event_type for event_type, _ in audit_events] == [ + "provider.credentials_reveal", + "provider.service_credentials_reveal", + ] + for _, audit_payload in audit_events: + assert audit_payload["username"] == "admin-test" + assert raw_llm_key not in repr(audit_payload) + assert raw_service_key not in repr(audit_payload) def test_admin_provider_update_never_writes_raw_api_key_to_storage(monkeypatch: pytest.MonkeyPatch): diff --git a/webui/src/api/provider.test.ts b/webui/src/api/provider.test.ts index 2a63e4f57..afbc51ca5 100644 --- a/webui/src/api/provider.test.ts +++ b/webui/src/api/provider.test.ts @@ -151,4 +151,13 @@ describe('providerAPI.revealCredentials', () => { expect(mockPost).toHaveBeenCalledWith('/api/provider/openai/credentials/reveal'); }); + + it('uses the explicit API service credential reveal endpoint', async () => { + mockPost.mockResolvedValue({ data: { api_key: 'service-key', has_credential: true } }); + + const { providerAPI } = await import('./provider'); + await providerAPI.revealServiceCredentials('threatbook-cn'); + + expect(mockPost).toHaveBeenCalledWith('/api/provider/threatbook-cn/service-credentials/reveal'); + }); }); diff --git a/webui/src/api/provider.ts b/webui/src/api/provider.ts index f127e70cc..cde35cd9d 100644 --- a/webui/src/api/provider.ts +++ b/webui/src/api/provider.ts @@ -122,6 +122,9 @@ export const providerAPI = { getServiceCredentials: (id: string) => client.get(`/api/provider/${id}/service-credentials`), + revealServiceCredentials: (id: string) => + client.post(`/api/provider/${id}/service-credentials/reveal`), + setServiceCredentials: (id: string, credentials: ProviderCredentialInput) => client.post<{ success: boolean; message: string }>(`/api/provider/${id}/service-credentials`, credentials) .then((response) => { diff --git a/webui/src/components/common/OnboardingModal.test.tsx b/webui/src/components/common/OnboardingModal.test.tsx index d99039991..c2d32e65b 100644 --- a/webui/src/components/common/OnboardingModal.test.tsx +++ b/webui/src/components/common/OnboardingModal.test.tsx @@ -11,7 +11,9 @@ const { clientPost, currentLanguage, defaultModelAPI, + mcpAPI, onboardingAPI, + providerAPI, sessionApi, } = vi.hoisted(() => ({ catalogAPI: { @@ -24,11 +26,18 @@ const { defaultModelAPI: { getResolved: vi.fn(), }, + mcpAPI: { + revealCredentials: vi.fn(), + }, onboardingAPI: { getStatus: vi.fn(), validate: vi.fn(), apply: vi.fn(), }, + providerAPI: { + revealCredentials: vi.fn(), + revealServiceCredentials: vi.fn(), + }, sessionApi: { create: vi.fn(), }, @@ -37,6 +46,11 @@ const { vi.mock('@/api/provider', () => ({ catalogAPI, defaultModelAPI, + providerAPI, +})); + +vi.mock('@/api/mcp', () => ({ + mcpAPI, })); vi.mock('@/api/onboarding', () => ({ @@ -148,6 +162,15 @@ describe('OnboardingModal', () => { vi.clearAllMocks(); currentLanguage.value = 'zh-CN'; defaultModelAPI.getResolved.mockRejectedValue(new Error('no default model')); + providerAPI.revealCredentials.mockResolvedValue({ + data: { has_credential: true, api_key: 'saved-model-key' }, + }); + providerAPI.revealServiceCredentials.mockResolvedValue({ + data: { has_credential: true, api_key: 'saved-intel-key' }, + }); + mcpAPI.revealCredentials.mockResolvedValue({ + data: { api_key: 'saved-intel-mcp-key' }, + }); onboardingAPI.getStatus.mockResolvedValue({ data: makeStatus(), }); @@ -369,7 +392,7 @@ describe('OnboardingModal', () => { expect(screen.getByText('onboarding.bootstrap.intelMcpCapability')).toBeInTheDocument(); }); - it('shows configured summaries and lets users enter edit mode', async () => { + it('shows configured keys masked in edit mode and reveals them on demand', async () => { const user = userEvent.setup(); onboardingAPI.getStatus.mockResolvedValue({ @@ -404,11 +427,60 @@ describe('OnboardingModal', () => { await user.click(screen.getByRole('button', { name: 'onboarding.bootstrap.editPrimary' })); expect(screen.getByRole('button', { name: 'onboarding.bootstrap.savePrimary' })).toBeInTheDocument(); + const modelKeyInput = screen.getByPlaceholderText('onboarding.bootstrap.modelKeyPlaceholder'); + expect(modelKeyInput).toHaveValue('************'); + expect(modelKeyInput).toHaveAttribute('type', 'password'); + expect(modelKeyInput).toHaveAttribute('readonly'); + expect(providerAPI.revealCredentials).not.toHaveBeenCalled(); + + await user.click(screen.getByTitle('onboarding.bootstrap.showKey')); + await waitFor(() => { + expect(providerAPI.revealCredentials).toHaveBeenCalledWith('threatbook-cn-llm'); + expect(modelKeyInput).toHaveValue('saved-model-key'); + }); + expect(modelKeyInput).toHaveAttribute('type', 'text'); + + await user.click(screen.getByTitle('onboarding.bootstrap.hideKey')); + expect(modelKeyInput).toHaveValue('************'); + expect(modelKeyInput).toHaveAttribute('type', 'password'); + + const providerSelect = screen.getByRole('combobox'); + await user.selectOptions(providerSelect, 'openai-compatible'); + expect(modelKeyInput).toHaveValue(''); + expect(modelKeyInput).not.toHaveAttribute('readonly'); + await user.selectOptions(providerSelect, 'threatbook-free'); + expect(modelKeyInput).toHaveValue('************'); + expect(modelKeyInput).toHaveAttribute('readonly'); await user.click(screen.getByRole('button', { name: 'onboarding.bootstrap.nextStep' })); expect(screen.getByText('onboarding.bootstrap.intelConfiguredHint')).toBeInTheDocument(); expect(screen.getAllByText('onboarding.bootstrap.intelConfiguredVerified').length).toBeGreaterThan(0); expect(screen.getByRole('button', { name: 'onboarding.bootstrap.editIntel' })).toBeInTheDocument(); + + await user.click(screen.getByRole('button', { name: 'onboarding.bootstrap.editIntel' })); + const intelKeyInput = screen.getByPlaceholderText('onboarding.bootstrap.intelKeyPlaceholder'); + expect(intelKeyInput).toHaveValue('************'); + expect(intelKeyInput).toHaveAttribute('type', 'password'); + expect(intelKeyInput).toHaveAttribute('readonly'); + expect(providerAPI.revealServiceCredentials).not.toHaveBeenCalled(); + + await user.click(screen.getByTitle('onboarding.bootstrap.showKey')); + await waitFor(() => { + expect(providerAPI.revealServiceCredentials).toHaveBeenCalledWith('threatbook-cn'); + expect(intelKeyInput).toHaveValue('saved-intel-key'); + }); + expect(intelKeyInput).toHaveAttribute('type', 'text'); + + await user.click(screen.getByTitle('onboarding.bootstrap.hideKey')); + expect(intelKeyInput).toHaveValue('************'); + expect(intelKeyInput).toHaveAttribute('type', 'password'); + + await user.click(screen.getByRole('button', { name: 'onboarding.bootstrap.intelRegionGlobal' })); + expect(intelKeyInput).toHaveValue(''); + expect(intelKeyInput).not.toHaveAttribute('readonly'); + await user.click(screen.getByRole('button', { name: 'onboarding.bootstrap.intelRegionChina' })); + expect(intelKeyInput).toHaveValue('************'); + expect(intelKeyInput).toHaveAttribute('readonly'); }); it('starts Rex onboarding after skipping intelligence setup', async () => { diff --git a/webui/src/components/common/OnboardingModal.tsx b/webui/src/components/common/OnboardingModal.tsx index f58f84eaf..2d504c972 100644 --- a/webui/src/components/common/OnboardingModal.tsx +++ b/webui/src/components/common/OnboardingModal.tsx @@ -1,10 +1,21 @@ import { useCallback, useEffect, useMemo, useState } from 'react'; import { useNavigate } from 'react-router-dom'; import { useTranslation } from 'react-i18next'; -import { AlertTriangle, ArrowRight, CheckCircle2, ExternalLink, X, XCircle } from 'lucide-react'; +import { + AlertTriangle, + ArrowRight, + CheckCircle2, + ExternalLink, + Eye, + EyeOff, + Loader2, + X, + XCircle, +} from 'lucide-react'; import { sessionApi } from '@/api/session'; import client from '@/api/client'; -import { catalogAPI, defaultModelAPI } from '@/api/provider'; +import { catalogAPI, defaultModelAPI, providerAPI } from '@/api/provider'; +import { mcpAPI } from '@/api/mcp'; import { onboardingAPI, type OnboardingApplyResponse, @@ -20,6 +31,7 @@ const MODEL_KEY_LINK = 'https://portal.agentflocks.com'; const THREATBOOK_PROVIDER_IDS = ['threatbook-cn-llm', 'threatbook-io-llm'] as const; const THREATBOOK_FREE_PROVIDER_OPTION = 'threatbook-free'; +const STORED_KEY_MASK = '************'; type SectionTone = 'success' | 'warning' | 'error'; type OnboardingStep = 'model' | 'intel'; @@ -226,11 +238,13 @@ function RegionChooser({ onChange, chinaLabel, globalLabel, + disabled = false, }: { value: OnboardingRegion; onChange: (value: OnboardingRegion) => void; chinaLabel: string; globalLabel: string; + disabled?: boolean; }) { return (
@@ -242,11 +256,12 @@ function RegionChooser({ key={candidate} type="button" onClick={() => onChange(candidate)} + disabled={disabled} className={`rounded-lg px-5 py-2 text-sm font-semibold transition-colors ${ value === candidate ? 'bg-green-50 text-green-700 ring-1 ring-green-200' : 'text-gray-500 hover:bg-gray-50 hover:text-gray-700' - }`} + } disabled:cursor-not-allowed disabled:opacity-50`} > {label} @@ -255,6 +270,57 @@ function RegionChooser({ ); } +function OnboardingSecretInput({ + value, + hasStoredKey, + visible, + revealing, + placeholder, + showLabel, + hideLabel, + onChange, + onToggleVisibility, +}: { + value: string; + hasStoredKey: boolean; + visible: boolean; + revealing: boolean; + placeholder: string; + showLabel: string; + hideLabel: string; + onChange: (value: string) => void; + onToggleVisibility: () => void; +}) { + const showingStoredMask = hasStoredKey && !visible; + + return ( +
+ onChange(event.target.value)} + placeholder={placeholder} + autoComplete="off" + className="w-full rounded-lg border border-gray-200 bg-white px-3 py-2 pr-10 text-xs transition-all placeholder-gray-300 focus:border-red-400 focus:outline-none focus:ring-2 focus:ring-red-400/50" + /> + +
+ ); +} + function SkipConfirmDialog({ t, target, @@ -327,6 +393,9 @@ export default function OnboardingModal({ onClose }: OnboardingModalProps) { const [primaryConfigured, setPrimaryConfigured] = useState(false); const [primaryStatus, setPrimaryStatus] = useState(null); const [primaryEditing, setPrimaryEditing] = useState(false); + const [primaryKeyVisible, setPrimaryKeyVisible] = useState(false); + const [primaryKeyRevealed, setPrimaryKeyRevealed] = useState(false); + const [primaryKeyRevealing, setPrimaryKeyRevealing] = useState(false); const [modelRegion, setModelRegion] = useState('cn'); const [modelSkipped, setModelSkipped] = useState(false); @@ -336,6 +405,9 @@ export default function OnboardingModal({ onClose }: OnboardingModalProps) { const [intelConfigured, setIntelConfigured] = useState(false); const [intelStatus, setIntelStatus] = useState(null); const [intelEditing, setIntelEditing] = useState(false); + const [intelKeyVisible, setIntelKeyVisible] = useState(false); + const [intelKeyRevealed, setIntelKeyRevealed] = useState(false); + const [intelKeyRevealing, setIntelKeyRevealing] = useState(false); const [intelSkipped, setIntelSkipped] = useState(false); const [intelRuntimeStatus, setIntelRuntimeStatus] = useState(null); @@ -562,6 +634,18 @@ export default function OnboardingModal({ onClose }: OnboardingModalProps) { const canSaveIntel = Boolean(intelApiKey.trim()); const showPrimaryConfiguredDetails = primaryConfigured && !primaryEditing; const showIntelConfiguredDetails = intelConfigured && !intelEditing; + const hasStoredPrimaryKey = Boolean( + primaryConfigured + && primaryEditing + && resolvedDefaultModel + && primaryProviderId === resolvedDefaultModel.providerId, + ); + const hasStoredIntelKey = Boolean( + intelConfigured + && intelEditing + && intelRuntimeStatus?.region === intelRegion + && (intelRuntimeStatus.api_configured || intelRuntimeStatus.mcp_configured), + ); const buildPrimaryPayload = (): OnboardingRequest => { if (primaryProviderIsThreatBook) { @@ -621,13 +705,79 @@ export default function OnboardingModal({ onClose }: OnboardingModalProps) { setPrimaryProviderId(providerId); } setPrimaryApiKey(''); + setPrimaryKeyVisible(false); + setPrimaryKeyRevealed(false); setPrimaryBaseUrl(''); setPrimaryStatus(null); - setPrimaryConfigured(false); setPrimaryEditing(true); setModelSkipped(false); }; + const handlePrimaryKeyVisibility = async () => { + if (primaryKeyVisible) { + setPrimaryKeyVisible(false); + return; + } + + if (hasStoredPrimaryKey && !primaryKeyRevealed) { + try { + setPrimaryKeyRevealing(true); + setPrimaryStatus(null); + const response = await providerAPI.revealCredentials(primaryProviderId); + if (!response.data.api_key) throw new Error(t('onboarding.bootstrap.revealModelKeyError')); + setPrimaryApiKey(response.data.api_key); + setPrimaryKeyRevealed(true); + } catch (err: any) { + setPrimaryStatus({ + tone: 'error', + message: err?.response?.data?.detail || err?.message || t('onboarding.bootstrap.revealModelKeyError'), + }); + return; + } finally { + setPrimaryKeyRevealing(false); + } + } + + setPrimaryKeyVisible(true); + }; + + const handleIntelKeyVisibility = async () => { + if (intelKeyVisible) { + setIntelKeyVisible(false); + return; + } + + if (hasStoredIntelKey && !intelKeyRevealed) { + try { + setIntelKeyRevealing(true); + setIntelStatus(null); + let apiKey: string | null | undefined; + if (intelRuntimeStatus?.api_configured) { + const serviceId = intelRuntimeStatus.api_service_id + || THREATBOOK_REGION_CONFIG[intelRegion].apiServiceId; + const response = await providerAPI.revealServiceCredentials(serviceId); + apiKey = response.data.api_key; + } else if (intelRuntimeStatus?.mcp_configured && intelRuntimeStatus.mcp_name) { + const response = await mcpAPI.revealCredentials(intelRuntimeStatus.mcp_name); + apiKey = response.data.api_key; + } + if (!apiKey) throw new Error(t('onboarding.bootstrap.revealIntelKeyError')); + setIntelApiKey(apiKey); + setIntelKeyRevealed(true); + } catch (err: any) { + setIntelStatus({ + tone: 'error', + message: err?.response?.data?.detail || err?.message || t('onboarding.bootstrap.revealIntelKeyError'), + }); + return; + } finally { + setIntelKeyRevealing(false); + } + } + + setIntelKeyVisible(true); + }; + const handleSavePrimary = async () => { if (!canSavePrimary) return; @@ -676,6 +826,9 @@ export default function OnboardingModal({ onClose }: OnboardingModalProps) { setModelRegion(regionForProvider(savedProviderId)); setPrimaryConfigured(true); setPrimaryEditing(false); + setPrimaryApiKey(''); + setPrimaryKeyVisible(false); + setPrimaryKeyRevealed(false); setModelSkipped(false); setHasLLM(true); @@ -729,6 +882,9 @@ export default function OnboardingModal({ onClose }: OnboardingModalProps) { setIntelRegion(payload.region); setIntelConfigured(true); setIntelEditing(false); + setIntelApiKey(''); + setIntelKeyVisible(false); + setIntelKeyRevealed(false); setIntelSkipped(false); setIntelStatus(buildSuccessStatus( validateData, @@ -806,6 +962,9 @@ export default function OnboardingModal({ onClose }: OnboardingModalProps) {
@@ -1064,24 +1237,32 @@ export default function OnboardingModal({ onClose }: OnboardingModalProps) { )} - { - setIntelApiKey(event.target.value); + hasStoredKey={hasStoredIntelKey} + visible={intelKeyVisible} + revealing={intelKeyRevealing} + onChange={(value) => { + setIntelApiKey(value); setIntelStatus(null); }} placeholder={t('onboarding.bootstrap.intelKeyPlaceholder')} - className="min-w-0 flex-1 rounded-lg border border-gray-200 bg-white px-3 py-2 text-xs transition-all placeholder-gray-300 focus:border-red-400 focus:outline-none focus:ring-2 focus:ring-red-400/50" + showLabel={t('onboarding.bootstrap.showKey')} + hideLabel={t('onboarding.bootstrap.hideKey')} + onToggleVisibility={handleIntelKeyVisibility} /> setStep('intel'), } : { label: starting ? t('onboarding.startingButton') : t('onboarding.startButton'), - disabled: starting || primarySaving || intelSaving, + disabled: starting || primarySaving || intelSaving || intelKeyRevealing, onClick: handleStart, }; @@ -1209,7 +1390,7 @@ export default function OnboardingModal({ onClose }: OnboardingModalProps) {
`threatbook_${capability}`)} compactResourceList minimal /> @@ -1289,7 +1290,7 @@ export default function OnboardingModal({ onClose }: OnboardingModalProps) { `threatbook_${capability}`)} compactResourceList onSwitchRegion={ intelStatus?.validation?.error_code === 'region_mismatch' diff --git a/webui/src/components/layout/Layout.test.tsx b/webui/src/components/layout/Layout.test.tsx index dd139bb72..3b837d04b 100644 --- a/webui/src/components/layout/Layout.test.tsx +++ b/webui/src/components/layout/Layout.test.tsx @@ -226,7 +226,7 @@ function makeOnboardingStatus(overrides: Record = {}) { mcp_name: 'threatbook_mcp', service_matrix: { cn: ['api', 'mcp'], - global: ['api', 'mcp'], + global: ['api'], }, }, ...overrides, diff --git a/webui/src/constants/threatbook.ts b/webui/src/constants/threatbook.ts index 68cb5ce91..ec0693006 100644 --- a/webui/src/constants/threatbook.ts +++ b/webui/src/constants/threatbook.ts @@ -3,27 +3,31 @@ export type ThreatBookRegion = 'cn' | 'global'; export const THREATBOOK_REGION_CONFIG: Record = { cn: { activationUrl: 'https://x.threatbook.com/flocks/activate', apiServiceId: 'threatbook-cn', + apiEndpoint: 'https://api.threatbook.cn', mcpEndpoint: 'https://mcp.threatbook.cn/mcp', }, global: { activationUrl: 'https://i.threatbook.io/flocks/activate', apiServiceId: 'threatbook-io', - mcpEndpoint: 'https://mcp.threatbook.io/mcp', + apiEndpoint: 'https://api.threatbook.io', }, }; +export const LEGACY_THREATBOOK_GLOBAL_MCP_ENDPOINT = 'https://mcp.threatbook.io/mcp'; + export function getDefaultThreatBookRegion(language: string | undefined): ThreatBookRegion { return language?.toLowerCase().startsWith('en') ? 'global' : 'cn'; } export function inferThreatBookRegionFromMcpUrl(url: string | undefined): ThreatBookRegion | null { const normalizedUrl = (url || '').trim().toLowerCase(); - if (normalizedUrl.startsWith(THREATBOOK_REGION_CONFIG.global.mcpEndpoint)) return 'global'; - if (normalizedUrl.startsWith(THREATBOOK_REGION_CONFIG.cn.mcpEndpoint)) return 'cn'; + if (normalizedUrl.startsWith(LEGACY_THREATBOOK_GLOBAL_MCP_ENDPOINT)) return 'global'; + if (normalizedUrl.startsWith(THREATBOOK_REGION_CONFIG.cn.mcpEndpoint || '')) return 'cn'; return null; } diff --git a/webui/src/locales/en-US/common.json b/webui/src/locales/en-US/common.json index 83076c91e..75ac4ca0a 100644 --- a/webui/src/locales/en-US/common.json +++ b/webui/src/locales/en-US/common.json @@ -195,7 +195,7 @@ "editIntel": "Edit configuration", "backToConfiguredDetails": "Back to current configuration", "intelPageTitle": "Configure ThreatBook intelligence services", - "intelPageDescription": "Enable free ThreatBook intelligence API and MCP capabilities for IOC lookup, threat intelligence analysis, and security operations tools in Rex.", + "intelPageDescription": "Enable the free ThreatBook intelligence API; the China region also includes MCP capabilities for IOC lookup, threat analysis, and security operations in Rex.", "intelRegionTitle": "Choose service region", "intelRegionHint": "Choose the region matching your API Key. Flocks will configure the available ThreatBook intelligence services automatically.", "intelRegionChina": "China region", @@ -212,7 +212,7 @@ "intelMcpLabel": "ThreatBook MCP", "saveIntel": "Save & Verify Intelligence", "intelChinaSuccess": "ThreatBook intelligence verified successfully. API and MCP have been configured.", - "intelGlobalSuccess": "ThreatBook intelligence verified successfully. International region capabilities have been configured.", + "intelGlobalSuccess": "ThreatBook intelligence verified successfully. The International API has been configured.", "intelConfiguredVerified": "{{region}} configured and verified", "intelConfiguredHint": "ThreatBook intelligence configuration was detected. Edit it to replace the key or re-verify.", "testing": "Verifying...", @@ -243,7 +243,7 @@ "skipModelTitle": "Skip model setup?", "skipModelDescription": "You can skip this step for now and continue configuring the default model later from the Models item in the left navigation.", "skipIntelTitle": "Skip ThreatBook intelligence?", - "skipIntelDescription": "You can skip this step for now. Later, open Tools from the left navigation, switch to the MCP tab, find ThreatBook MCP, and continue configuring ThreatBook intelligence API and MCP.", + "skipIntelDescription": "You can skip this step. Later, open Tools from the left navigation and configure threatbook-cn or threatbook-io on the API tab. China MCP setup remains available from ThreatBook MCP on the MCP tab.", "returnToConfig": "Return to setup", "confirmSkip": "Skip anyway", "summaryTitle": "Setup summary", diff --git a/webui/src/locales/en-US/tool.json b/webui/src/locales/en-US/tool.json index d7e5c6a1c..7cb9c48a3 100644 --- a/webui/src/locales/en-US/tool.json +++ b/webui/src/locales/en-US/tool.json @@ -387,37 +387,71 @@ "enableServer": "Enable Server", "hide": "Hide", "show": "Show", + "threatbookApi": { + "title": "Configure {{service}}", + "descriptions": { + "cn": "Enable the free ThreatBook China API service for file sample analysis and URL scanning in Rex.", + "global": "Enable the free ThreatBook International API service for IP, domain, URL, and file-hash intelligence in Rex." + }, + "region": "Service region", + "regionHint": "This tool is bound to a fixed region. Flocks automatically matches the key claim page and API endpoint.", + "regions": { "cn": "China", "global": "International" }, + "freeService": "Free ThreatBook API service", + "apiKey": "API Key", + "keyPlaceholder": "Paste the {{region}} ThreatBook API key", + "keyHint": "Paste the claimed key here. Flocks stores it securely and verifies the API service for this region.", + "keyRequired": "Enter an API key first.", + "revealFailed": "Failed to retrieve the saved API key. Try again.", + "keyConfigured": "Securely configured", + "claimFreeKey": "Claim free {{region}} key", + "endpoint": "API endpoint", + "endpointHint": "The endpoint is matched to this tool automatically and cannot be changed here.", + "saveAndVerify": "Save and verify", + "saving": "Verifying and saving...", + "saveSuccess": "{{region}} is configured and verified.", + "saveFailed": "Setup failed. Check the API key.", + "configuredTitle": "{{region}} configured", + "connected": "The configuration is verified and the API service is available.", + "saved": "The configuration is saved. You can test connectivity again.", + "edit": "Edit configuration", + "retest": "Test again", + "testing": "Testing...", + "testFailed": "Service test failed." + }, "threatbookMcp": { "title": "Configure ThreatBook MCP", - "description": "Enable the free ThreatBook intelligence MCP service for IOC lookup, threat analysis, and security operations in Rex.", + "description": "Enable the free ThreatBook China MCP service for IOC lookup, threat analysis, and security operations in Rex.", "region": "Service region", - "regionHint": "Choose the region for your API key. Flocks will match the claim page and MCP endpoint automatically.", + "regionHint": "ThreatBook MCP is currently available in China only. Flocks matches the claim page and MCP endpoint automatically.", "regions": { "cn": "China", "global": "International" }, "freeService": "Free ThreatBook MCP service", "apiKey": "API Key", - "keyPlaceholder": "Paste the ThreatBook API key for this region", + "keyPlaceholder": "Paste a ThreatBook China API key", "keyHint": "After claiming a key, paste it here. Flocks will build the service URL and store the key securely.", "keyRequired": "Enter the API key for the selected region.", "revealFailed": "Failed to retrieve the saved API key. Try again.", "keyConfigured": "Securely configured", - "claimFreeKey": "Claim free API key", + "claimFreeKey": "Claim free China key", "endpoint": "MCP endpoint", "endpointHint": "The endpoint is selected automatically. Your API key is stored separately and never needs to be appended manually.", "saveAndVerify": "Save and verify connection", "saving": "Verifying and saving...", "saveSuccess": "ThreatBook MCP is configured and verified.", "saveFailed": "ThreatBook MCP setup failed. Check the region and API key.", - "configuredTitle": "{{region}} configured", + "configuredTitle": "China configured", "connected": "The configuration is verified and the MCP service is connected.", "savedNotConnected": "The configuration is saved, but the MCP service is not connected.", "edit": "Edit configuration", "retest": "Test again", "testing": "Testing...", "testFailed": "Connection test failed.", - "loading": "Loading ThreatBook MCP configuration..." + "loading": "Loading ThreatBook MCP configuration...", + "legacyGlobalTitle": "Legacy International MCP configuration detected", + "legacyGlobalDescription": "ThreatBook MCP now supports China only. The old key will not be reused automatically; claim and enter a China key to migrate.", + "migrateToChina": "Switch to China" } }, diff --git a/webui/src/locales/zh-CN/common.json b/webui/src/locales/zh-CN/common.json index 08bcb5ce8..61aa2f33d 100644 --- a/webui/src/locales/zh-CN/common.json +++ b/webui/src/locales/zh-CN/common.json @@ -195,7 +195,7 @@ "editIntel": "编辑配置", "backToConfiguredDetails": "返回当前配置", "intelPageTitle": "配置微步情报服务", - "intelPageDescription": "免费启用微步情报API与MCP能力，为Rex提供IOC查询、威胁情报分析和安全运营工具能力。", + "intelPageDescription": "免费启用微步情报 API 能力；中国区同时提供 MCP 能力，为 Rex 提供 IOC 查询、威胁情报分析和安全运营工具能力。", "intelRegionTitle": "选择服务区域", "intelRegionHint": "选择 API Key 对应区域，系统会自动配置可用的微步情报服务。", "intelRegionChina": "中国区", @@ -212,7 +212,7 @@ "intelMcpLabel": "ThreatBook MCP", "saveIntel": "保存并验证微步情报", "intelChinaSuccess": "微步情报服务验证成功，已配置 ThreatBook API 和 MCP。", - "intelGlobalSuccess": "微步情报服务验证成功，已配置国际区能力。", + "intelGlobalSuccess": "微步情报服务验证成功，已配置国际区 API 能力。", "intelConfiguredVerified": "已配置{{region}}并验证通过", "intelConfiguredHint": "当前已检测到微步情报配置；如需更换 Key 或重新验证，可进入编辑配置。", "testing": "验证中...", @@ -243,7 +243,7 @@ "skipModelTitle": "跳过模型配置？", "skipModelDescription": "可以先跳过此步骤，稍后在左侧导航栏「模型清单」中继续完成默认模型配置。", "skipIntelTitle": "跳过微步情报服务？", - "skipIntelDescription": "可以先跳过此步骤，稍后在左侧导航栏「工具清单」页面的「MCP」页签中，找到 ThreatBook MCP 后继续完成微步情报 API 与 MCP 配置。", + "skipIntelDescription": "可以先跳过此步骤。稍后可在左侧导航栏「工具清单」的 API 页签中找到 threatbook-cn 或 threatbook-io 配置情报 API；中国区 MCP 可在 MCP 页签的 ThreatBook MCP 中继续配置。", "returnToConfig": "返回配置", "confirmSkip": "继续跳过", "summaryTitle": "配置小结", diff --git a/webui/src/locales/zh-CN/tool.json b/webui/src/locales/zh-CN/tool.json index efcb16cca..c58dcc794 100644 --- a/webui/src/locales/zh-CN/tool.json +++ b/webui/src/locales/zh-CN/tool.json @@ -387,37 +387,71 @@ "enableServer": "启用服务器", "hide": "隐藏", "show": "显示", + "threatbookApi": { + "title": "配置 {{service}}", + "descriptions": { + "cn": "免费启用微步中国区威胁情报 API 服务，为 Rex 提供文件样本分析与 URL 扫描能力。", + "global": "免费启用微步国际区威胁情报 API 服务，为 Rex 提供 IP、域名、URL 和文件哈希等威胁情报查询能力。" + }, + "region": "服务区域", + "regionHint": "当前工具已绑定固定服务区域，Key 领取入口和 API 服务地址会自动匹配。", + "regions": { "cn": "中国区", "global": "国际区" }, + "freeService": "免费 ThreatBook API 服务", + "apiKey": "API Key", + "keyPlaceholder": "粘贴{{region}} ThreatBook API Key", + "keyHint": "领取后将 Key 粘贴到此处，系统会安全存储并验证对应区域的 API 服务。", + "keyRequired": "请先填写 API Key。", + "revealFailed": "读取已保存的 API Key 失败，请重试。", + "keyConfigured": "已安全配置", + "claimFreeKey": "领取{{region}}免费 Key", + "endpoint": "API 服务地址", + "endpointHint": "服务地址由当前工具自动匹配，无需手动修改。", + "saveAndVerify": "保存并验证", + "saving": "正在验证并保存...", + "saveSuccess": "已配置{{region}}并验证通过。", + "saveFailed": "配置失败，请检查 API Key。", + "configuredTitle": "已配置{{region}}", + "connected": "配置已验证，API 服务当前可用。", + "saved": "配置已保存，可重新测试服务连通性。", + "edit": "编辑配置", + "retest": "重新测试", + "testing": "测试中...", + "testFailed": "服务测试失败。" + }, "threatbookMcp": { "title": "配置 ThreatBook MCP", - "description": "免费启用微步威胁情报 MCP 服务，为 Rex 提供 IOC 查询、威胁情报研判和安全分析能力。", + "description": "免费启用微步中国区威胁情报 MCP 服务，为 Rex 提供 IOC 查询、威胁情报研判和安全分析能力。", "region": "服务区域", - "regionHint": "请选择 API Key 所属区域，系统会自动匹配对应的领取入口和 MCP 服务地址。", + "regionHint": "ThreatBook MCP 当前仅提供中国区服务，系统会自动匹配领取入口和 MCP 服务地址。", "regions": { "cn": "中国区", "global": "国际区" }, "freeService": "免费 ThreatBook MCP 服务", "apiKey": "API Key", - "keyPlaceholder": "粘贴当前区域的 ThreatBook API Key", + "keyPlaceholder": "粘贴中国区 ThreatBook API Key", "keyHint": "领取后将 Key 粘贴到此处，系统会自动完成服务地址拼接和安全存储。", "keyRequired": "请先填写当前区域的 API Key。", "revealFailed": "读取已保存的 API Key 失败，请重试。", "keyConfigured": "已安全配置", - "claimFreeKey": "领取免费 API Key", + "claimFreeKey": "领取中国区免费 Key", "endpoint": "MCP 服务地址", "endpointHint": "服务地址由所选区域自动生成，API Key 将单独加密保存，无需手动拼接。", "saveAndVerify": "保存并验证连接", "saving": "正在验证并保存...", "saveSuccess": "ThreatBook MCP 已配置并验证通过。", "saveFailed": "ThreatBook MCP 配置失败，请检查区域和 API Key。", - "configuredTitle": "已配置{{region}}", + "configuredTitle": "已配置中国区", "connected": "配置已验证，MCP 服务当前已连接。", "savedNotConnected": "配置已保存，MCP 服务当前未连接。", "edit": "编辑配置", "retest": "重新测试", "testing": "测试中...", "testFailed": "连接测试失败。", - "loading": "正在读取 ThreatBook MCP 配置..." + "loading": "正在读取 ThreatBook MCP 配置...", + "legacyGlobalTitle": "检测到历史国际区 MCP 配置", + "legacyGlobalDescription": "ThreatBook MCP 当前仅支持中国区。原配置不会自动复用，请领取并填写中国区 Key 完成迁移。", + "migrateToChina": "改为中国区配置" } }, diff --git a/webui/src/pages/Tool/components/ServiceDetailPanel.test.tsx b/webui/src/pages/Tool/components/ServiceDetailPanel.test.tsx index 6e228c963..9913e8c64 100644 --- a/webui/src/pages/Tool/components/ServiceDetailPanel.test.tsx +++ b/webui/src/pages/Tool/components/ServiceDetailPanel.test.tsx @@ -134,13 +134,16 @@ vi.mock('react-i18next', () => ({ 'detail.threatbookMcp.saveSuccess': isChinese ? '配置成功' : 'Configured', 'detail.threatbookMcp.saveFailed': isChinese ? '配置失败' : 'Setup failed', 'detail.threatbookMcp.loading': isChinese ? '读取配置' : 'Loading configuration', - 'detail.threatbookMcp.configuredTitle': isChinese ? '已配置{{region}}' : '{{region}} configured', + 'detail.threatbookMcp.configuredTitle': isChinese ? '已配置中国区' : 'China configured', 'detail.threatbookMcp.connected': isChinese ? '当前已连接' : 'Connected', 'detail.threatbookMcp.savedNotConnected': isChinese ? '当前未连接' : 'Not connected', 'detail.threatbookMcp.keyConfigured': isChinese ? '已安全配置' : 'Securely configured', 'detail.threatbookMcp.edit': isChinese ? '编辑配置' : 'Edit configuration', 'detail.threatbookMcp.retest': isChinese ? '重新测试' : 'Test again', 'detail.threatbookMcp.testing': isChinese ? '测试中...' : 'Testing...', + 'detail.threatbookMcp.legacyGlobalTitle': isChinese ? '检测到历史国际区 MCP 配置' : 'Legacy International MCP configuration detected', + 'detail.threatbookMcp.legacyGlobalDescription': isChinese ? '需要迁移到中国区' : 'Migrate to China', + 'detail.threatbookMcp.migrateToChina': isChinese ? '改为中国区配置' : 'Switch to China', 'alert.connectionOk': '连接成功', }; return (translations[key] ?? key).replace('{{region}}', String(options?.region ?? '')); @@ -300,11 +303,11 @@ describe('MCPServerDetailPanel', () => { const link = await screen.findByRole('link', { name: '领取免费 API Key' }); expect(link).toHaveAttribute('href', 'https://x.threatbook.com/flocks/activate'); expect(link).toHaveAttribute('target', '_blank'); - expect(screen.getByRole('button', { name: '中国区' })).toBeInTheDocument(); + expect(screen.getByText('中国区')).toBeInTheDocument(); expect(screen.getByDisplayValue('https://mcp.threatbook.cn/mcp')).toBeInTheDocument(); }); - it('shows the international free key link for ThreatBook MCP in English', async () => { + it('keeps ThreatBook MCP on the China region in English', async () => { currentLanguage.value = 'en-US'; render( @@ -320,12 +323,13 @@ describe('MCPServerDetailPanel', () => { ); const link = await screen.findByRole('link', { name: 'Claim free API key' }); - expect(link).toHaveAttribute('href', 'https://i.threatbook.io/flocks/activate'); - expect(screen.getByRole('button', { name: 'International' })).toBeInTheDocument(); - expect(screen.getByDisplayValue('https://mcp.threatbook.io/mcp')).toBeInTheDocument(); + expect(link).toHaveAttribute('href', 'https://x.threatbook.com/flocks/activate'); + expect(screen.getByText('China')).toBeInTheDocument(); + expect(screen.queryByText('International')).not.toBeInTheDocument(); + expect(screen.getByDisplayValue('https://mcp.threatbook.cn/mcp')).toBeInTheDocument(); }); - it('sends the selected region and API key through the closed-loop setup action', async () => { + it('sends the China region and API key through the closed-loop setup action', async () => { const user = userEvent.setup(); render( @@ -341,13 +345,12 @@ describe('MCPServerDetailPanel', () => { />, ); - await user.click(await screen.findByRole('button', { name: '国际区' })); - await user.type(screen.getByPlaceholderText('粘贴当前区域的 ThreatBook API Key'), 'global-key'); + await user.type(await screen.findByPlaceholderText('粘贴当前区域的 ThreatBook API Key'), 'global-key'); await user.click(screen.getByRole('button', { name: '保存并验证连接' })); await waitFor(() => { expect(mcpAPI.configureThreatBook).toHaveBeenCalledWith('threatbook_mcp', { - region: 'global', + region: 'cn', api_key: 'global-key', }); }); @@ -446,17 +449,9 @@ describe('MCPServerDetailPanel', () => { expect(keyInput).toHaveAttribute('type', 'password'); expect(keyInput).toHaveAttribute('readonly'); - await user.click(screen.getByRole('button', { name: '国际区' })); - expect(keyInput).toHaveValue(''); - expect(keyInput).not.toHaveAttribute('readonly'); }); - it('locks region switching while the saved key is being revealed', async () => { - const user = userEvent.setup(); - let resolveReveal: ((value: { data: { api_key: string } }) => void) | undefined; - mcpAPI.revealCredentials.mockReturnValue(new Promise((resolve) => { - resolveReveal = resolve; - })); + it('offers migration instead of reusing a legacy international MCP key', async () => { mcpAPI.get.mockResolvedValue({ ...detailResponse, data: { @@ -464,7 +459,7 @@ describe('MCPServerDetailPanel', () => { name: 'threatbook_mcp', config: { type: 'sse', - url: 'https://mcp.threatbook.cn/mcp?apikey={secret:threatbook_mcp_key}', + url: 'https://mcp.threatbook.io/mcp?apikey={secret:threatbook_mcp_key}', }, }, }); @@ -488,15 +483,10 @@ describe('MCPServerDetailPanel', () => { />, ); - await user.click(await screen.findByRole('button', { name: '编辑配置' })); - await user.click(screen.getByTitle('显示')); - - expect(screen.getByRole('button', { name: '国际区' })).toBeDisabled(); - - resolveReveal?.({ data: { api_key: '312abcdef321' } }); - await waitFor(() => { - expect(screen.getByRole('button', { name: '国际区' })).toBeEnabled(); - }); + expect(await screen.findByText('检测到历史国际区 MCP 配置')).toBeInTheDocument(); + await userEvent.setup().click(screen.getByRole('button', { name: '改为中国区配置' })); + expect(screen.getByPlaceholderText('粘贴当前区域的 ThreatBook API Key')).toHaveValue(''); + expect(mcpAPI.revealCredentials).not.toHaveBeenCalled(); }); it('does not show the free key link for other MCP servers', async () => { diff --git a/webui/src/pages/Tool/components/ServiceDetailPanel.tsx b/webui/src/pages/Tool/components/ServiceDetailPanel.tsx index eb3686991..b3fe345cb 100644 --- a/webui/src/pages/Tool/components/ServiceDetailPanel.tsx +++ b/webui/src/pages/Tool/components/ServiceDetailPanel.tsx @@ -16,11 +16,16 @@ import { buildMCPConfigFromForm, buildMCPFormDataFromConfig, getMCPFormError, MC import type { MCPFormData, ConnStatus as MCPConnStatus } from '../ToolSheets'; import type { APIServiceCredentialField, APIServiceMetadata, ProviderCredentials } from '@/types'; import ThreatBookMCPConfigPanel from './ThreatBookMCPConfigPanel'; +import ThreatBookAPIConfigPanel from './ThreatBookAPIConfigPanel'; function isThreatBookMCPServer(serverName: string): boolean { return serverName.trim().toLowerCase() === 'threatbook_mcp'; } +function isThreatBookAPIService(serviceName: string): boolean { + return ['threatbook-cn', 'threatbook-io'].includes(serviceName.trim().toLowerCase()); +} + function KvRowValue({ value }: { value: string }) { const [showTooltip, setShowTooltip] = useState(false); @@ -810,6 +815,15 @@ export function APIServiceDetailPanel({
) : detailTab === 'overview' ? (
+ {isThreatBookAPIService(serviceName) ? ( + + ) : (
{(() => { const status = quickTesting @@ -930,6 +944,7 @@ export function APIServiceDetailPanel({
+ )}
+ +
+
+
+
+
+ {t("detail.threatbookApi.region")} +
+
{regionLabel}
+
+
+
+ {t("detail.threatbookApi.apiKey")} +
+
+ {credentials?.api_key_masked || + t("detail.threatbookApi.keyConfigured")} +
+
+
+
+ {t("detail.threatbookApi.endpoint")} +
+
+ {config.apiEndpoint} +
+
+
+
+ ) : ( +
+
+

+ {t("detail.threatbookApi.region")} +

+

+ {t("detail.threatbookApi.regionHint")} +

+
+ {regionLabel} +
+
+
+ {t("detail.threatbookApi.freeService")} +
+ +
+ + +

+ {t("detail.threatbookApi.endpointHint")} +

+
+ {result && } +
+ {isConfigured && ( + + )} + +
+
+ )} + {result && isConfigured && !editing && } +
+ ); +} diff --git a/webui/src/pages/Tool/components/ThreatBookMCPConfigPanel.tsx b/webui/src/pages/Tool/components/ThreatBookMCPConfigPanel.tsx index 594850451..e50ea58d3 100644 --- a/webui/src/pages/Tool/components/ThreatBookMCPConfigPanel.tsx +++ b/webui/src/pages/Tool/components/ThreatBookMCPConfigPanel.tsx @@ -1,5 +1,5 @@ -import { useEffect, useMemo, useState } from 'react'; -import { useTranslation } from 'react-i18next'; +import { useEffect, useMemo, useState } from "react"; +import { useTranslation } from "react-i18next"; import { AlertCircle, CheckCircle2, @@ -9,58 +9,63 @@ import { Loader2, Pencil, RefreshCw, -} from 'lucide-react'; -import { mcpAPI } from '@/api/mcp'; +} from "lucide-react"; +import { mcpAPI } from "@/api/mcp"; import { - getDefaultThreatBookRegion, inferThreatBookRegionFromMcpUrl, THREATBOOK_REGION_CONFIG, - type ThreatBookRegion, -} from '@/constants/threatbook'; -import type { MCPCredentials, MCPServer } from '@/types'; +} from "@/constants/threatbook"; +import type { MCPCredentials, MCPServer } from "@/types"; -interface ThreatBookMCPConfigPanelProps { +interface Props { serverName: string; - serverStatus: MCPServer['status']; + serverStatus: MCPServer["status"]; configUrl?: string; onConfigured: () => Promise; } -type OperationResult = { - success: boolean; - message: string; -} | null; +type OperationResult = { success: boolean; message: string } | null; + +function ResultMessage({ result }: { result: Exclude }) { + return ( +
+ {result.success ? ( + + ) : ( + + )} + {result.message} +
+ ); +} export default function ThreatBookMCPConfigPanel({ serverName, serverStatus, configUrl, onConfigured, -}: ThreatBookMCPConfigPanelProps) { - const { t, i18n } = useTranslation('tool'); +}: Props) { + const { t } = useTranslation("tool"); const configuredRegion = useMemo( () => inferThreatBookRegionFromMcpUrl(configUrl), [configUrl], ); - const languageDefaultRegion = useMemo( - () => getDefaultThreatBookRegion(i18n.resolvedLanguage || i18n.language), - [i18n.language, i18n.resolvedLanguage], - ); - const [region, setRegion] = useState(configuredRegion || languageDefaultRegion); const [credentials, setCredentials] = useState(null); - const [loadingCredentials, setLoadingCredentials] = useState(true); + const [loading, setLoading] = useState(true); const [editing, setEditing] = useState(true); - const [apiKey, setApiKey] = useState(''); + const [apiKey, setApiKey] = useState(""); const [showApiKey, setShowApiKey] = useState(false); const [storedKeyLoaded, setStoredKeyLoaded] = useState(false); - const [revealingKey, setRevealingKey] = useState(false); + const [revealing, setRevealing] = useState(false); const [saving, setSaving] = useState(false); const [testing, setTesting] = useState(false); const [result, setResult] = useState(null); const loadCredentials = async () => { try { - setLoadingCredentials(true); + setLoading(true); const response = await mcpAPI.getCredentials(serverName); setCredentials(response.data); setEditing(!(response.data.has_credential && configuredRegion)); @@ -68,48 +73,41 @@ export default function ThreatBookMCPConfigPanel({ setCredentials(null); setEditing(true); } finally { - setLoadingCredentials(false); + setLoading(false); } }; - useEffect(() => { - setRegion(configuredRegion || languageDefaultRegion); - }, [configuredRegion, languageDefaultRegion]); - useEffect(() => { void loadCredentials(); }, [serverName, configuredRegion]); - const regionConfig = THREATBOOK_REGION_CONFIG[region]; - const summaryRegion = configuredRegion || region; - const isConfigured = Boolean(credentials?.has_credential && configuredRegion); - const isConnected = serverStatus === 'connected'; - const hasStoredKeyForRegion = Boolean( - credentials?.has_credential && configuredRegion && configuredRegion === region, + const regionConfig = THREATBOOK_REGION_CONFIG.cn; + const isConfigured = Boolean( + credentials?.has_credential && configuredRegion === "cn", + ); + const isLegacyGlobal = Boolean( + credentials?.has_credential && configuredRegion === "global", ); - const showingStoredKeyMask = hasStoredKeyForRegion && !showApiKey; - const displayedApiKey = showingStoredKeyMask - ? credentials?.api_key_masked || '' + const showingStoredMask = isConfigured && !showApiKey; + const displayedApiKey = showingStoredMask + ? credentials?.api_key_masked || "" : apiKey; - const handleRegionChange = (nextRegion: ThreatBookRegion) => { - if (nextRegion === region) return; - setRegion(nextRegion); - setApiKey(''); + const resetEditor = () => { + setApiKey(""); setShowApiKey(false); setStoredKeyLoaded(false); setResult(null); }; - const handleApiKeyVisibility = async () => { + const handleVisibility = async () => { if (showApiKey) { setShowApiKey(false); return; } - - if (hasStoredKeyForRegion && !storedKeyLoaded) { + if (isConfigured && !storedKeyLoaded) { try { - setRevealingKey(true); + setRevealing(true); setResult(null); const response = await mcpAPI.revealCredentials(serverName); setApiKey(response.data.api_key); @@ -117,47 +115,54 @@ export default function ThreatBookMCPConfigPanel({ } catch (error: any) { setResult({ success: false, - message: error.response?.data?.detail || error.message || t('detail.threatbookMcp.revealFailed'), + message: + error.response?.data?.detail || + error.message || + t("detail.threatbookMcp.revealFailed"), }); return; } finally { - setRevealingKey(false); + setRevealing(false); } } - setShowApiKey(true); }; const handleSave = async () => { - const trimmedKey = apiKey.trim(); - if (!trimmedKey) { - setResult({ success: false, message: t('detail.threatbookMcp.keyRequired') }); + const value = apiKey.trim(); + if (!value) { + setResult({ + success: false, + message: t("detail.threatbookMcp.keyRequired"), + }); return; } - try { setSaving(true); setResult(null); const response = await mcpAPI.configureThreatBook(serverName, { - region, - api_key: trimmedKey, + region: "cn", + api_key: value, }); if (!response.data.success) { setResult({ success: false, message: response.data.message }); return; } - - setApiKey(''); - setShowApiKey(false); - setStoredKeyLoaded(false); - setResult({ success: true, message: t('detail.threatbookMcp.saveSuccess') }); + resetEditor(); + setResult({ + success: true, + message: t("detail.threatbookMcp.saveSuccess"), + }); await onConfigured(); await loadCredentials(); setEditing(false); } catch (error: any) { setResult({ success: false, - message: error.response?.data?.detail || error.message || t('detail.threatbookMcp.saveFailed'), + message: + error.response?.data?.detail || + error.message || + t("detail.threatbookMcp.saveFailed"), }); } finally { setSaving(false); @@ -177,18 +182,21 @@ export default function ThreatBookMCPConfigPanel({ } catch (error: any) { setResult({ success: false, - message: error.response?.data?.detail || error.message || t('detail.threatbookMcp.testFailed'), + message: + error.response?.data?.detail || + error.message || + t("detail.threatbookMcp.testFailed"), }); } finally { setTesting(false); } }; - if (loadingCredentials) { + if (loading) { return (
- {t('detail.threatbookMcp.loading')} + {t("detail.threatbookMcp.loading")}
); } @@ -196,225 +204,230 @@ export default function ThreatBookMCPConfigPanel({ return (
-

{t('detail.threatbookMcp.title')}

-

{t('detail.threatbookMcp.description')}

+

+ {t("detail.threatbookMcp.title")} +

+

+ {t("detail.threatbookMcp.description")} +

+ {isLegacyGlobal && ( +
+
+ +
+

+ {t("detail.threatbookMcp.legacyGlobalTitle")} +

+

+ {t("detail.threatbookMcp.legacyGlobalDescription")} +

+
+ +
+
+ )} + {isConfigured && !editing ? (
-
- -
+
+ +

- {t('detail.threatbookMcp.configuredTitle', { - region: t(`detail.threatbookMcp.regions.${summaryRegion}`), - })} + {t("detail.threatbookMcp.configuredTitle")}

- {isConnected - ? t('detail.threatbookMcp.connected') - : t('detail.threatbookMcp.savedNotConnected')} + {serverStatus === "connected" + ? t("detail.threatbookMcp.connected") + : t("detail.threatbookMcp.savedNotConnected")}

-
+
-
-
{t('detail.threatbookMcp.region')}
-
- {t(`detail.threatbookMcp.regions.${summaryRegion}`)} +
+ {t("detail.threatbookMcp.region")} +
+
+ {t("detail.threatbookMcp.regions.cn")}
-
{t('detail.threatbookMcp.apiKey')}
-
- {credentials?.api_key_masked || t('detail.threatbookMcp.keyConfigured')} +
+ {t("detail.threatbookMcp.apiKey")} +
+
+ {credentials?.api_key_masked || + t("detail.threatbookMcp.keyConfigured")}
-
{t('detail.threatbookMcp.endpoint')}
-
- {THREATBOOK_REGION_CONFIG[summaryRegion].mcpEndpoint} +
+ {t("detail.threatbookMcp.endpoint")} +
+
+ {regionConfig.mcpEndpoint}
- ) : ( + ) : !isLegacyGlobal || editing ? (
-
-
-

{t('detail.threatbookMcp.region')}

-

{t('detail.threatbookMcp.regionHint')}

-
-
- {(['cn', 'global'] as const).map((item) => ( - - ))} +
+

+ {t("detail.threatbookMcp.region")} +

+

+ {t("detail.threatbookMcp.regionHint")} +

+
+ {t("detail.threatbookMcp.regions.cn")}
-
- {t('detail.threatbookMcp.freeService')} + {t("detail.threatbookMcp.freeService")}
-
-
-
-
- - {result && ( -
- {result.success - ? - : } - {result.message} -
- )} - -
+ {result && } +
{isConfigured && ( )}
- )} - - {result && isConfigured && !editing && ( -
- {result.success - ? - : } - {result.message} -
- )} + ) : null} + {result && isConfigured && !editing && }
); } From 0e4fcfcc15f325b19f942c8566304849ccd75528 Mon Sep 17 00:00:00 2001 From: John Yin <10972267+john-yin2333@user.noreply.gitee.com> Date: Mon, 7 Sep 2026 10:16:32 +0800 Subject: [PATCH 46/63] Harden credential configuration flows --- flocks/mcp/utils.py | 68 ++++++- flocks/server/config_mutation.py | 64 +++++++ flocks/server/routes/channel.py | 2 + flocks/server/routes/config.py | 6 + flocks/server/routes/custom_provider.py | 5 + flocks/server/routes/default_model.py | 4 + flocks/server/routes/mcp.py | 142 ++++++++++++-- flocks/server/routes/model.py | 3 + flocks/server/routes/onboarding.py | 97 ++++++++-- flocks/server/routes/provider.py | 152 ++++++++++++--- flocks/server/routes/tool.py | 3 + flocks/tool/credential_context.py | 58 ++++++ flocks/tool/registry.py | 34 ++-- .../provider/test_api_service_credentials.py | 179 +++++++++++++++++- tests/provider/test_test_credentials.py | 11 ++ .../concurrency/test_config_mutation.py | 56 ++++++ .../routes/test_global_mutation_auth.py | 46 +++++ tests/server/routes/test_mcp_routes.py | 94 ++++++++- tests/server/routes/test_onboarding_routes.py | 82 ++++++++ tests/server/routes/test_tool_routes.py | 3 + ...test_credential_context_config_override.py | 20 ++ .../tool/test_failure_auto_disable_config.py | 22 +++ webui/src/api/provider.test.ts | 14 ++ webui/src/api/provider.ts | 11 ++ webui/src/components/layout/Layout.tsx | 10 +- webui/src/pages/Home/index.test.tsx | 3 + webui/src/pages/Home/index.tsx | 20 +- webui/src/pages/SecurityConfig/index.test.tsx | 4 + .../components/ServiceDetailPanelApi.test.tsx | 8 +- .../components/ThreatBookAPIConfigPanel.tsx | 20 +- 30 files changed, 1141 insertions(+), 100 deletions(-) create mode 100644 flocks/server/config_mutation.py create mode 100644 tests/server/concurrency/test_config_mutation.py diff --git a/flocks/mcp/utils.py b/flocks/mcp/utils.py index ef0df7c74..3ffefc20f 100644 --- a/flocks/mcp/utils.py +++ b/flocks/mcp/utils.py @@ -8,7 +8,7 @@ import re import hashlib from typing import Optional, Dict, Any -from urllib.parse import urlencode, urlparse, parse_qs, parse_qsl, urlunparse +from urllib.parse import unquote, urlencode, urlparse, parse_qs, parse_qsl, urlunparse _SENSITIVE_QUERY_PARAMS = frozenset({ "apikey", "api_key", "key", "token", "access_token", @@ -316,8 +316,6 @@ def extract_api_key_from_mcp_url(server_name: str, config: Dict[str, Any]) -> Di if not url or config.get("type") not in REMOTE_MCP_TYPES or "?" not in url: return dict(config) - from urllib.parse import unquote - base, _, raw_query = url.partition("?") fragment = "" if "#" in raw_query: @@ -445,6 +443,37 @@ def mask_sensitive_mcp_config_for_frontend( """Mask plain-text secrets before returning MCP config to the frontend.""" masked_config = dict(config) + url = config.get("url") + if isinstance(url, str) and "?" in url: + base, _, raw_query = url.partition("?") + fragment = "" + if "#" in raw_query: + raw_query, _, fragment = raw_query.partition("#") + + masked_parts: list[str] = [] + changed = False + for part in raw_query.split("&"): + if "=" not in part: + masked_parts.append(part) + continue + key_encoded, _, value_encoded = part.partition("=") + key = unquote(key_encoded) + value = unquote(value_encoded) + if ( + key.lower() in _SENSITIVE_QUERY_PARAMS + and not _is_secret_placeholder(value) + ): + masked_parts.append(f"{key_encoded}={MCP_MASKED_SECRET_VALUE}") + changed = True + else: + masked_parts.append(part) + + if changed: + masked_url = base + "?" + "&".join(masked_parts) + if fragment: + masked_url += "#" + fragment + masked_config["url"] = masked_url + auth_config = config.get("auth") if isinstance(auth_config, dict): auth_value = auth_config.get("value") @@ -478,6 +507,39 @@ def restore_masked_mcp_config_secrets( """Restore masked frontend sentinel values back to their previous secrets.""" restored_config = dict(updated_config) + previous_url = previous_config.get("url") + next_url = updated_config.get("url") + if ( + isinstance(previous_url, str) + and isinstance(next_url, str) + and "?" in previous_url + and "?" in next_url + ): + previous_parsed = urlparse(previous_url) + next_parsed = urlparse(next_url) + previous_values: Dict[str, str] = {} + for key, value in parse_qsl(previous_parsed.query, keep_blank_values=True): + if key.lower() in _SENSITIVE_QUERY_PARAMS: + previous_values[key.lower()] = value + + restored_query: list[tuple[str, str]] = [] + changed = False + for key, value in parse_qsl(next_parsed.query, keep_blank_values=True): + normalized_key = key.lower() + if ( + normalized_key in _SENSITIVE_QUERY_PARAMS + and value == MCP_MASKED_SECRET_VALUE + and normalized_key in previous_values + ): + value = previous_values[normalized_key] + changed = True + restored_query.append((key, value)) + + if changed: + restored_config["url"] = urlunparse( + next_parsed._replace(query=urlencode(restored_query)) + ) + previous_auth = previous_config.get("auth") next_auth = updated_config.get("auth") if ( diff --git a/flocks/server/config_mutation.py b/flocks/server/config_mutation.py new file mode 100644 index 000000000..10a36cff8 --- /dev/null +++ b/flocks/server/config_mutation.py @@ -0,0 +1,64 @@ +"""Serialization for global configuration transactions.""" + +from __future__ import annotations + +import asyncio +from contextlib import asynccontextmanager +from functools import wraps +from typing import Any, AsyncIterator, Awaitable, Callable, TypeVar, cast + + +T = TypeVar("T") + + +class AsyncReentrantLock: + """An asyncio lock that permits nested acquisition by the current task.""" + + def __init__(self) -> None: + self._lock = asyncio.Lock() + self._owner: asyncio.Task[Any] | None = None + self._depth = 0 + + async def acquire(self) -> None: + task = asyncio.current_task() + if task is None: + raise RuntimeError("Configuration mutations require an asyncio task") + if self._owner is task: + self._depth += 1 + return + await self._lock.acquire() + self._owner = task + self._depth = 1 + + def release(self) -> None: + task = asyncio.current_task() + if task is None or self._owner is not task: + raise RuntimeError("Configuration mutation lock released by a non-owner") + self._depth -= 1 + if self._depth == 0: + self._owner = None + self._lock.release() + + @asynccontextmanager + async def hold(self) -> AsyncIterator[None]: + await self.acquire() + try: + yield + finally: + self.release() + + +GLOBAL_CONFIG_MUTATION_LOCK = AsyncReentrantLock() + + +def serialized_config_mutation( + function: Callable[..., Awaitable[T]], +) -> Callable[..., Awaitable[T]]: + """Serialize an async route or helper that mutates global configuration.""" + + @wraps(function) + async def wrapped(*args: Any, **kwargs: Any) -> T: + async with GLOBAL_CONFIG_MUTATION_LOCK.hold(): + return await function(*args, **kwargs) + + return cast(Callable[..., Awaitable[T]], wrapped) diff --git a/flocks/server/routes/channel.py b/flocks/server/routes/channel.py index 5e872ca85..f701ff44c 100644 --- a/flocks/server/routes/channel.py +++ b/flocks/server/routes/channel.py @@ -15,6 +15,7 @@ from flocks.channel.gateway.manager import default_manager from flocks.channel.registry import default_registry from flocks.hooks.execution import ExecutionStopped, execute_with_hooks +from flocks.server.config_mutation import serialized_config_mutation from flocks.utils.log import Log router = APIRouter() @@ -534,6 +535,7 @@ def _append_telegram_allow_from(user_id: str) -> None: @router.post("/telegram/pair") +@serialized_config_mutation async def telegram_pair(req: TelegramPairRequest): """ Verify a Telegram pairing code. diff --git a/flocks/server/routes/config.py b/flocks/server/routes/config.py index 5920a2b65..c22364a34 100644 --- a/flocks/server/routes/config.py +++ b/flocks/server/routes/config.py @@ -29,6 +29,7 @@ from flocks.config.config import Config, GlobalConfig, ConfigInfo as ConfigInfoModel, UIConfig from flocks.config.config_writer import ConfigWriter from flocks.provider.provider import Provider +from flocks.server.config_mutation import serialized_config_mutation from flocks.utils.log import Log @@ -465,6 +466,7 @@ async def get_ui_display() -> UIDisplayResponse: @router.patch("/ui", response_model=UIDisplayResponse, summary="Update UI display preferences") +@serialized_config_mutation async def update_ui_config(request: UIConfigUpdateRequest) -> UIDisplayResponse: """Update visible WebUI display preferences.""" try: @@ -501,6 +503,7 @@ async def get_ui_favicon() -> FileResponse: @router.post("/ui/favicon", response_model=UIDisplayResponse, summary="Upload UI favicon") +@serialized_config_mutation async def upload_ui_favicon(file: UploadFile = File(...)) -> UIDisplayResponse: """Upload a custom favicon for visible WebUI branding.""" filename = Path(file.filename or "").name @@ -551,6 +554,7 @@ async def upload_ui_favicon(file: UploadFile = File(...)) -> UIDisplayResponse: @router.delete("/ui/favicon", response_model=UIDisplayResponse, summary="Reset UI favicon") +@serialized_config_mutation async def reset_ui_favicon() -> UIDisplayResponse: """Remove the uploaded favicon and fall back to the default bundled favicon.""" assets_dir = _ui_assets_dir() @@ -591,6 +595,7 @@ async def get_tool_failure_preference() -> ToolFailurePreference: response_model=ToolFailurePreference, summary="Update repeated tool-failure preference", ) +@serialized_config_mutation async def update_tool_failure_preference( request: ToolFailurePreference, ) -> ToolFailurePreference: @@ -630,6 +635,7 @@ async def get_config() -> Dict[str, Any]: @router.patch("/", summary="Update configuration") +@serialized_config_mutation async def update_config(config_data: Dict[str, Any]) -> Dict[str, Any]: """ Update configuration diff --git a/flocks/server/routes/custom_provider.py b/flocks/server/routes/custom_provider.py index 9e5bec81c..ddbc04af9 100644 --- a/flocks/server/routes/custom_provider.py +++ b/flocks/server/routes/custom_provider.py @@ -25,6 +25,7 @@ ProviderConfig, ) from flocks.provider.sdk.openai_compatible import OpenAICompatibleProvider +from flocks.server.config_mutation import serialized_config_mutation from flocks.utils.log import Log @@ -175,6 +176,7 @@ async def list_providers(): @router.post("/providers", response_model=ProviderResp, status_code=201) +@serialized_config_mutation async def create_provider(body: CreateProviderReq): """Create a custom OpenAI-compatible provider in flocks.json.""" pid = "custom-" + body.name.lower().replace(" ", "-").replace("_", "-") @@ -219,6 +221,7 @@ async def create_provider(body: CreateProviderReq): @router.delete("/providers/{provider_id}", status_code=204) +@serialized_config_mutation async def delete_provider(provider_id: str): """Delete a custom provider from flocks.json.""" raw = ConfigWriter.get_provider_raw(provider_id) @@ -278,6 +281,7 @@ async def list_models(provider_id: str): @router.post("/models/{provider_id}", response_model=ModelResp, status_code=201) +@serialized_config_mutation async def create_model(provider_id: str, body: CreateModelReq): """Add a model to a provider in flocks.json.""" raw = ConfigWriter.get_provider_raw(provider_id) @@ -337,6 +341,7 @@ async def create_model(provider_id: str, body: CreateModelReq): @router.delete("/models/{provider_id}/{model_id:path}", status_code=204) +@serialized_config_mutation async def delete_model(provider_id: str, model_id: str): """Remove a model from a provider in flocks.json.""" removed = ConfigWriter.remove_model(provider_id, model_id) diff --git a/flocks/server/routes/default_model.py b/flocks/server/routes/default_model.py index 20477eaa0..c5cf9c2ee 100644 --- a/flocks/server/routes/default_model.py +++ b/flocks/server/routes/default_model.py @@ -14,6 +14,7 @@ from flocks.provider.model_manager import get_model_manager from flocks.provider.provider import Provider from flocks.provider.types import DefaultModelConfig, ModelType +from flocks.server.config_mutation import serialized_config_mutation from flocks.utils.log import Log router = APIRouter() @@ -95,6 +96,7 @@ async def get_fallback_providers() -> FallbackProvidersConfig: summary="Replace runtime fallback models", description="Atomically replace the ordered fallback model configuration", ) +@serialized_config_mutation async def set_fallback_providers( body: FallbackProvidersConfig, ) -> FallbackProvidersConfig: @@ -226,6 +228,7 @@ async def get_default_model(model_type: ModelType) -> DefaultModelConfig: response_model=DefaultModelConfig, summary="Set default model for type", ) +@serialized_config_mutation async def set_default_model( model_type: ModelType, body: SetDefaultModelRequest ) -> DefaultModelConfig: @@ -244,6 +247,7 @@ async def set_default_model( status_code=status.HTTP_204_NO_CONTENT, summary="Delete default model for type", ) +@serialized_config_mutation async def delete_default_model(model_type: ModelType): """Remove default model setting for a model type.""" manager = get_model_manager() diff --git a/flocks/server/routes/mcp.py b/flocks/server/routes/mcp.py index 76e9ef8eb..ae23feeca 100644 --- a/flocks/server/routes/mcp.py +++ b/flocks/server/routes/mcp.py @@ -10,10 +10,14 @@ import asyncio import copy from typing import Dict, Optional, List, Any -from fastapi import APIRouter, HTTPException, Response +from urllib.parse import parse_qs, urlparse + +from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import JSONResponse from pydantic import BaseModel, Field +from flocks.audit import emit_audit_event +from flocks.auth.context import AuthUser from flocks.mcp import ( MCP, get_manager, @@ -44,6 +48,8 @@ from flocks.config.config_writer import ConfigWriter from flocks.security import get_secret_manager from flocks.security.secrets import get_mcp_secret_id +from flocks.server.auth import get_request_ip, get_request_user_agent, require_admin +from flocks.server.config_mutation import serialized_config_mutation from flocks.server.threatbook_regions import ( THREATBOOK_REGION_PRESETS, ThreatBookRegion, @@ -232,7 +238,11 @@ async def get_mcp_status(): description="Dynamically add a new Model Context Protocol (MCP) server to the system.", operation_id="mcp.add" ) -async def add_mcp_server(request: McpAddRequest): +@serialized_config_mutation +async def add_mcp_server( + request: McpAddRequest, + _admin: object = Depends(require_admin), +): """ Add a new MCP server and persist to both flocks.json and ``~/.flocks/plugins/tools/mcp/``. @@ -324,7 +334,10 @@ class ThreatBookMcpConfigureResponse(BaseModel): description="Test an MCP server connection without saving to configuration.", operation_id="mcp.test" ) -async def test_mcp_connection(request: McpTestRequest): +async def test_mcp_connection( + request: McpTestRequest, + _admin: object = Depends(require_admin), +): """ Test MCP server connectivity without persisting configuration. @@ -406,9 +419,11 @@ async def _restore_threatbook_mcp_setup( description="Validate a regional ThreatBook key before atomically saving and connecting MCP.", operation_id="mcp.threatbook_configure", ) +@serialized_config_mutation async def configure_threatbook_mcp( name: str, request: ThreatBookMcpConfigureRequest, + _admin: object = Depends(require_admin), ) -> ThreatBookMcpConfigureResponse: """Validate first, then persist the China endpoint and secret reference.""" if request.region != "cn": @@ -512,7 +527,11 @@ async def configure_threatbook_mcp( description="Test an existing MCP server using saved config merged with temporary overrides.", operation_id="mcp.test_existing" ) -async def test_existing_mcp_connection(name: str, request: McpUpdateRequest): +async def test_existing_mcp_connection( + name: str, + request: McpUpdateRequest, + _admin: object = Depends(require_admin), +): """Test a configured MCP server after merging temporary config overrides.""" temp_name = f"{name}__test__" import time @@ -566,7 +585,11 @@ async def test_existing_mcp_connection(name: str, request: McpUpdateRequest): description="Disconnect and remove an MCP server from the system and configuration.", operation_id="mcp.remove" ) -async def remove_mcp_server(name: str): +@serialized_config_mutation +async def remove_mcp_server( + name: str, + _admin: object = Depends(require_admin), +): """ Remove an MCP server. @@ -622,7 +645,12 @@ async def remove_mcp_server(name: str): description="Update and persist an existing MCP server configuration.", operation_id="mcp.update" ) -async def update_mcp_server(name: str, request: McpUpdateRequest): +@serialized_config_mutation +async def update_mcp_server( + name: str, + request: McpUpdateRequest, + _admin: object = Depends(require_admin), +): """Update an existing MCP server configuration and clear stale runtime state.""" try: existing_config = _load_raw_mcp_server_config(name) @@ -754,7 +782,10 @@ async def get_mcp_server_info(name: str): description="Connect to an MCP server.", operation_id="mcp.connect" ) -async def connect_mcp_server(name: str): +async def connect_mcp_server( + name: str, + _admin: object = Depends(require_admin), +): """Connect to an MCP server - returns true on success""" try: server_config = await _load_mcp_server_config(name) @@ -801,7 +832,10 @@ async def connect_mcp_server(name: str): description="Disconnect from an MCP server.", operation_id="mcp.disconnect" ) -async def disconnect_mcp_server(name: str): +async def disconnect_mcp_server( + name: str, + _admin: object = Depends(require_admin), +): """Disconnect from an MCP server - returns true on success""" try: success = await MCP.disconnect(name) @@ -820,7 +854,10 @@ async def disconnect_mcp_server(name: str): description="Start OAuth authentication flow for a Model Context Protocol (MCP) server.", operation_id="mcp.auth.start" ) -async def start_mcp_auth(name: str): +async def start_mcp_auth( + name: str, + _admin: object = Depends(require_admin), +): """ Start OAuth authentication flow @@ -839,7 +876,10 @@ async def start_mcp_auth(name: str): description="Remove OAuth credentials for an MCP server.", operation_id="mcp.auth.remove" ) -async def remove_mcp_auth(name: str): +async def remove_mcp_auth( + name: str, + _admin: object = Depends(require_admin), +): """Remove OAuth credentials - returns {"success": true}""" try: await McpAuth.remove(name) @@ -941,7 +981,10 @@ async def get_server_resources(name: str): description="Refresh tools from an MCP server.", operation_id="mcp.refresh" ) -async def refresh_mcp_tools(name: str): +async def refresh_mcp_tools( + name: str, + _admin: object = Depends(require_admin), +): """Refresh tools from a server - returns count of tools registered""" try: count = await MCP.refresh_tools(name) @@ -993,6 +1036,20 @@ def _get_mcp_credential(name: str) -> tuple[Optional[str], Optional[str]]: if api_key: secret_id = legacy_id + if not api_key: + raw_config = _load_raw_mcp_server_config(name) or {} + raw_url = raw_config.get("url") + if isinstance(raw_url, str): + query = parse_qs(urlparse(raw_url).query, keep_blank_values=True) + for key, values in query.items(): + if key.lower() not in {"apikey", "api_key", "api-key", "token", "access_token"}: + continue + candidate = values[0].strip() if values else "" + if candidate and not candidate.startswith("{secret:"): + api_key = candidate + secret_id = None + break + return (secret_id, api_key) if api_key else (None, None) @@ -1007,7 +1064,10 @@ def _mask_mcp_credential(api_key: str) -> str: summary="Get MCP server credentials (masked)", description="Get masked credential information for an MCP server." ) -async def get_mcp_credentials(name: str): +async def get_mcp_credentials( + name: str, + _admin: object = Depends(require_admin), +): """Get masked credential info for a server. Avoids duplicating an existing ``_mcp`` suffix and falls back to historical @@ -1032,14 +1092,39 @@ async def get_mcp_credentials(name: str): summary="Reveal MCP server credential", description="Reveal a stored MCP credential after an explicit user action." ) -async def reveal_mcp_credential(name: str, response: Response): +async def reveal_mcp_credential( + name: str, + request: Request, + response: Response, + admin: AuthUser = Depends(require_admin), +): """Return the full key only for an explicit reveal request.""" response.headers["Cache-Control"] = "no-store" response.headers["Pragma"] = "no-cache" try: - _, api_key = _get_mcp_credential(name) + secret_id, api_key = _get_mcp_credential(name) if not api_key: raise HTTPException(status_code=404, detail="No credentials found for this server") + try: + await emit_audit_event( + "mcp.credentials_reveal", + { + "action": "mcp_credentials_reveal", + "actor_id": admin.id, + "actor_name": admin.username, + "user_id": admin.id, + "username": admin.username, + "mcp_name": name, + "secret_id": secret_id, + "ip": get_request_ip(request), + "user_agent": get_request_user_agent(request), + }, + ) + except Exception as audit_error: + log.warn( + "mcp.credentials.reveal.audit_failed", + {"name": name, "error": str(audit_error)}, + ) return McpCredentialRevealResponse(api_key=api_key) except HTTPException: raise @@ -1054,7 +1139,12 @@ async def reveal_mcp_credential(name: str, response: Response): summary="Set MCP server credentials", description="Set authentication credentials for an MCP server." ) -async def set_mcp_credentials(name: str, request: McpCredentialRequest): +@serialized_config_mutation +async def set_mcp_credentials( + name: str, + request: McpCredentialRequest, + _admin: object = Depends(require_admin), +): """Set credentials for a server. Stores in .secret.json with flat KV format. @@ -1090,7 +1180,11 @@ async def set_mcp_credentials(name: str, request: McpCredentialRequest): summary="Delete MCP server credentials", description="Delete stored credentials for an MCP server." ) -async def delete_mcp_credentials(name: str): +@serialized_config_mutation +async def delete_mcp_credentials( + name: str, + _admin: object = Depends(require_admin), +): """Delete credentials for a server.""" try: secrets = get_secret_manager() @@ -1121,7 +1215,10 @@ async def delete_mcp_credentials(name: str): summary="Test MCP server credentials", description="Test if the stored credentials are valid by attempting connection." ) -async def test_mcp_credentials(name: str): +async def test_mcp_credentials( + name: str, + _admin: object = Depends(require_admin), +): """Test credentials by attempting connection""" try: import time @@ -1311,7 +1408,10 @@ async def get_catalog_configured(): description="Batch-configure all catalog entries that don't require API keys.", operation_id="mcp.catalog.auto_setup" ) -async def auto_setup_catalog(): +@serialized_config_mutation +async def auto_setup_catalog( + _admin: object = Depends(require_admin), +): """Batch write all no-secret catalog entries to flocks.json with enabled=false.""" try: catalog = McpCatalog.get() @@ -1354,7 +1454,11 @@ async def auto_setup_catalog(): description="Add an MCP server from the catalog to your configuration.", operation_id="mcp.catalog.install" ) -async def install_from_catalog(request: CatalogInstallRequest): +@serialized_config_mutation +async def install_from_catalog( + request: CatalogInstallRequest, + _admin: object = Depends(require_admin), +): """Install an MCP server from catalog into flocks.json. If credentials are provided, they are saved to .secret.json and the config diff --git a/flocks/server/routes/model.py b/flocks/server/routes/model.py index e4f0bff80..683041e92 100644 --- a/flocks/server/routes/model.py +++ b/flocks/server/routes/model.py @@ -20,6 +20,7 @@ ModelSetting, ModelType, ) +from flocks.server.config_mutation import serialized_config_mutation from flocks.utils.log import Log @@ -503,6 +504,7 @@ async def get_model_settings(provider_id: str, model_id: str): summary="Update model settings", description="Enable/disable a model or set default parameters", ) +@serialized_config_mutation async def update_model_settings( provider_id: str, model_id: str, body: UpdateModelSettingRequest ) -> ModelSetting: @@ -522,6 +524,7 @@ async def update_model_settings( summary="Delete model definition", description="Delete a model from a provider (removes from flocks.json and runtime)", ) +@serialized_config_mutation async def delete_model_definition( provider_id: str, model_id: str, ): diff --git a/flocks/server/routes/onboarding.py b/flocks/server/routes/onboarding.py index f52b0a376..49e1968dc 100644 --- a/flocks/server/routes/onboarding.py +++ b/flocks/server/routes/onboarding.py @@ -11,7 +11,7 @@ from contextlib import asynccontextmanager from typing import Any, Dict, List, Optional -from fastapi import APIRouter, HTTPException +from fastapi import APIRouter, Depends, HTTPException from pydantic import BaseModel, Field from flocks.config.config import Config @@ -21,6 +21,7 @@ from flocks.provider.provider import Provider from flocks.provider.types import ModelType from flocks.security import get_secret_manager +from flocks.server.config_mutation import serialized_config_mutation from flocks.server.routes.default_model import ( SetDefaultModelRequest, set_default_model, @@ -28,16 +29,22 @@ from flocks.server.routes.mcp import ( McpCredentialRequest, McpTestRequest, + _load_raw_mcp_server_config, + _restore_threatbook_mcp_setup, connect_mcp_server, get_mcp_credentials, set_mcp_credentials, test_mcp_connection, ) +from flocks.server.auth import require_admin from flocks.server.routes.provider import ( APIServiceUpdateRequest, ProviderCredentialRequest, TestCredentialRequest, + _get_api_service_secret_candidates, _get_inline_provider_api_key, + _load_api_service_metadata_data, + _test_provider_credentials_impl, get_service_credentials, set_provider_credentials, set_service_credentials, @@ -381,13 +388,37 @@ async def _temporary_config_and_secret_state(): @asynccontextmanager -async def _rollback_on_apply_failure(): +async def _rollback_on_apply_failure(mcp_name: Optional[str] = None): """Rollback onboarding writes if apply fails part-way through.""" config_snapshot, secret_snapshot = _snapshot_config_and_secret_state() + previous_mcp_config = ( + copy.deepcopy(_load_raw_mcp_server_config(mcp_name)) if mcp_name else None + ) + was_mcp_connected = False + if mcp_name: + try: + runtime_status = await MCP.status() + previous_status = runtime_status.get(mcp_name) + was_mcp_connected = bool( + previous_status is not None + and previous_status.status == McpStatus.CONNECTED + ) + except Exception: + was_mcp_connected = False try: yield except Exception: - await _restore_config_and_secret_state(config_snapshot, secret_snapshot) + if mcp_name: + await _restore_threatbook_mcp_setup( + mcp_name, + config_snapshot, + secret_snapshot, + previous_mcp_config, + was_mcp_connected, + ) + await _reload_runtime_state() + else: + await _restore_config_and_secret_state(config_snapshot, secret_snapshot) raise @@ -400,17 +431,47 @@ async def _test_provider_or_service_with_temp_credentials( provider_name: Optional[str] = None, service: bool = False, ) -> Dict[str, Any]: + from flocks.tool.credential_context import activate_credential_overrides + + body = TestCredentialRequest(model_id=model_id) if model_id else None + if service: + raw_service = ConfigWriter.get_api_service_raw(provider_id) or {} + metadata = _load_api_service_metadata_data(provider_id) or {} + secret_ids = _get_api_service_secret_candidates( + provider_id, + raw_service, + field_name="api_key", + ) + primary_secret_id = secret_ids[0] + config_override = { + **raw_service, + "apiKey": f"{{secret:{primary_secret_id}}}", + "enabled": True, + } + async with activate_credential_overrides( + secret_values={secret_id: api_key for secret_id in secret_ids}, + service_id=provider_id, + config_values=config_override, + ): + return await test_provider_credentials(provider_id, body) + + Provider._ensure_initialized() + if Provider.get(provider_id) is not None: + return await _test_provider_credentials_impl( + provider_id, + body, + api_key_override=api_key, + isolated_provider=True, + base_url_override=base_url, + ) + async with _temporary_config_and_secret_state(): request = ProviderCredentialRequest( api_key=api_key, base_url=base_url, provider_name=provider_name, ) - if service: - await set_service_credentials(provider_id, request) - else: - await set_provider_credentials(provider_id, request) - body = TestCredentialRequest(model_id=model_id) if model_id else None + await set_provider_credentials(provider_id, request) return await test_provider_credentials(provider_id, body) @@ -814,7 +875,11 @@ async def get_onboarding_status() -> OnboardingStatusResponse: summary="Validate onboarding configuration", description="Validate ThreatBook and/or third-party model configuration for onboarding.", ) -async def validate_onboarding(request: OnboardingValidateRequest) -> OnboardingValidateResponse: +@serialized_config_mutation +async def validate_onboarding( + request: OnboardingValidateRequest, + _admin: object = Depends(require_admin), +) -> OnboardingValidateResponse: return await _validate_onboarding_request(request) @@ -824,7 +889,15 @@ async def validate_onboarding(request: OnboardingValidateRequest) -> OnboardingV summary="Apply onboarding configuration", description="Persist onboarding configuration after validation succeeds.", ) -async def apply_onboarding(request: OnboardingValidateRequest) -> OnboardingApplyResponse: +@serialized_config_mutation +async def apply_onboarding( + request: OnboardingValidateRequest, + _admin: object = Depends(require_admin), +) -> OnboardingApplyResponse: + return await _apply_onboarding(request) + + +async def _apply_onboarding(request: OnboardingValidateRequest) -> OnboardingApplyResponse: validation = await _validate_onboarding_request(request) if not validation.can_apply: raise HTTPException(status_code=400, detail=validation.message or "Validation failed") @@ -837,7 +910,9 @@ async def apply_onboarding(request: OnboardingValidateRequest) -> OnboardingAppl default_model: Optional[Dict[str, str]] = None try: - async with _rollback_on_apply_failure(): + async with _rollback_on_apply_failure( + preset["threatbook_mcp_name"] if preset["requires_mcp"] else None + ): if threatbook_api_key: if request.use_threatbook_model: await set_provider_credentials( diff --git a/flocks/server/routes/provider.py b/flocks/server/routes/provider.py index b3dd043da..cdb375491 100644 --- a/flocks/server/routes/provider.py +++ b/flocks/server/routes/provider.py @@ -7,6 +7,7 @@ """ import asyncio +import copy import json import re import threading @@ -22,6 +23,7 @@ from flocks.provider.provider import Provider, ModelInfo as ProviderModelInfo from flocks.security.secrets import SecretManager from flocks.server.auth import get_request_ip, get_request_user_agent, require_admin +from flocks.server.config_mutation import serialized_config_mutation from flocks.config.config import Config from flocks.config.config_writer import ConfigWriter from flocks.storage.storage import Storage @@ -777,6 +779,7 @@ async def list_api_services_route(): summary="Update API service", description="Enable or disable an API service and all tools it exposes." ) +@serialized_config_mutation async def update_api_service_route( provider_id: str, request: Dict[str, Any] = Body(...), @@ -1247,6 +1250,10 @@ async def _write_api_service_status_cache(statuses: Dict[str, Any]) -> None: async def _save_api_service_status_if_configured(provider_id: str, response: Dict[str, Any]) -> None: """Persist API test status only when the provider is configured as an API service.""" + from flocks.tool.credential_context import is_temporary_credential_override_active + + if is_temporary_credential_override_active(): + return raw_service = ConfigWriter.get_api_service_raw(provider_id) if raw_service is None: return @@ -1395,6 +1402,7 @@ async def list_api_services() -> List[APIServiceSummary]: raise HTTPException(status_code=500, detail=str(e)) +@serialized_config_mutation async def update_api_service(provider_id: str, request: APIServiceUpdateRequest) -> APIServiceSummary: try: from flocks.tool.registry import ToolRegistry @@ -1498,6 +1506,7 @@ def _find_user_installed_tool_plugin_for(storage_key: str) -> Optional[tuple[str return None +@serialized_config_mutation async def delete_api_service(provider_id: str) -> Dict[str, Any]: """Delete an API service configuration and its stored credential. @@ -1796,6 +1805,7 @@ async def reveal_provider_credentials( summary="Set provider credentials", description="Set authentication credentials for a provider or API service." ) +@serialized_config_mutation async def set_provider_credentials( provider_id: str, request: ProviderCredentialRequest, @@ -1996,6 +2006,7 @@ async def set_provider_credentials( summary="Delete provider credentials", description="Delete stored credentials for a provider or API service." ) +@serialized_config_mutation async def delete_provider_credentials( provider_id: str, _admin: object = Depends(require_admin), @@ -2139,7 +2150,9 @@ async def get_service_credentials( username=field_values.get("username"), fields=safe_field_values or None, secret_ids=secret_ids or None, - has_credential=bool(any(value for value in field_values.values())), + has_credential=bool( + any(field_values.get(field_name) for field_name in sensitive_field_names) + ), ) except Exception as e: log.error("service.credentials.get.error", {"provider_id": provider_id, "error": str(e)}) @@ -2233,6 +2246,7 @@ async def reveal_service_credentials( summary="Set API service credentials", description="Set authentication credentials for an API service." ) +@serialized_config_mutation async def set_service_credentials( provider_id: str, request: ProviderCredentialRequest, @@ -2380,6 +2394,63 @@ async def set_service_credentials( raise HTTPException(status_code=500, detail=str(e)) +@router.post( + "/{provider_id}/service-credentials/configure", + response_model=Dict[str, Any], + summary="Validate and save API service credentials", + description="Validate credentials in an isolated request context and persist them only on success.", +) +@serialized_config_mutation +async def configure_service_credentials( + provider_id: str, + request: ProviderCredentialRequest, + _admin: object = Depends(require_admin), +): + """Test a service key without exposing it to other requests before saving.""" + if provider_id not in {"threatbook-cn", "threatbook-io"}: + raise HTTPException( + status_code=400, + detail="This configure flow is only available for ThreatBook API services", + ) + api_key = (request.api_key or "").strip() + if not api_key: + raise HTTPException(status_code=400, detail="API key required") + + raw_service = ConfigWriter.get_api_service_raw(provider_id) or {} + secret_ids = _get_api_service_secret_candidates( + provider_id, + raw_service, + field_name="api_key", + ) + if request.secret_id: + secret_ids.insert(0, request.secret_id) + secret_ids = list(dict.fromkeys(secret_ids)) + primary_secret_id = secret_ids[0] + config_override = { + **raw_service, + "apiKey": f"{{secret:{primary_secret_id}}}", + "enabled": True, + } + + from flocks.tool.credential_context import activate_credential_overrides + + async with activate_credential_overrides( + secret_values={secret_id: api_key for secret_id in secret_ids}, + service_id=provider_id, + config_values=config_override, + ): + validation = await _test_provider_credentials_impl( + provider_id, + api_key_override=api_key, + ) + + if not validation.get("success"): + return validation + + await set_service_credentials(provider_id, request) + return validation + + class TestCredentialRequest(BaseModel): """Optional request body for test-credentials""" model_id: Optional[str] = Field(None, description="Model to test with (uses first available if omitted)") @@ -2395,6 +2466,17 @@ async def test_provider_credentials( provider_id: str, body: Optional[TestCredentialRequest] = None, _admin: object = Depends(require_admin), +): + return await _test_provider_credentials_impl(provider_id, body) + + +async def _test_provider_credentials_impl( + provider_id: str, + body: Optional[TestCredentialRequest] = None, + *, + api_key_override: Optional[str] = None, + isolated_provider: bool = False, + base_url_override: Optional[str] = None, ): """Test credentials for a provider or API service by making a real API call""" from flocks.security import get_secret_manager @@ -2405,28 +2487,30 @@ async def test_provider_credentials( # test-credentials handles both LLM providers and API services. # Try _llm_key first (LLM provider), then all secret fields defined # in the credential schema (api_key, password, token, etc.). - secrets = get_secret_manager() - secret_id = f"{provider_id}_llm_key" - api_key = secrets.get(secret_id) + api_key = api_key_override if not api_key: - raw_service = ConfigWriter.get_api_service_raw(provider_id) or {} - secret_id = None - metadata = _load_api_service_metadata_data(provider_id) or {} - secret_field_names = _get_api_service_secret_field_names(provider_id, metadata) - if not secret_field_names: - secret_field_names = ["api_key"] - for field_name in secret_field_names: - for candidate in _get_api_service_secret_candidates( - provider_id, raw_service, field_name=field_name - ): - api_key = secrets.get(candidate) + secrets = get_secret_manager() + secret_id = f"{provider_id}_llm_key" + api_key = secrets.get(secret_id) + if not api_key: + raw_service = ConfigWriter.get_api_service_raw(provider_id) or {} + secret_id = None + metadata = _load_api_service_metadata_data(provider_id) or {} + secret_field_names = _get_api_service_secret_field_names(provider_id, metadata) + if not secret_field_names: + secret_field_names = ["api_key"] + for field_name in secret_field_names: + for candidate in _get_api_service_secret_candidates( + provider_id, raw_service, field_name=field_name + ): + api_key = secrets.get(candidate) + if api_key: + secret_id = candidate + break if api_key: - secret_id = candidate break - if api_key: - break - if not api_key: - api_key = _get_inline_provider_api_key(provider_id) + if not api_key: + api_key = _get_inline_provider_api_key(provider_id) if not api_key: response = { @@ -2443,14 +2527,19 @@ async def test_provider_credentials( # _load_dynamic_providers skips already-registered providers, so it is safe # to call multiple times. await _load_dynamic_providers() - # Apply config to ensure _config_models (user-defined models) are loaded - config = await Config.get() - await Provider.apply_config(config, provider_id=provider_id) + # Temporary onboarding probes use a shallow provider clone so candidate + # credentials never enter the process-wide Provider registry. + if not isolated_provider: + config = await Config.get() + await Provider.apply_config(config, provider_id=provider_id) provider = Provider.get(provider_id) if provider: from flocks.provider.provider import ProviderConfig, ChatMessage as ProviderChatMessage + if isolated_provider: + provider = copy.copy(provider) + # Always reconfigure with the freshest key from secret manager # to avoid stale keys from cached config or prior apply_config. effective_base_url = None @@ -2466,7 +2555,7 @@ async def test_provider_credentials( provider.configure(ProviderConfig( provider_id=provider_id, api_key=api_key, - base_url=effective_base_url, + base_url=base_url_override or effective_base_url, custom_settings=custom_settings, )) if hasattr(provider, '_client'): @@ -2551,11 +2640,10 @@ async def test_provider_credentials( # Try to test connectivity by calling a simple tool from flocks.tool.registry import ToolRegistry, ToolCategory, ToolInfo from flocks.server.routes.tool import _get_tool_source + from flocks.tool.credential_context import activate_credential_probe await ToolRegistry.init_async() - _set_api_service_tools_enabled(provider_id, True) - # If the plugin ships a `_test.yaml` with a `connectivity` block, # honour the declared (tool, params) probe. Tool failures (e.g. # wrong apikey) are surfaced as-is — that's the answer the test @@ -2569,7 +2657,8 @@ async def test_provider_credentials( "service": provider_id, "tool": spec.tool, "params": spec.params, }) try: - probe = await ToolRegistry.execute(tool_name=spec.tool, **spec.params) + async with activate_credential_probe(): + probe = await ToolRegistry.execute(tool_name=spec.tool, **spec.params) latency = int((time.time() - start) * 1000) response = { "success": probe.success, @@ -2603,8 +2692,6 @@ async def test_provider_credentials( }) for tool_info in all_tools: - if not tool_info.enabled: - continue source, source_name = _get_tool_source(tool_info) if source in ("api", "device") and source_name == provider_id: service_tools.append(tool_info) @@ -2878,7 +2965,8 @@ def _build_param_sets(tool_info) -> list[dict[str, object]]: "single_attempt": True, }) - result = await ToolRegistry.execute(tool_name=test_tool.name, **test_params) + async with activate_credential_probe(): + result = await ToolRegistry.execute(tool_name=test_tool.name, **test_params) latency = int((time.time() - start) * 1000) if result.success: @@ -2960,6 +3048,10 @@ def _build_param_sets(tool_info) -> list[dict[str, object]]: async def _save_api_service_status(provider_id: str, test_response: dict) -> None: """Persist a single API service test result into the status cache.""" + from flocks.tool.credential_context import is_temporary_credential_override_active + + if is_temporary_credential_override_active(): + return try: await Storage.init() cached = await Storage.read(_API_SERVICE_STATUS_KEY) or {} diff --git a/flocks/server/routes/tool.py b/flocks/server/routes/tool.py index 0fb3ac4bc..2d8b62798 100644 --- a/flocks/server/routes/tool.py +++ b/flocks/server/routes/tool.py @@ -11,6 +11,7 @@ from pydantic import BaseModel, Field from flocks.server.auth import require_admin +from flocks.server.config_mutation import serialized_config_mutation from flocks.server.routes._timing import log_route_timing from flocks.utils.log import Log from flocks.config.config_writer import ConfigWriter @@ -868,6 +869,7 @@ async def get_tool(tool_name: str): response_model=ToolInfoResponse, summary="Update tool settings", ) +@serialized_config_mutation async def update_tool( tool_name: str, request: ToolUpdateRequest, @@ -962,6 +964,7 @@ async def update_tool( response_model=ToolInfoResponse, summary="Reset a tool to its YAML/registration default", ) +@serialized_config_mutation async def reset_tool_setting(tool_name: str, _admin: object = Depends(require_admin)): """Remove the user setting for ``tool_name`` and restore the default. diff --git a/flocks/tool/credential_context.py b/flocks/tool/credential_context.py index d4b0dfe1e..61bf67da7 100644 --- a/flocks/tool/credential_context.py +++ b/flocks/tool/credential_context.py @@ -58,6 +58,15 @@ "device_verify_ssl_override", default=None ) +# Credential validation may need to invoke a disabled tool without changing +# the process-wide ToolRegistry state or persisting a temporary status result. +_credential_probe_active: ContextVar[bool] = ContextVar( + "credential_probe_active", default=False +) +_temporary_credential_override_active: ContextVar[bool] = ContextVar( + "temporary_credential_override_active", default=False +) + # --------------------------------------------------------------------------- # Internal result type for _build_overrides @@ -121,6 +130,55 @@ def get_verify_ssl_override() -> Optional[bool]: return _verify_ssl_override.get() +def is_credential_probe_active() -> bool: + """Return whether the current coroutine is validating temporary credentials.""" + return _credential_probe_active.get() + + +def is_temporary_credential_override_active() -> bool: + """Return whether temporary credentials are active in this coroutine.""" + return _temporary_credential_override_active.get() + + +@asynccontextmanager +async def activate_credential_probe() -> AsyncIterator[None]: + """Allow a connectivity probe to invoke a disabled tool without enabling it.""" + token = _credential_probe_active.set(True) + try: + yield + finally: + _credential_probe_active.reset(token) + + +@asynccontextmanager +async def activate_credential_overrides( + *, + secret_values: Dict[str, str], + service_id: Optional[str] = None, + config_values: Optional[Dict[str, Any]] = None, +) -> AsyncIterator[None]: + """Temporarily override credentials for one coroutine without persisting them.""" + t1 = _secret_override.set(dict(secret_values)) + t2 = _config_override.set(dict(config_values or {})) + t3 = _config_override_service.set(service_id) + t4 = _config_override_storage_key.set(service_id) + t5 = _verify_ssl_override.set( + bool((config_values or {}).get("verify_ssl", False)) + ) + t6 = _credential_probe_active.set(True) + t7 = _temporary_credential_override_active.set(True) + try: + yield + finally: + _secret_override.reset(t1) + _config_override.reset(t2) + _config_override_service.reset(t3) + _config_override_storage_key.reset(t4) + _verify_ssl_override.reset(t5) + _credential_probe_active.reset(t6) + _temporary_credential_override_active.reset(t7) + + # --------------------------------------------------------------------------- # Activation (called from ToolRegistry.execute) # --------------------------------------------------------------------------- diff --git a/flocks/tool/registry.py b/flocks/tool/registry.py index 5b902c9cf..8c789cdec 100644 --- a/flocks/tool/registry.py +++ b/flocks/tool/registry.py @@ -1105,6 +1105,9 @@ async def execute( device_id = None per_device_enabled = None + from flocks.tool.credential_context import is_credential_probe_active + + credential_probe_active = is_credential_probe_active() if tool.info.source == "device" and tool.info.provider: requested_device_id = kwargs.pop("device_id", None) try: @@ -1125,13 +1128,13 @@ async def execute( if resolution_error: return ToolResult(success=False, error=resolution_error) device_id = resolved_device_id - elif not tool.info.enabled: + elif not tool.info.enabled and not credential_probe_active: return ToolResult( success=False, error=f"Tool is disabled: {tool_name}" ) - if not tool.info.enabled: + if not tool.info.enabled and not credential_probe_active: return ToolResult( success=False, error=f"Tool is disabled: {tool_name}" @@ -1170,21 +1173,22 @@ async def execute( else: result = await tool.execute(ctx, **kwargs) - if result.success: - cls._reset_failure_state(tool_name) - else: - if await cls._failure_auto_disable_enabled(): - disabled = cls._record_failure(tool, kwargs, result.error) - else: + if not credential_probe_active: + if result.success: 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" - if result.error: - result.error = f"{result.error} ({suffix})" + else: + if await cls._failure_auto_disable_enabled(): + disabled = cls._record_failure(tool, kwargs, result.error) else: - result.error = suffix + 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" + if result.error: + result.error = f"{result.error} ({suffix})" + else: + result.error = suffix return result @classmethod diff --git a/tests/provider/test_api_service_credentials.py b/tests/provider/test_api_service_credentials.py index e42620b24..f6afdeb69 100644 --- a/tests/provider/test_api_service_credentials.py +++ b/tests/provider/test_api_service_credentials.py @@ -1,4 +1,5 @@ -from unittest.mock import MagicMock, patch +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -6,6 +7,130 @@ class TestAPIServiceCredentials: + @pytest.mark.asyncio + async def test_isolated_provider_probe_does_not_mutate_registered_instance(self): + from flocks.provider.provider import ProviderConfig + from flocks.server.routes import provider as provider_routes + + class FakeProvider: + def __init__(self): + self._config = ProviderConfig( + provider_id="openai", + api_key="saved-key", + base_url="https://saved.example/v1", + ) + self._base_url = "https://default.example/v1" + self._client = object() + + def configure(self, config): + self._config = config + self._client = None + + async def chat(self, *_args, **_kwargs): + return SimpleNamespace(content="Paris") + + shared_provider = FakeProvider() + original_client = shared_provider._client + + with ( + patch.object(provider_routes, "_ensure_provider_initialized", AsyncMock()), + patch.object(provider_routes, "_load_dynamic_providers", AsyncMock()), + patch.object(provider_routes.Provider, "get", return_value=shared_provider), + patch.object( + provider_routes.Provider, + "list_models", + return_value=[SimpleNamespace(id="gpt-test")], + ), + patch.object(provider_routes.Config, "get", side_effect=AssertionError("must not load config")), + ): + result = await provider_routes._test_provider_credentials_impl( + "openai", + provider_routes.TestCredentialRequest(model_id="gpt-test"), + api_key_override="candidate-key", + isolated_provider=True, + base_url_override="https://candidate.example/v1", + ) + + assert result["success"] is True + assert shared_provider._config.api_key == "saved-key" + assert shared_provider._config.base_url == "https://saved.example/v1" + assert shared_provider._client is original_client + + @pytest.mark.asyncio + async def test_temporary_api_probe_does_not_enable_tools_or_persist_status(self): + from flocks.server.routes import provider as provider_routes + from flocks.tool.credential_context import activate_credential_overrides + from flocks.tool.registry import ToolCategory, ToolInfo, ToolResult + + tool_info = ToolInfo( + name="threatbook_cn_probe", + description="Read-only connectivity probe", + category=ToolCategory.CUSTOM, + parameters=[], + enabled=False, + ) + + with ( + patch.object(provider_routes, "_ensure_provider_initialized", AsyncMock()), + patch.object(provider_routes, "_load_dynamic_providers", AsyncMock()), + patch.object(provider_routes.Provider, "get", return_value=None), + patch("flocks.tool.registry.ToolRegistry.init_async", AsyncMock()), + patch("flocks.tool.registry.ToolRegistry.list_tools", return_value=[tool_info]), + patch( + "flocks.tool.registry.ToolRegistry.get_dynamic_tools_by_module", + return_value={}, + ), + patch( + "flocks.tool.registry.ToolRegistry.execute", + AsyncMock(return_value=ToolResult(success=False, error="invalid key")), + ) as execute, + patch( + "flocks.server.routes.tool._get_tool_source", + return_value=("api", "threatbook-cn"), + ), + patch( + "flocks.tool.probe_loader.get_connectivity_spec", + return_value=None, + ), + patch.object(provider_routes, "_set_api_service_tools_enabled") as set_enabled, + patch.object(provider_routes.Storage, "write", AsyncMock()) as write_status, + ): + async with activate_credential_overrides( + secret_values={"threatbook_cn_api_key": "candidate-key"}, + service_id="threatbook-cn", + config_values={"enabled": True}, + ): + result = await provider_routes._test_provider_credentials_impl( + "threatbook-cn", + api_key_override="candidate-key", + ) + + assert result["success"] is False + execute.assert_awaited_once() + set_enabled.assert_not_called() + write_status.assert_not_awaited() + assert tool_info.enabled is False + + @pytest.mark.asyncio + async def test_base_url_alone_is_not_a_configured_credential(self): + from flocks.server.routes.provider import get_service_credentials + + mock_secrets = MagicMock() + mock_secrets.get.return_value = None + + with ( + patch("flocks.security.get_secret_manager", return_value=mock_secrets), + patch( + "flocks.config.config_writer.ConfigWriter.get_api_service_raw", + return_value={"base_url": "https://api.threatbook.cn"}, + ), + ): + result = await get_service_credentials("threatbook-cn") + + assert result.base_url == "https://api.threatbook.cn" + assert result.api_key_masked is None + assert result.has_credential is False + @pytest.mark.asyncio async def test_get_service_credentials_returns_base_url_and_username(self): from flocks.server.routes.provider import get_service_credentials @@ -69,6 +194,58 @@ async def test_set_service_credentials_persists_api_key_base_url_and_username(se ) assert result["success"] is True + @pytest.mark.asyncio + async def test_configure_service_credentials_does_not_save_failed_key(self): + from flocks.server.routes import provider as provider_routes + + validation = {"success": False, "message": "invalid key"} + test_credentials = AsyncMock(return_value=validation) + save_credentials = AsyncMock() + + with ( + patch.object(provider_routes, "_test_provider_credentials_impl", test_credentials), + patch.object(provider_routes, "set_service_credentials", save_credentials), + patch.object( + provider_routes.ConfigWriter, + "get_api_service_raw", + return_value={"enabled": True}, + ), + ): + result = await provider_routes.configure_service_credentials( + "threatbook-cn", + provider_routes.ProviderCredentialRequest(api_key="bad-key"), + ) + + assert result == validation + save_credentials.assert_not_awaited() + + @pytest.mark.asyncio + async def test_configure_service_credentials_saves_key_after_validation(self): + from flocks.server.routes import provider as provider_routes + + validation = {"success": True, "message": "connected", "latency_ms": 12} + test_credentials = AsyncMock(return_value=validation) + save_credentials = AsyncMock(return_value={"success": True}) + request = provider_routes.ProviderCredentialRequest(api_key="valid-key") + + with ( + patch.object(provider_routes, "_test_provider_credentials_impl", test_credentials), + patch.object(provider_routes, "set_service_credentials", save_credentials), + patch.object( + provider_routes.ConfigWriter, + "get_api_service_raw", + return_value={"enabled": True}, + ), + ): + result = await provider_routes.configure_service_credentials( + "threatbook-cn", + request, + ) + + assert result == validation + test_credentials.assert_awaited_once() + save_credentials.assert_awaited_once_with("threatbook-cn", request) + @pytest.mark.asyncio async def test_set_service_credentials_uses_metadata_secret_for_hyphenated_service(self): from flocks.server.routes.provider import ProviderCredentialRequest, set_service_credentials diff --git a/tests/provider/test_test_credentials.py b/tests/provider/test_test_credentials.py index e1687bf96..fdbf4e364 100644 --- a/tests/provider/test_test_credentials.py +++ b/tests/provider/test_test_credentials.py @@ -74,6 +74,7 @@ async def test_no_tools_returns_failure(self): # Setup: no tools match mock_tr.init = MagicMock() + mock_tr.init_async = AsyncMock() mock_tr.list_tools.return_value = [] mock_tr._dynamic_tools_by_module = {} @@ -108,6 +109,7 @@ async def test_tool_failure_returns_failure(self): mock_provider_cls.get.return_value = None mock_tr.init = MagicMock() + mock_tr.init_async = AsyncMock() mock_tr.list_tools.return_value = [tool_info] mock_tr._dynamic_tools_by_module = { "flocks.tool.generated.threatbook": ["threatbook_ip_query"], @@ -144,6 +146,7 @@ async def test_tool_success_returns_success(self): mock_provider_cls.get.return_value = None mock_tr.init = MagicMock() + mock_tr.init_async = AsyncMock() mock_tr.list_tools.return_value = [tool_info] mock_tr._dynamic_tools_by_module = { "flocks.tool.generated.threatbook": ["threatbook_ip_query"], @@ -205,6 +208,7 @@ async def test_service_prefers_lightweight_query_tool_over_file_upload(self): mock_provider_cls.get.return_value = None mock_tr.init = MagicMock() + mock_tr.init_async = AsyncMock() # Put upload first on purpose to prove sorting is stable. mock_tr.list_tools.return_value = [upload_tool, url_tool] mock_tr._dynamic_tools_by_module = { @@ -283,6 +287,7 @@ async def test_onesec_service_prefers_threat_probe_and_uses_enum_action(self): mock_provider_cls.get.return_value = None mock_tr.init = MagicMock() + mock_tr.init_async = AsyncMock() mock_tr.list_tools.return_value = [onesec_dns_tool, onesec_threat_tool] mock_tr._dynamic_tools_by_module = { "flocks.tool.generated.onesec": ["onesec_dns", "onesec_threat"], @@ -392,6 +397,7 @@ async def test_declared_manifest_probe_is_used_before_heuristic_tool_selection(s mock_provider_cls.get.return_value = None mock_tr.init = MagicMock() + mock_tr.init_async = AsyncMock() mock_tr.list_tools.return_value = [heuristic_tool] mock_tr._dynamic_tools_by_module = { "flocks.tool.generated.tdp_api": ["tdp_assets_domain_list"], @@ -465,6 +471,7 @@ async def test_login_probe_does_not_overmatch_business_tools(self): mock_provider_cls.get.return_value = None mock_tr.init = MagicMock() + mock_tr.init_async = AsyncMock() mock_tr.list_tools.return_value = [login_business_tool, ip_query_tool] mock_tr._dynamic_tools_by_module = { "flocks.tool.generated.tdp_api": [ @@ -535,6 +542,7 @@ async def test_single_attempt_only_to_avoid_account_lockout(self): mock_provider_cls.get.return_value = None mock_tr.init = MagicMock() + mock_tr.init_async = AsyncMock() mock_tr.list_tools.return_value = [assets_tool] mock_tr._dynamic_tools_by_module = { "flocks.tool.generated.qingteng": ["qingteng_assets"], @@ -627,6 +635,7 @@ async def test_action_dispatch_login_tool_uses_test_action(self): mock_provider_cls.get.return_value = None mock_tr.init = MagicMock() + mock_tr.init_async = AsyncMock() # Put the assets tool first to prove the sort ranking promotes # the `_login` action-dispatch tool above other groups. mock_tr.list_tools.return_value = [assets_tool, login_tool] @@ -677,6 +686,7 @@ async def test_service_metadata_secret_is_used_for_hyphenated_service(self): mock_provider_cls.get.return_value = None mock_tr.init = MagicMock() + mock_tr.init_async = AsyncMock() mock_tr.list_tools.return_value = [tool_info] mock_tr._dynamic_tools_by_module = { "flocks.tool.generated.threatbook_cn": ["threatbook_ip_query"], @@ -713,6 +723,7 @@ async def test_tool_exception_returns_failure(self): mock_provider_cls.get.return_value = None mock_tr.init = MagicMock() + mock_tr.init_async = AsyncMock() mock_tr.list_tools.return_value = [tool_info] mock_tr._dynamic_tools_by_module = { "flocks.tool.generated.threatbook": ["threatbook_ip_query"], diff --git a/tests/server/concurrency/test_config_mutation.py b/tests/server/concurrency/test_config_mutation.py new file mode 100644 index 000000000..c8ae61e43 --- /dev/null +++ b/tests/server/concurrency/test_config_mutation.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +import asyncio + +import pytest + +from flocks.server.config_mutation import serialized_config_mutation + + +@pytest.mark.asyncio +async def test_config_mutation_lock_is_reentrant_for_nested_routes(): + calls: list[str] = [] + + @serialized_config_mutation + async def inner() -> None: + calls.append("inner") + + @serialized_config_mutation + async def outer() -> None: + calls.append("outer") + await inner() + + await asyncio.wait_for(outer(), timeout=1) + + assert calls == ["outer", "inner"] + + +@pytest.mark.asyncio +async def test_config_mutation_lock_serializes_concurrent_tasks(): + first_entered = asyncio.Event() + release_first = asyncio.Event() + calls: list[str] = [] + + @serialized_config_mutation + async def mutate(name: str) -> None: + calls.append(f"{name}:start") + if name == "first": + first_entered.set() + await release_first.wait() + calls.append(f"{name}:end") + + first = asyncio.create_task(mutate("first")) + await first_entered.wait() + second = asyncio.create_task(mutate("second")) + await asyncio.sleep(0) + + assert calls == ["first:start"] + release_first.set() + await asyncio.gather(first, second) + + assert calls == [ + "first:start", + "first:end", + "second:start", + "second:end", + ] diff --git a/tests/server/routes/test_global_mutation_auth.py b/tests/server/routes/test_global_mutation_auth.py index 4f526516a..8ec9c425c 100644 --- a/tests/server/routes/test_global_mutation_auth.py +++ b/tests/server/routes/test_global_mutation_auth.py @@ -46,6 +46,7 @@ def test_provider_global_configuration_credentials_and_tests_require_admin(): client.get("/api/provider/example/service-credentials"), client.post("/api/provider/example/service-credentials/reveal"), client.post("/api/provider/example/service-credentials", json={"api_key": "secret"}), + client.post("/api/provider/example/service-credentials/configure", json={"api_key": "secret"}), client.post("/api/provider/example/test-credentials", json={}), client.patch("/api/provider/api-services/example", json={"enabled": False}), client.delete("/api/provider/api-services/example"), @@ -55,6 +56,51 @@ def test_provider_global_configuration_credentials_and_tests_require_admin(): assert [response.status_code for response in responses] == [403] * len(responses) +def test_onboarding_configuration_requires_admin(): + from flocks.server.routes.onboarding import router + + client = _user_client(router, prefix="/api/onboarding", role="member") + payload = { + "region": "cn", + "use_threatbook_model": True, + "threatbook_api_key": "secret", + } + + assert client.post("/api/onboarding/validate", json=payload).status_code == 403 + assert client.post("/api/onboarding/apply", json=payload).status_code == 403 + + +def test_mcp_credential_and_threatbook_configuration_require_admin(): + from flocks.server.routes.mcp import router + + client = _user_client(router, prefix="/api/mcp", role="member") + responses = [ + client.post("/api/mcp", json={}), + client.post("/api/mcp/test", json={}), + client.put("/api/mcp/threatbook_mcp", json={}), + client.delete("/api/mcp/threatbook_mcp"), + client.post("/api/mcp/threatbook_mcp/test", json={}), + client.post("/api/mcp/threatbook_mcp/connect"), + client.post("/api/mcp/threatbook_mcp/disconnect"), + client.post("/api/mcp/threatbook_mcp/auth"), + client.delete("/api/mcp/threatbook_mcp/auth"), + client.post("/api/mcp/threatbook_mcp/refresh"), + client.get("/api/mcp/threatbook_mcp/credentials"), + client.post("/api/mcp/threatbook_mcp/credentials/reveal"), + client.post("/api/mcp/threatbook_mcp/credentials", json={"api_key": "secret"}), + client.delete("/api/mcp/threatbook_mcp/credentials"), + client.post("/api/mcp/threatbook_mcp/test-credentials"), + client.post( + "/api/mcp/threatbook_mcp/threatbook-configure", + json={"region": "cn", "api_key": "secret"}, + ), + client.post("/api/mcp/catalog/auto-setup"), + client.post("/api/mcp/catalog/install", json={}), + ] + + assert [response.status_code for response in responses] == [403] * len(responses) + + def test_admin_reveal_returns_raw_llm_key_with_audit_and_no_store( monkeypatch: pytest.MonkeyPatch, ): diff --git a/tests/server/routes/test_mcp_routes.py b/tests/server/routes/test_mcp_routes.py index 6d3bb59a1..282b2e201 100644 --- a/tests/server/routes/test_mcp_routes.py +++ b/tests/server/routes/test_mcp_routes.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +from unittest.mock import AsyncMock import pytest from httpx import AsyncClient @@ -282,7 +283,7 @@ async def fake_config_get(cls): "get_mcp_server", lambda name: { "type": "remote", - "url": "https://example.com/mcp", + "url": "https://example.com/mcp?apikey=url-token&mode=full", "auth": { "type": "apikey", "location": "header", @@ -300,6 +301,8 @@ async def fake_config_get(cls): assert resp.status_code == 200, resp.text data = resp.json() + assert data["config"]["url"] == "https://example.com/mcp?apikey=***&mode=full" + assert "url-token" not in resp.text assert data["config"]["auth"]["value"] == "***" assert data["config"]["headers"]["Authorization"] == "***" assert data["config"]["headers"]["X-Client"] == "flocks" @@ -721,6 +724,59 @@ def set(self, key: str, value: str) -> None: ) assert stored_configs["demo-mcp"]["headers"]["X-Client"] == "flocks-web" + @pytest.mark.asyncio + async def test_update_mcp_server_restores_masked_url_query_secret( + self, client: AsyncClient, monkeypatch: pytest.MonkeyPatch + ): + stored_configs: dict[str, dict] = {} + saved_secrets: dict[str, str] = {} + + async def fake_status() -> dict[str, McpStatusInfo]: + return {} + + monkeypatch.setattr(mcp_routes.MCP, "status", fake_status) + monkeypatch.setattr( + mcp_routes.ConfigWriter, + "get_mcp_server", + lambda name: { + "type": "remote", + "url": "https://old.example.com/mcp?apikey=url-token&mode=full", + "enabled": False, + }, + ) + monkeypatch.setattr( + mcp_routes.ConfigWriter, + "add_mcp_server", + lambda name, config: stored_configs.__setitem__(name, config), + ) + monkeypatch.setattr(tool_loader, "save_mcp_config", lambda name, config: None) + + class SecretManagerStub: + def set(self, key: str, value: str) -> None: + saved_secrets[key] = value + + monkeypatch.setattr( + "flocks.security.get_secret_manager", + lambda: SecretManagerStub(), + ) + + resp = await client.put( + "/api/mcp/demo-mcp", + json={ + "config": { + "url": "https://new.example.com/mcp?apikey=***&mode=compact", + "enabled": False, + } + }, + ) + + assert resp.status_code == 200, resp.text + assert saved_secrets == {"demo-mcp_mcp_key": "url-token"} + assert ( + stored_configs["demo-mcp"]["url"] + == "https://new.example.com/mcp?apikey={secret:demo-mcp_mcp_key}&mode=compact" + ) + # --------------------------------------------------------------------- # should_reconnect: the contract for ``PUT /api/mcp/{name}`` is that # any save where the new config asks the server to be enabled AND the @@ -1217,12 +1273,48 @@ def get(self, secret_id: str): return "312abcdef321" if secret_id == "threatbook_mcp_key" else None monkeypatch.setattr(mcp_routes, "get_secret_manager", lambda: FakeSecrets()) + audit = AsyncMock() + monkeypatch.setattr(mcp_routes, "emit_audit_event", audit) resp = await client.post("/api/mcp/threatbook_mcp/credentials/reveal") assert resp.status_code == 200, resp.text assert resp.json() == {"api_key": "312abcdef321"} assert resp.headers["cache-control"] == "no-store" + audit.assert_awaited_once() + event_type, payload = audit.await_args.args + assert event_type == "mcp.credentials_reveal" + assert payload["mcp_name"] == "threatbook_mcp" + assert "312abcdef321" not in repr(payload) + + @pytest.mark.asyncio + async def test_mcp_credentials_support_historical_inline_url_key( + self, client: AsyncClient, monkeypatch: pytest.MonkeyPatch + ): + class FakeSecrets: + def get(self, secret_id: str): + return None + + monkeypatch.setattr(mcp_routes, "get_secret_manager", lambda: FakeSecrets()) + monkeypatch.setattr( + mcp_routes, + "_load_raw_mcp_server_config", + lambda _name: { + "type": "remote", + "url": "https://mcp.threatbook.cn/mcp?apikey=historical-key", + }, + ) + monkeypatch.setattr(mcp_routes, "emit_audit_event", AsyncMock()) + + masked = await client.get("/api/mcp/threatbook_mcp/credentials") + revealed = await client.post("/api/mcp/threatbook_mcp/credentials/reveal") + + assert masked.status_code == 200, masked.text + assert masked.json()["has_credential"] is True + assert masked.json()["api_key_masked"] == "************" + assert "historical-key" not in masked.text + assert revealed.status_code == 200, revealed.text + assert revealed.json() == {"api_key": "historical-key"} @pytest.mark.asyncio async def test_reveal_mcp_credentials_returns_not_found_when_missing( diff --git a/tests/server/routes/test_onboarding_routes.py b/tests/server/routes/test_onboarding_routes.py index 23cd49c12..88ece4903 100644 --- a/tests/server/routes/test_onboarding_routes.py +++ b/tests/server/routes/test_onboarding_routes.py @@ -1,5 +1,8 @@ from __future__ import annotations +from types import SimpleNamespace +from unittest.mock import AsyncMock + import pytest from flocks.server.routes import onboarding as onboarding_routes @@ -35,6 +38,85 @@ async def fake_status(): monkeypatch.setattr(onboarding_routes, "_build_threatbook_intel_status", fake_status) +@pytest.mark.asyncio +async def test_apply_rollback_restores_mcp_runtime(monkeypatch: pytest.MonkeyPatch): + config_snapshot = {"mcp": {"threatbook_mcp": {"url": "old"}}} + secret_snapshot = {"threatbook_mcp_key": "old-key"} + previous_config = {"type": "remote", "url": "old"} + restore = AsyncMock() + reload_runtime = AsyncMock() + + monkeypatch.setattr( + onboarding_routes, + "_snapshot_config_and_secret_state", + lambda: (config_snapshot, secret_snapshot), + ) + monkeypatch.setattr( + onboarding_routes, + "_load_raw_mcp_server_config", + lambda _name: previous_config, + ) + monkeypatch.setattr( + onboarding_routes.MCP, + "status", + AsyncMock( + return_value={ + "threatbook_mcp": SimpleNamespace( + status=onboarding_routes.McpStatus.CONNECTED + ), + } + ), + ) + monkeypatch.setattr(onboarding_routes, "_restore_threatbook_mcp_setup", restore) + monkeypatch.setattr(onboarding_routes, "_reload_runtime_state", reload_runtime) + + with pytest.raises(RuntimeError, match="apply failed"): + async with onboarding_routes._rollback_on_apply_failure("threatbook_mcp"): + raise RuntimeError("apply failed") + + restore.assert_awaited_once_with( + "threatbook_mcp", + config_snapshot, + secret_snapshot, + previous_config, + True, + ) + reload_runtime.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_registered_provider_validation_uses_isolated_instance( + monkeypatch: pytest.MonkeyPatch, +): + registered_provider = object() + probe = AsyncMock(return_value={"success": True}) + + monkeypatch.setattr(onboarding_routes.Provider, "_ensure_initialized", lambda: None) + monkeypatch.setattr( + onboarding_routes.Provider, + "get", + lambda provider_id: registered_provider if provider_id == "openai" else None, + ) + monkeypatch.setattr(onboarding_routes, "_test_provider_credentials_impl", probe) + + result = await onboarding_routes._test_provider_or_service_with_temp_credentials( + "openai", + "candidate-key", + model_id="gpt-4o", + base_url="https://example.test/v1", + ) + + assert result == {"success": True} + probe.assert_awaited_once() + _, body = probe.await_args.args + assert body.model_id == "gpt-4o" + assert probe.await_args.kwargs == { + "api_key_override": "candidate-key", + "isolated_provider": True, + "base_url_override": "https://example.test/v1", + } + + class TestOnboardingStatusRoutes: @pytest.mark.asyncio async def test_status_incomplete_when_no_default_model(self, client, monkeypatch: pytest.MonkeyPatch): diff --git a/tests/server/routes/test_tool_routes.py b/tests/server/routes/test_tool_routes.py index 036e8f68c..e7b0e419f 100644 --- a/tests/server/routes/test_tool_routes.py +++ b/tests/server/routes/test_tool_routes.py @@ -270,6 +270,7 @@ async def test_execute_rejects_message_outside_session( from flocks.server.routes import tool as tool_routes permission_ask = AsyncMock(return_value=None) + monkeypatch.setattr(tool_routes, "legacy_tool_permission_prompt_required", lambda: True) monkeypatch.setattr(tool_routes.PermissionNext, "ask", permission_ask) session_id, _ = await _create_session_and_message("owner-session") @@ -317,6 +318,7 @@ async def test_execute_uses_permission_flow_when_session_context_is_present( from flocks.server.routes import tool as tool_routes permission_ask = AsyncMock(return_value=None) + monkeypatch.setattr(tool_routes, "legacy_tool_permission_prompt_required", lambda: True) monkeypatch.setattr(tool_routes.PermissionNext, "ask", permission_ask) session_id, message_id = await _create_session_and_message("valid-session-context") @@ -367,6 +369,7 @@ async def test_batch_uses_actual_child_tool_name_for_permission_flow( from flocks.server.routes import tool as tool_routes permission_ask = AsyncMock(return_value=None) + monkeypatch.setattr(tool_routes, "legacy_tool_permission_prompt_required", lambda: True) monkeypatch.setattr(tool_routes.PermissionNext, "ask", permission_ask) session_id, message_id = await _create_session_and_message("valid-batch-session-context") diff --git a/tests/tool/test_credential_context_config_override.py b/tests/tool/test_credential_context_config_override.py index 20c89e89e..8054c3078 100644 --- a/tests/tool/test_credential_context_config_override.py +++ b/tests/tool/test_credential_context_config_override.py @@ -16,8 +16,11 @@ _config_override, _config_override_service, _config_override_storage_key, + activate_credential_overrides, activate_device_credentials, get_config_override, + is_credential_probe_active, + is_temporary_credential_override_active, ) _SAMPLE_CONFIG = {"base_url": "https://10.201.255.17", "enabled": True} @@ -139,6 +142,23 @@ def test_identical_service_and_storage_key(): assert get_config_override("other") is None +@pytest.mark.asyncio +async def test_temporary_credentials_scope_probe_state_to_current_coroutine(): + assert is_credential_probe_active() is False + assert is_temporary_credential_override_active() is False + + async with activate_credential_overrides( + secret_values={"demo_api_key": "candidate"}, + service_id="demo", + config_values={"enabled": True}, + ): + assert is_credential_probe_active() is True + assert is_temporary_credential_override_active() is True + + assert is_credential_probe_active() is False + assert is_temporary_credential_override_active() is False + + @pytest.mark.asyncio async def test_activate_preserves_legacy_fields_not_in_current_schema(monkeypatch): """Old device rows can keep using fields removed from a newer schema.""" diff --git a/tests/tool/test_failure_auto_disable_config.py b/tests/tool/test_failure_auto_disable_config.py index b54c654cb..08e16b1cb 100644 --- a/tests/tool/test_failure_auto_disable_config.py +++ b/tests/tool/test_failure_auto_disable_config.py @@ -7,6 +7,7 @@ from flocks.config.config import Config, ConfigInfo, ToolFailureConfig from flocks.config.config_writer import ConfigWriter +from flocks.tool.credential_context import activate_credential_probe from flocks.tool.registry import ( Tool, ToolCategory, @@ -128,3 +129,24 @@ async def test_config_can_turn_off_repeated_failure_auto_disable( assert "disabled" not in result.metadata assert tool.info.enabled is True assert ToolRegistry._failure_state == {} + + +@pytest.mark.asyncio +async def test_credential_probe_does_not_change_failure_tracking( + isolated_failure_tracking: Tool, +) -> None: + tool = isolated_failure_tracking + ToolRegistry._failure_state[tool.info.name] = { + "key": "existing", + "count": 1, + } + + async with activate_credential_probe(): + result = await ToolRegistry.execute(tool.info.name, query="candidate") + + assert result.success is False + assert result.metadata == {} + assert tool.info.enabled is True + assert ToolRegistry._failure_state == { + tool.info.name: {"key": "existing", "count": 1} + } diff --git a/webui/src/api/provider.test.ts b/webui/src/api/provider.test.ts index afbc51ca5..033a2726c 100644 --- a/webui/src/api/provider.test.ts +++ b/webui/src/api/provider.test.ts @@ -160,4 +160,18 @@ describe('providerAPI.revealCredentials', () => { expect(mockPost).toHaveBeenCalledWith('/api/provider/threatbook-cn/service-credentials/reveal'); }); + + it('validates and saves API service credentials through one request', async () => { + mockPost.mockResolvedValue({ data: { success: true, message: 'ok' } }); + + const { providerAPI } = await import('./provider'); + await providerAPI.configureServiceCredentials('threatbook-cn', { + api_key: 'new-key', + }); + + expect(mockPost).toHaveBeenCalledWith( + '/api/provider/threatbook-cn/service-credentials/configure', + { api_key: 'new-key' }, + ); + }); }); diff --git a/webui/src/api/provider.ts b/webui/src/api/provider.ts index cde35cd9d..56addd533 100644 --- a/webui/src/api/provider.ts +++ b/webui/src/api/provider.ts @@ -132,6 +132,17 @@ export const providerAPI = { return response; }), + configureServiceCredentials: (id: string, credentials: ProviderCredentialInput) => + client.post<{ + success: boolean; message: string; latency_ms?: number; error?: string; + }>(`/api/provider/${id}/service-credentials/configure`, credentials) + .then((response) => { + if (response.data.success) { + invalidateApiServicesListCache(); + } + return response; + }), + deleteCredentials: (id: string) => client.delete<{ success: boolean }>(`/api/provider/${id}/credentials`), diff --git a/webui/src/components/layout/Layout.tsx b/webui/src/components/layout/Layout.tsx index 0a50be13d..fde3eb296 100644 --- a/webui/src/components/layout/Layout.tsx +++ b/webui/src/components/layout/Layout.tsx @@ -377,7 +377,7 @@ export default function Layout() { }, [collapsed, sidebarWidth, updateSidebarWidth]); useEffect(() => { - if (!isHome) return undefined; + if (!isHome || user?.role !== 'admin') return undefined; let cancelled = false; onboardingAPI.getStatus() @@ -395,9 +395,13 @@ export default function Layout() { return () => { cancelled = true; }; - }, [isHome]); + }, [isHome, user?.role]); - const handleOpenOnboarding = useCallback(() => setShowOnboarding(true), []); + const handleOpenOnboarding = useCallback(() => { + if (user?.role === 'admin') { + setShowOnboarding(true); + } + }, [user?.role]); useEffect(() => { window.addEventListener('flocks:open-onboarding', handleOpenOnboarding); diff --git a/webui/src/pages/Home/index.test.tsx b/webui/src/pages/Home/index.test.tsx index 9be0d9f2b..2283bbbd3 100644 --- a/webui/src/pages/Home/index.test.tsx +++ b/webui/src/pages/Home/index.test.tsx @@ -114,6 +114,9 @@ describe('Home create WebUI contract page entry', () => { expect( screen.queryByRole('button', { name: 'createWebUIContractPage' }), ).not.toBeInTheDocument(); + expect( + screen.queryByRole('button', { name: 'getStarted' }), + ).not.toBeInTheDocument(); expect(createMock).not.toHaveBeenCalled(); }); diff --git a/webui/src/pages/Home/index.tsx b/webui/src/pages/Home/index.tsx index 4694ca106..29677717d 100644 --- a/webui/src/pages/Home/index.tsx +++ b/webui/src/pages/Home/index.tsx @@ -144,15 +144,17 @@ export default function Home() {
- + {user?.role === 'admin' ? ( + + ) : null} {canCreateWebUIContractPage ? (
From 2ddf55c9766e7fb279362059824aa97c4a6b5d51 Mon Sep 17 00:00:00 2001 From: John Yin <10972267+john-yin2333@user.noreply.gitee.com> Date: Mon, 7 Sep 2026 11:36:57 +0800 Subject: [PATCH 47/63] Allow all users to access onboarding --- flocks/server/routes/onboarding.py | 5 +---- .../routes/test_global_mutation_auth.py | 12 ++++++----- webui/src/components/layout/Layout.test.tsx | 20 +++++++++++++++++++ webui/src/components/layout/Layout.tsx | 10 +++------- webui/src/pages/Home/index.tsx | 20 +++++++++---------- 5 files changed, 40 insertions(+), 27 deletions(-) diff --git a/flocks/server/routes/onboarding.py b/flocks/server/routes/onboarding.py index 49e1968dc..a2505c3c7 100644 --- a/flocks/server/routes/onboarding.py +++ b/flocks/server/routes/onboarding.py @@ -11,7 +11,7 @@ from contextlib import asynccontextmanager from typing import Any, Dict, List, Optional -from fastapi import APIRouter, Depends, HTTPException +from fastapi import APIRouter, HTTPException from pydantic import BaseModel, Field from flocks.config.config import Config @@ -36,7 +36,6 @@ set_mcp_credentials, test_mcp_connection, ) -from flocks.server.auth import require_admin from flocks.server.routes.provider import ( APIServiceUpdateRequest, ProviderCredentialRequest, @@ -878,7 +877,6 @@ async def get_onboarding_status() -> OnboardingStatusResponse: @serialized_config_mutation async def validate_onboarding( request: OnboardingValidateRequest, - _admin: object = Depends(require_admin), ) -> OnboardingValidateResponse: return await _validate_onboarding_request(request) @@ -892,7 +890,6 @@ async def validate_onboarding( @serialized_config_mutation async def apply_onboarding( request: OnboardingValidateRequest, - _admin: object = Depends(require_admin), ) -> OnboardingApplyResponse: return await _apply_onboarding(request) diff --git a/tests/server/routes/test_global_mutation_auth.py b/tests/server/routes/test_global_mutation_auth.py index 8ec9c425c..583c6a0f4 100644 --- a/tests/server/routes/test_global_mutation_auth.py +++ b/tests/server/routes/test_global_mutation_auth.py @@ -56,18 +56,20 @@ def test_provider_global_configuration_credentials_and_tests_require_admin(): assert [response.status_code for response in responses] == [403] * len(responses) -def test_onboarding_configuration_requires_admin(): +def test_onboarding_configuration_allows_members(): from flocks.server.routes.onboarding import router client = _user_client(router, prefix="/api/onboarding", role="member") payload = { "region": "cn", - "use_threatbook_model": True, - "threatbook_api_key": "secret", + "use_threatbook_model": False, } - assert client.post("/api/onboarding/validate", json=payload).status_code == 403 - assert client.post("/api/onboarding/apply", json=payload).status_code == 403 + validate_response = client.post("/api/onboarding/validate", json=payload) + apply_response = client.post("/api/onboarding/apply", json=payload) + + assert validate_response.status_code == 200 + assert apply_response.status_code == 400 def test_mcp_credential_and_threatbook_configuration_require_admin(): diff --git a/webui/src/components/layout/Layout.test.tsx b/webui/src/components/layout/Layout.test.tsx index 3b837d04b..9e7347f06 100644 --- a/webui/src/components/layout/Layout.test.tsx +++ b/webui/src/components/layout/Layout.test.tsx @@ -399,6 +399,26 @@ describe('Layout onboarding entry', () => { expect(screen.queryByPlaceholderText('onboarding.bootstrap.modelKeyPlaceholder')).not.toBeInTheDocument(); }); + it('allows member users to open onboarding from the home entry', async () => { + const user = userEvent.setup(); + localStorage.setItem('flocks_onboarding_dismissed', 'true'); + useAuth.mockReturnValue({ + user: { + id: 'user-2', + username: 'member', + role: 'member', + status: 'active', + must_reset_password: false, + }, + logout: vi.fn(), + }); + + renderHomeWithLayout(); + await user.click(screen.getByRole('button', { name: 'getStarted' })); + + expect(await screen.findByText('onboarding.bootstrap.modelPageTitle')).toBeInTheDocument(); + }); + it('auto-opens onboarding from backend status even when the old dismissed flag exists', async () => { localStorage.setItem('flocks_onboarding_dismissed', 'true'); onboardingAPI.getStatus.mockResolvedValue({ diff --git a/webui/src/components/layout/Layout.tsx b/webui/src/components/layout/Layout.tsx index fde3eb296..0a50be13d 100644 --- a/webui/src/components/layout/Layout.tsx +++ b/webui/src/components/layout/Layout.tsx @@ -377,7 +377,7 @@ export default function Layout() { }, [collapsed, sidebarWidth, updateSidebarWidth]); useEffect(() => { - if (!isHome || user?.role !== 'admin') return undefined; + if (!isHome) return undefined; let cancelled = false; onboardingAPI.getStatus() @@ -395,13 +395,9 @@ export default function Layout() { return () => { cancelled = true; }; - }, [isHome, user?.role]); + }, [isHome]); - const handleOpenOnboarding = useCallback(() => { - if (user?.role === 'admin') { - setShowOnboarding(true); - } - }, [user?.role]); + const handleOpenOnboarding = useCallback(() => setShowOnboarding(true), []); useEffect(() => { window.addEventListener('flocks:open-onboarding', handleOpenOnboarding); diff --git a/webui/src/pages/Home/index.tsx b/webui/src/pages/Home/index.tsx index 29677717d..4694ca106 100644 --- a/webui/src/pages/Home/index.tsx +++ b/webui/src/pages/Home/index.tsx @@ -144,17 +144,15 @@ export default function Home() {
- {user?.role === 'admin' ? ( - - ) : null} + {canCreateWebUIContractPage ? (
+ {isThreatBookLLMProviderId(provider.id) && ( + + + {t('form.claimFreeKey')} + + )} {provider.id !== 'ollama' && providerAllowsEmptyApiKey(provider.id) && (

{t('form.apiKeyOptionalHint')}

)} From 41acdf7884af92d0ceb657dcf00d73582853727e Mon Sep 17 00:00:00 2001 From: John Yin <10972267+john-yin2333@user.noreply.gitee.com> Date: Mon, 7 Sep 2026 21:59:08 +0800 Subject: [PATCH 51/63] fix(soc): make dashboard metrics and tasks authoritative --- .../soc_ui/soc_dashboard/api/handlers.py | 459 +++++++++++++++++- .../soc_ui/soc_dashboard/api/routes.yaml | 5 + .../webuis/soc_ui/soc_dashboard/src/Page.tsx | 294 ++++++++--- .../stream_alert_denoise/workflow.json | 2 +- flocks/workflow/store.py | 223 +++++++++ tests/hub/test_soc_dashboard_schema.py | 234 +++++++++ tests/workflow/test_workflow_store.py | 101 ++++ .../utils/socDashboardPageRuntime.test.tsx | 104 ++++ 8 files changed, 1346 insertions(+), 76 deletions(-) diff --git a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/handlers.py b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/handlers.py index 240bd4286..77c974d44 100644 --- a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/handlers.py +++ b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/handlers.py @@ -904,9 +904,118 @@ def _empty_workflow_denoise_stats(): "seriesUnique": [], "timelineLabels": [], "timelineWindow": "", + "metricsAvailable": False, + "dataQuality": "legacy", + "coverageComplete": False, + "coverageStartedAt": 0, + "sourceCoverageRate": 0, + "invalidExecutionCount": 0, + "dataSource": "workflow.db.workflow_stats.call_count", } +def _get_workflow_metric_rollups(workflow_name, start_time, end_time): + if not WORKFLOW_DB.is_file(): + return None + start_ms = max(_safe_int(start_time), 0) * 1000 + end_ms = max(_safe_int(end_time), 0) * 1000 + try: + with sqlite3.connect(f"file:{WORKFLOW_DB}?mode=ro", uri=True, timeout=1.0) as conn: + conn.row_factory = sqlite3.Row + conn.execute("PRAGMA query_only = ON") + if not ( + _table_exists(conn, "workflow_metric_rollups") + and _table_exists(conn, "workflow_metric_meta") + ): + return None + meta = conn.execute( + "SELECT coverage_started_at, updated_at FROM workflow_metric_meta " + "WHERE workflow_id = ?", + (workflow_name,), + ).fetchone() + if meta is None: + return None + query = ( + "SELECT * FROM workflow_metric_rollups WHERE workflow_id = ?" + ) + query_params = [workflow_name] + if start_ms > 0 and end_ms > 0: + query += " AND bucket_start >= ? AND bucket_start <= ?" + query_params.extend((start_ms - (start_ms % 60000), end_ms)) + query += " ORDER BY bucket_start" + rows = conn.execute(query, query_params).fetchall() + except Exception: + return None + + result = _empty_workflow_denoise_stats() + source_counts = Counter() + for row in rows: + parsed_sources = _safe_json_object(row["source_counts"]) + if isinstance(parsed_sources, dict): + for key, value in parsed_sources.items(): + source_counts[_norm(key)] += max(_safe_int(value), 0) + raw_count = sum(max(_safe_int(row["raw_count"]), 0) for row in rows) + normalized_count = sum(max(_safe_int(row["normalized_count"]), 0) for row in rows) + after_filter_count = sum(max(_safe_int(row["after_filter_count"]), 0) for row in rows) + unique_count = sum(max(_safe_int(row["unique_count"]), 0) for row in rows) + filter_removed_count = sum(max(_safe_int(row["filter_removed_count"]), 0) for row in rows) + duplicate_count = sum(max(_safe_int(row["duplicate_count"]), 0) for row in rows) + success_count = sum(max(_safe_int(row["success_count"]), 0) for row in rows) + error_count = sum(max(_safe_int(row["error_count"]), 0) for row in rows) + invalid_count = sum(max(_safe_int(row["invalid_count"]), 0) for row in rows) + source_covered_count = sum(max(_safe_int(row["source_covered_count"]), 0) for row in rows) + coverage_started_at = max(_safe_int(meta["coverage_started_at"]), 0) + complete_window = not start_ms or coverage_started_at <= start_ms + quality = "complete" if complete_window and invalid_count == 0 else "partial" + + first_bucket = _safe_int(rows[0]["bucket_start"]) if rows else start_ms + last_bucket = _safe_int(rows[-1]["bucket_start"]) if rows else end_ms + bucket_start, bucket_seconds, bucket_count, labels, window = _timeline_spec( + [], + start_time or first_bucket // 1000, + end_time or max(last_bucket // 1000, first_bucket // 1000), + ) + series_raw = [0] * bucket_count + series_unique = [0] * bucket_count + for row in rows: + index = int(((_safe_int(row["bucket_start"]) // 1000) - bucket_start) / bucket_seconds) + if 0 <= index < bucket_count: + series_raw[index] += max(_safe_int(row["raw_count"]), 0) + series_unique[index] += max(_safe_int(row["unique_count"]), 0) + + result.update( + { + "callCount": success_count + error_count, + "successCount": success_count, + "errorCount": error_count, + "earliestStartedAt": first_bucket, + "latestStartedAt": last_bucket, + "rawCount": raw_count, + "normalizedCount": normalized_count, + "afterFilterCount": after_filter_count, + "uniqueCount": unique_count, + "filterRemovedCount": filter_removed_count, + "duplicateCount": duplicate_count, + "reducedCount": max(raw_count - unique_count, 0), + "reductionRate": _ratio(max(raw_count - unique_count, 0), raw_count), + "dedupRate": _ratio(duplicate_count, after_filter_count), + "sourceCounts": dict(source_counts), + "seriesRaw": series_raw, + "seriesUnique": series_unique, + "timelineLabels": labels, + "timelineWindow": window, + "metricsAvailable": True, + "dataQuality": quality, + "coverageComplete": complete_window, + "coverageStartedAt": coverage_started_at, + "sourceCoverageRate": _ratio(source_covered_count, normalized_count), + "invalidExecutionCount": invalid_count, + "dataSource": "workflow.db.workflow_metric_rollups", + } + ) + return result + + def _get_workflow_denoise_stats( workflow_name: str, start_time: int = 0, @@ -927,6 +1036,15 @@ def _get_workflow_denoise_stats( if not WORKFLOW_DB.is_file(): return empty + rollup_result = _get_workflow_metric_rollups(workflow_name, start_time, end_time) + if rollup_result is not None and rollup_result.get("coverageComplete"): + with _cache_lock: + _workflow_stats_cache[cache_key] = {"updatedAt": now, "value": rollup_result} + _workflow_stats_cache.move_to_end(cache_key) + while len(_workflow_stats_cache) > _WORKFLOW_CACHE_MAX: + _workflow_stats_cache.popitem(last=False) + return rollup_result + try: with sqlite3.connect(WORKFLOW_DB) as conn: sample_deltas = _workflow_stats_sample_deltas( @@ -990,6 +1108,22 @@ def _get_workflow_denoise_stats( result_dict["timelineLabels"] = labels result_dict["timelineWindow"] = window + if rollup_result is not None: + # A newly-created rollup cannot represent the part of a requested + # window that predates metric collection. Keep the legacy values + # visible until the selected window is fully covered, and expose + # the coverage state so the UI does not present them as exact. + result_dict.update( + { + "dataQuality": "legacy-partial", + "coverageComplete": False, + "coverageStartedAt": rollup_result.get("coverageStartedAt", 0), + "invalidExecutionCount": rollup_result.get("invalidExecutionCount", 0), + "sourceCoverageRate": rollup_result.get("sourceCoverageRate", 0), + "shadowMetricsAvailable": True, + } + ) + with _cache_lock: _workflow_stats_cache[cache_key] = {"updatedAt": now, "value": result_dict} _workflow_stats_cache.move_to_end(cache_key) @@ -1250,6 +1384,10 @@ async def get_task_center(ctx, request): return await asyncio.to_thread(_get_task_center, include_mock) +async def get_ai_tasks(ctx, request): + return await asyncio.to_thread(_get_ai_tasks) + + def _table_exists(conn, table_name): return bool( conn.execute( @@ -1912,6 +2050,271 @@ def _get_task_center(include_mock=False): } +def _workflow_task_metric(stats, key): + if key not in stats: + return None, "missing" + value = stats.get(key) + if isinstance(value, bool): + return None, "invalid" + try: + parsed = int(value) + except (TypeError, ValueError, OverflowError): + return None, "invalid" + if parsed < 0: + return None, "invalid" + return parsed, "complete" + + +def _workflow_task_metrics(output_text, status): + output = _safe_json_object(output_text) + stats = output.get("stats") if isinstance(output.get("stats"), dict) else {} + values = {} + qualities = [] + for field, key in ( + ("raw", "raw_count"), + ("normalized", "normalized_count"), + ("afterFilter", "after_filter_count"), + ("unique", "after_dedup_count"), + ): + values[field], quality = _workflow_task_metric(stats, key) + qualities.append(quality) + if all(quality == "complete" for quality in qualities): + quality = "complete" + elif "invalid" in qualities: + quality = "invalid" + elif status in WORKFLOW_RUNNING_STATUSES: + quality = "pending" + else: + quality = "missing" + return values, quality + + +def _workflow_task_input_count(inputs): + for key in ("_raw_alerts_count", "raw_count"): + value, quality = _workflow_task_metric(inputs, key) + if quality == "complete": + return value + for key in ("raw_alerts", "alerts"): + value = inputs.get(key) + if isinstance(value, list): + return len(value) + for key in ("syslog_message", "syslog", "alert"): + if inputs.get(key) not in (None, "", {}): + return 1 + return None + + +def _workflow_task_row(row, workflow_id, effective_status): + output_text = row["output_results"] + input_text = row["input_params"] + output = _safe_json_object(output_text) + inputs = _safe_json_object(input_text) + metrics = _workflow_execution_metrics(output_text, input_text) + counts, data_quality = _workflow_task_metrics(output_text, effective_status) + raw_count_source = "workflow_output" if counts["raw"] is not None else "pending" + if counts["raw"] is None: + input_count = _workflow_task_input_count(inputs) + if input_count is not None: + counts["raw"] = input_count + raw_count_source = "workflow_input" + preview = metrics["preview"] + stage = "triage" if workflow_id in TRIAGE_WORKFLOW_IDS else "denoise" + title = _workflow_latest_alert_name(workflow_id, output_text, input_text) + if stage == "denoise" and not preview: + raw_count = counts["raw"] + title = ( + f"降噪批次 · 原始 {raw_count} 条" + if raw_count is not None + else "降噪批次 · 原始条数待生成" + ) + session_id, message_id = _workflow_link_context(row) + total_steps = max(_workflow_node_count(workflow_id), _safe_int(row["step_count"])) + current_step = max(_safe_int(row["current_step_index"]), 0) + if effective_status == "running" and total_steps > 0: + progress = { + "mode": "steps", + "current": min(max(current_step, 1), total_steps), + "total": total_steps, + "percent": _ratio(min(max(current_step, 1), total_steps), total_steps), + "label": f"第 {min(max(current_step, 1), total_steps)}/{total_steps} 步", + } + else: + progress = { + "mode": "waiting" if effective_status in {"queued", "pending"} else "none", + "current": current_step, + "total": total_steps, + "percent": None, + "label": "等待调度" if effective_status in {"queued", "pending"} else "", + } + source_type = metrics["sourceType"] + return { + "taskId": f"workflow-execution:{row['id']}", + "workflowId": workflow_id, + "executionId": str(row["id"]), + "stage": stage, + "status": effective_status, + "startedAt": _safe_int(row["started_at"]), + "updatedAt": _safe_int(row["updated_at"]), + "finishedAt": _safe_int(row["finished_at"]), + "currentPhase": str(row["current_phase"] or ""), + "title": title, + "sourceType": source_type, + "srcIp": preview.get("sip") or preview.get("src_ip") or preview.get("net_real_src_ip"), + "dstIp": preview.get("dip") or preview.get("dst_ip") or preview.get("net_dest_ip"), + "counts": counts, + "dataQuality": data_quality, + "rawCountSource": raw_count_source, + "emptyBatch": effective_status in WORKFLOW_SUCCESS_STATUSES and counts["raw"] == 0, + "progress": progress, + "sessionId": session_id, + "messageId": message_id, + "error": str(row["error_message"] or ""), + "inputMode": str(output.get("input_mode") or inputs.get("input_mode") or ""), + } + + +def _get_ai_tasks(): + empty_summary = { + "active": 0, + "running": 0, + "waiting": 0, + "stale": 0, + "disabled": 0, + "returned": 0, + "truncated": False, + } + if not WORKFLOW_DB.is_file(): + return { + "generatedAt": datetime.now().isoformat(timespec="seconds"), + "connection": "unavailable", + "reason": "workflow_db_missing", + "summary": empty_summary, + "tasks": [], + } + workflow_ids = tuple(SOC_PINNED_WORKFLOW_NAMES) + try: + with sqlite3.connect(f"file:{WORKFLOW_DB}?mode=ro", uri=True, timeout=1.0) as conn: + conn.row_factory = sqlite3.Row + conn.execute("PRAGMA query_only = ON") + if not _table_exists(conn, "workflow_executions"): + return { + "generatedAt": datetime.now().isoformat(timespec="seconds"), + "connection": "unavailable", + "reason": "workflow_executions_missing", + "summary": empty_summary, + "tasks": [], + } + columns = { + row[1] + for row in conn.execute("PRAGMA table_info(workflow_executions)").fetchall() + } + updated_expr = "updated_at" if "updated_at" in columns else "started_at" + freshness_expr = f"COALESCE(NULLIF({updated_expr}, 0), started_at)" + latest_select = ", ".join( + [ + "id", + "workflow_id", + "status", + "started_at", + _workflow_execution_column_expr(columns, "finished_at", "0"), + _workflow_execution_column_expr(columns, "updated_at", "0"), + _workflow_execution_column_expr(columns, "current_phase", "''"), + _workflow_execution_column_expr(columns, "current_step_index", "0"), + _workflow_execution_column_expr(columns, "step_count", "0"), + _workflow_execution_column_expr(columns, "output_results", "'{}'"), + _workflow_execution_column_expr(columns, "input_params", "'{}'"), + _workflow_execution_column_expr(columns, "payload", "'{}'"), + _workflow_execution_column_expr(columns, "error_message", "''"), + ] + ) + now_ms = int(time.time() * 1000) + tasks = [] + summary = dict(empty_summary) + for workflow_id in workflow_ids: + trigger_state = _workflow_trigger_state(conn, workflow_id) + cutoff = now_ms - trigger_state["timeoutSeconds"] * 1000 + if trigger_state["hasConfig"] and not trigger_state["enabled"]: + summary["disabled"] += _safe_int( + conn.execute( + "SELECT COUNT(*) FROM workflow_executions " + "WHERE workflow_id = ? AND status IN ('running', 'queued', 'pending')", + (workflow_id,), + ).fetchone()[0] + ) + continue + active_count = _safe_int( + conn.execute( + "SELECT COUNT(*) FROM workflow_executions " + "WHERE workflow_id = ? AND status IN ('running', 'queued', 'pending') " + f"AND {freshness_expr} >= ?", + (workflow_id, cutoff), + ).fetchone()[0] + ) + running_count = _safe_int( + conn.execute( + "SELECT COUNT(*) FROM workflow_executions " + "WHERE workflow_id = ? AND status = 'running' " + f"AND {freshness_expr} >= ?", + (workflow_id, cutoff), + ).fetchone()[0] + ) + waiting_count = _safe_int( + conn.execute( + "SELECT COUNT(*) FROM workflow_executions " + "WHERE workflow_id = ? AND status IN ('queued', 'pending') " + f"AND {freshness_expr} >= ?", + (workflow_id, cutoff), + ).fetchone()[0] + ) + stale_count = _safe_int( + conn.execute( + "SELECT COUNT(*) FROM workflow_executions " + "WHERE workflow_id = ? AND status IN ('running', 'queued', 'pending') " + f"AND {freshness_expr} < ?", + (workflow_id, cutoff), + ).fetchone()[0] + ) + summary["active"] += active_count + summary["running"] += running_count + summary["waiting"] += waiting_count + summary["stale"] += stale_count + rows = conn.execute( + f"SELECT {latest_select} FROM workflow_executions " + "WHERE workflow_id = ? AND status IN ('running', 'queued', 'pending') " + f"AND {freshness_expr} >= ? " + f"ORDER BY CASE WHEN status = 'running' THEN 0 ELSE 1 END, {freshness_expr} DESC " + "LIMIT 50", + (workflow_id, cutoff), + ).fetchall() + for row in rows: + effective_status = str(row["status"] or "").lower() + tasks.append(_workflow_task_row(row, workflow_id, effective_status)) + tasks.sort( + key=lambda item: ( + 0 if item["status"] == "running" else 1, + -max(item["updatedAt"], item["startedAt"]), + ) + ) + tasks = tasks[:50] + summary["returned"] = len(tasks) + summary["truncated"] = summary["active"] > len(tasks) + return { + "generatedAt": datetime.now().isoformat(timespec="seconds"), + "connection": "online", + "reason": "", + "summary": summary, + "tasks": tasks, + } + except Exception as exc: + return { + "generatedAt": datetime.now().isoformat(timespec="seconds"), + "connection": "error", + "reason": str(exc), + "summary": empty_summary, + "tasks": [], + } + + def _get_activity(params): _ensure_sqlite_schema() _maybe_prune_activity() @@ -2472,31 +2875,49 @@ def _get_stats(params): timeline_labels = workflow_stats["timelineLabels"] or denoise.get("_timelineLabels", []) timeline_window = workflow_stats["timelineWindow"] or denoise.get("_timelineWindow", "") workflow_series_raw = workflow_stats["seriesRaw"] - if not workflow_series_raw and soc_unique_series: - workflow_series_raw = [0] * len(soc_unique_series) - processed_total = workflow_stats["callCount"] - reduced_count = max(processed_total - soc_unique_count, 0) + metrics_available = bool(workflow_stats.get("metricsAvailable")) + if metrics_available: + processed_total = workflow_stats["rawCount"] + normalized_total = workflow_stats["normalizedCount"] + after_filter_total = workflow_stats["afterFilterCount"] + unique_total = workflow_stats["uniqueCount"] + filter_removed_count = workflow_stats["filterRemovedCount"] + duplicate_count = workflow_stats["duplicateCount"] + workflow_series_unique = workflow_stats["seriesUnique"] + else: + processed_total = workflow_stats["callCount"] + normalized_total = processed_total + after_filter_total = processed_total + unique_total = soc_unique_count + filter_removed_count = 0 + duplicate_count = max(processed_total - soc_unique_count, 0) + workflow_series_unique = soc_unique_series + if not workflow_series_raw and workflow_series_unique: + workflow_series_raw = [0] * len(workflow_series_unique) + reduced_count = max(processed_total - unique_total, 0) reduction_rate = _ratio(reduced_count, processed_total) denoise.update( { "totalRaw": processed_total, - "totalNormalized": processed_total, - "afterFilter": processed_total, - "totalUnique": soc_unique_count, - "filterRemoved": 0, - "dedupRemoved": reduced_count, + "totalNormalized": normalized_total, + "afterFilter": after_filter_total, + "totalUnique": unique_total, + "filterRemoved": filter_removed_count, + "dedupRemoved": duplicate_count, "duplicates": reduced_count, "duplicateRate": reduction_rate, - "dedupRate": reduction_rate, - "uniqueRate": _ratio(min(soc_unique_count, processed_total), processed_total), - "files": processed_total, + "dedupRate": _ratio(duplicate_count, after_filter_total), + "uniqueRate": _ratio(min(unique_total, processed_total), processed_total), + "files": workflow_stats["callCount"], "sourceCounter": Counter(workflow_stats["sourceCounts"]), "seriesRaw": workflow_series_raw, - "seriesUnique": soc_unique_series, + "seriesUnique": workflow_series_unique, "_timelineLabels": timeline_labels, "_timelineWindow": timeline_window, - "workflowCallCount": processed_total, - "dataSource": "workflow.db.workflow_stats.call_count + soc.db.unique", + "workflowCallCount": workflow_stats["callCount"], + "socPersistedUnique": soc_unique_count, + "dataSource": workflow_stats.get("dataSource"), + "dataQuality": workflow_stats.get("dataQuality"), } ) triage = _read_triage(triage_files) @@ -2528,6 +2949,14 @@ def _get_stats(params): }, "workflowStatsDb": _display_path(WORKFLOW_DB), "workflowStats": workflow_stats, + "metricQuality": { + "status": workflow_stats.get("dataQuality", "legacy"), + "coverageComplete": workflow_stats.get("coverageComplete", False), + "coverageStartedAt": workflow_stats.get("coverageStartedAt", 0), + "sourceCoverageRate": workflow_stats.get("sourceCoverageRate", 0), + "invalidExecutionCount": workflow_stats.get("invalidExecutionCount", 0), + "metricsAvailable": metrics_available, + }, "sampleMode": sample_mode, "sampleFile": ", ".join(_source_label(path) for path in asset_files) if sample_mode else "", "assets": { diff --git a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/routes.yaml b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/routes.yaml index 020897392..db37b7a93 100644 --- a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/routes.yaml +++ b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/routes.yaml @@ -9,6 +9,11 @@ routes: handler: handlers.get_activity timeoutMs: 5000 description: Incremental alert denoise and triage activity + - method: GET + path: /ai-tasks + handler: handlers.get_ai_tasks + timeoutMs: 5000 + description: Authoritative SOC workflow execution task snapshot - method: GET path: /task-center handler: handlers.get_task_center diff --git a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx index c08e3d499..8361c0423 100644 --- a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx +++ b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx @@ -24,7 +24,7 @@ const EMPTY_STATS = { eventRange: { start: '', end: '', label: '', source: '' }, generatedAt: '', latencyMs: 0, - sourceStatus: { workflowRoot: '', denoise: [], triage: [], denoiseFiles: [], triageFiles: [], missing: [] }, + sourceStatus: { workflowRoot: '', denoise: [], triage: [], denoiseFiles: [], triageFiles: [], missing: [], metricQuality: {} }, denoise: { totalRaw: 0, totalNormalized: 0, @@ -382,6 +382,69 @@ function createTaskCenterState() { }; } +function createAiTasksState() { + return { + connection: 'initializing', + generatedAt: '', + reason: '', + summary: { + active: 0, + running: 0, + waiting: 0, + stale: 0, + disabled: 0, + returned: 0, + truncated: false, + }, + tasks: [], + }; +} + +function createMockAiTasksState() { + const now = Date.now(); + const tasks = [ + { + taskId: 'workflow:stream_alert_triage:mock-triage-run-002', + workflowId: 'stream_alert_triage', + executionId: 'mock-triage-run-002', + stage: 'triage', + status: 'running', + startedAt: now - 6200, + title: '远程命令执行攻击（Mock）', + sourceType: 'tdp', + counts: { raw: null }, + dataQuality: 'pending', + progress: { mode: 'steps', current: 2, total: 4, percent: 0.5, label: '2 / 4 步' }, + }, + { + taskId: 'workflow:stream_alert_denoise:mock-denoise-run-004', + workflowId: 'stream_alert_denoise', + executionId: 'mock-denoise-run-004', + stage: 'denoise', + status: 'queued', + startedAt: now - 18000, + title: '端口扫描聚类（Mock）', + sourceType: 'qingteng', + counts: { raw: 12 }, + dataQuality: 'complete', + progress: { mode: 'queued', current: 0, total: 4, percent: null, label: '等待执行' }, + }, + ]; + return { + ...createAiTasksState(), + connection: 'online', + generatedAt: new Date(now).toISOString(), + summary: { + ...createAiTasksState().summary, + active: tasks.length, + running: 1, + waiting: 1, + returned: tasks.length, + }, + tasks, + }; +} + function createMockTaskCenterState() { const now = Date.now(); const startedAt = now - 7 * 60 * 1000; @@ -713,12 +776,18 @@ function refreshLabel(value) { function mergeStats(raw) { const denoise = { ...EMPTY_STATS.denoise, ...((raw || {}).denoise || {}) }; - const processedTotal = Math.max(Number(denoise.totalRaw || 0), 0); - denoise.totalNormalized = processedTotal; + const sourceStatus = { + ...EMPTY_STATS.sourceStatus, + ...((raw || {}).sourceStatus || {}), + metricQuality: { + ...(EMPTY_STATS.sourceStatus.metricQuality || {}), + ...((raw || {}).sourceStatus?.metricQuality || {}), + }, + }; return { ...EMPTY_STATS, ...(raw || {}), - sourceStatus: { ...EMPTY_STATS.sourceStatus, ...((raw || {}).sourceStatus || {}) }, + sourceStatus, denoise, triage: { ...EMPTY_STATS.triage, ...((raw || {}).triage || {}) }, pipeline: { ...EMPTY_STATS.pipeline, ...((raw || {}).pipeline || {}) }, @@ -727,10 +796,7 @@ function mergeStats(raw) { dateRange: { ...EMPTY_STATS.dateRange, ...((raw || {}).dateRange || {}) }, eventRange: { ...EMPTY_STATS.eventRange, ...((raw || {}).eventRange || {}) }, timeline: { ...EMPTY_STATS.timeline, ...((raw || {}).timeline || {}) }, - sources: [ - { key: 'ndr', label: 'NDR', value: processedTotal, rate: processedTotal > 0 ? 1 : 0, active: processedTotal > 0 }, - { key: 'other', label: '其他接入', value: 0, rate: 0, active: false }, - ], + sources: Array.isArray((raw || {}).sources) ? raw.sources : [], }; } @@ -2414,56 +2480,75 @@ function CommandTaskCenterPanel({ taskCenter }) { ]); } -function CommandAiTaskPanel({ activity, timeFilter }) { - const tasks = buildEventQueueTasks(activity, timeFilter); - const filterTransitionKey = [timeFilter.mode, timeFilter.range, timeFilter.start, timeFilter.end].join('|'); - const visibleTasks = useAnimatedTaskWindow( - tasks.filter((task) => task.state !== 'completed'), - filterTransitionKey, - ); - const counts = { - processing: visibleTasks.filter((task) => task.state === 'processing').length, - waiting: visibleTasks.filter((task) => task.state === 'waiting').length, - }; - const queueCount = visibleTasks.length; - const banner = activity.connection === 'error' - ? '处理任务连接异常，正在重试' - : counts.processing - ? `AI 正在并行处理 ${counts.processing} 个任务` - : counts.waiting ? '最新 10 条待处理任务' : '等待新的降噪或研判任务'; +function WorkflowTaskProgress({ task }) { + const progress = task.progress || {}; + if (progress.mode !== 'steps' || !Number.isFinite(Number(progress.percent))) { + return h('div', { className: 'event-rail-progress indeterminate', 'aria-label': 'AI 任务执行中' }, [ + h('span', { className: 'event-rail-progress-track', key: 'track' }, [h('i', { key: 'fill' })]), + h('small', { key: 'label' }, progress.label || '执行中'), + ]); + } + const percent = Math.max(Math.min(Number(progress.percent), 1), 0); + return h('div', { className: 'event-rail-progress', 'aria-label': 'AI 任务步骤进度' }, [ + h('span', { className: 'event-rail-progress-track', style: { '--queue-progress': percent }, key: 'track' }, [h('i', { key: 'fill' })]), + h('small', { key: 'label' }, progress.label || `${Math.round(percent * 100)}%`), + ]); +} + +function CommandAiTaskPanel({ aiTasks }) { + const summary = { ...createAiTasksState().summary, ...(aiTasks.summary || {}) }; + const visibleTasks = (aiTasks.tasks || []).slice(0, EVENT_RAIL_TASK_LIMIT); + const banner = aiTasks.connection === 'error' || aiTasks.connection === 'unavailable' + ? '处理任务数据不可用，正在重试' + : summary.running && summary.waiting + ? `正在处理 ${summary.running} 个，等待 ${summary.waiting} 个` + : summary.running + ? `AI 正在处理 ${summary.running} 个任务` + : summary.waiting + ? `${summary.waiting} 个任务等待处理` + : summary.stale + ? `发现 ${summary.stale} 个失联任务，已移出活跃列表` + : '当前没有运行或等待中的 AI 任务'; return [ - h('div', { className: cx('event-update-banner', activity.connection === 'error' && 'warn'), key: 'banner' }, banner), + h('div', { className: cx('event-update-banner', ['error', 'unavailable'].includes(aiTasks.connection) && 'warn'), key: 'banner' }, banner), + summary.truncated + ? h('small', { className: 'event-rail-limit-note', key: 'limit' }, `共 ${summary.active} 条活跃任务，当前展示 ${visibleTasks.length} 条`) + : null, h('div', { className: 'event-rail-list', key: 'list' }, visibleTasks.length ? visibleTasks.map((task) => { - const event = task.event; - const sampleCount = Math.max(Number(task.denoise?.sampleCount || 1), 1); - const title = `${event?.alert?.threatName || '未知告警'}${sampleCount > 1 ? ` × ${sampleCount}` : ''}`; + const processing = task.status === 'running'; + const waiting = ['queued', 'pending'].includes(task.status); + const title = task.title || (task.stage === 'triage' ? '研判任务' : '降噪任务'); const stageLabel = task.stage === 'triage' - ? task.state === 'waiting' ? '待研判' : '智能研判' - : task.state === 'waiting' ? '待降噪' : '智能降噪'; - const stateLabel = task.state === 'processing' - ? '处理中' - : '等待处理'; - const detail = event?.triggerSource === 'workflow_execution' - ? task.stage === 'triage' - ? task.state === 'processing' ? '研判工作流处理中' : '研判工作流待处理' - : event.result?.isDuplicate ? '重复告警已收敛' : '降噪处理完成' - : task.state === 'processing' - ? task.stage === 'triage' ? '证据关联与结论生成中' : '特征提取与相似聚类中' - : '等待 AI 处理'; - const hasExecution = Boolean(workflowIdFromEvent(event) && executionIdFromWorkflowEvent(event)); + ? waiting ? '待研判' : '智能研判' + : waiting ? '待降噪' : '智能降噪'; + const stateLabel = processing ? '处理中' : '等待处理'; + const qualityDetail = task.dataQuality === 'invalid' + ? ' · 指标格式异常' + : task.dataQuality === 'missing' + ? ' · 原始条数未知' + : task.dataQuality === 'pending' && (task.counts?.raw === null || task.counts?.raw === undefined) + ? ' · 原始条数待生成' + : ''; + const rawDetail = task.stage === 'denoise' && task.counts?.raw !== null && task.counts?.raw !== undefined + ? ` · 原始 ${task.counts.raw} 条` + : ''; + const detail = task.stage === 'triage' + ? processing ? '研判工作流执行中' : '等待研判工作流调度' + : processing ? '降噪工作流执行中' : '等待降噪工作流调度'; + const hasExecution = Boolean(task.workflowId && task.executionId); const handleOpen = () => { - if (hasExecution) openWorkflowExecutionFromEvent(event); + if (hasExecution) openWorkflowExecution(task.workflowId, task.executionId); }; const handleKeyDown = (keyboardEvent) => { if (!hasExecution) return; if (keyboardEvent.key === 'Enter' || keyboardEvent.key === ' ') { keyboardEvent.preventDefault(); - openWorkflowExecutionFromEvent(event); + openWorkflowExecution(task.workflowId, task.executionId); } }; return h('article', { - className: cx('event-rail-item', `state-${task.state}`, `kind-${task.stage}`, `motion-${task.motion || 'stable'}`, hasExecution && 'clickable'), - key: task.key, + className: cx('event-rail-item', processing ? 'state-processing' : 'state-waiting', `kind-${task.stage}`, hasExecution && 'clickable'), + key: task.taskId, role: hasExecution ? 'button' : undefined, tabIndex: hasExecution ? 0 : undefined, title: hasExecution ? '打开执行详情' : undefined, @@ -2473,20 +2558,19 @@ function CommandAiTaskPanel({ activity, timeFilter }) { h('div', { className: 'event-rail-meta', key: 'meta' }, [ h('span', { className: cx('event-queue-kind', `kind-${task.stage}`), key: 'kind' }, stageLabel), h('span', { className: 'event-stage', key: 'stage' }, stateLabel), - h('time', { key: 'time' }, eventTimeLabel(event?.occurredAt)), + h('time', { key: 'time' }, eventTimeLabel(task.startedAt)), ]), h('strong', { title, key: 'title' }, title), - h('span', { title: eventEndpoint(event), key: 'endpoint' }, eventEndpoint(event)), - h('small', { key: 'result' }, hasExecution ? `${detail} · 查看执行` : detail), - task.state === 'processing' ? h(EventQueueProgress, { event, key: 'progress' }) : null, + h('span', { key: 'endpoint' }, [task.srcIp, task.dstIp].filter(Boolean).join(' → ') || '未提供网络端点'), + h('small', { key: 'result' }, `${detail}${rawDetail}${qualityDetail}${hasExecution ? ' · 查看执行' : ''}`), + processing ? h(WorkflowTaskProgress, { task, key: 'progress' }) : null, ]); - }) : h('div', { className: 'event-rail-empty' }, '等待新的降噪或研判任务')), + }) : h('div', { className: 'event-rail-empty' }, '暂无活跃的降噪或研判任务')), ]; } -function CommandEventRail({ activity, timeFilter, taskCenter, view, onViewChange, collapsed, onToggle, railWidth, onResizeStart, onResizeKeyDown }) { - const tasks = buildEventQueueTasks(activity, timeFilter); - const queueCount = tasks.filter((task) => task.state !== 'completed').length; +function CommandEventRail({ aiTasks, taskCenter, view, onViewChange, collapsed, onToggle, railWidth, onResizeStart, onResizeKeyDown }) { + const queueCount = Math.max(Number(aiTasks.summary?.active || 0), 0); const taskCenterCount = Number(taskCenter.sessionCount || 0); const content = collapsed ? [] : [ h('div', { className: 'event-rail-head rail-view-head', key: 'head' }, [ @@ -2512,7 +2596,7 @@ function CommandEventRail({ activity, timeFilter, taskCenter, view, onViewChange ]), view === 'taskCenter' ? h(CommandTaskCenterPanel, { taskCenter, key: 'taskCenterContent' }) - : h(CommandAiTaskPanel, { activity, timeFilter, key: 'aiTaskContent' }), + : h(CommandAiTaskPanel, { aiTasks, key: 'aiTaskContent' }), ]; return h('aside', { className: cx('command-event-rail', collapsed && 'collapsed') }, [ collapsed ? null : h('div', { @@ -2554,6 +2638,7 @@ export default function Page() { const [loading, setLoading] = useState(true); const [error, setError] = useState(''); const [activity, setActivity] = useState(createActivityState); + const [aiTasks, setAiTasks] = useState(createAiTasksState); const [taskCenter, setTaskCenter] = useState(createTaskCenterState); const activityCursor = useRef(''); const workflowProgressByFilter = useRef(new Map()); @@ -2740,6 +2825,50 @@ export default function Page() { }; }, [mockDashboardEnabled]); + useEffect(() => { + let stopped = false; + let timer = 0; + + const schedule = (delay) => { + if (!stopped) timer = window.setTimeout(() => void poll(), delay); + }; + + const poll = async () => { + if (stopped) return; + if (document.hidden) { + schedule(ACTIVITY_POLL_MS); + return; + } + try { + const response = await getApi().page.get('/ai-tasks', { params: {} }); + const payload = response.data || {}; + if (!stopped) { + setAiTasks({ + ...createAiTasksState(), + ...payload, + summary: { ...createAiTasksState().summary, ...(payload.summary || {}) }, + tasks: Array.isArray(payload.tasks) ? payload.tasks : [], + }); + } + } catch (aiTaskError) { + if (!stopped) { + setAiTasks((previous) => ({ + ...previous, + connection: 'error', + reason: aiTaskError instanceof Error ? aiTaskError.message : 'ai tasks api failed', + })); + } + } + schedule(ACTIVITY_POLL_MS); + }; + + void poll(); + return () => { + stopped = true; + window.clearTimeout(timer); + }; + }, []); + useEffect(() => { let stopped = false; let timer = 0; @@ -2787,7 +2916,8 @@ export default function Page() { : (payload.events || []); const incomingEvents = rawIncomingEvents.filter((event) => event?.stage !== 'denoise'); const workflowEvents = Array.isArray(payload.workflowEvents) ? payload.workflowEvents : []; - for (const workflowEvent of workflowEvents) { + const activeWorkflowEvents = workflowEvents.filter(isRunningWorkflowEvent); + for (const workflowEvent of activeWorkflowEvents) { const hasExecution = Boolean(workflowIdFromEvent(workflowEvent) && executionIdFromWorkflowEvent(workflowEvent)); if (hasExecution) { incomingEvents.push(workflowEvent); @@ -2819,9 +2949,9 @@ export default function Page() { workflowProgressByFilter.current.set(workflowFilterKey, { callCount, latestStartedAt }); } const incomingRecentEvents = bootstrap - ? [...(payload.recentEvents || []), ...workflowEvents] + ? [...(payload.recentEvents || []), ...activeWorkflowEvents] : workflowChanged - ? [...rawIncomingEvents, ...workflowEvents] + ? [...rawIncomingEvents, ...activeWorkflowEvents] : rawIncomingEvents; setActivity((previous) => enqueueActivity( previous, @@ -2924,6 +3054,14 @@ export default function Page() { ), [mockDashboardEnabled, taskCenter], ); + const displayAiTasks = useMemo( + () => ( + mockDashboardEnabled && !aiTasks.tasks.length + ? createMockAiTasksState() + : aiTasks + ), + [aiTasks, mockDashboardEnabled], + ); const displayActivityBusy = Boolean( displayActivity.denoise.current || displayActivity.triage.current @@ -2932,6 +3070,18 @@ export default function Page() { || displayActivity.batch?.receivedCount || displayActivity.batch?.triageUpdatedCount ); + const metricQuality = stats.sourceStatus?.metricQuality || {}; + const metricQualityWarning = !stats.generatedAt + ? '' + : metricQuality.metricsAvailable + ? metricQuality.status === 'partial' + ? `降噪指标为部分覆盖数据${metricQuality.coverageStartedAt ? `，完整采集始于 ${taskCenterTimeLabel(metricQuality.coverageStartedAt)}` : ''}` + : Number(metricQuality.invalidExecutionCount || 0) > 0 + ? `${metricQuality.invalidExecutionCount} 次降噪执行的指标格式异常，未计入统计` + : '' + : metricQuality.status === 'legacy-partial' + ? `降噪指标仍使用历史兼容口径；精确口径完整采集始于 ${taskCenterTimeLabel(metricQuality.coverageStartedAt)}` + : '降噪指标仍使用历史兼容口径；新指标链路产生数据后将自动切换'; return h('div', { className: cx('adtd-root command-root', displayActivityBusy && 'command-is-processing', eventRailCollapsed && 'event-rail-is-collapsed'), @@ -2953,6 +3103,7 @@ export default function Page() { activity: displayActivity, }), error ? h('div', { className: 'error-banner', key: 'error' }, `统计接口异常：${error}`) : null, + metricQualityWarning ? h('div', { className: 'quality-banner', key: 'quality' }, metricQualityWarning) : null, h('main', { className: cx('command-shell', eventRailCollapsed && 'event-rail-collapsed'), key: 'main', @@ -2963,8 +3114,7 @@ export default function Page() { ]), h(CommandEventRail, { key: 'events', - activity: displayActivity, - timeFilter, + aiTasks: displayAiTasks, taskCenter: displayTaskCenter, view: rightRailView, onViewChange: setRightRailView, @@ -2995,6 +3145,15 @@ const CSS = ` overflow-x: auto; } .adtd-root * { box-sizing: border-box; } +.quality-banner { + margin: 8px 0 0; + padding: 8px 12px; + border: 1px solid rgba(255,176,32,.34); + border-radius: 6px; + color: #f1c67d; + background: rgba(139,88,22,.16); + font-size: 11px; +} .adtd-header { position: relative; display: grid; @@ -5255,6 +5414,12 @@ const CSS = ` } .event-update-banner:before { content: "ⓘ"; margin-right: 7px; color: #6ba4fb; } .event-update-banner.warn { border-color: rgba(255,174,52,.42); color: #f1c67d; background: rgba(139,88,22,.2); } +.event-rail-limit-note { + display: block; + margin: 5px 14px 0; + color: rgba(170,222,255,.62); + font-size: 9px; +} .task-center-panel { min-height: 0; padding: 8px 14px 18px; @@ -5681,6 +5846,11 @@ const CSS = ` box-shadow: 0 0 10px #73e9ff; transform: translate(50%, -50%); } +.event-rail-progress.indeterminate .event-rail-progress-track i { + width: 38%; + transform: translateX(-120%); + animation: commandTaskIndeterminate 1.4s ease-in-out infinite; +} .event-rail-progress > small { color: #9a8eff; font-size: 9px; @@ -5694,6 +5864,10 @@ const CSS = ` font-size: 11px; } @keyframes commandFlow { to { stroke-dashoffset: -36; } } +@keyframes commandTaskIndeterminate { + 0% { transform: translateX(-120%); } + 100% { transform: translateX(340%); } +} @keyframes commandSpin { to { transform: rotate(360deg); } } @keyframes commandSpinReverse { to { transform: rotate(-360deg); } } @keyframes commandCorePulse { diff --git a/.flocks/flockshub/plugins/workflows/stream_alert_denoise/workflow.json b/.flocks/flockshub/plugins/workflows/stream_alert_denoise/workflow.json index ea1b4193e..c2c914faa 100644 --- a/.flocks/flockshub/plugins/workflows/stream_alert_denoise/workflow.json +++ b/.flocks/flockshub/plugins/workflows/stream_alert_denoise/workflow.json @@ -27,7 +27,7 @@ "id": "dedup_and_write", "type": "python", "description": "Dedup (terminal): URI normalization + MinHash LSH (128 perms, 5-gram). LSH state persisted to ~/.flocks/workspace/workflows/stream_alert_denoise/ (atomic write, file lock, FIFO LRU eviction). Each output alert = normalized fields + dedup_key + is_duplicate + _lsh_cluster_id. Appends enriched alerts to JSONL files under ~/.flocks/workspace/workflows/stream_alert_denoise//dedup_result_NNN.jsonl. Each new file begins with a header line {_type:file_header, created_at, ...}; max 10,000 alert records per file, auto-increments sequence number on rollover.", - "code": "\nimport os\nimport re\nimport sys\nimport gc as _gc_module\nimport json\nimport pickle\nimport hashlib\nimport datetime\nimport threading\nimport types\nfrom datasketch import MinHash, MinHashLSH\n\nIS_WINDOWS = sys.platform == 'win32'\nif IS_WINDOWS:\n import msvcrt # noqa: F401\nelse:\n import fcntl # noqa: F401\n\nMINHASH_SEED = 2024\nNUM_PERM = 128\nWORKFLOW_NAME = 'stream_alert_denoise'\nLSH_CLUSTER_WARN_THRESHOLD = 100000\n\n# ── Process-level in-memory LSH state cache ──────────────────────────────────\n# Previously the ~649 MB pickle was loaded from disk on EVERY syslog message,\n# causing linear memory growth (Python GC cannot free old objects fast enough\n# under high throughput). We now keep the live MinHashLSH + lsh_cache +\n# dedup_key_cache in sys.modules between exec() invocations. A\n# threading.Lock serialises concurrent workflow threads; the file lock\n# (fcntl/msvcrt) still guards cross-process disk writes.\n_MEM_CACHE_KEY = f'_flocks_lsh_cache_{WORKFLOW_NAME}'\nif _MEM_CACHE_KEY not in sys.modules:\n _m = types.ModuleType(_MEM_CACHE_KEY)\n _m.lsh_index = None\n _m.lsh_cache = {}\n _m.dedup_key_cache = {}\n _m.next_cluster_id = 0\n _m.threshold = None\n _m.state_mtime = 0.0\n _m.append_mtime = 0.0\n _m.initialized = False\n _m.permutations = None\n _m.lock = threading.Lock()\n sys.modules[_MEM_CACHE_KEY] = _m\n_mem = sys.modules[_MEM_CACHE_KEY]\nif not hasattr(_mem, 'append_mtime'):\n _mem.append_mtime = 0.0\nif not hasattr(_mem, 'permutations'):\n _mem.permutations = None\n\ndef normalize_uri(uri):\n uri = str(uri or '')\n uri = re.sub(r'\\d{4}-\\d{2}-\\d{2}', 'DATETIME', uri)\n uri = re.sub(r'[\\da-f]{8}-[\\da-f]{4}-[\\da-f]{4}-[\\da-f]{4}-[\\da-f]{12}', 'UUID', uri, flags=re.IGNORECASE)\n uri = re.sub(r'(\\.\\./)+', 'TRAVERSAL', uri)\n uri = re.sub(r'\\bNULL\\b', 'NULL_REPLACED', uri)\n uri = re.sub(r'chr\\$\\d+\\$\\|\\|chr\\$\\d+\\$', 'CHR_SEQUENCE', uri)\n uri = re.sub(r'\\b\\d+={1,2}\\d+\\b', 'NUMBER_COMPARISON', uri)\n uri = re.sub(r'\\b[a-fA-F0-9]{32}\\b', 'HEXADECIMAL CHARACTERS', uri)\n return uri\n\ndef gen_minhash(text, permutations):\n shingles = [text[i:i+5] for i in range(len(text) - 4)]\n m = MinHash(num_perm=NUM_PERM, seed=MINHASH_SEED, permutations=permutations)\n for s in shingles:\n m.update(s.encode('utf-8'))\n return m\n\ndef get_state_paths(threshold):\n from flocks.config import Config\n flocks_root = Config().get_global().data_dir.parent\n state_dir = str(flocks_root / 'workspace' / 'workflows' / WORKFLOW_NAME)\n os.makedirs(state_dir, exist_ok=True)\n base = os.path.join(state_dir, f'lsh_state_np{NUM_PERM}_th{int(threshold * 100)}')\n return base + '.pkl', base + '.lock', base + '.append.log'\n\ndef get_output_dir():\n from flocks.config import Config\n from pathlib import Path\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\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\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\ndef load_state(state_path, append_path, threshold, max_keys=100000):\n lsh_index = None\n lsh_cache = {}\n dedup_key_cache = {}\n next_cid = 0\n if os.path.exists(state_path) and os.path.getsize(state_path) > 0:\n try:\n with open(state_path, 'rb') as f:\n state = pickle.load(f)\n if state.get('num_perm') != NUM_PERM or state.get('threshold') != threshold:\n print('[dedup] state params mismatch, starting fresh')\n else:\n raw_lsh_cache = state['lsh_cache']\n seen_raw = state.get('dedup_key_cache', {})\n raw_dedup_cache = {k: None for k in seen_raw} if isinstance(seen_raw, set) else (dict(seen_raw) if isinstance(seen_raw, dict) else {})\n oversized = len(raw_lsh_cache) > max_keys\n if oversized:\n # Snapshot grew beyond limit (no eviction in old versions).\n # Rebuild LSH index from the NEWEST max_keys clusters only to\n # avoid loading millions of entries into RAM all at once.\n keep_cids = set(list(raw_lsh_cache.keys())[-max_keys:])\n lsh_cache = {cid: mh for cid, mh in raw_lsh_cache.items() if cid in keep_cids}\n lsh_index = MinHashLSH(threshold=threshold, num_perm=NUM_PERM)\n for cid, mh in lsh_cache.items():\n try: lsh_index.insert(cid, mh)\n except Exception: pass\n dedup_key_cache = dict(list(raw_dedup_cache.items())[-max_keys:])\n print(f'[dedup] snapshot truncated {len(raw_lsh_cache)}→{len(lsh_cache)} clusters '\n f'(was over limit {max_keys})')\n del raw_lsh_cache, raw_dedup_cache, state\n _gc_module.collect()\n else:\n lsh_index = state['lsh_index']\n lsh_cache = raw_lsh_cache\n dedup_key_cache = raw_dedup_cache\n next_cid = (max(lsh_cache.keys()) + 1) if lsh_cache else 0\n print(f'[dedup] loaded snapshot: {len(lsh_cache)} clusters, {len(dedup_key_cache)} dedup_keys, next_cid={next_cid}')\n except Exception as e:\n print(f'[dedup] failed to load snapshot ({e}), starting fresh')\n lsh_index = None\n lsh_cache = {}\n dedup_key_cache = {}\n next_cid = 0\n if lsh_index is None:\n lsh_index = MinHashLSH(threshold=threshold, num_perm=NUM_PERM)\n lsh_cache = {}\n dedup_key_cache = {}\n next_cid = 0\n # Replay incremental append-log: only first-seen clusters/keys are stored there.\n if os.path.exists(append_path) and os.path.getsize(append_path) > 0:\n replayed = 0\n try:\n with open(append_path, 'rb') as f:\n while True:\n try:\n rec = pickle.load(f)\n except EOFError:\n break\n except Exception as _re:\n print(f'[dedup] append-log truncated at record {replayed} ({_re}), stopping replay')\n break\n if rec[0] == 'c':\n _, cid, mh = rec\n if cid not in lsh_cache:\n try:\n lsh_index.insert(cid, mh)\n except Exception:\n pass\n lsh_cache[cid] = mh\n if cid + 1 > next_cid:\n next_cid = cid + 1\n elif rec[0] == 'k':\n dedup_key_cache[rec[1]] = None\n replayed += 1\n except Exception as _e:\n print(f'[dedup] failed to replay append-log ({_e})')\n if replayed:\n print(f'[dedup] replayed {replayed} append-log records ({len(lsh_cache)} clusters, {len(dedup_key_cache)} keys)')\n return lsh_index, lsh_cache, dedup_key_cache, next_cid\n\ndef evict_oldest(lsh_index, lsh_cache, dedup_key_cache, max_keys):\n evicted_keys = evicted_clusters = 0\n excess = len(dedup_key_cache) - max_keys\n if excess > 0:\n for k in list(dedup_key_cache.keys())[:excess]:\n del dedup_key_cache[k]\n evicted_keys = excess\n excess = len(lsh_cache) - max_keys\n if excess > 0:\n for cid in list(lsh_cache.keys())[:excess]:\n try: lsh_index.remove(cid)\n except (KeyError, ValueError): pass\n del lsh_cache[cid]\n evicted_clusters = excess\n return evicted_keys, evicted_clusters\n\ndef dump_state_atomic(state_path, lsh_index, lsh_cache, dedup_key_cache, threshold, next_cluster_id):\n tmp = state_path + '.tmp'\n try:\n state = {\n 'lsh_index': lsh_index, 'lsh_cache': lsh_cache,\n 'dedup_key_cache': dedup_key_cache, 'next_cluster_id': next_cluster_id,\n 'num_perm': NUM_PERM, 'threshold': threshold,\n }\n with open(tmp, 'wb') as f:\n pickle.dump(state, f); f.flush(); os.fsync(f.fileno())\n os.replace(tmp, state_path)\n print(f'[dedup] state saved: {len(lsh_cache)} clusters, {len(dedup_key_cache)} dedup_keys')\n except Exception as e:\n print(f'[dedup] failed to save state: {e}')\n if os.path.exists(tmp):\n try: os.remove(tmp)\n except Exception: pass\n\ndef append_deltas(append_path, new_clusters, new_keys):\n # Persist ONLY this run's first-seen clusters + dedup_keys (append-only, O(new)).\n if not new_clusters and not new_keys:\n return\n try:\n with open(append_path, 'ab') as f:\n for _cid, _mh in new_clusters:\n pickle.dump(('c', _cid, _mh), f)\n for _dk in new_keys:\n pickle.dump(('k', _dk), f)\n f.flush(); os.fsync(f.fileno())\n except Exception as e:\n print(f'[dedup] failed to append deltas: {e}')\n\ndef compact_state(state_path, append_path, lsh_index, lsh_cache, dedup_key_cache, threshold, next_cluster_id):\n # Fold the append-log back into a fresh full snapshot, then drop the log.\n dump_state_atomic(state_path, lsh_index, lsh_cache, dedup_key_cache, threshold, next_cluster_id)\n try:\n if os.path.exists(append_path):\n os.remove(append_path)\n except Exception as e:\n print(f'[dedup] failed to truncate append-log: {e}')\n\n# ── Main ──────────────────────────────────────────────────────────────────────\n\nfiltered_alerts = inputs.get('filtered_alerts', [])\ninput_mode = inputs.get('input_mode', 'unknown')\ndedup_enabled = inputs.get('dedup_enabled', True)\nthreshold = float(inputs.get('dedup_threshold', 0.7))\nstrict_fields = inputs.get('strict_fields', ['sip', 'dip'])\nlsh_fields = inputs.get('lsh_fields', ['req_http_url', 'req_body', 'rsp_body'])\nmax_len = int(inputs.get('max_field_len', 500))\nmax_dedup_keys = int(inputs.get('max_dedup_keys', 100000))\nif max_dedup_keys < 1:\n max_dedup_keys = 100000\nstats = dict(inputs.get('stats', {}))\n\nif _mem.permutations is None:\n _mem.permutations = MinHash(num_perm=NUM_PERM, seed=MINHASH_SEED).permutations\n_permutations = _mem.permutations\nstate_path, lock_path, append_path = get_state_paths(threshold)\nevicted_keys = evicted_clusters = 0\n\nwith _mem.lock:\n # ── Load or reuse in-memory LSH state ────────────────────────────────────\n if dedup_enabled:\n disk_state_mtime = os.path.getmtime(state_path) if os.path.exists(state_path) else 0.0\n disk_append_mtime = os.path.getmtime(append_path) if os.path.exists(append_path) else 0.0\n cache_stale = (\n not _mem.initialized\n or _mem.threshold != threshold\n or disk_state_mtime > _mem.state_mtime + 0.5\n or disk_append_mtime > _mem.append_mtime + 0.5\n )\n if cache_stale:\n print(f'[dedup] cache miss (initialized={_mem.initialized}, '\n f'stale_by={disk_state_mtime - _mem.state_mtime:.1f}s), loading from disk')\n # Drop old references before load so GC can immediately reclaim\n # the ~649 MB state rather than waiting for the next cycle.\n old_lsh = _mem.lsh_index\n old_cache = _mem.lsh_cache\n _mem.lsh_index = None\n _mem.lsh_cache = {}\n _mem.dedup_key_cache = {}\n del old_lsh, old_cache\n _gc_module.collect()\n lsh_index, lsh_cache, dedup_key_cache, next_cluster_id = load_state(state_path, append_path, threshold, max_dedup_keys)\n if lsh_index is None:\n lsh_index = MinHashLSH(threshold=threshold, num_perm=NUM_PERM)\n lsh_cache = {}\n dedup_key_cache = {}\n next_cluster_id = 0\n _mem.lsh_index = lsh_index\n _mem.lsh_cache = lsh_cache\n _mem.dedup_key_cache = dedup_key_cache\n _mem.next_cluster_id = next_cluster_id\n _mem.threshold = threshold\n _mem.state_mtime = disk_state_mtime\n _mem.append_mtime = disk_append_mtime\n _mem.initialized = True\n else:\n print(f'[dedup] cache hit: {len(_mem.lsh_cache)} clusters, {len(_mem.dedup_key_cache)} keys')\n lsh_index = _mem.lsh_index\n lsh_cache = _mem.lsh_cache\n dedup_key_cache = _mem.dedup_key_cache\n next_cluster_id = _mem.next_cluster_id\n else:\n lsh_index, lsh_cache, dedup_key_cache, next_cluster_id = None, {}, {}, 0\n\n _cid_box = [next_cluster_id]\n _new_clusters = []\n def query_most_similar(minhash):\n sim_keys = lsh_index.query(minhash)\n if sim_keys:\n candidates = sim_keys[:100]\n sims = [minhash.jaccard(lsh_cache[k]) for k in candidates]\n return candidates[sims.index(max(sims))]\n cluster_id = _cid_box[0]\n _cid_box[0] += 1\n lsh_index.insert(cluster_id, minhash)\n lsh_cache[cluster_id] = minhash\n _new_clusters.append((cluster_id, minhash))\n return cluster_id\n\n enriched = []\n _new_keys = []\n for alert in filtered_alerts:\n alert = dict(alert)\n text_strict = '. '.join(str(alert.get(f, ''))[:max_len] for f in strict_fields)\n text_lsh = normalize_uri('. '.join(str(alert.get(f, ''))[:max_len] for f in lsh_fields))\n\n if not dedup_enabled:\n dk = hashlib.md5(f'{text_strict}. {text_lsh}'.encode('utf-8')).hexdigest()\n alert['_lsh_cluster_id'] = None\n alert['dedup_key'] = dk\n alert['is_duplicate'] = dk in dedup_key_cache\n dedup_key_cache[dk] = None\n enriched.append(alert)\n continue\n\n mh = gen_minhash(text_lsh.lower(), _permutations)\n cluster_id = query_most_similar(mh)\n alert['_lsh_cluster_id'] = cluster_id\n\n dk = hashlib.md5(f'{text_strict}. {cluster_id}'.encode('utf-8')).hexdigest()\n already = dk in dedup_key_cache\n if already:\n del dedup_key_cache[dk]\n else:\n _new_keys.append(dk)\n dedup_key_cache[dk] = None\n alert['dedup_key'] = dk\n alert['is_duplicate'] = already\n enriched.append(alert)\n\n if dedup_enabled:\n evicted_keys, evicted_clusters = evict_oldest(lsh_index, lsh_cache, dedup_key_cache, max_dedup_keys)\n if evicted_keys or evicted_clusters:\n print(f'[dedup] LRU eviction: dropped {evicted_keys} keys, {evicted_clusters} clusters')\n if len(lsh_cache) > LSH_CLUSTER_WARN_THRESHOLD or len(dedup_key_cache) > LSH_CLUSTER_WARN_THRESHOLD:\n print(f'[dedup] WARNING: persisted state holds {len(lsh_cache)} clusters '\n f'and {len(dedup_key_cache)} dedup_keys (warn={LSH_CLUSTER_WARN_THRESHOLD})')\n # Write to disk under file lock to protect against concurrent processes\n _lock_fh = acquire_lock(lock_path)\n try:\n # Incremental persistence: append ONLY this run's newly-created clusters\n # and first-seen dedup_keys (O(new) per message instead of O(N) full pickle).\n append_deltas(append_path, _new_clusters, _new_keys)\n # Periodic compaction: when the append-log grows comparable to the snapshot,\n # fold it into a fresh full snapshot and truncate the log (bounds disk + replay).\n _snap_size = os.path.getsize(state_path) if os.path.exists(state_path) else 0\n _app_size = os.path.getsize(append_path) if os.path.exists(append_path) else 0\n if (_snap_size == 0 and lsh_cache) or _app_size > max(_snap_size, 4 * 1024 * 1024):\n compact_state(state_path, append_path, lsh_index, lsh_cache, dedup_key_cache, threshold, _cid_box[0])\n print(f'[dedup] compacted append-log into snapshot ({len(lsh_cache)} clusters, {len(dedup_key_cache)} keys)')\n finally:\n release_lock(_lock_fh)\n _mem.next_cluster_id = _cid_box[0]\n _mem.state_mtime = os.path.getmtime(state_path) if os.path.exists(state_path) else _mem.state_mtime\n _mem.append_mtime = os.path.getmtime(append_path) if os.path.exists(append_path) else 0.0\n\n# ── Unique alerts (first seen per dedup_key) ──────────────────────────────────\nseen_keys = {}\nunique_alerts = []\nfor a in enriched:\n k = a['dedup_key']\n if k not in seen_keys:\n seen_keys[k] = a\n unique_alerts.append(a)\n\ndup_count = len(enriched) - len(unique_alerts)\nprint(f'[dedup] input={len(filtered_alerts)}, enriched={len(enriched)}, unique={len(unique_alerts)}, duplicates={dup_count}')\n\nstats['after_dedup_count'] = len(enriched)\nstats['unique_key_count'] = len(unique_alerts)\nstats['dedup_removed_count'] = dup_count\nstats['dedup_ratio'] = round(dup_count / len(enriched), 4) if enriched else 0.0\nstats['dedup_state_persisted'] = bool(dedup_enabled)\nif dedup_enabled:\n stats['lsh_total_clusters'] = len(lsh_cache)\n stats['lsh_total_dedup_keys'] = len(dedup_key_cache)\n stats['lsh_max_dedup_keys'] = max_dedup_keys\n stats['lsh_evicted_keys'] = evicted_keys\n stats['lsh_evicted_clusters'] = evicted_clusters\n\nif dedup_enabled:\n summary = (\n f'stream_alert_denoise done: raw={stats.get(\"raw_count\", 0)}'\n f' -> normalized={stats.get(\"normalized_count\", 0)}'\n f' -> filtered={stats.get(\"after_filter_count\", 0)}'\n f' -> enriched={len(enriched)}, unique={len(unique_alerts)} (compression {stats[\"dedup_ratio\"]:.1%})'\n f' | clusters={len(lsh_cache)}, keys={len(dedup_key_cache)}, max={max_dedup_keys}'\n )\nelse:\n summary = (\n f'stream_alert_denoise done (dedup_enabled=False): '\n f'raw={stats.get(\"raw_count\", 0)}'\n f' -> filtered={stats.get(\"after_filter_count\", 0)}'\n f' -> enriched={len(enriched)}'\n )\nprint(f'[dedup] {summary}')\n\n# ── Write enriched alerts to JSONL (counter sidecar replaces O(N) scan) ───────\n# dedup_result_001.jsonl, 002.jsonl ... each starts with a file_header line.\n# A lightweight sidecar (.dedup_counter.json) tracks the active file seq/count\n# so we never scan the whole file on every execution.\nMAX_RECORDS_PER_FILE = 10000\n_JSONL_PREFIX = 'dedup_result'\n_COUNTER_FILE = '.dedup_counter.json'\n\ndef _get_counter(out_dir):\n path = os.path.join(out_dir, _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\ndef _set_counter(out_dir, seq, count):\n path = os.path.join(out_dir, _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\ndef _find_active_file(out_dir):\n seq, count = _get_counter(out_dir)\n if seq > 0:\n path = os.path.join(out_dir, f'{_JSONL_PREFIX}_{seq:03d}.jsonl')\n if os.path.exists(path):\n return path, count, seq\n # Sidecar missing/stale: one-time recovery scan\n import glob as _glob\n existing = sorted(_glob.glob(os.path.join(out_dir, _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(_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 _set_counter(out_dir, seq, count)\n return latest, count, seq\n\ndef _write_jsonl(out_dir, alerts, now):\n written = []\n active_path, active_count, seq = _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'{_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 }\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 _set_counter(out_dir, seq, active_count)\n return written\n\n# Persist ONLY genuinely first-seen alerts (cross-batch is_duplicate=False).\n# NOTE: unique_alerts is only batch-local dedup; in single-alert syslog streaming\n# it is always length 1, so filtering by is_duplicate is what actually drops repeats.\n_persisted_alerts = [a for a in enriched if not a.get('is_duplicate')]\n_now = datetime.datetime.now()\ntry:\n _out_dir = get_output_dir()\n _written_paths = _write_jsonl(_out_dir, _persisted_alerts, _now) if _persisted_alerts else []\n _out_path = _written_paths[-1] if _written_paths else ''\n print(f'[dedup] wrote {len(_persisted_alerts)} first-seen records (skipped {len(enriched)-len(_persisted_alerts)} duplicates) -> {_written_paths}')\n stats['output_path'] = _out_path\n stats['output_paths'] = _written_paths\n outputs['output_path'] = _out_path\n outputs['output_paths'] = _written_paths\nexcept Exception as _we:\n import traceback\n print(f'[dedup] WARNING: failed to write JSONL: {_we}\\n{traceback.format_exc()}')\n outputs['output_path'] = ''\n outputs['output_paths'] = []\n\n# ── Outputs ───────────────────────────────────────────────────────────────────\n# Strip large body/header fields from in-memory run result to reduce flocks\n# run-history memory footprint. Full data is already persisted to JSONL.\n_HEAVY_OUTPUT_FIELDS = {\n 'net_http_reqs_header', 'net_http_resp_header',\n 'net_http_resp_body', 'net_http_reqs_body',\n 'net_http_resp_line', 'net_http_reqs_line',\n 'net_http_reqs_cookie',\n}\ndef _slim_alert(a):\n return {k: v for k, v in a.items() if k not in _HEAVY_OUTPUT_FIELDS}\noutputs['enriched_alerts'] = [_slim_alert(a) for a in enriched]\noutputs['unique_alerts'] = [_slim_alert(a) for a in unique_alerts]\noutputs['stats'] = stats\noutputs['dedup_summary'] = summary\noutputs['input_mode'] = input_mode\n\nif enriched:\n outputs['dedup_key'] = enriched[0].get('dedup_key', '')\n outputs['is_duplicate'] = enriched[0].get('is_duplicate', False)\nelse:\n outputs['dedup_key'] = ''\n outputs['is_duplicate'] = False\n" + "code": "\nimport os\nimport re\nimport sys\nimport gc as _gc_module\nimport json\nimport pickle\nimport hashlib\nimport datetime\nimport threading\nimport types\nfrom datasketch import MinHash, MinHashLSH\n\nIS_WINDOWS = sys.platform == 'win32'\nif IS_WINDOWS:\n import msvcrt # noqa: F401\nelse:\n import fcntl # noqa: F401\n\nMINHASH_SEED = 2024\nNUM_PERM = 128\nWORKFLOW_NAME = 'stream_alert_denoise'\nLSH_CLUSTER_WARN_THRESHOLD = 100000\n\n# ── Process-level in-memory LSH state cache ──────────────────────────────────\n# Previously the ~649 MB pickle was loaded from disk on EVERY syslog message,\n# causing linear memory growth (Python GC cannot free old objects fast enough\n# under high throughput). We now keep the live MinHashLSH + lsh_cache +\n# dedup_key_cache in sys.modules between exec() invocations. A\n# threading.Lock serialises concurrent workflow threads; the file lock\n# (fcntl/msvcrt) still guards cross-process disk writes.\n_MEM_CACHE_KEY = f'_flocks_lsh_cache_{WORKFLOW_NAME}'\nif _MEM_CACHE_KEY not in sys.modules:\n _m = types.ModuleType(_MEM_CACHE_KEY)\n _m.lsh_index = None\n _m.lsh_cache = {}\n _m.dedup_key_cache = {}\n _m.next_cluster_id = 0\n _m.threshold = None\n _m.state_mtime = 0.0\n _m.append_mtime = 0.0\n _m.initialized = False\n _m.permutations = None\n _m.lock = threading.Lock()\n sys.modules[_MEM_CACHE_KEY] = _m\n_mem = sys.modules[_MEM_CACHE_KEY]\nif not hasattr(_mem, 'append_mtime'):\n _mem.append_mtime = 0.0\nif not hasattr(_mem, 'permutations'):\n _mem.permutations = None\n\ndef normalize_uri(uri):\n uri = str(uri or '')\n uri = re.sub(r'\\d{4}-\\d{2}-\\d{2}', 'DATETIME', uri)\n uri = re.sub(r'[\\da-f]{8}-[\\da-f]{4}-[\\da-f]{4}-[\\da-f]{4}-[\\da-f]{12}', 'UUID', uri, flags=re.IGNORECASE)\n uri = re.sub(r'(\\.\\./)+', 'TRAVERSAL', uri)\n uri = re.sub(r'\\bNULL\\b', 'NULL_REPLACED', uri)\n uri = re.sub(r'chr\\$\\d+\\$\\|\\|chr\\$\\d+\\$', 'CHR_SEQUENCE', uri)\n uri = re.sub(r'\\b\\d+={1,2}\\d+\\b', 'NUMBER_COMPARISON', uri)\n uri = re.sub(r'\\b[a-fA-F0-9]{32}\\b', 'HEXADECIMAL CHARACTERS', uri)\n return uri\n\ndef gen_minhash(text, permutations):\n shingles = [text[i:i+5] for i in range(len(text) - 4)]\n m = MinHash(num_perm=NUM_PERM, seed=MINHASH_SEED, permutations=permutations)\n for s in shingles:\n m.update(s.encode('utf-8'))\n return m\n\ndef get_state_paths(threshold):\n from flocks.config import Config\n flocks_root = Config().get_global().data_dir.parent\n state_dir = str(flocks_root / 'workspace' / 'workflows' / WORKFLOW_NAME)\n os.makedirs(state_dir, exist_ok=True)\n base = os.path.join(state_dir, f'lsh_state_np{NUM_PERM}_th{int(threshold * 100)}')\n return base + '.pkl', base + '.lock', base + '.append.log'\n\ndef get_output_dir():\n from flocks.config import Config\n from pathlib import Path\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\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\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\ndef load_state(state_path, append_path, threshold, max_keys=100000):\n lsh_index = None\n lsh_cache = {}\n dedup_key_cache = {}\n next_cid = 0\n if os.path.exists(state_path) and os.path.getsize(state_path) > 0:\n try:\n with open(state_path, 'rb') as f:\n state = pickle.load(f)\n if state.get('num_perm') != NUM_PERM or state.get('threshold') != threshold:\n print('[dedup] state params mismatch, starting fresh')\n else:\n raw_lsh_cache = state['lsh_cache']\n seen_raw = state.get('dedup_key_cache', {})\n raw_dedup_cache = {k: None for k in seen_raw} if isinstance(seen_raw, set) else (dict(seen_raw) if isinstance(seen_raw, dict) else {})\n oversized = len(raw_lsh_cache) > max_keys\n if oversized:\n # Snapshot grew beyond limit (no eviction in old versions).\n # Rebuild LSH index from the NEWEST max_keys clusters only to\n # avoid loading millions of entries into RAM all at once.\n keep_cids = set(list(raw_lsh_cache.keys())[-max_keys:])\n lsh_cache = {cid: mh for cid, mh in raw_lsh_cache.items() if cid in keep_cids}\n lsh_index = MinHashLSH(threshold=threshold, num_perm=NUM_PERM)\n for cid, mh in lsh_cache.items():\n try: lsh_index.insert(cid, mh)\n except Exception: pass\n dedup_key_cache = dict(list(raw_dedup_cache.items())[-max_keys:])\n print(f'[dedup] snapshot truncated {len(raw_lsh_cache)}→{len(lsh_cache)} clusters '\n f'(was over limit {max_keys})')\n del raw_lsh_cache, raw_dedup_cache, state\n _gc_module.collect()\n else:\n lsh_index = state['lsh_index']\n lsh_cache = raw_lsh_cache\n dedup_key_cache = raw_dedup_cache\n next_cid = (max(lsh_cache.keys()) + 1) if lsh_cache else 0\n print(f'[dedup] loaded snapshot: {len(lsh_cache)} clusters, {len(dedup_key_cache)} dedup_keys, next_cid={next_cid}')\n except Exception as e:\n print(f'[dedup] failed to load snapshot ({e}), starting fresh')\n lsh_index = None\n lsh_cache = {}\n dedup_key_cache = {}\n next_cid = 0\n if lsh_index is None:\n lsh_index = MinHashLSH(threshold=threshold, num_perm=NUM_PERM)\n lsh_cache = {}\n dedup_key_cache = {}\n next_cid = 0\n # Replay incremental append-log: only first-seen clusters/keys are stored there.\n if os.path.exists(append_path) and os.path.getsize(append_path) > 0:\n replayed = 0\n try:\n with open(append_path, 'rb') as f:\n while True:\n try:\n rec = pickle.load(f)\n except EOFError:\n break\n except Exception as _re:\n print(f'[dedup] append-log truncated at record {replayed} ({_re}), stopping replay')\n break\n if rec[0] == 'c':\n _, cid, mh = rec\n if cid not in lsh_cache:\n try:\n lsh_index.insert(cid, mh)\n except Exception:\n pass\n lsh_cache[cid] = mh\n if cid + 1 > next_cid:\n next_cid = cid + 1\n elif rec[0] == 'k':\n dedup_key_cache[rec[1]] = None\n replayed += 1\n except Exception as _e:\n print(f'[dedup] failed to replay append-log ({_e})')\n if replayed:\n print(f'[dedup] replayed {replayed} append-log records ({len(lsh_cache)} clusters, {len(dedup_key_cache)} keys)')\n return lsh_index, lsh_cache, dedup_key_cache, next_cid\n\ndef evict_oldest(lsh_index, lsh_cache, dedup_key_cache, max_keys):\n evicted_keys = evicted_clusters = 0\n excess = len(dedup_key_cache) - max_keys\n if excess > 0:\n for k in list(dedup_key_cache.keys())[:excess]:\n del dedup_key_cache[k]\n evicted_keys = excess\n excess = len(lsh_cache) - max_keys\n if excess > 0:\n for cid in list(lsh_cache.keys())[:excess]:\n try: lsh_index.remove(cid)\n except (KeyError, ValueError): pass\n del lsh_cache[cid]\n evicted_clusters = excess\n return evicted_keys, evicted_clusters\n\ndef dump_state_atomic(state_path, lsh_index, lsh_cache, dedup_key_cache, threshold, next_cluster_id):\n tmp = state_path + '.tmp'\n try:\n state = {\n 'lsh_index': lsh_index, 'lsh_cache': lsh_cache,\n 'dedup_key_cache': dedup_key_cache, 'next_cluster_id': next_cluster_id,\n 'num_perm': NUM_PERM, 'threshold': threshold,\n }\n with open(tmp, 'wb') as f:\n pickle.dump(state, f); f.flush(); os.fsync(f.fileno())\n os.replace(tmp, state_path)\n print(f'[dedup] state saved: {len(lsh_cache)} clusters, {len(dedup_key_cache)} dedup_keys')\n except Exception as e:\n print(f'[dedup] failed to save state: {e}')\n if os.path.exists(tmp):\n try: os.remove(tmp)\n except Exception: pass\n\ndef append_deltas(append_path, new_clusters, new_keys):\n # Persist ONLY this run's first-seen clusters + dedup_keys (append-only, O(new)).\n if not new_clusters and not new_keys:\n return\n try:\n with open(append_path, 'ab') as f:\n for _cid, _mh in new_clusters:\n pickle.dump(('c', _cid, _mh), f)\n for _dk in new_keys:\n pickle.dump(('k', _dk), f)\n f.flush(); os.fsync(f.fileno())\n except Exception as e:\n print(f'[dedup] failed to append deltas: {e}')\n\ndef compact_state(state_path, append_path, lsh_index, lsh_cache, dedup_key_cache, threshold, next_cluster_id):\n # Fold the append-log back into a fresh full snapshot, then drop the log.\n dump_state_atomic(state_path, lsh_index, lsh_cache, dedup_key_cache, threshold, next_cluster_id)\n try:\n if os.path.exists(append_path):\n os.remove(append_path)\n except Exception as e:\n print(f'[dedup] failed to truncate append-log: {e}')\n\n# ── Main ──────────────────────────────────────────────────────────────────────\n\nfiltered_alerts = inputs.get('filtered_alerts', [])\ninput_mode = inputs.get('input_mode', 'unknown')\ndedup_enabled = inputs.get('dedup_enabled', True)\nthreshold = float(inputs.get('dedup_threshold', 0.7))\nstrict_fields = inputs.get('strict_fields', ['sip', 'dip'])\nlsh_fields = inputs.get('lsh_fields', ['req_http_url', 'req_body', 'rsp_body'])\nmax_len = int(inputs.get('max_field_len', 500))\nmax_dedup_keys = int(inputs.get('max_dedup_keys', 100000))\nif max_dedup_keys < 1:\n max_dedup_keys = 100000\nstats = dict(inputs.get('stats', {}))\n\nif _mem.permutations is None:\n _mem.permutations = MinHash(num_perm=NUM_PERM, seed=MINHASH_SEED).permutations\n_permutations = _mem.permutations\nstate_path, lock_path, append_path = get_state_paths(threshold)\nevicted_keys = evicted_clusters = 0\n\nwith _mem.lock:\n # ── Load or reuse in-memory LSH state ────────────────────────────────────\n if dedup_enabled:\n disk_state_mtime = os.path.getmtime(state_path) if os.path.exists(state_path) else 0.0\n disk_append_mtime = os.path.getmtime(append_path) if os.path.exists(append_path) else 0.0\n cache_stale = (\n not _mem.initialized\n or _mem.threshold != threshold\n or disk_state_mtime > _mem.state_mtime + 0.5\n or disk_append_mtime > _mem.append_mtime + 0.5\n )\n if cache_stale:\n print(f'[dedup] cache miss (initialized={_mem.initialized}, '\n f'stale_by={disk_state_mtime - _mem.state_mtime:.1f}s), loading from disk')\n # Drop old references before load so GC can immediately reclaim\n # the ~649 MB state rather than waiting for the next cycle.\n old_lsh = _mem.lsh_index\n old_cache = _mem.lsh_cache\n _mem.lsh_index = None\n _mem.lsh_cache = {}\n _mem.dedup_key_cache = {}\n del old_lsh, old_cache\n _gc_module.collect()\n lsh_index, lsh_cache, dedup_key_cache, next_cluster_id = load_state(state_path, append_path, threshold, max_dedup_keys)\n if lsh_index is None:\n lsh_index = MinHashLSH(threshold=threshold, num_perm=NUM_PERM)\n lsh_cache = {}\n dedup_key_cache = {}\n next_cluster_id = 0\n _mem.lsh_index = lsh_index\n _mem.lsh_cache = lsh_cache\n _mem.dedup_key_cache = dedup_key_cache\n _mem.next_cluster_id = next_cluster_id\n _mem.threshold = threshold\n _mem.state_mtime = disk_state_mtime\n _mem.append_mtime = disk_append_mtime\n _mem.initialized = True\n else:\n print(f'[dedup] cache hit: {len(_mem.lsh_cache)} clusters, {len(_mem.dedup_key_cache)} keys')\n lsh_index = _mem.lsh_index\n lsh_cache = _mem.lsh_cache\n dedup_key_cache = _mem.dedup_key_cache\n next_cluster_id = _mem.next_cluster_id\n else:\n lsh_index, lsh_cache, dedup_key_cache, next_cluster_id = None, {}, {}, 0\n\n _cid_box = [next_cluster_id]\n _new_clusters = []\n def query_most_similar(minhash):\n sim_keys = lsh_index.query(minhash)\n if sim_keys:\n candidates = sim_keys[:100]\n sims = [minhash.jaccard(lsh_cache[k]) for k in candidates]\n return candidates[sims.index(max(sims))]\n cluster_id = _cid_box[0]\n _cid_box[0] += 1\n lsh_index.insert(cluster_id, minhash)\n lsh_cache[cluster_id] = minhash\n _new_clusters.append((cluster_id, minhash))\n return cluster_id\n\n enriched = []\n _new_keys = []\n for alert in filtered_alerts:\n alert = dict(alert)\n text_strict = '. '.join(str(alert.get(f, ''))[:max_len] for f in strict_fields)\n text_lsh = normalize_uri('. '.join(str(alert.get(f, ''))[:max_len] for f in lsh_fields))\n\n if not dedup_enabled:\n dk = hashlib.md5(f'{text_strict}. {text_lsh}'.encode('utf-8')).hexdigest()\n alert['_lsh_cluster_id'] = None\n alert['dedup_key'] = dk\n alert['is_duplicate'] = dk in dedup_key_cache\n dedup_key_cache[dk] = None\n enriched.append(alert)\n continue\n\n mh = gen_minhash(text_lsh.lower(), _permutations)\n cluster_id = query_most_similar(mh)\n alert['_lsh_cluster_id'] = cluster_id\n\n dk = hashlib.md5(f'{text_strict}. {cluster_id}'.encode('utf-8')).hexdigest()\n already = dk in dedup_key_cache\n if already:\n del dedup_key_cache[dk]\n else:\n _new_keys.append(dk)\n dedup_key_cache[dk] = None\n alert['dedup_key'] = dk\n alert['is_duplicate'] = already\n enriched.append(alert)\n\n if dedup_enabled:\n evicted_keys, evicted_clusters = evict_oldest(lsh_index, lsh_cache, dedup_key_cache, max_dedup_keys)\n if evicted_keys or evicted_clusters:\n print(f'[dedup] LRU eviction: dropped {evicted_keys} keys, {evicted_clusters} clusters')\n if len(lsh_cache) > LSH_CLUSTER_WARN_THRESHOLD or len(dedup_key_cache) > LSH_CLUSTER_WARN_THRESHOLD:\n print(f'[dedup] WARNING: persisted state holds {len(lsh_cache)} clusters '\n f'and {len(dedup_key_cache)} dedup_keys (warn={LSH_CLUSTER_WARN_THRESHOLD})')\n # Write to disk under file lock to protect against concurrent processes\n _lock_fh = acquire_lock(lock_path)\n try:\n # Incremental persistence: append ONLY this run's newly-created clusters\n # and first-seen dedup_keys (O(new) per message instead of O(N) full pickle).\n append_deltas(append_path, _new_clusters, _new_keys)\n # Periodic compaction: when the append-log grows comparable to the snapshot,\n # fold it into a fresh full snapshot and truncate the log (bounds disk + replay).\n _snap_size = os.path.getsize(state_path) if os.path.exists(state_path) else 0\n _app_size = os.path.getsize(append_path) if os.path.exists(append_path) else 0\n if (_snap_size == 0 and lsh_cache) or _app_size > max(_snap_size, 4 * 1024 * 1024):\n compact_state(state_path, append_path, lsh_index, lsh_cache, dedup_key_cache, threshold, _cid_box[0])\n print(f'[dedup] compacted append-log into snapshot ({len(lsh_cache)} clusters, {len(dedup_key_cache)} keys)')\n finally:\n release_lock(_lock_fh)\n _mem.next_cluster_id = _cid_box[0]\n _mem.state_mtime = os.path.getmtime(state_path) if os.path.exists(state_path) else _mem.state_mtime\n _mem.append_mtime = os.path.getmtime(append_path) if os.path.exists(append_path) else 0.0\n\n# ── Unique alerts (first seen per dedup_key) ──────────────────────────────────\nseen_keys = {}\nunique_alerts = []\nfor a in enriched:\n k = a['dedup_key']\n if k not in seen_keys:\n seen_keys[k] = a\n unique_alerts.append(a)\n\nfirst_seen_count = sum(1 for alert in enriched if not alert.get('is_duplicate'))\ndup_count = len(enriched) - first_seen_count\nprint(f'[dedup] input={len(filtered_alerts)}, enriched={len(enriched)}, unique={len(unique_alerts)}, duplicates={dup_count}')\n\nstats['metric_schema_version'] = 2\nstats['after_dedup_count'] = first_seen_count\nstats['unique_key_count'] = first_seen_count\nstats['dedup_removed_count'] = dup_count\nstats['dedup_ratio'] = round(dup_count / len(enriched), 4) if enriched else 0.0\nstats['dedup_state_persisted'] = bool(dedup_enabled)\nif dedup_enabled:\n stats['lsh_total_clusters'] = len(lsh_cache)\n stats['lsh_total_dedup_keys'] = len(dedup_key_cache)\n stats['lsh_max_dedup_keys'] = max_dedup_keys\n stats['lsh_evicted_keys'] = evicted_keys\n stats['lsh_evicted_clusters'] = evicted_clusters\n\nif dedup_enabled:\n summary = (\n f'stream_alert_denoise done: raw={stats.get(\"raw_count\", 0)}'\n f' -> normalized={stats.get(\"normalized_count\", 0)}'\n f' -> filtered={stats.get(\"after_filter_count\", 0)}'\n f' -> enriched={len(enriched)}, unique={len(unique_alerts)} (compression {stats[\"dedup_ratio\"]:.1%})'\n f' | clusters={len(lsh_cache)}, keys={len(dedup_key_cache)}, max={max_dedup_keys}'\n )\nelse:\n summary = (\n f'stream_alert_denoise done (dedup_enabled=False): '\n f'raw={stats.get(\"raw_count\", 0)}'\n f' -> filtered={stats.get(\"after_filter_count\", 0)}'\n f' -> enriched={len(enriched)}'\n )\nprint(f'[dedup] {summary}')\n\n# ── Write enriched alerts to JSONL (counter sidecar replaces O(N) scan) ───────\n# dedup_result_001.jsonl, 002.jsonl ... each starts with a file_header line.\n# A lightweight sidecar (.dedup_counter.json) tracks the active file seq/count\n# so we never scan the whole file on every execution.\nMAX_RECORDS_PER_FILE = 10000\n_JSONL_PREFIX = 'dedup_result'\n_COUNTER_FILE = '.dedup_counter.json'\n\ndef _get_counter(out_dir):\n path = os.path.join(out_dir, _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\ndef _set_counter(out_dir, seq, count):\n path = os.path.join(out_dir, _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\ndef _find_active_file(out_dir):\n seq, count = _get_counter(out_dir)\n if seq > 0:\n path = os.path.join(out_dir, f'{_JSONL_PREFIX}_{seq:03d}.jsonl')\n if os.path.exists(path):\n return path, count, seq\n # Sidecar missing/stale: one-time recovery scan\n import glob as _glob\n existing = sorted(_glob.glob(os.path.join(out_dir, _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(_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 _set_counter(out_dir, seq, count)\n return latest, count, seq\n\ndef _write_jsonl(out_dir, alerts, now):\n written = []\n active_path, active_count, seq = _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'{_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 }\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 _set_counter(out_dir, seq, active_count)\n return written\n\n# Persist ONLY genuinely first-seen alerts (cross-batch is_duplicate=False).\n# NOTE: unique_alerts is only batch-local dedup; in single-alert syslog streaming\n# it is always length 1, so filtering by is_duplicate is what actually drops repeats.\n_persisted_alerts = [a for a in enriched if not a.get('is_duplicate')]\n_now = datetime.datetime.now()\ntry:\n _out_dir = get_output_dir()\n _written_paths = _write_jsonl(_out_dir, _persisted_alerts, _now) if _persisted_alerts else []\n _out_path = _written_paths[-1] if _written_paths else ''\n print(f'[dedup] wrote {len(_persisted_alerts)} first-seen records (skipped {len(enriched)-len(_persisted_alerts)} duplicates) -> {_written_paths}')\n stats['output_path'] = _out_path\n stats['output_paths'] = _written_paths\n outputs['output_path'] = _out_path\n outputs['output_paths'] = _written_paths\nexcept Exception as _we:\n import traceback\n print(f'[dedup] WARNING: failed to write JSONL: {_we}\\n{traceback.format_exc()}')\n outputs['output_path'] = ''\n outputs['output_paths'] = []\n\n# ── Outputs ───────────────────────────────────────────────────────────────────\n# Strip large body/header fields from in-memory run result to reduce flocks\n# run-history memory footprint. Full data is already persisted to JSONL.\n_HEAVY_OUTPUT_FIELDS = {\n 'net_http_reqs_header', 'net_http_resp_header',\n 'net_http_resp_body', 'net_http_reqs_body',\n 'net_http_resp_line', 'net_http_reqs_line',\n 'net_http_reqs_cookie',\n}\ndef _slim_alert(a):\n return {k: v for k, v in a.items() if k not in _HEAVY_OUTPUT_FIELDS}\noutputs['enriched_alerts'] = [_slim_alert(a) for a in enriched]\noutputs['unique_alerts'] = [_slim_alert(a) for a in unique_alerts]\noutputs['stats'] = stats\noutputs['dedup_summary'] = summary\noutputs['input_mode'] = input_mode\n\nif enriched:\n outputs['dedup_key'] = enriched[0].get('dedup_key', '')\n outputs['is_duplicate'] = enriched[0].get('is_duplicate', False)\nelse:\n outputs['dedup_key'] = ''\n outputs['is_duplicate'] = False\n" } ], "edges": [ diff --git a/flocks/workflow/store.py b/flocks/workflow/store.py index 129d62131..95737c867 100644 --- a/flocks/workflow/store.py +++ b/flocks/workflow/store.py @@ -37,6 +37,12 @@ "workflow_syslog_config/", ) _WORKFLOW_PREFIXES = _WORKFLOW_KV_PREFIXES + _WORKFLOW_TABLE_PREFIXES +_SOC_DENOISE_WORKFLOW_ID = "stream_alert_denoise" +_METRIC_RETENTION_MS = 35 * 24 * 60 * 60 * 1000 +# Idempotency keys only need to cover realistic completion retries. Keeping +# this table bounded avoids growth proportional to high-volume syslog traffic. +_METRIC_CONTRIBUTION_KEEP = 100_000 +_METRIC_PRUNE_INTERVAL_MS = 5 * 60 * 1000 _EXECUTION_UPSERT_SQL = """ INSERT OR REPLACE INTO workflow_executions (id, workflow_id, status, current_phase, current_node_id, current_node_type, @@ -55,6 +61,7 @@ class WorkflowStore: _init_pid: Optional[int] = None _db_path: Optional[Path] = None _completion_lock: Optional[asyncio.Lock] = None + _last_metric_prune_at: int = 0 @classmethod def get_db_path(cls) -> Path: @@ -88,6 +95,7 @@ async def init(cls) -> None: cls._initialized = False cls._init_pid = None cls._completion_lock = None + cls._last_metric_prune_at = 0 await Storage._ensure_init() db_path.parent.mkdir(parents=True, exist_ok=True) @@ -151,6 +159,7 @@ async def close(cls) -> None: cls._init_pid = None cls._db_path = None cls._completion_lock = None + cls._last_metric_prune_at = 0 @classmethod async def _db(cls) -> aiosqlite.Connection: @@ -354,6 +363,182 @@ def _execution_row( ), ) + @classmethod + def _pipeline_metric_contribution(cls, exec_data: Dict[str, Any]) -> Optional[Dict[str, Any]]: + workflow_id = str(exec_data.get("workflowId") or "") + if workflow_id != _SOC_DENOISE_WORKFLOW_ID: + return None + output = exec_data.get("outputResults") + output = output if isinstance(output, dict) else {} + stats = output.get("stats") if isinstance(output.get("stats"), dict) else {} + status = str(exec_data.get("status") or "").lower() + success = status in {"success", "completed"} + error_count = 0 if success else 1 + invalid_count = 0 + + def metric_value(key: str) -> Optional[int]: + value = stats.get(key) + if isinstance(value, bool): + return None + try: + parsed = int(value) + except (TypeError, ValueError, OverflowError): + return None + return parsed if parsed >= 0 else None + + raw_count = metric_value("raw_count") + normalized_count = metric_value("normalized_count") + after_filter_count = metric_value("after_filter_count") + unique_count = metric_value("after_dedup_count") + schema_version = metric_value("metric_schema_version") or 0 + required = (raw_count, normalized_count, after_filter_count, unique_count) + valid = success and all(value is not None for value in required) + if valid and not ( + raw_count >= normalized_count >= after_filter_count >= unique_count >= 0 + ): + valid = False + if valid and schema_version < 2: + if raw_count == 1 and output.get("is_duplicate") is True: + unique_count = 0 + elif raw_count > 1: + valid = False + if success and not valid: + invalid_count = 1 + if not valid: + raw_count = normalized_count = after_filter_count = unique_count = 0 + + filter_removed_count = max(normalized_count - after_filter_count, 0) + duplicate_count = max(after_filter_count - unique_count, 0) + source_counts = {} + raw_source_counts = stats.get("normalize_type_counts") + if valid and isinstance(raw_source_counts, dict) and raw_source_counts.get("_type") != "dict": + for key, value in raw_source_counts.items(): + parsed = cls._as_int(value) + if parsed is not None and parsed > 0: + source_counts[str(key).strip().lower() or "unknown"] = parsed + source_covered_count = ( + normalized_count + if source_counts and sum(source_counts.values()) == normalized_count + else 0 + ) + started_at = cls._as_int(exec_data.get("startedAt")) or cls._now_ms() + bucket_start = started_at - (started_at % 60000) + return { + "execution_id": str(exec_data.get("id") or ""), + "workflow_id": workflow_id, + "bucket_start": bucket_start, + "raw_count": raw_count, + "normalized_count": normalized_count, + "after_filter_count": after_filter_count, + "unique_count": unique_count, + "filter_removed_count": filter_removed_count, + "duplicate_count": duplicate_count, + "source_counts": source_counts, + "source_covered_count": source_covered_count, + "success_count": 1 if success else 0, + "error_count": error_count, + "invalid_count": invalid_count, + "schema_version": schema_version, + } + + @classmethod + async def _record_pipeline_metric_contribution( + cls, + db: aiosqlite.Connection, + contribution: Optional[Dict[str, Any]], + ) -> None: + if not contribution or not contribution["execution_id"]: + return + now_ms = cls._now_ms() + cursor = await db.execute( + """ + INSERT OR IGNORE INTO workflow_metric_contributions + (execution_id, workflow_id, bucket_start, recorded_at) + VALUES (?, ?, ?, ?) + """, + ( + contribution["execution_id"], + contribution["workflow_id"], + contribution["bucket_start"], + now_ms, + ), + ) + if cursor.rowcount <= 0: + return + await db.execute( + """ + INSERT OR IGNORE INTO workflow_metric_meta + (workflow_id, coverage_started_at, updated_at) + VALUES (?, ?, ?) + """, + (contribution["workflow_id"], now_ms, now_ms), + ) + existing = await db.execute( + "SELECT source_counts FROM workflow_metric_rollups " + "WHERE workflow_id = ? AND bucket_start = ?", + (contribution["workflow_id"], contribution["bucket_start"]), + ) + existing_row = await existing.fetchone() + merged_sources = cls._json_loads(existing_row["source_counts"], {}) if existing_row else {} + if not isinstance(merged_sources, dict): + merged_sources = {} + for key, value in contribution["source_counts"].items(): + merged_sources[key] = max(cls._as_int(merged_sources.get(key)) or 0, 0) + value + await db.execute( + """ + INSERT INTO workflow_metric_rollups + (workflow_id, bucket_start, raw_count, normalized_count, + after_filter_count, unique_count, filter_removed_count, + duplicate_count, source_counts, source_covered_count, + success_count, error_count, invalid_count, schema_version, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(workflow_id, bucket_start) DO UPDATE SET + raw_count = raw_count + excluded.raw_count, + normalized_count = normalized_count + excluded.normalized_count, + after_filter_count = after_filter_count + excluded.after_filter_count, + unique_count = unique_count + excluded.unique_count, + filter_removed_count = filter_removed_count + excluded.filter_removed_count, + duplicate_count = duplicate_count + excluded.duplicate_count, + source_counts = excluded.source_counts, + source_covered_count = source_covered_count + excluded.source_covered_count, + success_count = success_count + excluded.success_count, + error_count = error_count + excluded.error_count, + invalid_count = invalid_count + excluded.invalid_count, + schema_version = MAX(schema_version, excluded.schema_version), + updated_at = excluded.updated_at + """, + ( + contribution["workflow_id"], + contribution["bucket_start"], + contribution["raw_count"], + contribution["normalized_count"], + contribution["after_filter_count"], + contribution["unique_count"], + contribution["filter_removed_count"], + contribution["duplicate_count"], + cls._json_dumps(merged_sources), + contribution["source_covered_count"], + contribution["success_count"], + contribution["error_count"], + contribution["invalid_count"], + contribution["schema_version"], + now_ms, + ), + ) + if now_ms - cls._last_metric_prune_at >= _METRIC_PRUNE_INTERVAL_MS: + cutoff = now_ms - _METRIC_RETENTION_MS + await db.execute( + "DELETE FROM workflow_metric_contributions WHERE execution_id IN (" + "SELECT execution_id FROM workflow_metric_contributions " + "ORDER BY recorded_at DESC LIMIT -1 OFFSET ?)", + (_METRIC_CONTRIBUTION_KEEP,), + ) + await db.execute( + "DELETE FROM workflow_metric_rollups WHERE bucket_start < ?", + (cutoff - (cutoff % 60000),), + ) + cls._last_metric_prune_at = now_ms + @classmethod async def upsert_execution(cls, exec_data: Dict[str, Any]) -> None: db = await cls._db() @@ -434,11 +619,13 @@ async def delete_executions_for_workflow(cls, workflow_id: str) -> int: @classmethod async def trim_executions(cls, workflow_id: str, *, keep: int) -> List[str]: + """Trim terminal history without deleting queued or running executions.""" db = await cls._db() async with db.execute( """ SELECT id FROM workflow_executions WHERE workflow_id = ? + AND status NOT IN ('running', 'queued', 'pending') ORDER BY started_at DESC, rowid DESC LIMIT -1 OFFSET ? """, @@ -509,6 +696,7 @@ async def complete_execution( """Atomically persist one final execution summary and its step batch.""" db = await cls._completion_db() exec_id, workflow_id, execution_row = cls._execution_row(exec_data) + metric_contribution = cls._pipeline_metric_contribution(exec_data) step_rows = cls._step_rows(exec_id, steps) lock = cls._completion_lock if lock is None: @@ -528,6 +716,7 @@ async def complete_execution( step_rows, ) await db.execute(_EXECUTION_UPSERT_SQL, execution_row) + await cls._record_pipeline_metric_contribution(db, metric_contribution) await db.commit() except BaseException: try: @@ -861,6 +1050,38 @@ async def kv_clear(cls, prefix: str) -> int: updated_at INTEGER ); +CREATE TABLE IF NOT EXISTS workflow_metric_rollups ( + workflow_id TEXT NOT NULL, + bucket_start INTEGER NOT NULL, + raw_count INTEGER NOT NULL DEFAULT 0, + normalized_count INTEGER NOT NULL DEFAULT 0, + after_filter_count INTEGER NOT NULL DEFAULT 0, + unique_count INTEGER NOT NULL DEFAULT 0, + filter_removed_count INTEGER NOT NULL DEFAULT 0, + duplicate_count INTEGER NOT NULL DEFAULT 0, + source_counts TEXT NOT NULL DEFAULT '{}', + source_covered_count INTEGER NOT NULL DEFAULT 0, + success_count INTEGER NOT NULL DEFAULT 0, + error_count INTEGER NOT NULL DEFAULT 0, + invalid_count INTEGER NOT NULL DEFAULT 0, + schema_version INTEGER NOT NULL DEFAULT 0, + updated_at INTEGER NOT NULL, + PRIMARY KEY (workflow_id, bucket_start) +); + +CREATE TABLE IF NOT EXISTS workflow_metric_contributions ( + execution_id TEXT PRIMARY KEY, + workflow_id TEXT NOT NULL, + bucket_start INTEGER NOT NULL, + recorded_at INTEGER NOT NULL +); + +CREATE TABLE IF NOT EXISTS workflow_metric_meta ( + workflow_id TEXT PRIMARY KEY, + coverage_started_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL +); + CREATE TABLE IF NOT EXISTS workflow_configs ( workflow_id TEXT NOT NULL, kind TEXT NOT NULL, @@ -884,4 +1105,6 @@ async def kv_clear(cls, prefix: str) -> int: "CREATE INDEX IF NOT EXISTS idx_workflow_executions_workflow_status ON workflow_executions(workflow_id, status)", "CREATE INDEX IF NOT EXISTS idx_workflow_executions_trigger ON workflow_executions(workflow_id, trigger_type, trigger_id)", "CREATE INDEX IF NOT EXISTS idx_workflow_execution_steps_exec_step ON workflow_execution_steps(exec_id, step_index)", + "CREATE INDEX IF NOT EXISTS idx_workflow_metric_rollups_workflow_bucket ON workflow_metric_rollups(workflow_id, bucket_start)", + "CREATE INDEX IF NOT EXISTS idx_workflow_metric_contributions_recorded ON workflow_metric_contributions(recorded_at)", ] diff --git a/tests/hub/test_soc_dashboard_schema.py b/tests/hub/test_soc_dashboard_schema.py index 6983b811b..deae91435 100644 --- a/tests/hub/test_soc_dashboard_schema.py +++ b/tests/hub/test_soc_dashboard_schema.py @@ -708,6 +708,240 @@ def test_soc_dashboard_activity_tolerates_empty_soc_db_with_workflow_events(tmp_ assert payload["workflowEvents"][0]["sessionId"] == "session-1" +def test_soc_dashboard_ai_tasks_use_authoritative_active_status(tmp_path: Path): + workflow_db = tmp_path / "workflow.db" + now_ms = int(datetime.now().timestamp() * 1000) + with sqlite3.connect(workflow_db) as conn: + conn.execute( + """ + CREATE TABLE workflow_executions ( + id TEXT PRIMARY KEY, + workflow_id TEXT NOT NULL, + status TEXT NOT NULL, + current_phase TEXT, + current_step_index INTEGER, + step_count INTEGER, + input_params TEXT NOT NULL DEFAULT '{}', + output_results TEXT NOT NULL DEFAULT '{}', + error_message TEXT, + started_at INTEGER NOT NULL, + finished_at INTEGER, + updated_at INTEGER, + payload TEXT NOT NULL DEFAULT '{}' + ) + """ + ) + rows = [ + ( + "triage-running", "stream_alert_triage", "running", "analysis", 2, 3, + "{}", json.dumps({"triage_results": [{"alert_name": "SSRF盲打探测"}]}), "", + now_ms - 2_000, None, now_ms, "{}", + ), + ( + "denoise-queued", "stream_alert_denoise", "queued", "queued", 0, 7, + json.dumps({"raw_alerts": [{"id": "a"}, {"id": "b"}]}), "{}", "", + now_ms - 1_000, None, 0, "{}", + ), + ( + "denoise-empty-completed", "stream_alert_denoise", "success", "completed", 7, 7, + "{}", + json.dumps({"stats": {"raw_count": 0, "normalized_count": 0, "after_filter_count": 0, "after_dedup_count": 0}}), + "", now_ms - 3_000, now_ms - 2_500, now_ms - 2_500, "{}", + ), + ( + "denoise-stale", "stream_alert_denoise", "running", "dedup", 3, 7, + "{}", "{}", "", now_ms - 3 * 60 * 60 * 1000, None, + now_ms - 3 * 60 * 60 * 1000, "{}", + ), + ] + conn.executemany( + "INSERT INTO workflow_executions VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + rows, + ) + conn.commit() + + handlers = _load_dashboard_handlers() + handlers.WORKFLOW_DB = workflow_db + + payload = handlers._get_ai_tasks() + + assert payload["connection"] == "online" + assert payload["summary"] == { + "active": 2, + "running": 1, + "waiting": 1, + "stale": 1, + "disabled": 0, + "returned": 2, + "truncated": False, + } + assert [task["executionId"] for task in payload["tasks"]] == [ + "triage-running", + "denoise-queued", + ] + assert payload["tasks"][0]["progress"]["label"] == "第 2/4 步" + assert payload["tasks"][1]["dataQuality"] == "pending" + assert payload["tasks"][1]["counts"]["raw"] == 2 + assert payload["tasks"][1]["rawCountSource"] == "workflow_input" + assert all(task["executionId"] != "denoise-empty-completed" for task in payload["tasks"]) + + +def test_soc_dashboard_ai_tasks_report_missing_workflow_database(tmp_path: Path): + handlers = _load_dashboard_handlers() + handlers.WORKFLOW_DB = tmp_path / "missing.db" + + payload = handlers._get_ai_tasks() + + assert payload["connection"] == "unavailable" + assert payload["reason"] == "workflow_db_missing" + assert payload["tasks"] == [] + + +def test_soc_dashboard_reads_persisted_denoise_metric_rollups(tmp_path: Path): + workflow_db = tmp_path / "workflow.db" + start_time = 1_800_000 + end_time = start_time + 3600 + start_ms = start_time * 1000 + with sqlite3.connect(workflow_db) as conn: + conn.execute( + "CREATE TABLE workflow_metric_meta " + "(workflow_id TEXT PRIMARY KEY, coverage_started_at INTEGER, updated_at INTEGER)" + ) + conn.execute( + """ + CREATE TABLE workflow_metric_rollups ( + workflow_id TEXT, + bucket_start INTEGER, + raw_count INTEGER, + normalized_count INTEGER, + after_filter_count INTEGER, + unique_count INTEGER, + filter_removed_count INTEGER, + duplicate_count INTEGER, + source_counts TEXT, + source_covered_count INTEGER, + success_count INTEGER, + error_count INTEGER, + invalid_count INTEGER, + schema_version INTEGER, + updated_at INTEGER, + PRIMARY KEY (workflow_id, bucket_start) + ) + """ + ) + conn.execute( + "INSERT INTO workflow_metric_meta VALUES (?, ?, ?)", + ("stream_alert_denoise", start_ms - 60_000, start_ms + 120_000), + ) + conn.executemany( + "INSERT INTO workflow_metric_rollups VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + [ + ("stream_alert_denoise", start_ms, 10, 10, 8, 3, 2, 5, '{"tdp": 10}', 10, 1, 0, 0, 2, start_ms), + ("stream_alert_denoise", start_ms + 60_000, 4, 4, 4, 1, 0, 3, '{"skyeye": 4}', 4, 1, 0, 0, 2, start_ms + 60_000), + ], + ) + conn.commit() + + handlers = _load_dashboard_handlers() + handlers.WORKFLOW_DB = workflow_db + + stats = handlers._get_workflow_denoise_stats( + "stream_alert_denoise", + start_time, + end_time, + force=True, + ) + + assert stats["metricsAvailable"] is True + assert stats["dataQuality"] == "complete" + assert stats["rawCount"] == 14 + assert stats["normalizedCount"] == 14 + assert stats["afterFilterCount"] == 12 + assert stats["uniqueCount"] == 4 + assert stats["filterRemovedCount"] == 2 + assert stats["duplicateCount"] == 8 + assert stats["sourceCounts"] == {"tdp": 10, "skyeye": 4} + assert stats["sourceCoverageRate"] == 1 + assert sum(stats["seriesRaw"]) == 14 + assert sum(stats["seriesUnique"]) == 4 + + +def test_soc_dashboard_does_not_replace_full_window_with_partial_rollup(tmp_path: Path): + workflow_db = tmp_path / "workflow.db" + start_time = 1_800_000 + end_time = start_time + 3600 + start_ms = start_time * 1000 + coverage_ms = start_ms + 180_000 + with sqlite3.connect(workflow_db) as conn: + conn.execute( + "CREATE TABLE workflow_metric_meta " + "(workflow_id TEXT PRIMARY KEY, coverage_started_at INTEGER, updated_at INTEGER)" + ) + conn.execute( + """ + CREATE TABLE workflow_metric_rollups ( + workflow_id TEXT, + bucket_start INTEGER, + raw_count INTEGER, + normalized_count INTEGER, + after_filter_count INTEGER, + unique_count INTEGER, + filter_removed_count INTEGER, + duplicate_count INTEGER, + source_counts TEXT, + source_covered_count INTEGER, + success_count INTEGER, + error_count INTEGER, + invalid_count INTEGER, + schema_version INTEGER, + updated_at INTEGER, + PRIMARY KEY (workflow_id, bucket_start) + ) + """ + ) + conn.execute( + "CREATE TABLE workflow_executions " + "(workflow_id TEXT, status TEXT, started_at INTEGER)" + ) + conn.execute( + "INSERT INTO workflow_metric_meta VALUES (?, ?, ?)", + ("stream_alert_denoise", coverage_ms, coverage_ms), + ) + conn.execute( + "INSERT INTO workflow_metric_rollups VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + ( + "stream_alert_denoise", coverage_ms, 2, 2, 2, 1, 0, 1, + '{"tdp": 2}', 2, 1, 0, 0, 2, coverage_ms, + ), + ) + conn.executemany( + "INSERT INTO workflow_executions VALUES (?, ?, ?)", + [ + ("stream_alert_denoise", "success", start_ms + offset) + for offset in (60_000, 120_000, 180_000) + ], + ) + conn.commit() + + handlers = _load_dashboard_handlers() + handlers.WORKFLOW_DB = workflow_db + + stats = handlers._get_workflow_denoise_stats( + "stream_alert_denoise", + start_time, + end_time, + force=True, + ) + + assert stats["metricsAvailable"] is False + assert stats["dataQuality"] == "legacy-partial" + assert stats["coverageComplete"] is False + assert stats["coverageStartedAt"] == coverage_ms + assert stats["callCount"] == 3 + assert stats["rawCount"] == 3 + assert stats["shadowMetricsAvailable"] is True + + def test_soc_dashboard_task_center_summarizes_tasks_and_workflows(tmp_path: Path, monkeypatch): tasks_db = tmp_path / "tasks.db" today_at_1100 = datetime.now().astimezone().replace( diff --git a/tests/workflow/test_workflow_store.py b/tests/workflow/test_workflow_store.py index b33234648..88f273efc 100644 --- a/tests/workflow/test_workflow_store.py +++ b/tests/workflow/test_workflow_store.py @@ -24,6 +24,7 @@ def _reset_state() -> None: WorkflowStore._init_pid = None WorkflowStore._db_path = None WorkflowStore._completion_lock = None + WorkflowStore._last_metric_prune_at = 0 @pytest.fixture(autouse=True) @@ -98,6 +99,106 @@ async def test_workflow_store_records_execution_steps_config_and_kv() -> None: assert await WorkflowStore.kv_list_keys("workflow_runtime/") == ["workflow_runtime/wf-1"] +@pytest.mark.asyncio +async def test_workflow_store_trim_keeps_active_executions() -> None: + await WorkflowStore.init() + for index, status in enumerate(("running", "queued", "pending", "success", "failed", "success")): + await WorkflowStore.upsert_execution( + { + "id": f"exec-{index}", + "workflowId": "wf-trim", + "status": status, + "startedAt": 100 + index, + } + ) + + trimmed = await WorkflowStore.trim_executions("wf-trim", keep=1) + remaining = await WorkflowStore.list_executions("wf-trim", limit=20) + + assert set(trimmed) == {"exec-3", "exec-4"} + assert {row["id"] for row in remaining} == {"exec-0", "exec-1", "exec-2", "exec-5"} + + +@pytest.mark.asyncio +async def test_workflow_store_rolls_up_denoise_metrics_idempotently() -> None: + await WorkflowStore.init() + now_ms = WorkflowStore._now_ms() + execution = { + "id": "denoise-exec-1", + "workflowId": "stream_alert_denoise", + "status": "success", + "startedAt": now_ms, + "finishedAt": now_ms + 1_000, + "outputResults": { + "stats": { + "metric_schema_version": 2, + "raw_count": 10, + "normalized_count": 10, + "after_filter_count": 6, + "after_dedup_count": 2, + "normalize_type_counts": {"tdp": 8, "skyeye": 2}, + } + }, + } + + await WorkflowStore.complete_execution(execution, []) + await WorkflowStore.complete_execution(execution, []) + + db = await WorkflowStore.raw_db() + async with db.execute( + "SELECT raw_count, normalized_count, after_filter_count, unique_count, " + "filter_removed_count, duplicate_count, source_counts, source_covered_count, " + "success_count, error_count, invalid_count, schema_version " + "FROM workflow_metric_rollups WHERE workflow_id = ?", + ("stream_alert_denoise",), + ) as cursor: + row = await cursor.fetchone() + async with db.execute("SELECT COUNT(*) AS total FROM workflow_metric_contributions") as cursor: + contribution_count = (await cursor.fetchone())["total"] + + assert row is not None + assert tuple(row[:6]) == (10, 10, 6, 2, 4, 4) + assert row["source_counts"] == '{"tdp": 8, "skyeye": 2}' + assert row["source_covered_count"] == 10 + assert tuple(row[8:12]) == (1, 0, 0, 2) + assert contribution_count == 1 + + +@pytest.mark.asyncio +async def test_workflow_store_marks_legacy_batch_metrics_invalid() -> None: + await WorkflowStore.init() + now_ms = WorkflowStore._now_ms() + + await WorkflowStore.complete_execution( + { + "id": "denoise-legacy-batch", + "workflowId": "stream_alert_denoise", + "status": "success", + "startedAt": now_ms, + "outputResults": { + "stats": { + "raw_count": 2, + "normalized_count": 2, + "after_filter_count": 2, + "after_dedup_count": 2, + } + }, + }, + [], + ) + + db = await WorkflowStore.raw_db() + async with db.execute( + "SELECT raw_count, success_count, invalid_count FROM workflow_metric_rollups " + "WHERE workflow_id = ?", + ("stream_alert_denoise",), + ) as cursor: + row = await cursor.fetchone() + + assert row is not None + assert tuple(row) == (0, 1, 1) + + @pytest.mark.asyncio async def test_workflow_store_increment_stats_is_atomic_for_concurrent_updates() -> None: await WorkflowStore.init() diff --git a/webui/src/utils/socDashboardPageRuntime.test.tsx b/webui/src/utils/socDashboardPageRuntime.test.tsx index 888fa93ff..a8170d7ba 100644 --- a/webui/src/utils/socDashboardPageRuntime.test.tsx +++ b/webui/src/utils/socDashboardPageRuntime.test.tsx @@ -50,6 +50,15 @@ describe('SOC dashboard contract page runtime', () => { }, }); } + if (path === '/ai-tasks') { + return Promise.resolve({ + data: { + connection: 'online', + summary: { active: 0, running: 0, waiting: 0, stale: 0 }, + tasks: [], + }, + }); + } if (path === '/task-center') { return Promise.resolve({ data: { scheduledTasks: [], workflows: [] } }); } @@ -73,6 +82,7 @@ describe('SOC dashboard contract page runtime', () => { await waitFor(() => { expect(pageGetMock).toHaveBeenCalledWith('/stats', expect.anything()); expect(pageGetMock).toHaveBeenCalledWith('/activity', expect.anything()); + expect(pageGetMock).toHaveBeenCalledWith('/ai-tasks', expect.anything()); expect(pageGetMock).toHaveBeenCalledWith('/task-center', expect.anything()); }); @@ -166,6 +176,15 @@ describe('SOC dashboard contract page runtime', () => { if (path === '/task-center') { return Promise.resolve({ data: { scheduledTasks: [], workflows: [] } }); } + if (path === '/ai-tasks') { + return Promise.resolve({ + data: { + connection: 'online', + summary: { active: 0, running: 0, waiting: 0, stale: 0 }, + tasks: [], + }, + }); + } return Promise.reject(new Error(`unexpected path: ${path}`)); }); @@ -194,6 +213,15 @@ describe('SOC dashboard contract page runtime', () => { }, }); } + if (path === '/ai-tasks') { + return Promise.resolve({ + data: { + connection: 'online', + summary: { active: 0, running: 0, waiting: 0, stale: 0 }, + tasks: [], + }, + }); + } if (path === '/task-center') { return Promise.resolve({ data: { @@ -280,6 +308,82 @@ describe('SOC dashboard contract page runtime', () => { expect(within(workflowStats).getByText('今日调用')).toBeInTheDocument(); }); + it('uses authoritative workflow task status instead of activity playback state', async () => { + const now = Date.now(); + pageGetMock.mockImplementation((path: string) => { + if (path === '/stats') return Promise.resolve({ data: {} }); + if (path === '/activity') { + return Promise.resolve({ + data: { + cursor: 'cursor', + events: [], + recentEvents: [], + workflowEvents: [ + { + eventId: 'workflow-execution:completed-history', + stage: 'denoise', + status: 'completed', + occurredAt: new Date(now).toISOString(), + triggerSource: 'workflow_execution', + workflowId: 'stream_alert_denoise', + alert: { id: 'history', threatName: '不应进入任务栏' }, + result: { isDuplicate: false, rawCount: 0 }, + }, + ], + batch: {}, + workflowStats: { callCount: 0, latestStartedAt: 0 }, + tokenUsage: { totalTokens: 0, todayTokens: 0, todayRequests: 0, dailySeries: [] }, + }, + }); + } + if (path === '/ai-tasks') { + return Promise.resolve({ + data: { + connection: 'online', + summary: { active: 2, running: 1, waiting: 1, stale: 0 }, + tasks: [ + { + taskId: 'workflow-execution:running-1', + workflowId: 'stream_alert_triage', + executionId: 'running-1', + stage: 'triage', + status: 'running', + startedAt: now, + title: 'SSRF盲打探测攻击结果未知', + counts: { raw: null }, + dataQuality: 'pending', + progress: { mode: 'steps', current: 2, total: 3, percent: 0.6667, label: '第 2/3 步' }, + }, + { + taskId: 'workflow-execution:queued-1', + workflowId: 'stream_alert_denoise', + executionId: 'queued-1', + stage: 'denoise', + status: 'queued', + startedAt: now - 1000, + title: '降噪批次 · 原始条数待生成', + counts: { raw: null }, + dataQuality: 'pending', + progress: { mode: 'waiting', percent: null, label: '等待调度' }, + }, + ], + }, + }); + } + if (path === '/task-center') return Promise.resolve({ data: { scheduledTasks: [], workflows: [] } }); + return Promise.reject(new Error(`unexpected path: ${path}`)); + }); + + render(); + + expect(await screen.findByText('正在处理 1 个，等待 1 个')).toBeInTheDocument(); + expect(screen.getByText('SSRF盲打探测攻击结果未知')).toBeInTheDocument(); + expect(screen.getByText('降噪批次 · 原始条数待生成')).toBeInTheDocument(); + expect(screen.getByText('第 2/3 步')).toBeInTheDocument(); + expect(screen.queryByText('不应进入任务栏')).not.toBeInTheDocument(); + expect(screen.queryByText('降噪处理完成')).not.toBeInTheDocument(); + }); + it('uses dashboard mock rows with the same workflow execution field shape as real task-center data', async () => { window.localStorage.setItem('soc-dashboard-mock-v1', '1'); From f0bb01eaea4d2eb77ef7cad88e1ce3e1a7d01130 Mon Sep 17 00:00:00 2001 From: John Yin <10972267+john-yin2333@user.noreply.gitee.com> Date: Mon, 7 Sep 2026 22:53:09 +0800 Subject: [PATCH 52/63] fix(soc): reject misleading dashboard zero metrics --- .../soc_ui/soc_dashboard/api/handlers.py | 156 +++++++++++++--- .../webuis/soc_ui/soc_dashboard/src/Page.tsx | 90 +++++++-- flocks/workflow/store.py | 95 +++++++++- tests/hub/test_soc_dashboard_schema.py | 151 ++++++++++++++- tests/workflow/test_workflow_store.py | 173 +++++++++++++++++- .../utils/socDashboardPageRuntime.test.tsx | 45 +++++ 6 files changed, 660 insertions(+), 50 deletions(-) diff --git a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/handlers.py b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/handlers.py index 77c974d44..e6ebf89ac 100644 --- a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/handlers.py +++ b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/handlers.py @@ -32,6 +32,7 @@ WORKFLOW_DB = Path.home() / ".flocks" / "data" / "workflow.db" WORKFLOW_SNAPSHOT_TABLE = "soc_dashboard_workflow_stats_samples" +WORKFLOW_METRIC_ROLLUP_SCHEMA_VERSION = 3 TASK_DB = Path.home() / ".flocks" / "data" / "tasks.db" USAGE_DB = Path.home() / ".flocks" / "data" / "flocks.db" SOC_PINNED_WORKFLOW_NAMES = { @@ -905,15 +906,50 @@ def _empty_workflow_denoise_stats(): "timelineLabels": [], "timelineWindow": "", "metricsAvailable": False, + "dataAvailable": True, "dataQuality": "legacy", + "unavailableReason": "", "coverageComplete": False, "coverageStartedAt": 0, "sourceCoverageRate": 0, + "sourceMetricsAvailable": False, "invalidExecutionCount": 0, + "unprocessedInputCount": 0, "dataSource": "workflow.db.workflow_stats.call_count", } +def _unavailable_workflow_denoise_stats(reason): + result = { + **_empty_workflow_denoise_stats(), + "dataAvailable": False, + "dataQuality": "unavailable", + "unavailableReason": str(reason or "workflow_metrics_unavailable"), + "dataSource": "unavailable", + } + for key in ( + "callCount", + "successCount", + "errorCount", + "earliestStartedAt", + "latestStartedAt", + "rawCount", + "normalizedCount", + "afterFilterCount", + "uniqueCount", + "filterRemovedCount", + "duplicateCount", + "reducedCount", + "reductionRate", + "dedupRate", + "sourceCoverageRate", + "invalidExecutionCount", + "unprocessedInputCount", + ): + result[key] = None + return result + + def _get_workflow_metric_rollups(workflow_name, start_time, end_time): if not WORKFLOW_DB.is_file(): return None @@ -935,10 +971,22 @@ def _get_workflow_metric_rollups(workflow_name, start_time, end_time): ).fetchone() if meta is None: return None + verified_row = conn.execute( + "SELECT MIN(bucket_start) FROM workflow_metric_rollups " + "WHERE workflow_id = ? AND schema_version >= ?", + (workflow_name, WORKFLOW_METRIC_ROLLUP_SCHEMA_VERSION), + ).fetchone() + verified_started_at = _safe_int(verified_row[0] if verified_row else 0) + if verified_started_at <= 0: + # Existing v2 rows predate input/output reconciliation and may + # contain false zero ingress counts. Keep them out of the exact + # path until a v3 contribution establishes verified coverage. + return None query = ( - "SELECT * FROM workflow_metric_rollups WHERE workflow_id = ?" + "SELECT * FROM workflow_metric_rollups " + "WHERE workflow_id = ? AND schema_version >= ?" ) - query_params = [workflow_name] + query_params = [workflow_name, WORKFLOW_METRIC_ROLLUP_SCHEMA_VERSION] if start_ms > 0 and end_ms > 0: query += " AND bucket_start >= ? AND bucket_start <= ?" query_params.extend((start_ms - (start_ms % 60000), end_ms)) @@ -964,9 +1012,19 @@ def _get_workflow_metric_rollups(workflow_name, start_time, end_time): error_count = sum(max(_safe_int(row["error_count"]), 0) for row in rows) invalid_count = sum(max(_safe_int(row["invalid_count"]), 0) for row in rows) source_covered_count = sum(max(_safe_int(row["source_covered_count"]), 0) for row in rows) - coverage_started_at = max(_safe_int(meta["coverage_started_at"]), 0) + coverage_started_at = max( + _safe_int(meta["coverage_started_at"]), + verified_started_at, + 0, + ) complete_window = not start_ms or coverage_started_at <= start_ms - quality = "complete" if complete_window and invalid_count == 0 else "partial" + source_complete = source_covered_count == raw_count + metrics_complete = invalid_count == 0 and error_count == 0 + quality = ( + "complete" + if complete_window and metrics_complete and source_complete + else "partial" + ) first_bucket = _safe_int(rows[0]["bucket_start"]) if rows else start_ms last_bucket = _safe_int(rows[-1]["bucket_start"]) if rows else end_ms @@ -1004,12 +1062,17 @@ def _get_workflow_metric_rollups(workflow_name, start_time, end_time): "seriesUnique": series_unique, "timelineLabels": labels, "timelineWindow": window, - "metricsAvailable": True, + "metricsAvailable": metrics_complete, + "sourceMetricsAvailable": source_complete, "dataQuality": quality, "coverageComplete": complete_window, "coverageStartedAt": coverage_started_at, - "sourceCoverageRate": _ratio(source_covered_count, normalized_count), + # Sources describe ingress volume on the dashboard. Use the raw + # denominator so failed/invalid executions cannot silently vanish + # from coverage while still contributing to the raw total. + "sourceCoverageRate": _ratio(min(source_covered_count, raw_count), raw_count), "invalidExecutionCount": invalid_count, + "unprocessedInputCount": max(raw_count - normalized_count, 0), "dataSource": "workflow.db.workflow_metric_rollups", } ) @@ -1032,9 +1095,8 @@ def _get_workflow_denoise_stats( _workflow_stats_cache.move_to_end(cache_key) return cached["value"] - empty = _empty_workflow_denoise_stats() if not WORKFLOW_DB.is_file(): - return empty + return _unavailable_workflow_denoise_stats("workflow_db_missing") rollup_result = _get_workflow_metric_rollups(workflow_name, start_time, end_time) if rollup_result is not None and rollup_result.get("coverageComplete"): @@ -1133,7 +1195,13 @@ def _get_workflow_denoise_stats( except Exception: with _cache_lock: cached = _workflow_stats_cache.get(cache_key) - return cached["value"] if cached else empty + if cached: + return { + **cached["value"], + "dataQuality": "stale", + "unavailableReason": "workflow_db_query_failed", + } + return _unavailable_workflow_denoise_stats("workflow_db_query_failed") def _get_workflow_progress( @@ -2090,17 +2158,31 @@ def _workflow_task_metrics(output_text, status): def _workflow_task_input_count(inputs): - for key in ("_raw_alerts_count", "raw_count"): + for key in ("_raw_alerts_count", "_alerts_count", "_alert_list_count", "raw_count"): value, quality = _workflow_task_metric(inputs, key) if quality == "complete": return value - for key in ("raw_alerts", "alerts"): + for key in ("raw_alerts", "alerts", "alert_list"): value = inputs.get(key) if isinstance(value, list): return len(value) - for key in ("syslog_message", "syslog", "alert"): - if inputs.get(key) not in (None, "", {}): - return 1 + if isinstance(value, dict) and value.get("_type") in {"list", "tuple", "set"}: + count, quality = _workflow_task_metric(value, "count") + if quality == "complete": + return count + if isinstance(value, dict) and isinstance(value.get("data"), list): + return len(value["data"]) + for key in ("syslog_message", "syslog"): + value = inputs.get(key) + if not isinstance(value, dict) or not value.get("message"): + continue + try: + json.loads(str(value["message"])) + except (TypeError, ValueError, json.JSONDecodeError): + continue + return 1 + if inputs.get("alert") not in (None, "", {}): + return 1 return None @@ -2117,13 +2199,17 @@ def _workflow_task_row(row, workflow_id, effective_status): if input_count is not None: counts["raw"] = input_count raw_count_source = "workflow_input" + if input_count == 0: + data_quality = "empty-input" preview = metrics["preview"] stage = "triage" if workflow_id in TRIAGE_WORKFLOW_IDS else "denoise" title = _workflow_latest_alert_name(workflow_id, output_text, input_text) if stage == "denoise" and not preview: raw_count = counts["raw"] title = ( - f"降噪批次 · 原始 {raw_count} 条" + "降噪批次 · 空输入" + if data_quality == "empty-input" + else f"降噪批次 · 原始 {raw_count} 条" if raw_count is not None else "降噪批次 · 原始条数待生成" ) @@ -2164,7 +2250,8 @@ def _workflow_task_row(row, workflow_id, effective_status): "counts": counts, "dataQuality": data_quality, "rawCountSource": raw_count_source, - "emptyBatch": effective_status in WORKFLOW_SUCCESS_STATUSES and counts["raw"] == 0, + "emptyBatch": counts["raw"] == 0, + "emptyInput": data_quality == "empty-input", "progress": progress, "sessionId": session_id, "messageId": message_id, @@ -2389,6 +2476,13 @@ def _get_activity(params): last_row_id = max(_safe_int(cursor.get("lastRowId")), 0) last_activity_id = max(_safe_int(cursor.get("lastActivityId")), 0) + previous_polled_at = max(_safe_int(cursor.get("polledAt")), 0) + current_polled_at = int(time.time() * 1000) + poll_window_ms = ( + max(current_polled_at - previous_polled_at, 1) + if previous_polled_at + else ACTIVITY_WINDOW_MS + ) if last_row_id > latest_row_id or last_activity_id > latest_activity_id: last_row_id = 0 last_activity_id = 0 @@ -2401,6 +2495,7 @@ def _get_activity(params): latest_row_id=latest_row_id, latest_activity_id=latest_activity_id, limit=limit, + window_ms=poll_window_ms, ) except Exception as exc: return { @@ -2448,9 +2543,11 @@ def _activity_response( workflow_stats=None, workflow_events=None, ): + generated_at = datetime.now().astimezone() + polled_at = int(generated_at.timestamp() * 1000) return { - "cursor": _encode_activity_cursor(last_row_id, last_activity_id), - "generatedAt": datetime.now().astimezone().isoformat(timespec="seconds"), + "cursor": _encode_activity_cursor(last_row_id, last_activity_id, polled_at), + "generatedAt": generated_at.isoformat(timespec="seconds"), "events": events, "recentEvents": recent_events or [], "overflowCount": 0, @@ -2477,11 +2574,12 @@ def _empty_activity_batch(): } -def _encode_activity_cursor(last_row_id, last_activity_id): +def _encode_activity_cursor(last_row_id, last_activity_id, polled_at=None): payload = json.dumps( { "lastRowId": max(_safe_int(last_row_id), 0), "lastActivityId": max(_safe_int(last_activity_id), 0), + "polledAt": max(_safe_int(polled_at), 0), }, separators=(",", ":"), ).encode("utf-8") @@ -2520,6 +2618,7 @@ def _activity_rows( latest_row_id, latest_activity_id, limit, + window_ms=ACTIVITY_WINDOW_MS, ): summary = _activity_insert_summary(conn, settings, last_row_id, latest_row_id) new_count = summary["receivedCount"] @@ -2561,11 +2660,11 @@ def _activity_rows( if new_count > ACTIVITY_SURGE_LIMIT else "burst" if new_count > ACTIVITY_NORMAL_LIMIT else "normal" ), - "windowMs": ACTIVITY_WINDOW_MS, + "windowMs": max(_safe_int(window_ms), 1), "triageUpdatedCount": updated_count, "sampledCount": sampled_count, "suppressedCount": max(new_count - sampled_count, 0), - "ratePerSecond": round(new_count / (ACTIVITY_WINDOW_MS / 1000), 1), + "ratePerSecond": round(new_count / (max(_safe_int(window_ms), 1) / 1000), 1), } return rows, max(new_count + updated_count - len(rows), 0), batch @@ -2869,7 +2968,7 @@ def _get_stats(params): range_end_time, force=force_refresh, ) - denoise = _read_denoise(denoise_files, workflow_stats.get("callCount", 0)) + denoise = _read_denoise(denoise_files, workflow_stats.get("callCount") or 0) soc_unique_count = denoise["totalUnique"] soc_unique_series = denoise["seriesUnique"] timeline_labels = workflow_stats["timelineLabels"] or denoise.get("_timelineLabels", []) @@ -2885,7 +2984,11 @@ def _get_stats(params): duplicate_count = workflow_stats["duplicateCount"] workflow_series_unique = workflow_stats["seriesUnique"] else: - processed_total = workflow_stats["callCount"] + # Legacy/unavailable data is kept only as an internal compatibility + # fallback. The quality contract tells the UI not to present it as an + # authoritative metric; unavailable fields themselves remain null in + # sourceStatus.workflowStats. + processed_total = max(_safe_int(workflow_stats.get("callCount")), 0) normalized_total = processed_total after_filter_total = processed_total unique_total = soc_unique_count @@ -2908,7 +3011,7 @@ def _get_stats(params): "duplicateRate": reduction_rate, "dedupRate": _ratio(duplicate_count, after_filter_total), "uniqueRate": _ratio(min(unique_total, processed_total), processed_total), - "files": workflow_stats["callCount"], + "files": workflow_stats.get("callCount"), "sourceCounter": Counter(workflow_stats["sourceCounts"]), "seriesRaw": workflow_series_raw, "seriesUnique": workflow_series_unique, @@ -2951,10 +3054,15 @@ def _get_stats(params): "workflowStats": workflow_stats, "metricQuality": { "status": workflow_stats.get("dataQuality", "legacy"), + "dataAvailable": workflow_stats.get("dataAvailable", True), + "unavailableReason": workflow_stats.get("unavailableReason", ""), "coverageComplete": workflow_stats.get("coverageComplete", False), "coverageStartedAt": workflow_stats.get("coverageStartedAt", 0), "sourceCoverageRate": workflow_stats.get("sourceCoverageRate", 0), + "sourceMetricsAvailable": workflow_stats.get("sourceMetricsAvailable", False), "invalidExecutionCount": workflow_stats.get("invalidExecutionCount", 0), + "unprocessedInputCount": workflow_stats.get("unprocessedInputCount", 0), + "errorExecutionCount": workflow_stats.get("errorCount", 0), "metricsAvailable": metrics_available, }, "sampleMode": sample_mode, diff --git a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx index 8361c0423..a441ffdda 100644 --- a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx +++ b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx @@ -775,7 +775,6 @@ function refreshLabel(value) { } function mergeStats(raw) { - const denoise = { ...EMPTY_STATS.denoise, ...((raw || {}).denoise || {}) }; const sourceStatus = { ...EMPTY_STATS.sourceStatus, ...((raw || {}).sourceStatus || {}), @@ -784,24 +783,52 @@ function mergeStats(raw) { ...((raw || {}).sourceStatus?.metricQuality || {}), }, }; + const metricQuality = sourceStatus.metricQuality || {}; + const denoiseMetricsUnavailable = metricQuality.metricsAvailable === false; + const sourceMetricsUnavailable = denoiseMetricsUnavailable + || metricQuality.sourceMetricsAvailable === false; + const denoise = { ...EMPTY_STATS.denoise, ...((raw || {}).denoise || {}) }; + if (denoiseMetricsUnavailable) { + for (const key of [ + 'totalRaw', 'totalNormalized', 'afterFilter', 'totalUnique', + 'filterRemoved', 'dedupRemoved', 'duplicates', 'duplicateRate', + 'dedupRate', 'uniqueRate', + ]) denoise[key] = null; + if (metricQuality.dataAvailable === false) denoise.files = null; + } + const pipeline = { ...EMPTY_STATS.pipeline, ...((raw || {}).pipeline || {}) }; + if (denoiseMetricsUnavailable) { + for (const key of ['raw', 'unique', 'reductionSaved', 'uniqueRate', 'coverageRate']) { + pipeline[key] = null; + } + } + const timeline = { ...EMPTY_STATS.timeline, ...((raw || {}).timeline || {}) }; + if (denoiseMetricsUnavailable) { + timeline.denoiseRaw = []; + timeline.denoiseUnique = []; + } + const sources = Array.isArray((raw || {}).sources) ? raw.sources : []; return { ...EMPTY_STATS, ...(raw || {}), sourceStatus, denoise, triage: { ...EMPTY_STATS.triage, ...((raw || {}).triage || {}) }, - pipeline: { ...EMPTY_STATS.pipeline, ...((raw || {}).pipeline || {}) }, + pipeline, closedLoop: { ...EMPTY_STATS.closedLoop, ...((raw || {}).closedLoop || {}) }, tokenUsage: { ...EMPTY_STATS.tokenUsage, ...((raw || {}).tokenUsage || {}) }, dateRange: { ...EMPTY_STATS.dateRange, ...((raw || {}).dateRange || {}) }, eventRange: { ...EMPTY_STATS.eventRange, ...((raw || {}).eventRange || {}) }, - timeline: { ...EMPTY_STATS.timeline, ...((raw || {}).timeline || {}) }, - sources: Array.isArray((raw || {}).sources) ? raw.sources : [], + timeline, + sources: sourceMetricsUnavailable + ? sources.map((item) => ({ ...item, value: null, rate: null, active: false })) + : sources, }; } function fullNumber(value) { - const n = Number(value || 0); + if (value === null || value === undefined || !Number.isFinite(Number(value))) return '--'; + const n = Number(value); return new Intl.NumberFormat('zh-CN').format(n); } @@ -837,7 +864,8 @@ function workflowDenoiseActivity(callCount, delta, generatedAt, workflowEvent) { } function compactNumber(value) { - const n = Number(value || 0); + if (value === null || value === undefined || !Number.isFinite(Number(value))) return '--'; + const n = Number(value); if (Math.abs(n) >= 100000000) return `${trim(n / 100000000)}亿`; if (Math.abs(n) >= 10000) return `${trim(n / 10000)}万`; return fullNumber(n); @@ -851,7 +879,8 @@ function formatTokenVolume(value) { function AnimatedNumber({ value, format, tag = 'span', className, duration = 900 }) { const { useEffect, useRef, useState } = getReact(); - const target = Number(value || 0); + const unavailable = value === null || value === undefined || !Number.isFinite(Number(value)); + const target = unavailable ? 0 : Number(value); const current = useRef(0); const [display, setDisplay] = useState(0); const formatter = format || ((number) => compactNumber(Math.round(number))); @@ -879,8 +908,8 @@ function AnimatedNumber({ value, format, tag = 'span', className, duration = 900 return h(tag, { className: cx('animated-number', className), - title: formatter(target), - }, formatter(display)); + title: unavailable ? '数据不可用' : formatter(target), + }, unavailable ? '--' : formatter(display)); } function trim(value) { @@ -888,7 +917,8 @@ function trim(value) { } function pct(value) { - return `${Math.round(Number(value || 0) * 1000) / 10}%`; + if (value === null || value === undefined || !Number.isFinite(Number(value))) return '--'; + return `${Math.round(Number(value) * 1000) / 10}%`; } function ratio(part, total) { @@ -980,7 +1010,9 @@ function SourceColumn({ stats }) { return h('div', { className: 'column left-col' }, [ h(Panel, { key: 'sources', title: '多源告警接入', meta: `${stats.denoise.files || 0} 个降噪批次` }, [ h('div', { className: 'source-list', key: 'list' }, stats.sources.map((item) => { - const width = `${Math.max(4, Math.round((item.rate || 0) * 100))}%`; + const width = item.value === null || item.value === undefined + ? '0%' + : `${Math.max(4, Math.round((item.rate || 0) * 100))}%`; return h('div', { className: 'source-row', key: item.key }, [ h('div', { className: cx('source-node', item.active && 'active'), key: 'node' }), h('div', { className: 'source-main', key: 'main' }, [ @@ -1996,7 +2028,7 @@ function CommandMetrics({ stats }) { return h('section', { className: 'command-metrics' }, [ h(CommandMetric, { label: '原始告警量', value: stats.denoise.totalRaw, sub: `${compactNumber(stats.denoise.totalUnique)} 条进入研判`, values: stats.timeline.denoiseRaw, color: '#2e72ff', key: 'raw' }), h(CommandMetric, { label: '安全事件量', value: stats.triage.attackTotal, sub: `${compactNumber(stats.triage.attackSuccess)} 条攻击成功`, values: stats.timeline.triageAttack, color: '#23ca8e', key: 'events' }), - h(CommandMetric, { label: '降噪率', value: stats.denoise.duplicateRate * 100, format: (value) => `${trim(value)}%`, sub: `${compactNumber(stats.denoise.duplicates)} 条告警已过滤/收敛`, values: stats.timeline.denoiseUnique, color: '#21d8a3', key: 'rate' }), + h(CommandMetric, { label: '降噪率', value: stats.denoise.duplicateRate === null ? null : stats.denoise.duplicateRate * 100, format: (value) => `${trim(value)}%`, sub: `${compactNumber(stats.denoise.duplicates)} 条告警已过滤/收敛`, values: stats.timeline.denoiseUnique, color: '#21d8a3', key: 'rate' }), h(TokenUsageMetric, { tokenUsage, key: 'tokens' }), ]); } @@ -2524,12 +2556,14 @@ function CommandAiTaskPanel({ aiTasks }) { const stateLabel = processing ? '处理中' : '等待处理'; const qualityDetail = task.dataQuality === 'invalid' ? ' · 指标格式异常' + : task.dataQuality === 'empty-input' + ? ' · 空输入任务' : task.dataQuality === 'missing' ? ' · 原始条数未知' : task.dataQuality === 'pending' && (task.counts?.raw === null || task.counts?.raw === undefined) ? ' · 原始条数待生成' : ''; - const rawDetail = task.stage === 'denoise' && task.counts?.raw !== null && task.counts?.raw !== undefined + const rawDetail = task.stage === 'denoise' && !task.emptyInput && task.counts?.raw !== null && task.counts?.raw !== undefined ? ` · 原始 ${task.counts.raw} 条` : ''; const detail = task.stage === 'triage' @@ -3071,17 +3105,37 @@ export default function Page() { || displayActivity.batch?.triageUpdatedCount ); const metricQuality = stats.sourceStatus?.metricQuality || {}; + const metricIssues = []; + if (Number(metricQuality.invalidExecutionCount || 0) > 0) { + metricIssues.push(`${metricQuality.invalidExecutionCount} 次指标格式异常`); + } + if (Number(metricQuality.errorExecutionCount || 0) > 0) { + metricIssues.push(`${metricQuality.errorExecutionCount} 次执行失败`); + } + if (Number(metricQuality.unprocessedInputCount || 0) > 0) { + metricIssues.push(`${metricQuality.unprocessedInputCount} 条输入未完成归一化`); + } + const sourceCoverageRate = Number(metricQuality.sourceCoverageRate); + if (metricQuality.metricsAvailable && Number.isFinite(sourceCoverageRate) && sourceCoverageRate < 1) { + metricIssues.push(`来源覆盖 ${Math.round(sourceCoverageRate * 1000) / 10}%`); + } const metricQualityWarning = !stats.generatedAt ? '' + : metricQuality.dataAvailable === false || metricQuality.status === 'unavailable' + ? '降噪统计数据源不可用，相关数字已隐藏；系统正在重试' + : metricQuality.status === 'stale' + ? '降噪统计查询失败，当前显示上次缓存数据，可能已过期' + : metricQuality.status === 'partial' && !metricQuality.metricsAvailable + ? `降噪指标校验未通过，相关数字已隐藏${metricIssues.length ? `：${metricIssues.join('；')}` : ''}` : metricQuality.metricsAvailable ? metricQuality.status === 'partial' - ? `降噪指标为部分覆盖数据${metricQuality.coverageStartedAt ? `，完整采集始于 ${taskCenterTimeLabel(metricQuality.coverageStartedAt)}` : ''}` - : Number(metricQuality.invalidExecutionCount || 0) > 0 - ? `${metricQuality.invalidExecutionCount} 次降噪执行的指标格式异常，未计入统计` + ? `降噪指标部分可用${metricIssues.length ? `：${metricIssues.join('；')}` : ''}${metricQuality.coverageStartedAt ? `；完整采集始于 ${taskCenterTimeLabel(metricQuality.coverageStartedAt)}` : ''}` + : metricIssues.length + ? `降噪指标存在异常：${metricIssues.join('；')}` : '' : metricQuality.status === 'legacy-partial' - ? `降噪指标仍使用历史兼容口径；精确口径完整采集始于 ${taskCenterTimeLabel(metricQuality.coverageStartedAt)}` - : '降噪指标仍使用历史兼容口径；新指标链路产生数据后将自动切换'; + ? `精确降噪指标尚未覆盖当前时间范围，相关数字已隐藏；完整采集始于 ${taskCenterTimeLabel(metricQuality.coverageStartedAt)}` + : '精确降噪指标尚不可用，相关数字已隐藏；新指标链路产生数据后将自动显示'; return h('div', { className: cx('adtd-root command-root', displayActivityBusy && 'command-is-processing', eventRailCollapsed && 'event-rail-is-collapsed'), diff --git a/flocks/workflow/store.py b/flocks/workflow/store.py index 95737c867..ef87b1969 100644 --- a/flocks/workflow/store.py +++ b/flocks/workflow/store.py @@ -38,6 +38,10 @@ ) _WORKFLOW_PREFIXES = _WORKFLOW_KV_PREFIXES + _WORKFLOW_TABLE_PREFIXES _SOC_DENOISE_WORKFLOW_ID = "stream_alert_denoise" +# Version 3 means the persisted contribution was validated against the +# independently counted workflow input. Older v2 rollups may contain false +# zero ingress values and must not be treated as authoritative by the UI. +_SOC_METRIC_ROLLUP_SCHEMA_VERSION = 3 _METRIC_RETENTION_MS = 35 * 24 * 60 * 60 * 1000 # Idempotency keys only need to cover realistic completion retries. Keeping # this table bounded avoids growth proportional to high-volume syslog traffic. @@ -376,6 +380,66 @@ def _pipeline_metric_contribution(cls, exec_data: Dict[str, Any]) -> Optional[Di error_count = 0 if success else 1 invalid_count = 0 + def sequence_count(value: Any) -> Optional[int]: + if isinstance(value, (list, tuple)): + return len(value) + if isinstance(value, dict): + if value.get("_type") in {"list", "tuple", "set"}: + count = cls._as_int(value.get("count")) + return count if count is not None and count >= 0 else None + data = value.get("data") + if isinstance(data, (list, tuple)): + return len(data) + if isinstance(value, str): + try: + return sequence_count(json.loads(value)) + except (TypeError, ValueError, json.JSONDecodeError): + return None + return None + + def input_alert_count() -> Optional[int]: + inputs = exec_data.get("inputParams") + if not isinstance(inputs, dict): + return None + + # Match the workflow's actual input priority. A syslog message is + # an alert only when its JSON payload can be decoded by the receive + # node; malformed/non-empty text must not be counted as accepted. + syslog_message = inputs.get("syslog_message") or inputs.get("syslog") + if isinstance(syslog_message, dict) and syslog_message.get("message"): + try: + json.loads(str(syslog_message["message"])) + except (TypeError, ValueError, json.JSONDecodeError): + pass + else: + return 1 + + known_empty = False + for key in ("raw_alerts", "alerts", "alert_list"): + marker = cls._as_int(inputs.get(f"_{key}_count")) + if marker is not None and marker >= 0: + if marker > 0: + return marker + known_empty = True + count = sequence_count(inputs.get(key)) + if count is not None: + if count > 0: + return count + known_empty = True + + # File inputs are intentionally unknown here: reading a user file + # while committing execution state would introduce I/O and TOCTOU + # races. Successful output metrics remain authoritative for them. + if inputs.get("alert_file"): + return None + return 0 if known_empty else None + + input_count = input_alert_count() + input_params = exec_data.get("inputParams") + has_unverified_file_input = ( + isinstance(input_params, dict) and bool(input_params.get("alert_file")) + ) + def metric_value(key: str) -> Optional[int]: value = stats.get(key) if isinstance(value, bool): @@ -390,14 +454,21 @@ def metric_value(key: str) -> Optional[int]: normalized_count = metric_value("normalized_count") after_filter_count = metric_value("after_filter_count") unique_count = metric_value("after_dedup_count") - schema_version = metric_value("metric_schema_version") or 0 + reported_schema_version = metric_value("metric_schema_version") or 0 required = (raw_count, normalized_count, after_filter_count, unique_count) valid = success and all(value is not None for value in required) if valid and not ( raw_count >= normalized_count >= after_filter_count >= unique_count >= 0 ): valid = False - if valid and schema_version < 2: + if valid and input_count is not None and raw_count != input_count: + valid = False + if valid and raw_count == 0 and input_count is None and has_unverified_file_input: + # A configured file cannot be safely re-read during persistence. + # Treat a zero output as unverifiable instead of claiming the file + # contained no alerts (it may have failed to load or changed). + valid = False + if valid and reported_schema_version < 2: if raw_count == 1 and output.get("is_duplicate") is True: unique_count = 0 elif raw_count > 1: @@ -405,7 +476,11 @@ def metric_value(key: str) -> Optional[int]: if success and not valid: invalid_count = 1 if not valid: - raw_count = normalized_count = after_filter_count = unique_count = 0 + # Preserve independently verifiable ingress volume even when a + # workflow fails or emits malformed/inconsistent stage metrics. + # Downstream stages remain zero because they were not verified. + raw_count = input_count or 0 + normalized_count = after_filter_count = unique_count = 0 filter_removed_count = max(normalized_count - after_filter_count, 0) duplicate_count = max(after_filter_count - unique_count, 0) @@ -438,7 +513,7 @@ def metric_value(key: str) -> Optional[int]: "success_count": 1 if success else 0, "error_count": error_count, "invalid_count": invalid_count, - "schema_version": schema_version, + "schema_version": _SOC_METRIC_ROLLUP_SCHEMA_VERSION, } @classmethod @@ -473,6 +548,18 @@ async def _record_pipeline_metric_contribution( """, (contribution["workflow_id"], now_ms, now_ms), ) + # A v2 bucket may already contain pre-reconciliation counts from the + # same minute. Do not let ON CONFLICT upgrade that mixed bucket to v3; + # replace only that obsolete derived bucket before adding verified data. + await db.execute( + "DELETE FROM workflow_metric_rollups " + "WHERE workflow_id = ? AND bucket_start = ? AND schema_version < ?", + ( + contribution["workflow_id"], + contribution["bucket_start"], + _SOC_METRIC_ROLLUP_SCHEMA_VERSION, + ), + ) existing = await db.execute( "SELECT source_counts FROM workflow_metric_rollups " "WHERE workflow_id = ? AND bucket_start = ?", diff --git a/tests/hub/test_soc_dashboard_schema.py b/tests/hub/test_soc_dashboard_schema.py index deae91435..2db851ade 100644 --- a/tests/hub/test_soc_dashboard_schema.py +++ b/tests/hub/test_soc_dashboard_schema.py @@ -797,6 +797,38 @@ def test_soc_dashboard_ai_tasks_report_missing_workflow_database(tmp_path: Path) assert payload["tasks"] == [] +def test_soc_dashboard_ai_task_input_count_supports_compacted_batches_and_empty_input(): + handlers = _load_dashboard_handlers() + + assert handlers._workflow_task_input_count( + {"alerts": {"_type": "list", "count": 2500}} + ) == 2500 + assert handlers._workflow_task_input_count({"_raw_alerts_count": 4000}) == 4000 + assert handlers._workflow_task_input_count({"alerts": []}) == 0 + assert handlers._workflow_task_input_count( + {"syslog_message": {"message": "not-json"}} + ) is None + + +def test_soc_dashboard_marks_missing_workflow_metrics_as_unavailable(tmp_path: Path): + handlers = _load_dashboard_handlers() + handlers.WORKFLOW_DB = tmp_path / "missing-workflow.db" + + stats = handlers._get_workflow_denoise_stats( + "stream_alert_denoise", + 1_800_000, + 1_803_600, + force=True, + ) + + assert stats["dataAvailable"] is False + assert stats["metricsAvailable"] is False + assert stats["dataQuality"] == "unavailable" + assert stats["unavailableReason"] == "workflow_db_missing" + assert stats["rawCount"] is None + assert stats["callCount"] is None + + def test_soc_dashboard_reads_persisted_denoise_metric_rollups(tmp_path: Path): workflow_db = tmp_path / "workflow.db" start_time = 1_800_000 @@ -836,8 +868,11 @@ def test_soc_dashboard_reads_persisted_denoise_metric_rollups(tmp_path: Path): conn.executemany( "INSERT INTO workflow_metric_rollups VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", [ - ("stream_alert_denoise", start_ms, 10, 10, 8, 3, 2, 5, '{"tdp": 10}', 10, 1, 0, 0, 2, start_ms), - ("stream_alert_denoise", start_ms + 60_000, 4, 4, 4, 1, 0, 3, '{"skyeye": 4}', 4, 1, 0, 0, 2, start_ms + 60_000), + # Pre-v3 rows were not reconciled with input volume and must + # not contaminate authoritative dashboard totals. + ("stream_alert_denoise", start_ms + 120_000, 999, 999, 999, 999, 0, 0, '{"tdp": 999}', 999, 1, 0, 0, 2, start_ms + 120_000), + ("stream_alert_denoise", start_ms, 10, 10, 8, 3, 2, 5, '{"tdp": 10}', 10, 1, 0, 0, 3, start_ms), + ("stream_alert_denoise", start_ms + 60_000, 4, 4, 4, 1, 0, 3, '{"skyeye": 4}', 4, 1, 0, 0, 3, start_ms + 60_000), ], ) conn.commit() @@ -853,6 +888,7 @@ def test_soc_dashboard_reads_persisted_denoise_metric_rollups(tmp_path: Path): ) assert stats["metricsAvailable"] is True + assert stats["sourceMetricsAvailable"] is True assert stats["dataQuality"] == "complete" assert stats["rawCount"] == 14 assert stats["normalizedCount"] == 14 @@ -866,6 +902,115 @@ def test_soc_dashboard_reads_persisted_denoise_metric_rollups(tmp_path: Path): assert sum(stats["seriesUnique"]) == 4 +def test_soc_dashboard_separates_core_metric_and_source_coverage_quality(tmp_path: Path): + workflow_db = tmp_path / "workflow.db" + start_time = 1_800_000 + end_time = start_time + 3600 + start_ms = start_time * 1000 + with sqlite3.connect(workflow_db) as conn: + conn.execute( + "CREATE TABLE workflow_metric_meta " + "(workflow_id TEXT PRIMARY KEY, coverage_started_at INTEGER, updated_at INTEGER)" + ) + conn.execute( + """ + CREATE TABLE workflow_metric_rollups ( + workflow_id TEXT, bucket_start INTEGER, raw_count INTEGER, + normalized_count INTEGER, after_filter_count INTEGER, unique_count INTEGER, + filter_removed_count INTEGER, duplicate_count INTEGER, source_counts TEXT, + source_covered_count INTEGER, success_count INTEGER, error_count INTEGER, + invalid_count INTEGER, schema_version INTEGER, updated_at INTEGER, + PRIMARY KEY (workflow_id, bucket_start) + ) + """ + ) + conn.execute( + "INSERT INTO workflow_metric_meta VALUES (?, ?, ?)", + ("stream_alert_denoise", start_ms - 60_000, start_ms), + ) + conn.execute( + "INSERT INTO workflow_metric_rollups VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + ( + "stream_alert_denoise", start_ms, 10, 6, 5, 4, 1, 1, + '{"tdp": 6}', 6, 1, 0, 0, 3, start_ms, + ), + ) + conn.commit() + + handlers = _load_dashboard_handlers() + handlers.WORKFLOW_DB = workflow_db + + stats = handlers._get_workflow_denoise_stats( + "stream_alert_denoise", start_time, end_time, force=True + ) + + assert stats["metricsAvailable"] is True + assert stats["sourceMetricsAvailable"] is False + assert stats["dataQuality"] == "partial" + assert stats["sourceCoverageRate"] == 0.6 + assert stats["unprocessedInputCount"] == 4 + + with sqlite3.connect(workflow_db) as conn: + conn.execute( + "UPDATE workflow_metric_rollups SET error_count = 1 " + "WHERE workflow_id = ?", + ("stream_alert_denoise",), + ) + conn.commit() + + failed_stats = handlers._get_workflow_denoise_stats( + "stream_alert_denoise", start_time, end_time, force=True + ) + + assert failed_stats["metricsAvailable"] is False + assert failed_stats["dataQuality"] == "partial" + + +def test_soc_dashboard_activity_cursor_tracks_actual_poll_window(): + handlers = _load_dashboard_handlers() + + encoded = handlers._encode_activity_cursor(12, 34, 5_000) + decoded = handlers._decode_activity_cursor(encoded) + + assert decoded == {"lastRowId": 12, "lastActivityId": 34, "polledAt": 5_000} + + +def test_soc_dashboard_activity_rate_uses_actual_poll_window(): + handlers = _load_dashboard_handlers() + settings = { + "table": "alerts", + "activity_table": "activity", + "record_column": "record_json", + "event_time_column": "event_time", + } + with sqlite3.connect(":memory:") as conn: + conn.row_factory = sqlite3.Row + conn.execute("CREATE TABLE alerts (record_json TEXT, event_time INTEGER)") + conn.execute( + "CREATE TABLE activity " + "(activity_id INTEGER PRIMARY KEY, alert_row_id INTEGER, event_time INTEGER, record_json TEXT)" + ) + conn.executemany( + "INSERT INTO alerts VALUES (?, ?)", + [(json.dumps({"id": index}), 100 + index) for index in range(4)], + ) + + _, _, batch = handlers._activity_rows( + conn, + settings, + last_row_id=0, + last_activity_id=0, + latest_row_id=4, + latest_activity_id=0, + limit=10, + window_ms=8_000, + ) + + assert batch["receivedCount"] == 4 + assert batch["windowMs"] == 8_000 + assert batch["ratePerSecond"] == 0.5 + + def test_soc_dashboard_does_not_replace_full_window_with_partial_rollup(tmp_path: Path): workflow_db = tmp_path / "workflow.db" start_time = 1_800_000 @@ -911,7 +1056,7 @@ def test_soc_dashboard_does_not_replace_full_window_with_partial_rollup(tmp_path "INSERT INTO workflow_metric_rollups VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", ( "stream_alert_denoise", coverage_ms, 2, 2, 2, 1, 0, 1, - '{"tdp": 2}', 2, 1, 0, 0, 2, coverage_ms, + '{"tdp": 2}', 2, 1, 0, 0, 3, coverage_ms, ), ) conn.executemany( diff --git a/tests/workflow/test_workflow_store.py b/tests/workflow/test_workflow_store.py index 88f273efc..f8a21978a 100644 --- a/tests/workflow/test_workflow_store.py +++ b/tests/workflow/test_workflow_store.py @@ -160,10 +160,181 @@ async def test_workflow_store_rolls_up_denoise_metrics_idempotently() -> None: assert tuple(row[:6]) == (10, 10, 6, 2, 4, 4) assert row["source_counts"] == '{"tdp": 8, "skyeye": 2}' assert row["source_covered_count"] == 10 - assert tuple(row[8:12]) == (1, 0, 0, 2) + assert tuple(row[8:12]) == (1, 0, 0, 3) assert contribution_count == 1 +@pytest.mark.asyncio +async def test_workflow_store_replaces_legacy_counts_in_the_first_verified_bucket() -> None: + await WorkflowStore.init() + now_ms = WorkflowStore._now_ms() + bucket_start = now_ms - (now_ms % 60_000) + db = await WorkflowStore.raw_db() + await db.execute( + """ + INSERT INTO workflow_metric_rollups + (workflow_id, bucket_start, raw_count, normalized_count, + after_filter_count, unique_count, filter_removed_count, + duplicate_count, source_counts, source_covered_count, + success_count, error_count, invalid_count, schema_version, updated_at) + VALUES (?, ?, 999, 999, 999, 999, 0, 0, ?, 999, 1, 0, 0, 2, ?) + """, + ("stream_alert_denoise", bucket_start, '{"tdp": 999}', now_ms), + ) + await db.commit() + + await WorkflowStore.complete_execution( + { + "id": "denoise-first-v3", + "workflowId": "stream_alert_denoise", + "status": "success", + "startedAt": now_ms, + "inputParams": {"alerts": [{"id": "a"}]}, + "outputResults": { + "stats": { + "metric_schema_version": 2, + "raw_count": 1, + "normalized_count": 1, + "after_filter_count": 1, + "after_dedup_count": 1, + "normalize_type_counts": {"tdp": 1}, + } + }, + }, + [], + ) + + async with db.execute( + "SELECT raw_count, source_counts, schema_version FROM workflow_metric_rollups " + "WHERE workflow_id = ? AND bucket_start = ?", + ("stream_alert_denoise", bucket_start), + ) as cursor: + row = await cursor.fetchone() + + assert row is not None + assert tuple(row) == (1, '{"tdp": 1}', 3) + + +@pytest.mark.asyncio +async def test_workflow_store_preserves_input_volume_when_metrics_are_invalid_or_failed() -> None: + await WorkflowStore.init() + now_ms = WorkflowStore._now_ms() + + await WorkflowStore.complete_execution( + { + "id": "denoise-invalid-output", + "workflowId": "stream_alert_denoise", + "status": "success", + "startedAt": now_ms, + "inputParams": {"alerts": [{"id": "a"}, {"id": "b"}, {"id": "c"}]}, + "outputResults": { + "stats": { + "metric_schema_version": 2, + "raw_count": 0, + "normalized_count": 0, + "after_filter_count": 0, + "after_dedup_count": 0, + } + }, + }, + [], + ) + await WorkflowStore.complete_execution( + { + "id": "denoise-failed-input", + "workflowId": "stream_alert_denoise", + "status": "failed", + "startedAt": now_ms, + "inputParams": {"alerts": {"_type": "list", "count": 1200}}, + "outputResults": {}, + }, + [], + ) + + db = await WorkflowStore.raw_db() + async with db.execute( + "SELECT raw_count, normalized_count, success_count, error_count, invalid_count " + "FROM workflow_metric_rollups WHERE workflow_id = ?", + ("stream_alert_denoise",), + ) as cursor: + row = await cursor.fetchone() + + assert row is not None + assert tuple(row) == (1203, 0, 1, 1, 1) + + +@pytest.mark.asyncio +async def test_workflow_store_does_not_count_malformed_syslog_as_accepted_alert() -> None: + await WorkflowStore.init() + now_ms = WorkflowStore._now_ms() + + await WorkflowStore.complete_execution( + { + "id": "denoise-malformed-syslog", + "workflowId": "stream_alert_denoise", + "status": "success", + "startedAt": now_ms, + "inputParams": {"syslog_message": {"message": "not-json"}}, + "outputResults": { + "stats": { + "metric_schema_version": 2, + "raw_count": 0, + "normalized_count": 0, + "after_filter_count": 0, + "after_dedup_count": 0, + } + }, + }, + [], + ) + + db = await WorkflowStore.raw_db() + async with db.execute( + "SELECT raw_count, invalid_count FROM workflow_metric_rollups WHERE workflow_id = ?", + ("stream_alert_denoise",), + ) as cursor: + row = await cursor.fetchone() + + assert row is not None + assert tuple(row) == (0, 0) + + +@pytest.mark.asyncio +async def test_workflow_store_does_not_accept_unverified_zero_for_file_input() -> None: + await WorkflowStore.init() + now_ms = WorkflowStore._now_ms() + + await WorkflowStore.complete_execution( + { + "id": "denoise-file-zero", + "workflowId": "stream_alert_denoise", + "status": "success", + "startedAt": now_ms, + "inputParams": {"alert_file": "/tmp/alerts.json"}, + "outputResults": { + "stats": { + "metric_schema_version": 2, + "raw_count": 0, + "normalized_count": 0, + "after_filter_count": 0, + "after_dedup_count": 0, + } + }, + }, + [], + ) + + db = await WorkflowStore.raw_db() + async with db.execute( + "SELECT raw_count, invalid_count FROM workflow_metric_rollups WHERE workflow_id = ?", + ("stream_alert_denoise",), + ) as cursor: + row = await cursor.fetchone() + + assert row is not None + assert tuple(row) == (0, 1) + + @pytest.mark.asyncio async def test_workflow_store_marks_legacy_batch_metrics_invalid() -> None: await WorkflowStore.init() diff --git a/webui/src/utils/socDashboardPageRuntime.test.tsx b/webui/src/utils/socDashboardPageRuntime.test.tsx index a8170d7ba..1ee12aa3d 100644 --- a/webui/src/utils/socDashboardPageRuntime.test.tsx +++ b/webui/src/utils/socDashboardPageRuntime.test.tsx @@ -89,6 +89,51 @@ describe('SOC dashboard contract page runtime', () => { expect(screen.getByText('Flocks AI 智能告警态势中心')).toBeInTheDocument(); }); + it('shows unavailable denoise metrics as dashes instead of false zeros', async () => { + pageGetMock.mockImplementation((path: string) => { + if (path === '/stats') { + return Promise.resolve({ + data: { + generatedAt: new Date().toISOString(), + denoise: { totalRaw: 0, totalUnique: 0, duplicateRate: 0, duplicates: 0 }, + pipeline: { raw: 0, unique: 0 }, + sources: [{ key: 'ndr', label: 'NDR', value: 0, rate: 0, active: false }], + sourceStatus: { + metricQuality: { + status: 'unavailable', + dataAvailable: false, + metricsAvailable: false, + unavailableReason: 'workflow_db_missing', + }, + }, + }, + }); + } + if (path === '/activity') { + return Promise.resolve({ + data: { + cursor: 'cursor', events: [], recentEvents: [], workflowEvents: [], batch: {}, + workflowStats: { callCount: null, latestStartedAt: null }, + }, + }); + } + if (path === '/ai-tasks') { + return Promise.resolve({ data: { connection: 'online', summary: {}, tasks: [] } }); + } + if (path === '/task-center') return Promise.resolve({ data: { scheduledTasks: [], workflows: [] } }); + return Promise.reject(new Error(`unexpected path: ${path}`)); + }); + + const { container } = render(); + + expect(await screen.findByText('降噪统计数据源不可用，相关数字已隐藏；系统正在重试')).toBeInTheDocument(); + const rawMetric = screen.getByText('原始告警量').closest('.command-metric') as HTMLElement; + expect(within(rawMetric).getByText('--')).toHaveAttribute('title', '数据不可用'); + const ndrSource = screen.getByText('NDR').closest('.command-source') as HTMLElement; + expect(within(ndrSource).getByText('--')).toBeInTheDocument(); + expect(container.querySelector('.command-source b')).toHaveAttribute('title', '数据不可用'); + }); + it('pauses task-center polling while the page is hidden', async () => { setDocumentHidden(true); From f235fcdae964904e4c27315501a19c3956f37b43 Mon Sep 17 00:00:00 2001 From: John Yin <10972267+john-yin2333@user.noreply.gitee.com> Date: Mon, 7 Sep 2026 23:49:26 +0800 Subject: [PATCH 53/63] fix(soc): harden dashboard metric quality contracts --- .../soc_ui/soc_dashboard/api/handlers.py | 314 +++++++++++++---- .../webuis/soc_ui/soc_dashboard/src/Page.tsx | 149 +++++--- .../soc_dashboard/src/severityValues.ts | 15 +- flocks/workflow/store.py | 60 ++-- tests/hub/test_soc_dashboard_schema.py | 319 ++++++++++++++++++ tests/workflow/test_workflow_store.py | 92 +++++ .../utils/socDashboardPageRuntime.test.tsx | 77 ++++- 7 files changed, 875 insertions(+), 151 deletions(-) diff --git a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/handlers.py b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/handlers.py index e6ebf89ac..69627de44 100644 --- a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/handlers.py +++ b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/api/handlers.py @@ -982,6 +982,15 @@ def _get_workflow_metric_rollups(workflow_name, start_time, end_time): # contain false zero ingress counts. Keep them out of the exact # path until a v3 contribution establishes verified coverage. return None + earliest_execution_ms = 0 + if _table_exists(conn, "workflow_executions"): + earliest_execution_row = conn.execute( + "SELECT MIN(started_at) FROM workflow_executions WHERE workflow_id = ?", + (workflow_name,), + ).fetchone() + earliest_execution_ms = _safe_int( + earliest_execution_row[0] if earliest_execution_row else 0 + ) query = ( "SELECT * FROM workflow_metric_rollups " "WHERE workflow_id = ? AND schema_version >= ?" @@ -1017,7 +1026,12 @@ def _get_workflow_metric_rollups(workflow_name, start_time, end_time): verified_started_at, 0, ) - complete_window = not start_ms or coverage_started_at <= start_ms + requested_start_ms = start_ms or earliest_execution_ms + # An unbounded query is complete only when execution history proves that + # verified rollups cover the workflow's first execution. Without that + # lower bound, treating recent v3 rows as all-time history would silently + # undercount older executions. + complete_window = requested_start_ms > 0 and coverage_started_at <= requested_start_ms source_complete = source_covered_count == raw_count metrics_complete = invalid_count == 0 and error_count == 0 quality = ( @@ -1089,15 +1103,17 @@ def _get_workflow_denoise_stats( now = time.time() cache_key = f"denoise:{workflow_name}:{start_time or 0}:{end_time or 0}" + # Source liveness is part of the quality contract. Do not let a fresh + # process-local cache make a disappeared database look healthy. + if not WORKFLOW_DB.is_file(): + return _unavailable_workflow_denoise_stats("workflow_db_missing") + with _cache_lock: cached = _workflow_stats_cache.get(cache_key) if not force and cached and now - float(cached.get("updatedAt") or 0) < _CACHE_TTL: _workflow_stats_cache.move_to_end(cache_key) return cached["value"] - if not WORKFLOW_DB.is_file(): - return _unavailable_workflow_denoise_stats("workflow_db_missing") - rollup_result = _get_workflow_metric_rollups(workflow_name, start_time, end_time) if rollup_result is not None and rollup_result.get("coverageComplete"): with _cache_lock: @@ -2147,7 +2163,19 @@ def _workflow_task_metrics(output_text, status): values[field], quality = _workflow_task_metric(stats, key) qualities.append(quality) if all(quality == "complete" for quality in qualities): - quality = "complete" + ordered = ( + values["raw"], + values["normalized"], + values["afterFilter"], + values["unique"], + ) + if ordered[0] >= ordered[1] >= ordered[2] >= ordered[3] >= 0: + quality = "complete" + else: + # Running executions may expose node-local intermediate output. + # Treat an impossible stage order as pending until completion; + # persisted finished output with the same shape is invalid. + quality = "pending" if status in WORKFLOW_RUNNING_STATUSES else "invalid" elif "invalid" in qualities: quality = "invalid" elif status in WORKFLOW_RUNNING_STATUSES: @@ -2158,29 +2186,57 @@ def _workflow_task_metrics(output_text, status): def _workflow_task_input_count(inputs): - for key in ("_raw_alerts_count", "_alerts_count", "_alert_list_count", "raw_count"): - value, quality = _workflow_task_metric(inputs, key) - if quality == "complete": - return value - for key in ("raw_alerts", "alerts", "alert_list"): - value = inputs.get(key) - if isinstance(value, list): - return len(value) - if isinstance(value, dict) and value.get("_type") in {"list", "tuple", "set"}: - count, quality = _workflow_task_metric(value, "count") - if quality == "complete": - return count - if isinstance(value, dict) and isinstance(value.get("data"), list): - return len(value["data"]) - for key in ("syslog_message", "syslog"): - value = inputs.get(key) - if not isinstance(value, dict) or not value.get("message"): - continue + # Match stream_alert_denoise's receive node exactly: a decodable syslog + # payload wins over every batch field. + value = inputs.get("syslog_message") or inputs.get("syslog") + if isinstance(value, dict) and value.get("message"): try: - json.loads(str(value["message"])) + parsed_syslog = json.loads(str(value["message"])) except (TypeError, ValueError, json.JSONDecodeError): + pass + else: + if isinstance(parsed_syslog, dict): + return 1 + + def sequence_count(value): + if isinstance(value, str): + try: + value = json.loads(value) + except (TypeError, ValueError, json.JSONDecodeError): + return None + if isinstance(value, list): + return len(value) + if isinstance(value, dict): + if value.get("_type") in {"list", "tuple", "set"}: + count, quality = _workflow_task_metric(value, "count") + return count if quality == "complete" else None + if "data" in value: + value = value.get("data") + return len(value) if isinstance(value, list) else (1 if value else 0) + return 1 if value else 0 + + # `alerts` shadows `alert_list` even when it is empty, matching + # inputs.get('alerts', inputs.get('alert_list', [])) in the workflow. + for key in ("alerts", "alert_list"): + marker_key = f"_{key}_count" + if key not in inputs and marker_key not in inputs: continue - return 1 + materialized_count = sequence_count(inputs.get(key)) if key in inputs else None + if materialized_count is not None: + return materialized_count + marker, quality = _workflow_task_metric(inputs, marker_key) + if quality == "complete": + return marker + return None + + # Compatibility fallbacks for older/compacted execution rows. They are + # considered only when no canonical workflow input is present. + for marker_key in ("_raw_alerts_count", "raw_count"): + marker, quality = _workflow_task_metric(inputs, marker_key) + if quality == "complete": + return marker + if "raw_alerts" in inputs: + return sequence_count(inputs.get("raw_alerts")) if inputs.get("alert") not in (None, "", {}): return 1 return None @@ -2193,14 +2249,24 @@ def _workflow_task_row(row, workflow_id, effective_status): inputs = _safe_json_object(input_text) metrics = _workflow_execution_metrics(output_text, input_text) counts, data_quality = _workflow_task_metrics(output_text, effective_status) - raw_count_source = "workflow_output" if counts["raw"] is not None else "pending" - if counts["raw"] is None: - input_count = _workflow_task_input_count(inputs) - if input_count is not None: - counts["raw"] = input_count - raw_count_source = "workflow_input" - if input_count == 0: - data_quality = "empty-input" + input_count = _workflow_task_input_count(inputs) + output_raw_count = counts["raw"] + raw_count_source = "workflow_output" if output_raw_count is not None else "pending" + if input_count is not None: + counts["raw"] = input_count + raw_count_source = "workflow_input" + if output_raw_count is not None and output_raw_count != input_count: + data_quality = ( + "pending" if effective_status in WORKFLOW_RUNNING_STATUSES else "invalid" + ) + elif input_count == 0 and data_quality in {"complete", "pending", "missing"}: + data_quality = "empty-input" + elif effective_status in WORKFLOW_RUNNING_STATUSES and output_raw_count == 0: + # A zero-initialized output is not proof that an active task received + # no alerts. Keep it unknown until input or completed output verifies it. + counts["raw"] = None + raw_count_source = "pending" + data_quality = "pending" preview = metrics["preview"] stage = "triage" if workflow_id in TRIAGE_WORKFLOW_IDS else "denoise" title = _workflow_latest_alert_name(workflow_id, output_text, input_text) @@ -2250,7 +2316,7 @@ def _workflow_task_row(row, workflow_id, effective_status): "counts": counts, "dataQuality": data_quality, "rawCountSource": raw_count_source, - "emptyBatch": counts["raw"] == 0, + "emptyBatch": data_quality == "empty-input", "emptyInput": data_quality == "empty-input", "progress": progress, "sessionId": session_id, @@ -2908,7 +2974,13 @@ async def get_stats(ctx, request): def _get_stats(params): - _ensure_sqlite_schema() + try: + _ensure_sqlite_schema() + except Exception: + # Source quality is resolved by the read-only query below. Keeping the + # endpoint alive lets the UI distinguish an unavailable SOC database + # from a healthy database whose selected window genuinely has zero rows. + pass time_window = _normalize_time_window( params.get("startTime"), params.get("endTime"), @@ -2942,6 +3014,16 @@ def _get_stats(params): cache_key, _stats_cache_ttl(range_start_time, range_end_time), ) + if cached is not None: + cached_status = cached.get("sourceStatus") or {} + cached_assets = cached_status.get("assets") or {} + cached_workflow = cached_status.get("metricQuality") or {} + source_liveness_changed = bool(cached_assets.get("exists")) != _active_source_exists() + workflow_liveness_changed = bool(cached_workflow.get("dataAvailable", True)) != bool( + WORKFLOW_DB.is_file() + ) + if source_liveness_changed or workflow_liveness_changed: + cached = None if cached is not None: return { **cached, @@ -2954,7 +3036,12 @@ def _get_stats(params): denoise_files, denoise_locations = [], [] triage_files, triage_locations = [], [] - asset_files = _find_asset_files(start_date, end_date, start_time, end_time) + asset_files, triage_quality = _find_asset_files_with_quality( + start_date, + end_date, + start_time, + end_time, + ) asset_denoise_files = [path for path in asset_files if _asset_file_role(path) == "denoise"] asset_triage_files = [path for path in asset_files if _asset_file_role(path) == "triage"] sample_mode = bool(asset_denoise_files or asset_triage_files) @@ -3033,6 +3120,65 @@ def _get_stats(params): available_dates = _available_asset_dates() date_range = _build_date_range(start_date, end_date, asset_files, available_dates) event_range = _build_event_range(date_range, denoise, triage) + triage_payload = _without_counters(triage) + pipeline_payload = dict(pipeline) + closed_loop_payload = dict(closed_loop) + attack_profile = _build_attack_profile(denoise, triage) + verdicts = [ + {"key": "attack_success", "label": "攻击成功", "value": triage["attackSuccess"], "color": "#ff4d6d"}, + {"key": "attack", "label": "攻击行为", "value": triage["attack"], "color": "#ffb020"}, + {"key": "attack_failed", "label": "攻击失败", "value": triage["attackFailed"], "color": "#2ee6a6"}, + {"key": "non_attack", "label": "非攻击", "value": triage["benign"], "color": "#58a6ff"}, + {"key": "unknown", "label": "未知", "value": triage["unknown"], "color": "#9b8cff"}, + ] + top_threat_types = _counter_items( + triage["threatTypeCounter"] or denoise["threatTypeCounter"], + 14, + ) + severity_levels = _counter_items( + _profile_counter(denoise, triage, "severityCounter"), + 8, + ) + risk_levels = _counter_items(triage["riskCounter"], 5) + triage_series_total = triage["seriesTotal"] + triage_series_attack = triage["seriesAttack"] + if not triage_quality["metricsAvailable"]: + for key in ( + "totalRecords", "batchTotal", "newTriaged", "cacheHit", "triageFailed", + "followersReused", "attackTotal", "attackSuccess", "attack", "attackFailed", + "benign", "unknown", "attackRate", "successRate", "cacheRate", "coverageRate", + "avgTriageMs", "headers", "files", "parseErrors", + ): + triage_payload[key] = None + for key in ( + "triageTotal", "attackTotal", "llmSaved", "workloadReuseRate", + "attackRate", "successRate", + ): + pipeline_payload[key] = None + closed_loop_payload = {key: None for key in closed_loop_payload} + verdicts = [{**item, "value": None} for item in verdicts] + attack_profile = [] + top_threat_types = [] + severity_levels = [] + risk_levels = [] + triage_series_total = [] + triage_series_attack = [] + + missing_sources = [ + item + for item in denoise_locations + triage_locations + if not item["exists"] or item["fileCount"] == 0 + ] + if not triage_quality["dataAvailable"]: + missing_sources.append( + { + "kind": "soc", + "path": _display_path(DEFAULT_SQLITE_DB), + "exists": DEFAULT_SQLITE_DB.is_file(), + "fileCount": 0, + "reason": triage_quality["unavailableReason"], + } + ) result = { "date": start_date, @@ -3052,6 +3198,7 @@ def _get_stats(params): }, "workflowStatsDb": _display_path(WORKFLOW_DB), "workflowStats": workflow_stats, + "triageQuality": triage_quality, "metricQuality": { "status": workflow_stats.get("dataQuality", "legacy"), "dataAvailable": workflow_stats.get("dataAvailable", True), @@ -3081,34 +3228,18 @@ def _get_stats(params): "triage": triage_locations, "denoiseFiles": [_file_brief(path) for path in denoise_files], "triageFiles": [_file_brief(path) for path in triage_files], - "missing": [] if sample_mode else [ - item - for item in denoise_locations + triage_locations - if not item["exists"] or item["fileCount"] == 0 - ], + "missing": missing_sources, }, "denoise": _without_counters(denoise), - "triage": _without_counters(triage), - "pipeline": pipeline, + "triage": triage_payload, + "pipeline": pipeline_payload, "sources": sources, - "closedLoop": closed_loop, - "attackProfile": _build_attack_profile(denoise, triage), - "verdicts": [ - {"key": "attack_success", "label": "攻击成功", "value": triage["attackSuccess"], "color": "#ff4d6d"}, - {"key": "attack", "label": "攻击行为", "value": triage["attack"], "color": "#ffb020"}, - {"key": "attack_failed", "label": "攻击失败", "value": triage["attackFailed"], "color": "#2ee6a6"}, - {"key": "non_attack", "label": "非攻击", "value": triage["benign"], "color": "#58a6ff"}, - {"key": "unknown", "label": "未知", "value": triage["unknown"], "color": "#9b8cff"}, - ], - "topThreatTypes": _counter_items( - triage["threatTypeCounter"] or denoise["threatTypeCounter"], - 14, - ), - "severityLevels": _counter_items( - _profile_counter(denoise, triage, "severityCounter"), - 8, - ), - "riskLevels": _counter_items(triage["riskCounter"], 5), + "closedLoop": closed_loop_payload, + "attackProfile": attack_profile, + "verdicts": verdicts, + "topThreatTypes": top_threat_types, + "severityLevels": severity_levels, + "riskLevels": risk_levels, "tokenUsage": _read_token_usage(), "timeline": { "labels": denoise.get("_timelineLabels") @@ -3117,8 +3248,8 @@ def _get_stats(params): or _timeline_window(start_date, end_date, len(denoise["seriesRaw"])), "denoiseRaw": denoise["seriesRaw"], "denoiseUnique": denoise["seriesUnique"], - "triageTotal": triage["seriesTotal"], - "triageAttack": triage["seriesAttack"], + "triageTotal": triage_series_total, + "triageAttack": triage_series_attack, }, } result["cacheHit"] = False @@ -3160,7 +3291,17 @@ def _date_span(start_date, end_date): def _find_asset_files(start_date, end_date, start_time=0, end_time=0): - return _find_sqlite_sources(start_date, end_date, start_time, end_time) + sources, _ = _find_sqlite_sources_with_quality( + start_date, + end_date, + start_time, + end_time, + ) + return sources + + +def _find_asset_files_with_quality(start_date, end_date, start_time=0, end_time=0): + return _find_sqlite_sources_with_quality(start_date, end_date, start_time, end_time) def _asset_file_date(path): @@ -3204,10 +3345,31 @@ def _active_source_exists(): def _find_sqlite_sources(start_date, end_date, start_time=0, end_time=0): + sources, _ = _find_sqlite_sources_with_quality( + start_date, + end_date, + start_time, + end_time, + ) + return sources + + +def _soc_source_quality(*, available, reason="", record_count=0): + return { + "status": "complete" if available else "unavailable", + "dataAvailable": bool(available), + "metricsAvailable": bool(available), + "unavailableReason": "" if available else str(reason or "soc_metrics_unavailable"), + "recordCount": max(_safe_int(record_count), 0) if available else None, + "dataSource": "soc.db.soc_dashboard_alert_facts" if available else "unavailable", + } + + +def _find_sqlite_sources_with_quality(start_date, end_date, start_time=0, end_time=0): settings = _sqlite_settings() db_path = settings["db_path"] if not db_path.is_file(): - return [] + return [], _soc_source_quality(available=False, reason="soc_db_missing") time_clause = "" query_params = [start_date, end_date] @@ -3223,9 +3385,24 @@ def _find_sqlite_sources(start_date, end_date, start_time=0, end_time=0): ) try: with sqlite3.connect(db_path) as conn: + conn.execute("PRAGMA query_only = ON") + required_tables = (DEFAULT_SQLITE_TABLE, FACTS_TABLE, META_TABLE) + if not all(_table_exists(conn, table_name) for table_name in required_tables): + return [], _soc_source_quality( + available=False, + reason="soc_dashboard_schema_unavailable", + ) + schema_row = conn.execute( + f"SELECT meta_value FROM {META_TABLE} WHERE meta_key='schema_version'" + ).fetchone() + if not schema_row or str(schema_row[0]) != SCHEMA_VERSION: + return [], _soc_source_quality( + available=False, + reason="soc_dashboard_schema_unavailable", + ) rows = conn.execute(query, query_params).fetchall() except Exception: - return [] + return [], _soc_source_quality(available=False, reason="soc_db_query_failed") sources = [] for asset_date, record_count in rows: @@ -3243,7 +3420,10 @@ def _find_sqlite_sources(start_date, end_date, start_time=0, end_time=0): end_time=end_time, ) ) - return sources + return sources, _soc_source_quality( + available=True, + record_count=sum(source.record_count for source in sources), + ) def _available_sqlite_dates(): diff --git a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx index a441ffdda..d14c473f4 100644 --- a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx +++ b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/Page.tsx @@ -24,56 +24,56 @@ const EMPTY_STATS = { eventRange: { start: '', end: '', label: '', source: '' }, generatedAt: '', latencyMs: 0, - sourceStatus: { workflowRoot: '', denoise: [], triage: [], denoiseFiles: [], triageFiles: [], missing: [], metricQuality: {} }, + sourceStatus: { workflowRoot: '', denoise: [], triage: [], denoiseFiles: [], triageFiles: [], missing: [], metricQuality: {}, triageQuality: {} }, denoise: { - totalRaw: 0, - totalNormalized: 0, - afterFilter: 0, - totalUnique: 0, - filterRemoved: 0, - dedupRemoved: 0, - duplicates: 0, - duplicateRate: 0, - dedupRate: 0, - uniqueRate: 0, - files: 0, - parseErrors: 0, + totalRaw: null, + totalNormalized: null, + afterFilter: null, + totalUnique: null, + filterRemoved: null, + dedupRemoved: null, + duplicates: null, + duplicateRate: null, + dedupRate: null, + uniqueRate: null, + files: null, + parseErrors: null, }, triage: { - totalRecords: 0, - newTriaged: 0, - cacheHit: 0, - triageFailed: 0, - followersReused: 0, - attackTotal: 0, - attackSuccess: 0, - attack: 0, - attackFailed: 0, - benign: 0, - unknown: 0, - attackRate: 0, - successRate: 0, - cacheRate: 0, - coverageRate: 0, - avgTriageMs: 0, - files: 0, - parseErrors: 0, + totalRecords: null, + newTriaged: null, + cacheHit: null, + triageFailed: null, + followersReused: null, + attackTotal: null, + attackSuccess: null, + attack: null, + attackFailed: null, + benign: null, + unknown: null, + attackRate: null, + successRate: null, + cacheRate: null, + coverageRate: null, + avgTriageMs: null, + files: null, + parseErrors: null, }, pipeline: { - raw: 0, - unique: 0, - triageTotal: 0, - attackTotal: 0, - reductionSaved: 0, - llmSaved: 0, - uniqueRate: 0, - workloadReuseRate: 0, - coverageRate: 0, - attackRate: 0, - successRate: 0, + raw: null, + unique: null, + triageTotal: null, + attackTotal: null, + reductionSaved: null, + llmSaved: null, + uniqueRate: null, + workloadReuseRate: null, + coverageRate: null, + attackRate: null, + successRate: null, }, sources: [], - closedLoop: { autoClosed: 0, resolved: 0, manualDecision: 0, pending: 0, resolutionRate: 0 }, + closedLoop: { autoClosed: null, resolved: null, manualDecision: null, pending: null, resolutionRate: null }, tokenUsage: { totalTokens: 0, todayTokens: 0, todayRequests: 0, dailySeries: [], dailyLabels: [], source: '' }, verdicts: [], attackProfile: [], @@ -782,9 +782,15 @@ function mergeStats(raw) { ...(EMPTY_STATS.sourceStatus.metricQuality || {}), ...((raw || {}).sourceStatus?.metricQuality || {}), }, + triageQuality: { + ...(EMPTY_STATS.sourceStatus.triageQuality || {}), + ...((raw || {}).sourceStatus?.triageQuality || {}), + }, }; const metricQuality = sourceStatus.metricQuality || {}; + const triageQuality = sourceStatus.triageQuality || {}; const denoiseMetricsUnavailable = metricQuality.metricsAvailable === false; + const triageMetricsUnavailable = triageQuality.metricsAvailable === false; const sourceMetricsUnavailable = denoiseMetricsUnavailable || metricQuality.sourceMetricsAvailable === false; const denoise = { ...EMPTY_STATS.denoise, ...((raw || {}).denoise || {}) }; @@ -796,26 +802,44 @@ function mergeStats(raw) { ]) denoise[key] = null; if (metricQuality.dataAvailable === false) denoise.files = null; } + const triage = { ...EMPTY_STATS.triage, ...((raw || {}).triage || {}) }; + if (triageMetricsUnavailable) { + for (const key of Object.keys(EMPTY_STATS.triage)) triage[key] = null; + } const pipeline = { ...EMPTY_STATS.pipeline, ...((raw || {}).pipeline || {}) }; if (denoiseMetricsUnavailable) { for (const key of ['raw', 'unique', 'reductionSaved', 'uniqueRate', 'coverageRate']) { pipeline[key] = null; } } + if (triageMetricsUnavailable) { + for (const key of [ + 'triageTotal', 'attackTotal', 'llmSaved', 'workloadReuseRate', + 'attackRate', 'successRate', + ]) pipeline[key] = null; + } const timeline = { ...EMPTY_STATS.timeline, ...((raw || {}).timeline || {}) }; if (denoiseMetricsUnavailable) { timeline.denoiseRaw = []; timeline.denoiseUnique = []; } + if (triageMetricsUnavailable) { + timeline.triageTotal = []; + timeline.triageAttack = []; + } const sources = Array.isArray((raw || {}).sources) ? raw.sources : []; + const closedLoop = { ...EMPTY_STATS.closedLoop, ...((raw || {}).closedLoop || {}) }; + if (triageMetricsUnavailable) { + for (const key of Object.keys(EMPTY_STATS.closedLoop)) closedLoop[key] = null; + } return { ...EMPTY_STATS, ...(raw || {}), sourceStatus, denoise, - triage: { ...EMPTY_STATS.triage, ...((raw || {}).triage || {}) }, + triage, pipeline, - closedLoop: { ...EMPTY_STATS.closedLoop, ...((raw || {}).closedLoop || {}) }, + closedLoop, tokenUsage: { ...EMPTY_STATS.tokenUsage, ...((raw || {}).tokenUsage || {}) }, dateRange: { ...EMPTY_STATS.dateRange, ...((raw || {}).dateRange || {}) }, eventRange: { ...EMPTY_STATS.eventRange, ...((raw || {}).eventRange || {}) }, @@ -823,6 +847,13 @@ function mergeStats(raw) { sources: sourceMetricsUnavailable ? sources.map((item) => ({ ...item, value: null, rate: null, active: false })) : sources, + verdicts: triageMetricsUnavailable + ? ((raw || {}).verdicts || []).map((item) => ({ ...item, value: null })) + : ((raw || {}).verdicts || []), + attackProfile: triageMetricsUnavailable ? [] : ((raw || {}).attackProfile || []), + topThreatTypes: triageMetricsUnavailable ? [] : ((raw || {}).topThreatTypes || []), + severityLevels: triageMetricsUnavailable ? [] : ((raw || {}).severityLevels || []), + riskLevels: triageMetricsUnavailable ? [] : ((raw || {}).riskLevels || []), }; } @@ -1860,6 +1891,9 @@ function laneLinkStatus(kind, event, peerLane) { } function triageContextText(stats) { + if (stats?.sourceStatus?.triageQuality?.metricsAvailable === false) { + return '窗口研判数据不可用'; + } const triage = stats?.triage || EMPTY_STATS.triage; const total = Math.max(Number(triage.totalRecords || 0), 0); const newTriaged = Math.max(Number(triage.newTriaged || 0), 0); @@ -1939,8 +1973,10 @@ function CommandGraph({ stats, activity }) { const activeSeverityTone = severityToneFor(activity.triage.current); const recentSeverityTone = severityToneFor(activity.triage.last); const activeSources = [...(stats.sources || [])].sort((a, b) => Number(b.value || 0) - Number(a.value || 0)).slice(0, 2); - while (activeSources.length < 2) activeSources.push({ key: `source-${activeSources.length}`, label: activeSources.length ? '备用数据源' : '告警数据源', value: 0 }); - const severities = severityRows(stats); + while (activeSources.length < 2) activeSources.push({ key: `source-${activeSources.length}`, label: activeSources.length ? '备用数据源' : '告警数据源', value: null }); + const triageMetricsAvailable = Boolean(stats.generatedAt) + && stats.sourceStatus?.triageQuality?.metricsAvailable !== false; + const severities = severityRows(stats, triageMetricsAvailable); return h('section', { className: cx('command-graph', denoiseActive && 'denoise-running', triageActive && 'triage-running', `load-${activity.mode}`) }, [ h(CommandConnections, { key: 'links' }), h('div', { className: 'source-stack', key: 'sources' }, activeSources.map((source) => h('div', { className: 'command-source', key: source.key }, [ @@ -2554,7 +2590,13 @@ function CommandAiTaskPanel({ aiTasks }) { ? waiting ? '待研判' : '智能研判' : waiting ? '待降噪' : '智能降噪'; const stateLabel = processing ? '处理中' : '等待处理'; - const qualityDetail = task.dataQuality === 'invalid' + const unverifiedZero = task.stage === 'denoise' + && Number(task.counts?.raw) === 0 + && !task.emptyInput + && task.rawCountSource !== 'workflow_input'; + const qualityDetail = unverifiedZero + ? ' · 原始条数待校验' + : task.dataQuality === 'invalid' ? ' · 指标格式异常' : task.dataQuality === 'empty-input' ? ' · 空输入任务' @@ -2563,7 +2605,7 @@ function CommandAiTaskPanel({ aiTasks }) { : task.dataQuality === 'pending' && (task.counts?.raw === null || task.counts?.raw === undefined) ? ' · 原始条数待生成' : ''; - const rawDetail = task.stage === 'denoise' && !task.emptyInput && task.counts?.raw !== null && task.counts?.raw !== undefined + const rawDetail = task.stage === 'denoise' && !task.emptyInput && !unverifiedZero && task.counts?.raw !== null && task.counts?.raw !== undefined ? ` · 原始 ${task.counts.raw} 条` : ''; const detail = task.stage === 'triage' @@ -3105,6 +3147,7 @@ export default function Page() { || displayActivity.batch?.triageUpdatedCount ); const metricQuality = stats.sourceStatus?.metricQuality || {}; + const triageQuality = stats.sourceStatus?.triageQuality || {}; const metricIssues = []; if (Number(metricQuality.invalidExecutionCount || 0) > 0) { metricIssues.push(`${metricQuality.invalidExecutionCount} 次指标格式异常`); @@ -3136,6 +3179,13 @@ export default function Page() { : metricQuality.status === 'legacy-partial' ? `精确降噪指标尚未覆盖当前时间范围，相关数字已隐藏；完整采集始于 ${taskCenterTimeLabel(metricQuality.coverageStartedAt)}` : '精确降噪指标尚不可用，相关数字已隐藏；新指标链路产生数据后将自动显示'; + const triageQualityWarning = !stats.generatedAt || triageQuality.metricsAvailable !== false + ? '' + : triageQuality.unavailableReason === 'soc_db_missing' + ? 'SOC 事件数据库不可用，研判、事件与闭环数字已隐藏；系统正在重试' + : triageQuality.unavailableReason === 'soc_dashboard_schema_unavailable' + ? 'SOC 统计结构尚未就绪，研判、事件与闭环数字已隐藏；系统正在重试' + : 'SOC 研判统计查询失败，研判、事件与闭环数字已隐藏；系统正在重试'; return h('div', { className: cx('adtd-root command-root', displayActivityBusy && 'command-is-processing', eventRailCollapsed && 'event-rail-is-collapsed'), @@ -3158,6 +3208,7 @@ export default function Page() { }), error ? h('div', { className: 'error-banner', key: 'error' }, `统计接口异常：${error}`) : null, metricQualityWarning ? h('div', { className: 'quality-banner', key: 'quality' }, metricQualityWarning) : null, + triageQualityWarning ? h('div', { className: 'quality-banner', key: 'triage-quality' }, triageQualityWarning) : null, h('main', { className: cx('command-shell', eventRailCollapsed && 'event-rail-collapsed'), key: 'main', diff --git a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/severityValues.ts b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/severityValues.ts index d2fcd6f2a..675e639f4 100644 --- a/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/severityValues.ts +++ b/.flocks/flockshub/plugins/webuis/soc_ui/soc_dashboard/src/severityValues.ts @@ -1,6 +1,6 @@ type SeverityItem = { label?: string; - value?: number; + value?: number | null; }; export function severityKey(value: unknown) { @@ -12,16 +12,19 @@ export function severityKey(value: unknown) { return ''; } -export function severityRows(stats: { severityLevels?: SeverityItem[] }) { +export function severityRows( + stats: { severityLevels?: SeverityItem[] }, + metricsAvailable = true, +) { const counts = { critical: 0, high: 0, medium: 0, low: 0 }; for (const item of stats.severityLevels || []) { const key = severityKey(item.label); if (key) counts[key] += Number(item.value || 0); } return [ - { key: 'critical', label: '严重', value: counts.critical, tone: 'critical' }, - { key: 'high', label: '高危', value: counts.high, tone: 'high' }, - { key: 'medium', label: '中危', value: counts.medium, tone: 'medium' }, - { key: 'low', label: '低危', value: counts.low, tone: 'low' }, + { key: 'critical', label: '严重', value: metricsAvailable ? counts.critical : null, tone: 'critical' }, + { key: 'high', label: '高危', value: metricsAvailable ? counts.high : null, tone: 'high' }, + { key: 'medium', label: '中危', value: metricsAvailable ? counts.medium : null, tone: 'medium' }, + { key: 'low', label: '低危', value: metricsAvailable ? counts.low : null, tone: 'low' }, ]; } diff --git a/flocks/workflow/store.py b/flocks/workflow/store.py index ef87b1969..746655258 100644 --- a/flocks/workflow/store.py +++ b/flocks/workflow/store.py @@ -381,21 +381,21 @@ def _pipeline_metric_contribution(cls, exec_data: Dict[str, Any]) -> Optional[Di invalid_count = 0 def sequence_count(value: Any) -> Optional[int]: - if isinstance(value, (list, tuple)): + if isinstance(value, str): + try: + value = json.loads(value) + except (TypeError, ValueError, json.JSONDecodeError): + return None + if isinstance(value, list): return len(value) if isinstance(value, dict): if value.get("_type") in {"list", "tuple", "set"}: count = cls._as_int(value.get("count")) return count if count is not None and count >= 0 else None - data = value.get("data") - if isinstance(data, (list, tuple)): - return len(data) - if isinstance(value, str): - try: - return sequence_count(json.loads(value)) - except (TypeError, ValueError, json.JSONDecodeError): - return None - return None + if "data" in value: + value = value.get("data") + return len(value) if isinstance(value, list) else (1 if value else 0) + return 1 if value else 0 def input_alert_count() -> Optional[int]: inputs = exec_data.get("inputParams") @@ -408,31 +408,41 @@ def input_alert_count() -> Optional[int]: syslog_message = inputs.get("syslog_message") or inputs.get("syslog") if isinstance(syslog_message, dict) and syslog_message.get("message"): try: - json.loads(str(syslog_message["message"])) + parsed_syslog = json.loads(str(syslog_message["message"])) except (TypeError, ValueError, json.JSONDecodeError): pass else: - return 1 - - known_empty = False - for key in ("raw_alerts", "alerts", "alert_list"): - marker = cls._as_int(inputs.get(f"_{key}_count")) + if isinstance(parsed_syslog, dict): + return 1 + + # `alerts` shadows `alert_list` even when empty, matching the + # receive node's inputs.get('alerts', inputs.get('alert_list', [])). + for key in ("alerts", "alert_list"): + marker_key = f"_{key}_count" + if key not in inputs and marker_key not in inputs: + continue + materialized_count = sequence_count(inputs.get(key)) if key in inputs else None + if materialized_count is not None: + return materialized_count + marker = cls._as_int(inputs.get(marker_key)) if marker is not None and marker >= 0: - if marker > 0: - return marker - known_empty = True - count = sequence_count(inputs.get(key)) - if count is not None: - if count > 0: - return count - known_empty = True + return marker + return None + + # Compatibility fallback for pre-v3/compacted rows which stored + # the accepted batch under raw_alerts rather than the API field. + marker = cls._as_int(inputs.get("_raw_alerts_count")) + if marker is not None and marker >= 0: + return marker + if "raw_alerts" in inputs: + return sequence_count(inputs.get("raw_alerts")) # File inputs are intentionally unknown here: reading a user file # while committing execution state would introduce I/O and TOCTOU # races. Successful output metrics remain authoritative for them. if inputs.get("alert_file"): return None - return 0 if known_empty else None + return None input_count = input_alert_count() input_params = exec_data.get("inputParams") diff --git a/tests/hub/test_soc_dashboard_schema.py b/tests/hub/test_soc_dashboard_schema.py index 2db851ade..53883a3c5 100644 --- a/tests/hub/test_soc_dashboard_schema.py +++ b/tests/hub/test_soc_dashboard_schema.py @@ -808,6 +808,95 @@ def test_soc_dashboard_ai_task_input_count_supports_compacted_batches_and_empty_ assert handlers._workflow_task_input_count( {"syslog_message": {"message": "not-json"}} ) is None + assert handlers._workflow_task_input_count( + {"syslog_message": {"message": "[]"}} + ) is None + assert handlers._workflow_task_input_count( + { + "syslog_message": {"message": json.dumps({"id": "syslog-alert"})}, + "alerts": [{"id": index} for index in range(5)], + } + ) == 1 + assert handlers._workflow_task_input_count( + {"_alerts_count": 0, "alert_list": [{"id": "shadowed"}]} + ) == 0 + assert handlers._workflow_task_input_count( + {"_alerts_count": 0, "alerts": [{"id": "materialized"}]} + ) == 1 + + +def test_soc_dashboard_ai_task_reconciles_zero_initialized_output_with_input(): + handlers = _load_dashboard_handlers() + row = { + "id": "running-zero-output", + "output_results": json.dumps( + { + "stats": { + "raw_count": 0, + "normalized_count": 0, + "after_filter_count": 0, + "after_dedup_count": 0, + } + } + ), + "input_params": json.dumps({"alerts": [{"id": "real-alert"}]}), + "started_at": 1, + "updated_at": 1, + "finished_at": 0, + "current_phase": "receive", + "current_step_index": 1, + "step_count": 5, + "payload": "{}", + "error_message": "", + } + + task = handlers._workflow_task_row( + row, + "stream_alert_denoise", + "running", + ) + + assert task["counts"]["raw"] == 1 + assert task["rawCountSource"] == "workflow_input" + assert task["dataQuality"] == "pending" + assert task["emptyBatch"] is False + assert task["emptyInput"] is False + + empty_row = { + **row, + "id": "running-empty-input", + "input_params": json.dumps({"alerts": []}), + } + empty_task = handlers._workflow_task_row( + empty_row, + "stream_alert_denoise", + "running", + ) + assert empty_task["counts"]["raw"] == 0 + assert empty_task["rawCountSource"] == "workflow_input" + assert empty_task["dataQuality"] == "empty-input" + assert empty_task["emptyBatch"] is True + assert empty_task["emptyInput"] is True + + +def test_soc_dashboard_ai_task_rejects_impossible_stage_counts(): + handlers = _load_dashboard_handlers() + output = json.dumps( + { + "stats": { + "raw_count": 1, + "normalized_count": 2, + "after_filter_count": 3, + "after_dedup_count": 4, + } + } + ) + + _, running_quality = handlers._workflow_task_metrics(output, "running") + _, completed_quality = handlers._workflow_task_metrics(output, "completed") + + assert running_quality == "pending" + assert completed_quality == "invalid" def test_soc_dashboard_marks_missing_workflow_metrics_as_unavailable(tmp_path: Path): @@ -829,6 +918,236 @@ def test_soc_dashboard_marks_missing_workflow_metrics_as_unavailable(tmp_path: P assert stats["callCount"] is None +def test_soc_dashboard_does_not_serve_fresh_cache_after_workflow_db_disappears( + tmp_path: Path, +): + workflow_db = tmp_path / "workflow.db" + start_time = 1_800_000 + start_ms = start_time * 1000 + with sqlite3.connect(workflow_db) as conn: + conn.execute( + "CREATE TABLE workflow_metric_meta " + "(workflow_id TEXT PRIMARY KEY, coverage_started_at INTEGER, updated_at INTEGER)" + ) + conn.execute( + """ + CREATE TABLE workflow_metric_rollups ( + workflow_id TEXT, bucket_start INTEGER, raw_count INTEGER, + normalized_count INTEGER, after_filter_count INTEGER, unique_count INTEGER, + filter_removed_count INTEGER, duplicate_count INTEGER, source_counts TEXT, + source_covered_count INTEGER, success_count INTEGER, error_count INTEGER, + invalid_count INTEGER, schema_version INTEGER, updated_at INTEGER, + PRIMARY KEY (workflow_id, bucket_start) + ) + """ + ) + conn.execute( + "INSERT INTO workflow_metric_meta VALUES (?, ?, ?)", + ("stream_alert_denoise", start_ms, start_ms), + ) + conn.execute( + "INSERT INTO workflow_metric_rollups VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + ( + "stream_alert_denoise", start_ms, 1, 1, 1, 1, 0, 0, + '{"tdp": 1}', 1, 1, 0, 0, 3, start_ms, + ), + ) + conn.commit() + + handlers = _load_dashboard_handlers() + handlers.WORKFLOW_DB = workflow_db + handlers._workflow_stats_cache.clear() + assert handlers._get_workflow_denoise_stats( + "stream_alert_denoise", start_time, start_time + 60, force=True + )["dataQuality"] == "complete" + + workflow_db.unlink() + stats = handlers._get_workflow_denoise_stats( + "stream_alert_denoise", start_time, start_time + 60 + ) + + assert stats["dataAvailable"] is False + assert stats["dataQuality"] == "unavailable" + assert stats["rawCount"] is None + + +def test_soc_dashboard_unbounded_rollup_requires_proven_history_start(tmp_path: Path): + workflow_db = tmp_path / "workflow.db" + bucket_ms = 2_000_000_000_000 + with sqlite3.connect(workflow_db) as conn: + conn.execute( + "CREATE TABLE workflow_metric_meta " + "(workflow_id TEXT PRIMARY KEY, coverage_started_at INTEGER, updated_at INTEGER)" + ) + conn.execute( + """ + CREATE TABLE workflow_metric_rollups ( + workflow_id TEXT, bucket_start INTEGER, raw_count INTEGER, + normalized_count INTEGER, after_filter_count INTEGER, unique_count INTEGER, + filter_removed_count INTEGER, duplicate_count INTEGER, source_counts TEXT, + source_covered_count INTEGER, success_count INTEGER, error_count INTEGER, + invalid_count INTEGER, schema_version INTEGER, updated_at INTEGER, + PRIMARY KEY (workflow_id, bucket_start) + ) + """ + ) + conn.execute( + "INSERT INTO workflow_metric_meta VALUES (?, ?, ?)", + ("stream_alert_denoise", bucket_ms, bucket_ms), + ) + conn.execute( + "INSERT INTO workflow_metric_rollups VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + ( + "stream_alert_denoise", bucket_ms, 7, 7, 6, 5, 1, 1, + '{"tdp": 7}', 7, 1, 0, 0, 3, bucket_ms, + ), + ) + conn.commit() + + handlers = _load_dashboard_handlers() + handlers.WORKFLOW_DB = workflow_db + + stats = handlers._get_workflow_metric_rollups("stream_alert_denoise", 0, 0) + + assert stats["coverageComplete"] is False + assert stats["dataQuality"] == "partial" + + +def test_soc_dashboard_hides_triage_metrics_when_soc_database_is_missing(tmp_path: Path): + workflow_db = tmp_path / "workflow.db" + start_time = 1_788_739_200 + start_ms = start_time * 1000 + with sqlite3.connect(workflow_db) as conn: + conn.execute( + "CREATE TABLE workflow_metric_meta " + "(workflow_id TEXT PRIMARY KEY, coverage_started_at INTEGER, updated_at INTEGER)" + ) + conn.execute( + """ + CREATE TABLE workflow_metric_rollups ( + workflow_id TEXT, bucket_start INTEGER, raw_count INTEGER, + normalized_count INTEGER, after_filter_count INTEGER, unique_count INTEGER, + filter_removed_count INTEGER, duplicate_count INTEGER, source_counts TEXT, + source_covered_count INTEGER, success_count INTEGER, error_count INTEGER, + invalid_count INTEGER, schema_version INTEGER, updated_at INTEGER, + PRIMARY KEY (workflow_id, bucket_start) + ) + """ + ) + conn.execute( + "INSERT INTO workflow_metric_meta VALUES (?, ?, ?)", + ("stream_alert_denoise", start_ms - 60_000, start_ms), + ) + conn.execute( + "INSERT INTO workflow_metric_rollups VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + ( + "stream_alert_denoise", start_ms, 10, 10, 8, 4, 2, 4, + '{"tdp": 10}', 10, 1, 0, 0, 3, start_ms, + ), + ) + conn.commit() + + handlers = _load_dashboard_handlers() + handlers.WORKFLOW_DB = workflow_db + handlers.DEFAULT_SQLITE_DB = tmp_path / "missing-soc.db" + handlers.USAGE_DB = tmp_path / "missing-usage.db" + handlers._schema_ready.clear() + handlers._stats_response_cache.clear() + handlers._workflow_stats_cache.clear() + + stats = handlers._get_stats( + { + "startTime": str(start_time), + "endTime": str(start_time + 3600), + "force": "true", + } + ) + + assert stats["sourceStatus"]["metricQuality"]["status"] == "complete" + assert stats["denoise"]["totalRaw"] == 10 + assert stats["sourceStatus"]["triageQuality"] == { + "status": "unavailable", + "dataAvailable": False, + "metricsAvailable": False, + "unavailableReason": "soc_db_missing", + "recordCount": None, + "dataSource": "unavailable", + } + assert stats["triage"]["totalRecords"] is None + assert stats["triage"]["attackTotal"] is None + assert stats["closedLoop"]["resolutionRate"] is None + assert stats["pipeline"]["attackRate"] is None + assert stats["timeline"]["triageTotal"] == [] + assert stats["sourceStatus"]["missing"][0]["reason"] == "soc_db_missing" + + +def test_soc_dashboard_keeps_real_zero_for_healthy_empty_soc_database(tmp_path: Path): + soc_db = tmp_path / "soc.db" + with sqlite3.connect(soc_db) as conn: + conn.execute( + """ + CREATE TABLE alert_records ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + record_json TEXT NOT NULL, + asset_date TEXT NOT NULL, + event_time INTEGER NOT NULL + ) + """ + ) + conn.commit() + handlers = _load_dashboard_handlers() + handlers.DEFAULT_SQLITE_DB = soc_db + handlers.WORKFLOW_DB = tmp_path / "missing-workflow.db" + handlers.USAGE_DB = tmp_path / "missing-usage.db" + handlers._schema_ready.clear() + handlers._stats_response_cache.clear() + handlers._workflow_stats_cache.clear() + + stats = handlers._get_stats( + {"startDate": "2026-09-07", "endDate": "2026-09-07", "force": "true"} + ) + + assert stats["sourceStatus"]["triageQuality"]["status"] == "complete" + assert stats["sourceStatus"]["triageQuality"]["metricsAvailable"] is True + assert stats["triage"]["totalRecords"] == 0 + assert stats["triage"]["attackTotal"] == 0 + assert stats["closedLoop"]["resolutionRate"] == 0 + + soc_db.unlink() + refreshed = handlers._get_stats( + {"startDate": "2026-09-07", "endDate": "2026-09-07"} + ) + assert refreshed["cacheHit"] is False + assert refreshed["sourceStatus"]["triageQuality"]["status"] == "unavailable" + assert refreshed["triage"]["totalRecords"] is None + + +def test_soc_dashboard_rejects_incomplete_soc_fact_schema(tmp_path: Path): + soc_db = tmp_path / "soc.db" + with sqlite3.connect(soc_db) as conn: + conn.execute( + "CREATE TABLE alert_records " + "(asset_date TEXT, event_time INTEGER, record_json TEXT)" + ) + conn.execute( + "CREATE TABLE soc_dashboard_alert_facts " + "(asset_date TEXT, event_time INTEGER)" + ) + conn.commit() + + handlers = _load_dashboard_handlers() + handlers.DEFAULT_SQLITE_DB = soc_db + + sources, quality = handlers._find_sqlite_sources_with_quality( + "2026-09-01", + "2026-09-07", + ) + + assert sources == [] + assert quality["metricsAvailable"] is False + assert quality["unavailableReason"] == "soc_dashboard_schema_unavailable" + + def test_soc_dashboard_reads_persisted_denoise_metric_rollups(tmp_path: Path): workflow_db = tmp_path / "workflow.db" start_time = 1_800_000 diff --git a/tests/workflow/test_workflow_store.py b/tests/workflow/test_workflow_store.py index f8a21978a..2fc1965b4 100644 --- a/tests/workflow/test_workflow_store.py +++ b/tests/workflow/test_workflow_store.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import json import os from pathlib import Path from unittest.mock import AsyncMock @@ -298,6 +299,97 @@ async def test_workflow_store_does_not_count_malformed_syslog_as_accepted_alert( assert row is not None assert tuple(row) == (0, 0) + await WorkflowStore.complete_execution( + { + "id": "denoise-non-object-syslog", + "workflowId": "stream_alert_denoise", + "status": "success", + "startedAt": now_ms, + "inputParams": {"syslog_message": {"message": "[]"}}, + "outputResults": { + "stats": { + "metric_schema_version": 2, + "raw_count": 0, + "normalized_count": 0, + "after_filter_count": 0, + "after_dedup_count": 0, + } + }, + }, + [], + ) + + async with db.execute( + "SELECT raw_count, invalid_count FROM workflow_metric_rollups WHERE workflow_id = ?", + ("stream_alert_denoise",), + ) as cursor: + row = await cursor.fetchone() + + assert row is not None + assert tuple(row) == (0, 0) + + +@pytest.mark.asyncio +async def test_workflow_store_input_reconciliation_matches_receive_node_priority() -> None: + await WorkflowStore.init() + now_ms = WorkflowStore._now_ms() + cases = [ + ( + "syslog-wins", + { + "syslog_message": {"message": json.dumps({"id": "single"})}, + "alerts": [{"id": index} for index in range(5)], + }, + 1, + ), + ( + "alerts-shadow-alert-list", + {"alerts": [], "alert_list": [{"id": "shadowed"}]}, + 0, + ), + ( + "canonical-alerts-win", + {"raw_alerts": [{"id": index} for index in range(5)], "alerts": [{"id": "api"}]}, + 1, + ), + ( + "materialized-alerts-win-over-stale-marker", + {"_alerts_count": 0, "alerts": [{"id": "materialized"}]}, + 1, + ), + ] + for execution_id, input_params, raw_count in cases: + await WorkflowStore.complete_execution( + { + "id": execution_id, + "workflowId": "stream_alert_denoise", + "status": "success", + "startedAt": now_ms, + "inputParams": input_params, + "outputResults": { + "stats": { + "metric_schema_version": 2, + "raw_count": raw_count, + "normalized_count": raw_count, + "after_filter_count": raw_count, + "after_dedup_count": raw_count, + } + }, + }, + [], + ) + + db = await WorkflowStore.raw_db() + async with db.execute( + "SELECT raw_count, success_count, invalid_count FROM workflow_metric_rollups " + "WHERE workflow_id = ?", + ("stream_alert_denoise",), + ) as cursor: + row = await cursor.fetchone() + + assert row is not None + assert tuple(row) == (3, 4, 0) + @pytest.mark.asyncio async def test_workflow_store_does_not_accept_unverified_zero_for_file_input() -> None: diff --git a/webui/src/utils/socDashboardPageRuntime.test.tsx b/webui/src/utils/socDashboardPageRuntime.test.tsx index 1ee12aa3d..091a35d14 100644 --- a/webui/src/utils/socDashboardPageRuntime.test.tsx +++ b/webui/src/utils/socDashboardPageRuntime.test.tsx @@ -134,6 +134,71 @@ describe('SOC dashboard contract page runtime', () => { expect(container.querySelector('.command-source b')).toHaveAttribute('title', '数据不可用'); }); + it('shows unavailable SOC triage metrics as dashes instead of false zeros', async () => { + pageGetMock.mockImplementation((path: string) => { + if (path === '/stats') { + return Promise.resolve({ + data: { + generatedAt: new Date().toISOString(), + denoise: { totalRaw: 10, totalUnique: 4, duplicateRate: 0.6, duplicates: 6 }, + triage: { + totalRecords: 0, + attackTotal: 0, + attackSuccess: 0, + benign: 0, + unknown: 0, + }, + pipeline: { attackRate: 0, successRate: 0 }, + closedLoop: { autoClosed: 0, manualDecision: 0, pending: 0, resolutionRate: 0 }, + severityLevels: [], + sourceStatus: { + metricQuality: { status: 'complete', metricsAvailable: true }, + triageQuality: { + status: 'unavailable', + dataAvailable: false, + metricsAvailable: false, + unavailableReason: 'soc_db_missing', + }, + }, + }, + }); + } + if (path === '/activity') { + return Promise.resolve({ + data: { + cursor: 'cursor', events: [], recentEvents: [], workflowEvents: [], batch: {}, + workflowStats: { callCount: 1, latestStartedAt: Date.now() }, + }, + }); + } + if (path === '/ai-tasks') { + return Promise.resolve({ data: { connection: 'online', summary: {}, tasks: [] } }); + } + if (path === '/task-center') return Promise.resolve({ data: { scheduledTasks: [], workflows: [] } }); + return Promise.reject(new Error(`unexpected path: ${path}`)); + }); + + render(); + + expect(await screen.findByText('SOC 事件数据库不可用，研判、事件与闭环数字已隐藏；系统正在重试')).toBeInTheDocument(); + const eventMetric = screen.getByText('安全事件量').closest('.command-metric') as HTMLElement; + expect(within(eventMetric).getByText('--')).toHaveAttribute('title', '数据不可用'); + const criticalSeverity = screen.getByText('严重').closest('.severity-node') as HTMLElement; + expect(within(criticalSeverity).getByText('--')).toHaveAttribute('title', '数据不可用'); + expect(screen.queryByText('SOC 事件数据库不可用，研判、事件与闭环数字已隐藏；系统正在重试')).toBeInTheDocument(); + }); + + it('renders unknown metrics as dashes on the loading frame', () => { + pageGetMock.mockImplementation(() => new Promise(() => {})); + + render(); + + const rawMetric = screen.getByText('原始告警量').closest('.command-metric') as HTMLElement; + const eventMetric = screen.getByText('安全事件量').closest('.command-metric') as HTMLElement; + expect(within(rawMetric).getByText('--')).toHaveAttribute('title', '数据不可用'); + expect(within(eventMetric).getByText('--')).toHaveAttribute('title', '数据不可用'); + }); + it('pauses task-center polling while the page is hidden', async () => { setDocumentHidden(true); @@ -406,9 +471,11 @@ describe('SOC dashboard contract page runtime', () => { stage: 'denoise', status: 'queued', startedAt: now - 1000, - title: '降噪批次 · 原始条数待生成', - counts: { raw: null }, - dataQuality: 'pending', + title: '降噪批次', + counts: { raw: 0 }, + dataQuality: 'complete', + rawCountSource: 'workflow_output', + emptyInput: false, progress: { mode: 'waiting', percent: null, label: '等待调度' }, }, ], @@ -423,7 +490,9 @@ describe('SOC dashboard contract page runtime', () => { expect(await screen.findByText('正在处理 1 个，等待 1 个')).toBeInTheDocument(); expect(screen.getByText('SSRF盲打探测攻击结果未知')).toBeInTheDocument(); - expect(screen.getByText('降噪批次 · 原始条数待生成')).toBeInTheDocument(); + expect(screen.getByText('降噪批次')).toBeInTheDocument(); + expect(screen.getByText(/原始条数待校验/)).toBeInTheDocument(); + expect(screen.queryByText(/原始 0 条/)).not.toBeInTheDocument(); expect(screen.getByText('第 2/3 步')).toBeInTheDocument(); expect(screen.queryByText('不应进入任务栏')).not.toBeInTheDocument(); expect(screen.queryByText('降噪处理完成')).not.toBeInTheDocument(); From 2798c55a3c624ac26426a2a002e132b04f5ae210 Mon Sep 17 00:00:00 2001 From: John Yin <10972267+john-yin2333@user.noreply.gitee.com> Date: Wed, 9 Sep 2026 10:23:03 +0800 Subject: [PATCH 54/63] fix workspace uploads with symlink root --- flocks/server/routes/workspace.py | 4 +- tests/workspace/test_workspace_routes.py | 60 ++++++++++++++++++++++++ 2 files changed, 63 insertions(+), 1 deletion(-) diff --git a/flocks/server/routes/workspace.py b/flocks/server/routes/workspace.py index c0c8cb871..a86bc2145 100644 --- a/flocks/server/routes/workspace.py +++ b/flocks/server/routes/workspace.py @@ -352,8 +352,10 @@ async def upload_files( files: List[UploadFile] = File(...), ): mgr = _get_manager() + workspace_root = _workspace_root(mgr) try: dest_dir = mgr.resolve_workspace_path(dest) if dest else mgr.get_workspace_dir() + relative_dest_dir = dest_dir.resolve().relative_to(workspace_root) except ValueError as e: raise HTTPException(status_code=400, detail=str(e)) dest_dir.mkdir(parents=True, exist_ok=True) @@ -411,7 +413,7 @@ async def upload_files( }) results.append({ "name": target.name, - "path": str(target.relative_to(mgr.get_workspace_dir())), + "path": str(relative_dest_dir / filename), "abs_path": str(target), "size": total, "is_text_file": is_text, diff --git a/tests/workspace/test_workspace_routes.py b/tests/workspace/test_workspace_routes.py index 4e1c6b674..80855bd30 100644 --- a/tests/workspace/test_workspace_routes.py +++ b/tests/workspace/test_workspace_routes.py @@ -392,6 +392,66 @@ def test_upload_to_nonexistent_dest_creates_it(self, workspace_client): assert r.status_code == 200 assert (_ws(workspace_client) / "new_folder" / "x.txt").exists() + @pytest.mark.parametrize("purpose", [None, "chat"]) + def test_upload_to_dest_when_workspace_root_is_symlink( + self, + workspace_client, + tmp_path: Path, + purpose: str | None, + ): + ws = _ws(workspace_client) + link = tmp_path / "workspace-link" + try: + link.symlink_to(ws, target_is_directory=True) + except OSError as exc: + pytest.skip(f"symlink unavailable on this platform: {exc}") + + from flocks.workspace.manager import WorkspaceManager + + WorkspaceManager.get_instance()._workspace_dir = link + url = "/api/workspace/upload?dest=uploads" + if purpose: + url += f"&purpose={purpose}" + + response = _client(workspace_client).post( + url, + files=[("files", ("report.pdf", b"report", "application/pdf"))], + ) + + assert response.status_code == 200 + result = response.json()["uploaded"][0] + assert result.get("error") is None + assert result["path"] == "uploads/report.pdf" + assert result["abs_path"] == str(ws / "uploads" / "report.pdf") + assert (ws / "uploads" / "report.pdf").read_bytes() == b"report" + + def test_upload_to_root_when_workspace_root_is_symlink( + self, + workspace_client, + tmp_path: Path, + ): + ws = _ws(workspace_client) + link = tmp_path / "workspace-link" + try: + link.symlink_to(ws, target_is_directory=True) + except OSError as exc: + pytest.skip(f"symlink unavailable on this platform: {exc}") + + from flocks.workspace.manager import WorkspaceManager + + WorkspaceManager.get_instance()._workspace_dir = link + response = _client(workspace_client).post( + "/api/workspace/upload", + files=[("files", ("root.txt", b"root", "text/plain"))], + ) + + assert response.status_code == 200 + result = response.json()["uploaded"][0] + assert result.get("error") is None + assert result["path"] == "root.txt" + assert result["abs_path"] == str(link / "root.txt") + assert (ws / "root.txt").read_bytes() == b"root" + def test_upload_overwrites_duplicate_file_without_chat_purpose(self, workspace_client): client = _client(workspace_client) first = client.post( From 75f5b2d1bbb3b6c455094aa7b4b4c6ebb62524a5 Mon Sep 17 00:00:00 2001 From: duguwanglong Date: Wed, 9 Sep 2026 10:53:43 +0800 Subject: [PATCH 55/63] feat(notifications): add scheduled token policy reminders Start a persistent one-month campaign on first service launch and remind users after service restarts and every Monday at 10:00 Asia/Shanghai. Deduplicate delivery per user with recoverable display reservations. Show policy notices before pending upgrade reminders, preserve active upgrades, and stop polling after campaign expiry. Share the reminder card UI and cover scheduling, display recovery, background tabs, and upgrade ordering with regression tests. Validation: 58 backend tests and 57 frontend tests passed; production build and Ruff checks passed. --- flocks/cli/main.py | 2 + flocks/config/config.py | 10 + flocks/notifications/token_policy.py | 240 ++++++++++++++++ flocks/server/app.py | 9 + flocks/server/routes/notifications.py | 16 ++ tests/server/routes/test_token_policy.py | 216 +++++++++++++++ webui/src/api/tokenPolicy.ts | 56 ++++ webui/src/components/common/NoticeCard.tsx | 65 +++++ .../common/TokenPolicyNotice.test.tsx | 44 +++ .../components/common/TokenPolicyNotice.tsx | 55 ++++ webui/src/components/common/UpdateModal.tsx | 48 +--- webui/src/components/layout/Layout.test.tsx | 106 +++++++ webui/src/components/layout/Layout.tsx | 38 ++- webui/src/hooks/useTokenPolicyNotice.test.tsx | 262 ++++++++++++++++++ webui/src/hooks/useTokenPolicyNotice.ts | 200 +++++++++++++ webui/src/locales/en-US/notification.json | 2 + webui/src/locales/zh-CN/notification.json | 2 + 17 files changed, 1326 insertions(+), 45 deletions(-) create mode 100644 flocks/notifications/token_policy.py create mode 100644 tests/server/routes/test_token_policy.py create mode 100644 webui/src/api/tokenPolicy.ts create mode 100644 webui/src/components/common/NoticeCard.tsx create mode 100644 webui/src/components/common/TokenPolicyNotice.test.tsx create mode 100644 webui/src/components/common/TokenPolicyNotice.tsx create mode 100644 webui/src/hooks/useTokenPolicyNotice.test.tsx create mode 100644 webui/src/hooks/useTokenPolicyNotice.ts diff --git a/flocks/cli/main.py b/flocks/cli/main.py index 0c8fc3bd3..38e352279 100644 --- a/flocks/cli/main.py +++ b/flocks/cli/main.py @@ -365,6 +365,8 @@ def serve( import uvicorn os.environ["_FLOCKS_SERVER_PORT"] = str(port) + # One id per service launch, inherited by all uvicorn workers/reload children. + os.environ["_FLOCKS_SERVICE_BOOT_ID"] = secrets_lib.token_hex(16) console.print(Panel(logo(), border_style="cyan")) console.print(f"[cyan]Starting server on:[/cyan] http://{host}:{port}") diff --git a/flocks/config/config.py b/flocks/config/config.py index c4c01a13c..fc4a88dc4 100644 --- a/flocks/config/config.py +++ b/flocks/config/config.py @@ -748,6 +748,13 @@ def normalize_fallback_provider_entries(value: Any) -> List[Dict[str, str]]: return normalized +class TokenPolicyNoticeConfig(BaseModel): + """Enable or disable the built-in, one-month token policy notice.""" + + model_config = {"extra": "forbid"} + enabled: bool = True + + class ConfigInfo(BaseModel): """ Main configuration schema @@ -757,6 +764,9 @@ class ConfigInfo(BaseModel): model_config = {"extra": "allow", "populate_by_name": True} # Allow extra fields for flexibility schema_: Optional[str] = Field(None, alias="$schema") + token_policy_notice: TokenPolicyNoticeConfig | None = Field( + default=None, alias="tokenPolicyNotice" + ) theme: Optional[str] = None keybinds: Optional[KeybindsConfig] = None log_level: Optional[str] = Field(None, alias="logLevel") diff --git a/flocks/notifications/token_policy.py b/flocks/notifications/token_policy.py new file mode 100644 index 000000000..d3f9415bb --- /dev/null +++ b/flocks/notifications/token_policy.py @@ -0,0 +1,240 @@ +"""Time-limited reminders with recoverable, per-user display reservations.""" + +from __future__ import annotations + +import calendar +import hashlib +import json +import os +from collections.abc import Callable +from datetime import UTC, datetime, timedelta, timezone +from multiprocessing import current_process, parent_process +from typing import Literal, TypeVar +from uuid import uuid4 + +from pydantic import BaseModel, Field + +from flocks.config.config import Config +from flocks.storage.storage import Storage + +CAMPAIGN_ID = "token-policy-change" +CAMPAIGN_KEY = f"notifications/token-policy-campaign/{CAMPAIGN_ID}" +LEASE_SECONDS = 15 +BEIJING = timezone(timedelta(hours=8)) +T = TypeVar("T") + + +def service_boot_id() -> str: + existing = os.environ.get("_FLOCKS_SERVICE_BOOT_ID") + if existing: + return existing + # `flocks serve` sets the id before starting workers. Direct uvicorn + # workers/reload children share their supervisor's multiprocessing authkey. + boot = hashlib.sha256(bytes(current_process().authkey)).hexdigest() if parent_process() else uuid4().hex + return os.environ.setdefault("_FLOCKS_SERVICE_BOOT_ID", boot) + + +def utc_now() -> datetime: + return datetime.now(UTC) + + +def month_after(start: datetime) -> datetime: + start = start.astimezone(BEIJING) + year, month = start.year + (start.month == 12), start.month % 12 + 1 + return start.replace(year=year, month=month, day=min(start.day, calendar.monthrange(year, month)[1])) + + +class Campaign(BaseModel): + starts_at: datetime + + @property + def expires_at(self) -> datetime: + return month_after(self.starts_at) + + +class PolicyNotice(BaseModel): + id: str = CAMPAIGN_ID + occurrence_id: str + expires_at: datetime + + +class PolicyStatus(BaseModel): + state: Literal["active", "finished", "disabled", "unavailable"] + server_now: datetime + notice: PolicyNotice | None = None + next_check_at: datetime | None = None + waiting_for_display: bool = False + lease_expires_at: datetime | None = None + + +class PolicyClaim(BaseModel): + occurrence_id: str = Field(pattern=r"^[a-f0-9]{64}$") + request_id: str = Field(pattern=r"^[a-zA-Z0-9-]{1,80}$") + + +class PolicyConfirmation(BaseModel): + confirmed: bool + + +def week_slot(now: datetime, start: datetime) -> tuple[str | None, datetime]: + local = now.astimezone(BEIJING) + monday = (local - timedelta(days=local.weekday())).replace(hour=10, minute=0, second=0, microsecond=0) + # Never replay the previous week's slot on Monday before 10. + slot = monday.isoformat() if start <= monday <= now else None + return slot, monday if now < monday else monday + timedelta(days=7) + + +async def _update_record(key: str, change: Callable[[dict], T]) -> T: + """Serialize campaign initialization and delivery transitions across workers.""" + await Storage.init(Storage.get_db_path()) + async with Storage.connect() as db: + try: + await db.execute("BEGIN IMMEDIATE") + async with db.execute("SELECT value FROM storage WHERE key = ?", (key,)) as cursor: + row = await cursor.fetchone() + record = json.loads(row[0]) if row else {} + before = json.dumps(record) + result = change(record) + after = json.dumps(record) + if after != before: + now = utc_now().isoformat() + await db.execute( + "INSERT INTO storage (key, value, type, created_at, updated_at) VALUES (?, ?, 'json', ?, ?) " + "ON CONFLICT(key) DO UPDATE SET value = excluded.value, updated_at = excluded.updated_at", + (key, after, now, now), + ) + await db.commit() + return result + except BaseException: + await db.rollback() + raise + + +class TokenPolicyService: + @staticmethod + async def initialize_campaign() -> Campaign: + def initialize(record: dict) -> Campaign: + record.setdefault("starts_at", utc_now().isoformat()) + return Campaign.model_validate(record) + + return await _update_record(CAMPAIGN_KEY, initialize) + + @staticmethod + async def load_campaign() -> Campaign | None: + """Read only: polling must never create or restart the campaign clock.""" + config = (await Config.get()).token_policy_notice + if config is not None and not config.enabled: + return None + return await Storage.get(CAMPAIGN_KEY, Campaign) + + @staticmethod + def state_key(user_id: str) -> str: + return f"notifications/token-policy/{CAMPAIGN_ID}/{user_id}" + + @staticmethod + def evaluate(campaign: Campaign, now: datetime, boot: str, record: dict) -> PolicyStatus: + result = PolicyStatus(state="active", server_now=now) + if now >= campaign.expires_at: + result.state = "finished" + return result + slot, next_monday = week_slot(now, campaign.starts_at) + result.next_check_at = min(next_monday, campaign.expires_at) + if record.get("shown_boot") == boot and (slot is None or record.get("shown_week") == slot): + return result + occurrence = hashlib.sha256(f"{CAMPAIGN_ID}:{boot}:{slot}".encode()).hexdigest() + result.notice = PolicyNotice(occurrence_id=occurrence, expires_at=campaign.expires_at) + return result + + @staticmethod + def active_lease(record: dict, status: PolicyStatus) -> dict | None: + lease = record.get("lease") + if ( + lease + and status.notice + and lease["occurrence_id"] == status.notice.occurrence_id + and datetime.fromisoformat(lease["expires_at"]) > status.server_now + ): + return lease + return None + + @staticmethod + def defer_to_lease(status: PolicyStatus, lease: dict) -> PolicyStatus: + status.notice = None + status.waiting_for_display = True + status.next_check_at = min(status.next_check_at, datetime.fromisoformat(lease["expires_at"])) + return status + + @classmethod + async def status_for_user(cls, user_id: str) -> PolicyStatus: + campaign = await cls.load_campaign() + if campaign is None: + config = (await Config.get()).token_policy_notice + state = "disabled" if config is not None and not config.enabled else "unavailable" + return PolicyStatus(state=state, server_now=utc_now()) + record = await Storage.get(cls.state_key(user_id)) or {} + status = cls.evaluate(campaign, utc_now(), service_boot_id(), record) + lease = cls.active_lease(record, status) + return cls.defer_to_lease(status, lease) if lease else status + + @classmethod + async def claim(cls, user_id: str, claim: PolicyClaim) -> PolicyStatus: + campaign = await cls.load_campaign() + if campaign is None: + return await cls.status_for_user(user_id) + + def reserve(record: dict) -> PolicyStatus: + now, boot = utc_now(), service_boot_id() + status = cls.evaluate(campaign, now, boot, record) + if not status.notice or status.notice.occurrence_id != claim.occurrence_id: + status.notice = None + return status + lease = cls.active_lease(record, status) + if lease and lease["request_id"] != claim.request_id: + return cls.defer_to_lease(status, lease) + if lease is None: + slot, _ = week_slot(now, campaign.starts_at) + lease = { + **claim.model_dump(), + "boot_id": boot, + "week_slot": slot, + "expires_at": min(now + timedelta(seconds=LEASE_SECONDS), campaign.expires_at).isoformat(), + } + record["lease"] = lease + # An HTTP retry keeps the original deadline; it cannot monopolize delivery. + status.lease_expires_at = datetime.fromisoformat(lease["expires_at"]) + return status + + return await _update_record(cls.state_key(user_id), reserve) + + @classmethod + async def confirm_display(cls, user_id: str, claim: PolicyClaim) -> PolicyConfirmation: + campaign = await cls.load_campaign() + if campaign is None: + return PolicyConfirmation(confirmed=False) + + def confirm(record: dict) -> PolicyConfirmation: + now, boot = utc_now(), service_boot_id() + if now >= campaign.expires_at: + return PolicyConfirmation(confirmed=False) + # A lost confirmation response can be retried after lease expiry. + if ( + record.get("shown_request") == claim.request_id + and record.get("shown_occurrence") == claim.occurrence_id + and record.get("shown_boot") == boot + ): + return PolicyConfirmation(confirmed=True) + status = cls.evaluate(campaign, now, boot, record) + lease = cls.active_lease(record, status) + if not lease or lease["request_id"] != claim.request_id or lease["occurrence_id"] != claim.occurrence_id: + return PolicyConfirmation(confirmed=False) + record.update( + shown_boot=boot, + shown_week=lease["week_slot"] or record.get("shown_week"), + shown_request=claim.request_id, + shown_occurrence=claim.occurrence_id, + shown_at=now.isoformat(), + ) + record.pop("lease", None) + return PolicyConfirmation(confirmed=True) + + return await _update_record(cls.state_key(user_id), confirm) diff --git a/flocks/server/app.py b/flocks/server/app.py index fce464789..d046aac0e 100644 --- a/flocks/server/app.py +++ b/flocks/server/app.py @@ -112,6 +112,9 @@ async def _runner() -> None: @asynccontextmanager async def lifespan(app: FastAPI): """Handle application lifecycle""" + from flocks.notifications.token_policy import service_boot_id + + service_boot_id() # Ensure file logging when server is started without CLI (e.g. uvicorn app:app) if Log._writer is None: await Log.init(print=False, dev=False, level=LogLevel.INFO) @@ -188,6 +191,12 @@ async def lifespan(app: FastAPI): # Initialize storage await _run_startup_phase(log, "storage.init", Storage.init) + from flocks.notifications.token_policy import TokenPolicyService + + try: + await _run_startup_phase(log, "notifications.token_policy.init", TokenPolicyService.initialize_campaign) + except Exception as exc: + log.warning("notifications.token_policy.init_failed", {"error": str(exc)}) log.info("storage.initialized") async def _recover_orphan_tool_parts() -> None: diff --git a/flocks/server/routes/notifications.py b/flocks/server/routes/notifications.py index 0792abd3b..d3e3395cd 100644 --- a/flocks/server/routes/notifications.py +++ b/flocks/server/routes/notifications.py @@ -14,10 +14,26 @@ NotificationService, ) from flocks.server.auth import require_user +from flocks.notifications.token_policy import PolicyClaim, PolicyConfirmation, PolicyStatus, TokenPolicyService router = APIRouter() +@router.get("/token-policy", response_model=PolicyStatus) +async def token_policy_status(request: Request) -> PolicyStatus: + return await TokenPolicyService.status_for_user(require_user(request).id) + + +@router.post("/token-policy/claim", response_model=PolicyStatus) +async def claim_token_policy(request: Request, claim: PolicyClaim) -> PolicyStatus: + return await TokenPolicyService.claim(require_user(request).id, claim) + + +@router.post("/token-policy/displayed", response_model=PolicyConfirmation) +async def confirm_token_policy_display(request: Request, claim: PolicyClaim) -> PolicyConfirmation: + return await TokenPolicyService.confirm_display(require_user(request).id, claim) + + @router.get( "/active", response_model=list[NotificationResponse], diff --git a/tests/server/routes/test_token_policy.py b/tests/server/routes/test_token_policy.py new file mode 100644 index 000000000..f1e2908a4 --- /dev/null +++ b/tests/server/routes/test_token_policy.py @@ -0,0 +1,216 @@ +from __future__ import annotations + +import asyncio +from datetime import datetime, timedelta +from types import SimpleNamespace + +import pytest +from pydantic import ValidationError + +from flocks.config.config import Config, ConfigInfo, TokenPolicyNoticeConfig +from flocks.notifications import token_policy as policy +from flocks.notifications.token_policy import PolicyClaim, TokenPolicyService +from flocks.storage.storage import Storage + + +def at(value: str) -> datetime: + return datetime.fromisoformat(value) + + +@pytest.fixture +async def clock(monkeypatch): + current = [at("2026-09-09T09:00:00+08:00")] + monkeypatch.setattr(policy, "utc_now", lambda: current[0]) + monkeypatch.setenv("_FLOCKS_SERVICE_BOOT_ID", "boot-1") + await TokenPolicyService.initialize_campaign() + return current + + +async def reserve(user="alice", request="request-1"): + status = await TokenPolicyService.status_for_user(user) + assert status.notice is not None + claim = PolicyClaim(occurrence_id=status.notice.occurrence_id, request_id=request) + return await TokenPolicyService.claim(user, claim), claim + + +async def display(user="alice", request="request-1"): + status, claim = await reserve(user, request) + assert status.notice + assert (await TokenPolicyService.confirm_display(user, claim)).confirmed + return status + + +async def test_startup_starts_one_persistent_month(clock, monkeypatch): + initial = await TokenPolicyService.load_campaign() + assert initial.starts_at == clock[0] + assert initial.expires_at == at("2026-10-09T09:00:00+08:00") + await display() + assert (await TokenPolicyService.status_for_user("alice")).notice is None + clock[0] += timedelta(days=1) + monkeypatch.setenv("_FLOCKS_SERVICE_BOOT_ID", "boot-2") + assert (await TokenPolicyService.initialize_campaign()).starts_at == initial.starts_at + await display(request="restart") + assert (await TokenPolicyService.status_for_user("alice")).notice is None + assert (await TokenPolicyService.status_for_user("bob")).notice + + +async def test_query_never_initializes_campaign(clock): + await Storage.delete(policy.CAMPAIGN_KEY) + assert (await TokenPolicyService.status_for_user("alice")).state == "unavailable" + assert await Storage.get(policy.CAMPAIGN_KEY) is None + + +async def test_monday_ten_and_late_login_once(clock): + await display() + clock[0] = at("2026-09-14T09:59:59+08:00") + status = await TokenPolicyService.status_for_user("alice") + assert status.notice is None + assert status.next_check_at == at("2026-09-14T10:00:00+08:00") + clock[0] += timedelta(seconds=1) + assert (await TokenPolicyService.status_for_user("alice")).notice + clock[0] = at("2026-09-15T18:00:00+08:00") + await display(request="week-2") + assert (await TokenPolicyService.status_for_user("alice")).notice is None + + +async def test_restart_and_weekly_merge_without_missed_week_queue(clock, monkeypatch): + await display() + clock[0] = at("2026-09-28T10:00:00+08:00") + monkeypatch.setenv("_FLOCKS_SERVICE_BOOT_ID", "boot-2") + await display(request="combined") + assert (await TokenPolicyService.status_for_user("alice")).notice is None + + +async def test_previous_week_not_replayed_monday_before_ten(clock): + await display() + clock[0] = at("2026-09-21T09:59:00+08:00") + assert (await TokenPolicyService.status_for_user("alice")).notice is None + + +async def test_expiry_stops_restart_claims_and_confirmation(clock, monkeypatch): + status, claim = await reserve() + clock[0] = status.notice.expires_at + monkeypatch.setenv("_FLOCKS_SERVICE_BOOT_ID", "boot-2") + status = await TokenPolicyService.status_for_user("alice") + assert status.state == "finished" + assert status.next_check_at is None + assert (await TokenPolicyService.claim("alice", claim)).notice is None + assert not (await TokenPolicyService.confirm_display("alice", claim)).confirmed + + +async def test_concurrent_tabs_only_one_lease_then_persistent_display(clock): + candidate = (await TokenPolicyService.status_for_user("alice")).notice + claims = [PolicyClaim(occurrence_id=candidate.occurrence_id, request_id=f"tab-{i}") for i in range(4)] + results = await asyncio.gather(*(TokenPolicyService.claim("alice", claim) for claim in claims)) + winners = [i for i, result in enumerate(results) if result.notice] + assert len(winners) == 1 + waiting = await TokenPolicyService.status_for_user("alice") + assert waiting.waiting_for_display + assert waiting.next_check_at == clock[0] + timedelta(seconds=policy.LEASE_SECONDS) + winner = claims[winners[0]] + assert (await TokenPolicyService.confirm_display("alice", winner)).confirmed + status = await TokenPolicyService.status_for_user("alice") + assert not status.waiting_for_display + assert status.notice is None + + +async def test_lost_claim_or_refresh_recovers_after_lease_expiry(clock): + _, abandoned = await reserve(request="old-page") + record = await Storage.get(TokenPolicyService.state_key("alice")) + assert "shown_at" not in record + clock[0] += timedelta(seconds=policy.LEASE_SECONDS) + _, replacement = await reserve(request="new-page") + assert not (await TokenPolicyService.confirm_display("alice", abandoned)).confirmed + assert (await TokenPolicyService.confirm_display("alice", replacement)).confirmed + + +async def test_lease_retry_does_not_extend_deadline_and_confirmation_is_idempotent(clock): + initial, claim = await reserve() + clock[0] += timedelta(seconds=2) + assert (await TokenPolicyService.claim("alice", claim)).lease_expires_at == initial.lease_expires_at + assert (await TokenPolicyService.confirm_display("alice", claim)).confirmed + clock[0] += timedelta(seconds=policy.LEASE_SECONDS) + assert (await TokenPolicyService.confirm_display("alice", claim)).confirmed + assert (await TokenPolicyService.status_for_user("alice")).notice is None + + +async def test_stale_boot_confirmation_and_claim_cannot_consume_new_boot(clock, monkeypatch): + _, old_claim = await reserve() + monkeypatch.setenv("_FLOCKS_SERVICE_BOOT_ID", "boot-2") + assert not (await TokenPolicyService.confirm_display("alice", old_claim)).confirmed + assert (await TokenPolicyService.claim("alice", old_claim)).notice is None + await display(request="new-boot") + + +async def test_stale_week_confirmation_cannot_consume_new_week(clock): + clock[0] = at("2026-09-14T09:59:55+08:00") + _, old_claim = await reserve() + clock[0] = at("2026-09-14T10:00:00+08:00") + assert not (await TokenPolicyService.confirm_display("alice", old_claim)).confirmed + await display(request="monday") + + +def test_calendar_month_end_and_configuration_scope(): + assert policy.month_after(at("2026-01-31T10:00:00+08:00")) == at("2026-02-28T10:00:00+08:00") + assert policy.month_after(at("2026-12-31T10:00:00+08:00")) == at("2027-01-31T10:00:00+08:00") + with pytest.raises(ValidationError): + TokenPolicyNoticeConfig(startsAt="2026-09-09T10:00:00+08:00") + + +def test_workers_share_supervisor_boot_id(monkeypatch): + monkeypatch.delenv("_FLOCKS_SERVICE_BOOT_ID", raising=False) + monkeypatch.setattr(policy, "parent_process", lambda: object()) + monkeypatch.setattr(policy, "current_process", lambda: SimpleNamespace(authkey=b"supervisor-one")) + first = policy.service_boot_id() + monkeypatch.delenv("_FLOCKS_SERVICE_BOOT_ID") + assert policy.service_boot_id() == first + monkeypatch.delenv("_FLOCKS_SERVICE_BOOT_ID") + monkeypatch.setattr(policy, "current_process", lambda: SimpleNamespace(authkey=b"supervisor-two")) + assert policy.service_boot_id() != first + + +def test_unrelated_config_preserves_campaign_switch(): + merged = Config.merge_config_concat_arrays( + ConfigInfo(tokenPolicyNotice={"enabled": False}), ConfigInfo(theme="dark") + ) + assert merged.token_policy_notice.enabled is False + + +async def test_disabled_campaign_rejects_inflight_delivery(clock, monkeypatch): + _, claim = await reserve() + + async def disabled_config(): + return ConfigInfo(tokenPolicyNotice={"enabled": False}) + + monkeypatch.setattr(Config, "get", disabled_config) + assert (await TokenPolicyService.status_for_user("alice")).state == "disabled" + assert not (await TokenPolicyService.confirm_display("alice", claim)).confirmed + assert (await TokenPolicyService.claim("alice", claim)).notice is None + + +async def test_routes_auth_and_validation(client, clock): + response = await client.get("/api/notifications/token-policy") + assert response.status_code == 200 + body = {"occurrence_id": response.json()["notice"]["occurrence_id"], "request_id": "browser-1"} + response = await client.post("/api/notifications/token-policy/claim", json=body) + assert response.status_code == 200 + assert response.json()["lease_expires_at"] + response = await client.post("/api/notifications/token-policy/displayed", json=body) + assert response.json() == {"confirmed": True} + assert (await client.get("/api/notifications/token-policy")).json()["notice"] is None + for action in ["claim", "displayed"]: + response = await client.post( + f"/api/notifications/token-policy/{action}", json={"occurrence_id": "invalid", "request_id": "browser-1"} + ) + assert response.status_code == 422 + from flocks.auth.service import AuthService + + if not await AuthService.has_users(): + await AuthService.bootstrap_admin(username="admin", password="Password123!") + for method, path in [ + ("GET", "/api/notifications/token-policy"), + ("POST", "/api/notifications/token-policy/claim"), + ("POST", "/api/notifications/token-policy/displayed"), + ]: + response = await client.request(method, path, headers={"sec-fetch-mode": "cors"}, json=body) + assert response.status_code == 401 diff --git a/webui/src/api/tokenPolicy.ts b/webui/src/api/tokenPolicy.ts new file mode 100644 index 000000000..55a5c9f81 --- /dev/null +++ b/webui/src/api/tokenPolicy.ts @@ -0,0 +1,56 @@ +import client from './client'; + +export interface TokenPolicyNotice { + id: string; + occurrence_id: string; + expires_at: string; +} + +export interface TokenPolicyStatus { + state: 'active' | 'finished' | 'disabled' | 'unavailable'; + notice: TokenPolicyNotice | null; + server_now: string; + next_check_at: string | null; + waiting_for_display: boolean; + lease_expires_at: string | null; +} + +export async function confirmTokenPolicyDisplay( + occurrenceId: string, + requestId: string, + signal?: AbortSignal, +): Promise { + const { data } = await client.post<{ confirmed: boolean }>( + '/api/notifications/token-policy/displayed', + { + occurrence_id: occurrenceId, + request_id: requestId, + }, + { timeout: 3000, signal }, + ); + return data.confirmed; +} + +export async function getTokenPolicyStatus(signal?: AbortSignal): Promise { + const { data } = await client.get('/api/notifications/token-policy', { + timeout: 3000, + signal, + }); + return data; +} + +export async function claimTokenPolicy( + occurrenceId: string, + requestId: string, + signal?: AbortSignal, +): Promise { + const { data } = await client.post( + '/api/notifications/token-policy/claim', + { + occurrence_id: occurrenceId, + request_id: requestId, + }, + { timeout: 3000, signal }, + ); + return data; +} diff --git a/webui/src/components/common/NoticeCard.tsx b/webui/src/components/common/NoticeCard.tsx new file mode 100644 index 000000000..094ea3f42 --- /dev/null +++ b/webui/src/components/common/NoticeCard.tsx @@ -0,0 +1,65 @@ +import { useId, type ButtonHTMLAttributes, type ReactNode } from 'react'; +import { X, type LucideIcon } from 'lucide-react'; + +interface NoticeCardProps { + title: ReactNode; + subtitle?: ReactNode; + icon: LucideIcon; + closeLabel: string; + onClose: () => void; + children: ReactNode; +} + +/** Shared appearance for automatic reminders; upgrade progress stays separate. */ +export default function NoticeCard({ + title, + subtitle, + icon: Icon, + closeLabel, + onClose, + children, +}: NoticeCardProps) { + const titleId = useId(); + return ( +
+
event.stopPropagation()} + > +
+
+ + +
+
+ {title} +
+ {subtitle &&
{subtitle}
} +
+
+ +
+ {children} +
+
+ ); +} + +export function NoticePrimaryAction({ className = '', ...props }: ButtonHTMLAttributes) { + return ( +
) : ( -
-
e.stopPropagation()} - > -
-
- - - -
-
- {info?.has_update ? t('newVersionTitle') : modalTitle} -
- {latestDisplayVersion && ( -
{formatUpdateVersion(latestDisplayVersion)}
- )} -
-
- -
- +
{info?.has_update ? ( <> @@ -448,17 +432,13 @@ export default function UpdateModal({ )} {canUpgrade && info?.has_update && info.update_allowed !== false && ( - + )}
-
-
+ ), document.body, ); diff --git a/webui/src/components/layout/Layout.test.tsx b/webui/src/components/layout/Layout.test.tsx index 9e7347f06..2bc280920 100644 --- a/webui/src/components/layout/Layout.test.tsx +++ b/webui/src/components/layout/Layout.test.tsx @@ -6,6 +6,13 @@ import { MemoryRouter, Route, Routes, useLocation } from 'react-router-dom'; import Layout from './Layout'; import Home from '@/pages/Home'; import { UPDATE_DISMISSED_KEY } from '@/utils/updateDismissal'; +import { getTokenPolicyStatus, claimTokenPolicy, confirmTokenPolicyDisplay } from '@/api/tokenPolicy'; + +vi.mock('@/api/tokenPolicy', () => ({ + getTokenPolicyStatus: vi.fn(), + claimTokenPolicy: vi.fn(), + confirmTokenPolicyDisplay: vi.fn(), +})); const { catalogAPI, @@ -157,6 +164,7 @@ vi.mock('@/components/common/LanguageSwitcher', () => ({ vi.mock('@/components/common/UpdateModal', () => ({ UPDATE_DISMISSED_KEY: 'update-dismissed', default: (props: Record) => { + React.useEffect(() => { (props.onPresented as (() => void) | undefined)?.(); }, [props.onPresented]); updateModalMock(props); return
; }, @@ -285,6 +293,9 @@ describe('Layout onboarding entry', () => { beforeEach(() => { vi.clearAllMocks(); vi.useRealTimers(); + vi.mocked(confirmTokenPolicyDisplay).mockResolvedValue(true); + vi.mocked(getTokenPolicyStatus).mockResolvedValue({ state: 'active', waiting_for_display: false, lease_expires_at: null, notice: null, server_now: new Date().toISOString(), next_check_at: null }); + vi.mocked(claimTokenPolicy).mockResolvedValue({ state: 'active', waiting_for_display: false, lease_expires_at: null, notice: null, server_now: new Date().toISOString(), next_check_at: null }); localStorage.clear(); productNameContextValue.productName = 'Flocks'; productNameContextValue.proProductName = 'Flocks Pro'; @@ -383,6 +394,101 @@ describe('Layout onboarding entry', () => { sessionApi.create.mockResolvedValue({ id: 'session-1' }); }); + it.each(['gotIt', 'close'])('shows the policy before a pending update, then resumes the update on %s', async (action) => { + const policyStatus = { + state: 'active' as const, waiting_for_display: false, lease_expires_at: new Date(Date.now() + 15000).toISOString(), + notice: { id: 'policy', occurrence_id: 'a'.repeat(64), expires_at: new Date(Date.now() + 86400000).toISOString() }, + server_now: new Date().toISOString(), + next_check_at: null, + }; + const pending = deferred(); + vi.mocked(getTokenPolicyStatus).mockReturnValueOnce(pending.promise); + vi.mocked(claimTokenPolicy).mockResolvedValueOnce(policyStatus); + checkUpdate.mockResolvedValue({ has_update: true, current_version: '0.2.0', latest_version: '0.3.0', error: null }); + renderHomeWithLayout(); + await waitFor(() => expect(checkUpdate).toHaveBeenCalled()); + expect(screen.queryByRole('dialog', { name: 'update-modal' })).not.toBeInTheDocument(); + await act(async () => pending.resolve(policyStatus)); + expect(await screen.findByRole('dialog', { name: 'tokenPolicyTitle' })).toBeInTheDocument(); + expect(screen.queryByRole('dialog', { name: 'update-modal' })).not.toBeInTheDocument(); + await userEvent.setup().click(screen.getByRole('button', { name: action })); + expect(await screen.findByRole('dialog', { name: 'update-modal' })).toBeInTheDocument(); + expect(ackNotification).not.toHaveBeenCalled(); + expect(localStorage.getItem(UPDATE_DISMISSED_KEY)).toBeNull(); + }); + + it('does not replace an already presented upgrade when a policy becomes due', async () => { + checkUpdate.mockResolvedValue({ has_update: true, current_version: '0.2.0', latest_version: '0.3.0', error: null }); + renderHomeWithLayout(); + const update = await screen.findByRole('dialog', { name: 'update-modal' }); + const policyStatus = { + state: 'active' as const, waiting_for_display: false, lease_expires_at: new Date(Date.now() + 15000).toISOString(), + notice: { id: 'policy', occurrence_id: 'a'.repeat(64), expires_at: new Date(Date.now() + 86400000).toISOString() }, + server_now: new Date().toISOString(), next_check_at: null, + }; + vi.mocked(getTokenPolicyStatus).mockResolvedValue(policyStatus); + vi.mocked(claimTokenPolicy).mockResolvedValue(policyStatus); + fireEvent.focus(window); + await flushEffects(); + expect(claimTokenPolicy).not.toHaveBeenCalled(); + expect(screen.getByRole('dialog', { name: 'update-modal' })).toBe(update); + const props = updateModalMock.mock.calls.at(-1)![0] as unknown as { onClose: () => void }; + act(() => props.onClose()); + expect(await screen.findByRole('dialog', { name: 'tokenPolicyTitle' })).toBeInTheDocument(); + }); + + it('keeps background upgrades pending and shows policy first on return', async () => { + const visibility = vi.spyOn(document, 'visibilityState', 'get').mockReturnValue('hidden'); + const policyStatus = { + state: 'active' as const, waiting_for_display: false, lease_expires_at: new Date(Date.now() + 15000).toISOString(), + notice: { id: 'policy', occurrence_id: 'a'.repeat(64), expires_at: new Date(Date.now() + 86400000).toISOString() }, + server_now: new Date().toISOString(), next_check_at: null, + }; + vi.mocked(getTokenPolicyStatus).mockResolvedValue(policyStatus); + vi.mocked(claimTokenPolicy).mockResolvedValue(policyStatus); + checkUpdate.mockResolvedValue({ has_update: true, current_version: '0.2.0', latest_version: '0.3.0', error: null }); + try { + renderHomeWithLayout(); + await waitFor(() => expect(checkUpdate).toHaveBeenCalled()); + expect(screen.queryByRole('dialog', { name: 'update-modal' })).not.toBeInTheDocument(); + expect(getTokenPolicyStatus).not.toHaveBeenCalled(); + visibility.mockReturnValue('visible'); + fireEvent(document, new Event('visibilitychange')); + expect(await screen.findByRole('dialog', { name: 'tokenPolicyTitle' })).toBeInTheDocument(); + expect(screen.queryByRole('dialog', { name: 'update-modal' })).not.toBeInTheDocument(); + await userEvent.setup().click(screen.getByRole('button', { name: 'gotIt' })); + expect(await screen.findByRole('dialog', { name: 'update-modal' })).toBeInTheDocument(); + } finally { + visibility.mockRestore(); + } + }); + + it('waits for an abandoned reservation to expire before allowing an upgrade', async () => { + const policyStatus = { + state: 'active' as const, waiting_for_display: false, lease_expires_at: new Date(Date.now() + 15000).toISOString(), + notice: { id: 'policy', occurrence_id: 'a'.repeat(64), expires_at: new Date(Date.now() + 86400000).toISOString() }, + server_now: new Date().toISOString(), next_check_at: null, + }; + vi.mocked(getTokenPolicyStatus).mockResolvedValueOnce({ + ...policyStatus, notice: null, lease_expires_at: null, waiting_for_display: true, + next_check_at: new Date(Date.now() + 600).toISOString(), + }).mockResolvedValue(policyStatus); + vi.mocked(claimTokenPolicy).mockResolvedValue(policyStatus); + checkUpdate.mockResolvedValue({ has_update: true, current_version: '0.2.0', latest_version: '0.3.0', error: null }); + renderHomeWithLayout(); + await waitFor(() => expect(checkUpdate).toHaveBeenCalled()); + expect(screen.queryByRole('dialog', { name: 'update-modal' })).not.toBeInTheDocument(); + expect(await screen.findByRole('dialog', { name: 'tokenPolicyTitle' })).toBeInTheDocument(); + expect(screen.queryByRole('dialog', { name: 'update-modal' })).not.toBeInTheDocument(); + }); + + it('still displays an upgrade when the policy status request fails', async () => { + vi.mocked(getTokenPolicyStatus).mockRejectedValue(new Error('timeout')); + checkUpdate.mockResolvedValue({ has_update: true, current_version: '0.2.0', latest_version: '0.3.0', error: null }); + renderHomeWithLayout(); + expect(await screen.findByRole('dialog', { name: 'update-modal' })).toBeInTheDocument(); + }); + it('opens onboarding from the home entry and shows configured details for an existing default model', async () => { const user = userEvent.setup(); localStorage.setItem('flocks_onboarding_dismissed', 'true'); diff --git a/webui/src/components/layout/Layout.tsx b/webui/src/components/layout/Layout.tsx index 0a50be13d..24a988a3c 100644 --- a/webui/src/components/layout/Layout.tsx +++ b/webui/src/components/layout/Layout.tsx @@ -30,6 +30,7 @@ import { useState, useEffect, useCallback, useMemo, useRef, lazy, Suspense } fro import type { ComponentType, CSSProperties, KeyboardEvent as ReactKeyboardEvent, PointerEvent as ReactPointerEvent } from 'react'; import { useTranslation } from 'react-i18next'; import { onboardingAPI } from '@/api/onboarding'; +import { useTokenPolicyNotice } from '@/hooks/useTokenPolicyNotice'; // Modals are only rendered after the user clicks/triggers them; pulling them // into the eager Layout chunk costs ~1.7k LOC + i18n keys + lucide icons that // the home page never needs. @@ -142,6 +143,7 @@ function saveSocDashboardTitle(title: string): void { const OnboardingModal = lazyLayoutComponent(() => import('@/components/common/OnboardingModal')); const UpdateModal = lazyLayoutComponent(() => import('@/components/common/UpdateModal'), ['update']); const NotificationModal = lazyLayoutComponent(() => import('@/components/common/NotificationModal'), ['notification']); +const TokenPolicyNotice = lazyLayoutComponent(() => import('@/components/common/TokenPolicyNotice'), ['notification']); import { checkUpdate, type VersionInfo } from '@/api/update'; import { consoleUpgradeApi } from '@/api/consoleUpgrade'; import { @@ -257,7 +259,17 @@ export default function Layout() { const accountMenuRef = useRef(null); const isHome = location.pathname === '/'; const [showOnboarding, setShowOnboarding] = useState(false); - const [showUpdate, setShowUpdate] = useState(false); + const [updateState, setUpdateState] = useState<'idle' | 'pending' | 'visible'>('idle'); + const tokenPolicy = useTokenPolicyNotice(user?.id, showOnboarding || updateState === 'visible'); + const handleUpdatePresented = useCallback(() => { + if (document.visibilityState === 'visible') setUpdateState('visible'); + }, []); + const requestUpdate = useCallback(() => { + setUpdateState((state) => state === 'visible' ? state : 'pending'); + }, []); + const closeUpdate = useCallback(() => { + setUpdateState('idle'); + }, []); const { t, i18n } = useTranslation('nav'); const { t: tWebUIContractPage } = useTranslation('webuiContractPage'); const { t: tAuth } = useTranslation('auth'); @@ -438,7 +450,7 @@ export default function Layout() { && !isUpdateDismissed(info, localStorage.getItem(UPDATE_DISMISSED_KEY)) ) { lastPromptedVersionRef.current = updateDismissalKey; - setShowUpdate(true); + requestUpdate(); } return; } @@ -453,7 +465,7 @@ export default function Layout() { checkingUpdateRef.current = false; setHasCompletedUpdateCheck(true); } - }, [canManageUpdates, flocksproStatusReady, i18n.language, isFlocksproActive]); + }, [canManageUpdates, flocksproStatusReady, i18n.language, isFlocksproActive, requestUpdate]); useEffect(() => { if (!flocksproStatusReady || !canManageUpdates) return undefined; @@ -620,7 +632,7 @@ export default function Layout() { const allNotifications = updateNotification ? [...notifications, updateNotification].sort((a, b) => a.priority - b.priority) : notifications; - const visibleNotifications = backendNotificationsReady && updateNotificationReady && !showOnboarding && !showUpdate && allNotifications.length > 0 + const visibleNotifications = tokenPolicy.ready && !tokenPolicy.notice && backendNotificationsReady && updateNotificationReady && !showOnboarding && updateState === 'idle' && allNotifications.length > 0 ? allNotifications : []; @@ -885,8 +897,8 @@ export default function Layout() { const openManualUpdateCheck = useCallback(() => { setAccountMenuOpen(false); setUpdateInfo(null); - setShowUpdate(true); - }, []); + requestUpdate(); + }, [requestUpdate]); return (
@@ -899,14 +911,18 @@ export default function Layout() { onClose={() => setShowOnboarding(false)} /> )} - {showUpdate && ( + {tokenPolicy.notice && !showOnboarding && updateState !== 'visible' && ( + + )} + {updateState !== 'idle' && (updateState === 'visible' || (tokenPolicy.ready && !tokenPolicy.notice && !showOnboarding)) && ( setShowUpdate(false)} - onDismiss={() => setShowUpdate(false)} + onPresented={handleUpdatePresented} + onClose={closeUpdate} + onDismiss={closeUpdate} /> )} {visibleNotifications.length > 0 && ( @@ -1008,7 +1024,7 @@ export default function Layout() { {hasVisibleProductUpdate && ( +
+
+
+ ( + + {children} + + ), + }} + > + {t('tokenPolicyBody')} + +
+
+
+ {t('gotIt')} +
+
- , + , document.body, ); } diff --git a/webui/src/locales/en-US/notification.json b/webui/src/locales/en-US/notification.json index 5998a64e6..1176e768d 100644 --- a/webui/src/locales/en-US/notification.json +++ b/webui/src/locales/en-US/notification.json @@ -1,7 +1,7 @@ { "title": "Flocks Notifications", - "tokenPolicyTitle": "Model token policy update", - "tokenPolicyBody": "The model token policy has changed. Please review the update.", + "tokenPolicyTitle": "Flocks Community Announcement | October Token Support Policy Changes", + "tokenPolicyBody": "Dear Flocks community,\n\nThank you for using Flocks, sharing feedback, and contributing since its launch on March 31. Alongside this community release, we are announcing the token support policy that will take effect in October so you can plan your usage and budget.\n\nTo support ongoing platform maintenance, reliability improvements, and long-term development of security operations use cases, Flocks is introducing monetary billing and self-service top-ups while retaining daily free support.\n\nThe current free token support policy remains in effect throughout September. Starting in October:\n\n1. **Daily free credit remains available.** Each person receives CNY 10 in free model usage credit per day. Consumption is calculated using the selected model’s pricing. Prices vary by model, so CNY 10 corresponds to different token amounts; it is not uniformly equivalent to 10 million tokens.\n2. **Temporary quota increases will be discontinued.** All additional daily token allowances previously granted to individual users will be removed, including the previous daily allowances of 30–50 million tokens.\n3. **The shared free token pool will be discontinued.** The previous shared free pool will no longer be available.\n4. **Self-service top-ups are available.** For usage beyond your daily free credit, visit [https://portal.agentflocks.com](https://portal.agentflocks.com), sign in with your X community account, and top up. The platform is already live.\n5. **Enterprise requests will be handled by sales.** Enterprise users should contact ThreatBook sales to discuss model usage and related services.\n\nThese changes concern token support and purchasing on the ThreatBook model platform. They will affect future allowances and usage costs for community members who use temporary increases or rely on the shared pool. We are providing advance notice and welcome specific questions in the group.\n\nIf you have scheduled tasks such as continuous inspections or alert analysis, review your models, task frequency, and daily consumption before October, and plan your credit and budget accordingly. If you encounter any issues while using or adjusting the service, post in this group or contact the group owner.\n\nFor enterprise purchasing, please contact ThreatBook sales.\n\nProject: [https://github.com/AgentFlocks/flocks](https://github.com/AgentFlocks/flocks) \nDocumentation: [https://agentflocks.github.io/flocks-docs/](https://agentflocks.github.io/flocks-docs/) \nTop-ups: [https://portal.agentflocks.com](https://portal.agentflocks.com)\n\nThank you for your continued support. We look forward to working with you to keep improving the security operations use cases already running successfully.\n\nThe Flocks Team \nSeptember 9, 2026", "subtitle": "Benefits and version highlights are collected here.", "close": "Close notification", "gotIt": "Got it", diff --git a/webui/src/locales/zh-CN/notification.json b/webui/src/locales/zh-CN/notification.json index 70a5c3521..a630fb634 100644 --- a/webui/src/locales/zh-CN/notification.json +++ b/webui/src/locales/zh-CN/notification.json @@ -1,7 +1,7 @@ { "title": "Flocks 通知", - "tokenPolicyTitle": "模型 Token 策略变更", - "tokenPolicyBody": "模型token策略发生变更，请确认", + "tokenPolicyTitle": "Flocks社区公告｜10月Token支持策略调整", + "tokenPolicyBody": "各位Flocks社区伙伴：\n\n感谢大家自3月31日发布以来的使用、反馈与共建。配合本次社区版本更新，我们提前同步10月起的Token奖励策略，方便大家安排后续使用和预算。\n\n为支持平台持续维护、稳定性建设和安全运营场景的长期迭代，Flocks在保留每日免费支持的基础上，引入按金额计费和自助充值机制。\n\n**9月内，现行免费Token支持策略持续有效，9月30日旧策略全部失效。10月1日起，调整如下：**\n\n1. **每天仍有免费额度。** 每人每天免费赠送10元模型调用额度，按所选模型的价格计算消耗。不同模型计价不同，10元对应的Token数量也不同，不统一折算为1000万Token。\n2. **统一取消临时加额。** 原先为个别用户单独增加的每日临时Token额度全部取消，包括原每日3000万～5000万额度。\n3. **取消公共免费Token池。** 原公共免费池机制不再保留。\n4. **支持自助充值。** 如需超出每日免费额度的调用，可前往[https://portal.agentflocks.com](https://portal.agentflocks.com)，使用X社区账号登录后自助充值。平台已上线。\n5. **企业需求由销售对接。** 企业用户请联系微步销售，沟通模型用量及相关服务需求。\n\n本次调整涉及微步模型平台的Token支持与购买方式。对使用临时加额或依赖公共池的伙伴，这会改变后续的可用额度和使用成本，我们提前和大家说明，也欢迎大家在群里反馈具体问题。\n\n如果你已配置持续巡检、告警研判等定时任务，建议在10月前梳理所用模型、任务频率和日常消耗，提前安排额度及预算。使用或调整中遇到问题，可在本群或戳群主反馈；\n\n企业采购需求请联系微步销售。\n\n项目地址：[https://github.com/AgentFlocks/flocks](https://github.com/AgentFlocks/flocks) \n使用文档：[https://agentflocks.github.io/flocks-docs/](https://agentflocks.github.io/flocks-docs/) \n充值入口：[https://portal.agentflocks.com](https://portal.agentflocks.com)\n\n感谢大家一直以来的支持。期待和大家一起，把已经跑通的安全运营场景持续用好、不断完善。\n\nFlocks团队 \n2026年9月9日", "subtitle": "福利提醒和版本变化都在这里。", "close": "关闭通知", "gotIt": "知道了", From 28aaf95bced5c0b5446c255e0934a72bd6142f3d Mon Sep 17 00:00:00 2001 From: duguwanglong Date: Wed, 9 Sep 2026 16:10:54 +0800 Subject: [PATCH 58/63] fix(notifications): retire token trial notice and render release notes as markdown Remove the built-in token free-period extension announcement. Render version update bodies with GitHub-flavored Markdown and safe external links. Add regression coverage for the retired notice and Markdown rendering while preserving per-user acknowledgement tests. Validation: 8 backend tests and 37 frontend tests passed; frontend build passed. --- flocks/notifications/service.py | 36 +---------------- .../routes/test_notifications_routes.py | 40 ++++++++++++++----- .../common/NotificationModal.test.tsx | 38 ++++++++++++++++++ .../components/common/NotificationModal.tsx | 19 ++++++++- 4 files changed, 87 insertions(+), 46 deletions(-) create mode 100644 webui/src/components/common/NotificationModal.test.tsx diff --git a/flocks/notifications/service.py b/flocks/notifications/service.py index 3cf3a02cd..a967bc75b 100644 --- a/flocks/notifications/service.py +++ b/flocks/notifications/service.py @@ -84,41 +84,7 @@ class NotificationAckStatus(BaseModel): acknowledged: bool -DEFAULT_NOTIFICATIONS: tuple[NotificationConfig, ...] = ( - NotificationConfig( - id="token-free-period-extended-2026-04", - kind="benefit", - priority=10, - starts_at="2026-03-30T00:00:00+08:00", - locales={ - "zh-CN": NotificationContent( - title="Token 免费期已延长", - summary="福利已自动生效，无需额外操作。", - body=( - "为了让你有更充足的时间体验 Flocks，我们已延长 token 免费使用期。" - ), - highlights=[ - "3月30日-4月29日注册的老用户，授权自动延期至60天", - "4月29日之后注册的新用户，依旧默认30天注册授权", - ], - primaryAction=NotificationAction(label="知道了"), - ), - "en-US": NotificationContent( - title="Token free period extended", - summary="The benefit is active automatically. No action is required.", - body=( - "We have extended the free token period so you have more time " - "to experience Flocks." - ), - highlights=[ - "Existing users who registered between March 30 and April 29 will have their authorization automatically extended to 60 days", - "New users who register after April 29 will still receive the default 30-day trial authorization", - ], - primaryAction=NotificationAction(label="Got it"), - ), - }, - ), -) +DEFAULT_NOTIFICATIONS: tuple[NotificationConfig, ...] = () class NotificationService: diff --git a/tests/server/routes/test_notifications_routes.py b/tests/server/routes/test_notifications_routes.py index a0400e801..a7d1fa75c 100644 --- a/tests/server/routes/test_notifications_routes.py +++ b/tests/server/routes/test_notifications_routes.py @@ -6,6 +6,26 @@ from flocks.notifications.service import NotificationService +@pytest.fixture +def builtin_notice(monkeypatch): + from flocks.notifications import service as notification_service + + monkeypatch.setattr(notification_service, "DEFAULT_NOTIFICATIONS", ( + notification_service.NotificationConfig( + id="test-benefit", + kind="benefit", + locales={"en-US": notification_service.NotificationContent(title="Test benefit")}, + ), + )) + + +@pytest.mark.asyncio +async def test_retired_token_notice_is_not_active(client: AsyncClient): + response = await client.get("/api/notifications/active", params={"locale": "zh-CN"}) + assert response.status_code == 200, response.text + assert "token-free-period-extended-2026-04" not in {item["id"] for item in response.json()} + + @pytest.mark.asyncio async def test_notifications_require_browser_login(client: AsyncClient): from flocks.auth.service import AuthService @@ -22,17 +42,17 @@ async def test_notifications_require_browser_login(client: AsyncClient): @pytest.mark.asyncio -async def test_active_notifications_and_dismiss_forever(client: AsyncClient): +async def test_active_notifications_and_dismiss_forever(client: AsyncClient, builtin_notice): response = await client.get( "/api/notifications/active", params={"locale": "zh-CN"}, ) assert response.status_code == 200, response.text items = response.json() - assert [item["id"] for item in items] == ["token-free-period-extended-2026-04"] + assert [item["id"] for item in items] == ["test-benefit"] assert items[0]["kind"] == "benefit" - ack_response = await client.post("/api/notifications/token-free-period-extended-2026-04/ack") + ack_response = await client.post("/api/notifications/test-benefit/ack") assert ack_response.status_code == 200, ack_response.text response = await client.get( @@ -51,10 +71,10 @@ async def test_active_notifications_and_dismiss_forever(client: AsyncClient): @pytest.mark.asyncio -async def test_notification_ack_is_per_user(): +async def test_notification_ack_is_per_user(builtin_notice): await NotificationService.acknowledge( user_id="user-a", - notification_id="token-free-period-extended-2026-04", + notification_id="test-benefit", ) user_a_items = await NotificationService.list_active( @@ -66,18 +86,18 @@ async def test_notification_ack_is_per_user(): locale="en-US", ) - assert "token-free-period-extended-2026-04" not in {item.id for item in user_a_items} - assert "token-free-period-extended-2026-04" in {item.id for item in user_b_items} + assert "test-benefit" not in {item.id for item in user_a_items} + assert "test-benefit" in {item.id for item in user_b_items} @pytest.mark.asyncio -async def test_config_notification_overrides_builtin(monkeypatch): +async def test_config_notification_overrides_builtin(monkeypatch, builtin_notice): from flocks.notifications import service as notification_service async def fake_load_config_notifications(): return [ notification_service.NotificationConfig( - id="token-free-period-extended-2026-04", + id="test-benefit", enabled=False, priority=999, locales={ @@ -98,7 +118,7 @@ async def fake_load_config_notifications(): user_id="user-a", locale="zh-CN", ) - assert "token-free-period-extended-2026-04" not in {item.id for item in items} + assert "test-benefit" not in {item.id for item in items} @pytest.mark.asyncio diff --git a/webui/src/components/common/NotificationModal.test.tsx b/webui/src/components/common/NotificationModal.test.tsx new file mode 100644 index 000000000..24824731a --- /dev/null +++ b/webui/src/components/common/NotificationModal.test.tsx @@ -0,0 +1,38 @@ +import { render, screen } from '@testing-library/react'; +import { describe, expect, it, vi } from 'vitest'; +import NotificationModal from './NotificationModal'; + +vi.mock('react-i18next', () => ({ + useTranslation: () => ({ t: (key: string) => key }), +})); + +describe('NotificationModal', () => { + it('renders release notes as Markdown with safe links', () => { + render( + alert(1)', + highlights: [], + priority: 100, + }]} + onAcknowledge={vi.fn()} + onClose={vi.fn()} + onDismissForever={vi.fn()} + />, + ); + + expect(screen.getByRole('heading', { level: 3, name: '记忆与自进化' })).toBeInTheDocument(); + expect(screen.getByRole('listitem')).toHaveTextContent('支持 Dream 和 /dream'); + expect(screen.getByText('Dream').tagName).toBe('STRONG'); + expect(screen.getByText('/dream').tagName).toBe('CODE'); + const link = screen.getByRole('link', { name: '项目地址' }); + expect(link).toHaveAttribute('href', 'https://github.com/AgentFlocks/flocks'); + expect(link).toHaveAttribute('target', '_blank'); + expect(link).toHaveAttribute('rel', 'noopener noreferrer'); + expect(screen.getByText('危险链接')).not.toHaveAttribute('href', 'javascript:alert(1)'); + expect(document.querySelector('script')).toBeNull(); + }); +}); diff --git a/webui/src/components/common/NotificationModal.tsx b/webui/src/components/common/NotificationModal.tsx index 9d700f8dc..a3516e0c3 100644 --- a/webui/src/components/common/NotificationModal.tsx +++ b/webui/src/components/common/NotificationModal.tsx @@ -1,5 +1,7 @@ import { createPortal } from 'react-dom'; import { useEffect } from 'react'; +import ReactMarkdown from 'react-markdown'; +import remarkGfm from 'remark-gfm'; import { Bell, CheckCircle, @@ -144,7 +146,22 @@ export default function NotificationModal({
{notification.body && ( -

{notification.body}

+ notification.kind === 'whats_new' ? ( +
+ ( + {children} + ), + }} + > + {notification.body} + +
+ ) : ( +

{notification.body}

+ ) )} {notification.highlights.length > 0 && ( From 0a3ebeb856c6f3403b7fb4fbedafef003107abd1 Mon Sep 17 00:00:00 2001 From: John Yin <10972267+john-yin2333@user.noreply.gitee.com> Date: Wed, 9 Sep 2026 16:13:59 +0800 Subject: [PATCH 59/63] Fix shell status and persist command results --- flocks/server/routes/session.py | 7 +- flocks/session/message.py | 7 +- flocks/session/runner.py | 138 +++++++++++++++--- .../routes/test_session_status_by_id.py | 106 ++++++++++++++ tests/session/test_runner_shell_hook.py | 80 +++++++++- 5 files changed, 309 insertions(+), 29 deletions(-) diff --git a/flocks/server/routes/session.py b/flocks/server/routes/session.py index 9f006665b..ed7abb8f1 100644 --- a/flocks/server/routes/session.py +++ b/flocks/server/routes/session.py @@ -783,7 +783,10 @@ async def _build_session_runtime_status(session: SessionModel) -> SessionRuntime # starts. Treat that window (and inter-prompt queue hand-offs) as queued # so API clients never mistake accepted work for completion. if runtime_status.type == "idle": - if SessionLoop.is_running(session.id): + if ( + Session.has_active_operations(session.id) + or SessionLoop.is_running(session.id) + ): runtime_status = SessionStatusBusy() elif ( _is_prompt_chain_active(session.id) @@ -5016,6 +5019,7 @@ class ShellRequest(BaseModel): async def run_shell_command(sessionID: str, request: ShellRequest, http_request: Request): """Run shell command""" from flocks.hooks.execution import ExecutionStopped + from flocks.server.routes.event import publish_event from flocks.session.runner import SessionRunner current_user = require_user(http_request) @@ -5038,6 +5042,7 @@ async def run_shell_command(sessionID: str, request: ShellRequest, http_request: agent=request.agent, command=request.command, model=model, + event_publish_callback=publish_event, ) except SessionNotFoundError as exc: raise HTTPException( diff --git a/flocks/session/message.py b/flocks/session/message.py index f3ac17bef..4ae9c3afd 100644 --- a/flocks/session/message.py +++ b/flocks/session/message.py @@ -2469,7 +2469,12 @@ async def update(cls, session_id: str, message_id: str, **updates) -> Optional[M # Update timestamp time_data = message.time if hasattr(message, 'time') else message.model_dump().get("time", {}) if isinstance(time_data, dict): - patch["time"] = {**time_data, "updated": int(datetime.now().timestamp() * 1000)} + requested_time = patch.get("time") + patch["time"] = { + **time_data, + **(requested_time if isinstance(requested_time, dict) else {}), + "updated": int(datetime.now().timestamp() * 1000), + } updated = message.model_copy(update=patch) messages[msg_index] = updated diff --git a/flocks/session/runner.py b/flocks/session/runner.py index 0bdf81cf9..02391feda 100644 --- a/flocks/session/runner.py +++ b/flocks/session/runner.py @@ -27,7 +27,15 @@ 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, TextPart +from flocks.session.message import ( + Message, + MessageInfo, + MessageRole, + TextPart, + ToolPart, + ToolStateCompleted, + ToolStateRunning, +) from flocks.session.prompt import SessionPrompt, SystemPromptBlock, TurnPromptContext from flocks.session.core.status import SessionStatus, SessionStatusRetry, SessionStatusBusy from flocks.session.core.defaults import ( @@ -1132,6 +1140,9 @@ async def shell( agent: str, command: str, model: Optional[Dict[str, str]] = None, + event_publish_callback: Optional[ + Callable[[str, Dict[str, Any]], Awaitable[None]] + ] = None, ) -> Dict[str, Any]: """ Execute a shell command in session context. @@ -1151,24 +1162,89 @@ async def shell( cwd = session.directory or os.getcwd() + async def _publish(event_type: str, payload: Dict[str, Any]) -> None: + if event_publish_callback is None: + return + try: + await event_publish_callback(event_type, payload) + except Exception as exc: + log.debug("runner.shell.publish_failed", { + "session_id": session_id, + "event_type": event_type, + "error": str(exc), + }) + async def _effect( execution_command: str = command, execution_cwd: str = cwd, ) -> Dict[str, Any]: + started_at_ms = int(time.time() * 1000) + user_part_id = Identifier.create("part") user_msg = await Message.create( session_id=session_id, role=MessageRole.USER, content="The following tool was executed by the user", agent=agent, + time={"created": started_at_ms}, + part_id=user_part_id, ) + assistant_part_id = Identifier.create("part") assistant_msg = await Message.create( session_id=session_id, role=MessageRole.ASSISTANT, content="", agent=agent, parent_id=user_msg.id, + providerID="builtin", + modelID="shell", + mode=agent, + path={"cwd": execution_cwd, "root": execution_cwd}, + time={"created": started_at_ms}, + part_id=assistant_part_id, + ) + + call_id = Identifier.create("call") + tool_part_id = Identifier.create("part") + tool_input = { + "command": execution_command, + "workdir": execution_cwd, + } + running_part = ToolPart( + id=tool_part_id, + messageID=assistant_msg.id, + sessionID=session_id, + callID=call_id, + tool="bash", + metadata=None, + state=ToolStateRunning( + input=tool_input, + title="Shell", + metadata={}, + time={"start": started_at_ms}, + ), ) + await Message.store_part(session_id, assistant_msg.id, running_part) + + await _publish("message.updated", { + "info": user_msg.model_dump(mode="json", by_alias=True), + }) + await _publish("message.part.updated", { + "part": { + "id": user_part_id, + "messageID": user_msg.id, + "sessionID": session_id, + "type": "text", + "text": "The following tool was executed by the user", + "time": {"start": started_at_ms}, + }, + }) + await _publish("message.updated", { + "info": assistant_msg.model_dump(mode="json", by_alias=True), + }) + await _publish("message.part.updated", { + "part": running_part.model_dump(mode="json", by_alias=True), + }) start_time = asyncio.get_event_loop().time() try: @@ -1196,6 +1272,7 @@ async def _effect( exit_code = -1 end_time = asyncio.get_event_loop().time() + finished_at_ms = int(time.time() * 1000) log.info("runner.shell", { "session_id": session_id, @@ -1204,25 +1281,48 @@ async def _effect( "duration_ms": int((end_time - start_time) * 1000), }) + completed_part = ToolPart( + id=tool_part_id, + messageID=assistant_msg.id, + sessionID=session_id, + callID=call_id, + tool="bash", + metadata=None, + state=ToolStateCompleted( + input=tool_input, + output=output, + title="Shell", + metadata={"exitCode": exit_code}, + time={"start": started_at_ms, "end": finished_at_ms}, + attachments=None, + ), + ) + stored_part = await Message.store_part( + session_id, + assistant_msg.id, + completed_part, + ) + updated_assistant = await Message.update( + session_id, + assistant_msg.id, + finish="stop", + time={"completed": finished_at_ms}, + ) + if updated_assistant is None: + raise RuntimeError( + f"Failed to finalize shell message {assistant_msg.id}" + ) + + await _publish("message.part.updated", { + "part": stored_part.model_dump(mode="json", by_alias=True), + }) + await _publish("message.updated", { + "info": updated_assistant.model_dump(mode="json", by_alias=True), + }) + return { - "info": { - "id": assistant_msg.id, - "sessionID": session_id, - "role": "assistant", - "agent": agent, - }, - "parts": [{ - "id": Identifier.create("part"), - "messageID": assistant_msg.id, - "sessionID": session_id, - "type": "tool", - "tool": "bash", - "state": { - "status": "completed", - "input": {"command": execution_command}, - "output": output, - }, - }], + "info": updated_assistant.model_dump(mode="json", by_alias=True), + "parts": [stored_part.model_dump(mode="json", by_alias=True)], } from flocks.session.tool_execution import ( diff --git a/tests/server/routes/test_session_status_by_id.py b/tests/server/routes/test_session_status_by_id.py index 39fb62efb..af98701d1 100644 --- a/tests/server/routes/test_session_status_by_id.py +++ b/tests/server/routes/test_session_status_by_id.py @@ -2,7 +2,9 @@ from __future__ import annotations +import asyncio from collections.abc import Callable, Iterator +from types import SimpleNamespace import pytest from fastapi import status @@ -121,6 +123,110 @@ def _hold_background_work(coro, **_kwargs) -> None: assert status_response.json()["isProcessing"] is True +@pytest.mark.asyncio +async def test_shell_is_busy_while_active_and_idle_after_completion( + client, + session_id: str, + monkeypatch, +) -> None: + started = asyncio.Event() + release = asyncio.Event() + + async def _hold_shell(**_kwargs): + started.set() + await release.wait() + return {"info": {"id": "msg_shell"}, "parts": []} + + monkeypatch.setattr( + "flocks.session.runner.SessionRunner.shell", + _hold_shell, + ) + + shell_request = asyncio.create_task( + client.post( + f"/api/session/{session_id}/shell", + json={"agent": "rex", "command": "printf done"}, + ) + ) + await asyncio.wait_for(started.wait(), timeout=1) + + try: + running_response = await client.get(f"/api/session/{session_id}/status") + finally: + release.set() + + shell_response = await asyncio.wait_for(shell_request, timeout=1) + completed_response = await client.get(f"/api/session/{session_id}/status") + + assert running_response.status_code == status.HTTP_200_OK + assert running_response.json()["status"] == {"type": "busy"} + assert running_response.json()["isProcessing"] is True + assert running_response.json()["pendingPromptCount"] == 0 + assert shell_response.status_code == status.HTTP_200_OK + assert completed_response.status_code == status.HTTP_200_OK + assert completed_response.json()["status"] == {"type": "idle"} + assert completed_response.json()["isProcessing"] is False + assert completed_response.json()["pendingPromptCount"] == 0 + assert not session_routes.Session.has_active_operations(session_id) + + +@pytest.mark.asyncio +async def test_shell_response_is_available_after_message_cache_reload( + client, + session_id: str, + monkeypatch, +) -> None: + from flocks.session.message import Message + + process = SimpleNamespace( + communicate=lambda: asyncio.sleep( + 0, + result=(b"PR741_SHELL_PERSISTENCE_OK", b""), + ), + returncode=0, + ) + + async def _create_subprocess(*_args, **_kwargs): + return process + + async def _run_lifecycle(_payload, effect, **_kwargs): + return await effect() + + monkeypatch.setattr( + "flocks.session.runner.asyncio.create_subprocess_shell", + _create_subprocess, + ) + monkeypatch.setattr( + "flocks.session.tool_execution.run_tool_execution_lifecycle", + _run_lifecycle, + ) + + shell_response = await client.post( + f"/api/session/{session_id}/shell", + json={"agent": "rex", "command": "printf PR741_SHELL_PERSISTENCE_OK"}, + ) + assert shell_response.status_code == status.HTTP_200_OK + assistant_id = shell_response.json()["info"]["id"] + + Message.invalidate_cache(session_id) + history_response = await client.get(f"/api/session/{session_id}/message") + + assert history_response.status_code == status.HTTP_200_OK + assistant = next( + item + for item in history_response.json() + if item["info"]["id"] == assistant_id + ) + assert assistant["info"]["finish"] == "stop" + tool_part = next(part for part in assistant["parts"] if part["type"] == "tool") + assert tool_part["tool"] == "bash" + assert tool_part["state"]["status"] == "completed" + assert tool_part["state"]["input"]["command"] == ( + "printf PR741_SHELL_PERSISTENCE_OK" + ) + assert tool_part["state"]["output"] == "PR741_SHELL_PERSISTENCE_OK" + + @pytest.mark.asyncio async def test_status_by_id_prefers_busy_and_counts_queued_prompts(client, session_id: str) -> None: from flocks.session.interaction_queue import InteractionQueue diff --git a/tests/session/test_runner_shell_hook.py b/tests/session/test_runner_shell_hook.py index 62e256cfd..4f5dca2f4 100644 --- a/tests/session/test_runner_shell_hook.py +++ b/tests/session/test_runner_shell_hook.py @@ -7,6 +7,7 @@ import pytest from flocks.hooks.pipeline import HookBase, HookPipeline +from flocks.session.message import Message, ToolPart from flocks.session.runner import SessionRunner from flocks.session.tool_execution import build_session_tool_execution_payload @@ -39,15 +40,11 @@ async def tool_before(self, ctx): create_process = AsyncMock(return_value=process) monkeypatch.setattr( "flocks.session.runner.Session.get_by_id", - AsyncMock(return_value=SimpleNamespace(directory=str(tmp_path))), - ) - monkeypatch.setattr( - "flocks.session.runner.Message.create", AsyncMock( - side_effect=[ - SimpleNamespace(id="msg_user"), - SimpleNamespace(id="msg_assistant"), - ] + return_value=SimpleNamespace( + directory=str(tmp_path), + project_id="project_shell_hook", + ) ), ) monkeypatch.setattr( @@ -77,6 +74,73 @@ async def tool_before(self, ctx): assert result["parts"][0]["state"]["output"] == "ok\n" +@pytest.mark.asyncio +async def test_session_shell_persists_completed_tool_and_publishes_events( + monkeypatch: pytest.MonkeyPatch, + tmp_path, +) -> None: + """The HTTP result and reloaded message history share one terminal tool part.""" + + process = SimpleNamespace( + communicate=AsyncMock(return_value=(b"PERSISTED\n", b"")), + returncode=0, + ) + monkeypatch.setattr( + "flocks.session.runner.Session.get_by_id", + AsyncMock( + return_value=SimpleNamespace( + directory=str(tmp_path), + project_id="project_shell_persistence", + ) + ), + ) + monkeypatch.setattr( + "flocks.session.runner.asyncio.create_subprocess_shell", + AsyncMock(return_value=process), + ) + publish_event = AsyncMock() + + result = await SessionRunner.shell( + session_id="ses_shell_persistence", + agent="rex", + command="printf PERSISTED", + event_publish_callback=publish_event, + ) + + assistant_id = result["info"]["id"] + Message.invalidate_cache("ses_shell_persistence") + restored = await Message.get_with_parts_lazy( + "ses_shell_persistence", + assistant_id, + ) + + assert restored is not None + assert restored.info.finish == "stop" + assert restored.info.time["completed"] >= restored.info.time["created"] + tool_parts = [part for part in restored.parts if isinstance(part, ToolPart)] + assert len(tool_parts) == 1 + tool_part = tool_parts[0] + assert tool_part.tool == "bash" + assert tool_part.state.status == "completed" + assert tool_part.state.input == { + "command": "printf PERSISTED", + "workdir": str(tmp_path), + } + assert tool_part.state.output == "PERSISTED\n" + assert tool_part.state.metadata["exitCode"] == 0 + assert result["parts"] == [tool_part.model_dump(mode="json", by_alias=True)] + + event_names = [call.args[0] for call in publish_event.await_args_list] + assert event_names.count("message.updated") == 3 + assert event_names.count("message.part.updated") == 3 + final_part_event = publish_event.await_args_list[-2].args + final_message_event = publish_event.await_args_list[-1].args + assert final_part_event[0] == "message.part.updated" + assert final_part_event[1]["part"]["state"]["status"] == "completed" + assert final_message_event[0] == "message.updated" + assert final_message_event[1]["info"]["finish"] == "stop" + + @pytest.mark.asyncio async def test_tool_execution_payload_falls_back_to_session_owner_subject() -> None: payload = await build_session_tool_execution_payload( From 1ad1de1dc7b0f7895dd428fa8814643920ccb95e Mon Sep 17 00:00:00 2001 From: duguwanglong Date: Wed, 9 Sep 2026 17:27:45 +0800 Subject: [PATCH 60/63] fix(notifications): revise October model usage allowance announcement Update Chinese and English copy to retain 10 million free tokens per day, explain self-service top-ups, and clarify the end of shared pools and temporary allowances on October 1. Preserve the highlighted effective date, correct the top-up link, and update announcement assertions. Validation: all 6 TokenPolicyNotice tests passed. --- webui/src/components/common/TokenPolicyNotice.test.tsx | 7 ++++--- webui/src/locales/en-US/notification.json | 4 ++-- webui/src/locales/zh-CN/notification.json | 4 ++-- 3 files changed, 8 insertions(+), 7 deletions(-) diff --git a/webui/src/components/common/TokenPolicyNotice.test.tsx b/webui/src/components/common/TokenPolicyNotice.test.tsx index 501bfa4f8..6bee07320 100644 --- a/webui/src/components/common/TokenPolicyNotice.test.tsx +++ b/webui/src/components/common/TokenPolicyNotice.test.tsx @@ -46,12 +46,13 @@ it('records presentation before a user immediately closes the card', () => { expect(events).toEqual(['display', 'close']); }); -it('renders the full announcement as a modal with five policy items and safe links', () => { +it('renders the full announcement as a modal with three policy items and safe links', () => { render(); expect(screen.getByRole('dialog')).toHaveAttribute('aria-modal', 'true'); expect(screen.getByRole('region')).toHaveClass('overflow-y-auto'); - expect(screen.getAllByRole('listitem')).toHaveLength(5); - expect(screen.getByText(/不统一折算为1000万Token/)).toBeInTheDocument(); + expect(screen.getAllByRole('listitem')).toHaveLength(3); + expect(screen.getByText(/每人每天仍然免费赠送1000w Token/)).toBeInTheDocument(); + expect(screen.getByText('10月1日起，调整如下：').tagName).toBe('STRONG'); expect(screen.getByText(/2026年9月9日/)).toBeInTheDocument(); const links = screen.getAllByRole('link'); expect(links.map((link) => link.getAttribute('href'))).toEqual([ diff --git a/webui/src/locales/en-US/notification.json b/webui/src/locales/en-US/notification.json index 1176e768d..426ad82b0 100644 --- a/webui/src/locales/en-US/notification.json +++ b/webui/src/locales/en-US/notification.json @@ -1,7 +1,7 @@ { "title": "Flocks Notifications", - "tokenPolicyTitle": "Flocks Community Announcement | October Token Support Policy Changes", - "tokenPolicyBody": "Dear Flocks community,\n\nThank you for using Flocks, sharing feedback, and contributing since its launch on March 31. Alongside this community release, we are announcing the token support policy that will take effect in October so you can plan your usage and budget.\n\nTo support ongoing platform maintenance, reliability improvements, and long-term development of security operations use cases, Flocks is introducing monetary billing and self-service top-ups while retaining daily free support.\n\nThe current free token support policy remains in effect throughout September. Starting in October:\n\n1. **Daily free credit remains available.** Each person receives CNY 10 in free model usage credit per day. Consumption is calculated using the selected model’s pricing. Prices vary by model, so CNY 10 corresponds to different token amounts; it is not uniformly equivalent to 10 million tokens.\n2. **Temporary quota increases will be discontinued.** All additional daily token allowances previously granted to individual users will be removed, including the previous daily allowances of 30–50 million tokens.\n3. **The shared free token pool will be discontinued.** The previous shared free pool will no longer be available.\n4. **Self-service top-ups are available.** For usage beyond your daily free credit, visit [https://portal.agentflocks.com](https://portal.agentflocks.com), sign in with your X community account, and top up. The platform is already live.\n5. **Enterprise requests will be handled by sales.** Enterprise users should contact ThreatBook sales to discuss model usage and related services.\n\nThese changes concern token support and purchasing on the ThreatBook model platform. They will affect future allowances and usage costs for community members who use temporary increases or rely on the shared pool. We are providing advance notice and welcome specific questions in the group.\n\nIf you have scheduled tasks such as continuous inspections or alert analysis, review your models, task frequency, and daily consumption before October, and plan your credit and budget accordingly. If you encounter any issues while using or adjusting the service, post in this group or contact the group owner.\n\nFor enterprise purchasing, please contact ThreatBook sales.\n\nProject: [https://github.com/AgentFlocks/flocks](https://github.com/AgentFlocks/flocks) \nDocumentation: [https://agentflocks.github.io/flocks-docs/](https://agentflocks.github.io/flocks-docs/) \nTop-ups: [https://portal.agentflocks.com](https://portal.agentflocks.com)\n\nThank you for your continued support. We look forward to working with you to keep improving the security operations use cases already running successfully.\n\nThe Flocks Team \nSeptember 9, 2026", + "tokenPolicyTitle": "Flocks Community Announcement | Model Usage Allowances in October", + "tokenPolicyBody": "Dear Flocks community,\n\nThank you for your continued support of Flocks and for sharing your experience and feedback with the community. To support ongoing platform maintenance and the long-term use of more security operations scenarios, we will adjust model usage allowances on October 1, introducing monetary billing and self-service top-ups.\n\nThe current free token policy remains in effect throughout September.\n\n**Starting October 1, the changes are as follows:**\n\n1. **Daily free usage remains available.** Each person will continue to receive 10 million tokens of free model usage per day.\n2. **Self-service top-ups are available when you need more usage.** Sign in to [https://portal.agentflocks.com](https://portal.agentflocks.com) with your X community account to top up for usage beyond your daily free allowance. Enterprise users can also contact ThreatBook sales to discuss their team's future model usage and related services.\n3. **Existing additional support will also change.** Starting October 1, the shared free token pool and temporary allowances previously granted to individual users, including daily allowances above 30 million tokens, will be discontinued. The new allowance rules will apply uniformly.\n\nThese changes concern token support and purchasing on the ThreatBook model platform. We understand that they will affect some community members' existing usage plans. If you run scheduled tasks such as inspections or alert analysis, review model pricing, task frequency, and daily consumption to prepare your allowance and budget in advance. For those using temporary increases or relying on the shared pool, this may change future available allowances and usage costs. We are therefore providing about a month's notice. Please share any specific issues you encounter in the group or privately with the group owner.\n\nFor enterprise purchasing, please contact ThreatBook sales.\n\nThank you for your continued support. We look forward to working with you to keep improving the security operations use cases already running successfully.\n\nProject: [https://github.com/AgentFlocks/flocks](https://github.com/AgentFlocks/flocks) \nDocumentation: [https://agentflocks.github.io/flocks-docs/](https://agentflocks.github.io/flocks-docs/) \nTop-ups: [https://portal.agentflocks.com](https://portal.agentflocks.com)\n\nThe Flocks Team \nSeptember 9, 2026", "subtitle": "Benefits and version highlights are collected here.", "close": "Close notification", "gotIt": "Got it", diff --git a/webui/src/locales/zh-CN/notification.json b/webui/src/locales/zh-CN/notification.json index a630fb634..c4fcbff98 100644 --- a/webui/src/locales/zh-CN/notification.json +++ b/webui/src/locales/zh-CN/notification.json @@ -1,7 +1,7 @@ { "title": "Flocks 通知", - "tokenPolicyTitle": "Flocks社区公告｜10月Token支持策略调整", - "tokenPolicyBody": "各位Flocks社区伙伴：\n\n感谢大家自3月31日发布以来的使用、反馈与共建。配合本次社区版本更新，我们提前同步10月起的Token奖励策略，方便大家安排后续使用和预算。\n\n为支持平台持续维护、稳定性建设和安全运营场景的长期迭代，Flocks在保留每日免费支持的基础上，引入按金额计费和自助充值机制。\n\n**9月内，现行免费Token支持策略持续有效，9月30日旧策略全部失效。10月1日起，调整如下：**\n\n1. **每天仍有免费额度。** 每人每天免费赠送10元模型调用额度，按所选模型的价格计算消耗。不同模型计价不同，10元对应的Token数量也不同，不统一折算为1000万Token。\n2. **统一取消临时加额。** 原先为个别用户单独增加的每日临时Token额度全部取消，包括原每日3000万～5000万额度。\n3. **取消公共免费Token池。** 原公共免费池机制不再保留。\n4. **支持自助充值。** 如需超出每日免费额度的调用，可前往[https://portal.agentflocks.com](https://portal.agentflocks.com)，使用X社区账号登录后自助充值。平台已上线。\n5. **企业需求由销售对接。** 企业用户请联系微步销售，沟通模型用量及相关服务需求。\n\n本次调整涉及微步模型平台的Token支持与购买方式。对使用临时加额或依赖公共池的伙伴，这会改变后续的可用额度和使用成本，我们提前和大家说明，也欢迎大家在群里反馈具体问题。\n\n如果你已配置持续巡检、告警研判等定时任务，建议在10月前梳理所用模型、任务频率和日常消耗，提前安排额度及预算。使用或调整中遇到问题，可在本群或戳群主反馈；\n\n企业采购需求请联系微步销售。\n\n项目地址：[https://github.com/AgentFlocks/flocks](https://github.com/AgentFlocks/flocks) \n使用文档：[https://agentflocks.github.io/flocks-docs/](https://agentflocks.github.io/flocks-docs/) \n充值入口：[https://portal.agentflocks.com](https://portal.agentflocks.com)\n\n感谢大家一直以来的支持。期待和大家一起，把已经跑通的安全运营场景持续用好、不断完善。\n\nFlocks团队 \n2026年9月9日", + "tokenPolicyTitle": "Flocks社区公告｜关于10月的模型使用额度", + "tokenPolicyBody": "各位Flocks社区伙伴：\n\n感谢你一路以来对Flocks的支持，也感谢大家把自己的经验和使用反馈带回社区。为了支持平台持续维护、让更多安全运营场景长期用起来，我们将在10月1日调整模型额度的支持方式，引入按金额计费和自助充值机制。\n\n9月内，现行免费Token策略持续有效。\n\n**10月1日起，调整如下：**\n\n1. **每天仍有免费额度。** 每人每天仍然免费赠送1000w Token模型调用额度。\n2. **需要更多用量时，可以自主充值。** 使用X社区账号登录[https://portal.agentflocks.com](https://portal.agentflocks.com)，即可自助充值，满足超出每日免费额度的使用需求。企业用户也可以联系微步销售，沟通团队后续的模型用量及相关服务需求。\n3. **原有额外支持将同步调整。** 10月1日起，公共免费Token池及此前单独增加的临时额度（包括每日3000万以上Token额度）将停止提供，统一按新的额度规则使用。\n\n本次调整涉及微步模型平台的Token支持与购买方式。我们理解，这会影响部分伙伴现有的使用安排。如果你正在运行巡检、告警研判等定时任务，建议评估了解对应模型的价格、任务频率和日常消耗，提前做好额度和预算准备。对使用临时加额或依赖公共池的伙伴，这可能会改变后续的可用额度和使用成本，所以提前一个月左右和大家说明，也欢迎大家在群里或私戳群主，反馈使用过程中遇到的具体问题。\n\n企业采购需求请联系微步销售。\n\n感谢大家一直以来的支持。期待和大家一起，把已经跑通的安全运营场景持续用好、不断完善。\n\n项目地址：[https://github.com/AgentFlocks/flocks](https://github.com/AgentFlocks/flocks) \n使用文档：[https://agentflocks.github.io/flocks-docs/](https://agentflocks.github.io/flocks-docs/) \n充值入口：[https://portal.agentflocks.com](https://portal.agentflocks.com)\n\nFlocks团队 \n2026年9月9日", "subtitle": "福利提醒和版本变化都在这里。", "close": "关闭通知", "gotIt": "知道了", From 866856283b287ecf2d2217c6827293c1434caf27 Mon Sep 17 00:00:00 2001 From: stephamie7 <1223696150@qq.com> Date: Wed, 9 Sep 2026 17:33:56 +0800 Subject: [PATCH 61/63] chore/update-version-2026.9.9 --- pyproject.toml | 2 +- uv.lock | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 92400a074..bc211e7c0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "flocks" -version = "v2026.8.17" +version = "v2026.9.9" description = "AI-Native SecOps platform with multi-agent collaboration" authors = [ {name = "Flocks Team", email = "team@example.com"} diff --git a/uv.lock b/uv.lock index 4f6e68d7a..5a4498bf1 100644 --- a/uv.lock +++ b/uv.lock @@ -553,7 +553,7 @@ wheels = [ [[package]] name = "flocks" -version = "2026.8.17" +version = "2026.9.9" source = { editable = "." } dependencies = [ { name = "aiofiles" }, From 072ee3aa4bde619d65b69f4fac80bec787a4e4a9 Mon Sep 17 00:00:00 2001 From: duguwanglong Date: Wed, 9 Sep 2026 17:40:24 +0800 Subject: [PATCH 62/63] fix(notifications): simplify token policy announcement copy Remove the sentence describing the scope of ThreatBook model platform token support and purchasing from both Chinese and English announcements. Validation: all 6 TokenPolicyNotice tests passed. --- webui/src/locales/en-US/notification.json | 2 +- webui/src/locales/zh-CN/notification.json | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/webui/src/locales/en-US/notification.json b/webui/src/locales/en-US/notification.json index 426ad82b0..3d2612777 100644 --- a/webui/src/locales/en-US/notification.json +++ b/webui/src/locales/en-US/notification.json @@ -1,7 +1,7 @@ { "title": "Flocks Notifications", "tokenPolicyTitle": "Flocks Community Announcement | Model Usage Allowances in October", - "tokenPolicyBody": "Dear Flocks community,\n\nThank you for your continued support of Flocks and for sharing your experience and feedback with the community. To support ongoing platform maintenance and the long-term use of more security operations scenarios, we will adjust model usage allowances on October 1, introducing monetary billing and self-service top-ups.\n\nThe current free token policy remains in effect throughout September.\n\n**Starting October 1, the changes are as follows:**\n\n1. **Daily free usage remains available.** Each person will continue to receive 10 million tokens of free model usage per day.\n2. **Self-service top-ups are available when you need more usage.** Sign in to [https://portal.agentflocks.com](https://portal.agentflocks.com) with your X community account to top up for usage beyond your daily free allowance. Enterprise users can also contact ThreatBook sales to discuss their team's future model usage and related services.\n3. **Existing additional support will also change.** Starting October 1, the shared free token pool and temporary allowances previously granted to individual users, including daily allowances above 30 million tokens, will be discontinued. The new allowance rules will apply uniformly.\n\nThese changes concern token support and purchasing on the ThreatBook model platform. We understand that they will affect some community members' existing usage plans. If you run scheduled tasks such as inspections or alert analysis, review model pricing, task frequency, and daily consumption to prepare your allowance and budget in advance. For those using temporary increases or relying on the shared pool, this may change future available allowances and usage costs. We are therefore providing about a month's notice. Please share any specific issues you encounter in the group or privately with the group owner.\n\nFor enterprise purchasing, please contact ThreatBook sales.\n\nThank you for your continued support. We look forward to working with you to keep improving the security operations use cases already running successfully.\n\nProject: [https://github.com/AgentFlocks/flocks](https://github.com/AgentFlocks/flocks) \nDocumentation: [https://agentflocks.github.io/flocks-docs/](https://agentflocks.github.io/flocks-docs/) \nTop-ups: [https://portal.agentflocks.com](https://portal.agentflocks.com)\n\nThe Flocks Team \nSeptember 9, 2026", + "tokenPolicyBody": "Dear Flocks community,\n\nThank you for your continued support of Flocks and for sharing your experience and feedback with the community. To support ongoing platform maintenance and the long-term use of more security operations scenarios, we will adjust model usage allowances on October 1, introducing monetary billing and self-service top-ups.\n\nThe current free token policy remains in effect throughout September.\n\n**Starting October 1, the changes are as follows:**\n\n1. **Daily free usage remains available.** Each person will continue to receive 10 million tokens of free model usage per day.\n2. **Self-service top-ups are available when you need more usage.** Sign in to [https://portal.agentflocks.com](https://portal.agentflocks.com) with your X community account to top up for usage beyond your daily free allowance. Enterprise users can also contact ThreatBook sales to discuss their team's future model usage and related services.\n3. **Existing additional support will also change.** Starting October 1, the shared free token pool and temporary allowances previously granted to individual users, including daily allowances above 30 million tokens, will be discontinued. The new allowance rules will apply uniformly.\n\nWe understand that they will affect some community members' existing usage plans. If you run scheduled tasks such as inspections or alert analysis, review model pricing, task frequency, and daily consumption to prepare your allowance and budget in advance. For those using temporary increases or relying on the shared pool, this may change future available allowances and usage costs. We are therefore providing about a month's notice. Please share any specific issues you encounter in the group or privately with the group owner.\n\nFor enterprise purchasing, please contact ThreatBook sales.\n\nThank you for your continued support. We look forward to working with you to keep improving the security operations use cases already running successfully.\n\nProject: [https://github.com/AgentFlocks/flocks](https://github.com/AgentFlocks/flocks) \nDocumentation: [https://agentflocks.github.io/flocks-docs/](https://agentflocks.github.io/flocks-docs/) \nTop-ups: [https://portal.agentflocks.com](https://portal.agentflocks.com)\n\nThe Flocks Team \nSeptember 9, 2026", "subtitle": "Benefits and version highlights are collected here.", "close": "Close notification", "gotIt": "Got it", diff --git a/webui/src/locales/zh-CN/notification.json b/webui/src/locales/zh-CN/notification.json index c4fcbff98..f1656883b 100644 --- a/webui/src/locales/zh-CN/notification.json +++ b/webui/src/locales/zh-CN/notification.json @@ -1,7 +1,7 @@ { "title": "Flocks 通知", "tokenPolicyTitle": "Flocks社区公告｜关于10月的模型使用额度", - "tokenPolicyBody": "各位Flocks社区伙伴：\n\n感谢你一路以来对Flocks的支持，也感谢大家把自己的经验和使用反馈带回社区。为了支持平台持续维护、让更多安全运营场景长期用起来，我们将在10月1日调整模型额度的支持方式，引入按金额计费和自助充值机制。\n\n9月内，现行免费Token策略持续有效。\n\n**10月1日起，调整如下：**\n\n1. **每天仍有免费额度。** 每人每天仍然免费赠送1000w Token模型调用额度。\n2. **需要更多用量时，可以自主充值。** 使用X社区账号登录[https://portal.agentflocks.com](https://portal.agentflocks.com)，即可自助充值，满足超出每日免费额度的使用需求。企业用户也可以联系微步销售，沟通团队后续的模型用量及相关服务需求。\n3. **原有额外支持将同步调整。** 10月1日起，公共免费Token池及此前单独增加的临时额度（包括每日3000万以上Token额度）将停止提供，统一按新的额度规则使用。\n\n本次调整涉及微步模型平台的Token支持与购买方式。我们理解，这会影响部分伙伴现有的使用安排。如果你正在运行巡检、告警研判等定时任务，建议评估了解对应模型的价格、任务频率和日常消耗，提前做好额度和预算准备。对使用临时加额或依赖公共池的伙伴，这可能会改变后续的可用额度和使用成本，所以提前一个月左右和大家说明，也欢迎大家在群里或私戳群主，反馈使用过程中遇到的具体问题。\n\n企业采购需求请联系微步销售。\n\n感谢大家一直以来的支持。期待和大家一起，把已经跑通的安全运营场景持续用好、不断完善。\n\n项目地址：[https://github.com/AgentFlocks/flocks](https://github.com/AgentFlocks/flocks) \n使用文档：[https://agentflocks.github.io/flocks-docs/](https://agentflocks.github.io/flocks-docs/) \n充值入口：[https://portal.agentflocks.com](https://portal.agentflocks.com)\n\nFlocks团队 \n2026年9月9日", + "tokenPolicyBody": "各位Flocks社区伙伴：\n\n感谢你一路以来对Flocks的支持，也感谢大家把自己的经验和使用反馈带回社区。为了支持平台持续维护、让更多安全运营场景长期用起来，我们将在10月1日调整模型额度的支持方式，引入按金额计费和自助充值机制。\n\n9月内，现行免费Token策略持续有效。\n\n**10月1日起，调整如下：**\n\n1. **每天仍有免费额度。** 每人每天仍然免费赠送1000w Token模型调用额度。\n2. **需要更多用量时，可以自主充值。** 使用X社区账号登录[https://portal.agentflocks.com](https://portal.agentflocks.com)，即可自助充值，满足超出每日免费额度的使用需求。企业用户也可以联系微步销售，沟通团队后续的模型用量及相关服务需求。\n3. **原有额外支持将同步调整。** 10月1日起，公共免费Token池及此前单独增加的临时额度（包括每日3000万以上Token额度）将停止提供，统一按新的额度规则使用。\n\n我们理解，这会影响部分伙伴现有的使用安排。如果你正在运行巡检、告警研判等定时任务，建议评估了解对应模型的价格、任务频率和日常消耗，提前做好额度和预算准备。对使用临时加额或依赖公共池的伙伴，这可能会改变后续的可用额度和使用成本，所以提前一个月左右和大家说明，也欢迎大家在群里或私戳群主，反馈使用过程中遇到的具体问题。\n\n企业采购需求请联系微步销售。\n\n感谢大家一直以来的支持。期待和大家一起，把已经跑通的安全运营场景持续用好、不断完善。\n\n项目地址：[https://github.com/AgentFlocks/flocks](https://github.com/AgentFlocks/flocks) \n使用文档：[https://agentflocks.github.io/flocks-docs/](https://agentflocks.github.io/flocks-docs/) \n充值入口：[https://portal.agentflocks.com](https://portal.agentflocks.com)\n\nFlocks团队 \n2026年9月9日", "subtitle": "福利提醒和版本变化都在这里。", "close": "关闭通知", "gotIt": "知道了", From 023e5a6b2253e938c29d0029530d8246b4e16867 Mon Sep 17 00:00:00 2001 From: stephamie7 <1223696150@qq.com> Date: Wed, 9 Sep 2026 17:52:20 +0800 Subject: [PATCH 63/63] fix(webui): name update modal mock for hooks lint --- webui/src/components/layout/Layout.test.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/webui/src/components/layout/Layout.test.tsx b/webui/src/components/layout/Layout.test.tsx index 2bc280920..881f250c6 100644 --- a/webui/src/components/layout/Layout.test.tsx +++ b/webui/src/components/layout/Layout.test.tsx @@ -163,7 +163,7 @@ vi.mock('@/components/common/LanguageSwitcher', () => ({ vi.mock('@/components/common/UpdateModal', () => ({ UPDATE_DISMISSED_KEY: 'update-dismissed', - default: (props: Record) => { + default: function MockUpdateModal(props: Record) { React.useEffect(() => { (props.onPresented as (() => void) | undefined)?.(); }, [props.onPresented]); updateModalMock(props); return
;