From f7882be54691dd38e0c4c18bea70462f7587f2bc Mon Sep 17 00:00:00 2001 From: Ed1s0nZ Date: Wed, 16 Sep 2026 17:52:58 +0800 Subject: [PATCH] feat: manage task process lifetimes and preserve turn history --- .github/workflows/process-isolation.yml | 49 +++ README.md | 1 - README_CN.md | 1 - cmd/server/main.go | 36 ++ config.example.yaml | 9 + docs/en-US/contributing-guide.md | 10 +- docs/en-US/tool-execution-governance.md | 30 ++ docs/zh-CN/MULTI_AGENT_EINO.md | 2 +- docs/zh-CN/contributing-guide.md | 10 +- docs/zh-CN/tool-execution-governance.md | 59 ++++ go.mod | 1 - internal/app/app.go | 3 + internal/config/config.go | 10 + internal/handler/agent.go | 31 +- internal/handler/batch_queue_executor.go | 19 +- internal/handler/eino_single_agent.go | 36 +- internal/handler/multi_agent.go | 34 +- internal/handler/task_lifecycle.go | 62 ++++ internal/handler/task_manager.go | 268 ++++++++++++--- internal/handler/task_process_cleanup_test.go | 225 +++++++++++++ internal/handler/workflow_integration.go | 26 +- internal/mcp/execution_ownership_test.go | 76 +++++ internal/mcp/execution_service.go | 72 +++- internal/mcp/external_manager.go | 14 + .../multiagent/eino_model_facing_trace.go | 32 ++ internal/multiagent/eino_turn_history.go | 196 +++++++++++ internal/multiagent/eino_turn_history_test.go | 254 ++++++++++++++ .../multiagent/eino_turn_loop_bridge_test.go | 3 + internal/multiagent/eino_turn_loop_runtime.go | 15 +- .../multiagent/eino_turn_loop_runtime_test.go | 3 + internal/processguard/cgroup_linux_test.go | 120 +++++++ internal/processguard/check.go | 52 +++ internal/processguard/group_unix.go | 180 ++++++++++ internal/processguard/guard.go | 89 +++++ internal/processguard/guard_test.go | 160 +++++++++ internal/processguard/job_windows_test.go | 97 ++++++ internal/processguard/platform_linux.go | 313 ++++++++++++++++++ internal/processguard/platform_other.go | 31 ++ internal/processguard/platform_windows.go | 183 ++++++++++ internal/processguard/watchdog.go | 182 ++++++++++ internal/runlease/scope.go | 126 +++++++ internal/runlease/scope_test.go | 70 ++++ internal/security/executor.go | 153 ++------- internal/security/executor_test.go | 8 +- internal/security/procattr_unix.go | 11 + internal/security/procattr_windows.go | 15 +- internal/security/process_scope.go | 220 ++++++++++++ internal/security/process_scope_test.go | 198 +++++++++++ internal/security/shell_execute_stream.go | 55 +-- internal/security/shell_session.go | 99 ++++-- web/static/i18n/en-US.json | 3 + web/static/i18n/zh-CN.json | 3 + web/static/js/monitor.js | 7 +- web/static/js/tasks.js | 18 +- 54 files changed, 3650 insertions(+), 330 deletions(-) create mode 100644 .github/workflows/process-isolation.yml create mode 100644 internal/handler/task_lifecycle.go create mode 100644 internal/handler/task_process_cleanup_test.go create mode 100644 internal/mcp/execution_ownership_test.go create mode 100644 internal/multiagent/eino_turn_history.go create mode 100644 internal/multiagent/eino_turn_history_test.go create mode 100644 internal/processguard/cgroup_linux_test.go create mode 100644 internal/processguard/check.go create mode 100644 internal/processguard/group_unix.go create mode 100644 internal/processguard/guard.go create mode 100644 internal/processguard/guard_test.go create mode 100644 internal/processguard/job_windows_test.go create mode 100644 internal/processguard/platform_linux.go create mode 100644 internal/processguard/platform_other.go create mode 100644 internal/processguard/platform_windows.go create mode 100644 internal/processguard/watchdog.go create mode 100644 internal/runlease/scope.go create mode 100644 internal/runlease/scope_test.go create mode 100644 internal/security/process_scope.go create mode 100644 internal/security/process_scope_test.go diff --git a/.github/workflows/process-isolation.yml b/.github/workflows/process-isolation.yml new file mode 100644 index 00000000..82782ccb --- /dev/null +++ b/.github/workflows/process-isolation.yml @@ -0,0 +1,49 @@ +name: Process isolation +on: + push: + paths: ['internal/processguard/**', 'internal/runlease/**', 'internal/security/**', 'internal/handler/task*', 'internal/mcp/**', 'go.mod', 'go.sum', '.github/workflows/process-isolation.yml'] + pull_request: + paths: ['internal/processguard/**', 'internal/runlease/**', 'internal/security/**', 'internal/handler/task*', 'internal/mcp/**', 'go.mod', 'go.sum', '.github/workflows/process-isolation.yml'] +permissions: + contents: read +jobs: + lifecycle: + strategy: + matrix: + os: [ubuntu-latest, macos-latest, windows-latest] + runs-on: ${{ matrix.os }} + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-go@v5 + with: + go-version-file: go.mod + - run: go test -race ./internal/processguard ./internal/runlease + - name: Test Linux cgroup isolation + if: runner.os == 'Linux' + shell: bash + run: | + set -euo pipefail + case "$(docker info --format '{{.Architecture}}')" in + aarch64|arm64) fixture_arch=arm64 ;; + x86_64|amd64) fixture_arch=amd64 ;; + *) echo 'Unsupported Docker architecture' >&2; exit 1 ;; + esac + fixture_dir="$(mktemp -d)" + trap 'rm -rf "$fixture_dir"' EXIT + CGO_ENABLED=0 GOOS=linux GOARCH="$fixture_arch" go test ./internal/processguard -c -o "$fixture_dir/processguard.test" + docker run --rm --user 0 --privileged --cgroupns=private --network none --init \ + --mount "type=bind,src=$fixture_dir/processguard.test,dst=/fixture,readonly" \ + --entrypoint /bin/sh "${CSAI_TEST_IMAGE:-debian:bookworm-slim}" -c ' + set -eu + mkdir /sys/fs/cgroup/infra /sys/fs/cgroup/csai-fixture + # cat may already have exited when its PID is visited. + for pid in $(cat /sys/fs/cgroup/cgroup.procs); do + echo "$pid" > /sys/fs/cgroup/infra/cgroup.procs 2>/dev/null || true + done + echo "+cpu +memory +pids" > /sys/fs/cgroup/cgroup.subtree_control + echo "+cpu +memory +pids" > /sys/fs/cgroup/csai-fixture/cgroup.subtree_control + export CSAI_TEST_CGROUP_ROOT=/sys/fs/cgroup/csai-fixture + exec /fixture -test.v -test.timeout=60s + ' + - if: runner.os == 'Linux' + run: go test ./internal/security ./internal/handler ./internal/mcp diff --git a/README.md b/README.md index c92a43d0..4d1854e7 100644 --- a/README.md +++ b/README.md @@ -328,7 +328,6 @@ CyberStrikeAI/ ├── agents/ # Multi-agent Markdown (orchestrator.md + sub-agent *.md) ├── docs/ # Topic docs (deployment, config, security, API, knowledge base, C2, WebShell, etc.) ├── images/ # Docs screenshots & diagrams -├── scripts/ # Repository maintenance checks, including documentation validation ├── config.yaml # Runtime configuration ├── run.sh # Convenience launcher └── README*.md diff --git a/README_CN.md b/README_CN.md index c38c3eda..58549fc8 100644 --- a/README_CN.md +++ b/README_CN.md @@ -326,7 +326,6 @@ CyberStrikeAI/ ├── agents/ # 多代理 Markdown（orchestrator.md + 子代理 *.md） ├── docs/ # 专题文档（部署、配置、安全、API、知识库、C2、WebShell 等） ├── images/ # 文档配图 -├── scripts/ # 仓库维护检查，包括文档校验 ├── config.yaml # 运行配置 ├── run.sh # 启动脚本 └── README*.md diff --git a/cmd/server/main.go b/cmd/server/main.go index c14ad2a5..898abf2a 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -6,6 +6,7 @@ import ( "cyberstrike-ai/internal/config" "cyberstrike-ai/internal/database" "cyberstrike-ai/internal/logger" + "cyberstrike-ai/internal/processguard" "cyberstrike-ai/internal/security" "cyberstrike-ai/internal/termout" "flag" @@ -14,6 +15,7 @@ import ( "os/signal" "strings" "syscall" + "time" "go.uber.org/zap" "golang.org/x/term" @@ -24,6 +26,7 @@ func main() { var httpsBootstrap = flag.Bool("https", false, "Enable HTTPS for the main site; uses an in-memory self-signed certificate when no cert/key is configured") var httpBootstrap = flag.Bool("http", false, "Force plain HTTP for the main site, overriding TLS settings in the configuration file") var resetAdminPassword = flag.Bool("reset-admin-password", false, "Interactively reset the built-in admin password and exit") + checkIsolation := flag.Bool("check-process-isolation", false, "Probe task containment and cleanup, then exit without starting services") flag.Parse() // 环境变量兼容（便于 systemd/docker 等不传参场景） @@ -62,6 +65,22 @@ func main() { termout.PrintConfigCreated() } + if *checkIsolation { + if err := configureProcessIsolation(cfg); err != nil { + fmt.Fprintln(os.Stderr, err) + os.Exit(1) + } + checkCtx, cancelCheck := context.WithTimeout(context.Background(), 15*time.Second) + backend, checkErr := processguard.Check(checkCtx) + cancelCheck() + if checkErr != nil { + fmt.Fprintln(os.Stderr, checkErr) + os.Exit(1) + } + fmt.Printf("{\"checked\":true,\"backend\":%q}\n", backend) + return + } + if *resetAdminPassword { if err := runResetAdminPassword(cfg); err != nil { fmt.Fprintf(os.Stderr, "Failed to reset admin password: %v\n", err) @@ -109,6 +128,18 @@ func main() { }) defer log.Sync() + if err := configureProcessIsolation(cfg); err != nil { + log.Fatal("进程隔离初始化失败", "error", err) + } + + probeCtx, probeCancel := context.WithTimeout(context.Background(), 15*time.Second) + backend, probeErr := processguard.Check(probeCtx) + probeCancel() + if probeErr != nil { + log.Fatal("进程隔离启动检查失败", "error", probeErr) + } + log.Info("任务进程隔离已就绪", zap.String("backend", backend)) + // 创建可取消的根 context，用于优雅关闭 ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -210,3 +241,8 @@ func readHiddenPassword(prompt string) (string, error) { } return string(password), nil } + +func configureProcessIsolation(cfg *config.Config) error { + isolation := cfg.Security.ProcessIsolation + return processguard.Configure(processguard.Options{Mode: isolation.Mode, CgroupRoot: isolation.CgroupRoot, MaxProcesses: isolation.MaxProcesses, MemoryMaxBytes: isolation.MemoryMaxBytes, CPUQuotaMicros: isolation.CPUQuotaMicros}) +} diff --git a/config.example.yaml b/config.example.yaml index a499f56b..24e1964a 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -364,6 +364,15 @@ database: # 系统会从该目录加载所有 .yaml 格式的工具配置文件 # 推荐方式：在 tools/ 目录下为每个工具创建独立的配置文件 security: + # auto: native Job Object on Windows; Unix watchdog unless cgroup_root is set. + # Linux production: configure a systemd service with Delegate=cpu memory pids; set + # mode: required, cgroup_root: auto. Explicit isolation errors fail closed. + process_isolation: + mode: auto + cgroup_root: "" + max_processes: 256 + memory_max_bytes: 2147483648 + cpu_quota_micros: 0 # 0=no CPU cap; 100000=one core (100 ms period) tools_dir: tools # 工具配置文件目录（相对于配置文件所在目录） # 工具描述模式：加载 tools 下工具时，暴露给 AI/API 使用的描述来源 # short - 优先使用 short_description（简短描述，省 token），为空时用 description diff --git a/docs/en-US/contributing-guide.md b/docs/en-US/contributing-guide.md index 9019d5f8..52f3505b 100644 --- a/docs/en-US/contributing-guide.md +++ b/docs/en-US/contributing-guide.md @@ -103,15 +103,7 @@ Update: - `docs/zh-CN/README.md` - `docs/en-US/README.md` -Before submitting documentation changes, run: - -```bash -python3 scripts/check-docs.py -``` - -The check verifies local links, fenced code blocks, bilingual filename parity, locale index coverage, and the Go version documented by the root READMEs. Keep versioned examples derived from authoritative files such as `go.mod` and `config.example.yaml` whenever possible. - -The same command runs automatically in the `Documentation` GitHub Actions workflow when documentation, `go.mod`, or the checker itself changes. +Before submitting documentation changes, verify local links, fenced code blocks, bilingual filename parity, locale index coverage, and the Go version documented by the root READMEs. Keep versioned examples derived from authoritative files such as `go.mod` and `config.example.yaml` whenever possible. ## Review Focus diff --git a/docs/en-US/tool-execution-governance.md b/docs/en-US/tool-execution-governance.md index 5ca06b80..1e193866 100644 --- a/docs/en-US/tool-execution-governance.md +++ b/docs/en-US/tool-execution-governance.md @@ -219,3 +219,33 @@ Expected behavior: - CyberStrikeAI cannot control how a remote external MCP server collects output internally; it caps results after they enter CyberStrikeAI and protects calls with concurrency limits and circuit breakers. - Oversized tool output is spilled to local `tmp/reduction/.../trunc/` (or `reduction_root_dir`) before truncation; the bounded result includes an absolute path for `read_file`. + +## Task-owned process lifetime + +Each task has a distinct `runId`. Local commands and detached MCP workers retain ownership through context values. Managed `exec ... &` and Eino background execution return promptly but cannot outlive task cleanup. Unowned background launches are rejected. Cleanup seals process and worker admission, cancels workers, terminates processes and waits for reaping. Interrupt-and-continue retains ownership; SSE disconnection does not end the task. Failed cleanup retains the conversation slot and reports `cleanup_failed`; a 15-second sweep retries it. + +### Kernel containment and crash recovery + +- **Linux with a delegated cgroup v2 root:** each task gets a cgroup with process, memory and optional CPU limits. Go uses `clone3(CLONE_INTO_CGROUP)` to assign membership at creation. A separate guardian uses `cgroup.kill` on owner pipe EOF and checks `populated=0`. Startup exclusively locks the delegated root and recovers stale task groups. `setsid` does not escape this scope. +- **Windows:** a guardian joins a Job Object before task commands are created. Commands inherit the Job atomically through the parent-process attribute. Kill-on-close, explicit termination and active-process accounting cover cleanup, with process/memory/CPU limits. +- **macOS and Unix without configured cgroups:** a separate process-group guardian acknowledges registration before a gated child executes user code. Owner death triggers EOF cleanup. This fallback cannot contain deliberate `setsid` escapes and is not strong kernel isolation. + +These are lifecycle controls, not a sandbox against code with the same privileges as the supervisor. Protect cgroup control files, process handles and the host/guardian using appropriate identities or containers. Container deployments should use an init process to reap orphans. Required mode fails closed when containment is unavailable. + +### Deployment and verification + +For Linux production, configure a dedicated systemd service with `Delegate=cpu memory pids`, `KillMode=control-group` and `Restart=on-failure`. Set `security.process_isolation.mode: required` and `security.process_isolation.cgroup_root: auto` in the existing application configuration. Linux requires kernel 5.14+, cgroup v2, clone3 permission and delegation of CPU/memory/pids controllers. `cgroup_root: auto` resolves the dedicated systemd unit root; never use the entire host hierarchy root. Settings take effect after restart. Windows required mode uses an empty cgroup root. macOS rejects required mode. + +Startup probes actual process creation and cleanup before accepting requests. Run the probe without starting HTTP/MCP services: + +```sh +./cyberstrike-ai --check-process-isolation -config config.yaml +``` + +The result reports `cgroup_v2`, `windows_job`, or `process_group_watchdog`; task APIs expose the actual `isolationBackend`. Systemd's `KillMode=control-group` and restart policy add service-wide recovery. + +Ordinary remote MCP cancellation is a notification, not a shutdown receipt. Adapters can implement `ExternalCancellationConfirmer` using server-side state, leases or cancellation receipts. Without acknowledgement, execution records persist as `orphaned`, and after local cleanup task history reports `cleanup_unconfirmed` instead of claiming success. Workers which have not returned remain tracked as cleanup failures. + +Run `go test -race ./internal/processguard ./internal/mcp ./internal/handler ./internal/security`. The `Process isolation` GitHub Actions workflow also runs the Linux cgroup integration test with commands embedded in `.github/workflows/process-isolation.yml`. The Linux integration fixture uses a disposable private-cgroup container with no host cgroup or Docker socket mounts. Tests cover owner SIGKILL, launch admission, setsid, limits, stale-group recovery and remote cancellation semantics. Windows tests are included in cross-platform CI; cross-compilation is not a Windows runtime test. + +References: [cgroup v2](https://cdn.kernel.org/doc/html/latest/admin-guide/cgroup-v2.html), [Job Objects](https://learn.microsoft.com/en-us/windows/win32/procthread/job-objects), [MCP cancellation](https://modelcontextprotocol.io/specification/2025-11-25/basic/utilities/cancellation). diff --git a/docs/zh-CN/MULTI_AGENT_EINO.md b/docs/zh-CN/MULTI_AGENT_EINO.md index b7b92841..31e4b9c9 100644 --- a/docs/zh-CN/MULTI_AGENT_EINO.md +++ b/docs/zh-CN/MULTI_AGENT_EINO.md @@ -11,7 +11,7 @@ | 项 | 说明 | |----|------| -| 依赖与代理 | `go.mod` 直接依赖 `github.com/cloudwego/eino`、`eino-ext/.../openai`；`go.mod` 注释与 `scripts/bootstrap-go.sh` 指导 **GOPROXY**（如 `https://goproxy.cn,direct`）。 | +| 依赖与代理 | `go.mod` 直接依赖 `github.com/cloudwego/eino`、`eino-ext/.../openai`；`go.mod` 注释指导 **GOPROXY**（如 `https://goproxy.cn,direct`）。 | | 配置 | `config.yaml` → `agent.max_iterations` 为全局 ReAct 上限（主/子代理统一）；`multi_agent`：`enabled`、`robot_use_multi_agent`、`sub_agents`（含可选 `bind_role`）、`eino_skills`、`eino_middleware` 等；结构体见 `internal/config/config.go`。 | | Markdown 子代理 / 主代理 | 在 `agents_dir` 下放 `*.md`。**子代理**：供 Deep `task` 与 `supervisor` `transfer`。**主代理（按模式分离）**：`orchestrator.md`（或 `kind: orchestrator` 的**单个**其他 .md）→ **Deep**；固定名 `orchestrator-plan-execute.md` → **plan_execute**；固定名 `orchestrator-supervisor.md` → **supervisor**。正文优先于 YAML：`multi_agent.orchestrator_instruction`、`orchestrator_instruction_plan_execute`、`orchestrator_instruction_supervisor`；plan_execute / supervisor **不会**回退到 Deep 的 `orchestrator_instruction`。皆空时 plan_execute / supervisor 使用代码内置默认提示。管理：**Agents → Agent管理**；API：`/api/multi-agent/markdown-agents*`。 | | MCP 桥 | `internal/einomcp`：`ToolsFromDefinitions` + 会话 ID 持有者，执行走 `Agent.ExecuteMCPToolForConversation`。 | diff --git a/docs/zh-CN/contributing-guide.md b/docs/zh-CN/contributing-guide.md index fe02ae3c..5c3134df 100644 --- a/docs/zh-CN/contributing-guide.md +++ b/docs/zh-CN/contributing-guide.md @@ -103,15 +103,7 @@ docs/en-US/ - `docs/zh-CN/README.md` - `docs/en-US/README.md` -提交文档变更前运行： - -```bash -python3 scripts/check-docs.py -``` - -该检查会验证本地链接、代码块闭合、中英文文件名对齐、语言导航覆盖率，以及根 README 中的 Go 版本是否与 `go.mod` 一致。版本化示例应尽量从 `go.mod`、`config.example.yaml` 等权威文件派生，避免手工同步。 - -当文档、`go.mod` 或检查脚本变更时，GitHub Actions 中的 `Documentation` 工作流会自动运行同一条命令。 +提交文档变更前，请检查本地链接、代码块闭合、中英文文件名对齐、语言导航覆盖率，以及根 README 中的 Go 版本是否与 `go.mod` 一致。版本化示例应尽量从 `go.mod`、`config.example.yaml` 等权威文件派生，避免手工同步。 ## Review 关注点 diff --git a/docs/zh-CN/tool-execution-governance.md b/docs/zh-CN/tool-execution-governance.md index 71a99765..04d284ae 100644 --- a/docs/zh-CN/tool-execution-governance.md +++ b/docs/zh-CN/tool-execution-governance.md @@ -221,3 +221,62 @@ PY - 外部 MCP 的远端 server 内部如何采集输出不由 CyberStrikeAI 控制；CyberStrikeAI 会在结果进入本系统后统一兜底、限并发和熔断。 - 超长工具输出会在截断前写入本地 `tmp/reduction/.../trunc/`（或 `reduction_root_dir`），bounded result 中包含可 `read_file` 的绝对路径。 + +## 本轮任务的进程生命周期 + +每轮任务有独立 `runId`，本地进程与异步 MCP worker 都通过 context 保留归属。取消 context、工具返回、SSE 断线不等于资源已回收。 + +- `exec ... &` 与 Eino 后台执行立即返回，但仍属于本轮任务；无任务归属的后台启动被拒绝。前台工具退出时清理遗留子进程，跨工具调用运行的任务应使用显式后台入口。 +- 完成、失败、超时收尾、用户停止、服务正常关闭都会封闭启动入口、取消 worker、终止进程并等待回收。“中断并继续”保留本轮归属。 +- Unix 首先给进程组 3 秒退出宽限期，再强制终止并验证；内核隔离、守护进程退出与 worker 收尾各有有界等待。清理期间仍占用本轮任务；失败保留 `cleanup_failed`、进程组/worker ID，每 15 秒巡检重试。 +- 工具 worker 的登记与任务关闭互斥；即使 MCP 使用 `WithoutCancel`，任务也会等待其退出。旧轮次的延迟回调无法接管新轮次的进程或提交新 worker。 + +### 操作系统隔离和崩溃回收 + +| 平台 | 实现 | 崩溃回收与限制 | +| --- | --- | --- | +| Linux，已配置 cgroup v2 | 每轮创建独立 cgroup，使用 `clone3(CLONE_INTO_CGROUP)` 在创建时归组；配置进程数、内存、CPU 限额 | 独立守护进程在主程序管道 EOF 后写 `cgroup.kill`，验证 `populated=0`；启动时独占委派根并回收遗留任务组。`setsid` 不会脱离 cgroup | +| Windows | 每轮 Job Object，守护进程先加入 Job；任务通过原子父进程属性继承 Job，避免启动后再分配的竞态 | `KILL_ON_JOB_CLOSE`、主动 `TerminateJobObject` 和活跃进程数验证；进程数、内存与可选 CPU 限额 | +| macOS／未配置 cgroup 的 Unix | 独立进程组和守护进程，子进程在登记确认前通过管道等待 | 主程序被强杀后，管道 EOF 触发整组清理；无法限制主动 `setsid` 逃逸，不能作为强隔离部署 | + +守护进程自身异常退出时，仍存活的宿主会终止相应资源；IPC 有超时。Linux 的遗留回收依据独占目录与随机任务标识，不重放历史 PID。 + +这属于生命周期隔离，不是针对同权限恶意代码的完整安全沙箱。能够修改 cgroup、取得外部父进程句柄或杀死宿主与守护进程的程序仍需要容器/不同操作系统身份及权限策略限制。Linux 容器部署应使用 init 回收孤儿进程。`required` 模式在缺少相应内核能力或权限时拒绝运行，不会悄悄降级。 + +### 配置与上线 + +默认 `auto` 在 Windows 使用 Job Object；Unix 未指定 cgroup 根时使用进程组守护。Linux 生产部署需配置独立的 systemd 服务，设置 `Delegate=cpu memory pids`、`KillMode=control-group` 和 `Restart=on-failure`，将以下严格隔离配置 **合并到现有配置**： + +```yaml +security: + process_isolation: + mode: required + cgroup_root: auto + max_processes: 256 + memory_max_bytes: 2147483648 + cpu_quota_micros: 200000 +``` + +需要 Linux 5.14+、cgroup v2、允许 `clone3`，以及 `Delegate=cpu memory pids`。`cgroup_root: auto` 用于 systemd 的独立委派单元；也可指定受控的绝对路径，禁止使用整台主机的层级根。配置改变需要重启。Windows 留空 `cgroup_root`，可以使用 `required`；macOS 的 `required` 会明确报错。 + +服务启动时会真正创建并回收一个探测进程，验证权限和内核支持；也可在不启动 HTTP/MCP 服务的情况下单独运行： + +```sh +./cyberstrike-ai --check-process-isolation -config config.yaml +``` + +输出 `cgroup_v2`、`windows_job` 或 `process_group_watchdog`。运行/历史任务 API 的 `isolationBackend` 字段提供每轮实际采用的后端。systemd 的 `KillMode=control-group` 和失败重启为整个服务额外兜底；它与每任务 cgroup 配合使用。 + +### 远端 MCP 取消 + +普通 MCP 取消通知不提供远端退出回执。适配器可实现 `ExternalCancellationConfirmer`，使用服务端任务状态、租约或取消回执确认结束。没有确认时，不再标成“已终止”：工具记录持久化为 `orphaned` 并说明原因；本地 worker 和进程清理完成后，任务历史保留 `cleanup_unconfirmed`，响应明确提示远端状态待确认。它不是自动重试成功，也不是远端零残留保证。仍未返回的本地 worker 则保持 `cleanup_failed` 并继续追踪。 + +### 回归验证 + +```sh +go test -race ./internal/processguard ./internal/mcp ./internal/handler ./internal/security +``` + +`Process isolation` GitHub Actions 工作流还会执行 Linux cgroup 集成测试，命令直接保存在 `.github/workflows/process-isolation.yml` 中。该测试使用一次性、私有 cgroup 命名空间的 Docker 容器，不挂载宿主 cgroup 或 Docker socket。测试覆盖宿主 `SIGKILL`、启动登记竞态、`setsid`、资源限额、启动时遗留回收、worker 取消和远端未确认状态。Windows Job Object 测试已纳入跨平台 CI；交叉编译不等于 Windows 实机测试。 + +参考：[Go os/exec](https://pkg.go.dev/os/exec)、[Linux cgroup v2](https://cdn.kernel.org/doc/html/latest/admin-guide/cgroup-v2.html)、[Windows Job Objects](https://learn.microsoft.com/en-us/windows/win32/procthread/job-objects)、[MCP cancellation](https://modelcontextprotocol.io/specification/2025-11-25/basic/utilities/cancellation)。 diff --git a/go.mod b/go.mod index 006cf5ba..894b3864 100644 --- a/go.mod +++ b/go.mod @@ -1,7 +1,6 @@ module cyberstrike-ai // 若 go mod download 超时，可执行: go env -w GOPROXY=https://goproxy.cn,direct -// 或使用 scripts/bootstrap-go.sh go 1.25.0 diff --git a/internal/app/app.go b/internal/app/app.go index 2ca1417e..1200c94d 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -745,6 +745,9 @@ func (a *App) RunWithContext(ctx context.Context) error { // Shutdown 关闭应用 func (a *App) Shutdown() { + if a.agentHandler != nil { + a.agentHandler.ShutdownTasks() + } shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second) _ = einoobserve.ShutdownOtel(shutdownCtx) shutdownCancel() diff --git a/internal/config/config.go b/internal/config/config.go index 4fb988a4..283ad029 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -1073,7 +1073,17 @@ type SpaceSearchConfig struct { BaseURL string `yaml:"base_url,omitempty" json:"base_url,omitempty"` } +type ProcessIsolationConfig struct { + Mode string `yaml:"mode" json:"mode"` + CgroupRoot string `yaml:"cgroup_root" json:"cgroup_root"` + MaxProcesses int `yaml:"max_processes" json:"max_processes"` + MemoryMaxBytes int64 `yaml:"memory_max_bytes" json:"memory_max_bytes"` + CPUQuotaMicros int64 `yaml:"cpu_quota_micros" json:"cpu_quota_micros"` +} + type SecurityConfig struct { + ProcessIsolation ProcessIsolationConfig `yaml:"process_isolation,omitempty" json:"process_isolation"` + Tools []ToolConfig `yaml:"tools,omitempty"` // 向后兼容：支持在主配置文件中定义工具 ToolsDir string `yaml:"tools_dir,omitempty"` // 工具配置文件目录（新方式） ToolDescriptionMode string `yaml:"tool_description_mode,omitempty"` // 工具描述模式: "short" | "full"，默认 short diff --git a/internal/handler/agent.go b/internal/handler/agent.go index bae168e6..6f270e6b 100644 --- a/internal/handler/agent.go +++ b/internal/handler/agent.go @@ -221,13 +221,19 @@ func (h *AgentHandler) CancelRunningTaskForConversation(conversationID string) { if h == nil || conversationID == "" || h.tasks == nil { return } - h.cancelRunningMCPToolsForConversation(conversationID) - h.tasks.AbortActiveEinoExecute(conversationID, "") - if ok, err := h.tasks.CancelTask(conversationID, ErrTaskCancelled); ok { - h.logger.Info("已取消会话运行中任务", zap.String("conversationId", conversationID)) - } else if err != nil { - h.logger.Warn("取消会话运行中任务失败", zap.String("conversationId", conversationID), zap.Error(err)) + ok, err := h.tasks.CancelTask(conversationID, ErrTaskCancelled) + if !ok { + h.cancelRunningMCPToolsForConversation(conversationID) + h.tasks.AbortActiveEinoExecute(conversationID, "") } + if h.logger != nil { + if err != nil { + h.logger.Warn("取消会话运行中任务失败", zap.String("conversationId", conversationID), zap.Error(err)) + } else if ok { + h.logger.Info("已取消会话运行中任务", zap.String("conversationId", conversationID)) + } + } + } // ConversationTaskRuntimeState exposes the authoritative live state and start @@ -893,15 +899,24 @@ func (h *AgentHandler) ProcessMessageForRobot(ctx context.Context, platform stri taskCtx, cancelWithCause := context.WithCancelCause(ctx) defer cancelWithCause(nil) taskStatus := "completed" + var taskRunID string defer func() { - h.tasks.FinishTask(conversationID, taskStatus) + if taskRunID == "" { + return + } + if cleanupErr := h.tasks.FinishTaskRun(conversationID, taskRunID, taskStatus); cleanupErr != nil { + err = errors.Join(err, cleanupErr) + } }() - if _, err := h.tasks.StartTask(conversationID, message, cancelWithCause); err != nil { + if startedTask, err := h.tasks.StartTask(conversationID, message, cancelWithCause); err != nil { if errors.Is(err, ErrTaskAlreadyRunning) { return "", conversationID, fmt.Errorf("当前会话已有任务正在执行中，请稍后再试") } return "", conversationID, fmt.Errorf("无法启动任务: %w", err) + } else { + taskRunID = startedTask.RunID } + taskCtx = h.tasks.BindProcessScope(taskCtx, conversationID, taskRunID) progressCallback := h.createProgressCallback(taskCtx, cancelWithCause, conversationID, assistantMessageID, nil) robotMode := config.NormalizeAgentMode(agentMode) diff --git a/internal/handler/batch_queue_executor.go b/internal/handler/batch_queue_executor.go index 7f1ae0be..8d93fc14 100644 --- a/internal/handler/batch_queue_executor.go +++ b/internal/handler/batch_queue_executor.go @@ -171,19 +171,14 @@ func (h *AgentHandler) executeOneBatchSubTask(queueID string, queue *BatchTaskQu taskCtx, timeoutCancel := context.WithTimeout(baseCtx, 6*time.Hour) registered := false + var taskRunID string finishStatus := "completed" defer func() { h.batchTaskManager.SetTaskCancel(queueID, task.ID, nil) timeoutCancel() if registered { - if h.taskEventBus != nil { - ev := StreamEvent{Type: "done", Message: "", Data: map[string]interface{}{"conversationId": conversationID}} - if b, err := json.Marshal(ev); err == nil { - h.taskEventBus.Publish(conversationID, append(append([]byte("data: "), b...), '\n', '\n')) - } - } - h.tasks.FinishTask(conversationID, finishStatus) + h.tasks.FinishTaskRun(conversationID, taskRunID, finishStatus) } cancelWithCause(nil) }() @@ -204,7 +199,7 @@ func (h *AgentHandler) executeOneBatchSubTask(queueID string, queue *BatchTaskQu h.taskEventBus.Publish(conversationID, line) } - if _, err := h.tasks.StartTask(conversationID, task.Message, cancelWithCause); err != nil { + if startedTask, err := h.tasks.StartTask(conversationID, task.Message, cancelWithCause); err != nil { h.logger.Warn("批量队列子任务注册会话运行状态失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), @@ -216,7 +211,11 @@ func (h *AgentHandler) executeOneBatchSubTask(queueID string, queue *BatchTaskQu } h.batchTaskManager.UpdateTaskStatus(queueID, task.ID, BatchTaskStatusFailed, "", failMsg) return + } else { + taskRunID = startedTask.RunID } + baseCtx = h.tasks.BindProcessScope(baseCtx, conversationID, taskRunID) + taskCtx = h.tasks.BindProcessScope(taskCtx, conversationID, taskRunID) registered = true h.batchTaskManager.SetTaskCancel(queueID, task.ID, timeoutCancel) @@ -342,6 +341,10 @@ func (h *AgentHandler) executeOneBatchSubTask(queueID string, queue *BatchTaskQu } } + if cleanupErr := h.tasks.FinishTaskRun(conversationID, taskRunID, finishStatus); cleanupErr != nil { + h.batchTaskManager.UpdateTaskStatusWithConversationID(queueID, task.ID, BatchTaskStatusFailed, resText, cleanupErr.Error(), conversationID) + return + } if !decision.Finalizable { h.batchTaskManager.UpdateTaskStatusWithConversationID(queueID, task.ID, BatchTaskStatusFailed, resText, finalizationCheckMessage(decision), conversationID) return diff --git a/internal/handler/eino_single_agent.go b/internal/handler/eino_single_agent.go index 7e4c7f5d..88ea6691 100644 --- a/internal/handler/eino_single_agent.go +++ b/internal/handler/eino_single_agent.go @@ -130,9 +130,10 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) { // 仅在成功 StartTask 后再 FinishTask。若 StartTask 因 ErrTaskAlreadyRunning 失败仍 defer FinishTask， // 会误删其他连接上正在运行的同会话任务，导致「第一次拦截、第二次却放行」。 taskOwned := false + var taskRunID string defer func() { if taskOwned { - h.tasks.FinishTask(conversationID, taskStatus) + h.tasks.FinishTaskRun(conversationID, taskRunID, taskStatus) } }() @@ -165,7 +166,7 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) { baseCtx, cancelWithCause = context.WithCancelCause(detachedAgentContext(c.Request.Context())) taskCtx, timeoutCancel := context.WithTimeout(baseCtx, 600*time.Minute) - if _, err := h.tasks.StartTask(conversationID, req.Message, cancelWithCause); err != nil { + if startedTask, err := h.tasks.StartTask(conversationID, req.Message, cancelWithCause); err != nil { var errorMsg string if errors.Is(err, ErrTaskAlreadyRunning) { errorMsg = "⚠️ 当前会话已有任务正在执行中，请等待当前任务完成或点击「停止任务」后再尝试。" @@ -183,8 +184,13 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) { sendEvent("done", "", map[string]interface{}{"conversationId": conversationID}) timeoutCancel() return + } else { + taskRunID = startedTask.RunID } + baseCtx = h.tasks.BindProcessScope(baseCtx, conversationID, taskRunID) + taskCtx = h.tasks.BindProcessScope(taskCtx, conversationID, taskRunID) taskOwned = true + sendEvent = h.taskFinishingEventSender(sendEvent, conversationID, taskRunID, func() string { return taskStatus }) var cumulativeMCPExecutionIDs []string // 同一请求内分段续跑时，主代理 iteration 事件按偏移累计，避免 UI 出现「第3轮 → 第1轮」回跳。 @@ -454,18 +460,31 @@ func (h *AgentHandler) EinoSingleAgentLoop(c *gin.Context) { defer cancelWithCause(nil) taskCtx, timeoutCancel := context.WithTimeout(baseCtx, 600*time.Minute) defer timeoutCancel() + jsonTask, startErr := h.tasks.StartTask(prep.ConversationID, req.Message, cancelWithCause) + if startErr != nil { + c.JSON(http.StatusConflict, gin.H{"error": startErr.Error()}) + return + } + taskCtx = h.tasks.BindProcessScope(taskCtx, prep.ConversationID, jsonTask.RunID) + taskCtx = mcp.WithMCPConversationID(taskCtx, prep.ConversationID) + taskCtx = mcp.WithToolRunRegistry(taskCtx, h.tasks) + taskCtx = mcp.WithEinoExecuteRunRegistry(taskCtx, h.tasks) + jsonTaskStatus := "failed" + defer func() { _ = h.tasks.FinishTaskRun(prep.ConversationID, jsonTask.RunID, jsonTaskStatus) }() + respond := h.taskFinishingJSONResponder(c, prep.ConversationID, jsonTask.RunID, func() string { return jsonTaskStatus }) + progressCallback := h.createProgressCallback(taskCtx, cancelWithCause, prep.ConversationID, prep.AssistantMessageID, progressCallbackRaw) taskCtx = multiagent.WithHITLToolInterceptor(taskCtx, func(ctx context.Context, toolName, arguments string) (string, error) { return h.interceptHITLForEinoTool(ctx, cancelWithCause, prep.ConversationID, prep.AssistantMessageID, nil, toolName, arguments) }) if h.config == nil { - c.JSON(http.StatusInternalServerError, gin.H{"error": "服务器配置未加载"}) + respond(http.StatusInternalServerError, gin.H{"error": "服务器配置未加载"}) return } runCfg, _, err := h.configForAIChannel(req.AIChannelID) if err != nil { - c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + respond(http.StatusBadRequest, gin.H{"error": err.Error()}) return } @@ -498,7 +517,7 @@ func (h *AgentHandler) EinoSingleAgentLoop(c *gin.Context) { if shouldPersistEinoAgentTraceAfterRunError(baseCtx) { h.persistEinoAgentTraceForResume(prep.ConversationID, result) } - c.JSON(http.StatusInternalServerError, gin.H{"error": runErr.Error()}) + respond(http.StatusInternalServerError, gin.H{"error": runErr.Error()}) return } mw := &h.config.MultiAgent.EinoMiddleware @@ -525,7 +544,12 @@ func (h *AgentHandler) EinoSingleAgentLoop(c *gin.Context) { if !decision.Finalizable { responseText = finalizationBlockedMessage(decision) } - c.JSON(http.StatusOK, gin.H{ + + jsonTaskStatus = decision.Status + if jsonTaskStatus == "" { + jsonTaskStatus = "completed" + } + respond(http.StatusOK, gin.H{ "response": responseText, "conversationId": prep.ConversationID, "mcpExecutionIds": result.MCPExecutionIDs, diff --git a/internal/handler/multi_agent.go b/internal/handler/multi_agent.go index 5c4bf392..6e7a86bd 100644 --- a/internal/handler/multi_agent.go +++ b/internal/handler/multi_agent.go @@ -147,9 +147,10 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) { taskStatus := "completed" // 仅在成功 StartTask 后再 FinishTask；避免「任务已存在」分支 return 时误删正在运行的同会话任务。 taskOwned := false + var taskRunID string defer func() { if taskOwned { - h.tasks.FinishTask(conversationID, taskStatus) + h.tasks.FinishTaskRun(conversationID, taskRunID, taskStatus) } }() @@ -172,7 +173,7 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) { baseCtx, cancelWithCause = context.WithCancelCause(detachedAgentContext(c.Request.Context())) taskCtx, timeoutCancel := context.WithTimeout(baseCtx, 600*time.Minute) - if _, err := h.tasks.StartTask(conversationID, req.Message, cancelWithCause); err != nil { + if startedTask, err := h.tasks.StartTask(conversationID, req.Message, cancelWithCause); err != nil { var errorMsg string if errors.Is(err, ErrTaskAlreadyRunning) { errorMsg = "⚠️ 当前会话已有任务正在执行中，请等待当前任务完成或点击「停止任务」后再尝试。" @@ -190,8 +191,13 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) { sendEvent("done", "", map[string]interface{}{"conversationId": conversationID}) timeoutCancel() return + } else { + taskRunID = startedTask.RunID } + baseCtx = h.tasks.BindProcessScope(baseCtx, conversationID, taskRunID) + taskCtx = h.tasks.BindProcessScope(taskCtx, conversationID, taskRunID) taskOwned = true + sendEvent = h.taskFinishingEventSender(sendEvent, conversationID, taskRunID, func() string { return taskStatus }) // 同一 HTTP 流内多段 Run（如中断并继续）合并 MCP execution id，供最终 response / 库表与工具芯片展示完整列表 var cumulativeMCPExecutionIDs []string @@ -468,13 +474,26 @@ func (h *AgentHandler) MultiAgentLoop(c *gin.Context) { defer cancelWithCause(nil) taskCtx, timeoutCancel := context.WithTimeout(baseCtx, 600*time.Minute) defer timeoutCancel() + jsonTask, startErr := h.tasks.StartTask(prep.ConversationID, req.Message, cancelWithCause) + if startErr != nil { + c.JSON(http.StatusConflict, gin.H{"error": startErr.Error()}) + return + } + taskCtx = h.tasks.BindProcessScope(taskCtx, prep.ConversationID, jsonTask.RunID) + taskCtx = mcp.WithMCPConversationID(taskCtx, prep.ConversationID) + taskCtx = mcp.WithToolRunRegistry(taskCtx, h.tasks) + taskCtx = mcp.WithEinoExecuteRunRegistry(taskCtx, h.tasks) + jsonTaskStatus := "failed" + defer func() { _ = h.tasks.FinishTaskRun(prep.ConversationID, jsonTask.RunID, jsonTaskStatus) }() + respond := h.taskFinishingJSONResponder(c, prep.ConversationID, jsonTask.RunID, func() string { return jsonTaskStatus }) + progressCallback := h.createProgressCallback(taskCtx, cancelWithCause, prep.ConversationID, prep.AssistantMessageID, nil) taskCtx = multiagent.WithHITLToolInterceptor(taskCtx, func(ctx context.Context, toolName, arguments string) (string, error) { return h.interceptHITLForEinoTool(ctx, cancelWithCause, prep.ConversationID, prep.AssistantMessageID, nil, toolName, arguments) }) runCfg, _, err := h.configForAIChannel(req.AIChannelID) if err != nil { - c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + respond(http.StatusBadRequest, gin.H{"error": err.Error()}) return } @@ -522,7 +541,7 @@ func (h *AgentHandler) MultiAgentLoop(c *gin.Context) { } errData := multiagent.EinoClientRunErrorFields(runErr) errData["error"] = errMsg - c.JSON(http.StatusInternalServerError, errData) + respond(http.StatusInternalServerError, errData) return } mw := &h.config.MultiAgent.EinoMiddleware @@ -552,7 +571,12 @@ func (h *AgentHandler) MultiAgentLoop(c *gin.Context) { if !decision.Finalizable { responseText = finalizationBlockedMessage(decision) } - c.JSON(http.StatusOK, ChatResponse{ + + jsonTaskStatus = decision.Status + if jsonTaskStatus == "" { + jsonTaskStatus = "completed" + } + respond(http.StatusOK, ChatResponse{ Response: responseText, MCPExecutionIDs: result.MCPExecutionIDs, ConversationID: prep.ConversationID, diff --git a/internal/handler/task_lifecycle.go b/internal/handler/task_lifecycle.go new file mode 100644 index 00000000..a9a75091 --- /dev/null +++ b/internal/handler/task_lifecycle.go @@ -0,0 +1,62 @@ +package handler + +import ( + "cyberstrike-ai/internal/runlease" + "errors" + "github.com/gin-gonic/gin" + "go.uber.org/zap" + "net/http" +) + +// ShutdownTasks stops local work before shared MCP clients and databases close. +func (h *AgentHandler) ShutdownTasks() { + if h != nil && h.tasks != nil { + h.tasks.Shutdown() + if h.logger != nil { + for _, task := range h.tasks.GetActiveTasks() { + h.logger.Warn("服务关闭时任务仍有未完成的清理", zap.String("runId", task.RunID), zap.String("cleanupError", task.CleanupError)) + } + } + } +} + +// taskFinishingEventSender makes the visible done event follow local cleanup. +func (h *AgentHandler) taskFinishingEventSender(send func(string, string, interface{}), conversationID, runID string, status func() string) func(string, string, interface{}) { + return func(eventType, message string, data interface{}) { + if eventType == "done" { + if err := h.tasks.FinishTaskRun(conversationID, runID, status()); err != nil { + if h.logger != nil { + h.logger.Warn(taskCleanupMessage(err), zap.String("runId", runID), zap.Error(err)) + } + send("error", taskCleanupMessage(err)+": "+err.Error(), map[string]interface{}{"errorType": taskCleanupStatus(err)}) + data = map[string]interface{}{"conversationId": conversationID, "runId": runID, "status": taskCleanupStatus(err), "cleanupError": err.Error()} + } + } + send(eventType, message, data) + } +} + +// taskFinishingJSONResponder applies the same ordering to successful and failed +// non-streaming requests. No response claims completion before cleanup returns. +func (h *AgentHandler) taskFinishingJSONResponder(c *gin.Context, conversationID, runID string, status func() string) func(int, interface{}) { + return func(code int, payload interface{}) { + if err := h.tasks.FinishTaskRun(conversationID, runID, status()); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error(), "status": taskCleanupStatus(err), "conversationId": conversationID}) + return + } + c.JSON(code, payload) + } +} + +func taskCleanupStatus(err error) string { + if errors.Is(err, runlease.ErrUnconfirmed) { + return "cleanup_unconfirmed" + } + return "cleanup_failed" +} +func taskCleanupMessage(err error) string { + if errors.Is(err, runlease.ErrUnconfirmed) { + return "本地执行已结束，远端 MCP 停止状态待确认" + } + return "任务资源清理失败，将自动重试" +} diff --git a/internal/handler/task_manager.go b/internal/handler/task_manager.go index 24d35b40..c55c5c44 100644 --- a/internal/handler/task_manager.go +++ b/internal/handler/task_manager.go @@ -2,13 +2,17 @@ package handler import ( "context" + "encoding/json" "errors" + "fmt" "sort" "strings" "sync" "time" "cyberstrike-ai/internal/multiagent" + "cyberstrike-ai/internal/runlease" + "cyberstrike-ai/internal/security" ) // ErrTaskCancelled 用户取消任务的错误 @@ -26,12 +30,20 @@ func shouldPersistEinoAgentTraceAfterRunError(baseCtx context.Context) bool { // AgentTask 描述正在运行的Agent任务 type AgentTask struct { - ConversationID string `json:"conversationId"` - Title string `json:"title,omitempty"` - Message string `json:"message,omitempty"` - StartedAt time.Time `json:"startedAt"` - Status string `json:"status"` - CancellingAt time.Time `json:"-"` // 进入 cancelling 状态的时间，用于清理长时间卡住的任务 + RunID string `json:"runId"` + CleanupError string `json:"cleanupError,omitempty"` + processes *security.ProcessScope + workers *runlease.Scope + IsolationBackend string `json:"isolationBackend,omitempty"` + finishing chan struct{} + stopping chan struct{} + finalStatus string + ConversationID string `json:"conversationId"` + Title string `json:"title,omitempty"` + Message string `json:"message,omitempty"` + StartedAt time.Time `json:"startedAt"` + Status string `json:"status"` + CancellingAt time.Time `json:"-"` // 进入 cancelling 状态的时间，用于清理长时间卡住的任务 // ActiveMCPExecutionID 当前正在执行的 MCP 工具 executionId（仅内存，供「中断并继续」= 仅掐当前工具） ActiveMCPExecutionID string `json:"-"` @@ -297,12 +309,15 @@ func (m *AgentTaskManager) ActiveMCPExecutionID(conversationID string) string { // CompletedTask 已完成的任务（用于历史记录） type CompletedTask struct { - ConversationID string `json:"conversationId"` - Title string `json:"title,omitempty"` - Message string `json:"message,omitempty"` - StartedAt time.Time `json:"startedAt"` - CompletedAt time.Time `json:"completedAt"` - Status string `json:"status"` + CleanupError string `json:"cleanupError,omitempty"` + IsolationBackend string `json:"isolationBackend,omitempty"` + RunID string `json:"runId"` + ConversationID string `json:"conversationId"` + Title string `json:"title,omitempty"` + Message string `json:"message,omitempty"` + StartedAt time.Time `json:"startedAt"` + CompletedAt time.Time `json:"completedAt"` + Status string `json:"status"` } // AgentTaskManager 管理正在运行的Agent任务 @@ -315,6 +330,9 @@ type AgentTaskManager struct { eventBus *TaskEventBus // 可选：任务结束时关闭镜像 SSE 订阅 // toolCanceler 在用户整轮停止任务或会话结束时终止该会话仍在运行的 MCP 工具（非「中断并继续」）。 toolCanceler func(conversationID string) + shuttingDown bool + shutdown chan struct{} + shutdownOnce sync.Once } const ( @@ -330,6 +348,7 @@ const ( func NewAgentTaskManager() *AgentTaskManager { m := &AgentTaskManager{ tasks: make(map[string]*AgentTask), + shutdown: make(chan struct{}), completedTasks: make([]*CompletedTask, 0), maxHistorySize: 50, // 最多保留50条历史记录 historyRetention: 24 * time.Hour, // 保留24小时 @@ -368,6 +387,7 @@ func (m *AgentTaskManager) GetTaskSnapshot(conversationID string) *AgentTask { return nil } snapshot := *task + snapshot.IsolationBackend = task.processes.IsolationBackend() return &snapshot } @@ -375,16 +395,26 @@ func (m *AgentTaskManager) GetTaskSnapshot(conversationID string) *AgentTask { func (m *AgentTaskManager) runStuckCancellingCleanup() { ticker := time.NewTicker(cleanupInterval) defer ticker.Stop() - for range ticker.C { - m.cleanupStuckCancelling() + for { + select { + case <-m.shutdown: + return + case <-ticker.C: + m.cleanupStuckCancelling() + } } } func (m *AgentTaskManager) cleanupStuckCancelling() { m.mu.Lock() - var toFinish []string + type pendingFinish struct{ id, runID, status string } + var toFinish []pendingFinish now := time.Now() for id, task := range m.tasks { + if task.Status == "cleanup_failed" { + toFinish = append(toFinish, pendingFinish{id, task.RunID, task.finalStatus}) + continue + } if task.Status != "cancelling" { continue } @@ -400,11 +430,11 @@ func (m *AgentTaskManager) cleanupStuckCancelling() { continue } } - toFinish = append(toFinish, id) + toFinish = append(toFinish, pendingFinish{id, task.RunID, "cancelled"}) } m.mu.Unlock() - for _, id := range toFinish { - m.FinishTask(id, "cancelled") + for _, pending := range toFinish { + _ = m.FinishTaskRun(pending.id, pending.runID, pending.status) } } @@ -413,11 +443,16 @@ func (m *AgentTaskManager) StartTask(conversationID, message string, cancel cont m.mu.Lock() defer m.mu.Unlock() + if m.shuttingDown { + return nil, errors.New("task manager is shutting down") + } if _, exists := m.tasks[conversationID]; exists { return nil, ErrTaskAlreadyRunning } + scope := security.NewProcessScope() task := &AgentTask{ + RunID: scope.ID, processes: scope, workers: runlease.New(), ConversationID: conversationID, Message: message, StartedAt: time.Now(), @@ -444,7 +479,7 @@ func (m *AgentTaskManager) CancelTask(conversationID string, cause error) (bool, } // 如果已经处于取消流程，视为成功（幂等），避免前端重复点击报「未找到任务」 - if task.Status == "cancelling" { + if task.Status == "cancelling" || task.finishing != nil { m.mu.Unlock() return true, nil } @@ -466,6 +501,13 @@ func (m *AgentTaskManager) CancelTask(conversationID string, cause error) (bool, interruptPush := task.agentTurnLoopInterrupt interruptNote := task.InterruptContinueNote runtimeCancel := task.agentRuntimeCancel + activeExecuteCancel := task.activeEinoExecuteCancel + if !errors.Is(cause, multiagent.ErrInterruptContinue) { + task.processes.Seal() + task.workers.Seal() + task.stopping = make(chan struct{}) + defer close(task.stopping) + } var toolCanceler func(string) if errors.Is(cause, ErrTaskCancelled) { toolCanceler = m.toolCanceler @@ -494,6 +536,13 @@ func (m *AgentTaskManager) CancelTask(conversationID string, cause error) (bool, if toolCanceler != nil { toolCanceler(conversationID) } + if !errors.Is(cause, multiagent.ErrInterruptContinue) { + task.workers.Cancel() + if activeExecuteCancel != nil { + activeExecuteCancel() + } + return true, task.processes.Close() + } return true, nil } @@ -507,54 +556,172 @@ func (m *AgentTaskManager) UpdateTaskStatus(conversationID string, status string return } - if status != "" { - task.Status = status + if task.finishing != nil || task.Status == "cleanup_failed" { + return + } + switch status { + case "completed", "cancelled", "failed", "timeout": + task.finalStatus = status + task.Status = "cleaning" + task.processes.Seal() + task.workers.Seal() + default: + if status != "" { + task.Status = status + } } } -// FinishTask 完成任务并从管理器中移除 +// BindProcessScope snapshots ownership once at task start. Continuations must +// derive from this context, never resolve ownership again using conversation ID. +func (m *AgentTaskManager) BindProcessScope(ctx context.Context, conversationID, runID string) context.Context { + m.mu.RLock() + defer m.mu.RUnlock() + if task := m.tasks[conversationID]; task != nil && task.RunID == runID { + return runlease.WithScope(security.WithProcessScope(ctx, task.processes), task.workers) + } + // Fail closed if a task disappeared before its execution context was bound. + scope := security.NewProcessScope() + scope.Seal() + workers := runlease.New() + workers.Seal() + return runlease.WithScope(security.WithProcessScope(ctx, scope), workers) +} + +// FinishTask is retained for callers that operate on the current task. Owners +// use FinishTaskRun, so a delayed defer cannot finish a newer conversation run. func (m *AgentTaskManager) FinishTask(conversationID string, finalStatus string) { + m.mu.RLock() + task := m.tasks[conversationID] + m.mu.RUnlock() + if task != nil { + _ = m.FinishTaskRun(conversationID, task.RunID, finalStatus) + } +} + +func (m *AgentTaskManager) FinishTaskRun(conversationID, runID, finalStatus string) error { m.mu.Lock() - task, exists := m.tasks[conversationID] - if !exists { + task := m.tasks[conversationID] + if task == nil || task.RunID != runID { m.mu.Unlock() - return + return nil } - - if finalStatus != "" { - task.Status = finalStatus + if task.stopping != nil { + select { + case <-task.stopping: + default: + stopping := task.stopping + m.mu.Unlock() + <-stopping + return m.FinishTaskRun(conversationID, runID, finalStatus) + } } + if task.finishing != nil { + done := task.finishing + m.mu.Unlock() + <-done + m.mu.RLock() + cleanupError := task.CleanupError + m.mu.RUnlock() + if cleanupError != "" { + return errors.New(cleanupError) + } + return nil + } + done := make(chan struct{}) + task.finishing = done + task.finalStatus = finalStatus + task.Status = "cleaning" + task.processes.Seal() + task.workers.Seal() toolCanceler := m.toolCanceler - activeEinoExecuteCancel := task.activeEinoExecuteCancel - - // 保存到历史记录 - completedTask := &CompletedTask{ - ConversationID: task.ConversationID, - Message: task.Message, - StartedAt: task.StartedAt, - CompletedAt: time.Now(), - Status: finalStatus, - } - - // 添加到历史记录 - m.completedTasks = append(m.completedTasks, completedTask) - - // 清理过期和过多的历史记录 - m.cleanupHistory() - - // 从运行任务中移除 - delete(m.tasks, conversationID) + activeCancel := task.activeEinoExecuteCancel + cancel := task.cancel bus := m.eventBus m.mu.Unlock() + + // Keep the conversation occupied throughout cleanup, including callbacks. + if cancel != nil { + cancel(nil) + } if toolCanceler != nil { toolCanceler(conversationID) } - if activeEinoExecuteCancel != nil { - activeEinoExecuteCancel() + if activeCancel != nil { + activeCancel() } - if bus != nil { + task.workers.Cancel() + processErr := task.processes.Close() + waitCtx, waitCancel := context.WithTimeout(context.Background(), 3*time.Second) + workerErr := task.workers.Wait(waitCtx) + waitCancel() + cleanupErr := errors.Join(processErr, workerErr) + if processErr != nil && errors.Is(workerErr, runlease.ErrUnconfirmed) { + // A simultaneous local failure must not be labelled as local completion. + cleanupErr = fmt.Errorf("local cleanup: %w; remote state: %v", processErr, workerErr) + } + // The local worker has returned, but remote notification cancellation is + // not an acknowledgement. Preserve an actionable history/tool status. + unconfirmed := processErr == nil && errors.Is(workerErr, runlease.ErrUnconfirmed) + if unconfirmed { + finalStatus = "cleanup_unconfirmed" + } + cleanupMessage := "" + if cleanupErr != nil { + cleanupMessage = cleanupErr.Error() + } + if (cleanupErr == nil || unconfirmed) && bus != nil { + // Subscribers must receive completion only after local processes are reaped. + payload, _ := json.Marshal(StreamEvent{Type: "done", Data: map[string]interface{}{"conversationId": conversationID, "runId": runID, "status": finalStatus, "cleanupError": cleanupMessage}}) + bus.Publish(conversationID, append(append([]byte("data: "), payload...), '\n', '\n')) bus.CloseConversation(conversationID) } + + m.mu.Lock() + defer m.mu.Unlock() + defer close(done) + if cleanupErr != nil && !unconfirmed { + task.Status = "cleanup_failed" + task.CleanupError = cleanupErr.Error() + task.finishing = nil + return cleanupErr + } + task.CleanupError = "" + if unconfirmed { + task.CleanupError = cleanupErr.Error() + } + task.Status = finalStatus + m.completedTasks = append(m.completedTasks, &CompletedTask{ + RunID: task.RunID, CleanupError: task.CleanupError, IsolationBackend: task.processes.IsolationBackend(), ConversationID: task.ConversationID, Message: task.Message, + StartedAt: task.StartedAt, CompletedAt: time.Now(), Status: finalStatus, + }) + m.cleanupHistory() + delete(m.tasks, conversationID) + return cleanupErr +} + +// Shutdown rejects new tasks before cancelling and reaping existing task jobs. +func (m *AgentTaskManager) Shutdown() { + m.mu.Lock() + m.shuttingDown = true + m.shutdownOnce.Do(func() { close(m.shutdown) }) + tasks := make([]*AgentTask, 0, len(m.tasks)) + for _, task := range m.tasks { + tasks = append(tasks, task) + task.processes.Seal() + task.workers.Seal() + } + m.mu.Unlock() + var wg sync.WaitGroup + for _, task := range tasks { + wg.Add(1) + go func(task *AgentTask) { + defer wg.Done() + _, _ = m.CancelTask(task.ConversationID, ErrTaskCancelled) + _ = m.FinishTaskRun(task.ConversationID, task.RunID, "cancelled") + }(task) + } + wg.Wait() } // cleanupHistory 清理过期的历史记录 @@ -589,6 +756,7 @@ func (m *AgentTaskManager) GetActiveTasks() []*AgentTask { result := make([]*AgentTask, 0, len(m.tasks)) for _, task := range m.tasks { result = append(result, &AgentTask{ + RunID: task.RunID, CleanupError: task.CleanupError, IsolationBackend: task.processes.IsolationBackend(), ConversationID: task.ConversationID, Message: task.Message, StartedAt: task.StartedAt, diff --git a/internal/handler/task_process_cleanup_test.go b/internal/handler/task_process_cleanup_test.go new file mode 100644 index 00000000..f758edd2 --- /dev/null +++ b/internal/handler/task_process_cleanup_test.go @@ -0,0 +1,225 @@ +package handler + +import ( + "context" + "errors" + "os/exec" + "runtime" + "sync" + "testing" + "time" + + "cyberstrike-ai/internal/mcp" + "cyberstrike-ai/internal/runlease" + "cyberstrike-ai/internal/security" +) + +func TestTaskCleanupWaitsBeforeReleasingConversation(t *testing.T) { + manager := NewAgentTaskManager() + task, _ := manager.StartTask("conv", "old", func(error) {}) + entered, release := make(chan struct{}), make(chan struct{}) + manager.SetToolCanceler(func(string) { close(entered); <-release }) + done := make(chan error, 1) + go func() { done <- manager.FinishTaskRun("conv", task.RunID, "completed") }() + <-entered + if status := manager.GetTaskSnapshot("conv").Status; status != "cleaning" { + t.Errorf("status = %s", status) + } + if _, err := manager.StartTask("conv", "new", nil); !errors.Is(err, ErrTaskAlreadyRunning) { + t.Errorf("new task admitted during cleanup: %v", err) + } + ctx := manager.BindProcessScope(context.Background(), "conv", task.RunID) + if _, err := security.StartShellSessionContext(ctx, exec.Command("unused-command")); !errors.Is(err, security.ErrProcessScopeClosed) { + t.Errorf("late process admitted: %v", err) + } + close(release) + if err := <-done; err != nil { + t.Fatal(err) + } + manager.SetToolCanceler(nil) + next, err := manager.StartTask("conv", "new", nil) + if err != nil { + t.Fatal(err) + } + defer manager.FinishTask("conv", "completed") + if next.RunID == task.RunID { + t.Fatal("run identity reused") + } + _ = manager.FinishTaskRun("conv", task.RunID, "cancelled") + if manager.GetTaskSnapshot("conv").RunID != next.RunID { + t.Fatal("old defer removed new task") + } + // A delayed worker keeps its original closed scope, even after a new run starts. + if _, err := security.StartManagedBackground(ctx, "sh", "sleep 300", ""); !errors.Is(err, security.ErrProcessScopeClosed) { + t.Fatalf("old context borrowed new task: %v", err) + } +} + +func TestTaskFinishAndShutdownReapBackgroundProcesses(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("Unix shell") + } + for _, shutdown := range []bool{false, true} { + name := "finish" + if shutdown { + name = "shutdown" + } + t.Run(name, func(t *testing.T) { + manager := NewAgentTaskManager() + task, _ := manager.StartTask("conv", "job", nil) + ctx := manager.BindProcessScope(context.Background(), "conv", task.RunID) + session, err := security.StartManagedBackground(ctx, "sh", "sleep 300", "") + if err != nil { + t.Fatal(err) + } + if shutdown { + manager.Shutdown() + } else if err := manager.FinishTaskRun("conv", task.RunID, "completed"); err != nil { + t.Fatal(err) + } + done := make(chan struct{}) + go func() { _ = session.Wait(); close(done) }() + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("task ended before process exited") + } + if manager.GetTaskSnapshot("conv") != nil { + t.Fatal("finished task still active") + } + if shutdown { + if _, err := manager.StartTask("new", "job", nil); err == nil { + t.Fatal("shutdown admitted a new task") + } + } + }) + } +} + +func TestTaskDonePublishedAfterCleanup(t *testing.T) { + manager := NewAgentTaskManager() + bus := NewTaskEventBus() + manager.SetTaskEventBus(bus) + task, _ := manager.StartTask("conv", "job", nil) + _, events := bus.Subscribe("conv") + if err := manager.FinishTaskRun("conv", task.RunID, "completed"); err != nil { + t.Fatal(err) + } + if event, ok := <-events; !ok || len(event) == 0 { + t.Fatal("subscriber closed without done event") + } + if _, ok := <-events; ok { + t.Fatal("subscriber not closed after completion") + } +} + +func TestTaskCleanupFailureRetainsOwnershipAndRetries(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("Unix shell") + } + manager := NewAgentTaskManager() + task, _ := manager.StartTask("conv", "job", nil) + ctx := manager.BindProcessScope(context.Background(), "conv", task.RunID) + // Simulate an executor which has not yet reaped its direct child. + session, err := security.StartShellSessionContext(ctx, exec.Command("sh", "-c", "sleep 300")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { session.Terminate(); _ = session.Wait(); manager.Shutdown() }) + if err := manager.FinishTaskRun("conv", task.RunID, "completed"); err == nil { + t.Fatal("unreaped process reported as cleaned up") + } + snapshot := manager.GetTaskSnapshot("conv") + if snapshot == nil || snapshot.Status != "cleanup_failed" || snapshot.CleanupError == "" { + t.Fatalf("missing actionable cleanup state: %+v", snapshot) + } + if len(manager.GetCompletedTasks()) != 0 { + t.Fatal("cleanup failure recorded as completed") + } + if _, err := manager.StartTask("conv", "new", nil); !errors.Is(err, ErrTaskAlreadyRunning) { + t.Fatal("cleanup failure released conversation") + } + _ = session.Wait() + manager.cleanupStuckCancelling() + if manager.GetTaskSnapshot("conv") != nil { + t.Fatal("cleanup retry did not finish reaped task") + } +} + +func TestTaskFinishWaitsForCancellationCallbacks(t *testing.T) { + manager := NewAgentTaskManager() + task, _ := manager.StartTask("conv", "job", nil) + entered, release := make(chan struct{}), make(chan struct{}) + var once sync.Once + manager.SetToolCanceler(func(string) { once.Do(func() { close(entered); <-release }) }) + cancelled := make(chan struct{}) + go func() { _, _ = manager.CancelTask("conv", ErrTaskCancelled); close(cancelled) }() + <-entered + finished := make(chan struct{}) + go func() { _ = manager.FinishTaskRun("conv", task.RunID, "cancelled"); close(finished) }() + select { + case <-finished: + t.Error("task finished while old cancellation callbacks could still affect new run") + case <-time.After(30 * time.Millisecond): + } + close(release) + <-cancelled + <-finished + manager.Shutdown() +} + +func TestTaskWaitsForDetachedMCPWorker(t *testing.T) { + manager := NewAgentTaskManager() + defer manager.Shutdown() + task, _ := manager.StartTask("conv", "job", nil) + ctx := manager.BindProcessScope(context.Background(), "conv", task.RunID) + service := mcp.NewExecutionService(nil, nil) + entered, cancelled, release := make(chan struct{}), make(chan struct{}), make(chan struct{}) + _, err := service.Submit(ctx, mcp.ExecutionRequest{Run: func(ctx context.Context) (*mcp.ToolResult, error) { + close(entered) + <-ctx.Done() + close(cancelled) + <-release + return nil, ctx.Err() + }}) + if err != nil { + t.Fatal(err) + } + <-entered + done := make(chan error, 1) + go func() { done <- manager.FinishTaskRun("conv", task.RunID, "completed") }() + <-cancelled + if manager.GetTaskSnapshot("conv") == nil { + t.Error("task released while detached worker still running") + } + close(release) + if err = <-done; err != nil { + t.Fatal(err) + } +} + +func TestTaskReportsUnconfirmedRemoteCleanup(t *testing.T) { + manager := NewAgentTaskManager() + defer manager.Shutdown() + task, _ := manager.StartTask("conv", "job", nil) + ctx := manager.BindProcessScope(context.Background(), "conv", task.RunID) + service := mcp.NewExecutionService(nil, nil) + entered := make(chan struct{}) + _, err := service.Submit(ctx, mcp.ExecutionRequest{Remote: true, Run: func(ctx context.Context) (*mcp.ToolResult, error) { + close(entered) + <-ctx.Done() + return nil, ctx.Err() + }}) + if err != nil { + t.Fatal(err) + } + <-entered + err = manager.FinishTaskRun("conv", task.RunID, "completed") + if !errors.Is(err, runlease.ErrUnconfirmed) { + t.Fatalf("remote cancellation reported as verified: %v", err) + } + history := manager.GetCompletedTasks() + if len(history) != 1 || history[0].Status != "cleanup_unconfirmed" || history[0].CleanupError == "" { + t.Fatalf("missing retained warning: %+v", history) + } +} diff --git a/internal/handler/workflow_integration.go b/internal/handler/workflow_integration.go index bdbdc894..67ff0cfa 100644 --- a/internal/handler/workflow_integration.go +++ b/internal/handler/workflow_integration.go @@ -55,9 +55,10 @@ func (h *AgentHandler) runRoleWorkflowStreamIfBound( taskStatus := "completed" taskOwned := false + var taskRunID string defer func() { if taskOwned { - h.tasks.FinishTask(conversationID, taskStatus) + h.tasks.FinishTaskRun(conversationID, taskRunID, taskStatus) } }() @@ -69,7 +70,7 @@ func (h *AgentHandler) runRoleWorkflowStreamIfBound( taskCtx, timeoutCancel := context.WithTimeout(baseCtx, 600*time.Minute) defer timeoutCancel() - if _, err := h.tasks.StartTask(conversationID, userMessage, cancelWithCause); err != nil { + if startedTask, err := h.tasks.StartTask(conversationID, userMessage, cancelWithCause); err != nil { var errorMsg string if errors.Is(err, ErrTaskAlreadyRunning) { errorMsg = "⚠️ 当前会话已有任务正在执行中，请等待当前任务完成或点击「停止任务」后再尝试。" @@ -86,8 +87,13 @@ func (h *AgentHandler) runRoleWorkflowStreamIfBound( } sendEvent("done", "", map[string]interface{}{"conversationId": conversationID}) return true + } else { + taskRunID = startedTask.RunID } + baseCtx = h.tasks.BindProcessScope(baseCtx, conversationID, taskRunID) + taskCtx = h.tasks.BindProcessScope(taskCtx, conversationID, taskRunID) taskOwned = true + sendEvent = h.taskFinishingEventSender(sendEvent, conversationID, taskRunID, func() string { return taskStatus }) progress := h.createProgressCallback(taskCtx, cancelWithCause, conversationID, assistantMessageID, sendEvent) result, err := workflowrunner.RunRoleBoundWorkflow(taskCtx, workflowrunner.RunArgs{ @@ -202,9 +208,10 @@ func (h *AgentHandler) runRoleWorkflowJSONIfBound(c *gin.Context, req *ChatReque taskStatus := "completed" taskOwned := false + var taskRunID string defer func() { if taskOwned { - h.tasks.FinishTask(conversationID, taskStatus) + h.tasks.FinishTaskRun(conversationID, taskRunID, taskStatus) } }() @@ -213,7 +220,7 @@ func (h *AgentHandler) runRoleWorkflowJSONIfBound(c *gin.Context, req *ChatReque taskCtx, timeoutCancel := context.WithTimeout(baseCtx, 600*time.Minute) defer timeoutCancel() - if _, err := h.tasks.StartTask(conversationID, userMessage, cancelWithCause); err != nil { + if startedTask, err := h.tasks.StartTask(conversationID, userMessage, cancelWithCause); err != nil { if errors.Is(err, ErrTaskAlreadyRunning) { c.JSON(http.StatusConflict, gin.H{ "error": "⚠️ 当前会话已有任务正在执行中，请等待当前任务完成或点击「停止任务」后再尝试。", @@ -224,8 +231,13 @@ func (h *AgentHandler) runRoleWorkflowJSONIfBound(c *gin.Context, req *ChatReque c.JSON(http.StatusInternalServerError, gin.H{"error": "❌ 无法启动任务: " + err.Error()}) } return true + } else { + taskRunID = startedTask.RunID } + baseCtx = h.tasks.BindProcessScope(baseCtx, conversationID, taskRunID) + taskCtx = h.tasks.BindProcessScope(taskCtx, conversationID, taskRunID) taskOwned = true + respond := h.taskFinishingJSONResponder(c, conversationID, taskRunID, func() string { return taskStatus }) progress := h.createProgressCallback(taskCtx, cancelWithCause, conversationID, assistantMessageID, nil) result, err := workflowrunner.RunRoleBoundWorkflow(taskCtx, workflowrunner.RunArgs{ @@ -253,7 +265,7 @@ func (h *AgentHandler) runRoleWorkflowJSONIfBound(c *gin.Context, req *ChatReque _ = h.appendAssistantMessageNotice(assistantMessageID, cancelMsg) _ = h.db.AddProcessDetail(assistantMessageID, conversationID, "cancelled", cancelMsg, nil) } - c.JSON(http.StatusOK, gin.H{ + respond(http.StatusOK, gin.H{ "status": "cancelled", "message": cancelMsg, "conversationId": conversationID, @@ -265,7 +277,7 @@ func (h *AgentHandler) runRoleWorkflowJSONIfBound(c *gin.Context, req *ChatReque if assistantMessageID != "" { _, _ = h.db.Exec("UPDATE messages SET content = ?, updated_at = ? WHERE id = ?", errMsg, time.Now(), assistantMessageID) } - c.JSON(http.StatusInternalServerError, gin.H{"error": errMsg, "conversationId": conversationID}) + respond(http.StatusInternalServerError, gin.H{"error": errMsg, "conversationId": conversationID}) return true } decision := h.finalizeCandidateForDeliveryWithPolicy( @@ -283,7 +295,7 @@ func (h *AgentHandler) runRoleWorkflowJSONIfBound(c *gin.Context, req *ChatReque responseText = finalizationBlockedMessage(decision) taskStatus = decision.Status } - c.JSON(http.StatusOK, gin.H{ + respond(http.StatusOK, gin.H{ "response": responseText, "conversationId": prep.ConversationID, "assistantMessageId": prep.AssistantMessageID, diff --git a/internal/mcp/execution_ownership_test.go b/internal/mcp/execution_ownership_test.go new file mode 100644 index 00000000..5b7a60ac --- /dev/null +++ b/internal/mcp/execution_ownership_test.go @@ -0,0 +1,76 @@ +package mcp + +import ( + "context" + "cyberstrike-ai/internal/runlease" + "errors" + "testing" + "time" +) + +func TestExecutionOwnedAfterContextDetached(t *testing.T) { + scope := runlease.New() + parent, cancel := context.WithCancel(runlease.WithScope(context.Background(), scope)) + service := NewExecutionService(nil, nil) + entered := make(chan struct{}) + handle, err := service.Submit(parent, ExecutionRequest{Run: func(ctx context.Context) (*ToolResult, error) { close(entered); <-ctx.Done(); return nil, ctx.Err() }}) + if err != nil { + t.Fatal(err) + } + <-entered + cancel() + snapshot, _ := service.Get(handle.ID) + if snapshot.Execution.Status != ToolExecutionStatusRunning { + t.Fatal("caller cancellation ended detached worker") + } + scope.Cancel() + deadline, stop := context.WithTimeout(context.Background(), time.Second) + defer stop() + if err = scope.Wait(deadline); err != nil { + t.Fatal(err) + } + snapshot, _ = service.Get(handle.ID) + if snapshot.Execution.Status != ToolExecutionStatusCancelled { + t.Fatalf("unexpected state: %s", snapshot.Execution.Status) + } + if _, err = service.Submit(parent, ExecutionRequest{Run: func(context.Context) (*ToolResult, error) { t.Error("closed task executed tool"); return nil, nil }}); !errors.Is(err, runlease.ErrClosed) { + t.Fatalf("late submit: %v", err) + } +} +func TestRemoteCancellationRequiresAcknowledgement(t *testing.T) { + for _, confirm := range []bool{false, true} { + name := "unconfirmed" + if confirm { + name = "confirmed" + } + t.Run(name, func(t *testing.T) { + scope := runlease.New() + ctx := runlease.WithScope(context.Background(), scope) + service := NewExecutionService(nil, nil) + entered := make(chan struct{}) + req := ExecutionRequest{Remote: true, Run: func(ctx context.Context) (*ToolResult, error) { close(entered); <-ctx.Done(); return nil, ctx.Err() }} + if confirm { + req.ConfirmCancellation = func(context.Context) error { return nil } + } + handle, err := service.Submit(ctx, req) + if err != nil { + t.Fatal(err) + } + <-entered + scope.Cancel() + deadline, stop := context.WithTimeout(context.Background(), time.Second) + defer stop() + err = scope.Wait(deadline) + snapshot, _ := service.Get(handle.ID) + if confirm { + if err != nil || snapshot.Execution.Status != ToolExecutionStatusCancelled { + t.Fatalf("confirmed: %v %+v", err, snapshot.Execution) + } + } else { + if !errors.Is(err, runlease.ErrUnconfirmed) || snapshot.Execution.Status != ToolExecutionStatusOrphaned { + t.Fatalf("notification treated as confirmation: %v %+v", err, snapshot.Execution) + } + } + }) + } +} diff --git a/internal/mcp/execution_service.go b/internal/mcp/execution_service.go index 3e1a63e8..5e858603 100644 --- a/internal/mcp/execution_service.go +++ b/internal/mcp/execution_service.go @@ -9,6 +9,7 @@ import ( "time" "cyberstrike-ai/internal/authctx" + "cyberstrike-ai/internal/runlease" "github.com/google/uuid" "go.uber.org/zap" @@ -37,15 +38,18 @@ type ExecutionPreRunFunc func(context.Context, *ToolExecution) (func(), error) type ExecutionDoneFunc func(*ToolExecution) type ExecutionRequest struct { - ID string - ToolName string - Arguments map[string]interface{} - ConversationID string - OwnerUserID string - HardTimeout time.Duration - PreRun ExecutionPreRunFunc - Run ExecutionRunFunc - OnDone ExecutionDoneFunc + Remote bool + // A remote adapter may positively confirm server-side cancellation. + ConfirmCancellation func(context.Context) error + ID string + ToolName string + Arguments map[string]interface{} + ConversationID string + OwnerUserID string + HardTimeout time.Duration + PreRun ExecutionPreRunFunc + Run ExecutionRunFunc + OnDone ExecutionDoneFunc } type ExecutionHandle struct { @@ -57,13 +61,17 @@ type ExecutionSnapshot struct { } type executionEntry struct { - exec *ToolExecution - cancel context.CancelFunc - done chan struct{} - preRun ExecutionPreRunFunc - run ExecutionRunFunc - result *ToolResult - err error + releaseLease func() + remote bool + runStarted bool + confirmCancellation func(context.Context) error + exec *ToolExecution + cancel context.CancelFunc + done chan struct{} + preRun ExecutionPreRunFunc + run ExecutionRunFunc + result *ToolResult + err error } // ExecutionService keeps Eino-facing tool calls synchronous while moving the @@ -151,12 +159,19 @@ func (s *ExecutionService) Submit(ctx context.Context, req ExecutionRequest) (*E } else { runCtx, cancel = context.WithCancel(runCtx) } - entry := &executionEntry{exec: exec, cancel: cancel, done: make(chan struct{}), preRun: req.PreRun, run: req.Run} + releaseLease, leaseErr := runlease.FromContext(ctx).Register(id, cancel) + if leaseErr != nil { + cancel() + return nil, leaseErr + } + entry := &executionEntry{exec: exec, cancel: cancel, done: make(chan struct{}), preRun: req.PreRun, run: req.Run, + releaseLease: releaseLease, remote: req.Remote, confirmCancellation: req.ConfirmCancellation} s.mu.Lock() if _, exists := s.entries[id]; exists { s.mu.Unlock() cancel() + releaseLease() return nil, fmt.Errorf("execution already exists: %s", id) } s.entries[id] = entry @@ -188,8 +203,15 @@ func (s *ExecutionService) runWorker(ctx context.Context, entry *executionEntry, entry.cancel() notifyToolRunEnd(ctx, id) close(entry.done) + if entry.releaseLease != nil { + entry.releaseLease() + } }() + if ctx.Err() != nil { + s.finishEntry(ctx, entry, nil, ctx.Err(), onDone) + return + } if entry.preRun != nil { var preErr error release, preErr = entry.preRun(ctx, cloneToolExecution(entry.exec)) @@ -198,7 +220,12 @@ func (s *ExecutionService) runWorker(ctx context.Context, entry *executionEntry, return } } + if ctx.Err() != nil { + s.finishEntry(ctx, entry, nil, ctx.Err(), onDone) + return + } s.markEntryRunning(entry) + entry.runStarted = true result, err := entryResultRecover(ctx, entry.exec.ToolName, s.logger, func() (*ToolResult, error) { return nilSafeRun(ctx, entry) @@ -229,6 +256,12 @@ func (s *ExecutionService) finishEntry(ctx context.Context, entry *executionEntr if errors.As(err, &blockedErr) { result, err = blockedErr.result, nil } + cancellationUnconfirmed := entry.remote && entry.runStarted && ctx.Err() != nil && err != nil + if cancellationUnconfirmed && entry.confirmCancellation != nil { + confirmCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + cancellationUnconfirmed = entry.confirmCancellation(confirmCtx) != nil + cancel() + } cancelledWithUserNote := s.applyAbortUserNoteToCancelledToolResult(id, &result, &err) now := time.Now() @@ -287,6 +320,11 @@ func (s *ExecutionService) finishEntry(ctx context.Context, entry *executionEntr } entry.exec.Result = result } + if cancellationUnconfirmed { + entry.exec.Status = ToolExecutionStatusOrphaned + entry.exec.Error = "取消已请求，但远端 MCP 未确认执行已停止" + runlease.FromContext(ctx).MarkUnconfirmed(id, entry.exec.Error) + } finalExec := cloneToolExecution(entry.exec) s.mu.Unlock() diff --git a/internal/mcp/external_manager.go b/internal/mcp/external_manager.go index c4707b65..f708b3d1 100644 --- a/internal/mcp/external_manager.go +++ b/internal/mcp/external_manager.go @@ -706,6 +706,13 @@ func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args var client ExternalMCPClient var blockedByGuard bool handle, err := m.executionService.Submit(ctx, ExecutionRequest{ + ConfirmCancellation: func(confirmCtx context.Context) error { + if confirmer, ok := client.(ExternalCancellationConfirmer); ok { + return confirmer.ConfirmToolCancellation(confirmCtx, actualToolName, args) + } + return fmt.Errorf("external MCP client has no cancellation acknowledgement") + }, + Remote: true, ToolName: toolName, Arguments: args, ConversationID: MCPConversationIDFromContext(ctx), @@ -1649,3 +1656,10 @@ func (m *ExternalMCPManager) StopAll() { } m.refreshWg.Wait() } + +// ExternalCancellationConfirmer is an optional adapter contract for MCP +// servers with server-side cancellation receipts or lease/task status APIs. +// Ordinary notifications/cancelled must never be treated as confirmation. +type ExternalCancellationConfirmer interface { + ConfirmToolCancellation(context.Context, string, map[string]interface{}) error +} diff --git a/internal/multiagent/eino_model_facing_trace.go b/internal/multiagent/eino_model_facing_trace.go index 33d8d011..13cee221 100644 --- a/internal/multiagent/eino_model_facing_trace.go +++ b/internal/multiagent/eino_model_facing_trace.go @@ -93,6 +93,7 @@ func (m *modelFacingTraceMiddleware) BeforeModelRewriteState( ) (context.Context, *adk.ChatModelAgentState, error) { if m.holder != nil && state != nil { m.holder.storeFromState(state) + captureEinoTurnHistory(ctx, state.Messages) } return ctx, state, nil } @@ -119,6 +120,37 @@ func (m *agenticModelFacingTraceMiddleware) BeforeModelRewriteState( ) (context.Context, *adk.TypedChatModelAgentState[*schema.AgenticMessage], error) { if m.holder != nil && state != nil { m.holder.storeFromAgenticState(state) + captureEinoTurnHistory(ctx, AgenticMessagesToEino(state.Messages)) } return ctx, state, nil } + +// Capture completed output separately from the model-input trace: changing +// Snapshot's meaning would affect last_react_input persistence and retries. +func (m *modelFacingTraceMiddleware) AfterModelRewriteState(ctx context.Context, state *adk.ChatModelAgentState, _ *adk.ModelContext) (context.Context, *adk.ChatModelAgentState, error) { + if state != nil { + captureEinoTurnHistory(ctx, state.Messages) + } + return ctx, state, nil +} + +func (m *agenticModelFacingTraceMiddleware) AfterModelRewriteState(ctx context.Context, state *adk.TypedChatModelAgentState[*schema.AgenticMessage], _ *adk.TypedModelContext[*schema.AgenticMessage]) (context.Context, *adk.TypedChatModelAgentState[*schema.AgenticMessage], error) { + if state != nil { + captureEinoTurnHistory(ctx, AgenticMessagesToEino(state.Messages)) + } + return ctx, state, nil +} + +func (m *modelFacingTraceMiddleware) BeforeAgent(ctx context.Context, runCtx *adk.ChatModelAgentContext) (context.Context, *adk.ChatModelAgentContext, error) { + if runCtx != nil { + ctx = context.WithValue(ctx, einoTurnInstructionKey{}, runCtx.Instruction) + } + return ctx, runCtx, nil +} + +func (m *agenticModelFacingTraceMiddleware) BeforeAgent(ctx context.Context, runCtx *adk.ChatModelAgentContext) (context.Context, *adk.ChatModelAgentContext, error) { + if runCtx != nil { + ctx = context.WithValue(ctx, einoTurnInstructionKey{}, runCtx.Instruction) + } + return ctx, runCtx, nil +} diff --git a/internal/multiagent/eino_turn_history.go b/internal/multiagent/eino_turn_history.go new file mode 100644 index 00000000..7451f67d --- /dev/null +++ b/internal/multiagent/eino_turn_history.go @@ -0,0 +1,196 @@ +package multiagent + +import ( + "context" + "strings" + "sync" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" +) + +type einoTurnHistoryKey struct{} +type einoTurnInstructionKey struct{} + +// Owned by one TurnLoop, never shared between conversations. Model state is +// authoritative after compaction; events are a fallback for agents without the +// trace middleware and supply tool results completed after the last snapshot. +type einoTurnHistory struct { + mu sync.Mutex + messages []*schema.Message + modelState bool + pending map[string]bool + events []*schema.Message +} + +func (h *einoTurnHistory) begin(messages []*schema.Message) { + h.mu.Lock() + defer h.mu.Unlock() + h.messages = cloneSchemaMessages(messages) + h.modelState = false + h.events = nil +} + +func captureEinoTurnHistory(ctx context.Context, messages []*schema.Message) { + h, _ := ctx.Value(einoTurnHistoryKey{}).(*einoTurnHistory) + if h == nil || len(messages) == 0 { + return + } + h.mu.Lock() + defer h.mu.Unlock() + // Remove only the instruction known to be regenerated on Run. Other system + // content may contain durable context and must not be indiscriminately dropped. + instruction, _ := ctx.Value(einoTurnInstructionKey{}).(string) + h.messages = nil + for _, msg := range cloneSchemaMessages(messages) { + if msg.Role == schema.System && instruction != "" { + if msg.Content == instruction { + continue + } + msg.Content = strings.TrimPrefix(msg.Content, instruction+"\n\n") + } + h.messages = append(h.messages, msg) + } + h.modelState = true + h.pending = make(map[string]bool) + for _, msg := range h.messages { + for _, call := range msg.ToolCalls { + h.pending[call.ID] = true + } + if msg.Role == schema.Tool { + delete(h.pending, msg.ToolCallID) + } + } + // Snapshots already contain completed results. Release raw event payloads as + // compaction advances instead of retaining another full transcript for long-running turns. + for i, msg := range h.events { + if msg != nil && (msg.Role != schema.Tool || !h.pending[msg.ToolCallID]) { + h.events[i] = nil + } + } +} + +func (h *einoTurnHistory) nextInput() []*schema.Message { + h.mu.Lock() + defer h.mu.Unlock() + messages := cloneSchemaMessages(h.messages) + if !h.modelState { + messages = append(messages, cloneSchemaMessages(h.events)...) + } else { + // Never resurrect events discarded by summarization. Only pending calls in + // the authoritative state may acquire results from the event stream. + results := make(map[string]*schema.Message) + for _, msg := range h.events { + if msg != nil && msg.Role == schema.Tool { + results[msg.ToolCallID] = msg + } + } + var merged []*schema.Message + for i := 0; i < len(messages); i++ { + msg := messages[i] + merged = append(merged, msg) + if msg.Role != schema.Assistant || len(msg.ToolCalls) == 0 { + continue + } + present := make(map[string]bool) + for i+1 < len(messages) && messages[i+1].Role == schema.Tool { + i++ + merged = append(merged, messages[i]) + present[messages[i].ToolCallID] = true + } + for _, call := range msg.ToolCalls { + if !present[call.ID] && results[call.ID] != nil { + merged = append(merged, cloneSchemaMessages([]*schema.Message{results[call.ID]})...) + } + } + } + messages = merged + } + // Cancellation may leave a partial parallel tool batch. Explicit unknown + // results keep the protocol valid without claiming an unfinished call succeeded. + _, state, _ := newToolPairReconcilerMiddleware(nil, "turn_loop_continue").BeforeModelRewriteState( + context.Background(), &adk.ChatModelAgentState{Messages: messages}, nil) + return state.Messages +} + +type einoTurnEventHandler func(context.Context, *adk.TurnContext[EinoTurnLoopItem, *schema.Message], *adk.AsyncIterator[*adk.AgentEvent]) error + +func (h *einoTurnHistory) wrapEvents(handler einoTurnEventHandler) einoTurnEventHandler { + return func(ctx context.Context, tc *adk.TurnContext[EinoTurnLoopItem, *schema.Message], events *adk.AsyncIterator[*adk.AgentEvent]) error { + iter, gen := adk.NewAsyncIteratorPair[*adk.AgentEvent]() + done := make(chan struct{}) + go func() { + defer close(done) + defer gen.Close() + var streams sync.WaitGroup + defer streams.Wait() + for { + ev, ok := events.Next() + if !ok { + return + } + if ev != nil && ev.Output != nil && ev.Output.MessageOutput != nil { + mv := ev.Output.MessageOutput + h.mu.Lock() + index := len(h.events) + h.events = append(h.events, nil) + h.mu.Unlock() + save := func(msg *schema.Message) { + if msg == nil { + return + } + h.mu.Lock() + if !h.modelState || (msg.Role == schema.Tool && h.pending[msg.ToolCallID]) { + h.events[index] = cloneSchemaMessages([]*schema.Message{msg})[0] + } + h.mu.Unlock() + } + if mv.IsStreaming && mv.MessageStream != nil { + copies := mv.MessageStream.Copy(2) + // Copy the event as well: the framework can retain its original event. + eventCopy, outputCopy, variantCopy := *ev, *ev.Output, *mv + variantCopy.MessageStream = copies[0] + outputCopy.MessageOutput = &variantCopy + eventCopy.Output = &outputCopy + ev = &eventCopy + streams.Add(1) + go func() { + defer streams.Done() + defer copies[1].Close() + msg, err := (&adk.MessageVariant{IsStreaming: true, MessageStream: copies[1]}).GetMessage() + if err == nil { + save(msg) + } // incomplete streams are not completed history + }() + } else { + save(mv.Message) + } + } + gen.Send(ev) + } + }() + var err error + if handler != nil { + err = handler(ctx, tc, iter) + } + // Drain even if the UI bridge returned early on voluntary cancellation. + // The next GenInput must not race asynchronous event/stream consumers. + for { + ev, ok := iter.Next() + if !ok { + break + } + if ev != nil && ev.Output != nil && ev.Output.MessageOutput != nil { + mv := ev.Output.MessageOutput + if mv.IsStreaming && mv.MessageStream != nil { + mv.MessageStream.Close() + } + } + if err == nil && ev != nil && ev.Err != nil && !isEinoVoluntaryCancelErr(ev.Err) { + err = ev.Err + } + } + <-done + return err + } +} diff --git a/internal/multiagent/eino_turn_history_test.go b/internal/multiagent/eino_turn_history_test.go new file mode 100644 index 00000000..6eb17f36 --- /dev/null +++ b/internal/multiagent/eino_turn_history_test.go @@ -0,0 +1,254 @@ +package multiagent + +import ( + "context" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +type historyTool struct{ calls atomic.Int32 } + +func (h *historyTool) Info(context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{Name: "history_tool", Desc: "Record a completed test operation", ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{})}, nil +} +func (h *historyTool) InvokableRun(context.Context, string, ...tool.Option) (string, error) { + h.calls.Add(1) + return "completed-tool-evidence", nil +} + +type historyModel struct { + mu sync.Mutex + inputs [][]*schema.Message + started chan int + releases [2]chan struct{} +} + +func (m *historyModel) WithTools([]*schema.ToolInfo) (model.ToolCallingChatModel, error) { + return m, nil +} +func (m *historyModel) Generate(ctx context.Context, input []*schema.Message, _ ...model.Option) (*schema.Message, error) { + m.mu.Lock() + m.inputs = append(m.inputs, cloneSchemaMessages(input)) + n := len(m.inputs) + m.mu.Unlock() + m.started <- n + if n == 1 { + return schema.AssistantMessage("work started", []schema.ToolCall{{ID: "completed-call", Type: "function", Function: schema.FunctionCall{Name: "history_tool", Arguments: "{}"}}}), nil + } + if n == 2 || n == 3 { + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-m.releases[n-2]: + } + } + return schema.AssistantMessage("completed-response", nil), nil +} +func (m *historyModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + msg, err := m.Generate(ctx, input, opts...) + if err != nil { + return nil, err + } + return schema.StreamReaderFromArray([]*schema.Message{msg}), nil +} + +type historyCompactor struct { + adk.BaseChatModelAgentMiddleware +} + +func (*historyCompactor) BeforeModelRewriteState(ctx context.Context, state *adk.ChatModelAgentState, _ *adk.ModelContext) (context.Context, *adk.ChatModelAgentState, error) { + hasResult := false + for _, m := range state.Messages { + hasResult = hasResult || m.Role == schema.Tool + } + if !hasResult { + return ctx, state, nil + } + out := *state + out.Messages = nil + for _, m := range state.Messages { + if m.Content == "old-verbose-history" { + summary := schema.UserMessage("compressed-progress-summary") + summary.Extra = map[string]any{"_eino_adk_summarization_content_type": "summary"} + out.Messages = append(out.Messages, summary) + } else { + out.Messages = append(out.Messages, m) + } + } + return ctx, &out, nil +} + +func TestEinoTurnHistoryRetainsCompletedWorkAcrossInterrupts(t *testing.T) { + for _, safe := range []bool{false, true} { + for _, stream := range []bool{false, true} { + name := "timeout" + if safe { + name = "safe" + } + if stream { + name += "/stream" + } + t.Run(name, func(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + m := &historyModel{started: make(chan int, 8), releases: [2]chan struct{}{make(chan struct{}), make(chan struct{})}} + operation := &historyTool{} + agent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "history-agent", Instruction: "stable-agent-instruction", Model: m, + ToolsConfig: adk.ToolsConfig{ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{operation}}}, + Handlers: []adk.ChatModelAgentMiddleware{&historyCompactor{}, newSystemMessageNormalizerMiddleware(nil, "test"), newModelFacingTraceMiddleware(newModelFacingTraceHolder())}, + }) + if err != nil { + t.Fatal(err) + } + timeout := 20 * time.Millisecond + if safe { + timeout = time.Second + } + runtime := NewEinoTurnLoopRuntime(EinoTurnLoopRuntimeConfig{Agent: agent, EnableStreaming: stream, InterruptTimeout: timeout, InitialMessages: []*schema.Message{schema.UserMessage("original-task"), schema.SystemMessage("durable-system-context"), schema.AssistantMessage("old-verbose-history", nil)}}) + runtime.Run(ctx) + waitCall := func(want int) { + t.Helper() + select { + case n := <-m.started: + if n != want { + t.Fatalf("call %d, want %d", n, want) + } + case <-ctx.Done(): + t.Fatal("model call timed out") + } + } + waitCall(1) + waitCall(2) + if !runtime.PushInterruptContinue("first-supplement") { + t.Fatal("push rejected") + } + if safe { + close(m.releases[0]) + } + waitCall(3) + if !runtime.PushInterruptContinue("second-supplement") { + t.Fatal("push rejected") + } + if safe { + close(m.releases[1]) + } + waitCall(4) + runtime.StopWhenIdle() + if state := runtime.Wait(); state.ExitReason != nil { + t.Fatal(state.ExitReason) + } + if operation.calls.Load() != 1 { + t.Fatalf("tool executed %d times", operation.calls.Load()) + } + m.mu.Lock() + defer m.mu.Unlock() + for _, i := range []int{2, 3} { + input := m.inputs[i] + for _, marker := range []string{"original-task", "compressed-progress-summary", "completed-tool-evidence", "first-supplement", "durable-system-context", "stable-agent-instruction"} { + count := 0 + for _, msg := range input { + count += strings.Count(msg.Content, marker) + } + if count != 1 { + t.Errorf("call %d: %q occurs %d times", i+1, marker, count) + } + } + for _, msg := range input { + if strings.Contains(msg.Content, "old-verbose-history") { + t.Error("compacted history resurrected") + } + } + if input[len(input)-1].Role != schema.User { + t.Error("supplement must be last user message") + } + if safe { + count := 0 + for _, msg := range input { + if msg.Content == "completed-response" { + count++ + } + } + if count != i-1 { + t.Errorf("completed responses=%d, want %d", count, i-1) + } + } + } + if !strings.Contains(m.inputs[3][len(m.inputs[3])-1].Content, "second-supplement") { + t.Error("second supplement lost") + } + }) + } + } +} + +func TestEinoTurnHistoryPendingToolBatch(t *testing.T) { + h := &einoTurnHistory{} + ctx := context.WithValue(context.Background(), einoTurnHistoryKey{}, h) + calls := []schema.ToolCall{{ID: "done", Function: schema.FunctionCall{Name: "tool"}}, {ID: "pending", Function: schema.FunctionCall{Name: "tool"}}} + captureEinoTurnHistory(ctx, []*schema.Message{schema.UserMessage("summary"), schema.AssistantMessage("", calls)}) + h.events = []*schema.Message{schema.AssistantMessage("discarded-old-output", nil), schema.ToolMessage("actual-result", "done")} + got := h.nextInput() + if len(got) != 4 || got[2].Content != "actual-result" || got[3].Content != patchedMissingToolResult { + t.Fatalf("bad reconciled messages: %#v", got) + } + if got[2].ToolCallID != "done" || got[3].ToolCallID != "pending" { + t.Fatal("tool IDs lost") + } +} + +func TestEinoTurnHistoryAgenticSnapshotAndIsolation(t *testing.T) { + first, second := &einoTurnHistory{}, &einoTurnHistory{} + first.begin([]*schema.Message{schema.UserMessage("first-task")}) + second.begin([]*schema.Message{schema.UserMessage("second-task")}) + ctx := context.WithValue(context.Background(), einoTurnHistoryKey{}, first) + mw := newAgenticModelFacingTraceMiddleware(newModelFacingTraceHolder()) + ctx, _, err := mw.BeforeAgent(ctx, &adk.ChatModelAgentContext{Instruction: "agent-instruction"}) + if err != nil { + t.Fatal(err) + } + state := &adk.TypedChatModelAgentState[*schema.AgenticMessage]{Messages: EinoMessagesToAgentic([]*schema.Message{ + schema.SystemMessage("agent-instruction\n\nsystem-summary"), schema.UserMessage("compacted-first-task"), + })} + if _, _, err = mw.BeforeModelRewriteState(ctx, state, nil); err != nil { + t.Fatal(err) + } + state.Messages = append(state.Messages, EinoMessagesToAgentic([]*schema.Message{schema.AssistantMessage("finished-step", nil)})[0]) + if _, _, err = mw.AfterModelRewriteState(ctx, state, nil); err != nil { + t.Fatal(err) + } + got := first.nextInput() + if len(got) != 3 || got[0].Content != "system-summary" || got[2].Content != "finished-step" { + t.Fatalf("agentic state lost: %#v", got) + } + other := second.nextInput() + if len(other) != 1 || other[0].Content != "second-task" { + t.Fatalf("conversation leaked: %#v", other) + } +} + +func TestEinoTurnHistoryFallbackKeepsStreamedOutput(t *testing.T) { + h := &einoTurnHistory{} + h.begin([]*schema.Message{schema.UserMessage("initial-task")}) + events, gen := adk.NewAsyncIteratorPair[*adk.AgentEvent]() + gen.Send(&adk.AgentEvent{Output: &adk.AgentOutput{MessageOutput: &adk.MessageVariant{ + IsStreaming: true, MessageStream: schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("completed-", nil), schema.AssistantMessage("stream", nil)}), + }}}) + gen.Close() + if err := h.wrapEvents(nil)(context.Background(), nil, events); err != nil { + t.Fatal(err) + } + got := h.nextInput() + if len(got) != 2 || got[1].Content != "completed-stream" { + t.Fatalf("stream history lost: %#v", got) + } +} diff --git a/internal/multiagent/eino_turn_loop_bridge_test.go b/internal/multiagent/eino_turn_loop_bridge_test.go index c20f8e66..01522fbe 100644 --- a/internal/multiagent/eino_turn_loop_bridge_test.go +++ b/internal/multiagent/eino_turn_loop_bridge_test.go @@ -102,6 +102,9 @@ func TestRunEinoADKAgentLoopUsesTurnLoopInterruptPush(t *testing.T) { t.Fatalf("model calls = %d, want at least 2", len(inputs)) } last := inputs[len(inputs)-1] + if len(last) < 2 || last[0].Content != "initial task" { + t.Fatalf("initial task lost: %#v", last) + } if len(last) == 0 || last[len(last)-1].Role != schema.User || last[len(last)-1].Content == "initial task" { t.Fatalf("last model input = %#v, want interrupt supplement turn", last) } diff --git a/internal/multiagent/eino_turn_loop_runtime.go b/internal/multiagent/eino_turn_loop_runtime.go index 4ee38743..02559f5e 100644 --- a/internal/multiagent/eino_turn_loop_runtime.go +++ b/internal/multiagent/eino_turn_loop_runtime.go @@ -49,6 +49,7 @@ func NewEinoTurnLoopRuntime(cfg EinoTurnLoopRuntimeConfig) *EinoTurnLoopRuntime } enableStreaming := cfg.EnableStreaming prepareAgent := cfg.PrepareAgent + history := &einoTurnHistory{} if prepareAgent == nil { prepareAgent = func(context.Context, *adk.TurnLoop[EinoTurnLoopItem, *schema.Message], []EinoTurnLoopItem) (adk.Agent, error) { return cfg.Agent, nil @@ -58,9 +59,11 @@ func NewEinoTurnLoopRuntime(cfg EinoTurnLoopRuntimeConfig) *EinoTurnLoopRuntime Store: cfg.Store, CheckpointID: cfg.CheckpointID, GenInput: func(ctx context.Context, _ *adk.TurnLoop[EinoTurnLoopItem, *schema.Message], items []EinoTurnLoopItem) (*adk.GenInputResult[EinoTurnLoopItem, *schema.Message], error) { - msgs := mergeEinoTurnLoopMessages(items) + msgs := append(history.nextInput(), mergeEinoTurnLoopMessages(items)...) + history = &einoTurnHistory{} + history.begin(msgs) return &adk.GenInputResult[EinoTurnLoopItem, *schema.Message]{ - RunCtx: ctx, + RunCtx: context.WithValue(ctx, einoTurnHistoryKey{}, history), Input: &adk.AgentInput{ Messages: msgs, EnableStreaming: enableStreaming, @@ -74,13 +77,15 @@ func NewEinoTurnLoopRuntime(cfg EinoTurnLoopRuntimeConfig) *EinoTurnLoopRuntime consumed = append(consumed, newItems...) remaining := append([]EinoTurnLoopItem(nil), unhandledItems...) return &adk.GenResumeResult[EinoTurnLoopItem, *schema.Message]{ - RunCtx: ctx, + RunCtx: context.WithValue(ctx, einoTurnHistoryKey{}, history), Consumed: consumed, Remaining: remaining, }, nil }, - PrepareAgent: prepareAgent, - OnAgentEvents: cfg.OnAgentEvents, + PrepareAgent: prepareAgent, + OnAgentEvents: func(ctx context.Context, tc *adk.TurnContext[EinoTurnLoopItem, *schema.Message], events *adk.AsyncIterator[*adk.AgentEvent]) error { + return history.wrapEvents(cfg.OnAgentEvents)(ctx, tc, events) + }, }) if len(cfg.InitialMessages) > 0 { loop.Push(EinoTurnLoopItem{Kind: "initial", Messages: cloneSchemaMessages(cfg.InitialMessages)}) diff --git a/internal/multiagent/eino_turn_loop_runtime_test.go b/internal/multiagent/eino_turn_loop_runtime_test.go index be637fbd..c18c0718 100644 --- a/internal/multiagent/eino_turn_loop_runtime_test.go +++ b/internal/multiagent/eino_turn_loop_runtime_test.go @@ -183,6 +183,9 @@ func TestEinoTurnLoopRuntimePushInterruptStartsNextTurn(t *testing.T) { t.Fatalf("first input = %q, want initial task", got) } lastInput := inputs[len(inputs)-1] + if len(lastInput) < 2 || lastInput[0].Content != "initial task" { + t.Fatalf("initial history lost after preempt: %#v", lastInput) + } if len(lastInput) == 0 || !strings.Contains(lastInput[len(lastInput)-1].Content, "focus on ssh") { t.Fatalf("last input = %#v, want interrupt note", lastInput) } diff --git a/internal/processguard/cgroup_linux_test.go b/internal/processguard/cgroup_linux_test.go new file mode 100644 index 00000000..8b4de155 --- /dev/null +++ b/internal/processguard/cgroup_linux_test.go @@ -0,0 +1,120 @@ +//go:build linux + +package processguard + +import ( + "context" + "fmt" + "os" + "os/exec" + "path/filepath" + "strings" + "syscall" + "testing" + "time" +) + +func TestCgroupContainsSetsidAndAppliesLimits(t *testing.T) { + opts := testOptions() + if opts.CgroupRoot == "" { + t.Skip("set CSAI_TEST_CGROUP_ROOT to a delegated cgroup v2 root") + } + opts.Mode = "required" + opts.CPUQuotaMicros = 50000 + id := testID() + g, err := NewWithOptions(id, opts) + if err != nil { + t.Fatal(err) + } + defer closeTestGroup(t, g) + file := filepath.Join(t.TempDir(), "escaped") + cmd, err := startTestCommand(g, fmt.Sprintf("setsid sh -c 'echo $$ > %s; exec sleep 300' /dev/null 2>&1 &", file)) + if err != nil { + t.Fatal(err) + } + reaped := make(chan struct{}) + go func() { _ = cmd.Wait(); close(reaped) }() + pid := readPID(t, file) + <-reaped // The launching shell is gone; the cgroup must still own setsid descendants. + root := filepath.Join(opts.CgroupRoot, "task-"+id) + for name, want := range map[string]string{"pids.max": "64", "memory.max": "268435456", "cpu.max": "50000 100000"} { + data, err := os.ReadFile(filepath.Join(root, name)) + if err != nil || strings.TrimSpace(string(data)) != want { + t.Fatalf("%s=%s err=%v", name, data, err) + } + } + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err = g.Close(ctx); err != nil { + t.Fatal(err) + } + <-reaped + waitGone(t, pid) + if _, err = os.Stat(root); !os.IsNotExist(err) { + t.Fatalf("cgroup retained after cleanup: %v", err) + } +} + +func TestCgroupStartupDelegationAndRecovery(t *testing.T) { + opts := testOptions() + if opts.CgroupRoot == "" { + t.Skip("requires delegated cgroup fixture") + } + before, err := os.ReadFile("/proc/self/cgroup") + if err != nil { + t.Fatal(err) + } + original := "" + for _, line := range strings.Split(string(before), "\n") { + if strings.HasPrefix(line, "0::") { + original = filepath.Join("/sys/fs/cgroup", strings.TrimPrefix(line, "0::")) + } + } + root := filepath.Join(opts.CgroupRoot, "startup-fixture") + if err = os.Mkdir(root, 0700); err != nil { + t.Fatal(err) + } + if err = os.WriteFile(filepath.Join(root, "cgroup.procs"), []byte(fmt.Sprint(os.Getpid())), 0600); err != nil { + t.Fatal(err) + } + defer func() { + _ = os.WriteFile(filepath.Join(original, "cgroup.procs"), []byte(fmt.Sprint(os.Getpid())), 0600) + if rootLock != nil { + _ = rootLock.Close() + rootLock = nil + } + _ = removeCgroupTree(root) + }() + stale := filepath.Join(root, "task-"+testID()) + if err = os.Mkdir(stale, 0700); err != nil { + t.Fatal(err) + } + dir, err := os.Open(stale) + if err != nil { + t.Fatal(err) + } + defer dir.Close() + cmd := exec.Command("sh", "-c", "exec sleep 300") + cmd.SysProcAttr = &syscall.SysProcAttr{UseCgroupFD: true, CgroupFD: int(dir.Fd()), Setsid: true} + if err = cmd.Start(); err != nil { + t.Fatal(err) + } + reaped := make(chan struct{}) + go func() { _ = cmd.Wait(); close(reaped) }() + opts.CgroupRoot = root + opts.Mode = "required" + if err = configurePlatform(&opts); err != nil { + _ = cmd.Process.Kill() + <-reaped + t.Fatal(err) + } + <-reaped + waitGone(t, cmd.Process.Pid) + if _, err = os.Stat(stale); !os.IsNotExist(err) { + t.Fatalf("stale task cgroup was not removed: %v", err) + } + data, err := os.ReadFile(filepath.Join(root, "cgroup.subtree_control")) + if err != nil || !strings.Contains(string(data), "memory") { + t.Fatalf("delegation not enabled: %s %v", data, err) + } +} diff --git a/internal/processguard/check.go b/internal/processguard/check.go new file mode 100644 index 00000000..de1bbb2d --- /dev/null +++ b/internal/processguard/check.go @@ -0,0 +1,52 @@ +package processguard + +import ( + "context" + "crypto/rand" + "errors" + "fmt" + "os" + "os/exec" +) + +// Check exercises the real creation path, including clone3/Job inheritance, +// watchdog readiness, admission and cleanup. It does not start the HTTP server. +func Check(ctx context.Context) (backend string, err error) { + var id [16]byte + if _, err = rand.Read(id[:]); err != nil { + return "", err + } + name := fmt.Sprintf("%x-%x-%x-%x-%x", id[:4], id[4:6], id[6:8], id[8:10], id[10:]) + g, err := New(name) + if err != nil { + return "", err + } + defer func() { err = errors.Join(err, g.Close(ctx)) }() + exe, err := os.Executable() + if err != nil { + return "", err + } + cmd := exec.CommandContext(ctx, exe, "-h") + configureGuardian(cmd) + launch, err := g.Prepare(cmd) + if err != nil { + return "", err + } + defer launch.Dispose() + if err = cmd.Start(); err != nil { + return "", err + } + if err = launch.Commit(); err != nil { + _ = cmd.Process.Kill() + _ = cmd.Wait() + return "", err + } + err = cmd.Wait() + if err != nil { + return "", err + } + if err = g.Release(cmd.Process.Pid); err != nil { + return "", err + } + return g.Name(), nil +} diff --git a/internal/processguard/group_unix.go b/internal/processguard/group_unix.go new file mode 100644 index 00000000..baa6afc3 --- /dev/null +++ b/internal/processguard/group_unix.go @@ -0,0 +1,180 @@ +//go:build !windows + +package processguard + +import ( + "context" + "encoding/base64" + "encoding/json" + "fmt" + "io" + "os" + "os/exec" + "strconv" + "sync" + "syscall" + "time" +) + +type unixGroup struct { + mu sync.Mutex + watcher *watchdog + pids map[int]struct{} + closed bool +} + +func newUnixGroup() (*unixGroup, error) { + g := &unixGroup{pids: make(map[int]struct{})} + w, err := startWatchdog(watchRequest{Name: "process_group"}, func() { + g.mu.Lock() + defer g.mu.Unlock() + g.closed = true + for pid := range g.pids { + _ = syscall.Kill(-pid, syscall.SIGKILL) + } + }) + if err != nil { + return nil, err + } + g.watcher = w + return g, nil +} +func (g *unixGroup) Name() string { return "process_group_watchdog" } +func configureGuardian(cmd *exec.Cmd) { cmd.SysProcAttr = &syscall.SysProcAttr{Setsid: true} } + +type childSpec struct { + Path string + Args []string +} + +func (g *unixGroup) Prepare(cmd *exec.Cmd) (*Launch, error) { + g.mu.Lock() + defer g.mu.Unlock() + if g.closed { + return nil, fmt.Errorf("process group is closed") + } + // A dead guardian rejects subsequent launches before user code is executed. + if _, err := g.watcher.send(watchRequest{Op: "ping"}); err != nil { + return nil, err + } + read, write, err := os.Pipe() + if err != nil { + return nil, err + } + spec, _ := json.Marshal(childSpec{Path: cmd.Path, Args: cmd.Args}) + exe, err := os.Executable() + if err != nil { + read.Close() + write.Close() + return nil, err + } + fd := 3 + len(cmd.ExtraFiles) + cmd.ExtraFiles = append(cmd.ExtraFiles, read) + cmd.Path = exe + cmd.Args = []string{exe, childArg, strconv.Itoa(fd), base64.RawStdEncoding.EncodeToString(spec)} + return &Launch{Dispose: func() { read.Close(); write.Close() }, Commit: func() error { + g.mu.Lock() + defer g.mu.Unlock() + if g.closed { + return fmt.Errorf("process group is closed") + } + pid := cmd.Process.Pid + if _, err := g.watcher.send(watchRequest{Op: "add", PID: pid}); err != nil { + return err + } + g.pids[pid] = struct{}{} + _, err := write.Write([]byte{1}) + return err + }}, nil +} +func (g *unixGroup) Release(pid int) error { + g.mu.Lock() + defer g.mu.Unlock() + if _, ok := g.pids[pid]; !ok { + return nil + } + select { + case <-g.watcher.done: + delete(g.pids, pid) + return nil + default: + } + if _, err := g.watcher.send(watchRequest{Op: "release", PID: pid}); err != nil { + return err + } + delete(g.pids, pid) + return nil +} +func (g *unixGroup) Close(ctx context.Context) error { + g.mu.Lock() + defer g.mu.Unlock() + g.closed = true + for pid := range g.pids { + _ = syscall.Kill(-pid, syscall.SIGKILL) + } + for { + for pid := range g.pids { + if syscall.Kill(-pid, 0) == syscall.ESRCH { + delete(g.pids, pid) + } + } + if len(g.pids) == 0 { + return g.watcher.close() + } + select { + case <-ctx.Done(): + return ctx.Err() + case <-time.After(10 * time.Millisecond): + } + } +} +func gatedChildMain(args []string) error { + if len(args) != 2 { + return fmt.Errorf("invalid internal launch") + } + fd, err := strconv.Atoi(args[0]) + if err != nil || fd < 3 { + return fmt.Errorf("invalid launch gate") + } + gate := os.NewFile(uintptr(fd), "launch-gate") + var token [1]byte + if _, err = io.ReadFull(gate, token[:]); err != nil { + return fmt.Errorf("owner exited before launch: %w", err) + } + gate.Close() + if token[0] != 1 { + return fmt.Errorf("invalid launch token") + } + b, err := base64.RawStdEncoding.DecodeString(args[1]) + if err != nil { + return err + } + var spec childSpec + if err = json.Unmarshal(b, &spec); err != nil { + return err + } + return syscall.Exec(spec.Path, spec.Args, os.Environ()) +} +func groupGuardian(dec *json.Decoder, enc *json.Encoder) error { + pids := make(map[int]struct{}) + return serveGuardian(dec, enc, func(req watchRequest) error { + switch req.Op { + case "ping": + case "add": + if req.PID <= 1 { + return fmt.Errorf("invalid PID") + } + pids[req.PID] = struct{}{} + case "release": + delete(pids, req.PID) + default: + return fmt.Errorf("unknown guardian command") + } + return nil + }, func() error { + for pid := range pids { + _ = syscall.Kill(-pid, syscall.SIGKILL) + } + return nil + }) +} diff --git a/internal/processguard/guard.go b/internal/processguard/guard.go new file mode 100644 index 00000000..0fd303a4 --- /dev/null +++ b/internal/processguard/guard.go @@ -0,0 +1,89 @@ +// Package processguard provides OS containment and out-of-process crash cleanup. +// It is intentionally independent of the Agent/MCP packages so it can be +// cross-compiled and exercised without starting the application. +package processguard + +import ( + "context" + "fmt" + "os/exec" + "sync" +) + +type Options struct { + Mode string `yaml:"mode" json:"mode"` // auto, required, process_group + CgroupRoot string `yaml:"cgroup_root" json:"cgroup_root"` + MaxProcesses int `yaml:"max_processes" json:"max_processes"` + MemoryMaxBytes int64 `yaml:"memory_max_bytes" json:"memory_max_bytes"` + CPUQuotaMicros int64 `yaml:"cpu_quota_micros" json:"cpu_quota_micros"` // per 100000 us +} + +// Prepared commands must call Commit after Start and always call Dispose. +// Commit releases the Unix fallback launch gate only after watchdog ownership +// is acknowledged. Strong backends assign containment atomically at creation. +type Launch struct { + Commit func() error + Dispose func() +} + +type Group interface { + Name() string + Prepare(*exec.Cmd) (*Launch, error) + Release(int) error + Close(context.Context) error +} + +var configured = struct { + sync.RWMutex + opts Options +}{opts: Options{Mode: "auto", MaxProcesses: 256, MemoryMaxBytes: 2 << 30}} + +func normalize(o Options) (Options, error) { + if o.Mode == "" { + o.Mode = "auto" + } + if o.Mode != "auto" && o.Mode != "required" && o.Mode != "process_group" { + return o, fmt.Errorf("invalid process isolation mode %q", o.Mode) + } + if o.MaxProcesses == 0 { + o.MaxProcesses = 256 + } + if o.MemoryMaxBytes == 0 { + o.MemoryMaxBytes = 2 << 30 + } + if o.MaxProcesses < 1 || o.MaxProcesses > 65535 || o.MemoryMaxBytes < 0 || o.CPUQuotaMicros < 0 { + return o, fmt.Errorf("invalid process isolation resource limits") + } + return o, nil +} + +// Configure validates deployment before accepting any tasks. An explicit root +// or required mode fails closed; it never silently falls back after an error. +func Configure(o Options) error { + var err error + o, err = normalize(o) + if err != nil { + return err + } + if err = configurePlatform(&o); err != nil { + return err + } + configured.Lock() + configured.opts = o + configured.Unlock() + return nil +} +func New(id string) (Group, error) { + configured.RLock() + o := configured.opts + configured.RUnlock() + return NewWithOptions(id, o) +} +func NewWithOptions(id string, o Options) (Group, error) { + var err error + o, err = normalize(o) + if err != nil { + return nil, err + } + return newPlatformGroup(id, o) +} diff --git a/internal/processguard/guard_test.go b/internal/processguard/guard_test.go new file mode 100644 index 00000000..b7a4c4e3 --- /dev/null +++ b/internal/processguard/guard_test.go @@ -0,0 +1,160 @@ +//go:build !windows + +package processguard + +import ( + "context" + "fmt" + "os" + "os/exec" + "path/filepath" + "runtime" + "strconv" + "strings" + "syscall" + "testing" + "time" +) + +func testID() string { + return fmt.Sprintf("%08x-1111-4111-8111-%012x", os.Getpid(), uint64(time.Now().UnixNano())&0xffffffffffff) +} +func testOptions() Options { + return Options{CgroupRoot: os.Getenv("CSAI_TEST_CGROUP_ROOT"), MaxProcesses: 64, MemoryMaxBytes: 256 << 20} +} +func closeTestGroup(t *testing.T, g Group) { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := g.Close(ctx); err != nil { + t.Error(err) + } +} +func startTestCommand(g Group, command string) (*exec.Cmd, error) { + cmd := exec.Command("sh", "-c", command) + cmd.SysProcAttr = &syscall.SysProcAttr{Setsid: true} + launch, err := g.Prepare(cmd) + if err != nil { + return nil, err + } + defer launch.Dispose() + if err = cmd.Start(); err != nil { + return nil, err + } + if err = launch.Commit(); err != nil { + _ = cmd.Process.Kill() + _ = cmd.Wait() + return nil, err + } + return cmd, nil +} +func readPID(t *testing.T, path string) int { + t.Helper() + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + b, err := os.ReadFile(path) + if err == nil { + pid, err := strconv.Atoi(strings.TrimSpace(string(b))) + if err == nil && pid > 0 { + return pid + } + } + time.Sleep(10 * time.Millisecond) + } + t.Fatalf("no PID written to %s", path) + return 0 +} +func waitGone(t *testing.T, pid int) { + t.Helper() + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + if syscall.Kill(pid, 0) == syscall.ESRCH { + return + } + time.Sleep(10 * time.Millisecond) + } + t.Fatalf("PID %d survived cleanup", pid) +} +func TestGuardianReapsAfterOwnerSIGKILL(t *testing.T) { + pidPath := filepath.Join(t.TempDir(), "pid") + owner := exec.Command(os.Args[0], "-test.run=^TestGuardianOwnerHelper$") + owner.Env = append(os.Environ(), "CSAI_GUARD_TEST_OWNER=1", "CSAI_GUARD_TEST_PID="+pidPath) + if err := owner.Start(); err != nil { + t.Fatal(err) + } + defer func() { _ = owner.Process.Kill(); _ = owner.Wait() }() + pid := readPID(t, pidPath) + if err := owner.Process.Kill(); err != nil { + t.Fatal(err) + } + _ = owner.Wait() + waitGone(t, pid) +} +func TestGuardianOwnerHelper(t *testing.T) { + if os.Getenv("CSAI_GUARD_TEST_OWNER") != "1" { + t.Skip("subprocess helper") + } + g, err := NewWithOptions(testID(), testOptions()) + if err != nil { + t.Fatal(err) + } + command := fmt.Sprintf("echo $$ > %q; exec sleep 300", os.Getenv("CSAI_GUARD_TEST_PID")) + cmd, err := startTestCommand(g, command) + if err != nil { + t.Fatal(err) + } + go cmd.Wait() + select {} +} +func TestGroupCloseAndAdmission(t *testing.T) { + g, err := NewWithOptions(testID(), testOptions()) + if err != nil { + t.Fatal(err) + } + defer closeTestGroup(t, g) + cmd, err := startTestCommand(g, "exec sleep 300") + if err != nil { + t.Fatal(err) + } + reaped := make(chan struct{}) + go func() { _ = cmd.Wait(); close(reaped) }() + closeTestGroup(t, g) + <-reaped + waitGone(t, cmd.Process.Pid) + if _, err = g.Prepare(exec.Command("sh", "-c", "true")); err == nil { + t.Fatal("closed containment accepted a command") + } +} +func TestLaunchGateOwnerDisappearsBeforeCommit(t *testing.T) { + if runtime.GOOS == "linux" && testOptions().CgroupRoot != "" { + t.Skip("cgroup assignment is atomic without a gate") + } + g, err := NewWithOptions(testID(), Options{}) + if err != nil { + t.Fatal(err) + } + defer closeTestGroup(t, g) + file := filepath.Join(t.TempDir(), "should-not-exist") + cmd := exec.Command("sh", "-c", fmt.Sprintf("echo escaped > %q", file)) + cmd.SysProcAttr = &syscall.SysProcAttr{Setsid: true} + launch, err := g.Prepare(cmd) + if err != nil { + t.Fatal(err) + } + if err = cmd.Start(); err != nil { + t.Fatal(err) + } + launch.Dispose() // simulate owner crashing before watchdog registration + _ = cmd.Wait() + if _, err = os.Stat(file); !os.IsNotExist(err) { + t.Fatal("unregistered child executed user code") + } +} +func TestRequiredIsolationFailsClosed(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("Windows has Job Objects") + } + if _, err := NewWithOptions(testID(), Options{Mode: "required"}); err == nil { + t.Fatal("required isolation silently downgraded") + } +} diff --git a/internal/processguard/job_windows_test.go b/internal/processguard/job_windows_test.go new file mode 100644 index 00000000..5973a71b --- /dev/null +++ b/internal/processguard/job_windows_test.go @@ -0,0 +1,97 @@ +//go:build windows + +package processguard + +import ( + "context" + "fmt" + "golang.org/x/sys/windows" + "os" + "os/exec" + "path/filepath" + "strconv" + "strings" + "testing" + "time" +) + +func TestWindowsJobOwnerHelper(t *testing.T) { + if os.Getenv("CSAI_JOB_OWNER") != "1" { + t.Skip("subprocess helper") + } + g, err := NewWithOptions(fmt.Sprintf("test-%d-%d", os.Getpid(), time.Now().UnixNano()), Options{Mode: "required"}) + if err != nil { + t.Fatal(err) + } + cmd := exec.Command(os.Args[0], "-test.run=^TestWindowsJobPayload$") + cmd.Env = append(os.Environ(), "CSAI_JOB_PAYLOAD=1") + launch, err := g.Prepare(cmd) + if err != nil { + t.Fatal(err) + } + defer launch.Dispose() + if err = cmd.Start(); err != nil { + t.Fatal(err) + } + if err = launch.Commit(); err != nil { + t.Fatal(err) + } + go cmd.Wait() + select {} +} +func TestWindowsJobPayload(t *testing.T) { + if os.Getenv("CSAI_JOB_PAYLOAD") != "1" { + t.Skip("subprocess helper") + } + if err := os.WriteFile(os.Getenv("CSAI_JOB_PIDFILE"), []byte(strconv.Itoa(os.Getpid())), 0600); err != nil { + t.Fatal(err) + } + time.Sleep(300 * time.Second) +} +func TestWindowsJobReapsAfterOwnerKilled(t *testing.T) { + file := filepath.Join(t.TempDir(), "pid") + owner := exec.Command(os.Args[0], "-test.run=^TestWindowsJobOwnerHelper$") + owner.Env = append(os.Environ(), "CSAI_JOB_OWNER=1", "CSAI_JOB_PIDFILE="+file) + if err := owner.Start(); err != nil { + t.Fatal(err) + } + defer func() { _ = owner.Process.Kill(); _ = owner.Wait() }() + var pid int + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + data, _ := os.ReadFile(file) + pid, _ = strconv.Atoi(strings.TrimSpace(string(data))) + if pid > 0 { + break + } + time.Sleep(10 * time.Millisecond) + } + if pid == 0 { + t.Fatal("job child did not start") + } + handle, err := windows.OpenProcess(windows.SYNCHRONIZE, false, uint32(pid)) + if err != nil { + t.Fatal(err) + } + defer windows.CloseHandle(handle) + _ = owner.Process.Kill() + _ = owner.Wait() + event, err := windows.WaitForSingleObject(handle, 5000) + if err != nil || event != windows.WAIT_OBJECT_0 { + t.Fatalf("child survived owner death: %d %v", event, err) + } +} +func TestWindowsJobClose(t *testing.T) { + g, err := NewWithOptions(fmt.Sprintf("test-%d-%d", os.Getpid(), time.Now().UnixNano()), Options{Mode: "required", CPUQuotaMicros: 100000}) + if err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err = g.Close(ctx); err != nil { + t.Fatal(err) + } + if _, err = g.Prepare(exec.Command("cmd.exe", "/c", "exit")); err == nil { + t.Fatal("closed job admitted a process") + } +} diff --git a/internal/processguard/platform_linux.go b/internal/processguard/platform_linux.go new file mode 100644 index 00000000..9e27f9b5 --- /dev/null +++ b/internal/processguard/platform_linux.go @@ -0,0 +1,313 @@ +//go:build linux + +package processguard + +import ( + "context" + "crypto/sha256" + "encoding/json" + "errors" + "fmt" + "os" + "os/exec" + "path/filepath" + "strconv" + "strings" + "sync" + "syscall" + "time" + + "golang.org/x/sys/unix" +) + +var rootLock *os.File // retained until server exit; never inherited by commands + +func configurePlatform(o *Options) error { + if o.CgroupRoot == "" { + if o.Mode == "required" { + return fmt.Errorf("required isolation needs security.process_isolation.cgroup_root") + } + return nil + } + if o.Mode == "process_group" { + return fmt.Errorf("cgroup_root cannot be combined with process_group mode") + } + if o.CgroupRoot == "auto" { + data, err := os.ReadFile("/proc/self/cgroup") + if err != nil { + return err + } + for _, line := range strings.Split(string(data), "\n") { + if strings.HasPrefix(line, "0::") { + o.CgroupRoot = filepath.Join("/sys/fs/cgroup", strings.TrimPrefix(line, "0::")) + break + } + } + } + root, err := validateRoot(o.CgroupRoot) + if err != nil { + return err + } + o.CgroupRoot = root + // An exclusive host-side lock prevents one server's recovery sweep from + // killing tasks owned by another server using the same delegated root. + hash := sha256.Sum256([]byte(root)) + lockPath := filepath.Join(os.TempDir(), fmt.Sprintf("cyberstrike-cgroup-%d-%x.lock", os.Getuid(), hash[:12])) + fd, err := unix.Open(lockPath, unix.O_CREAT|unix.O_RDWR|unix.O_CLOEXEC|unix.O_NOFOLLOW, 0600) + if err != nil { + return err + } + lock := os.NewFile(uintptr(fd), lockPath) + if err = unix.Flock(fd, unix.LOCK_EX|unix.LOCK_NB); err != nil { + lock.Close() + return fmt.Errorf("cgroup root is already owned: %w", err) + } + success := false + defer func() { + if !success { + lock.Close() + } + }() + // cgroup v2 requires the delegated parent to have no processes before + // domain controllers can be enabled. Move only this server, never outsiders. + data, err := os.ReadFile(filepath.Join(root, "cgroup.procs")) + if err != nil { + return err + } + for _, pid := range strings.Fields(string(data)) { + if pid != strconv.Itoa(os.Getpid()) { + return fmt.Errorf("delegated root contains another process %s", pid) + } + } + if len(strings.Fields(string(data))) > 0 { + supervisor := filepath.Join(root, "supervisor") + if err = os.Mkdir(supervisor, 0700); err != nil && !os.IsExist(err) { + return err + } + if err = os.WriteFile(filepath.Join(supervisor, "cgroup.procs"), []byte(strconv.Itoa(os.Getpid())), 0600); err != nil { + return err + } + } + if err = os.WriteFile(filepath.Join(root, "cgroup.subtree_control"), []byte("+cpu +memory +pids"), 0600); err != nil { + return fmt.Errorf("delegate cpu, memory and pids controllers: %w", err) + } + // Recover only our names under the exclusively owned root. No PID replay. + entries, err := os.ReadDir(root) + if err != nil { + return err + } + for _, entry := range entries { + if entry.IsDir() && validTaskName(entry.Name()) { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + err = killAndRemoveCgroup(ctx, filepath.Join(root, entry.Name())) + cancel() + if err != nil { + return fmt.Errorf("recover %s: %w", entry.Name(), err) + } + } + } + rootLock = lock + success = true + return nil +} + +func validateRoot(root string) (string, error) { + if !filepath.IsAbs(root) { + return "", fmt.Errorf("cgroup root must be absolute") + } + root = filepath.Clean(root) + resolved, err := filepath.EvalSymlinks(root) + if err != nil { + return "", err + } + if root != resolved || root == "/sys/fs/cgroup" || root == "/" { + return "", fmt.Errorf("use a dedicated delegated cgroup, not the hierarchy root or a symlink") + } + var st unix.Statfs_t + if err = unix.Statfs(root, &st); err != nil { + return "", err + } + if st.Type != unix.CGROUP2_SUPER_MAGIC { + return "", fmt.Errorf("%s is not cgroup v2", root) + } + return root, nil +} +func validTaskName(name string) bool { + if !strings.HasPrefix(name, "task-") || len(name) != 41 { + return false + } + for _, c := range name[5:] { + if !(c >= '0' && c <= '9' || c >= 'a' && c <= 'f' || c == '-') { + return false + } + } + return true +} + +type cgroupGroup struct { + mu sync.Mutex + path string + dir *os.File + watcher *watchdog + closed bool +} + +func newPlatformGroup(id string, o Options) (Group, error) { + if o.CgroupRoot == "" { + if o.Mode == "required" { + return nil, fmt.Errorf("required isolation has no delegated cgroup root") + } + return newUnixGroup() + } + root, err := validateRoot(o.CgroupRoot) + if err != nil { + return nil, err + } + name := "task-" + id + if !validTaskName(name) { + return nil, fmt.Errorf("invalid task run ID") + } + path := filepath.Join(root, name) + if err = os.Mkdir(path, 0700); err != nil { + return nil, err + } + success := false + defer func() { + if !success { + _ = os.Remove(path) + } + }() + limits := map[string]string{"pids.max": strconv.Itoa(o.MaxProcesses), "memory.max": strconv.FormatInt(o.MemoryMaxBytes, 10), "memory.oom.group": "1"} + if o.CPUQuotaMicros > 0 { + limits["cpu.max"] = fmt.Sprintf("%d 100000", o.CPUQuotaMicros) + } + for file, value := range limits { + if err = os.WriteFile(filepath.Join(path, file), []byte(value), 0600); err != nil { + return nil, fmt.Errorf("set %s: %w", file, err) + } + } + if _, err = os.Stat(filepath.Join(path, "cgroup.kill")); err != nil { + return nil, fmt.Errorf("cgroup.kill requires Linux 5.14+: %w", err) + } + dir, err := os.Open(path) + if err != nil { + return nil, err + } + g := &cgroupGroup{path: path, dir: dir} + w, err := startWatchdog(watchRequest{Name: "cgroup", Path: path}, func() { + // A guardian crash is also fail-closed while the owner is still alive. + _ = os.WriteFile(filepath.Join(path, "cgroup.kill"), []byte("1"), 0600) + }) + if err != nil { + dir.Close() + return nil, err + } + g.watcher = w + success = true + return g, nil +} +func (g *cgroupGroup) Name() string { return "cgroup_v2" } +func (g *cgroupGroup) Prepare(cmd *exec.Cmd) (*Launch, error) { + g.mu.Lock() + defer g.mu.Unlock() + if g.closed { + return nil, fmt.Errorf("cgroup is closed") + } + if _, err := g.watcher.send(watchRequest{Op: "ping"}); err != nil { + return nil, err + } + if cmd.SysProcAttr == nil { + cmd.SysProcAttr = &syscall.SysProcAttr{} + } + // clone3(CLONE_INTO_CGROUP), not a racy write of a newly started PID. + cmd.SysProcAttr.UseCgroupFD = true + cmd.SysProcAttr.CgroupFD = int(g.dir.Fd()) + return &Launch{Commit: func() error { return nil }, Dispose: func() {}}, nil +} +func (g *cgroupGroup) Release(pid int) error { return nil } +func (g *cgroupGroup) Close(ctx context.Context) error { + g.mu.Lock() + defer g.mu.Unlock() + g.closed = true + if g.dir == nil { + return nil + } + if err := killAndRemoveCgroup(ctx, g.path); err != nil { + return err + } + watchErr := g.watcher.close() + err := errors.Join(watchErr, g.dir.Close()) + g.dir = nil + return err +} +func killAndRemoveCgroup(ctx context.Context, path string) error { + if err := os.WriteFile(filepath.Join(path, "cgroup.kill"), []byte("1"), 0600); err != nil { + if os.IsNotExist(err) { + return nil + } + return err + } + for { + data, err := os.ReadFile(filepath.Join(path, "cgroup.events")) + if os.IsNotExist(err) { + return nil + } + if err != nil { + return err + } + if strings.Contains(string(data), "populated 0") { + return removeCgroupTree(path) + } + select { + case <-ctx.Done(): + return ctx.Err() + case <-time.After(10 * time.Millisecond): + } + } +} +func removeCgroupTree(path string) error { + entries, err := os.ReadDir(path) + if os.IsNotExist(err) { + return nil + } + if err != nil { + return err + } + for _, e := range entries { + if e.IsDir() { + if err = removeCgroupTree(filepath.Join(path, e.Name())); err != nil { + return err + } + } + } + err = os.Remove(path) + if os.IsNotExist(err) { + return nil + } + return err +} +func guardianMain(dec *json.Decoder, enc *json.Encoder) error { + var req watchRequest + if err := dec.Decode(&req); err != nil { + return err + } + if req.Name == "process_group" { + return groupGuardian(dec, enc) + } + if req.Name != "cgroup" || !validTaskName(filepath.Base(req.Path)) { + return fmt.Errorf("invalid cgroup guardian") + } + if _, err := validateRoot(req.Path); err != nil { + return err + } + return serveGuardian(dec, enc, func(r watchRequest) error { + if r.Op != "ping" { + return fmt.Errorf("unknown command") + } + return nil + }, func() error { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + return killAndRemoveCgroup(ctx, req.Path) + }) +} diff --git a/internal/processguard/platform_other.go b/internal/processguard/platform_other.go new file mode 100644 index 00000000..1bd6e5f8 --- /dev/null +++ b/internal/processguard/platform_other.go @@ -0,0 +1,31 @@ +//go:build !linux && !windows + +package processguard + +import ( + "encoding/json" + "fmt" +) + +func configurePlatform(o *Options) error { + if o.Mode == "required" || o.CgroupRoot != "" { + return fmt.Errorf("kernel task containment is unavailable on this OS; use a Linux cgroup deployment") + } + return nil +} +func newPlatformGroup(id string, o Options) (Group, error) { + if err := configurePlatform(&o); err != nil { + return nil, err + } + return newUnixGroup() +} +func guardianMain(dec *json.Decoder, enc *json.Encoder) error { + var req watchRequest + if err := dec.Decode(&req); err != nil { + return err + } + if req.Name != "process_group" { + return fmt.Errorf("unsupported guardian") + } + return groupGuardian(dec, enc) +} diff --git a/internal/processguard/platform_windows.go b/internal/processguard/platform_windows.go new file mode 100644 index 00000000..bb4e136e --- /dev/null +++ b/internal/processguard/platform_windows.go @@ -0,0 +1,183 @@ +//go:build windows + +package processguard + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "os" + "os/exec" + "runtime" + "sync" + "syscall" + "time" + "unsafe" + + "golang.org/x/sys/windows" +) + +func configurePlatform(o *Options) error { + if o.CgroupRoot != "" { + return fmt.Errorf("cgroups are Linux-only") + } + if o.Mode == "process_group" { + return fmt.Errorf("Windows tasks require Job Object containment") + } + return nil +} +func configureGuardian(cmd *exec.Cmd) { + cmd.SysProcAttr = &syscall.SysProcAttr{CreationFlags: syscall.CREATE_NEW_PROCESS_GROUP} +} +func gatedChildMain(args []string) error { + return fmt.Errorf("Unix launch gates are unavailable on Windows") +} + +type jobGroup struct { + mu sync.Mutex + job windows.Handle + parent windows.Handle + watcher *watchdog + closed bool + broken bool +} + +func newPlatformGroup(id string, o Options) (Group, error) { + if err := configurePlatform(&o); err != nil { + return nil, err + } + name := "Local\\CyberStrikeAI-" + id + g := &jobGroup{} + w, err := startWatchdog(watchRequest{Name: name, Options: o}, func() { + g.mu.Lock() + defer g.mu.Unlock() + g.broken = true + if !g.closed && g.job != 0 { + _ = windows.TerminateJobObject(g.job, 1) + } + }) + if err != nil { + return nil, err + } + fail := func(err error) (Group, error) { w.close(); return nil, err } + namePtr, err := windows.UTF16PtrFromString(name) + if err != nil { + return fail(err) + } + proc := windows.NewLazySystemDLL("kernel32.dll").NewProc("OpenJobObjectW") + h, _, callErr := proc.Call(0x0004|0x0008, 0, uintptr(unsafe.Pointer(namePtr))) + if h == 0 { + return fail(callErr) + } + parent, err := windows.OpenProcess(windows.PROCESS_CREATE_PROCESS|windows.PROCESS_DUP_HANDLE, false, uint32(w.cmd.Process.Pid)) + if err != nil { + windows.CloseHandle(windows.Handle(h)) + return fail(err) + } + g.mu.Lock() + defer g.mu.Unlock() + if g.broken { + windows.CloseHandle(parent) + windows.CloseHandle(windows.Handle(h)) + return fail(fmt.Errorf("job guardian exited during setup")) + } + g.job = windows.Handle(h) + g.parent = parent + g.watcher = w + return g, nil +} +func (g *jobGroup) Name() string { return "windows_job" } +func (g *jobGroup) Prepare(cmd *exec.Cmd) (*Launch, error) { + g.mu.Lock() + defer g.mu.Unlock() + if g.closed || g.broken { + return nil, fmt.Errorf("job is closed") + } + if _, err := g.watcher.send(watchRequest{Op: "ping"}); err != nil { + return nil, err + } + if cmd.SysProcAttr == nil { + cmd.SysProcAttr = &syscall.SysProcAttr{} + } + // Windows inherits the job at CreateProcess time from this parent. The + // guardian joined the job BEFORE acknowledging readiness, closing the + // Start-then-Assign race and its suspended-process crash window. + cmd.SysProcAttr.ParentProcess = syscall.Handle(g.parent) + return &Launch{Commit: func() error { return nil }, Dispose: func() {}}, nil +} +func (g *jobGroup) Release(pid int) error { return nil } +func (g *jobGroup) Close(ctx context.Context) error { + g.mu.Lock() + defer g.mu.Unlock() + if g.closed { + return nil + } + if err := windows.TerminateJobObject(g.job, 1); err != nil { + return err + } + type accounting struct { + TotalUser, TotalKernel, PeriodUser, PeriodKernel int64 + PageFaults, TotalProcesses, ActiveProcesses, Terminated uint32 + } + for { + var info accounting + if err := windows.QueryInformationJobObject(g.job, windows.JobObjectBasicAccountingInformation, uintptr(unsafe.Pointer(&info)), uint32(unsafe.Sizeof(info)), nil); err != nil { + return err + } + if info.ActiveProcesses == 0 { + break + } + select { + case <-ctx.Done(): + return ctx.Err() + case <-time.After(10 * time.Millisecond): + } + } + g.closed = true + return errors.Join(g.watcher.close(), windows.CloseHandle(g.parent), windows.CloseHandle(g.job)) +} +func guardianMain(dec *json.Decoder, enc *json.Encoder) error { + var req watchRequest + if err := dec.Decode(&req); err != nil { + return err + } + name, err := windows.UTF16PtrFromString(req.Name) + if err != nil { + return err + } + job, err := windows.CreateJobObject(nil, name) + if err != nil { + return err + } + defer windows.CloseHandle(job) + limits := windows.JOBOBJECT_EXTENDED_LIMIT_INFORMATION{} + limits.BasicLimitInformation.LimitFlags = windows.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE | windows.JOB_OBJECT_LIMIT_ACTIVE_PROCESS | windows.JOB_OBJECT_LIMIT_JOB_MEMORY + limits.BasicLimitInformation.ActiveProcessLimit = uint32(req.Options.MaxProcesses + 1) + limits.JobMemoryLimit = uintptr(req.Options.MemoryMaxBytes) + if _, err = windows.SetInformationJobObject(job, windows.JobObjectExtendedLimitInformation, uintptr(unsafe.Pointer(&limits)), uint32(unsafe.Sizeof(limits))); err != nil { + return err + } + if req.Options.CPUQuotaMicros > 0 { + rate := req.Options.CPUQuotaMicros / 10 / int64(runtime.NumCPU()) + if rate < 1 { + rate = 1 + } + if rate > 10000 { + rate = 10000 + } + cpu := struct{ Flags, Rate uint32 }{Flags: 1 | 4, Rate: uint32(rate)} + if _, err = windows.SetInformationJobObject(job, windows.JobObjectCpuRateControlInformation, uintptr(unsafe.Pointer(&cpu)), uint32(unsafe.Sizeof(cpu))); err != nil { + return err + } + } + if err = windows.AssignProcessToJobObject(job, windows.CurrentProcess()); err != nil { + return err + } + return serveGuardian(dec, enc, func(req watchRequest) error { + if req.Op != "ping" { + return fmt.Errorf("unknown guardian command") + } + return nil + }, func() error { return windows.TerminateJobObject(job, uint32(os.Getpid())) }) +} diff --git a/internal/processguard/watchdog.go b/internal/processguard/watchdog.go new file mode 100644 index 00000000..ec92eceb --- /dev/null +++ b/internal/processguard/watchdog.go @@ -0,0 +1,182 @@ +package processguard + +import ( + "encoding/json" + "fmt" + "io" + "os" + "os/exec" + "sync" + "time" +) + +const guardianArg = "--cyberstrike-internal-process-guardian" +const childArg = "--cyberstrike-internal-process-child" + +type watchRequest struct { + Op string + PID int + Path string + Name string + Options Options +} +type watchReply struct { + PID int + Error string +} + +type watchdog struct { + mu sync.Mutex + cmd *exec.Cmd + input *os.File + output *os.File + encoder *json.Encoder + decoder *json.Decoder + done chan struct{} + failed error +} + +// The re-exec modes run before application configuration, listeners or MCP +// initialization. Stdin is a private pipe; no network control port is opened. +func init() { + if len(os.Args) < 2 { + return + } + switch os.Args[1] { + case guardianArg: + err := guardianMain(json.NewDecoder(os.Stdin), json.NewEncoder(os.Stdout)) + if err != nil { + _ = json.NewEncoder(os.Stdout).Encode(watchReply{Error: err.Error()}) + fmt.Fprintln(os.Stderr, err) + os.Exit(1) + } + os.Exit(0) + case childArg: + if err := gatedChildMain(os.Args[2:]); err != nil { + fmt.Fprintln(os.Stderr, err) + os.Exit(1) + } + os.Exit(0) + } +} + +func startWatchdog(req watchRequest, onExit func()) (*watchdog, error) { + exe, err := os.Executable() + if err != nil { + return nil, err + } + cmd := exec.Command(exe, guardianArg) + configureGuardian(cmd) + in, err := cmd.StdinPipe() + if err != nil { + return nil, err + } + out, err := cmd.StdoutPipe() + if err != nil { + in.Close() + return nil, err + } + // No inherited stderr pipe that could keep a caller's output reader alive. + if err = cmd.Start(); err != nil { + in.Close() + out.Close() + return nil, err + } + w := &watchdog{cmd: cmd, input: in.(*os.File), output: out.(*os.File), done: make(chan struct{})} + w.encoder = json.NewEncoder(w.input) + w.decoder = json.NewDecoder(w.output) + // The guardian only exits after EOF or failure; exit invalidates all RPCs. + go func() { + _ = cmd.Wait() + close(w.done) + if onExit != nil { + onExit() + } + }() + req.Op = "init" + if _, err = w.send(req); err != nil { + w.close() + return nil, err + } + return w, nil +} + +func (w *watchdog) send(req watchRequest) (watchReply, error) { + w.mu.Lock() + defer w.mu.Unlock() + if w.failed != nil { + return watchReply{}, w.failed + } + type response struct { + reply watchReply + err error + } + result := make(chan response, 1) + // Pipe deadlines are not supported by every Windows pipe implementation. + // On timeout kill the helper and close both ends to release this goroutine. + go func() { + if err := w.encoder.Encode(req); err != nil { + result <- response{err: err} + return + } + var reply watchReply + err := w.decoder.Decode(&reply) + if err == nil && reply.Error != "" { + err = fmt.Errorf("process guardian: %s", reply.Error) + } + result <- response{reply, err} + }() + select { + case r := <-result: + w.failed = r.err + return r.reply, r.err + case <-time.After(3 * time.Second): + _ = w.cmd.Process.Kill() + _ = w.input.Close() + _ = w.output.Close() + w.failed = fmt.Errorf("process guardian acknowledgement timed out") + return watchReply{}, w.failed + } +} + +func (w *watchdog) close() error { + w.mu.Lock() + _ = w.input.Close() + w.mu.Unlock() + select { + case <-w.done: + case <-time.After(3 * time.Second): + _ = w.cmd.Process.Kill() + select { + case <-w.done: + case <-time.After(3 * time.Second): + return fmt.Errorf("process guardian did not exit") + } + } + _ = w.output.Close() + return nil +} + +func serveGuardian(dec *json.Decoder, enc *json.Encoder, apply func(watchRequest) error, cleanup func() error) error { + defer cleanup() + if err := enc.Encode(watchReply{PID: os.Getpid()}); err != nil { + return err + } + for { + var req watchRequest + if err := dec.Decode(&req); err != nil { + if err == io.EOF { + return nil + } + return err + } + err := apply(req) + reply := watchReply{PID: os.Getpid()} + if err != nil { + reply.Error = err.Error() + } + if err := enc.Encode(reply); err != nil { + return err + } + } +} diff --git a/internal/runlease/scope.go b/internal/runlease/scope.go new file mode 100644 index 00000000..69145cc0 --- /dev/null +++ b/internal/runlease/scope.go @@ -0,0 +1,126 @@ +// Package runlease binds asynchronous tool workers to a task run even when +// their contexts detach from a per-call timeout or an SSE connection. +package runlease + +import ( + "context" + "errors" + "fmt" + "sort" + "strings" + "sync" +) + +var ErrClosed = errors.New("task is ending; new tool executions are not allowed") +var ErrUnconfirmed = errors.New("remote cancellation is unconfirmed") + +// Bound detached worker fan-out independently of OS process limits. +const MaxTaskWorkers = 256 + +type contextKey struct{} +type Scope struct { + mu sync.Mutex + sealed bool + workers map[string]context.CancelFunc + unconfirmed map[string]string + changed chan struct{} +} + +func New() *Scope { + return &Scope{workers: make(map[string]context.CancelFunc), unconfirmed: make(map[string]string), changed: make(chan struct{})} +} +func WithScope(ctx context.Context, s *Scope) context.Context { + return context.WithValue(ctx, contextKey{}, s) +} +func FromContext(ctx context.Context) *Scope { + if ctx == nil { + return nil + } + s, _ := ctx.Value(contextKey{}).(*Scope) + return s +} +func (s *Scope) notify() { close(s.changed); s.changed = make(chan struct{}) } +func (s *Scope) Register(id string, cancel context.CancelFunc) (func(), error) { + if s == nil { + return func() {}, nil + } + s.mu.Lock() + defer s.mu.Unlock() + if s.sealed { + return nil, ErrClosed + } + if len(s.workers) >= MaxTaskWorkers { + return nil, fmt.Errorf("task worker limit reached (%d)", MaxTaskWorkers) + } + if _, ok := s.workers[id]; ok { + return nil, fmt.Errorf("duplicate worker %s", id) + } + s.workers[id] = cancel + var once sync.Once + return func() { once.Do(func() { s.mu.Lock(); delete(s.workers, id); s.notify(); s.mu.Unlock() }) }, nil +} +func (s *Scope) Seal() { + if s == nil { + return + } + s.mu.Lock() + s.sealed = true + s.mu.Unlock() +} +func (s *Scope) Cancel() { + if s == nil { + return + } + s.mu.Lock() + s.sealed = true + cs := make([]context.CancelFunc, 0, len(s.workers)) + for _, c := range s.workers { + cs = append(cs, c) + } + s.mu.Unlock() + for _, c := range cs { + if c != nil { + c() + } + } +} +func (s *Scope) MarkUnconfirmed(id, message string) { + if s == nil { + return + } + s.mu.Lock() + s.unconfirmed[id] = message + s.notify() + s.mu.Unlock() +} +func (s *Scope) Wait(ctx context.Context) error { + if s == nil { + return nil + } + for { + s.mu.Lock() + pending := make([]string, 0, len(s.workers)) + for id := range s.workers { + pending = append(pending, id) + } + uncertain := make([]string, 0, len(s.unconfirmed)) + for id, msg := range s.unconfirmed { + uncertain = append(uncertain, id+": "+msg) + } + changed := s.changed + s.mu.Unlock() + if len(pending) == 0 { + if len(uncertain) > 0 { + sort.Strings(uncertain) + return fmt.Errorf("%w: %s", ErrUnconfirmed, strings.Join(uncertain, "; ")) + } + return nil + } + select { + case <-changed: + case <-ctx.Done(): + sort.Strings(pending) + return fmt.Errorf("tool workers still running %v: %w", pending, ctx.Err()) + } + } +} diff --git a/internal/runlease/scope_test.go b/internal/runlease/scope_test.go new file mode 100644 index 00000000..7cc67200 --- /dev/null +++ b/internal/runlease/scope_test.go @@ -0,0 +1,70 @@ +package runlease + +import ( + "context" + "errors" + "sync" + "testing" + "time" +) + +func TestConcurrentAdmissionAndCancellation(t *testing.T) { + scope := New() + ctx := WithScope(context.Background(), scope) + if FromContext(context.WithoutCancel(ctx)) != scope { + t.Fatal("detachment lost task ownership") + } + var wg sync.WaitGroup + for i := 0; i < 64; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + worker, cancel := context.WithCancel(context.Background()) + defer cancel() + release, err := scope.Register(string(rune('a'+i)), cancel) + if errors.Is(err, ErrClosed) { + return + } + if err != nil { + t.Error(err) + return + } + <-worker.Done() + release() + }(i) + } + scope.Cancel() + wg.Wait() + wait, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + if err := scope.Wait(wait); err != nil { + t.Fatal(err) + } +} +func TestDetachedWorkerCapacityReleasedOnCompletion(t *testing.T) { + scope := New() + releases := make([]func(), 0, MaxTaskWorkers) + for i := 0; i < MaxTaskWorkers; i++ { + release, err := scope.Register(string(rune(i)), func() {}) + if err != nil { + t.Fatal(err) + } + releases = append(releases, release) + } + if _, err := scope.Register("overflow", func() {}); err == nil { + t.Fatal("unbounded worker admission") + } + releases[0]() + release, err := scope.Register("replacement", func() {}) + if err != nil { + t.Fatal(err) + } + release() + for _, release := range releases { + release() + } + scope.Cancel() + if err = scope.Wait(context.Background()); err != nil { + t.Fatal(err) + } +} diff --git a/internal/security/executor.go b/internal/security/executor.go index a1f6ae97..0cf0e75a 100644 --- a/internal/security/executor.go +++ b/internal/security/executor.go @@ -1,7 +1,6 @@ package security import ( - "bufio" "context" "encoding/json" "fmt" @@ -9,7 +8,6 @@ import ( "os" "os/exec" "runtime" - "strconv" "strings" "sync" "time" @@ -836,128 +834,13 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int zap.Bool("isBackground", isBackground), ) - // 如果是后台命令，使用特殊处理来获取实际的后台进程PID if isBackground { - // 移除命令末尾的 & 符号 - commandWithoutAmpersand := strings.TrimSuffix(strings.TrimSpace(command), "&") - commandWithoutAmpersand = strings.TrimSpace(commandWithoutAmpersand) - - // 构建新命令：后台作业重定向标准流后 echo $pid（与 RedirectBackgroundJobStdio 一致）。 - pidCommand := RedirectBackgroundJobStdio(commandWithoutAmpersand+" &") + " pid=$!; echo $pid" - - // 创建新命令来获取PID - var pidCmd *exec.Cmd - if workDir != "" { - pidCmd = exec.CommandContext(ctx, shell, "-c", pidCommand) - pidCmd.Dir = workDir - } else { - pidCmd = exec.CommandContext(ctx, shell, "-c", pidCommand) - } - ConfigureShellCmdForAgentExecute(pidCmd) - - // 获取stdout管道 - stdout, err := pidCmd.StdoutPipe() + job := strings.TrimSpace(strings.TrimSuffix(strings.TrimSpace(command), "&")) + session, err := StartManagedBackground(ctx, shell, job, workDir) if err != nil { - e.logger.Error("创建stdout管道失败", - zap.String("command", command), - zap.Error(err), - ) - // 如果创建管道失败，使用shell进程的PID作为fallback - if err := pidCmd.Start(); err != nil { - return &mcp.ToolResult{ - Content: []mcp.Content{ - { - Type: "text", - Text: fmt.Sprintf("后台命令启动失败: %v", err), - }, - }, - IsError: true, - }, nil - } - pid := pidCmd.Process.Pid - go pidCmd.Wait() // 在后台等待，避免僵尸进程 - return &mcp.ToolResult{ - Content: []mcp.Content{ - { - Type: "text", - Text: fmt.Sprintf("后台命令已启动\n命令: %s\n进程ID: %d (可能不准确，获取PID失败)\n\n注意: 后台进程将继续运行，不会等待其完成。", command, pid), - }, - }, - IsError: false, - }, nil + return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: fmt.Sprintf("后台命令启动失败: %v", err)}}, IsError: true}, nil } - - // 启动命令 - if err := pidCmd.Start(); err != nil { - stdout.Close() - e.logger.Error("后台命令启动失败", - zap.String("command", command), - zap.Error(err), - ) - return &mcp.ToolResult{ - Content: []mcp.Content{ - { - Type: "text", - Text: fmt.Sprintf("后台命令启动失败: %v", err), - }, - }, - IsError: true, - }, nil - } - - // 读取第一行输出（PID） - reader := bufio.NewReader(stdout) - pidLine, err := reader.ReadString('\n') - stdout.Close() - - var actualPid int - if err != nil && err != io.EOF { - e.logger.Warn("读取后台进程PID失败", - zap.String("command", command), - zap.Error(err), - ) - // 如果读取失败，使用shell进程的PID - actualPid = pidCmd.Process.Pid - } else { - // 解析PID - pidStr := strings.TrimSpace(pidLine) - if parsedPid, err := strconv.Atoi(pidStr); err == nil { - actualPid = parsedPid - } else { - e.logger.Warn("解析后台进程PID失败", - zap.String("command", command), - zap.String("pidLine", pidStr), - zap.Error(err), - ) - // 如果解析失败，使用shell进程的PID - actualPid = pidCmd.Process.Pid - } - } - - // 在goroutine中等待shell进程，避免僵尸进程 - go func() { - if err := pidCmd.Wait(); err != nil { - e.logger.Debug("后台命令shell进程执行完成", - zap.String("command", command), - zap.Error(err), - ) - } - }() - - e.logger.Info("后台命令已启动", - zap.String("command", command), - zap.Int("actualPid", actualPid), - ) - - return &mcp.ToolResult{ - Content: []mcp.Content{ - { - Type: "text", - Text: fmt.Sprintf("后台命令已启动\n命令: %s\n进程ID: %d\n\n注意: 后台进程将继续运行，不会等待其完成。", command, actualPid), - }, - }, - IsError: false, - }, nil + return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: fmt.Sprintf("后台命令已启动\n命令: %s\n进程组ID: %d\n\n后台进程由本轮任务托管，任务结束时自动清理。", command, session.rootPID)}}}, nil } // 非后台命令：等待输出 @@ -1041,7 +924,7 @@ func combinedOutputCancellableWithLimit(ctx context.Context, cmd *exec.Cmd, maxB cmd.Stdout = stdoutBuf cmd.Stderr = stderrBuf - session, err := StartShellSession(cmd) + session, err := StartShellSessionContext(ctx, cmd) if err != nil { return "", err } @@ -1248,7 +1131,7 @@ func streamCommandOutput(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallba _ = stdoutPipe.Close() return "", err } - session, err := StartShellSession(cmd) + session, err := StartShellSessionContext(ctx, cmd) if err != nil { _ = stdoutPipe.Close() _ = stderrPipe.Close() @@ -1265,6 +1148,8 @@ func streamCommandOutput(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallba }() defer close(stopWatch) + readStop := make(chan struct{}) + defer close(readStop) chunks := make(chan string, 64) var wg sync.WaitGroup readFn := func(r io.Reader) { @@ -1273,7 +1158,11 @@ func streamCommandOutput(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallba for { n, readErr := r.Read(buf) if n > 0 { - chunks <- string(buf[:n]) + select { + case chunks <- string(buf[:n]): + case <-readStop: + return + } } if readErr != nil { return @@ -1422,24 +1311,24 @@ func runCommandWithPTY(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback } _ = prepareShellCmdSession(cmd) - ptmx, err := pty.Start(cmd) + var ptmx *os.File + session, err := startShellSessionContext(ctx, cmd, func() error { + var startErr error + ptmx, startErr = pty.Start(cmd) + return startErr + }) if err != nil { return "", err } defer func() { _ = ptmx.Close() }() - rootPID := 0 - if cmd.Process != nil { - rootPID = cmd.Process.Pid - } - // ctx 取消时尽快终止子进程 done := make(chan struct{}) go func() { select { case <-ctx.Done(): _ = ptmx.Close() // 触发读退出 - terminateProcessGroup(rootPID, cmd) + session.Terminate() case <-done: } }() @@ -1484,7 +1373,7 @@ func runCommandWithPTY(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback } flush() - waitErr := cmd.Wait() + waitErr := session.Wait() return finalizeBoundedOutput(outBuilder, maxBytes, tee), waitErr } diff --git a/internal/security/executor_test.go b/internal/security/executor_test.go index 4b62889a..96bb65ff 100644 --- a/internal/security/executor_test.go +++ b/internal/security/executor_test.go @@ -54,7 +54,13 @@ func TestExecuteSystemCommand_BackgroundDoesNotBlockOnChildStdout(t *testing.T) executor, _ := setupTestExecutor(t) // 子进程先向 stdout 写无换行字符再长时间 sleep；若与 echo $pid 共享管道且未重定向子进程 stdout， // ReadString('\n') 会阻塞到子进程退出。后台包装须将子进程标准流与 PID 行分离。 - ctx, cancel := context.WithTimeout(context.Background(), 4*time.Second) + scope := NewProcessScope() + t.Cleanup(func() { + if err := scope.Close(); err != nil { + t.Error(err) + } + }) + ctx, cancel := context.WithTimeout(WithProcessScope(context.Background(), scope), 4*time.Second) defer cancel() args := map[string]interface{}{ "command": `(sh -c 'printf x; sleep 120') &`, diff --git a/internal/security/procattr_unix.go b/internal/security/procattr_unix.go index 8f516ec8..9969a2c4 100644 --- a/internal/security/procattr_unix.go +++ b/internal/security/procattr_unix.go @@ -39,3 +39,14 @@ func terminateProcessGroup(rootPID int, cmd *exec.Cmd) { func terminateCmdTree(cmd *exec.Cmd) { terminateProcessGroup(0, cmd) } + +// stopProcessGroup gives the whole job a grace period to release resources. +func stopProcessGroup(pid int, cmd *exec.Cmd) { + if pid > 0 { + _ = syscall.Kill(-pid, syscall.SIGTERM) + } +} + +func processGroupExists(pid int) bool { + return pid > 0 && syscall.Kill(-pid, 0) != syscall.ESRCH +} diff --git a/internal/security/procattr_windows.go b/internal/security/procattr_windows.go index af7da8c1..c7bbdf1b 100644 --- a/internal/security/procattr_windows.go +++ b/internal/security/procattr_windows.go @@ -3,9 +3,11 @@ package security import ( + "context" "os/exec" "strconv" "syscall" + "time" ) func prepareShellCmdSession(cmd *exec.Cmd) error { @@ -29,7 +31,9 @@ func terminateProcessGroup(rootPID int, cmd *exec.Cmd) { if pid <= 0 { return } - tk := exec.Command("taskkill", "/F", "/T", "/PID", strconv.Itoa(pid)) + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + tk := exec.CommandContext(ctx, "taskkill", "/F", "/T", "/PID", strconv.Itoa(pid)) if err := tk.Run(); err != nil { if cmd != nil && cmd.Process != nil { _ = cmd.Process.Kill() @@ -41,3 +45,12 @@ func terminateProcessGroup(rootPID int, cmd *exec.Cmd) { func terminateCmdTree(cmd *exec.Cmd) { terminateProcessGroup(0, cmd) } + +func stopProcessGroup(pid int, cmd *exec.Cmd) { + // Windows has no portable SIGTERM equivalent for arbitrary console jobs. + terminateProcessGroup(pid, cmd) +} + +// Windows taskkill /T is best effort; unlike a Unix PGID it has no persistent +// group handle to query after the root exits. Job Objects are needed for that. +func processGroupExists(pid int) bool { return false } diff --git a/internal/security/process_scope.go b/internal/security/process_scope.go new file mode 100644 index 00000000..48a7ca2c --- /dev/null +++ b/internal/security/process_scope.go @@ -0,0 +1,220 @@ +package security + +import ( + "context" + "errors" + "fmt" + "os/exec" + "sync" + "time" + + "cyberstrike-ai/internal/processguard" + "github.com/google/uuid" +) + +var ErrProcessScopeClosed = errors.New("task is ending; new processes are not allowed") +var ErrBackgroundNeedsTask = errors.New("background commands require a managed task") + +type processScopeKey struct{} + +// ProcessScope owns local commands for one task run, including background work. +// Ownership is carried by context values, so MCP's WithoutCancel retains it. +// Start and Seal serialize under the same lock: no process can escape cleanup +// by starting between the final snapshot and task completion. +type ProcessScope struct { + ID string + mu sync.Mutex + closed bool + sessions map[*ShellSession]struct{} + guard processguard.Group + guardErr error + closeMu sync.Mutex +} + +func NewProcessScope() *ProcessScope { + return &ProcessScope{ID: uuid.NewString(), sessions: make(map[*ShellSession]struct{})} +} + +func WithProcessScope(ctx context.Context, scope *ProcessScope) context.Context { + return context.WithValue(ctx, processScopeKey{}, scope) +} + +func ProcessScopeFromContext(ctx context.Context) *ProcessScope { + if ctx == nil { + return nil + } + scope, _ := ctx.Value(processScopeKey{}).(*ProcessScope) + return scope +} + +func (s *ProcessScope) Seal() { + if s == nil { + return + } + s.mu.Lock() + s.closed = true + s.mu.Unlock() +} + +func startShellSessionContext(ctx context.Context, cmd *exec.Cmd, start func() error) (*ShellSession, error) { + scope := ProcessScopeFromContext(ctx) + if scope != nil { + scope.mu.Lock() + defer scope.mu.Unlock() + if scope.closed { + return nil, ErrProcessScopeClosed + } + } + if err := ctx.Err(); err != nil { + return nil, err + } + if err := prepareShellCmdSession(cmd); err != nil { + return nil, err + } + + var launch *processguard.Launch + if scope != nil { + if scope.guard == nil && scope.guardErr == nil { + scope.guard, scope.guardErr = processguard.New(scope.ID) + } + if scope.guardErr != nil { + return nil, scope.guardErr + } + var err error + launch, err = scope.guard.Prepare(cmd) + if err != nil { + return nil, err + } + defer launch.Dispose() + } + // Bound Go's output-copy goroutines when descendants inherit a pipe. + if cmd.WaitDelay == 0 { + cmd.WaitDelay = 2 * time.Second + } + if err := start(); err != nil { + return nil, err + } + if launch != nil { + if err := launch.Commit(); err != nil { + terminateProcessGroup(cmd.Process.Pid, cmd) + _ = cmd.Wait() + if scope != nil { + _ = scope.guard.Release(cmd.Process.Pid) + } + return nil, err + } + } + session := &ShellSession{Cmd: cmd, rootPID: cmd.Process.Pid, scope: scope, done: make(chan struct{})} + if scope != nil { + scope.sessions[session] = struct{}{} + } + return session, nil +} + +// Close seals the scope, asks every process group to exit, then escalates to +// SIGKILL. It waits for command reaping, with one shared deadline, not N timeouts. +// Failed entries remain owned, permitting a later Close to retry cleanup. +func (s *ProcessScope) Close() error { + if s == nil { + return nil + } + s.closeMu.Lock() + defer s.closeMu.Unlock() + s.mu.Lock() + s.closed = true + sessions := make([]*ShellSession, 0, len(s.sessions)) + for session := range s.sessions { + sessions = append(sessions, session) + } + s.mu.Unlock() + if len(sessions) == 0 { + return s.closeGuard() + } + for _, session := range sessions { + session.signal(false) + } + if waitShellSessions(sessions, 3*time.Second) { + return s.closeGuard() + } + for _, session := range sessions { + session.Terminate() + } + guardErr := s.closeGuard() + if waitShellSessions(sessions, 3*time.Second) { + return guardErr + } + remaining := make([]int, 0, len(sessions)) + for _, session := range sessions { + if !session.tryComplete() { + remaining = append(remaining, session.rootPID) + } + } + return fmt.Errorf("task %s: process cleanup timed out (process groups %v)", s.ID, remaining) +} + +func waitShellSessions(sessions []*ShellSession, timeout time.Duration) bool { + deadline := time.NewTimer(timeout) + defer deadline.Stop() + ticker := time.NewTicker(10 * time.Millisecond) + defer ticker.Stop() + for { + complete := true + for _, session := range sessions { + if !session.tryComplete() { + complete = false + } + } + if complete { + return true + } + select { + case <-deadline.C: + return false + case <-ticker.C: + } + } + +} + +// StartManagedBackground returns promptly while retaining task ownership. The +// shell executes the job in the foreground internally, keeping a waitable root +// alive; tool completion must not cancel the job's lifetime. +func StartManagedBackground(ctx context.Context, shell, command, dir string) (*ShellSession, error) { + if ProcessScopeFromContext(ctx) == nil { + return nil, ErrBackgroundNeedsTask + } + cmd := exec.Command(shell, "-c", PrepareShellCommandForExecute(command)) + cmd.Dir = dir + ConfigureShellCmdForAgentExecute(cmd) + // Nil output streams use /dev/null; background output cannot hold tool pipes. + session, err := StartShellSessionContext(ctx, cmd) + if err != nil { + return nil, err + } + go func() { _ = session.Wait() }() + return session, nil +} + +func (s *ProcessScope) closeGuard() error { + if s.guard == nil { + return nil + } + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + return s.guard.Close(ctx) +} + +func (s *ProcessScope) IsolationBackend() string { + if s == nil { + return "none" + } + s.mu.Lock() + defer s.mu.Unlock() + if s.guard != nil { + return s.guard.Name() + } + if s.guardErr != nil { + return "unavailable" + } + return "pending" +} diff --git a/internal/security/process_scope_test.go b/internal/security/process_scope_test.go new file mode 100644 index 00000000..b13146a0 --- /dev/null +++ b/internal/security/process_scope_test.go @@ -0,0 +1,198 @@ +//go:build !windows + +package security + +import ( + "context" + "errors" + "fmt" + "io" + "os" + "os/exec" + "path/filepath" + "strconv" + "strings" + "sync" + "syscall" + "testing" + "time" + + "github.com/cloudwego/eino/adk/filesystem" +) + +func readTestPID(t *testing.T, path string) int { + t.Helper() + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + b, err := os.ReadFile(path) + if err == nil { + if pid, err := strconv.Atoi(strings.TrimSpace(string(b))); err == nil && pid > 0 { + return pid + } + } + time.Sleep(10 * time.Millisecond) + } + t.Fatalf("process did not write PID to %s", path) + return 0 +} + +func requireProcessGone(t *testing.T, pid int) { + t.Helper() + deadline := time.Now().Add(time.Second) + for time.Now().Before(deadline) { + if syscall.Kill(pid, 0) == syscall.ESRCH { + return + } + time.Sleep(10 * time.Millisecond) + } + t.Fatalf("process %d survived task cleanup", pid) +} + +func TestProcessScope_BackgroundSurvivesToolButEndsWithTask(t *testing.T) { + executor, _ := setupTestExecutor(t) + scope := NewProcessScope() + t.Cleanup(func() { _ = scope.Close() }) + taskCtx := WithProcessScope(context.Background(), scope) + ctx, cancel := context.WithCancel(context.WithoutCancel(taskCtx)) + defer cancel() + pidFile := filepath.Join(t.TempDir(), "pid") + result, err := executor.executeSystemCommand(ctx, map[string]interface{}{ + "command": fmt.Sprintf("echo $$ > %q; sleep 300 &", pidFile), + }) + if err != nil || result.IsError { + t.Fatalf("background launch: %v, %+v", err, result) + } + pid := readTestPID(t, pidFile) + cancel() // MCP completes and cancels its per-tool context. + if err := syscall.Kill(pid, 0); err != nil { + t.Fatalf("tool completion killed task background process: %v", err) + } + if err := scope.Close(); err != nil { + t.Fatal(err) + } + requireProcessGone(t, pid) + if _, err := StartManagedBackground(taskCtx, "sh", "sleep 300", ""); !errors.Is(err, ErrProcessScopeClosed) { + t.Fatalf("closed task accepted a new process: %v", err) + } +} + +func TestProcessScope_EinoBackgroundReturnsPromptlyAndIsOwned(t *testing.T) { + for _, useFlag := range []bool{false, true} { + t.Run(fmt.Sprint(useFlag), func(t *testing.T) { + scope := NewProcessScope() + t.Cleanup(func() { _ = scope.Close() }) + ctx := WithProcessScope(context.Background(), scope) + pidFile := filepath.Join(t.TempDir(), "pid") + command := fmt.Sprintf("echo $$ > %q; sleep 300", pidFile) + if !useFlag { + command += " &" + } + stream, err := NewEinoStreamingShell().ExecuteStreaming(ctx, &filesystem.ExecuteRequest{Command: command, RunInBackendGround: useFlag}) + if err != nil { + t.Fatal(err) + } + defer stream.Close() + done := make(chan error, 1) + go func() { + for { + _, err := stream.Recv() + if err != nil { + done <- err + return + } + } + }() + select { + case err := <-done: + if !errors.Is(err, io.EOF) { + t.Fatal(err) + } + case <-time.After(time.Second): + t.Fatal("background launch waited for job completion") + } + pid := readTestPID(t, pidFile) + if err := scope.Close(); err != nil { + t.Fatal(err) + } + requireProcessGone(t, pid) + }) + } +} + +func TestProcessScope_ForceKillsIgnoringTERMAndGrandchild(t *testing.T) { + scope := NewProcessScope() + t.Cleanup(func() { _ = scope.Close() }) + ctx := WithProcessScope(context.Background(), scope) + pidFile := filepath.Join(t.TempDir(), "child") + session, err := StartManagedBackground(ctx, "sh", fmt.Sprintf("trap '' TERM; sleep 300 & echo $! > %q; wait", pidFile), "") + if err != nil { + t.Fatal(err) + } + childPID := readTestPID(t, pidFile) + if err := scope.Close(); err != nil { + t.Fatal(err) + } + requireProcessGone(t, session.rootPID) + requireProcessGone(t, childPID) + if session.Cmd.ProcessState == nil { + t.Fatal("root process was not reaped") + } +} + +func TestProcessScope_ConcurrentStartAndClose(t *testing.T) { + scope := NewProcessScope() + ctx := WithProcessScope(context.Background(), scope) + t.Cleanup(func() { _ = scope.Close() }) + var wg sync.WaitGroup + var mu sync.Mutex + var sessions []*ShellSession + begin := make(chan struct{}) + for i := 0; i < 24; i++ { + wg.Add(1) + go func() { + defer wg.Done() + <-begin + session, err := StartManagedBackground(ctx, "sh", "sleep 300", "") + if err != nil { + if !errors.Is(err, ErrProcessScopeClosed) { + t.Errorf("start: %v", err) + } + return + } + mu.Lock() + sessions = append(sessions, session) + mu.Unlock() + }() + } + close(begin) + if err := scope.Close(); err != nil { + t.Fatal(err) + } + wg.Wait() + for _, session := range sessions { + requireProcessGone(t, session.rootPID) + } +} + +func TestProcessScope_UnmanagedBackgroundRejected(t *testing.T) { + if _, err := StartManagedBackground(context.Background(), "sh", "sleep 300", ""); !errors.Is(err, ErrBackgroundNeedsTask) { + t.Fatal(err) + } +} + +func TestProcessScope_ForegroundExitKillsLeftoverChild(t *testing.T) { + scope := NewProcessScope() + t.Cleanup(func() { _ = scope.Close() }) + ctx := WithProcessScope(context.Background(), scope) + pidFile := filepath.Join(t.TempDir(), "child") + // A shell that exits with a redirected child must not lose that child. + cmd := exec.CommandContext(ctx, "sh", "-c", fmt.Sprintf("sleep 300 /dev/null 2>&1 & echo $! > %q", pidFile)) + if _, err := combinedOutputCancellable(ctx, cmd); err != nil { + t.Fatal(err) + } + pid := readTestPID(t, pidFile) + if err := scope.Close(); err != nil { + t.Fatal(err) + } + requireProcessGone(t, pid) +} diff --git a/internal/security/shell_execute_stream.go b/internal/security/shell_execute_stream.go index 02c5cb74..2ced0996 100644 --- a/internal/security/shell_execute_stream.go +++ b/internal/security/shell_execute_stream.go @@ -6,6 +6,7 @@ import ( "fmt" "io" "os/exec" + "strings" "sync" "github.com/cloudwego/eino/adk/filesystem" @@ -49,7 +50,7 @@ func (s *EinoStreamingShell) ExecuteStreaming(ctx context.Context, input *filesy } sr, w := schema.Pipe[*filesystem.ExecuteResponse](100) - if input.RunInBackendGround { + if input.RunInBackendGround || IsBackgroundShellCommand(input.Command) { go runShellInBackground(ctx, input.Command, w) return sr, nil } @@ -60,45 +61,18 @@ func (s *EinoStreamingShell) ExecuteStreaming(ctx context.Context, input *filesy func runShellInBackground(ctx context.Context, command string, w *schema.StreamWriter[*filesystem.ExecuteResponse]) { defer w.Close() - command = PrepareShellCommandForExecute(command) - cmd := exec.CommandContext(ctx, "/bin/sh", "-c", command) - applyDefaultTerminalEnv(cmd) - attachNonInteractiveStdin(cmd) - stdout, err := cmd.StdoutPipe() + command = strings.TrimSpace(command) + if IsBackgroundShellCommand(command) { + command = strings.TrimSpace(strings.TrimSuffix(command, "&")) + } + session, err := StartManagedBackground(ctx, "/bin/sh", command, "") if err != nil { - _ = w.Send(nil, fmt.Errorf("failed to create stdout pipe: %w", err)) + _ = w.Send(nil, err) return } - stderr, err := cmd.StderrPipe() - if err != nil { - _ = stdout.Close() - _ = w.Send(nil, fmt.Errorf("failed to create stderr pipe: %w", err)) - return - } - session, err := StartShellSession(cmd) - if err != nil { - _ = stdout.Close() - _ = stderr.Close() - _ = w.Send(nil, fmt.Errorf("failed to start command: %w", err)) - return - } - - done := make(chan struct{}) - go func() { - drainShellPipes(stdout, stderr) - _ = session.Wait() - close(done) - }() - - select { - case <-done: - case <-ctx.Done(): - TerminateShellCmdSession(session) - } - exitCode := 0 _ = w.Send(&filesystem.ExecuteResponse{ - Output: "command started in background\n", + Output: fmt.Sprintf("command started in background (process group %d); cleaned up when this task ends\n", session.rootPID), ExitCode: &exitCode, }, nil) } @@ -136,7 +110,7 @@ func streamShellForeground(ctx context.Context, command string, w *schema.Stream _ = w.Send(nil, fmt.Errorf("failed to create stderr pipe: %w", err)) return } - session, err := StartShellSession(cmd) + session, err := StartShellSessionContext(ctx, cmd) if err != nil { _ = stdoutPipe.Close() _ = stderrPipe.Close() @@ -154,6 +128,8 @@ func streamShellForeground(ctx context.Context, command string, w *schema.Stream }() defer close(stopWatch) + readStop := make(chan struct{}) + defer close(readStop) chunks := make(chan string, 64) var wg sync.WaitGroup readFn := func(r io.Reader) { @@ -162,7 +138,11 @@ func streamShellForeground(ctx context.Context, command string, w *schema.Stream for { n, readErr := r.Read(buf) if n > 0 { - chunks <- string(buf[:n]) + select { + case chunks <- string(buf[:n]): + case <-readStop: + return + } } if readErr != nil { return @@ -186,6 +166,7 @@ func streamShellForeground(ctx context.Context, command string, w *schema.Stream hadOutput = true if w.Send(&filesystem.ExecuteResponse{Output: chunk}, nil) { TerminateShellCmdSession(session) + go func() { _ = session.Wait() }() return } } diff --git a/internal/security/shell_session.go b/internal/security/shell_session.go index 72cb15e1..a01b50c8 100644 --- a/internal/security/shell_session.go +++ b/internal/security/shell_session.go @@ -1,47 +1,102 @@ package security -import "os/exec" +import ( + "context" + "os/exec" + "sync" + "time" +) -// ShellSession 在 Start 时记录根 shell 的进程组 ID，取消/超时时可杀整组（即使 cmd.Process 已失效）。 +// ShellSession caches the process group while its command is alive. Signals +// and Wait completion synchronize to avoid signalling already-released sessions. type ShellSession struct { - Cmd *exec.Cmd - rootPID int + Cmd *exec.Cmd + rootPID int + scope *ProcessScope + done chan struct{} + waitOnce sync.Once + waitErr error + signalMu sync.Mutex + finished bool + waited bool } -// StartShellSession 配置独立进程组并启动 shell，缓存 rootPID（Unix 下即 PGID）。 func StartShellSession(cmd *exec.Cmd) (*ShellSession, error) { - if err := prepareShellCmdSession(cmd); err != nil { - return nil, err - } - if err := cmd.Start(); err != nil { - return nil, err - } - pid := 0 - if cmd.Process != nil { - pid = cmd.Process.Pid - } - return &ShellSession{Cmd: cmd, rootPID: pid}, nil + return StartShellSessionContext(context.Background(), cmd) +} + +func StartShellSessionContext(ctx context.Context, cmd *exec.Cmd) (*ShellSession, error) { + return startShellSessionContext(ctx, cmd, cmd.Start) } -// Wait 等待 shell 退出。 func (s *ShellSession) Wait() error { if s == nil || s.Cmd == nil { return nil } - return s.Cmd.Wait() + s.waitOnce.Do(func() { + s.waitErr = s.Cmd.Wait() + s.signalMu.Lock() + s.waited = true + terminateProcessGroup(s.rootPID, s.Cmd) + s.signalMu.Unlock() + // Usually the group disappears immediately. Retain ownership if the + // kernel cannot confirm exit; task cleanup will retry and report it. + deadline := time.Now().Add(time.Second) + for !s.tryComplete() && time.Now().Before(deadline) { + time.Sleep(10 * time.Millisecond) + } + }) + return s.waitErr } -// Terminate 终止 shell 及其进程组。 -func (s *ShellSession) Terminate() { +func (s *ShellSession) signal(force bool) { if s == nil { return } - terminateProcessGroup(s.rootPID, s.Cmd) + s.signalMu.Lock() + defer s.signalMu.Unlock() + if s.finished { + return + } + if force { + terminateProcessGroup(s.rootPID, s.Cmd) + } else { + stopProcessGroup(s.rootPID, s.Cmd) + } } -// TerminateShellSession 终止由 StartShellSession 启动的会话。 +func (s *ShellSession) Terminate() { s.signal(true) } + func TerminateShellSession(session *ShellSession) { if session != nil { session.Terminate() } } + +// tryComplete confirms group exit after Wait reaped the direct child. Never +// release ownership merely because a signal was sent successfully. +func (s *ShellSession) tryComplete() bool { + s.signalMu.Lock() + defer s.signalMu.Unlock() + if s.finished { + return true + } + if !s.waited || processGroupExists(s.rootPID) { + return false + } + s.finished = true + if s.scope != nil { + s.scope.mu.Lock() + if s.scope.guard != nil { + if err := s.scope.guard.Release(s.rootPID); err != nil { + s.scope.mu.Unlock() + s.finished = false + return false + } + } + delete(s.scope.sessions, s) + s.scope.mu.Unlock() + } + close(s.done) + return true +} diff --git a/web/static/i18n/en-US.json b/web/static/i18n/en-US.json index 7f44c275..52b1228e 100644 --- a/web/static/i18n/en-US.json +++ b/web/static/i18n/en-US.json @@ -1107,6 +1107,9 @@ "autoRefresh": "Auto refresh", "historyHint": "Tip: Completed task history available. Check \"Show history\" to view.", "statusRunning": "Running", + "statusCleanupUnconfirmed": "Remote cancellation unconfirmed", + "statusCleaning": "Cleaning up", + "statusCleanupFailed": "Cleanup failed (retrying)", "statusCancelling": "Cancelling", "statusFailed": "Failed", "statusTimeout": "Timeout", diff --git a/web/static/i18n/zh-CN.json b/web/static/i18n/zh-CN.json index 9841e46f..311a89e5 100644 --- a/web/static/i18n/zh-CN.json +++ b/web/static/i18n/zh-CN.json @@ -1095,6 +1095,9 @@ "autoRefresh": "自动刷新", "historyHint": "提示：有已完成的任务历史，请勾选\"显示历史记录\"查看", "statusRunning": "执行中", + "statusCleanupUnconfirmed": "远端停止状态待确认", + "statusCleaning": "清理中", + "statusCleanupFailed": "清理异常（自动重试）", "statusCancelling": "取消中", "statusFailed": "执行失败", "statusTimeout": "执行超时", diff --git a/web/static/js/monitor.js b/web/static/js/monitor.js index 7a7cf6fd..e3630b6d 100644 --- a/web/static/js/monitor.js +++ b/web/static/js/monitor.js @@ -3,7 +3,7 @@ const progressTaskState = new Map(); let userInterruptModalPending = null; let activeTaskInterval = null; const ACTIVE_TASK_REFRESH_INTERVAL = 2000; // 运行态与审批态需要及时自刷新 -const TASK_FINAL_STATUSES = new Set(['failed', 'timeout', 'cancelled', 'completed']); +const TASK_FINAL_STATUSES = new Set(['failed', 'timeout', 'cancelled', 'completed', 'cleanup_unconfirmed']); const hitlInterruptToolItemMap = new Map(); let activeTasksLoadPromise = null; let activeTasksVisualSignature = ''; @@ -7153,13 +7153,16 @@ function renderActiveTasks(tasks) { const statusMap = { 'running': _t('tasks.statusRunning'), 'cancelling': _t('tasks.statusCancelling'), + 'cleaning': _t('tasks.statusCleaning'), + 'cleanup_failed': _t('tasks.statusCleanupFailed'), + 'cleanup_unconfirmed': _t('tasks.statusCleanupUnconfirmed'), 'failed': _t('tasks.statusFailed'), 'timeout': _t('tasks.statusTimeout'), 'cancelled': _t('tasks.statusCancelled'), 'completed': _t('tasks.statusCompleted') }; const statusText = statusMap[task.status] || _t('tasks.statusRunning'); - const isFinalStatus = ['failed', 'timeout', 'cancelled', 'completed'].includes(task.status); + const isFinalStatus = ['failed', 'timeout', 'cancelled', 'completed', 'cleanup_unconfirmed'].includes(task.status); const taskDisplayName = getActiveTaskDisplayName(task); const stopTaskBtnText = _t('tasks.stopTask'); diff --git a/web/static/js/tasks.js b/web/static/js/tasks.js index b312eb5c..c77b483a 100644 --- a/web/static/js/tasks.js +++ b/web/static/js/tasks.js @@ -187,7 +187,7 @@ function updateCompletedTasksHistory(currentTasks) { const exists = tasksState.completedTasksHistory.some(t => t.conversationId === task.conversationId); if (!exists) { // 如果任务状态不是最终状态，标记为completed - const finalStatus = ['completed', 'failed', 'timeout', 'cancelled'].includes(task.status) + const finalStatus = ['completed', 'failed', 'timeout', 'cancelled', 'cleanup_unconfirmed'].includes(task.status) ? task.status : 'completed'; @@ -309,7 +309,7 @@ function updateTaskStats(tasks) { tasks.forEach(task => { if (task.status === 'running') { stats.running++; - } else if (task.status === 'cancelling') { + } else if (['cancelling', 'cleaning', 'cleanup_failed'].includes(task.status)) { stats.cancelling++; } else if (task.status === 'completed') { stats.completed++; @@ -364,7 +364,7 @@ function filterAndSortTasks() { if (statusFilter === 'active') { // 仅运行中的任务（不包括历史） filtered = tasksState.allTasks.filter(task => - task.status === 'running' || task.status === 'cancelling' + ['running', 'cancelling', 'cleaning', 'cleanup_failed'].includes(task.status) ); } else if (statusFilter === 'history') { // 仅历史记录 @@ -445,7 +445,7 @@ function updateTaskDurations() { const status = item.dataset.status; const durationEl = item.querySelector('.task-duration'); - if (durationEl && startedAt && (status === 'running' || status === 'cancelling')) { + if (durationEl && startedAt && (['running', 'cancelling', 'cleaning', 'cleanup_failed'].includes(status))) { durationEl.textContent = calculateDuration(startedAt); } }); @@ -470,6 +470,9 @@ function renderTasks(tasks) { // 状态映射 const statusMap = { 'running': { text: _t('tasks.statusRunning'), class: 'task-status-running' }, + 'cleaning': { text: _t('tasks.statusCleaning'), class: 'task-status-cancelling' }, + 'cleanup_unconfirmed': { text: _t('tasks.statusCleanupUnconfirmed'), class: 'task-status-failed' }, + 'cleanup_failed': { text: _t('tasks.statusCleanupFailed'), class: 'task-status-failed' }, 'cancelling': { text: _t('tasks.statusCancelling'), class: 'task-status-cancelling' }, 'failed': { text: _t('tasks.statusFailed'), class: 'task-status-failed' }, 'timeout': { text: _t('tasks.statusTimeout'), class: 'task-status-timeout' }, @@ -530,10 +533,10 @@ function renderTaskItem(task, statusMap, isHistory = false) { : ''; const status = statusMap[task.status] || { text: task.status, class: 'task-status-unknown' }; - const isFinalStatus = ['failed', 'timeout', 'cancelled', 'completed'].includes(task.status); - const canCancel = !isFinalStatus && task.status !== 'cancelling' && !isHistory; + const isFinalStatus = ['failed', 'timeout', 'cancelled', 'completed', 'cleanup_unconfirmed'].includes(task.status); + const canCancel = !isFinalStatus && !['cancelling', 'cleaning'].includes(task.status) && !isHistory; const isSelected = tasksState.selectedTasks.has(task.conversationId); - const duration = (task.status === 'running' || task.status === 'cancelling') + const duration = (['running', 'cancelling', 'cleaning', 'cleanup_failed'].includes(task.status)) ? calculateDuration(task.startedAt) : ''; @@ -561,6 +564,7 @@ function renderTaskItem(task, statusMap, isHistory = false) { ${task.conversationId ? `` : ''} + ${task.cleanupError ? `
${escapeHtml(task.cleanupError)}
` : ''} ${task.conversationId ? `
` + _t('tasks.conversationIdLabel') + `: