mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-09-17 23:22:27 +02:00
feat: manage task process lifetimes and preserve turn history
This commit is contained in:
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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 "任务资源清理失败,将自动重试"
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
})
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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())) })
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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())
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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') &`,
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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 }
|
||||
|
||||
@@ -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"
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,47 +1,102 @@
|
||||
package security
|
||||
|
||||
import "os/exec"
|
||||
import (
|
||||
"context"
|
||||
"os/exec"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// ShellSession 在 Start 时记录根 shell 的进程组 ID,取消/超时时可杀整组(即使 cmd.Process 已失效)。
|
||||
// ShellSession caches the process group while its command is alive. Signals
|
||||
// and Wait completion synchronize to avoid signalling already-released sessions.
|
||||
type ShellSession struct {
|
||||
Cmd *exec.Cmd
|
||||
rootPID int
|
||||
Cmd *exec.Cmd
|
||||
rootPID int
|
||||
scope *ProcessScope
|
||||
done chan struct{}
|
||||
waitOnce sync.Once
|
||||
waitErr error
|
||||
signalMu sync.Mutex
|
||||
finished bool
|
||||
waited bool
|
||||
}
|
||||
|
||||
// StartShellSession 配置独立进程组并启动 shell,缓存 rootPID(Unix 下即 PGID)。
|
||||
func StartShellSession(cmd *exec.Cmd) (*ShellSession, error) {
|
||||
if err := prepareShellCmdSession(cmd); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := cmd.Start(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
pid := 0
|
||||
if cmd.Process != nil {
|
||||
pid = cmd.Process.Pid
|
||||
}
|
||||
return &ShellSession{Cmd: cmd, rootPID: pid}, nil
|
||||
return StartShellSessionContext(context.Background(), cmd)
|
||||
}
|
||||
|
||||
func StartShellSessionContext(ctx context.Context, cmd *exec.Cmd) (*ShellSession, error) {
|
||||
return startShellSessionContext(ctx, cmd, cmd.Start)
|
||||
}
|
||||
|
||||
// Wait 等待 shell 退出。
|
||||
func (s *ShellSession) Wait() error {
|
||||
if s == nil || s.Cmd == nil {
|
||||
return nil
|
||||
}
|
||||
return s.Cmd.Wait()
|
||||
s.waitOnce.Do(func() {
|
||||
s.waitErr = s.Cmd.Wait()
|
||||
s.signalMu.Lock()
|
||||
s.waited = true
|
||||
terminateProcessGroup(s.rootPID, s.Cmd)
|
||||
s.signalMu.Unlock()
|
||||
// Usually the group disappears immediately. Retain ownership if the
|
||||
// kernel cannot confirm exit; task cleanup will retry and report it.
|
||||
deadline := time.Now().Add(time.Second)
|
||||
for !s.tryComplete() && time.Now().Before(deadline) {
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
})
|
||||
return s.waitErr
|
||||
}
|
||||
|
||||
// Terminate 终止 shell 及其进程组。
|
||||
func (s *ShellSession) Terminate() {
|
||||
func (s *ShellSession) signal(force bool) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
terminateProcessGroup(s.rootPID, s.Cmd)
|
||||
s.signalMu.Lock()
|
||||
defer s.signalMu.Unlock()
|
||||
if s.finished {
|
||||
return
|
||||
}
|
||||
if force {
|
||||
terminateProcessGroup(s.rootPID, s.Cmd)
|
||||
} else {
|
||||
stopProcessGroup(s.rootPID, s.Cmd)
|
||||
}
|
||||
}
|
||||
|
||||
// TerminateShellSession 终止由 StartShellSession 启动的会话。
|
||||
func (s *ShellSession) Terminate() { s.signal(true) }
|
||||
|
||||
func TerminateShellSession(session *ShellSession) {
|
||||
if session != nil {
|
||||
session.Terminate()
|
||||
}
|
||||
}
|
||||
|
||||
// tryComplete confirms group exit after Wait reaped the direct child. Never
|
||||
// release ownership merely because a signal was sent successfully.
|
||||
func (s *ShellSession) tryComplete() bool {
|
||||
s.signalMu.Lock()
|
||||
defer s.signalMu.Unlock()
|
||||
if s.finished {
|
||||
return true
|
||||
}
|
||||
if !s.waited || processGroupExists(s.rootPID) {
|
||||
return false
|
||||
}
|
||||
s.finished = true
|
||||
if s.scope != nil {
|
||||
s.scope.mu.Lock()
|
||||
if s.scope.guard != nil {
|
||||
if err := s.scope.guard.Release(s.rootPID); err != nil {
|
||||
s.scope.mu.Unlock()
|
||||
s.finished = false
|
||||
return false
|
||||
}
|
||||
}
|
||||
delete(s.scope.sessions, s)
|
||||
s.scope.mu.Unlock()
|
||||
}
|
||||
close(s.done)
|
||||
return true
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user