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/acp/agent.py b/flocks/acp/agent.py index 3025fa1e9..b5ec12b32 100644 --- a/flocks/acp/agent.py +++ b/flocks/acp/agent.py @@ -94,6 +94,31 @@ def to_locations(tool_name: str, input_data: Dict[str, Any]) -> List[ToolCallLoc return [] +def edit_diff_text( + input_data: Dict[str, Any], + metadata: Any, +) -> tuple[str, str]: + """Extract edit text for ACP diff rendering from result metadata or edits[].""" + if isinstance(metadata, dict): + filediff = metadata.get("filediff") + if isinstance(filediff, dict): + before = filediff.get("before") + after = filediff.get("after") + if isinstance(before, str) and isinstance(after, str): + return before, after + + edits = input_data.get("edits") + if isinstance(edits, list) and len(edits) == 1 and isinstance(edits[0], dict): + old_text = edits[0].get("oldString", "") + new_text = edits[0].get("newString", "") + return ( + old_text if isinstance(old_text, str) else "", + new_text if isinstance(new_text, str) else "", + ) + + return "", "" + + def parse_uri(uri: str) -> Dict[str, Any]: """ Parse URI into file or text content @@ -503,8 +528,10 @@ async def _handle_tool_part_update(self, session_id: str, part: Dict[str, Any]) # Add diff content for edit tools if kind == "edit": file_path = input_data.get("filePath", "") - old_text = input_data.get("oldString", "") - new_text = input_data.get("newString", input_data.get("content", "")) + old_text, new_text = edit_diff_text( + input_data, + state.get("metadata"), + ) content.append({ "type": "diff", "path": file_path, @@ -854,8 +881,10 @@ async def _process_message(self, message: Dict[str, Any]) -> None: if kind == "edit": file_path = input_data.get("filePath", "") - old_text = input_data.get("oldString", "") - new_text = input_data.get("newString", input_data.get("content", "")) + old_text, new_text = edit_diff_text( + input_data, + state.get("metadata"), + ) content.append({ "type": "diff", "path": file_path, diff --git a/flocks/agent/agents/hephaestus/prompt_builder.py b/flocks/agent/agents/hephaestus/prompt_builder.py index 9617affcb..0d9f7f842 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, ) @@ -39,6 +38,8 @@ def build_hephaestus_prompt( 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, @@ -62,7 +63,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 +246,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..98e183526 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, ) @@ -49,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) @@ -58,12 +58,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 +139,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 +273,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/flocks/agent/registry.py b/flocks/agent/registry.py index ee1b5ca32..75bcf50ef 100644 --- a/flocks/agent/registry.py +++ b/flocks/agent/registry.py @@ -40,6 +40,7 @@ ) import flocks.agent.delegatable_settings as delegatable_settings from flocks.agent.toolset import agent_declares_tool +from flocks.agent.tool_permissions import permission_items_to_ruleset from flocks.agent.prompt_utils import categorize_tools from flocks.agent.agent_factory import ( scan_and_load, @@ -172,6 +173,7 @@ def _storage_custom_agent_to_info(agent_data: Dict[str, Any]) -> Optional[AgentI native=False, hidden=agent_data.get("hidden", False), delegatable=agent_data.get("delegatable"), + permission=permission_items_to_ruleset(agent_data.get("permission")), tools=agent_data.get("tools", []), tags=agent_data.get("tags", []), ) diff --git a/flocks/agent/tool_permissions.py b/flocks/agent/tool_permissions.py new file mode 100644 index 000000000..83fec2c19 --- /dev/null +++ b/flocks/agent/tool_permissions.py @@ -0,0 +1,86 @@ +"""Helpers for keeping agent tool selections and permission rules aligned.""" + +from __future__ import annotations + +from typing import Any, Dict, List + +from flocks.permission.rule import PermissionLevel, PermissionRule, PermissionScope + +QUESTION_TOOL_NAME = "question" +TOOLS_MANAGED_PERMISSION_SOURCE = "agent_tools" + + +def normalize_permission_items(value: Any) -> List[Dict[str, Any]]: + """Return stored permission rules in API/storage dict form.""" + if not isinstance(value, list): + return [] + + normalized: List[Dict[str, Any]] = [] + for item in value: + if not isinstance(item, dict): + continue + + permission = item.get("permission") + action = item.get("action") or item.get("level") + action = getattr(action, "value", action) + pattern = item.get("pattern") or "*" + if not permission or action not in {"allow", "ask", "deny"}: + continue + + rule = dict(item) + rule["permission"] = str(permission) + rule["action"] = str(action) + rule["pattern"] = str(pattern) + normalized.append(rule) + + return normalized + + +def sync_question_permission_with_tools( + permissions: Any, + tools: List[str], +) -> List[Dict[str, Any]]: + """Make the Agent tools checkbox an effective question allow/deny toggle.""" + normalized = normalize_permission_items(permissions) + without_managed_question = [ + rule + for rule in normalized + if not ( + rule.get("permission") == QUESTION_TOOL_NAME + and rule.get("pattern", "*") == "*" + and rule.get("source") == TOOLS_MANAGED_PERMISSION_SOURCE + ) + ] + + if QUESTION_TOOL_NAME in set(tools): + return without_managed_question + + if any( + rule.get("permission") == QUESTION_TOOL_NAME + and rule.get("action") == "deny" + and rule.get("pattern", "*") == "*" + for rule in without_managed_question + ): + return without_managed_question + + return without_managed_question + [ + { + "permission": QUESTION_TOOL_NAME, + "action": "deny", + "pattern": "*", + "source": TOOLS_MANAGED_PERMISSION_SOURCE, + } + ] + + +def permission_items_to_ruleset(value: Any) -> List[PermissionRule]: + """Convert stored permission dicts to the internal PermissionRule shape.""" + return [ + PermissionRule( + permission=item["permission"], + level=PermissionLevel(item["action"]), + scope=PermissionScope.PATTERN, + pattern=item.get("pattern") or "*", + ) + for item in normalize_permission_items(value) + ] 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..c4c01a13c 100644 --- a/flocks/config/config.py +++ b/flocks/config/config.py @@ -29,22 +29,41 @@ 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", +} 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: - return PermissionAction.DENY - return existing if existing is not None else incoming + 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(): + if pattern in merged: + merged[pattern] = _merge_permission_action( + merged[pattern], action, + ) + else: + merged[pattern] = action + return merged + + 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: @@ -58,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 @@ -79,11 +123,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 @@ -121,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): @@ -158,7 +209,6 @@ class CommandConfig(BaseModel): description: Optional[str] = None agent: Optional[str] = None model: Optional[str] = None - subtask: Optional[bool] = None # ==================== Provider Configuration ==================== @@ -750,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 @@ -1116,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: @@ -1454,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/ingest/kafka/manager.py b/flocks/ingest/kafka/manager.py index 1b752be8a..5f7846431 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", @@ -123,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. @@ -446,10 +461,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 +707,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: @@ -748,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, @@ -767,15 +787,10 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: input_params=summarized_inputs, ) 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", step_compactor=lambda step: _compact_step_for_kafka_storage( step, input_key=input_key, @@ -784,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, @@ -845,9 +861,16 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: } ) finally: - await cleanup_workflow_tool_context(tool_context) + steps = step_recorder.take_steps() + if tool_context is not None: + 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..63e1d1b18 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: @@ -620,13 +621,7 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: input_params=summarized_inputs, ) 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", - ) + step_recorder = ExecutionStepRecorder() start_time = time.time() trigger_meta = mapped_inputs.get("_flocks", {}).get("trigger", {}) tool_context = None @@ -692,9 +687,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/permission/helpers.py b/flocks/permission/helpers.py index 60206ebd7..104e8562c 100644 --- a/flocks/permission/helpers.py +++ b/flocks/permission/helpers.py @@ -6,6 +6,55 @@ Ruleset = List[PermissionRule] +_LEGACY_PERMISSION_NAMES = { + "task": "delegate_task", + "todowrite": "todo", + "todoread": "todo", +} + + +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(): + if pattern in merged: + merged[pattern] = _merge_legacy_permission( + merged[pattern], action, + ) + else: + merged[pattern] = action + return merged + + 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 _canonicalize_permission_config(config: Dict[str, Any]) -> 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) + else: + canonical[name] = value + return canonical + def from_config(permission_config: Union[Dict[str, Any], BaseModel]) -> Ruleset: """ @@ -22,6 +71,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/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/server/routes/agent.py b/flocks/server/routes/agent.py index 5df1fa9d4..b537da16c 100644 --- a/flocks/server/routes/agent.py +++ b/flocks/server/routes/agent.py @@ -29,6 +29,11 @@ import flocks.agent.delegatable_settings as delegatable_settings from flocks.agent.registry import Agent +from flocks.agent.tool_permissions import ( + normalize_permission_items, + permission_items_to_ruleset, + sync_question_permission_with_tools, +) from flocks.agent.agent import AgentInfo as AgentInfoModel, AgentModel as AgentModelConfig from flocks.agent.agent_factory import find_yaml_agent, read_yaml_agent, update_yaml_agent, delete_yaml_agent from flocks.utils.log import Log @@ -92,6 +97,7 @@ def agent_to_response( delegatable_override: Optional[bool] = None, skills: Optional[List[str]] = None, tools: Optional[List[str]] = None, + permission: Optional[List[Dict[str, Any]]] = None, ) -> AgentResponse: """Convert internal AgentInfo to API response format.""" delegatable = ( @@ -124,7 +130,7 @@ def agent_to_response( topP=agent.top_p, temperature=temperature_override if temperature_override is not None else agent.temperature, color=agent.color, - permission=[], + permission=permission or [], model=model_info, prompt=agent.prompt, options=agent.options, @@ -155,6 +161,8 @@ def _agent_data_to_info(agent_data: Dict[str, Any]) -> AgentInfoModel: temperature=agent_data.get("temperature"), color=agent_data.get("color"), mode=mode, + permission=permission_items_to_ruleset(agent_data.get("permission")), + tools=agent_data.get("tools", []), model=AgentModelConfig( model_id=model_data["modelID"], provider_id=model_data["providerID"], @@ -185,7 +193,7 @@ def _custom_agent_data_to_response(agent_data: Dict[str, Any]) -> AgentResponse: model=model_info, native=agent_data.get("native", False), hidden=agent_data.get("hidden", False), - permission=[], + permission=normalize_permission_items(agent_data.get("permission")), options={}, delegatable=delegatable, skills=agent_data.get("skills", []), @@ -208,8 +216,8 @@ def _load_delegatable_overrides() -> Dict[str, bool]: return delegatable_settings.load_overrides() -async def _load_custom_agent_extras(name: str) -> tuple[List[str], List[str]]: - """Load skills/tools list for an agent from storage. +async def _load_custom_agent_extras(name: str) -> tuple[List[str], List[str], List[Dict[str, Any]]]: + """Load skills/tools/permission overlay for an agent from storage. Works for both full Storage-based custom agents and YAML agents with a skills/tools overlay (written by the YAML update path). @@ -218,10 +226,14 @@ async def _load_custom_agent_extras(name: str) -> tuple[List[str], List[str]]: try: data = await Storage.read(f"agent/custom/{name}") if not isinstance(data, dict): - return [], [] - return data.get("skills", []), data.get("tools", []) + return [], [], [] + return ( + data.get("skills", []), + data.get("tools", []), + normalize_permission_items(data.get("permission")), + ) except Exception: - return [], [] + return [], [], [] def _get_all_tool_names() -> List[str]: @@ -257,8 +269,9 @@ async def _build_single_agent_response( if agent.native: tools = _compute_native_agent_tools(agent, all_tool_names) skills: List[str] = [] + permission: List[Dict[str, Any]] = [] else: - skills, tools = await _load_custom_agent_extras(agent.name) + skills, tools, permission = await _load_custom_agent_extras(agent.name) override = overrides.get(agent.name, {}) model_override = {k: override[k] for k in ("modelID", "providerID") if k in override} or None temperature_override = override.get("temperature") @@ -269,6 +282,7 @@ async def _build_single_agent_response( delegatable_override=delegatable_overrides.get(agent.name), skills=skills, tools=tools, + permission=permission, ) @@ -359,6 +373,7 @@ class AgentCreateRequest(BaseModel): mode: str = Field("primary", description="Agent mode") model: Optional[AgentModelInfo] = Field(None, description="Preferred model") delegatable: Optional[bool] = Field(None, description="Whether this agent can be delegated to") + permission: List[Dict[str, Any]] = Field(default_factory=list, description="Permission rules") skills: List[str] = Field(default_factory=list, description="Enabled skill names") tools: List[str] = Field(default_factory=list, description="Enabled tool names") @@ -373,6 +388,7 @@ class AgentUpdateRequest(BaseModel): color: Optional[str] = Field(None, description="Color") model: Optional[AgentModelInfo] = Field(None, description="Preferred model") delegatable: Optional[bool] = Field(None, description="Whether this agent can be delegated to") + permission: Optional[List[Dict[str, Any]]] = Field(None, description="Permission rules") skills: Optional[List[str]] = Field(None, description="Enabled skill names") tools: Optional[List[str]] = Field(None, description="Enabled tool names") @@ -402,6 +418,11 @@ async def create_agent(req: AgentCreateRequest): if existing: raise HTTPException(status_code=409, detail=f"Agent {req.name} already exists") + permission = ( + sync_question_permission_with_tools(req.permission, req.tools) + if "tools" in req.model_fields_set + else normalize_permission_items(req.permission) + ) agent_data: Dict[str, Any] = { "name": req.name, "name_cn": req.nameCn, @@ -415,6 +436,7 @@ async def create_agent(req: AgentCreateRequest): "delegatable": req.delegatable if req.delegatable is not None else req.mode != "primary", "native": False, "hidden": False, + "permission": permission, "skills": req.skills, "tools": req.tools, } @@ -468,10 +490,16 @@ async def update_agent(name: str, req: AgentUpdateRequest): if req.delegatable is not None: agent_data["delegatable"] = req.delegatable delegatable_settings.forget_override(name) + if req.permission is not None: + agent_data["permission"] = normalize_permission_items(req.permission) if req.skills is not None: agent_data["skills"] = req.skills if req.tools is not None: agent_data["tools"] = req.tools + agent_data["permission"] = sync_question_permission_with_tools( + agent_data.get("permission"), + req.tools, + ) await Storage.write(agent_key, agent_data) @@ -508,12 +536,18 @@ async def update_agent(name: str, req: AgentUpdateRequest): # Persist skills/tools overlay for YAML agents in Storage. # The entry intentionally omits "name" so it is not mistaken # for a full Storage-based custom agent on subsequent updates. - if req.skills is not None or req.tools is not None: + if req.skills is not None or req.tools is not None or req.permission is not None: extras: Dict[str, Any] = agent_data if isinstance(agent_data, dict) else {} + if req.permission is not None: + extras["permission"] = normalize_permission_items(req.permission) if req.skills is not None: extras["skills"] = req.skills if req.tools is not None: extras["tools"] = req.tools + extras["permission"] = sync_question_permission_with_tools( + extras.get("permission"), + req.tools, + ) await Storage.write(agent_key, extras) # Sync: apply updates to the in-memory AgentInfo cache @@ -538,6 +572,11 @@ async def update_agent(name: str, req: AgentUpdateRequest): ) if req.delegatable is not None: agent.delegatable = req.delegatable + if req.tools is not None: + agent.tools = req.tools + agent.permission = permission_items_to_ruleset(extras.get("permission")) + elif req.permission is not None: + agent.permission = permission_items_to_ruleset(extras.get("permission")) overrides = await _load_model_overrides() delegatable_overrides = _load_delegatable_overrides() all_tool_names = await _get_all_tool_names_async() 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/server/routes/workflow.py b/flocks/server/routes/workflow.py index b838c543a..dfda6b13c 100644 --- a/flocks/server/routes/workflow.py +++ b/flocks/server/routes/workflow.py @@ -55,13 +55,12 @@ compact_step_for_storage, create_execution_record, derive_loop_progress, + ExecutionProgressWriter, + ExecutionStepRecorder, 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, - workflow_execution_step_prefix as _workflow_execution_step_prefix, ) from flocks.workflow.io import load_workflow, dump_workflow from flocks.workflow.store import WorkflowStore @@ -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 @@ -146,6 +144,7 @@ class ActiveWorkflowExecution: workflow_id: str task: asyncio.Task[Any] cancel_event: threading.Event + progress_writer: ExecutionProgressWriter _active_workflow_executions: Dict[str, ActiveWorkflowExecution] = {} @@ -1122,12 +1121,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 - loop = asyncio.get_running_loop() + step_recorder = ExecutionStepRecorder() pending_step_index: Optional[int] = None pending_step: Optional[Dict[str, Any]] = None execution_summary: Dict[str, Any] = { @@ -1150,21 +1149,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,161 +1169,123 @@ def _on_step_start(_run_id, step_index, node, _inputs): "error": "Run cancelled before node completed", } ) - _write_progress( - { - "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: try: - asyncio.run_coroutine_threadsafe( - record_execution_step(exec_id, step_count, step_dict), - loop, - ).result(timeout=5) + 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: - log.warning( - "workflow.execution_step.write_failed", + execution_error = exc + + duration = time.time() - start_time + current_data = dict(execution_summary) + 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]) + 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( { - "exec_id": exec_id, - "step_index": step_count, - "error": str(exc), - }, + "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), + } ) - if step_count % _PROGRESS_FLUSH_EVERY_STEPS == 0: - _write_progress( + else: + current_data.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, + "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), } ) - 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", + 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, - "step_index": pending_step_index, - "error": str(exc), + "status": current_data["status"], + "duration": duration, + }, + ) + else: + log.error( + "workflow.execute.error", + { + "id": workflow_id, + "exec_id": exec_id, + "error": str(execution_error), }, ) - - 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, - ) - - 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) - if status_value == "cancelled" and not final_history: - await _flush_pending_step() - final_steps = result.steps - if pending_step_index is not None: - 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), - } - ) - - 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) - 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, - "currentPhase": "cancelled" if cancel_event.is_set() else "error", - "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), - }, - ) finally: _active_workflow_executions.pop(exec_id, None) @@ -1521,19 +1467,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 +1572,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: @@ -1761,6 +1695,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( @@ -1770,6 +1705,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}", @@ -1778,6 +1714,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 @@ -1826,13 +1763,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", { @@ -2455,9 +2392,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 +2477,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/server/routes/workspace.py b/flocks/server/routes/workspace.py index 197f03242..0775a641a 100644 --- a/flocks/server/routes/workspace.py +++ b/flocks/server/routes/workspace.py @@ -1,7 +1,7 @@ """ Workspace routes -Generic file-manager API for ~/.flocks/workspace/ plus a read-only view +Generic file-manager API for ~/.flocks/workspace/ plus a managed view of the agent memory directory (~/.flocks/data/memory/). All `path` query/body parameters are **relative** to the respective root. @@ -26,9 +26,10 @@ POST /api/workspace/move move / rename POST /api/workspace/reveal open containing folder in system file manager -Memory view (read-only, points to data/memory/) +Memory view (points to data/memory/) GET /api/workspace/memory/list list memory files GET /api/workspace/memory/file read memory file content + PUT /api/workspace/memory/file update editable memory files GET /api/workspace/memory/preview preview single memory file inline GET /api/workspace/memory/download download single memory file @@ -50,10 +51,12 @@ from pathlib import Path from typing import List, Optional, Literal -from fastapi import APIRouter, HTTPException, Query, UploadFile, File +from fastapi import APIRouter, HTTPException, Query, UploadFile, File, Request, status from fastapi.responses import FileResponse, StreamingResponse from pydantic import BaseModel +from flocks.project.project import Project +from flocks.server.auth import get_optional_user, require_user from flocks.workspace.manager import WorkspaceManager from flocks.workspace.models import WorkspaceNode, WorkspaceStats from flocks.utils.log import Log @@ -98,6 +101,11 @@ def _get_manager() -> WorkspaceManager: return mgr +def _workspace_root(mgr: WorkspaceManager) -> Path: + """Return the canonical workspace root used for relative path rendering.""" + return mgr.get_workspace_dir().resolve() + + def _is_allowed_upload_filename(filename: str) -> bool: return Path(filename).suffix.lower() in _ALLOWED_UPLOAD_EXTENSIONS @@ -247,15 +255,16 @@ async def list_tree( depth: int = Query(2, ge=1, le=5, description="Tree depth"), ): mgr = _get_manager() + workspace_root = _workspace_root(mgr) try: - base = mgr.resolve_workspace_path(path) if path else mgr.get_workspace_dir() + base = mgr.resolve_workspace_path(path) if path else workspace_root except ValueError as e: raise HTTPException(status_code=400, detail=str(e)) if not base.exists(): raise HTTPException(status_code=404, detail=f"Path not found: {path}") if not base.is_dir(): raise HTTPException(status_code=400, detail=f"Not a directory: {path}") - return await asyncio.to_thread(_build_tree_sync, base, mgr.get_workspace_dir(), depth) + return await asyncio.to_thread(_build_tree_sync, base, workspace_root, depth) @router.get("/list", response_model=List[WorkspaceNode], summary="List directory") @@ -263,15 +272,16 @@ async def list_dir( path: str = Query("", description="Relative path from workspace root"), ): mgr = _get_manager() + workspace_root = _workspace_root(mgr) try: - base = mgr.resolve_workspace_path(path) if path else mgr.get_workspace_dir() + base = mgr.resolve_workspace_path(path) if path else workspace_root except ValueError as e: raise HTTPException(status_code=400, detail=str(e)) if not base.exists(): raise HTTPException(status_code=404, detail=f"Path not found: {path}") if not base.is_dir(): raise HTTPException(status_code=400, detail=f"Not a directory: {path}") - return await asyncio.to_thread(_list_dir_sync, base, mgr.get_workspace_dir()) + return await asyncio.to_thread(_list_dir_sync, base, workspace_root) class DirCreateRequest(BaseModel): @@ -585,9 +595,39 @@ def _memory_child_sort_key(path: Path) -> tuple[int, str]: return (0 if path.is_dir() else 1, path.name.casefold()) -def _build_memory_node_sync(path: Path, memory_dir: Path) -> WorkspaceNode: +def _memory_path_parts(rel_path: str) -> list[str]: + return [part for part in rel_path.replace("\\", "/").split("/") if part] + + +def _is_editable_memory_path( + rel_path: str, + writable_project_ids: Optional[set[str]] = None, +) -> bool: + """Return whether the current user may edit a memory file path.""" + path_parts = _memory_path_parts(rel_path) + if len(path_parts) == 1: + return path_parts[0] in {"USER.md", "MEMORY.md"} + if len(path_parts) == 2: + return path_parts[0] == "daily" and path_parts[1].endswith(".md") + if len(path_parts) == 3: + return ( + path_parts[0] == "projects" + and path_parts[2] == "MEMORY.md" + and writable_project_ids is not None + and path_parts[1] in writable_project_ids + ) + return False + + +def _build_memory_node_sync( + path: Path, + memory_dir: Path, + writable_project_ids: Optional[set[str]], +) -> WorkspaceNode: """Build one recursive node for the Memory tree.""" node = _node_from_path(path, memory_dir) + if node.type == "file": + node.editable = _is_editable_memory_path(node.path, writable_project_ids) if node.type == "directory": children = ( child @@ -595,13 +635,16 @@ def _build_memory_node_sync(path: Path, memory_dir: Path) -> WorkspaceNode: if not child.is_symlink() ) node.children = [ - _build_memory_node_sync(child, memory_dir) + _build_memory_node_sync(child, memory_dir, writable_project_ids) for child in sorted(children, key=_memory_child_sort_key) ] return node -def _list_memory_sync(memory_dir: Path) -> List[WorkspaceNode]: +def _list_memory_sync( + memory_dir: Path, + writable_project_ids: Optional[set[str]], +) -> List[WorkspaceNode]: """Build the USER/global/daily/project Memory hierarchy.""" children = ( child @@ -615,16 +658,21 @@ def _list_memory_sync(memory_dir: Path) -> List[WorkspaceNode]: child.name.casefold(), ), ) - return [_build_memory_node_sync(child, memory_dir) for child in ordered] + return [ + _build_memory_node_sync(child, memory_dir, writable_project_ids) + for child in ordered + ] @router.get("/memory/list", response_model=List[WorkspaceNode], summary="List memory files") -async def list_memory(): +async def list_memory(request: Request): mgr = _get_manager() memory_dir = mgr.get_memory_dir() if not memory_dir.exists(): return [] - return await asyncio.to_thread(_list_memory_sync, memory_dir) + user = get_optional_user(request) + writable_project_ids = Project.registered_project_ids(user.id) if user else set() + return await asyncio.to_thread(_list_memory_sync, memory_dir, writable_project_ids) @router.get("/memory/file", summary="Read memory file content") @@ -664,12 +712,19 @@ async def read_memory_file( @router.put("/memory/file", summary="Write memory file content") -async def write_memory_file(body: FileWriteRequest): +async def write_memory_file(request: Request, body: FileWriteRequest): + user = require_user(request) mgr = _get_manager() try: target = mgr.resolve_memory_path(body.path) except ValueError as e: raise HTTPException(status_code=400, detail=str(e)) + writable_project_ids = Project.registered_project_ids(user.id) + if not _is_editable_memory_path(body.path, writable_project_ids): + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="Memory file is read-only", + ) target.parent.mkdir(parents=True, exist_ok=True) try: target.write_text(body.content, encoding="utf-8") 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..f1b716e05 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,17 +309,41 @@ 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.project import Project + + try: + config = await Config.get() + config_instructions = tuple(config.instructions or ()) + 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, - tool_revision=ToolRegistry.revision(), + turn_context=TurnPromptContext( + worktree=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, @@ -441,8 +465,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 +523,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/lifecycle/retry.py b/flocks/session/lifecycle/retry.py index ae7bf77f3..fdf048145 100644 --- a/flocks/session/lifecycle/retry.py +++ b/flocks/session/lifecycle/retry.py @@ -135,6 +135,9 @@ def retryable(error: Dict[str, Any]) -> Optional[str]: """ error_name = error.get("name", "") error_data = error.get("data", {}) + + if error_name == "StreamToolArgumentsTruncatedError": + return "Model output was truncated while generating tool arguments" # Check if it's an APIError with isRetryable flag if error_name == "APIError": 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.py b/flocks/session/prompt.py index 75efef655..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 @@ -37,6 +37,7 @@ SystemPromptCache = Dict[str, Any] AsyncPromptFactory = Callable[[], Awaitable[Optional[str]]] StringPromptFactory = Callable[[], Optional[str]] +AsyncPromptLoader = AsyncPromptFactory # Prompt template directory (same structure as Flocks) @@ -139,13 +140,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 +822,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 +884,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 +894,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 +1016,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 +1024,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 +1035,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 +1097,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 +1193,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 +1249,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 +1317,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 +1364,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 +1428,116 @@ 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] = (), + 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. + + 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, 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/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/runner.py b/flocks/session/runner.py index b07a35ed3..0bdf81cf9 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, @@ -590,18 +590,35 @@ async def _list_callable_tool_infos_for_turn( messages: List[MessageInfo], ) -> Tuple[List[Any], Dict[str, Any]]: execution_mode = self._execution_mode_from_messages(messages) + from flocks.agent.tool_permissions import QUESTION_TOOL_NAME + from flocks.permission.next import PermissionNext + + declared_tool_names = getattr(agent, "tools", None) + permission_ruleset = getattr(self, "_turn_permission_ruleset", None) + if permission_ruleset is None: + permission_ruleset = self._permission_ruleset_for_agent(agent) result = await list_session_callable_tool_infos( session_id=self.session.id, - declared_tool_names=getattr(agent, "tools", None), + declared_tool_names=declared_tool_names, agent=agent.name, step=self._step, event_publish_callback=self.callbacks.event_publish_callback, ) - tool_infos = [ - tool_info - for tool_info in result.tool_infos - if is_tool_allowed(execution_mode, tool_info.name) - ] + permission_denied_tool_names: List[str] = [] + tool_infos = [] + for tool_info in result.tool_infos: + if not is_tool_allowed(execution_mode, tool_info.name): + continue + if tool_info.name == QUESTION_TOOL_NAME: + denied_by_permission = PermissionNext.evaluate_request( + tool_info.name, + ["*"], + permission_ruleset, + ) == "deny" + if denied_by_permission: + permission_denied_tool_names.append(tool_info.name) + continue + tool_infos.append(tool_info) if ( execution_mode == SessionExecutionMode.PLAN and all(tool_info.name != "plan_exit" for tool_info in tool_infos) @@ -611,6 +628,7 @@ async def _list_callable_tool_infos_for_turn( tool_infos.append(plan_exit.info) metadata = dict(result.metadata) metadata["executionMode"] = execution_mode.value + metadata["permissionDeniedToolNames"] = sorted(permission_denied_tool_names) metadata["modeAllowedToolNames"] = sorted( tool_info.name for tool_info in tool_infos ) @@ -1338,6 +1356,8 @@ def classify_failover_error(error: Dict[str, Any]) -> FailoverDecision: "policy violation", )): return FailoverDecision(True, "content_policy") + if error_name == "StreamToolArgumentsTruncatedError": + return FailoverDecision(True, "stream_truncated") if error_name == "JSONDecodeError" or any( pattern in lowered for pattern in ( "malformed response", "invalid response", "empty choices", @@ -1360,12 +1380,18 @@ def _deferred_failure_result( assistant_message_id: Optional[str], decision: FailoverDecision, attempts: int, + allow_fallback_override: Optional[bool] = None, ) -> StepResult: state = LlmAttemptState( received_chunk=self._attempt_state.received_chunk, observable_output_started=self._attempt_state.observable_output_started, tool_execution_started=self._attempt_state.tool_execution_started, ) + allow_fallback = ( + allow_fallback_override + if allow_fallback_override is not None + else decision.eligible and state.replay_safe + ) return StepResult( action="stop", error=message, @@ -1374,11 +1400,15 @@ def _deferred_failure_result( error_data=error_data, assistant_message_id=assistant_message_id, reason=decision.reason, - allow_fallback=decision.eligible and state.replay_safe, + allow_fallback=allow_fallback, attempt_state=state, attempts=attempts, ), ) + + @staticmethod + def _is_stream_tool_arguments_truncated_error(error: Dict[str, Any]) -> bool: + return error.get("name") == "StreamToolArgumentsTruncatedError" async def _process_step( self, @@ -1418,6 +1448,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: @@ -1509,24 +1540,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 +1565,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) @@ -1659,34 +1643,96 @@ async def device_asset_prompt_factory() -> Optional[str]: # Disable tools when max steps reached tools = [] - # Create assistant message (will be reused across retries) - assistant_msg = await Message.create( - session_id=self.session.id, - role=MessageRole.ASSISTANT, - content="", - agent=agent.name, - model_id=self.model_id, - provider_id=self.provider_id, - parent_id=last_user.id, - ) - - # Publish assistant message SSE event so frontends can show the message card - if self.callbacks.event_publish_callback: - import time as _time + async def _publish_assistant_created(msg: MessageInfo) -> None: + if not self.callbacks.event_publish_callback: + return await self.callbacks.event_publish_callback("message.updated", { "info": { - "id": assistant_msg.id, + "id": msg.id, "sessionID": self.session.id, "role": "assistant", - "time": {"created": int(_time.time() * 1000)}, + "time": {"created": int(time.time() * 1000)}, "parentID": last_user.id, "modelID": self.model_id, "providerID": self.provider_id, "agent": agent.name, "mode": agent.name, - "tokens": {"input": 0, "output": 0, "reasoning": 0, "cache": {"read": 0, "write": 0}}, + "tokens": { + "input": 0, + "output": 0, + "reasoning": 0, + "cache": {"read": 0, "write": 0}, + }, } }) + + async def _create_attempt_assistant_message(*, publish: bool = True) -> MessageInfo: + msg = await Message.create( + session_id=self.session.id, + role=MessageRole.ASSISTANT, + content="", + agent=agent.name, + model_id=self.model_id, + provider_id=self.provider_id, + parent_id=last_user.id, + ) + if publish: + await _publish_assistant_created(msg) + return msg + + async def _replace_assistant_message_for_replay(reason: str) -> bool: + nonlocal assistant_msg + previous_msg = assistant_msg + try: + next_msg = await _create_attempt_assistant_message(publish=False) + except Exception as exc: + log.error("runner.step.replay_message_create_failed", { + "session_id": self.session.id, + "previous_message_id": previous_msg.id, + "reason": reason, + "error": str(exc), + }) + return False + + try: + deleted = await Message.delete(self.session.id, previous_msg.id) + except Exception as exc: + deleted = False + log.error("runner.step.replay_message_delete_failed", { + "session_id": self.session.id, + "previous_message_id": previous_msg.id, + "next_message_id": next_msg.id, + "reason": reason, + "error": str(exc), + }) + if not deleted: + try: + await Message.delete(self.session.id, next_msg.id) + except Exception as exc: + log.debug("runner.step.replay_message_cleanup_failed", { + "session_id": self.session.id, + "message_id": next_msg.id, + "error": str(exc), + }) + return False + + if self.callbacks.event_publish_callback: + await self.callbacks.event_publish_callback("message.removed", { + "sessionID": self.session.id, + "messageID": previous_msg.id, + }) + await _publish_assistant_created(next_msg) + assistant_msg = next_msg + log.info("runner.step.replay_message_replaced", { + "session_id": self.session.id, + "previous_message_id": previous_msg.id, + "next_message_id": next_msg.id, + "reason": reason, + }) + return True + + # Create assistant message for the first attempt. + assistant_msg = await _create_attempt_assistant_message() # Retry loop matching Flocks' SessionProcessor.process() # MAX_ERROR_RETRIES caps exception-based retries so a permanently-failing @@ -1870,6 +1916,9 @@ async def device_asset_prompt_factory() -> Optional[str]: # Check if retryable retry_message = SessionRetry.retryable(error_dict) failover_decision = self.classify_failover_error(error_dict) + is_stream_tool_args_truncated = ( + self._is_stream_tool_arguments_truncated_error(error_dict) + ) retry_limit = MAX_ERROR_RETRIES will_retry = retry_message is not None and error_attempt <= retry_limit retry_blocked_by_tool_execution = ( @@ -1883,7 +1932,18 @@ async def device_asset_prompt_factory() -> Optional[str]: elif self._defer_step_errors and not self._attempt_state.replay_safe: # Retrying after text/reasoning/tool activity can duplicate # visible output or execute a tool twice. - will_retry = False + will_retry = will_retry and is_stream_tool_args_truncated + + if will_retry and is_stream_tool_args_truncated: + # A truncated tool-argument stream already created a + # partial assistant message (and usually a tool part). The + # retry is only safe if that partial attempt can be removed + # before the next provider call. + replaced = await _replace_assistant_message_for_replay( + reason="stream_tool_arguments_truncated" + ) + if not replaced: + will_retry = False if will_retry: # Error is retryable and we have budget left @@ -1915,6 +1975,8 @@ async def device_asset_prompt_factory() -> Optional[str]: # Wait before retry await SessionRetry.sleep(delay_ms, self._abort) + + self._attempt_state = LlmAttemptState() # Continue to next retry attempt continue @@ -1944,12 +2006,19 @@ async def device_asset_prompt_factory() -> Optional[str]: error_dict["data"]["displayMessage"] = CONNECTION_ERROR_DISPLAY_MESSAGE if self._defer_step_errors: + allow_fallback_override = None + if is_stream_tool_args_truncated: + allow_fallback_override = ( + failover_decision.eligible + and not retry_blocked_by_tool_execution + ) return self._deferred_failure_result( message=final_error_message, error_data=error_dict, assistant_message_id=assistant_msg.id, decision=failover_decision, attempts=error_attempt, + allow_fallback_override=allow_fallback_override, ) if self.callbacks.on_error: @@ -2100,6 +2169,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 +2338,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( @@ -2421,6 +2624,17 @@ def _exception_to_error_dict(self, exception: Exception) -> Dict[str, Any]: } } + if type(exception).__name__ == "StreamToolArgumentsTruncatedError": + error_dict["data"].update({ + "isRetryable": True, + "streamToolArgumentsTruncated": True, + "toolCallID": getattr(exception, "tool_call_id", None), + "toolName": getattr(exception, "tool_name", None), + "finishReason": getattr(exception, "finish_reason", None), + "argumentsLength": getattr(exception, "arguments_len", None), + "argumentsPreview": getattr(exception, "arguments_preview", None), + }) + transport_exception = _find_retryable_transport_exception(exception) if transport_exception is not None: transport_type = type(transport_exception).__name__ @@ -2662,14 +2876,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 +2903,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 +2930,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. @@ -3637,7 +3875,11 @@ async def _flush_reasoning_rewriter() -> None: "agent": agent.name, }) - await tool_accumulator.flush_remaining(stream_finish_reason) + try: + await tool_accumulator.flush_remaining(stream_finish_reason) + except Exception: + await processor.drain_parallel_tool_calls() + raise if stream_text_rewriter is not None: trailing_text = stream_text_rewriter.flush() @@ -3843,13 +4085,52 @@ 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, + ) + metadata = dict(getattr(request, "metadata", None) or {}) + deny_only_preflight = metadata.get("reason") == "question_tool" tool_metadata = get_tool_catalog_metadata(str(getattr(request, "permission", "") or "")) if self.callbacks.event_publish_callback: @@ -3858,30 +4139,40 @@ 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 deny_only_preflight: + 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 "") metadata.setdefault("sessionID", self.session.id) 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/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/session/streaming/stream_events.py b/flocks/session/streaming/stream_events.py index 607054027..64d96b8d5 100644 --- a/flocks/session/streaming/stream_events.py +++ b/flocks/session/streaming/stream_events.py @@ -18,6 +18,7 @@ "tool-input-start", "tool-input-delta", "tool-input-end", + "tool-input-error", "tool-call", "tool-result", "tool-error", @@ -82,6 +83,15 @@ class ToolInputEndEvent(BaseStreamEvent): id: str +class ToolInputErrorEvent(BaseStreamEvent): + """Tool input failed before a runnable tool call was produced.""" + type: Literal["tool-input-error"] = "tool-input-error" + id: str + tool_name: str + input: Dict[str, Any] = Field(default_factory=dict) + error: str + + class ToolCallEvent(BaseStreamEvent): """Tool call request (ready to execute)""" type: Literal["tool-call"] = "tool-call" @@ -152,6 +162,7 @@ class FinishEvent(BaseStreamEvent): ToolInputStartEvent | ToolInputDeltaEvent | ToolInputEndEvent | + ToolInputErrorEvent | ToolCallEvent | ToolResultEvent | ToolErrorEvent | @@ -184,6 +195,7 @@ def event_from_dict(data: Dict[str, Any]) -> StreamEvent: "tool-input-start": ToolInputStartEvent, "tool-input-delta": ToolInputDeltaEvent, "tool-input-end": ToolInputEndEvent, + "tool-input-error": ToolInputErrorEvent, "tool-call": ToolCallEvent, "tool-result": ToolResultEvent, "tool-error": ToolErrorEvent, diff --git a/flocks/session/streaming/stream_processor.py b/flocks/session/streaming/stream_processor.py index dc4e4b967..c03940a5d 100644 --- a/flocks/session/streaming/stream_processor.py +++ b/flocks/session/streaming/stream_processor.py @@ -38,6 +38,7 @@ TextDeltaEvent, TextEndEvent, ToolInputStartEvent, + ToolInputErrorEvent, ) from flocks.tool.registry import ToolRegistry, ToolContext, ToolResult from flocks.permission import PermissionNext @@ -218,6 +219,9 @@ async def process_event(self, event: StreamEvent) -> None: elif event_type == "tool-input-end": pass # Input is complete + + elif event_type == "tool-input-error": + await self._handle_tool_input_error(event) elif event_type == "tool-call": if self._should_run_tool_call_parallel(event): @@ -471,6 +475,69 @@ async def _handle_tool_input_start(self, event: ToolInputStartEvent) -> None: }) except Exception as e: log.error("stream.tool_input_start.store_part_failed", {"error": str(e)}) + + async def _handle_tool_input_error(self, event: ToolInputErrorEvent) -> None: + """Mark an input-generation failure without executing a tool.""" + if event.id in self.tool_calls: + tool_state = self.tool_calls[event.id] + part_id = tool_state.part_id + else: + part_id = Identifier.create("part") + tool_state = ToolCallState( + id=event.id, + name=event.tool_name, + input=event.input, + part_id=part_id, + status="pending", + ) + self.tool_calls[event.id] = tool_state + + tool_state.name = event.tool_name + tool_state.input = event.input + tool_state.status = "error" + tool_state.error = event.error + + tool_error_time = int(datetime.now().timestamp() * 1000) + error_state = ToolStateError( + status="error", + input=event.input, + error=event.error, + time={"start": tool_error_time, "end": tool_error_time}, + ) + error_part = ToolPart( + id=part_id, + sessionID=self.session_id, + messageID=self.assistant_message.id, + type="tool", + callID=event.id, + tool=event.tool_name, + state=error_state, + ) + await Message.store_part(self.session_id, self.assistant_message.id, error_part) + + if self.event_publish_callback: + await self.event_publish_callback("message.part.updated", { + "part": { + "id": part_id, + "messageID": self.assistant_message.id, + "sessionID": self.session_id, + "type": "tool", + "callID": event.id, + "tool": event.tool_name, + "state": { + "status": "error", + "input": event.input, + "error": event.error, + "time": {"start": tool_error_time, "end": tool_error_time}, + }, + }, + }) + + log.warn("stream.tool_input.error", { + "tool_call_id": event.id, + "tool_name": event.tool_name, + "error": event.error, + }) async def _handle_tool_call(self, event: ToolCallEvent) -> None: """ @@ -1778,7 +1845,7 @@ def _parse_dsml_text_tool_calls(self, text: str) -> list[dict]: re.DOTALL | re.IGNORECASE, ): body = match.group(1).strip() - if not body or not body[:1] in "{[": + if not body or body[:1] not in "{[": continue try: diff --git a/flocks/session/streaming/tool_accumulator.py b/flocks/session/streaming/tool_accumulator.py index 6e266e95d..bfd6536de 100644 --- a/flocks/session/streaming/tool_accumulator.py +++ b/flocks/session/streaming/tool_accumulator.py @@ -20,12 +20,37 @@ from flocks.tool.registry import ToolRegistry from flocks.session.streaming.stream_events import ( ToolInputStartEvent, + ToolInputErrorEvent, ToolCallEvent, ) log = Log.create(service="tool_accumulator") +class StreamToolArgumentsTruncatedError(RuntimeError): + """Raised when the stream ends before tool arguments form valid JSON.""" + + def __init__( + self, + *, + tool_call_id: str, + tool_name: str, + finish_reason: str, + arguments_len: int, + arguments_preview: str, + ) -> None: + self.tool_call_id = tool_call_id + self.tool_name = tool_name + self.finish_reason = finish_reason + self.arguments_len = arguments_len + self.arguments_preview = arguments_preview + super().__init__( + "Model output was truncated while generating tool arguments " + f"for '{tool_name}' (finish_reason='{finish_reason}', " + f"{arguments_len} chars). The tool was not executed." + ) + + class ToolCallAccumulator: """Accumulates streamed tool-call JSON fragments and dispatches execution. @@ -121,13 +146,52 @@ async def flush_remaining( ) -> None: """Process any tool calls still in the accumulator after the stream ends.""" is_truncated = stream_finish_reason in ("length", "max_tokens") + truncation_error: StreamToolArgumentsTruncatedError | None = None for tc_id, tc_data in list(self._accumulator.items()): - if tc_data.get("completed"): + if tc_data.get("completed") or tc_data.get("failed"): continue accumulated_args = tc_data.get("arguments_str", "") tool_name = tc_data.get("name", "") - if not (accumulated_args and tool_name): + if not tool_name: + continue + + if is_truncated: + if accumulated_args: + detail = ( + f"Tool arguments for '{tool_name}' cut off at " + f"{len(accumulated_args)} chars." + ) + else: + detail = f"Tool arguments for '{tool_name}' were not completed." + error_msg = ( + f"Output was truncated (finish_reason='{stream_finish_reason}'). " + f"{detail} The tool was not executed." + ) + await self._processor.process_event( + ToolInputErrorEvent( + id=tc_id, + tool_name=tool_name, + input={ + "tool": tool_name, + "arguments_preview": accumulated_args[:500], + "finish_reason": stream_finish_reason, + }, + error=error_msg, + ) + ) + tc_data["failed"] = True + if truncation_error is None: + truncation_error = StreamToolArgumentsTruncatedError( + tool_call_id=tc_id, + tool_name=tool_name, + finish_reason=str(stream_finish_reason), + arguments_len=len(accumulated_args), + arguments_preview=accumulated_args[:500], + ) + continue + + if not accumulated_args: continue arguments, ok = _parse_json_robust(accumulated_args) @@ -151,17 +215,10 @@ async def flush_remaining( continue # All strategies failed — redirect to invalid tool - if is_truncated: - error_msg = ( - f"Output was truncated (finish_reason='{stream_finish_reason}'). " - f"Tool arguments for '{tool_name}' cut off at {len(accumulated_args)} chars. " - f"Please reduce the content size or split the operation." - ) - else: - error_msg = ( - f"Failed to parse tool arguments ({len(accumulated_args)} chars). " - f"Please ensure valid JSON with balanced braces/brackets." - ) + error_msg = ( + f"Failed to parse tool arguments ({len(accumulated_args)} chars). " + f"Please ensure valid JSON with balanced braces/brackets." + ) await self._processor.process_event( ToolCallEvent( @@ -175,6 +232,9 @@ async def flush_remaining( ) ) + if truncation_error is not None: + raise truncation_error + # ------------------------------------------------------------------ # Internal helpers # ------------------------------------------------------------------ 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/file/edit.py b/flocks/tool/file/edit.py index 7eb48ac4d..ca29c325c 100644 --- a/flocks/tool/file/edit.py +++ b/flocks/tool/file/edit.py @@ -1,7 +1,5 @@ -""" -Edit Tool - File editing with batch exact replacements +"""Edit Tool - File editing with one or more targeted replacements. -Supports both legacy single-edit arguments and pi-style edits[] batch edits. All edits in one call are matched against the same original file snapshot. """ @@ -29,18 +27,19 @@ DESCRIPTION = """Edit a single existing file using targeted text replacement. -Use this tool for one or more changes within the same file. +Provide every replacement in the required `edits` array: +`{"filePath":"path/to/file","edits":[{"oldString":"current text","newString":"replacement text"}]}` Do not use this tool when a dedicated tool is a better fit: - Create a new file -> `write` - Modify, add, delete, or move multiple files in one coordinated change -> `apply_patch` Usage notes: -- Prefer `edits` for one or more disjoint replacements. +- `edits` must contain at least one replacement. - Every `edits[].oldString` is matched against the original file snapshot. - Each `oldString` must be unique, and edits must not overlap. -- The tool preserves the file's encoding and line-ending style. -- Legacy `oldString` / `newString` / `replaceAll` remains supported.""" +- To replace repeated text, provide enough surrounding context to target each occurrence uniquely. +- The tool preserves the file's encoding and line-ending style.""" def normalize_line_endings(text: str) -> str: @@ -311,72 +310,28 @@ def _get_no_change_error(filepath: str, total_edits: int) -> str: return f"No changes made to {filepath}. The replacements produced identical content." -def _prepare_batch_edits( +def _prepare_edits( filepath: str, - edits: Optional[List[Dict[str, Any]]], - old_string: Optional[str], - new_string: Optional[str], + edits: List[Dict[str, Any]], ) -> tuple[Optional[List[Dict[str, str]]], Optional[str]]: - """Return normalized batch edits or a validation error.""" - if edits is not None: - if old_string is not None or new_string is not None: - return None, "Use either edits or oldString/newString, not both." - if not isinstance(edits, list) or not edits: - return None, "edits must contain at least one replacement." - - prepared: List[Dict[str, str]] = [] - for index, edit in enumerate(edits): - if not isinstance(edit, dict): - return None, f"edits[{index}] must be an object." - current_old = edit.get("oldString") - current_new = edit.get("newString") - if not isinstance(current_old, str) or not isinstance(current_new, str): - return None, f"edits[{index}] must include string oldString and newString." - if current_old == "": - return None, _get_empty_old_string_error(filepath, index, len(edits)) - if current_old == current_new: - return None, f"edits[{index}].oldString and newString must be different." - prepared.append({"oldString": current_old, "newString": current_new}) - return prepared, None - - if old_string is None or new_string is None: - return None, "Provide edits or legacy oldString/newString arguments." - if old_string != "" and old_string == new_string: - return None, "oldString and newString must be different" - return [{"oldString": old_string, "newString": new_string}], None - - -def _apply_replace_all( - normalized_content: str, - old_string: str, - new_string: str, - filepath: str, -) -> tuple[str, str]: - """Apply a legacy replaceAll operation without mutating untouched content.""" - normalized_old = normalize_line_endings(old_string) - normalized_new = normalize_line_endings(new_string) - content_index = _build_fuzzy_text_index(normalized_content) - found, _, _, used_fuzzy = _fuzzy_find_text( - normalized_content, - normalized_old, - content_index, - ) - if not found: - raise ValueError(_get_not_found_error(filepath, 0, 1)) - - if not used_fuzzy: - if normalized_old == "": - raise ValueError(_get_empty_old_string_error(filepath, 0, 1)) - new_content = normalized_content.replace(normalized_old, normalized_new) - else: - fuzzy_spans = _find_fuzzy_spans(content_index, normalized_old) - new_content = normalized_content - for match_start, match_end in reversed(fuzzy_spans): - new_content = new_content[:match_start] + normalized_new + new_content[match_end:] - - if new_content == normalized_content: - raise ValueError(_get_no_change_error(filepath, 1)) - return normalized_content, new_content + """Validate and normalize the required edits array.""" + if not isinstance(edits, list) or not edits: + return None, "edits must contain at least one replacement." + + prepared: List[Dict[str, str]] = [] + for index, edit in enumerate(edits): + if not isinstance(edit, dict): + return None, f"edits[{index}] must be an object." + current_old = edit.get("oldString") + current_new = edit.get("newString") + if not isinstance(current_old, str) or not isinstance(current_new, str): + return None, f"edits[{index}] must include string oldString and newString." + if current_old == "": + return None, _get_empty_old_string_error(filepath, index, len(edits)) + if current_old == current_new: + return None, f"edits[{index}].oldString and newString must be different." + prepared.append({"oldString": current_old, "newString": current_new}) + return prepared, None def _apply_edits_to_normalized_content( @@ -463,12 +418,14 @@ def _apply_edits_to_normalized_content( name="edits", type=ParameterType.ARRAY, description=( - "One or more targeted replacements. Each edits[].oldString is matched " - "against the original file, not incrementally." + "Required non-empty list of targeted replacements. Each oldString is " + "matched against the original file snapshot, must be unique, and must " + "not overlap another replacement." ), - required=False, + required=True, json_schema={ "type": "array", + "minItems": 1, "items": { "type": "object", "properties": { @@ -490,36 +447,14 @@ def _apply_edits_to_normalized_content( }, }, ), - ToolParameter( - name="oldString", - type=ParameterType.STRING, - description="Legacy single-edit old text. Use edits[] for new callers.", - required=False, - ), - ToolParameter( - name="newString", - type=ParameterType.STRING, - description="Legacy single-edit replacement text. Use edits[] for new callers.", - required=False, - ), - ToolParameter( - name="replaceAll", - type=ParameterType.BOOLEAN, - description="Legacy single-edit option to replace every occurrence of oldString.", - required=False, - default=False, - ), ], ) async def edit_tool( ctx: ToolContext, filePath: str, - edits: Optional[List[Dict[str, Any]]] = None, - oldString: Optional[str] = None, - newString: Optional[str] = None, - replaceAll: bool = False, + edits: List[Dict[str, Any]], ) -> ToolResult: - """Edit a file with legacy single-edit or pi-style edits[] semantics.""" + """Edit an existing file using one or more targeted replacements.""" if not filePath: return ToolResult(success=False, error="filePath is required") @@ -572,54 +507,7 @@ async def edit_tool( title = resolution.display_path - if oldString == "" and edits is None: - if newString is None: - return ToolResult(success=False, error="newString is required when oldString is empty", title=title) - if ctx.agent == "self-improve" and Path(filepath).name == "SKILL.md": - from flocks.memory.evolution.skill_guard import ( - validate_evolution_skill_write, - ) - - skill_path = Path(filepath) - if skill_path.exists(): - evolution_error = ( - "Read the existing managed Skill and use a precise edit" - ) - else: - evolution_error = await validate_evolution_skill_write( - skill_path, - newString, - exists=False, - ) - if evolution_error: - return ToolResult( - success=False, - error=evolution_error, - title=title, - ) - diff = trim_diff(generate_diff(filepath, "", newString)) - parent_dir = os.path.dirname(filepath) - if parent_dir and not os.path.exists(parent_dir): - os.makedirs(parent_dir, exist_ok=True) - - try: - with open(filepath, "w", encoding="utf-8", newline="") as file_handle: - file_handle.write(newString) - except Exception as error: - return ToolResult( - success=False, - error=f"Failed to write file: {str(error)}", - title=title, - ) - - return ToolResult( - success=True, - output="Edit applied successfully. If you need to make additional edits to this file, use the Read tool first to get the current file content.", - title=title, - metadata={"diff": diff, "diagnostics": {}}, - ) - - prepared_edits, validation_error = _prepare_batch_edits(filepath, edits, oldString, newString) + prepared_edits, validation_error = _prepare_edits(filepath, edits) if validation_error: return ToolResult(success=False, error=validation_error, title=title) assert prepared_edits is not None @@ -648,22 +536,11 @@ async def edit_tool( normalized_content_old = normalize_line_endings(content_without_bom) try: - if replaceAll: - if edits is not None: - raise ValueError("replaceAll is only supported with legacy oldString/newString arguments.") - assert oldString is not None and newString is not None - base_content, normalized_content_new = _apply_replace_all( - normalized_content_old, - oldString, - newString, - filepath, - ) - else: - base_content, normalized_content_new = _apply_edits_to_normalized_content( - normalized_content_old, - prepared_edits, - filepath, - ) + base_content, normalized_content_new = _apply_edits_to_normalized_content( + normalized_content_old, + prepared_edits, + filepath, + ) except ValueError as error: return ToolResult(success=False, error=str(error), title=title) diff --git a/flocks/tool/system/question.py b/flocks/tool/system/question.py index e83af3201..e97447d3b 100644 --- a/flocks/tool/system/question.py +++ b/flocks/tool/system/question.py @@ -422,6 +422,21 @@ async def question_tool( error="No valid questions provided" ) + try: + await ctx.ask( + "question", + ["*"], + metadata={"reason": "question_tool"}, + ) + except (PermissionError, asyncio.TimeoutError) as e: + error = str(e) or "question" + if not error.lower().startswith("permission denied"): + error = f"Permission denied: {error}" + return ToolResult( + success=False, + error=error, + ) + channel_result = await _send_channel_question_if_applicable(ctx, normalized_questions) if channel_result is not None: return channel_result diff --git a/flocks/tool/task/run_workflow.py b/flocks/tool/task/run_workflow.py index 9b78e5123..7c50452b7 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,16 @@ 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 = ExecutionStepRecorder() + progress_writer: Optional[ExecutionProgressWriter] = None 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 +605,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 +627,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 +638,17 @@ 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) - 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) + 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 - 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 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 ctx.abort.is_set(): + progress_update["currentPhase"] = "cancelling" 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 +658,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": step_recorder.step_count, + "step_count": step_recorder.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 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 +694,7 @@ async def _flush_pending_step() -> None: canonical_workflow_id, input_params=workflow_inputs, ) + progress_writer = ExecutionProgressWriter(tracked_execution) # Update metadata to show workflow is running _emit_metadata( @@ -890,7 +814,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 +830,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 = step_recorder.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 +859,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 +881,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 +920,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 +959,34 @@ async def _flush_pending_step() -> None: "error": error_msg, }, ) + terminal_status = "cancelled" if ctx.abort.is_set() else "error" + final_step_count = step_recorder.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 +996,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 +1012,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/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/execution_store.py b/flocks/workflow/execution_store.py index 9d0ae1c79..6503ee981 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,67 +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] = {} - -# 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) - if lock is None: - lock = asyncio.Lock() - _trim_locks[workflow_id] = 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: @@ -432,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. """ @@ -453,26 +393,17 @@ 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", 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.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 +429,112 @@ 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), - }, + 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 + + +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 _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 + async def _flush(self) -> None: 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 + 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): + 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( @@ -622,7 +617,7 @@ async def create_execution_record( input_params: Optional[Dict[str, Any]] = None, exec_id: Optional[str] = None, ) -> Dict[str, Any]: - """Create and persist a running workflow execution record. + """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 @@ -645,18 +640,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 +664,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) - await _update_workflow_stats(workflow_id, success, float(duration)) + + await WorkflowStore.complete_execution( + compact_execution_summary(summary_data), + prepared_steps, + ) + + 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 @@ -676,7 +717,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( @@ -686,6 +727,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: @@ -694,80 +748,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 - - # 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, - *, - 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 9db0bdd4b..9d348d5db 100644 --- a/flocks/workflow/poller_manager.py +++ b/flocks/workflow/poller_manager.py @@ -448,15 +448,12 @@ 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_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", + exec_data = await create_execution_record( + workflow_id, + input_params=inputs, ) + exec_id = str(exec_data["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) @@ -555,9 +552,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..129d62131 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 @@ -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: @@ -44,8 +51,10 @@ 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 @classmethod def get_db_path(cls) -> Path: @@ -57,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", { @@ -70,11 +78,16 @@ async def init(cls) -> None: "new_db_path": str(db_path), }, ) - if cls._conn: - await cls._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) @@ -90,9 +103,16 @@ 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 + cls._completion_lock = asyncio.Lock() await cls._migrate_legacy_kv() try: @@ -101,7 +121,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 @@ -120,10 +143,14 @@ 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 + cls._completion_lock = None @classmethod async def _db(cls) -> aiosqlite.Connection: @@ -137,6 +164,19 @@ 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 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] + + @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) @@ -280,21 +320,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, @@ -316,6 +353,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 @@ -354,7 +397,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), @@ -396,7 +439,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)), @@ -408,6 +451,26 @@ async def trim_executions(cls, workflow_id: str, *, keep: int) -> List[str]: await db.commit() return exec_ids + @classmethod + def _step_rows( + cls, + exec_id: str, + steps: Iterable[Tuple[int, Dict[str, Any]]], + ) -> List[Tuple[Any, ...]]: + return [ + ( + 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 + ] + @classmethod async def record_step( cls, @@ -415,26 +478,71 @@ 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 = cls._step_rows(exec_id, 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 (?, ?, ?, ?, ?, ?, ?, ?) """, - ( - 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), - ), + rows, ) await db.commit() + @classmethod + async def complete_execution( + cls, + exec_data: Dict[str, Any], + steps: Iterable[Tuple[int, Dict[str, Any]]], + ) -> None: + """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) + step_rows = cls._step_rows(exec_id, steps) + lock = cls._completion_lock + if lock is None: + lock = asyncio.Lock() + cls._completion_lock = lock + + async with lock: + try: + await db.execute("BEGIN IMMEDIATE") + 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(_EXECUTION_UPSERT_SQL, execution_row) + await db.commit() + 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 async def list_steps( cls, 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/flocks/workflow/triggers/runtime.py b/flocks/workflow/triggers/runtime.py index 5e1a5b3da..2d88fb1e4 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, @@ -250,6 +250,7 @@ async def _execute_workflow_effect( input_params=mapped_inputs, ) exec_id = exec_data["id"] + step_recorder = ExecutionStepRecorder() started_at = time.time() tool_context = None try: @@ -266,10 +267,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 +283,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 +296,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 +317,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/flocks/workspace/models.py b/flocks/workspace/models.py index b5e0bee49..c29563d77 100644 --- a/flocks/workspace/models.py +++ b/flocks/workspace/models.py @@ -14,6 +14,7 @@ class WorkspaceNode(BaseModel): size: Optional[int] = None modified_at: Optional[float] = None is_text_file: bool = False + editable: bool = False children: Optional[list["WorkspaceNode"]] = None diff --git a/tests/agent/test_prompt_builders.py b/tests/agent/test_prompt_builders.py new file mode 100644 index 000000000..fc65ba067 --- /dev/null +++ b/tests/agent/test_prompt_builders.py @@ -0,0 +1,53 @@ +"""Direct tests for Rex and Hephaestus prompt builders.""" + +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 + + +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 + + +@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/config/test_config.py b/tests/config/test_config.py index 713448883..85ca64e89 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) @@ -153,6 +154,82 @@ 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.ALLOW, + "explore": PermissionAction.DENY, + } + 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_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/ingest/test_kafka_manager.py b/tests/ingest/test_kafka_manager.py index a3c2b0432..7e3a19043 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,19 @@ 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 + ): 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 +594,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 +615,15 @@ 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 [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 @@ -635,11 +635,15 @@ 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 + ): 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 @@ -682,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, @@ -691,10 +743,14 @@ 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 + ): 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..29750f055 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,18 @@ 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 + ): 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 +389,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 +426,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[0][0] == "exec-syslog" - assert recorded_steps[0][1] == 1 - assert recorded_steps[0][2]["node_id"] == "receive_alert" + 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/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"] diff --git a/tests/sandbox/test_sandbox_file_tools.py b/tests/sandbox/test_sandbox_file_tools.py index 2d1f07667..86a83f1cd 100644 --- a/tests/sandbox/test_sandbox_file_tools.py +++ b/tests/sandbox/test_sandbox_file_tools.py @@ -99,8 +99,7 @@ async def test_file_tools_allow_only_host_memory_root_in_sandbox( "edit", ctx=ctx, filePath=str(memory_file), - oldString="old fact", - newString="new fact", + edits=[{"oldString": "old fact", "newString": "new fact"}], ) assert read_result.success @@ -175,22 +174,34 @@ async def test_sandbox_self_improve_can_manage_only_marked_host_skills( "edit", ctx=ctx, filePath=str(managed_path), - oldString="Initial workflow.", - newString="Improved workflow.", + edits=[ + { + "oldString": "Initial workflow.", + "newString": "Improved workflow.", + } + ], ) unmanaged_result = await ToolRegistry.execute( "edit", ctx=ctx, filePath=str(unmanaged_path), - oldString="Manual workflow.", - newString="Changed workflow.", + edits=[ + { + "oldString": "Manual workflow.", + "newString": "Changed workflow.", + } + ], ) project_skill_result = await ToolRegistry.execute( "edit", ctx=ctx, filePath=str(project_skill_path), - oldString="", - newString=managed_content.replace("managed-skill", "project-skill"), + edits=[ + { + "oldString": "managed-skill", + "newString": "project-skill", + } + ], ) assert create_result.success @@ -204,7 +215,7 @@ async def test_sandbox_self_improve_can_manage_only_marked_host_skills( assert "existing managed Skills" in (unmanaged_result.error or "") assert unmanaged_path.read_text(encoding="utf-8") == unmanaged_content assert not project_skill_result.success - assert "outside the self-improve user root" in (project_skill_result.error or "") + assert "not found" in (project_skill_result.error or "") assert not project_skill_path.exists() @@ -231,8 +242,12 @@ async def test_sandbox_agent_cannot_write_or_edit_daily_memory( "edit", ctx=ctx, filePath=str(daily_file), - oldString="lifecycle entry", - newString="replacement", + edits=[ + { + "oldString": "lifecycle entry", + "newString": "replacement", + } + ], ) assert not write_result.success @@ -268,8 +283,7 @@ async def test_edit_tool_rejects_path_outside_sandbox() -> None: "edit", ctx=ctx, filePath=outside_file, - oldString="hello", - newString="world", + edits=[{"oldString": "hello", "newString": "world"}], ) assert not result.success assert "Path escapes sandbox workspace" in (result.error or "") diff --git a/tests/server/routes/test_agent_routes.py b/tests/server/routes/test_agent_routes.py index 4a22e83d7..2e091d87c 100644 --- a/tests/server/routes/test_agent_routes.py +++ b/tests/server/routes/test_agent_routes.py @@ -173,6 +173,93 @@ async def test_created_agent_survives_registry_reload(self, client: AsyncClient) assert list_resp.status_code == status.HTTP_200_OK assert "test-agent" in [agent["name"] for agent in list_resp.json()] + @pytest.mark.asyncio + async def test_create_agent_without_tools_field_keeps_permission_unchanged( + self, + client: AsyncClient, + ): + """Older clients that omit tools do not implicitly disable question.""" + payload = {k: v for k, v in _AGENT_PAYLOAD.items() if k != "tools"} + payload["name"] = "legacy-create-agent" + + resp = await client.post("/api/agent", json=payload) + + assert resp.status_code == status.HTTP_200_OK + assert resp.json()["permission"] == [] + + @pytest.mark.asyncio + async def test_create_agent_without_question_tool_adds_persistent_deny(self, client: AsyncClient): + """Unchecking the question tool persists a deny rule for the always-load tool.""" + from flocks.agent.registry import Agent + from flocks.storage.storage import Storage + + payload = { + **_AGENT_PAYLOAD, + "name": "no-question-agent", + "tools": ["tool_search"], + } + + resp = await client.post("/api/agent", json=payload) + assert resp.status_code == status.HTTP_200_OK + assert resp.json()["permission"] == [{ + "permission": "question", + "action": "deny", + "pattern": "*", + "source": "agent_tools", + }] + + stored = await Storage.read("agent/custom/no-question-agent") + assert stored["permission"] == resp.json()["permission"] + + Agent._custom_agents.clear() + Agent.invalidate_cache() + + agent = await Agent.get("no-question-agent") + assert agent is not None + assert any( + rule.permission == "question" and rule.level.value == "deny" + for rule in agent.permission + ) + + get_resp = await client.get("/api/agent/no-question-agent") + assert get_resp.status_code == status.HTTP_200_OK + assert get_resp.json()["permission"] == resp.json()["permission"] + + @pytest.mark.asyncio + async def test_update_agent_question_tool_toggle_adds_and_removes_managed_deny( + self, + client: AsyncClient, + ): + """The Tools checkbox controls only the system-managed question deny rule.""" + payload = { + **_AGENT_PAYLOAD, + "name": "question-toggle-agent", + "tools": ["question", "tool_search"], + } + + create_resp = await client.post("/api/agent", json=payload) + assert create_resp.status_code == status.HTTP_200_OK + assert create_resp.json()["permission"] == [] + + disable_resp = await client.put( + "/api/agent/question-toggle-agent", + json={"tools": ["tool_search"]}, + ) + assert disable_resp.status_code == status.HTTP_200_OK + assert disable_resp.json()["permission"] == [{ + "permission": "question", + "action": "deny", + "pattern": "*", + "source": "agent_tools", + }] + + enable_resp = await client.put( + "/api/agent/question-toggle-agent", + json={"tools": ["question", "tool_search"]}, + ) + assert enable_resp.status_code == status.HTTP_200_OK + assert enable_resp.json()["permission"] == [] + @pytest.mark.asyncio async def test_create_subagent_defaults_to_delegatable(self, client: AsyncClient): """Sub-agents default to delegatable=true when the field is omitted.""" diff --git a/tests/server/routes/test_workflow_run_route.py b/tests/server/routes/test_workflow_run_route.py index 588fe227b..69c213656 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 @@ -40,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] = [] @@ -104,9 +116,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 +217,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,35 +236,44 @@ 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=[], last_node_id="node-1", steps=1, ) - ) - record_result = AsyncMock(return_value=None) - storage_read = AsyncMock( - return_value={ - "id": "exec-1", - "workflowId": "wf-1", - "currentNodeType": "tool", - "executionLog": [], - } - ) + run_mock = Mock(side_effect=run_workflow_mock) + record_result = AsyncMock(return_value=None) + upsert_execution = AsyncMock(return_value=None) 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 = _progress_writer("exec-1") await workflow_module._run_workflow_execution_task( workflow_id="wf-1", @@ -264,13 +281,391 @@ async def test_run_workflow_execution_task_reuses_existing_mcp_without_reinit( 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 + 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() + 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 +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.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 = _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": []}, + 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 + 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 = _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": []}, + 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": "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 = _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": []}, + 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 = _progress_writer("exec-storage-error") + + 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 = _progress_writer("exec-blocked-progress") + 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 = _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() + 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 diff --git a/tests/session/test_auto_model_failover.py b/tests/session/test_auto_model_failover.py index b1d655b1f..6cf432e65 100644 --- a/tests/session/test_auto_model_failover.py +++ b/tests/session/test_auto_model_failover.py @@ -192,6 +192,78 @@ async def test_auto_runner_uses_standard_retry_policy( assert result.failure is not None assert call_llm.await_count == expected_calls + assert result.failure.allow_fallback is True + + +@pytest.mark.asyncio +async def test_retry_exhausted_safe_api_error_switches_to_fallback(monkeypatch): + ctx = _ctx() + last_user = SimpleNamespace(id="msg_user", agent="rex", role="user") + events = [] + calls = [] + create_count = 0 + + provider = MagicMock() + provider.is_configured.return_value = True + + async def create_message(**_kwargs): + nonlocal create_count + create_count += 1 + return SimpleNamespace(id=f"msg_assistant_{create_count}") + + async def call_llm(runner, *_args, **_kwargs): + calls.append((runner.provider_id, runner.model_id)) + if runner.provider_id == "primary": + failure = RuntimeError("Provider HTTP 500") + failure.status_code = 500 + raise failure + return StepResult(action="stop", content="recovered") + + async def publish(event, payload): + events.append((event, payload)) + + monkeypatch.setattr( + "flocks.session.runner.Agent.get", + AsyncMock(return_value=SimpleNamespace( + name="rex", + steps=None, + mode="primary", + prompt="", + tools=[], + )), + ) + monkeypatch.setattr("flocks.session.runner.Provider.get", lambda _provider_id: provider) + monkeypatch.setattr("flocks.session.runner.Provider.apply_config", AsyncMock()) + monkeypatch.setattr( + "flocks.session.runner.SessionPrompt.build_system_prompts", + AsyncMock(return_value=[]), + ) + monkeypatch.setattr(SessionRunner, "_build_callable_tool_schema", AsyncMock(return_value=[])) + monkeypatch.setattr( + SessionRunner, + "_to_chat_messages", + AsyncMock(return_value=[SimpleNamespace(role="user", content="hello")]), + ) + monkeypatch.setattr(Message, "get_text_content", AsyncMock(return_value="hello")) + monkeypatch.setattr(Message, "parts", AsyncMock(return_value=[])) + monkeypatch.setattr(Message, "create", create_message) + monkeypatch.setattr(Message, "delete", AsyncMock(return_value=True)) + monkeypatch.setattr(Message, "update", AsyncMock()) + monkeypatch.setattr(SessionRunner, "_call_llm", call_llm) + monkeypatch.setattr("flocks.session.runner.SessionRetry.sleep", AsyncMock()) + + result = await SessionLoop._process_step_with_failover( + ctx, + LoopCallbacks(event_publish_callback=publish), + [last_user], + last_user, + ) + + assert result.content == "recovered" + assert calls == [("primary", "primary-model")] * 6 + [("fallback", "fallback-model")] + assert (ctx.provider_id, ctx.model_id) == ("fallback", "fallback-model") + assert any(event == "message.removed" for event, _ in events) + assert any(event == "session.model.fallback" for event, _ in events) @pytest.mark.asyncio diff --git a/tests/session/test_context_usage.py b/tests/session/test_context_usage.py index 2bc9db71a..524bead19 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, MagicMock import pytest from flocks.session import context_usage +from flocks.session.prompt import SystemPromptBlock, TurnPromptContext def _message( @@ -275,12 +277,12 @@ 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( type="tool", - tool="delegate_task", + tool="task", state=SimpleNamespace(input={}, output="d" * 40, time={"start": 4}), ), SimpleNamespace( @@ -288,20 +290,77 @@ 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 + + +@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), + ) + 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: "/ambient-worktree", + ) + 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, + ) + worktree_for_directory.assert_called_once_with("/workspace/project") diff --git a/tests/session/test_execution_mode.py b/tests/session/test_execution_mode.py index 0ef33c1de..b933911d5 100644 --- a/tests/session/test_execution_mode.py +++ b/tests/session/test_execution_mode.py @@ -1,4 +1,5 @@ from types import SimpleNamespace +from unittest.mock import AsyncMock import pytest from fastapi import HTTPException @@ -406,6 +407,7 @@ async def test_runner_filters_tools_with_message_mode(monkeypatch) -> None: runner._step = 1 runner.callbacks = SimpleNamespace(event_publish_callback=None) agent = SimpleNamespace( + name="rex", tools=[ "read", "bash", @@ -479,3 +481,240 @@ async def list_tools(**_kwargs): "task", "write", ] + + +@pytest.mark.asyncio +async def test_runner_hides_always_load_question_when_permission_denied(monkeypatch) -> None: + from flocks.permission.helpers import from_config + from flocks.session.runner import SessionRunner + + runner = object.__new__(SessionRunner) + runner.session = SimpleNamespace(id="session-question-deny") + runner._step = 1 + runner.callbacks = SimpleNamespace(event_publish_callback=None) + runner._turn_permission_ruleset = from_config({"question": "deny"}) + agent = SimpleNamespace(name="custom-agent", tools=["request_ledger"]) + + result = SimpleNamespace( + tool_infos=[ + SimpleNamespace(name="question"), + SimpleNamespace(name="tool_search"), + SimpleNamespace(name="request_ledger"), + ], + metadata={}, + ) + + async def list_tools(**_kwargs): + return result + + monkeypatch.setattr( + "flocks.session.runner.list_session_callable_tool_infos", + list_tools, + ) + + messages = [ + SimpleNamespace( + role="user", + executionMode=SessionExecutionMode.BUILD, + ) + ] + + tools, metadata = await runner._list_callable_tool_infos_for_turn( + agent, + messages, + ) + + assert [tool.name for tool in tools] == ["tool_search", "request_ledger"] + assert metadata["permissionDeniedToolNames"] == ["question"] + assert metadata["modeAllowedToolNames"] == ["request_ledger", "tool_search"] + + +@pytest.mark.asyncio +async def test_runner_keeps_always_load_question_without_permission_denial(monkeypatch) -> None: + from flocks.session.runner import SessionRunner + + runner = object.__new__(SessionRunner) + runner.session = SimpleNamespace(id="session-question-allowed") + runner._step = 1 + runner.callbacks = SimpleNamespace(event_publish_callback=None) + runner._turn_permission_ruleset = [] + agent = SimpleNamespace(name="custom-agent", tools=["question", "request_ledger"]) + + result = SimpleNamespace( + tool_infos=[ + SimpleNamespace(name="question"), + SimpleNamespace(name="tool_search"), + SimpleNamespace(name="request_ledger"), + ], + metadata={}, + ) + + async def list_tools(**_kwargs): + return result + + monkeypatch.setattr( + "flocks.session.runner.list_session_callable_tool_infos", + list_tools, + ) + + messages = [ + SimpleNamespace( + role="user", + executionMode=SessionExecutionMode.BUILD, + ) + ] + + tools, metadata = await runner._list_callable_tool_infos_for_turn( + agent, + messages, + ) + + assert [tool.name for tool in tools] == [ + "question", + "tool_search", + "request_ledger", + ] + assert metadata["permissionDeniedToolNames"] == [] + + +@pytest.mark.asyncio +async def test_runner_keeps_question_visible_when_only_tool_list_omits_it(monkeypatch) -> None: + from flocks.session.runner import SessionRunner + + runner = object.__new__(SessionRunner) + runner.session = SimpleNamespace(id="session-question-unchecked") + runner._step = 1 + runner.callbacks = SimpleNamespace(event_publish_callback=None) + runner._turn_permission_ruleset = [] + agent = SimpleNamespace(name="custom-agent", tools=["request_ledger"]) + + result = SimpleNamespace( + tool_infos=[ + SimpleNamespace(name="question"), + SimpleNamespace(name="tool_search"), + SimpleNamespace(name="request_ledger"), + ], + metadata={}, + ) + + async def list_tools(**_kwargs): + return result + + monkeypatch.setattr( + "flocks.session.runner.list_session_callable_tool_infos", + list_tools, + ) + + messages = [ + SimpleNamespace( + role="user", + executionMode=SessionExecutionMode.BUILD, + ) + ] + + tools, metadata = await runner._list_callable_tool_infos_for_turn( + agent, + messages, + ) + + assert [tool.name for tool in tools] == [ + "question", + "tool_search", + "request_ledger", + ] + assert metadata["permissionDeniedToolNames"] == [] + + +@pytest.mark.asyncio +async def test_runner_keeps_non_question_tools_visible_when_permission_denied(monkeypatch) -> None: + from flocks.permission.helpers import from_config + from flocks.session.runner import SessionRunner + + runner = object.__new__(SessionRunner) + runner.session = SimpleNamespace(id="session-bash-deny") + runner._step = 1 + runner.callbacks = SimpleNamespace(event_publish_callback=None) + runner._turn_permission_ruleset = from_config({"bash": "deny"}) + agent = SimpleNamespace(name="custom-agent", tools=["question", "bash"]) + + result = SimpleNamespace( + tool_infos=[ + SimpleNamespace(name="question"), + SimpleNamespace(name="bash"), + ], + metadata={}, + ) + + async def list_tools(**_kwargs): + return result + + monkeypatch.setattr( + "flocks.session.runner.list_session_callable_tool_infos", + list_tools, + ) + + messages = [ + SimpleNamespace( + role="user", + executionMode=SessionExecutionMode.BUILD, + ) + ] + + tools, metadata = await runner._list_callable_tool_infos_for_turn( + agent, + messages, + ) + + assert [tool.name for tool in tools] == ["question", "bash"] + assert metadata["permissionDeniedToolNames"] == [] + + +@pytest.mark.asyncio +async def test_question_preflight_ask_permission_does_not_prompt() -> None: + from flocks.permission.helpers import from_config + from flocks.session.runner import SessionRunner + + runner = object.__new__(SessionRunner) + runner.session = SimpleNamespace(id="session-question-ask") + runner._step = 1 + runner.callbacks = SimpleNamespace( + event_publish_callback=None, + on_permission_request=AsyncMock(return_value=False), + ) + runner._turn_permission_ruleset = from_config({"question": "ask"}) + + await runner._handle_permission(SimpleNamespace( + permission="question", + patterns=["*"], + always=[], + metadata={"reason": "question_tool"}, + message_id="msg_1", + )) + + runner.callbacks.on_permission_request.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_question_preflight_deny_permission_still_blocks() -> None: + from flocks.permission.helpers import from_config + from flocks.session.runner import SessionRunner + + runner = object.__new__(SessionRunner) + runner.session = SimpleNamespace(id="session-question-deny") + runner._step = 1 + runner.callbacks = SimpleNamespace( + event_publish_callback=None, + on_permission_request=AsyncMock(return_value=True), + ) + runner._turn_permission_ruleset = from_config({"question": "deny"}) + + with pytest.raises(PermissionError, match="Permission denied: question"): + await runner._handle_permission(SimpleNamespace( + permission="question", + patterns=["*"], + always=[], + metadata={"reason": "question_tool"}, + message_id="msg_1", + )) + + runner.callbacks.on_permission_request.assert_not_awaited() 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_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_retry.py b/tests/session/test_retry.py index 8fc07d29e..78ac43abd 100644 --- a/tests/session/test_retry.py +++ b/tests/session/test_retry.py @@ -50,6 +50,14 @@ def test_api_error_retryable_generic(self): result = SessionRetry.retryable(error) assert result == "Internal server error" + def test_stream_tool_arguments_truncated_error_is_retryable(self): + error = { + "name": "StreamToolArgumentsTruncatedError", + "data": {"message": "tool arguments were truncated"}, + } + result = SessionRetry.retryable(error) + assert result == "Model output was truncated while generating tool arguments" + def test_json_message_too_many_requests(self): import json msg = json.dumps({"type": "error", "error": {"type": "too_many_requests"}}) diff --git a/tests/session/test_runner_step.py b/tests/session/test_runner_step.py index efee61f53..849bbccd0 100644 --- a/tests/session/test_runner_step.py +++ b/tests/session/test_runner_step.py @@ -33,7 +33,13 @@ StepResult, ToolCall, ) -from flocks.session.prompt import SessionPrompt, get_prompt_flocks_config_guard +from flocks.session.streaming.tool_accumulator import StreamToolArgumentsTruncatedError +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 @@ -261,6 +267,24 @@ def test_connection_error_exception_is_retryable(self): assert result["data"]["isRetryable"] is True assert result["data"]["displayMessage"] == runner_mod.CONNECTION_ERROR_DISPLAY_MESSAGE + def test_stream_tool_arguments_truncated_exception_is_retryable(self): + runner = _make_runner() + exc = StreamToolArgumentsTruncatedError( + tool_call_id="call_trunc", + tool_name="write", + finish_reason="length", + arguments_len=42, + arguments_preview='{"path":', + ) + + result = runner._exception_to_error_dict(exc) + + assert result["name"] == "StreamToolArgumentsTruncatedError" + assert result["data"]["isRetryable"] is True + assert result["data"]["streamToolArgumentsTruncated"] is True + assert result["data"]["toolCallID"] == "call_trunc" + assert result["data"]["toolName"] == "write" + def test_incomplete_chunked_read_exception_is_retryable_connection_error(self): runner = _make_runner() exc = Exception( @@ -629,6 +653,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 = {} @@ -641,14 +706,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 +729,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 +740,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 +763,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 +787,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 +814,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 +831,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 +845,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 +864,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 +883,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 +895,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 +917,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 +928,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 +937,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 +949,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 +963,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 +981,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 +998,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 +1011,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 +1025,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 +1037,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 +1171,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 +1192,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 +1204,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 +1540,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 +2320,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 +2378,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")) @@ -2322,6 +2396,212 @@ async def fake_to_chat_messages(_messages, _system_prompts): # noqa: ANN001 assert result.content == "done" +@pytest.mark.asyncio +async def test_process_step_retries_truncated_tool_arguments_with_fresh_message(monkeypatch): + runner = _make_runner("ses_runner_stream_tool_args_retry") + events = [] + + async def capture_event(event_type, data): + events.append((event_type, data)) + + runner.callbacks = RunnerCallbacks( + on_error=AsyncMock(), + event_publish_callback=capture_event, + ) + + last_user = UserMessageInfo( + id="msg_user_stream_tool_args_retry", + sessionID=runner.session.id, + role="user", + time={"created": 1_000}, + agent="rex", + model={"providerID": "anthropic", "modelID": "claude-sonnet"}, + ) + agent = SimpleNamespace(name="rex", steps=None, mode="primary", prompt="", tools=[]) + provider = MagicMock() + provider.is_configured.return_value = True + assistant_1 = SimpleNamespace(id="msg_assistant_truncated_1") + assistant_2 = SimpleNamespace(id="msg_assistant_truncated_2") + create_mock = AsyncMock(side_effect=[assistant_1, assistant_2]) + delete_mock = AsyncMock(return_value=True) + update_mock = AsyncMock(return_value=None) + call_ids = [] + + async def fake_call_llm(*_args, **kwargs): + call_ids.append(kwargs["assistant_msg"].id) + if len(call_ids) == 1: + runner._attempt_state.observable_output_started = True + raise StreamToolArgumentsTruncatedError( + tool_call_id="call_trunc", + tool_name="write", + finish_reason="length", + arguments_len=42, + arguments_preview='{"path":', + ) + assert runner._attempt_state.observable_output_started is False + return StepResult(action="stop", content="done") + + 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_prompt_blocks", AsyncMock(return_value=[])) + monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) + monkeypatch.setattr( + runner, + "_to_chat_messages", + AsyncMock(return_value=[SimpleNamespace(role="user", content="hi")]), + ) + monkeypatch.setattr(runner_mod.Message, "get_text_content", AsyncMock(return_value="hi")) + monkeypatch.setattr(runner_mod.Message, "parts", AsyncMock(return_value=[])) + monkeypatch.setattr(runner_mod.Message, "create", create_mock) + monkeypatch.setattr(runner_mod.Message, "delete", delete_mock) + monkeypatch.setattr(runner_mod.Message, "update", update_mock) + monkeypatch.setattr(runner_mod.SessionRetry, "sleep", AsyncMock(return_value=None)) + monkeypatch.setattr(runner, "_call_llm", fake_call_llm) + + result = await runner._process_step([last_user], last_user) + + assert result.action == "stop" + assert result.content == "done" + assert call_ids == [assistant_1.id, assistant_2.id] + delete_mock.assert_awaited_once_with(runner.session.id, assistant_1.id) + runner.callbacks.on_error.assert_not_awaited() + assert ("message.removed", {"sessionID": runner.session.id, "messageID": assistant_1.id}) in events + assert events[-1][0] == "message.updated" + assert events[-1][1]["info"]["id"] == assistant_2.id + assert update_mock.await_args_list[-1].args[1] == assistant_2.id + assert update_mock.await_args_list[-1].kwargs["finish"] == "stop" + + +@pytest.mark.asyncio +async def test_process_step_does_not_retry_truncated_tool_arguments_after_tool_started(monkeypatch): + runner = _make_runner("ses_runner_stream_tool_args_no_retry_after_tool") + runner.callbacks = RunnerCallbacks(on_error=AsyncMock()) + + last_user = UserMessageInfo( + id="msg_user_stream_tool_args_no_retry_after_tool", + sessionID=runner.session.id, + role="user", + time={"created": 1_000}, + agent="rex", + model={"providerID": "anthropic", "modelID": "claude-sonnet"}, + ) + agent = SimpleNamespace(name="rex", steps=None, mode="primary", prompt="", tools=[]) + provider = MagicMock() + provider.is_configured.return_value = True + assistant = SimpleNamespace(id="msg_assistant_no_retry_after_tool") + update_mock = AsyncMock(return_value=None) + delete_mock = AsyncMock(return_value=True) + call_count = 0 + + async def fake_call_llm(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + runner._attempt_state.observable_output_started = True + runner._attempt_state.tool_execution_started = True + raise StreamToolArgumentsTruncatedError( + tool_call_id="call_trunc", + tool_name="write", + finish_reason="length", + arguments_len=42, + arguments_preview='{"path":', + ) + + 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_prompt_blocks", AsyncMock(return_value=[])) + monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) + monkeypatch.setattr( + runner, + "_to_chat_messages", + AsyncMock(return_value=[SimpleNamespace(role="user", content="hi")]), + ) + monkeypatch.setattr(runner_mod.Message, "get_text_content", AsyncMock(return_value="hi")) + monkeypatch.setattr(runner_mod.Message, "parts", AsyncMock(return_value=[])) + monkeypatch.setattr(runner_mod.Message, "store_part", AsyncMock(return_value=None)) + monkeypatch.setattr(runner_mod.Message, "create", AsyncMock(return_value=assistant)) + monkeypatch.setattr(runner_mod.Message, "delete", delete_mock) + monkeypatch.setattr(runner_mod.Message, "update", update_mock) + monkeypatch.setattr(runner_mod.SessionRetry, "sleep", AsyncMock(return_value=None)) + monkeypatch.setattr(runner, "_call_llm", fake_call_llm) + + result = await runner._process_step([last_user], last_user) + + assert call_count == 1 + assert result.action == "stop" + assert "truncated" in result.error.lower() + delete_mock.assert_not_awaited() + runner.callbacks.on_error.assert_awaited_once() + assert update_mock.await_args_list[-1].args[1] == assistant.id + assert update_mock.await_args_list[-1].kwargs["finish"] == "error" + + +@pytest.mark.asyncio +async def test_process_step_allows_fallback_after_truncated_tool_argument_retries_exhausted(monkeypatch): + runner = SessionRunner( + session=_make_session("ses_runner_stream_tool_args_fallback"), + provider_id="primary", + model_id="primary-model", + defer_step_errors=True, + failover_available=True, + ) + + last_user = UserMessageInfo( + id="msg_user_stream_tool_args_fallback", + sessionID=runner.session.id, + role="user", + time={"created": 1_000}, + agent="rex", + model={"providerID": "primary", "modelID": "primary-model"}, + ) + agent = SimpleNamespace(name="rex", steps=None, mode="primary", prompt="", tools=[]) + provider = MagicMock() + provider.is_configured.return_value = True + create_count = 0 + + async def create_message(**_kwargs): + nonlocal create_count + create_count += 1 + return SimpleNamespace(id=f"msg_assistant_fallback_{create_count}") + + async def fake_call_llm(*_args, **_kwargs): + runner._attempt_state.observable_output_started = True + raise StreamToolArgumentsTruncatedError( + tool_call_id="call_trunc", + tool_name="write", + finish_reason="length", + arguments_len=42, + arguments_preview='{"path":', + ) + + 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_prompt_blocks", AsyncMock(return_value=[])) + monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) + monkeypatch.setattr( + runner, + "_to_chat_messages", + AsyncMock(return_value=[SimpleNamespace(role="user", content="hi")]), + ) + monkeypatch.setattr(runner_mod.Message, "get_text_content", AsyncMock(return_value="hi")) + monkeypatch.setattr(runner_mod.Message, "parts", AsyncMock(return_value=[])) + monkeypatch.setattr(runner_mod.Message, "create", create_message) + monkeypatch.setattr(runner_mod.Message, "delete", AsyncMock(return_value=True)) + monkeypatch.setattr(runner_mod.Message, "update", AsyncMock(return_value=None)) + monkeypatch.setattr(runner_mod.SessionRetry, "sleep", AsyncMock(return_value=None)) + monkeypatch.setattr(runner, "_call_llm", fake_call_llm) + + result = await runner._process_step([last_user], last_user) + + assert result.failure is not None + assert result.failure.reason == "stream_truncated" + assert result.failure.allow_fallback is True + assert result.failure.attempt_state.observable_output_started is True + assert result.failure.attempt_state.tool_execution_started is False + + @pytest.mark.asyncio async def test_process_step_limits_connection_error_retries(monkeypatch): runner = _make_runner("ses_runner_connection_error") @@ -2352,7 +2632,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 +2708,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 +2941,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 +2988,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 +2997,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 +3017,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 +3050,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 +3075,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 +3092,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 +3120,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 +3155,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 +3217,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 +3269,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 +3317,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 +3366,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 +3420,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_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/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) diff --git a/tests/session/test_stream_processor.py b/tests/session/test_stream_processor.py index 62b21e51b..b770f579c 100644 --- a/tests/session/test_stream_processor.py +++ b/tests/session/test_stream_processor.py @@ -28,6 +28,7 @@ TextEndEvent, TextStartEvent, ToolCallEvent, + ToolInputErrorEvent, ToolInputStartEvent, ) from flocks.session.message import MessageRole, ToolStateError @@ -495,6 +496,46 @@ async def test_tool_call_state_created(self): assert "tc_002" in proc.tool_calls assert proc.tool_calls["tc_002"].name == "read_file" + @pytest.mark.asyncio + async def test_tool_input_error_updates_pending_part_without_execution(self): + event_callback = AsyncMock() + proc = _make_processor(event_callback=event_callback) + execute_mock = AsyncMock(return_value=ToolResult(success=True, output="should not run")) + + with ( + patch("flocks.session.streaming.stream_processor.Message.store_part", new=AsyncMock()) as mock_store, + patch( + "flocks.session.streaming.stream_processor.ToolRegistry.execute", + new=execute_mock, + ), + ): + await proc.process_event(ToolInputStartEvent(id="tc_trunc", tool_name="write")) + await proc.process_event( + ToolInputErrorEvent( + id="tc_trunc", + tool_name="write", + input={"arguments_preview": '{"path": "/tmp/f"', "finish_reason": "length"}, + error="Output was truncated while generating tool arguments.", + ) + ) + + execute_mock.assert_not_awaited() + state = proc.tool_calls["tc_trunc"] + assert state.status == "error" + assert state.name == "write" + assert "truncated" in state.error + + completed_part = mock_store.await_args_list[-1].args[2] + assert completed_part.tool == "write" + assert completed_part.state.status == "error" + assert completed_part.state.input["finish_reason"] == "length" + assert "truncated" in completed_part.state.error + + published_part = event_callback.await_args_list[-1].args[1]["part"] + assert published_part["tool"] == "write" + assert published_part["state"]["status"] == "error" + assert published_part["state"]["input"]["finish_reason"] == "length" + # --------------------------------------------------------------------------- # ToolCall execution diff --git a/tests/session/test_tool_accumulator.py b/tests/session/test_tool_accumulator.py index 0ef761f6f..ad90eb162 100644 --- a/tests/session/test_tool_accumulator.py +++ b/tests/session/test_tool_accumulator.py @@ -15,8 +15,14 @@ from unittest.mock import AsyncMock, MagicMock, patch -from flocks.session.streaming.tool_accumulator import ToolCallAccumulator -from flocks.session.streaming.stream_events import ToolInputStartEvent, ToolCallEvent +from flocks.session.streaming.tool_accumulator import ( + StreamToolArgumentsTruncatedError, + ToolCallAccumulator, +) +from flocks.session.streaming.stream_events import ( + ToolInputErrorEvent, + ToolInputStartEvent, +) # --------------------------------------------------------------------------- @@ -236,7 +242,7 @@ async def test_flush_invalid_json_sends_invalid_tool(self): assert any(e.tool_name == "invalid" for e in tool_call_events) @pytest.mark.asyncio - async def test_flush_with_length_finish_reason_mentions_truncated(self): + async def test_flush_with_length_finish_reason_marks_input_error_and_raises(self): acc, proc = _make_accumulator() with patch("flocks.session.streaming.tool_accumulator.ToolRegistry") as mock_reg: mock_reg.get_schema.return_value = None @@ -248,15 +254,167 @@ async def test_flush_with_length_finish_reason_mentions_truncated(self): "arguments_str": '{"path": "/tmp/f", "content": "abc', "completed": False, } - await acc.flush_remaining(stream_finish_reason="length") + with pytest.raises(StreamToolArgumentsTruncatedError) as exc_info: + await acc.flush_remaining(stream_finish_reason="length") + + assert exc_info.value.tool_call_id == "call_trunc" + assert exc_info.value.tool_name == "write_file" + assert exc_info.value.finish_reason == "length" + + input_error_events = [ + c.args[0] for c in proc.process_event.call_args_list + if c.args[0].type == "tool-input-error" + ] + assert len(input_error_events) == 1 + assert isinstance(input_error_events[0], ToolInputErrorEvent) + assert input_error_events[0].tool_name == "write_file" + assert "truncated" in input_error_events[0].error.lower() tool_call_events = [ c.args[0] for c in proc.process_event.call_args_list - if c.args[0].type == "tool-call" and c.args[0].tool_name == "invalid" + if c.args[0].type == "tool-call" + ] + assert tool_call_events == [] + + @pytest.mark.asyncio + async def test_flush_with_truncated_repairable_json_does_not_execute_tool(self): + acc, proc = _make_accumulator() + with patch("flocks.session.streaming.tool_accumulator.ToolRegistry") as mock_reg: + mock_reg.get_schema.return_value = None + mock_reg.get.return_value = MagicMock() + acc._accumulator["call_repairable"] = { + "id": "call_repairable", + "name": "write_file", + "arguments_str": '{"path": "/tmp/f", "content": "abc', + "completed": False, + } + + with pytest.raises(StreamToolArgumentsTruncatedError): + await acc.flush_remaining(stream_finish_reason="max_tokens") + + event_types = [c.args[0].type for c in proc.process_event.call_args_list] + assert "tool-input-error" in event_types + assert "tool-call" not in event_types + + @pytest.mark.asyncio + async def test_flush_with_truncated_valid_json_missing_required_does_not_execute_tool(self): + acc, proc = _make_accumulator() + schema = MagicMock() + schema.required = ["path", "content"] + + with patch("flocks.session.streaming.tool_accumulator.ToolRegistry") as mock_reg: + mock_reg.get_schema.return_value = schema + mock_reg.get.return_value = MagicMock() + + await acc.feed_chunk(_make_chunk( + tc_id="call_missing_required", + name="write_file", + arguments='{"path": "/tmp/f"}', + )) + + with pytest.raises(StreamToolArgumentsTruncatedError) as exc_info: + await acc.flush_remaining(stream_finish_reason="length") + + assert exc_info.value.tool_call_id == "call_missing_required" + event_types = [c.args[0].type for c in proc.process_event.call_args_list] + assert event_types == ["tool-input-start", "tool-input-error"] + + @pytest.mark.asyncio + async def test_flush_with_truncated_name_only_tool_marks_input_error(self): + acc, proc = _make_accumulator() + with patch("flocks.session.streaming.tool_accumulator.ToolRegistry") as mock_reg: + mock_reg.get_schema.return_value = None + + await acc.feed_chunk(_make_chunk( + tc_id="call_name_only", + name="write_file", + )) + + with pytest.raises(StreamToolArgumentsTruncatedError) as exc_info: + await acc.flush_remaining(stream_finish_reason="max_tokens") + + assert exc_info.value.tool_call_id == "call_name_only" + events = [c.args[0] for c in proc.process_event.call_args_list] + assert [event.type for event in events] == [ + "tool-input-start", + "tool-input-error", + ] + assert events[-1].input["arguments_preview"] == "" + + @pytest.mark.asyncio + async def test_flush_with_multiple_truncated_tools_marks_all_errors(self): + acc, proc = _make_accumulator() + with patch("flocks.session.streaming.tool_accumulator.ToolRegistry") as mock_reg: + mock_reg.get_schema.return_value = None + mock_reg.get.return_value = MagicMock() + + await acc.feed_chunk(_make_chunk( + index=0, + tc_id="call_done", + name="tool_a", + arguments='{"value": 1}', + )) + await acc.feed_chunk(_make_chunk( + index=1, + tc_id="call_trunc_b", + name="tool_b", + arguments='{"value": "unfinished', + )) + await acc.feed_chunk(_make_chunk( + index=2, + tc_id="call_trunc_c", + name="tool_c", + arguments='{"value": "also unfinished', + )) + + with pytest.raises(StreamToolArgumentsTruncatedError) as exc_info: + await acc.flush_remaining(stream_finish_reason="length") + + assert exc_info.value.tool_call_id == "call_trunc_b" + + events = [c.args[0] for c in proc.process_event.call_args_list] + tool_call_events = [event for event in events if event.type == "tool-call"] + input_error_events = [ + event for event in events if event.type == "tool-input-error" + ] + + assert [event.tool_call_id for event in tool_call_events] == ["call_done"] + assert [event.id for event in input_error_events] == [ + "call_trunc_b", + "call_trunc_c", + ] + assert all("truncated" in event.error.lower() for event in input_error_events) + + @pytest.mark.asyncio + async def test_flush_with_truncation_does_not_execute_later_valid_pending_tool(self): + acc, proc = _make_accumulator() + with patch("flocks.session.streaming.tool_accumulator.ToolRegistry") as mock_reg: + mock_reg.get_schema.return_value = None + mock_reg.get.return_value = MagicMock() + acc._accumulator["call_trunc"] = { + "id": "call_trunc", + "name": "tool_b", + "arguments_str": '{"value": "unfinished', + "completed": False, + "input_started": True, + } + acc._accumulator["call_valid_pending"] = { + "id": "call_valid_pending", + "name": "tool_c", + "arguments_str": '{"value": 1}', + "completed": False, + "input_started": True, + } + + with pytest.raises(StreamToolArgumentsTruncatedError): + await acc.flush_remaining(stream_finish_reason="length") + + events = [c.args[0] for c in proc.process_event.call_args_list] + assert [event.type for event in events] == [ + "tool-input-error", + "tool-input-error", ] - if tool_call_events: - error_msg = tool_call_events[0].input.get("error", "") - assert "truncated" in error_msg.lower() or "length" in error_msg.lower() + assert [event.id for event in events] == ["call_trunc", "call_valid_pending"] @pytest.mark.asyncio async def test_flush_empty_accumulator_does_nothing(self): 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_memory_file_write.py b/tests/tool/test_memory_file_write.py index 42d62e4bb..9f69df106 100644 --- a/tests/tool/test_memory_file_write.py +++ b/tests/tool/test_memory_file_write.py @@ -157,8 +157,7 @@ async def approve(_request) -> None: result = await edit_tool( ctx, filePath=str(memory_path), - oldString="existing", - newString="replacement", + edits=[{"oldString": "existing", "newString": "replacement"}], ) assert result.success is False diff --git a/tests/tool/test_question_channel.py b/tests/tool/test_question_channel.py index 963335f24..2d25d975f 100644 --- a/tests/tool/test_question_channel.py +++ b/tests/tool/test_question_channel.py @@ -106,6 +106,35 @@ async def fake_handler(_session_id: str, questions: list[dict]) -> list[list[str assert captured_questions[0]["custom"] is False +@pytest.mark.asyncio +async def test_question_tool_respects_permission_denial_before_waiting( + monkeypatch: pytest.MonkeyPatch, +) -> None: + handler = AsyncMock(return_value=[["yes"]]) + channel_send = AsyncMock(return_value=None) + + async def deny_question(_request): + raise PermissionError("Permission denied: question") + + monkeypatch.setattr(question_module, "_question_handler", handler) + monkeypatch.setattr(question_module, "_send_channel_question_if_applicable", channel_send) + + result = await question_module.question_tool( + ToolContext( + session_id="ses_question_denied", + message_id="msg_1", + call_id="call_1", + permission_callback=deny_question, + ), + questions=[{"question": "Continue?", "type": "confirm"}], + ) + + assert result.success is False + assert "Permission denied" in (result.error or "") + handler.assert_not_awaited() + channel_send.assert_not_awaited() + + @pytest.mark.asyncio async def test_question_tool_sends_plain_text_for_channel_session() -> None: binding = SimpleNamespace( 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..32457de57 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", ] @@ -329,8 +329,7 @@ async def test_edit_string_replacement(self, tool_context, temp_dir): "edit", ctx=tool_context, filePath=filepath, - oldString="Hello", - newString="Hi" + edits=[{"oldString": "Hello", "newString": "Hi"}], ) assert result.success @@ -351,8 +350,12 @@ async def test_edit_multiline_replacement(self, tool_context, temp_dir): "edit", ctx=tool_context, filePath=filepath, - oldString="def foo():\n return 1", - newString="def foo():\n return 42" + edits=[ + { + "oldString": "def foo():\n return 1", + "newString": "def foo():\n return 42", + } + ], ) assert result.success @@ -362,9 +365,9 @@ async def test_edit_multiline_replacement(self, tool_context, temp_dir): assert "return 42" in content @pytest.mark.asyncio - async def test_edit_replace_all(self, tool_context, temp_dir): - """Test replace all occurrences""" - filepath = os.path.join(temp_dir, "edit_replace_all.txt") + async def test_edit_rejects_non_unique_target(self, tool_context, temp_dir): + """Test that each targeted replacement must be unique.""" + filepath = os.path.join(temp_dir, "edit_duplicate.txt") with open(filepath, 'w') as f: f.write("foo bar foo baz foo\n") @@ -372,32 +375,50 @@ async def test_edit_replace_all(self, tool_context, temp_dir): "edit", ctx=tool_context, filePath=filepath, - oldString="foo", - newString="qux", - replaceAll=True + edits=[{"oldString": "foo", "newString": "qux"}], ) - - assert result.success - + + assert not result.success + assert "must be unique" in (result.error or "") + with open(filepath, 'r') as f: - content = f.read() - assert "foo" not in content - assert content.count("qux") == 3 + assert f.read() == "foo bar foo baz foo\n" def test_edit_schema_supports_batch_edits(self): - """Test edit schema exposes pi-style edits[] plus legacy compatibility.""" + """Test edit schema exposes only the required edits[] interface.""" tool = ToolRegistry.get("edit") assert tool is not None schema = tool.info.get_schema() edits_prop = schema.properties["edits"] + assert set(schema.properties) == {"filePath", "edits"} assert edits_prop["type"] == "array" + assert edits_prop["minItems"] == 1 assert edits_prop["items"]["type"] == "object" assert edits_prop["items"]["required"] == ["oldString", "newString"] assert "oldString" in edits_prop["items"]["properties"] assert "newString" in edits_prop["items"]["properties"] - assert schema.required == ["filePath"] + assert schema.required == ["filePath", "edits"] + + @pytest.mark.asyncio + async def test_edit_rejects_top_level_replacement_fields(self, tool_context, temp_dir): + """Test that replacements must be supplied through edits[].""" + filepath = os.path.join(temp_dir, "edit_invalid_shape.txt") + with open(filepath, "w", encoding="utf-8") as f: + f.write("Hello\n") + + result = await ToolRegistry.execute( + "edit", + ctx=tool_context, + filePath=filepath, + oldString="Hello", + newString="Hi", + ) + + assert not result.success + assert "unknown parameters" in (result.error or "") + assert "edits" in (result.error or "") class TestBashTool: @@ -801,13 +822,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 @@ -1447,8 +1468,7 @@ async def test_edit_not_found(self, tool_context, temp_dir): "edit", ctx=tool_context, filePath=filepath, - oldString="nonexistent", - newString="replacement" + edits=[{"oldString": "nonexistent", "newString": "replacement"}], ) assert not result.success @@ -1457,24 +1477,20 @@ async def test_edit_not_found(self, tool_context, temp_dir): assert "slightly larger unique snippet" in result.error.lower() @pytest.mark.asyncio - async def test_edit_create_new_file(self, tool_context, temp_dir): - """Test edit creates new file with empty oldString""" + async def test_edit_requires_existing_file(self, tool_context, temp_dir): + """Test that edit does not create a missing file.""" filepath = os.path.join(temp_dir, "new_edit_file.txt") - content = "New file content\n" result = await ToolRegistry.execute( "edit", ctx=tool_context, filePath=filepath, - oldString="", - newString=content + edits=[{"oldString": "missing", "newString": "replacement"}], ) - assert result.success - assert os.path.exists(filepath) - - with open(filepath, 'r') as f: - assert f.read() == content + assert not result.success + assert "not found" in (result.error or "").lower() + assert not os.path.exists(filepath) @pytest.mark.asyncio async def test_edit_multi_snapshot_semantics(self, tool_context, temp_dir): @@ -1573,8 +1589,7 @@ async def test_edit_fuzzy_match_preserves_unedited_unicode(self, tool_context, t "edit", ctx=tool_context, filePath=filepath, - oldString="Don't stop", - newString="Do not stop", + edits=[{"oldString": "Don't stop", "newString": "Do not stop"}], ) assert result.success 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_execution_store_compact.py b/tests/workflow/test_execution_store_compact.py index f5d02a555..ea8c3d5cd 100644 --- a/tests/workflow/test_execution_store_compact.py +++ b/tests/workflow/test_execution_store_compact.py @@ -16,6 +16,9 @@ """ from __future__ import annotations + +import asyncio +from types import SimpleNamespace from typing import Any, Dict, List from unittest.mock import AsyncMock, patch @@ -25,7 +28,8 @@ DEFAULT_GENERIC_SEQUENCE_THRESHOLD, DEFAULT_LARGE_LIST_KEYS, DEFAULT_MAX_INLINE_COLLECTION_BYTES, - _trim_execution_history, + ExecutionProgressWriter, + ExecutionStepRecorder, compact_history_for_storage, compact_execution_summary, compact_outputs_for_storage, @@ -40,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 ─────────────────────────────────────────────── @@ -299,11 +308,153 @@ 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" +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() + + 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_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: - record_step = AsyncMock(return_value=None) - upsert_execution = AsyncMock(return_value=None) - update_stats = AsyncMock(return_value=None) + 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", @@ -315,29 +466,139 @@ async def test_record_execution_result_backfills_execution_log_steps() -> None: ], } - def raise_create_task(coro, *args, **kwargs): # noqa: ANN001, ARG001 - coro.close() - 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("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)), + 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) - 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] + 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 + 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=None) + increment_stats = AsyncMock(return_value=None) + trim_executions = AsyncMock(return_value=[]) + record_audit = AsyncMock(return_value=None) + explicit_steps = [ + (2, {"node_id": "step-2", "outputs": {"ok": True}}), + (1, {"node_id": "step-1", "outputs": {"ok": True}}), + ] + exec_data = { + "id": "exec-trigger", + "workflowId": "wf-trigger", + "status": "success", + "duration": 0.01, + "executionLog": [], + "stepCount": 2, + } + + 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), + ): + await record_execution_result( + "wf-trigger", + "exec-trigger", + exec_data, + steps=explicit_steps, + ) + + summary, persisted_steps = complete_execution.await_args.args + assert complete_execution.await_args.kwargs == {} + assert summary["executionLog"] == [] + assert persisted_steps == explicit_steps + 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=[]) + + 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")) + + 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: @@ -425,62 +686,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"]) - - 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"]) +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), + 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() - - -@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) + 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) + + 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_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/tests/workflow/test_poller_manager.py b/tests/workflow/test_poller_manager.py index 7da394808..cb10353ec 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 @@ -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 { @@ -173,17 +172,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 +236,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 +248,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[0][0] == "exec-1" - assert recorded_steps[0][1] == 1 - assert recorded_steps[0][2]["node_id"] == "load" + 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 @@ -372,8 +373,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 *, @@ -409,7 +412,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 @@ -424,7 +430,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_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 dabf9bb7e..7f05ce780 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,27 @@ 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"]) + create_record.assert_awaited_once_with( + "wf-trigger", + input_params={"message": "hello"}, + ) + assert result["executionLog"] == [] + assert result["stepCount"] == 1 + record_result.assert_awaited_once() + 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 bca4f3ada..b33234648 100644 --- a/tests/workflow/test_workflow_store.py +++ b/tests/workflow/test_workflow_store.py @@ -1,7 +1,9 @@ from __future__ import annotations import asyncio +import os from pathlib import Path +from unittest.mock import AsyncMock import pytest @@ -18,8 +20,10 @@ 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 @pytest.fixture(autouse=True) @@ -74,8 +78,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 +117,207 @@ 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_and_summary_with_one_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) + + 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}}), + ], + ) + + 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"] + assert await WorkflowStore.get_stats("wf-complete") is None + + +@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_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) + 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, + ) + 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) + ] + assert await WorkflowStore.get_stats("wf-trigger") is None + + +@pytest.mark.asyncio +async def test_complete_execution_rolls_back_partial_transaction( + monkeypatch: pytest.MonkeyPatch, +) -> None: + await WorkflowStore.init() + db = await WorkflowStore.raw_completion_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}})], + ) + + 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 + + +@pytest.mark.asyncio +async def test_complete_execution_rolls_back_cancelled_transaction( + monkeypatch: pytest.MonkeyPatch, +) -> None: + await WorkflowStore.init() + db = await WorkflowStore.raw_completion_db() + original_commit = db.commit + + async def cancel_commit() -> None: + raise asyncio.CancelledError + + monkeypatch.setattr(db, "commit", cancel_commit) + + with pytest.raises(asyncio.CancelledError): + await WorkflowStore.complete_execution( + { + "id": "exec-cancelled-commit", + "workflowId": "wf-cancelled-commit", + "status": "success", + "startedAt": 1, + "finishedAt": 2, + "executionLog": [], + "stepCount": 1, + }, + [(1, {"node_id": "node-1", "outputs": {"ok": True}})], + ) + + 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_pid_change_drops_inherited_connections_without_closing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + await WorkflowStore.init() + 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 + + 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() diff --git a/tests/workspace/test_workspace_models.py b/tests/workspace/test_workspace_models.py index 1b0a6da4c..91c0e3671 100644 --- a/tests/workspace/test_workspace_models.py +++ b/tests/workspace/test_workspace_models.py @@ -15,6 +15,7 @@ def test_file_node_defaults(self): assert node.size is None assert node.modified_at is None assert node.is_text_file is False + assert node.editable is False assert node.children is None def test_directory_node_with_children(self): @@ -37,10 +38,12 @@ def test_file_node_full_fields(self): size=204800, modified_at=1741900000.0, is_text_file=False, + editable=True, ) assert node.size == 204800 assert node.modified_at == pytest.approx(1741900000.0) assert node.is_text_file is False + assert node.editable is True def test_invalid_type_raises(self): with pytest.raises(Exception): diff --git a/tests/workspace/test_workspace_routes.py b/tests/workspace/test_workspace_routes.py index 80843bf1c..1fb938997 100644 --- a/tests/workspace/test_workspace_routes.py +++ b/tests/workspace/test_workspace_routes.py @@ -49,6 +49,7 @@ def workspace_client(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv("FLOCKS_WORKSPACE_DIR", str(ws)) monkeypatch.setenv("FLOCKS_DATA_DIR", str(data)) + monkeypatch.setenv("FLOCKS_ROOT", str(tmp_path / ".flocks")) # Reset both singletons so they re-read env vars from flocks.workspace.manager import WorkspaceManager @@ -58,9 +59,21 @@ def workspace_client(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): # Build a minimal FastAPI app with only the workspace router from fastapi import FastAPI + from flocks.auth.context import AuthUser from flocks.server.routes.workspace import router app = FastAPI() + + @app.middleware("http") + async def inject_auth_user(request, call_next): + request.state.auth_user = AuthUser( + id="usr_workspace", + username="workspace", + role="member", + status="active", + ) + return await call_next(request) + app.include_router(router, prefix="/api/workspace") client = TestClient(app, raise_server_exceptions=True) @@ -125,6 +138,33 @@ def test_list_file_path_returns_400(self, workspace_client): r = _client(workspace_client).get("/api/workspace/list?path=file.txt") assert r.status_code == 400 + def test_list_subdir_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 + + manager = WorkspaceManager.get_instance() + manager._workspace_dir = link + (ws / "outputs" / "report.txt").write_text("ok") + + root = _client(workspace_client).get("/api/workspace/list") + r = _client(workspace_client).get("/api/workspace/list?path=outputs") + + assert root.status_code == 200 + assert any(item["name"] == "outputs" and item["path"] == "outputs" for item in root.json()) + assert r.status_code == 200 + assert any( + item["name"] == "report.txt" + and item["path"] == "outputs/report.txt" + and item["type"] == "file" + for item in r.json() + ) + class TestDirTree: def test_tree_root(self, workspace_client): @@ -150,6 +190,30 @@ def test_tree_nonexistent_returns_404(self, workspace_client): r = _client(workspace_client).get("/api/workspace/tree?path=nope") assert r.status_code == 404 + def test_tree_subdir_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 + + manager = WorkspaceManager.get_instance() + manager._workspace_dir = link + (ws / "outputs" / "nested").mkdir() + + root = _client(workspace_client).get("/api/workspace/tree?depth=1") + r = _client(workspace_client).get("/api/workspace/tree?path=outputs&depth=1") + + assert root.status_code == 200 + assert root.json()["path"] == "" + assert r.status_code == 200 + data = r.json() + assert data["path"] == "outputs" + assert any(child["path"] == "outputs/nested" for child in data["children"]) + class TestDirCreate: def test_create_new_dir(self, workspace_client): @@ -697,10 +761,21 @@ def test_list_memory_without_files_returns_empty_daily_tree(self, workspace_clie assert daily["type"] == "directory" assert daily["children"] == [] - def test_list_memory_with_files(self, workspace_client): + def test_list_memory_with_files(self, workspace_client, monkeypatch): + from flocks.project.project import Project + + monkeypatch.setattr( + Project, + "registered_project_ids", + classmethod(lambda cls, owner_id: {"prj_example"}), + ) mem = _mem(workspace_client) (mem / "USER.md").write_text("# User") (mem / "MEMORY.md").write_text("# Memory") + (mem / "SHORT_MEMORY.md").write_text("# Short") + (mem / "bak.txt").write_text("backup") + (mem / "archive").mkdir() + (mem / "archive" / "2026-08-18.md").write_text("# Archived") (mem / "daily").mkdir() (mem / "daily" / "2026-03-14.md").write_text("## Daily") (mem / "projects" / "prj_example").mkdir(parents=True) @@ -716,17 +791,29 @@ def test_list_memory_with_files(self, workspace_client): "MEMORY.md", "projects", "daily", + "archive", + "bak.txt", + "SHORT_MEMORY.md", ] nodes = {node["name"]: node for node in r.json()} - assert set(nodes) == {"USER.md", "MEMORY.md", "daily", "projects"} + assert set(nodes) == {"USER.md", "MEMORY.md", "SHORT_MEMORY.md", "archive", "bak.txt", "daily", "projects"} assert nodes["USER.md"]["type"] == "file" + assert nodes["USER.md"]["editable"] is True assert nodes["MEMORY.md"]["type"] == "file" + assert nodes["MEMORY.md"]["editable"] is True + assert nodes["SHORT_MEMORY.md"]["editable"] is False + assert nodes["bak.txt"]["editable"] is False + assert nodes["archive"]["editable"] is False + assert nodes["archive"]["children"][0]["path"] == "archive/2026-08-18.md" + assert nodes["archive"]["children"][0]["editable"] is False assert nodes["daily"]["type"] == "directory" assert nodes["daily"]["children"][0]["path"] == "daily/2026-03-14.md" + assert nodes["daily"]["children"][0]["editable"] is True assert nodes["projects"]["type"] == "directory" project = nodes["projects"]["children"][0] assert project["path"] == "projects/prj_example" assert project["children"][0]["path"] == "projects/prj_example/MEMORY.md" + assert project["children"][0]["editable"] is True def test_read_memory_file(self, workspace_client): mem = _mem(workspace_client) @@ -834,6 +921,40 @@ def test_write_memory_file(self, workspace_client): } assert target.read_text() == "new content" + def test_write_non_editable_memory_file_returns_403(self, workspace_client): + mem = _mem(workspace_client) + target = mem / "SHORT_MEMORY.md" + target.write_text("old content") + + r = _client(workspace_client).put( + "/api/workspace/memory/file", + json={"path": "SHORT_MEMORY.md", "content": "tampered"}, + ) + + assert r.status_code == 403 + assert target.read_text() == "old content" + + def test_write_unreadable_project_memory_returns_403(self, workspace_client, monkeypatch): + from flocks.project.project import Project + + monkeypatch.setattr( + Project, + "registered_project_ids", + classmethod(lambda cls, owner_id: {"prj_allowed"}), + ) + mem = _mem(workspace_client) + target = mem / "projects" / "prj_stale" / "MEMORY.md" + target.parent.mkdir(parents=True) + target.write_text("old content") + + r = _client(workspace_client).put( + "/api/workspace/memory/file", + json={"path": "projects/prj_stale/MEMORY.md", "content": "tampered"}, + ) + + assert r.status_code == 403 + assert target.read_text() == "old content" + def test_write_memory_traversal_rejected(self, workspace_client): r = _client(workspace_client).put( "/api/workspace/memory/file", diff --git a/tui/flocks/acp/agent.ts b/tui/flocks/acp/agent.ts index 8440bbcae..0600bc214 100644 --- a/tui/flocks/acp/agent.ts +++ b/tui/flocks/acp/agent.ts @@ -61,6 +61,31 @@ function parseTodoEntries(rawOutput: string, rawMetadata: unknown): Todo.Info[] return undefined } +function editDiffText(rawInput: unknown, rawMetadata: unknown): { oldText: string; newText: string } { + if (rawMetadata && typeof rawMetadata === "object") { + const filediff = (rawMetadata as Record)["filediff"] + if (filediff && typeof filediff === "object") { + const value = filediff as Record + if (typeof value["before"] === "string" && typeof value["after"] === "string") { + return { oldText: value["before"], newText: value["after"] } + } + } + } + + if (rawInput && typeof rawInput === "object") { + const edits = (rawInput as Record)["edits"] + if (Array.isArray(edits) && edits.length === 1 && edits[0] && typeof edits[0] === "object") { + const edit = edits[0] as Record + return { + oldText: typeof edit["oldString"] === "string" ? edit["oldString"] : "", + newText: typeof edit["newString"] === "string" ? edit["newString"] : "", + } + } + } + + return { oldText: "", newText: "" } +} + export namespace ACP { const log = Log.create({ service: "acp-agent" }) @@ -285,13 +310,7 @@ export namespace ACP { if (kind === "edit") { const input = part.state.input const filePath = typeof input["filePath"] === "string" ? input["filePath"] : "" - const oldText = typeof input["oldString"] === "string" ? input["oldString"] : "" - const newText = - typeof input["newString"] === "string" - ? input["newString"] - : typeof input["content"] === "string" - ? input["content"] - : "" + const { oldText, newText } = editDiffText(input, part.state.metadata) content.push({ type: "diff", path: filePath, @@ -622,13 +641,7 @@ export namespace ACP { if (kind === "edit") { const input = part.state.input const filePath = typeof input["filePath"] === "string" ? input["filePath"] : "" - const oldText = typeof input["oldString"] === "string" ? input["oldString"] : "" - const newText = - typeof input["newString"] === "string" - ? input["newString"] - : typeof input["content"] === "string" - ? input["content"] - : "" + const { oldText, newText } = editDiffText(input, part.state.metadata) content.push({ type: "diff", path: filePath, 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..06059b7e6 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)" : ""} @@ -2086,7 +2086,7 @@ function Edit(props: ToolProps) { - Edit {normalizePath(props.input.filePath!)} {input({ replaceAll: props.input.replaceAll })} + Edit {normalizePath(props.input.filePath!)} 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("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 976f1cd51..4d02feeca 100644 --- a/tui/flocks/command/index.ts +++ b/tui/flocks/command/index.ts @@ -26,11 +26,11 @@ 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 template: z.promise(z.string()).or(z.string()), - subtask: z.boolean().optional(), hints: z.array(z.string()), }) .meta({ @@ -40,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) @@ -84,10 +88,10 @@ export namespace Command { agent: command.agent, model: command.model, description: command.description, + subtask: command.subtask, get template() { return command.template }, - subtask: command.subtask, hints: hints(command.template), } } diff --git a/tui/flocks/config/config.test.ts b/tui/flocks/config/config.test.ts new file mode 100644 index 000000000..3683ad309 --- /dev/null +++ b/tui/flocks/config/config.test.ts @@ -0,0 +1,41 @@ +import { describe, expect, test } from "bun:test" +import { Config } from "./config" + +describe("permission aliases", () => { + test("maps task to delegate_task with deny precedence", () => { + expect( + Config.Permission.parse({ + delegate_task: "allow", + task: { explore: "deny" }, + }), + ).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 5085330ed..c9b259860 100644 --- a/tui/flocks/config/config.ts +++ b/tui/flocks/config/config.ts @@ -501,16 +501,30 @@ 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 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 (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 => { @@ -519,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 @@ -536,7 +553,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(), 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 f2f3331e5..c9cbd3e61 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, @@ -346,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({ @@ -466,12 +466,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", - }) - } } } @@ -581,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 befddd369..18d267c59 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" @@ -43,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 @@ -200,16 +200,6 @@ export namespace SessionPrompt { .meta({ ref: "AgentPartInput", }), - MessageV2.SubtaskPart.omit({ - messageID: true, - sessionID: true, - }) - .partial({ - id: true, - }) - .meta({ - ref: "SubtaskPartInput", - }), ]), ), }) @@ -344,7 +334,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 +342,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 +368,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 +1016,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 +1035,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, }, @@ -1653,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 @@ -1742,30 +1726,8 @@ 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 ?? [])] + const isSubtask = Command.shouldDelegate(command, agent.mode) await Plugin.trigger( "command.execute.before", @@ -1777,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: userModel, - agent: userAgent, - 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, @@ -1817,15 +1809,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 +1834,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 deleted file mode 100644 index 676c4d8dc..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 the Task tool 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. -- 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 7a0e5fd5c..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 the Task tool 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. -- 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. - -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] - - -user: What is the codebase structure? -assistant: [Uses the Task tool] - - -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. - - 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/edit.ts b/tui/flocks/tool/edit.ts index 57124501a..daba4439b 100644 --- a/tui/flocks/tool/edit.ts +++ b/tui/flocks/tool/edit.ts @@ -27,17 +27,25 @@ export const EditTool = Tool.define("edit", { description: DESCRIPTION, parameters: z.object({ filePath: z.string().describe("The absolute path to the file to modify"), - oldString: z.string().describe("The text to replace"), - newString: z.string().describe("The text to replace it with (must be different from oldString)"), - replaceAll: z.boolean().optional().describe("Replace all occurrences of oldString (default false)"), + edits: z + .array( + z.object({ + oldString: z.string().min(1).describe("Unique text to replace in the original file"), + newString: z.string().describe("Replacement text, which must differ from oldString"), + }), + ) + .min(1) + .describe("One or more non-overlapping replacements matched against the original file"), }), async execute(params, ctx) { if (!params.filePath) { throw new Error("filePath is required") } - if (params.oldString === params.newString) { - throw new Error("oldString and newString must be different") + for (const [index, edit] of params.edits.entries()) { + if (edit.oldString === edit.newString) { + throw new Error(`edits[${index}].oldString and newString must be different`) + } } const filePath = path.isAbsolute(params.filePath) ? params.filePath : path.join(Instance.directory, params.filePath) @@ -47,33 +55,13 @@ export const EditTool = Tool.define("edit", { let contentOld = "" let contentNew = "" await FileTime.withLock(filePath, async () => { - if (params.oldString === "") { - contentNew = params.newString - diff = trimDiff(createTwoFilesPatch(filePath, filePath, contentOld, contentNew)) - await ctx.ask({ - permission: "edit", - patterns: [path.relative(Instance.worktree, filePath)], - always: ["*"], - metadata: { - filepath: filePath, - diff, - }, - }) - await Bun.write(filePath, params.newString) - await Bus.publish(File.Event.Edited, { - file: filePath, - }) - FileTime.read(ctx.sessionID, filePath) - return - } - const file = Bun.file(filePath) const stats = await file.stat().catch(() => {}) if (!stats) throw new Error(`File ${filePath} not found`) if (stats.isDirectory()) throw new Error(`Path is a directory, not a file: ${filePath}`) await FileTime.assert(ctx.sessionID, filePath) contentOld = await file.text() - contentNew = replace(contentOld, params.oldString, params.newString, params.replaceAll) + contentNew = replaceEdits(contentOld, params.edits) diff = trimDiff( createTwoFilesPatch(filePath, filePath, normalizeLineEndings(contentOld), normalizeLineEndings(contentNew)), @@ -472,8 +460,7 @@ export const EscapeNormalizedReplacer: Replacer = function* (content, find) { } export const MultiOccurrenceReplacer: Replacer = function* (content, find) { - // This replacer yields all exact matches, allowing the replace function - // to handle multiple occurrences based on replaceAll parameter + // Yield exact matches so the caller can reject ambiguous replacement targets. let startIndex = 0 while (true) { @@ -605,9 +592,9 @@ export function trimDiff(diff: string): string { return trimmedLines.join("\n") } -export function replace(content: string, oldString: string, newString: string, replaceAll = false): string { - if (oldString === newString) { - throw new Error("oldString and newString must be different") +function findUniqueMatch(content: string, oldString: string): { index: number; search: string } { + if (oldString === "") { + throw new Error("oldString must not be empty") } let notFound = true @@ -627,12 +614,8 @@ export function replace(content: string, oldString: string, newString: string, r const index = content.indexOf(search) if (index === -1) continue notFound = false - if (replaceAll) { - return content.replaceAll(search, newString) - } - const lastIndex = content.lastIndexOf(search) - if (index !== lastIndex) continue - return content.substring(0, index) + newString + content.substring(index + search.length) + if (index !== content.lastIndexOf(search)) continue + return { index, search } } } @@ -643,3 +626,51 @@ export function replace(content: string, oldString: string, newString: string, r "Found multiple matches for oldString. Provide more surrounding lines in oldString to identify the correct match.", ) } + +export function replace(content: string, oldString: string, newString: string): string { + if (oldString === newString) { + throw new Error("oldString and newString must be different") + } + + const match = findUniqueMatch(content, oldString) + return content.substring(0, match.index) + newString + content.substring(match.index + match.search.length) +} + +export function replaceEdits( + content: string, + edits: Array<{ oldString: string; newString: string }>, +): string { + if (edits.length === 0) { + throw new Error("edits must contain at least one replacement") + } + + const matches = edits.map((edit, editIndex) => { + if (edit.oldString === edit.newString) { + throw new Error(`edits[${editIndex}].oldString and newString must be different`) + } + const match = findUniqueMatch(content, edit.oldString) + return { + editIndex, + start: match.index, + end: match.index + match.search.length, + newString: edit.newString, + } + }) + + matches.sort((a, b) => a.start - b.start) + for (let index = 1; index < matches.length; index++) { + const previous = matches[index - 1] + const current = matches[index] + if (previous.end > current.start) { + throw new Error( + `edits[${previous.editIndex}] and edits[${current.editIndex}] overlap; merge them into one edit or target disjoint regions`, + ) + } + } + + let result = content + for (const match of matches.reverse()) { + result = result.substring(0, match.start) + match.newString + result.substring(match.end) + } + return result +} diff --git a/tui/flocks/tool/edit.txt b/tui/flocks/tool/edit.txt index 69c26127e..7f5be02d0 100644 --- a/tui/flocks/tool/edit.txt +++ b/tui/flocks/tool/edit.txt @@ -1,10 +1,10 @@ -Performs exact string replacements in files. +Performs one or more targeted string replacements in an existing file. Usage: -- You must use your `Read` tool at least once before editing a file. CRITICAL: After each successful edit, the file content changes. You MUST re-read the file with `Read` before making any further edits to that file, otherwise oldString will not match the updated content and the edit will fail. -- When editing text from Read tool output, ensure you preserve the exact indentation (tabs/spaces) as it appears AFTER the line number prefix. The line number prefix format is: NNNNN| (5-digit zero-padded line number, then pipe, then a single space). Everything after that "| " is the actual file content to match. Never include any part of the line number prefix in the oldString or newString. For example, if Read outputs "00042| return x", the actual content to use in oldString is " return x" (4 spaces + "return x"). +- Pass every replacement inside the required `edits` array: `{"filePath":"path/to/file","edits":[{"oldString":"current text","newString":"replacement text"}]}`. +- You must use your `Read` tool at least once before editing a file. After each successful edit, re-read the file before making another edit call because the file contents have changed. +- Each `edits[].oldString` is matched against the original file snapshot. It must be non-empty, unique, and must not overlap another replacement in the same call. +- When editing text from Read tool output, preserve the exact indentation after the `NNNNN| ` line-number prefix. Never include the prefix itself in `oldString` or `newString`. - ALWAYS prefer editing existing files in the codebase. NEVER write new files unless explicitly required. - Only use emojis if the user explicitly requests it. Avoid adding emojis to files unless asked. -- The edit will FAIL if `oldString` is not found in the file with an error "oldString not found in content". -- The edit will FAIL if `oldString` is found multiple times in the file with an error "oldString found multiple times and requires more code context to uniquely identify the intended match". Either provide a larger string with more surrounding context to make it unique or use `replaceAll` to change every instance of `oldString`. -- Use `replaceAll` for replacing and renaming strings across the file. This parameter is useful if you want to rename a variable for instance. +- If a target is missing or appears more than once, provide a larger unique snippet from the current file contents. 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[]; diff --git a/webui/src/api/workspace.ts b/webui/src/api/workspace.ts index 79edf843a..edc8fce5d 100644 --- a/webui/src/api/workspace.ts +++ b/webui/src/api/workspace.ts @@ -9,6 +9,7 @@ export interface WorkspaceNode { size?: number; modified_at?: number; is_text_file?: boolean; + editable?: boolean; children?: WorkspaceNode[]; project_name?: string; project_worktree?: string; diff --git a/webui/src/components/common/SessionChat.test.ts b/webui/src/components/common/SessionChat.test.ts index 23026c37e..43eb8b090 100644 --- a/webui/src/components/common/SessionChat.test.ts +++ b/webui/src/components/common/SessionChat.test.ts @@ -1771,13 +1771,14 @@ describe('getThinkingFirstSentence', () => { }); describe('process group duration', () => { - it('uses the full wall-clock range and the current time for an active step', () => { + it('sums actual process intervals without counting gaps', () => { const parts = [ - { id: 'reason', type: 'reasoning', time: { start: 1_000, end: 2_000 } }, - { id: 'tool', type: 'tool', state: { status: 'running', time: { start: 2_500 } } }, + { id: 'reason', type: 'reasoning', time: { start: 0, end: 5_000 } }, + { id: 'summary', type: 'text', text: '继续检查', time: { start: 8_000, end: 9_000 } }, + { id: 'tool', type: 'tool', state: { status: 'completed', time: { start: 20_000, end: 28_000 } } }, ] as Message['parts']; - expect(getProcessGroupDurationMs(parts, 5_000)).toBe(4_000); + expect(getProcessGroupDurationMs(parts)).toBe(14_000); expect(formatProcessDuration(500)).toBe('1s'); expect(formatProcessDuration(7_600)).toBe('7s'); expect(formatProcessDuration(260_900)).toBe('4m20s'); @@ -1785,6 +1786,49 @@ describe('process group duration', () => { .toBeNull(); }); + it('uses current time only for unfinished active intervals', () => { + const parts = [ + { id: 'done', type: 'tool', state: { status: 'completed', time: { start: 0, end: 5_000 } } }, + { id: 'stale-text', type: 'text', text: 'missed final update', time: { start: 8_000 } }, + { id: 'running', type: 'tool', state: { status: 'running', time: { start: 20_000 } } }, + { id: 'future-text', type: 'text', text: '等待下一步', time: { start: 40_000 } }, + ] as Message['parts']; + + expect(getProcessGroupDurationMs(parts, 22_000)).toBe(7_000); + expect(getProcessGroupDurationMs(parts, 23_000)).toBe(8_000); + expect(getProcessGroupDurationMs(parts, 23_000, 'running')).toBe(8_000); + }); + + it('merges overlapping intervals so parallel process steps are not double counted', () => { + const parts = [ + { id: 'tool-a', type: 'tool', state: { status: 'completed', time: { start: 0, end: 10_000 } } }, + { id: 'tool-b', type: 'tool', state: { status: 'completed', time: { start: 5_000, end: 15_000 } } }, + { id: 'reason', type: 'reasoning', time: { start: 20_000, end: 22_000 } }, + ] as Message['parts']; + + expect(getProcessGroupDurationMs(parts)).toBe(17_000); + }); + + it('ignores invalid or incomplete historical intervals', () => { + const parts = [ + { id: 'missing-start', type: 'reasoning', time: { end: 5_000 } }, + { id: 'missing-end', type: 'tool', state: { status: 'completed', time: { start: 10_000 } } }, + { id: 'negative', type: 'text', text: 'bad clock', time: { start: 20_000, end: 19_000 } }, + { id: 'zero', type: 'tool', state: { status: 'completed', time: { start: 25_000, end: 25_000 } } }, + { id: 'valid', type: 'reasoning', time: { start: 30_000, end: 32_000 } }, + ] as unknown as Message['parts']; + + expect(getProcessGroupDurationMs(parts)).toBe(2_000); + }); + + it('keeps zero-duration completed intervals displayable', () => { + const parts = [ + { id: 'instant', type: 'tool', state: { status: 'completed', time: { start: 1_000, end: 1_000 } } }, + ] as Message['parts']; + + expect(getProcessGroupDurationMs(parts)).toBe(0); + }); + it('updates the displayed duration while the last process step is active', () => { vi.useFakeTimers(); vi.setSystemTime(5_000); @@ -1815,6 +1859,54 @@ describe('process group duration', () => { vi.useRealTimers(); } }); + + it('continues timing while the active assistant output streams after a process group', () => { + vi.useFakeTimers(); + vi.setSystemTime(5_000); + try { + render(React.createElement(ChatMessageBubble, { + message: makeMessage({ + id: 'assistant-active-output-duration', + role: 'assistant', + parts: [ + { + id: 'reason-active-output', + type: 'reasoning', + text: '先读取上下文', + time: { start: 0, end: 1_000 }, + }, + { + id: 'tool-active-output', + type: 'tool', + tool: 'todo', + state: { + status: 'completed', + input: { todos: ['检查'] }, + output: { ok: true }, + time: { start: 1_000, end: 1_000 }, + }, + }, + { + id: 'text-active-output', + type: 'text', + text: '继续输出中...', + time: { start: 1_000 }, + }, + ] as Message['parts'], + }), + isActive: true, + collapseIntermediateSteps: true, + })); + + expect(screen.getByText('查看 2 个步骤')).toBeInTheDocument(); + expect(screen.getAllByText('继续输出中...').length).toBeGreaterThan(0); + expect(screen.getByTestId('chat-process-duration')).toHaveTextContent('已处理 5s'); + act(() => vi.advanceTimersByTime(1_000)); + expect(screen.getByTestId('chat-process-duration')).toHaveTextContent('已处理 6s'); + } finally { + vi.useRealTimers(); + } + }); }); describe('ChatMessageBubble reasoning streaming', () => { @@ -3479,6 +3571,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, { @@ -5644,6 +5760,27 @@ describe('areChatMessagePartsRenderEqual', () => { )).toBe(false); }); + it('detects top-level part time updates used by reasoning and text parts', () => { + expect(areChatMessagePartsRenderEqual( + [ + { + id: 'reason-1', + type: 'reasoning', + text: '检查上下文', + time: { start: 1_000 }, + } as Message['parts'][number], + ], + [ + { + id: 'reason-1', + type: 'reasoning', + text: '检查上下文', + time: { start: 1_000, end: 2_500 }, + } as Message['parts'][number], + ], + )).toBe(false); + }); + it('keeps skipping rerenders when semantically identical parts are recreated', () => { expect(areChatMessagePartsRenderEqual( [ diff --git a/webui/src/components/common/SessionChat.tsx b/webui/src/components/common/SessionChat.tsx index 04a207432..00817a265 100644 --- a/webui/src/components/common/SessionChat.tsx +++ b/webui/src/components/common/SessionChat.tsx @@ -311,24 +311,81 @@ function getProcessPartTime(part: MessagePart): { start: number; end?: number } return part.type === 'tool' ? part.state?.time : part.time; } +type ProcessTimeInterval = { + start: number; + end: number; +}; + +function getProcessPartInterval(part: MessagePart, activeNowMs?: number): ProcessTimeInterval | null { + const time = getProcessPartTime(part); + if (!time || !Number.isFinite(time.start)) return null; + + const end = Number.isFinite(time.end) ? time.end : activeNowMs; + if (end === undefined || !Number.isFinite(end) || end < time.start) return null; + + return { start: time.start, end }; +} + +function sumMergedProcessIntervals(intervals: ProcessTimeInterval[]): number { + if (intervals.length === 0) return 0; + + const sorted = [...intervals].sort((a, b) => a.start - b.start); + let total = 0; + let currentStart = sorted[0].start; + let currentEnd = sorted[0].end; + + for (let index = 1; index < sorted.length; index += 1) { + const interval = sorted[index]; + if (interval.start <= currentEnd) { + currentEnd = Math.max(currentEnd, interval.end); + continue; + } + + total += currentEnd - currentStart; + currentStart = interval.start; + currentEnd = interval.end; + } + + return total + currentEnd - currentStart; +} + export function getProcessGroupDurationMs( parts: readonly MessagePart[], activeNowMs?: number, + activePartId?: string, ): number | null { - let firstStart = Number.POSITIVE_INFINITY; - let lastEnd = Number.NEGATIVE_INFINITY; - - for (const part of parts) { - const time = getProcessPartTime(part); - if (!time || !Number.isFinite(time.start)) continue; - const end = Number.isFinite(time.end) ? time.end : activeNowMs; - if (end === undefined || !Number.isFinite(end)) continue; - firstStart = Math.min(firstStart, time.start); - lastEnd = Math.max(lastEnd, end); + let activeIntervalIndex = -1; + if (activeNowMs !== undefined) { + const startIndex = activePartId + ? parts.findIndex((part) => part.id === activePartId) + : parts.length - 1; + const endIndex = activePartId ? startIndex : 0; + for (let index = startIndex; index >= endIndex; index -= 1) { + const part = parts[index]; + if (!part) continue; + const time = getProcessPartTime(part); + if ( + !!time + && Number.isFinite(time.start) + && !Number.isFinite(time.end) + && time.start <= activeNowMs + && (part.type !== 'tool' || isActiveToolPart(part)) + ) { + activeIntervalIndex = index; + break; + } + } } - if (!Number.isFinite(firstStart) || !Number.isFinite(lastEnd)) return null; - return Math.max(0, lastEnd - firstStart); + const intervals = parts + .map((part, index) => getProcessPartInterval( + part, + index === activeIntervalIndex ? activeNowMs : undefined, + )) + .filter((interval): interval is ProcessTimeInterval => interval !== null); + + if (intervals.length === 0) return null; + return sumMergedProcessIntervals(intervals); } export function formatProcessDuration(durationMs: number): string { @@ -5336,13 +5393,24 @@ function ChatMessageBubbleInner({ })()} ); - const renderProcessGroup = (group: Array<{ part: MessagePart; index: number }>, groupIndex: number) => { + const renderProcessGroup = ( + group: Array<{ part: MessagePart; index: number }>, + groupIndex: number, + activeTimingPart?: MessagePart, + ) => { const processGroupOpen = processGroupsDefaultOpen || (processGroupsOpenWhileActive && isActive); const processGroupKey = `${message.id}:process:${groupIndex}`; - const processGroupActive = isActive && group.some(({ part }) => part === activeTailPart); + const groupParts = group.map(({ part }) => part); + const groupActivePart = groupParts.find((part) => part === activeTailPart); + const timedActivePart = groupActivePart || activeTimingPart; + const durationParts = timedActivePart && !groupParts.includes(timedActivePart) + ? [...groupParts, timedActivePart] + : groupParts; + const processGroupActive = isActive && !!timedActivePart; const processDurationMs = getProcessGroupDurationMs( - group.map(({ part }) => part), + durationParts, processGroupActive ? processElapsedClock : undefined, + processGroupActive ? timedActivePart?.id : undefined, ); const hasStoredOpenState = !!processGroupOpenState && Object.prototype.hasOwnProperty.call(processGroupOpenState, processGroupKey); @@ -5390,9 +5458,9 @@ function ChatMessageBubbleInner({ const lastIntermediateProcessIndex = displayParts.reduce((lastIndex, part, index) => ( isIntermediateProcessPart(part) ? index : lastIndex ), -1); - const flushProcessGroup = () => { + const flushProcessGroup = (activeTimingPart?: MessagePart) => { if (processGroup.length === 0) return; - nodes.push(renderProcessGroup(processGroup, processGroupIndex)); + nodes.push(renderProcessGroup(processGroup, processGroupIndex, activeTimingPart)); processGroup = []; processGroupIndex += 1; }; @@ -5403,7 +5471,7 @@ function ChatMessageBubbleInner({ } if (!isRenderableDisplayPart(part)) return; if (isPendingQuestionToolPart(part) || isRenderableTextPart(part)) { - flushProcessGroup(); + flushProcessGroup(isActive && part === activeTailPart && isRenderableTextPart(part) ? part : undefined); } nodes.push(renderPart(part, index)); }); @@ -5448,11 +5516,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 (
diff --git a/webui/src/components/common/sessionChatRenderEquality.ts b/webui/src/components/common/sessionChatRenderEquality.ts index 222477bed..b6c809054 100644 --- a/webui/src/components/common/sessionChatRenderEquality.ts +++ b/webui/src/components/common/sessionChatRenderEquality.ts @@ -57,6 +57,9 @@ export function areChatMessagePartsRenderEqual( || prevPart.mime !== nextPart.mime || prevPart.filename !== nextPart.filename || prevPart.url !== nextPart.url + || prevPart.time?.start !== nextPart.time?.start + || prevPart.time?.end !== nextPart.time?.end + || prevPart.time?.compacted !== nextPart.time?.compacted || prevPart.image?.url !== nextPart.image?.url || prevPart.image?.alt !== nextPart.image?.alt ) { diff --git a/webui/src/pages/Home/index.test.tsx b/webui/src/pages/Home/index.test.tsx index 97b1a0adf..9be0d9f2b 100644 --- a/webui/src/pages/Home/index.test.tsx +++ b/webui/src/pages/Home/index.test.tsx @@ -4,16 +4,20 @@ import userEvent from '@testing-library/user-event'; import { MemoryRouter } from 'react-router-dom'; import Home from './index'; -const { createMock, navigateMock, toastErrorMock, useAuthMock, useStatsMock } = vi.hoisted(() => ({ - createMock: vi.fn(), - navigateMock: vi.fn(), - toastErrorMock: vi.fn(), - useAuthMock: vi.fn(), - useStatsMock: vi.fn(), -})); +const { createMock, navigateMock, toastErrorMock, useAuthMock, useStatsMock } = + vi.hoisted(() => ({ + createMock: vi.fn(), + navigateMock: vi.fn(), + toastErrorMock: vi.fn(), + useAuthMock: vi.fn(), + useStatsMock: vi.fn(), + })); vi.mock('react-router-dom', async () => { - const actual = await vi.importActual('react-router-dom'); + const actual = + await vi.importActual( + 'react-router-dom', + ); return { ...actual, useNavigate: () => navigateMock, @@ -75,7 +79,9 @@ describe('Home create WebUI contract page entry', () => { , ); - await user.click(screen.getByRole('button', { name: 'createWebUIContractPage' })); + await user.click( + screen.getByRole('button', { name: 'createWebUIContractPage' }), + ); await waitFor(() => { expect(createMock).toHaveBeenCalledWith({ @@ -105,7 +111,9 @@ describe('Home create WebUI contract page entry', () => { , ); - expect(screen.queryByRole('button', { name: 'createWebUIContractPage' })).not.toBeInTheDocument(); + expect( + screen.queryByRole('button', { name: 'createWebUIContractPage' }), + ).not.toBeInTheDocument(); expect(createMock).not.toHaveBeenCalled(); }); @@ -122,7 +130,75 @@ describe('Home create WebUI contract page entry', () => { , ); - expect(screen.getByText('stats.loadErrorHint.authExpired')).toBeInTheDocument(); + expect( + screen.getByText('stats.loadErrorHint.authExpired'), + ).toBeInTheDocument(); expect(screen.queryByText(/backend is running/i)).not.toBeInTheDocument(); }); + + it('links dashboard stat cards to their matching pages', () => { + useStatsMock.mockReturnValue({ + stats: { + agents: { total: 11 }, + workflows: { total: 7 }, + skills: { total: 14 }, + tools: { total: 202 }, + tasks: { week: 3, scheduledActive: 2 }, + models: { total: 39 }, + system: { status: 'healthy', message: 'healthy' }, + }, + loading: false, + error: null, + }); + + render( + + + , + ); + + expect( + screen.getByRole('link', { name: 'stats.agentCount' }), + ).toHaveAttribute('href', '/agents'); + expect(screen.getByRole('link', { name: 'stats.agentCount' })).toHaveClass( + 'hover:border-purple-200', + ); + expect( + screen.getByRole('link', { name: 'stats.workflowCount' }), + ).toHaveAttribute('href', '/workflows'); + expect( + screen.getByRole('link', { name: 'stats.workflowCount' }), + ).toHaveClass('hover:border-teal-200'); + expect( + screen.getByRole('link', { name: 'stats.skillCount' }), + ).toHaveAttribute('href', '/skills'); + expect(screen.getByRole('link', { name: 'stats.skillCount' })).toHaveClass( + 'hover:border-green-200', + ); + expect( + screen.getByRole('link', { name: 'stats.toolCount' }), + ).toHaveAttribute('href', '/tools'); + expect(screen.getByRole('link', { name: 'stats.toolCount' })).toHaveClass( + 'hover:border-orange-200', + ); + expect( + screen.getByRole('link', { name: 'stats.weeklyTasks' }), + ).toHaveAttribute('href', '/tasks'); + expect(screen.getByRole('link', { name: 'stats.weeklyTasks' })).toHaveClass( + 'hover:border-amber-200', + ); + expect( + screen.getByRole('link', { name: 'stats.activeScheduled' }), + ).toHaveAttribute('href', '/tasks'); + expect( + screen.getByRole('link', { name: 'stats.activeScheduled' }), + ).toHaveClass('hover:border-violet-200'); + expect( + screen.getByRole('link', { name: 'stats.modelCount' }), + ).toHaveAttribute('href', '/models'); + expect(screen.getByRole('link', { name: 'stats.modelCount' })).toHaveClass( + 'hover:border-pink-200', + ); + expect(screen.getByText('stats.systemStatus').closest('a')).toBeNull(); + }); }); diff --git a/webui/src/pages/Home/index.tsx b/webui/src/pages/Home/index.tsx index f2d6e0c79..4694ca106 100644 --- a/webui/src/pages/Home/index.tsx +++ b/webui/src/pages/Home/index.tsx @@ -1,7 +1,7 @@ -import { - Bot, - Zap, - Sparkles, +import { + Bot, + Zap, + Sparkles, Github, ChevronDown, ChevronRight, @@ -17,6 +17,7 @@ import { Loader2, } from 'lucide-react'; import { useCallback, useState } from 'react'; +import type { ReactNode } from 'react'; import { Link, useNavigate } from 'react-router-dom'; import { useTranslation } from 'react-i18next'; import { useStats } from '@/hooks/useStats'; @@ -29,6 +30,51 @@ import { useProductName } from '@/contexts/ProductNameContext'; const GITHUB_URL = 'https://github.com/AgentFlocks/flocks'; const GITEE_URL = 'https://gitee.com/flocks/flocks'; const GITEE_LOGO_URL = `${import.meta.env.BASE_URL}gitee-logo.png`; +const STAT_CARD_CLASS = 'bg-white rounded-xl p-6 border border-gray-200'; +const STAT_LINK_CARD_CLASS = `${STAT_CARD_CLASS} block transition-all duration-200 hover:shadow-md focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-offset-2`; + +function StatCard({ + title, + value, + icon, + to, + interactiveClassName, + valueClassName = 'text-gray-900', + children, +}: { + title: string; + value: ReactNode; + icon: ReactNode; + to?: string; + interactiveClassName?: string; + valueClassName?: string; + children?: ReactNode; +}) { + const content = ( + <> +
+ {title} + {icon} +
+
{value}
+ {children} + + ); + + if (to) { + return ( + + {content} + + ); + } + + return
{content}
; +} export default function Home() { const { stats, loading, error } = useStats(); @@ -39,18 +85,30 @@ export default function Home() { const { productName } = useProductName(); const canCreateWebUIContractPage = user?.role === 'admin'; const [isRepoMenuOpen, setIsRepoMenuOpen] = useState(false); - const [creatingWebUIContractPageSession, setCreatingWebUIContractPageSession] = useState(false); - const statsErrorHint = error ? t(`stats.loadErrorHint.${error.message}`, { defaultValue: t('stats.loadErrorHint.unavailable') }) : ''; + const [ + creatingWebUIContractPageSession, + setCreatingWebUIContractPageSession, + ] = useState(false); + const statsErrorHint = error + ? t(`stats.loadErrorHint.${error.message}`, { + defaultValue: t('stats.loadErrorHint.unavailable'), + }) + : ''; const handleCreateWebUIContractPage = useCallback(async () => { if (creatingWebUIContractPageSession) return; setCreatingWebUIContractPageSession(true); try { - const session = await sessionApi.create({ title: t('createWebUIContractPageSessionTitle') }); + const session = await sessionApi.create({ + title: t('createWebUIContractPageSessionTitle'), + }); const message = t('createWebUIContractPageInitialMessage'); - navigate(`/sessions?session=${session.id}&message=${encodeURIComponent(message)}`); + navigate( + `/sessions?session=${session.id}&message=${encodeURIComponent(message)}`, + ); } catch (err: unknown) { - const detail = err instanceof Error ? err.message : t('createWebUIContractPageError'); + const detail = + err instanceof Error ? err.message : t('createWebUIContractPageError'); toast.error(t('createWebUIContractPageError'), detail); } finally { setCreatingWebUIContractPageSession(false); @@ -87,7 +145,9 @@ export default function Home() {
@@ -159,8 +223,12 @@ export default function Home() {
-

{t('quickActions.sessions.title')}

-

{t('quickActions.sessions.description')}

+

+ {t('quickActions.sessions.title')} +

+

+ {t('quickActions.sessions.description')} +

@@ -173,8 +241,12 @@ export default function Home() {
-

{t('quickActions.workflows.title')}

-

{t('quickActions.workflows.description')}

+

+ {t('quickActions.workflows.title')} +

+

+ {t('quickActions.workflows.description')} +

@@ -187,8 +259,12 @@ export default function Home() {
-

{t('quickActions.agents.title')}

-

{t('quickActions.agents.description')}

+

+ {t('quickActions.agents.title')} +

+

+ {t('quickActions.agents.description')} +

@@ -205,98 +281,100 @@ export default function Home() {
- {t('stats.abnormal')} - {statsErrorHint} + + {t('stats.abnormal')} + + + {statsErrorHint} +
)}
-
-
- {t('stats.agentCount')} - -
-
- {stats?.agents.total ?? 0} -
-
+ } + to="/agents" + interactiveClassName="hover:border-purple-200 hover:shadow-purple-50 focus-visible:ring-purple-500" + /> -
-
- {t('stats.workflowCount')} - -
-
- {stats?.workflows.total ?? 0} -
-
+ } + to="/workflows" + interactiveClassName="hover:border-teal-200 hover:shadow-teal-50 focus-visible:ring-teal-500" + /> -
-
- {t('stats.skillCount')} - -
-
- {stats?.skills.total ?? 0} -
-
+ } + to="/skills" + interactiveClassName="hover:border-green-200 hover:shadow-green-50 focus-visible:ring-green-500" + /> -
-
- {t('stats.toolCount')} - -
-
- {stats?.tools.total ?? 0} -
-
+ } + to="/tools" + interactiveClassName="hover:border-orange-200 hover:shadow-orange-50 focus-visible:ring-orange-500" + /> -
-
- {t('stats.weeklyTasks')} - -
-
- {stats?.tasks.week ?? 0} -
-
+ } + to="/tasks" + interactiveClassName="hover:border-amber-200 hover:shadow-amber-50 focus-visible:ring-amber-500" + /> -
-
- {t('stats.activeScheduled')} - -
-
- {stats?.tasks.scheduledActive ?? 0} -
-
+ } + to="/tasks" + interactiveClassName="hover:border-violet-200 hover:shadow-violet-50 focus-visible:ring-violet-500" + /> -
-
- {t('stats.modelCount')} - -
-
- {stats?.models.total ?? 0} -
-
+ } + to="/models" + interactiveClassName="hover:border-pink-200 hover:shadow-pink-50 focus-visible:ring-pink-500" + /> -
-
- {t('stats.systemStatus')} - -
-
- {stats?.system.status === 'healthy' ? t('stats.normal') : t('stats.abnormal')} -
-
- {stats?.system.status ? t(`stats.statusMessage.${stats.system.status}`) : ''} -
-
+ + } + > +
+ {stats?.system.status + ? t(`stats.statusMessage.${stats.system.status}`) + : ''} +
+
)} diff --git a/webui/src/pages/Workspace/index.test.tsx b/webui/src/pages/Workspace/index.test.tsx index aa3d03005..7e930f774 100644 --- a/webui/src/pages/Workspace/index.test.tsx +++ b/webui/src/pages/Workspace/index.test.tsx @@ -196,7 +196,7 @@ function directory(name: string, path: string) { }; } -function file(name: string, path: string, isTextFile = true) { +function file(name: string, path: string, isTextFile = true, editable = isTextFile) { return { name, path, @@ -204,6 +204,7 @@ function file(name: string, path: string, isTextFile = true) { size: 24, modified_at: 1710000000, is_text_file: isTextFile, + editable, }; } @@ -452,7 +453,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 +497,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', true, false), + file('bak.txt', 'bak.txt', true, false), + { + ...directory('archive', 'archive'), + children: [file('2026-08-18.md', 'archive/2026-08-18.md', true, false)], + }, ], }); + mocks.readMemoryFile.mockResolvedValue({ + data: { + path: 'bak.txt', + content: 'backup memory', + truncated: false, + }, + }); const user = userEvent.setup(); renderWithRouter(); @@ -526,11 +537,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..970748c9e 100644 --- a/webui/src/pages/Workspace/index.tsx +++ b/webui/src/pages/Workspace/index.tsx @@ -1594,11 +1594,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 +1621,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 +1915,7 @@ function MemoryTab() { {formatBytes(selected.size ?? 0)} {formatDate(selected.modified_at)} - {selected.is_text_file && !editing && !truncated && contentState === 'ready' && ( + {selected.is_text_file && selected.editable === true && !editing && !truncated && contentState === 'ready' && (