feat: manage task process lifetimes and preserve turn history

This commit is contained in:
Ed1s0nZ
2026-09-16 17:52:58 +08:00
parent fd1c13a43d
commit f7882be546
54 changed files with 3650 additions and 330 deletions
+3
View File
@@ -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()
+10
View File
@@ -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
+23 -8
View File
@@ -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)
+11 -8
View File
@@ -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
+30 -6
View File
@@ -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,
+29 -5
View File
@@ -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,
+62
View File
@@ -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 "任务资源清理失败,将自动重试"
}
+218 -50
View File
@@ -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,
@@ -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)
}
}
+19 -7
View File
@@ -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,
+76
View File
@@ -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)
}
}
})
}
}
+55 -17
View File
@@ -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()
+14
View File
@@ -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
}
@@ -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
}
+196
View File
@@ -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
}
}
@@ -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)
}
}
@@ -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)
}
+10 -5
View File
@@ -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)})
@@ -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)
}
+120
View File
@@ -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 >/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)
}
}
+52
View File
@@ -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
}
+180
View File
@@ -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
})
}
+89
View File
@@ -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)
}
+160
View File
@@ -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")
}
}
+97
View File
@@ -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")
}
}
+313
View File
@@ -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)
})
}
+31
View File
@@ -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)
}
+183
View File
@@ -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())) })
}
+182
View File
@@ -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
}
}
}
+126
View File
@@ -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())
}
}
}
+70
View File
@@ -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)
}
}
+21 -132
View File
@@ -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
}
+7 -1
View File
@@ -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') &`,
+11
View File
@@ -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
}
+14 -1
View File
@@ -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 }
+220
View File
@@ -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"
}
+198
View File
@@ -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 >/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)
}
+18 -37
View File
@@ -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
}
}
+77 -22
View File
@@ -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,缓存 rootPIDUnix 下即 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
}