mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-08-10 13:10:34 +02:00
Add files via upload
This commit is contained in:
@@ -110,6 +110,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
|||||||
|
|
||||||
// 创建安全工具执行器
|
// 创建安全工具执行器
|
||||||
executor := security.NewExecutor(&cfg.Security, mcpServer, log.Logger)
|
executor := security.NewExecutor(&cfg.Security, mcpServer, log.Logger)
|
||||||
|
executor.SetShellNoOutputTimeoutSeconds(cfg.Agent.ShellNoOutputTimeoutSeconds)
|
||||||
|
|
||||||
// 注册工具
|
// 注册工具
|
||||||
executor.RegisterTools(mcpServer)
|
executor.RegisterTools(mcpServer)
|
||||||
@@ -333,6 +334,8 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
|||||||
monitorHandler.SetAudit(auditSvc)
|
monitorHandler.SetAudit(auditSvc)
|
||||||
monitorHandler.SetMonitorRetention(monitorRetention)
|
monitorHandler.SetMonitorRetention(monitorRetention)
|
||||||
monitorHandler.SetExternalMCPManager(externalMCPMgr) // 设置外部MCP管理器,以便获取外部MCP执行记录
|
monitorHandler.SetExternalMCPManager(externalMCPMgr) // 设置外部MCP管理器,以便获取外部MCP执行记录
|
||||||
|
monitorHandler.SetTaskManager(agentHandler.TaskManager())
|
||||||
|
monitorHandler.SetAgentHandler(agentHandler)
|
||||||
notificationHandler := handler.NewNotificationHandler(db, agentHandler, log.Logger)
|
notificationHandler := handler.NewNotificationHandler(db, agentHandler, log.Logger)
|
||||||
groupHandler := handler.NewGroupHandler(db, log.Logger)
|
groupHandler := handler.NewGroupHandler(db, log.Logger)
|
||||||
authHandler := handler.NewAuthHandler(authManager, cfg, configPath, log.Logger)
|
authHandler := handler.NewAuthHandler(authManager, cfg, configPath, log.Logger)
|
||||||
|
|||||||
@@ -196,6 +196,16 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
|||||||
pendingByID[tc.ToolCallID] = tc
|
pendingByID[tc.ToolCallID] = tc
|
||||||
pendingQueueByAgent[tc.EinoAgent] = append(pendingQueueByAgent[tc.EinoAgent], tc.ToolCallID)
|
pendingQueueByAgent[tc.EinoAgent] = append(pendingQueueByAgent[tc.EinoAgent], tc.ToolCallID)
|
||||||
}
|
}
|
||||||
|
markPendingWithMonitor := func(tc toolCallPendingInfo) {
|
||||||
|
markPending(tc)
|
||||||
|
beginEinoADKFilesystemToolMonitor(
|
||||||
|
args.FilesystemMonitorAgent,
|
||||||
|
args.FilesystemMonitorRecord,
|
||||||
|
args.MCPExecutionBinder,
|
||||||
|
tc.ToolCallID,
|
||||||
|
tc.ToolName,
|
||||||
|
)
|
||||||
|
}
|
||||||
popNextPendingForAgent := func(agentName string) (toolCallPendingInfo, bool) {
|
popNextPendingForAgent := func(agentName string) (toolCallPendingInfo, bool) {
|
||||||
pendingMu.Lock()
|
pendingMu.Lock()
|
||||||
defer pendingMu.Unlock()
|
defer pendingMu.Unlock()
|
||||||
@@ -331,7 +341,7 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
|||||||
toolCallID = tid
|
toolCallID = tid
|
||||||
}
|
}
|
||||||
recordPendingExecuteStdoutDup(toolName, content, isErr)
|
recordPendingExecuteStdoutDup(toolName, content, isErr)
|
||||||
recordEinoADKFilesystemToolMonitor(args.FilesystemMonitorAgent, args.FilesystemMonitorRecord, toolName, toolCallID, runAccumulatedMsgs, content, isErr)
|
recordEinoADKFilesystemToolMonitor(args.FilesystemMonitorAgent, args.FilesystemMonitorRecord, args.MCPExecutionBinder, toolName, toolCallID, runAccumulatedMsgs, content, isErr)
|
||||||
if args.FilesystemMonitorAgent != nil && args.MCPExecutionBinder != nil {
|
if args.FilesystemMonitorAgent != nil && args.MCPExecutionBinder != nil {
|
||||||
if execID := args.MCPExecutionBinder.ExecutionID(toolCallID); execID != "" {
|
if execID := args.MCPExecutionBinder.ExecutionID(toolCallID); execID != "" {
|
||||||
args.FilesystemMonitorAgent.UpdateMCPExecutionDisplayResult(execID, content)
|
args.FilesystemMonitorAgent.UpdateMCPExecutionDisplayResult(execID, content)
|
||||||
@@ -974,7 +984,7 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
|||||||
if merged := mergeStreamingToolCallFragments(toolStreamFragments); len(merged) > 0 {
|
if merged := mergeStreamingToolCallFragments(toolStreamFragments); len(merged) > 0 {
|
||||||
lastToolChunk = mergeMessageToolCalls(&schema.Message{ToolCalls: merged})
|
lastToolChunk = mergeMessageToolCalls(&schema.Message{ToolCalls: merged})
|
||||||
}
|
}
|
||||||
tryEmitToolCallsOnce(lastToolChunk, ev.AgentName, orchestratorName, conversationID, orchMode, progress, toolEmitSeen, subAgentToolStep, mainAgentToolStep, markPending)
|
tryEmitToolCallsOnce(lastToolChunk, ev.AgentName, orchestratorName, conversationID, orchMode, progress, toolEmitSeen, subAgentToolStep, mainAgentToolStep, markPendingWithMonitor)
|
||||||
// 流式路径此前只把 tool_calls 推给进度 UI,未写入 runAccumulatedMsgs;落库后 loadHistory→RepairOrphan 会删掉全部 tool 结果,表现为「续跑/下轮失忆」。
|
// 流式路径此前只把 tool_calls 推给进度 UI,未写入 runAccumulatedMsgs;落库后 loadHistory→RepairOrphan 会删掉全部 tool 结果,表现为「续跑/下轮失忆」。
|
||||||
if lastToolChunk != nil && len(lastToolChunk.ToolCalls) > 0 {
|
if lastToolChunk != nil && len(lastToolChunk.ToolCalls) > 0 {
|
||||||
runAccumulatedMsgs = append(runAccumulatedMsgs, schema.AssistantMessage("", lastToolChunk.ToolCalls))
|
runAccumulatedMsgs = append(runAccumulatedMsgs, schema.AssistantMessage("", lastToolChunk.ToolCalls))
|
||||||
@@ -1009,7 +1019,7 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
|||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
runAccumulatedMsgs = append(runAccumulatedMsgs, msg)
|
runAccumulatedMsgs = append(runAccumulatedMsgs, msg)
|
||||||
tryEmitToolCallsOnce(mergeMessageToolCalls(msg), ev.AgentName, orchestratorName, conversationID, orchMode, progress, toolEmitSeen, subAgentToolStep, mainAgentToolStep, markPending)
|
tryEmitToolCallsOnce(mergeMessageToolCalls(msg), ev.AgentName, orchestratorName, conversationID, orchMode, progress, toolEmitSeen, subAgentToolStep, mainAgentToolStep, markPendingWithMonitor)
|
||||||
|
|
||||||
if mv.Role == schema.Assistant {
|
if mv.Role == schema.Assistant {
|
||||||
if progress != nil && strings.TrimSpace(msg.ReasoningContent) != "" {
|
if progress != nil && strings.TrimSpace(msg.ReasoningContent) != "" {
|
||||||
|
|||||||
@@ -7,11 +7,25 @@ import (
|
|||||||
"cyberstrike-ai/internal/einomcp"
|
"cyberstrike-ai/internal/einomcp"
|
||||||
)
|
)
|
||||||
|
|
||||||
// newEinoExecuteMonitorCallback 在 Eino filesystem execute 结束时写入 MCP 监控库并 recorder(executionId),
|
// newEinoExecuteMonitorCallbacks 在 Eino filesystem execute 开始/结束时写入 MCP 监控库并 recorder(executionId),
|
||||||
// 与 CallTool 路径一致,供助手消息展示「渗透测试详情」芯片。
|
// 与 CallTool 路径一致,使监控页能展示「执行中」状态。
|
||||||
func newEinoExecuteMonitorCallback(ag *agent.Agent, recorder einomcp.ExecutionRecorder) func(toolCallID, command, stdout string, success bool, invokeErr error) {
|
func newEinoExecuteMonitorCallbacks(ag *agent.Agent, recorder einomcp.ExecutionRecorder) (
|
||||||
return func(toolCallID, command, stdout string, success bool, invokeErr error) {
|
begin func(toolCallID, command string) string,
|
||||||
if ag == nil || recorder == nil {
|
finish func(executionID, toolCallID, command, stdout string, success bool, invokeErr error),
|
||||||
|
) {
|
||||||
|
begin = func(toolCallID, command string) string {
|
||||||
|
if ag == nil {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
args := map[string]interface{}{"command": command}
|
||||||
|
id := ag.BeginLocalToolExecution("execute", args)
|
||||||
|
if id != "" && recorder != nil {
|
||||||
|
recorder(id, toolCallID)
|
||||||
|
}
|
||||||
|
return id
|
||||||
|
}
|
||||||
|
finish = func(executionID, toolCallID, command, stdout string, success bool, invokeErr error) {
|
||||||
|
if ag == nil {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
var err error
|
var err error
|
||||||
@@ -23,9 +37,10 @@ func newEinoExecuteMonitorCallback(ag *agent.Agent, recorder einomcp.ExecutionRe
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
args := map[string]interface{}{"command": command}
|
args := map[string]interface{}{"command": command}
|
||||||
id := ag.RecordLocalToolExecution("execute", args, stdout, err)
|
id := ag.FinishLocalToolExecution(executionID, "execute", args, stdout, err)
|
||||||
if id != "" {
|
if id != "" && recorder != nil && executionID == "" {
|
||||||
recorder(id, toolCallID)
|
recorder(id, toolCallID)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
return begin, finish
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -63,8 +63,11 @@ type einoStreamingShellWrap struct {
|
|||||||
outputChunk func(toolName, toolCallID, chunk string)
|
outputChunk func(toolName, toolCallID, chunk string)
|
||||||
// toolTimeoutMinutes 与 agent.tool_timeout_minutes 对齐;>0 时对单次 execute 套用 context 超时(与 MCP 工具经 executeToolViaMCP 行为一致)。0 表示仅依赖上层 ctx(如整任务 10h 上限)。
|
// toolTimeoutMinutes 与 agent.tool_timeout_minutes 对齐;>0 时对单次 execute 套用 context 超时(与 MCP 工具经 executeToolViaMCP 行为一致)。0 表示仅依赖上层 ctx(如整任务 10h 上限)。
|
||||||
toolTimeoutMinutes int
|
toolTimeoutMinutes int
|
||||||
// recordMonitor 在 execute 流结束后写入 tool_executions 并 recorder(executionId),使「渗透测试详情」与常规 MCP 一致。
|
// shellNoOutputTimeoutSec:无任何输出时的空闲秒数;0=关闭。
|
||||||
recordMonitor func(toolCallID, command, stdout string, success bool, invokeErr error)
|
shellNoOutputTimeoutSec int
|
||||||
|
// beginMonitor 在 execute 开始时写入 running 状态;finishMonitor 在流结束后更新为 completed/failed。
|
||||||
|
beginMonitor func(toolCallID, command string) string
|
||||||
|
finishMonitor func(executionID, toolCallID, command, stdout string, success bool, invokeErr error)
|
||||||
}
|
}
|
||||||
|
|
||||||
func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) {
|
func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) {
|
||||||
@@ -76,15 +79,26 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
|||||||
}
|
}
|
||||||
req := *input
|
req := *input
|
||||||
userCmd := strings.TrimSpace(req.Command)
|
userCmd := strings.TrimSpace(req.Command)
|
||||||
|
tid := strings.TrimSpace(compose.GetToolCallID(ctx))
|
||||||
|
agentTag := strings.TrimSpace(w.einoAgentName)
|
||||||
if security.IsBackgroundShellCommand(req.Command) && !req.RunInBackendGround {
|
if security.IsBackgroundShellCommand(req.Command) && !req.RunInBackendGround {
|
||||||
req.RunInBackendGround = true
|
req.RunInBackendGround = true
|
||||||
}
|
}
|
||||||
req.Command = prependPythonUnbufferedEnv(req.Command)
|
req.Command = security.PrepareNonInteractiveShellCommand(prependPythonUnbufferedEnv(req.Command))
|
||||||
tid := strings.TrimSpace(compose.GetToolCallID(ctx))
|
|
||||||
agentTag := strings.TrimSpace(w.einoAgentName)
|
|
||||||
convID := mcp.MCPConversationIDFromContext(ctx)
|
convID := mcp.MCPConversationIDFromContext(ctx)
|
||||||
execReg := mcp.EinoExecuteRunRegistryFromContext(ctx)
|
execReg := mcp.EinoExecuteRunRegistryFromContext(ctx)
|
||||||
|
|
||||||
|
var monitorExecID string
|
||||||
|
if w.beginMonitor != nil {
|
||||||
|
monitorExecID = w.beginMonitor(tid, userCmd)
|
||||||
|
}
|
||||||
|
if monitorExecID != "" && convID != "" {
|
||||||
|
if toolReg := mcp.ToolRunRegistryFromContext(ctx); toolReg != nil {
|
||||||
|
toolReg.RegisterRunningTool(convID, monitorExecID)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
toolRunReg := mcp.ToolRunRegistryFromContext(ctx)
|
||||||
|
|
||||||
execCtx, execCancel := context.WithCancel(ctx)
|
execCtx, execCancel := context.WithCancel(ctx)
|
||||||
var timeoutCancel context.CancelFunc
|
var timeoutCancel context.CancelFunc
|
||||||
if w.toolTimeoutMinutes > 0 {
|
if w.toolTimeoutMinutes > 0 {
|
||||||
@@ -104,23 +118,23 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
|||||||
}
|
}
|
||||||
if einoExecuteRecvErrIsToolTimeout(err, execCtx) {
|
if einoExecuteRecvErrIsToolTimeout(err, execCtx) {
|
||||||
hint := "\n\n" + einoExecuteTimeoutUserHint() + "\n"
|
hint := "\n\n" + einoExecuteTimeoutUserHint() + "\n"
|
||||||
if w.recordMonitor != nil {
|
if w.finishMonitor != nil {
|
||||||
w.recordMonitor(tid, userCmd, hint, false, context.DeadlineExceeded)
|
w.finishMonitor(monitorExecID, tid, userCmd, hint, false, context.DeadlineExceeded)
|
||||||
}
|
}
|
||||||
if w.invokeNotify != nil && tid != "" {
|
if w.invokeNotify != nil && tid != "" {
|
||||||
w.invokeNotify.Fire(tid, "execute", agentTag, false, hint, context.DeadlineExceeded)
|
w.invokeNotify.Fire(tid, "execute", agentTag, false, hint, context.DeadlineExceeded)
|
||||||
}
|
}
|
||||||
return schema.StreamReaderFromArray([]*filesystem.ExecuteResponse{{Output: hint}}), nil
|
return schema.StreamReaderFromArray([]*filesystem.ExecuteResponse{{Output: hint}}), nil
|
||||||
}
|
}
|
||||||
if w.recordMonitor != nil {
|
if w.finishMonitor != nil {
|
||||||
w.recordMonitor(tid, userCmd, "", false, err)
|
w.finishMonitor(monitorExecID, tid, userCmd, "", false, err)
|
||||||
}
|
}
|
||||||
if w.invokeNotify != nil && tid != "" {
|
if w.invokeNotify != nil && tid != "" {
|
||||||
w.invokeNotify.Fire(tid, "execute", agentTag, false, "", err)
|
w.invokeNotify.Fire(tid, "execute", agentTag, false, "", err)
|
||||||
}
|
}
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
if sr == nil || w.invokeNotify == nil {
|
if sr == nil {
|
||||||
if timeoutCancel != nil {
|
if timeoutCancel != nil {
|
||||||
timeoutCancel()
|
timeoutCancel()
|
||||||
}
|
}
|
||||||
@@ -132,7 +146,7 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
|||||||
|
|
||||||
outR, outW := schema.Pipe[*filesystem.ExecuteResponse](32)
|
outR, outW := schema.Pipe[*filesystem.ExecuteResponse](32)
|
||||||
|
|
||||||
go func(inner *schema.StreamReader[*filesystem.ExecuteResponse], command string, cancel context.CancelFunc, timeoutCleanup context.CancelFunc, tctx context.Context, conversationID string, reg mcp.EinoExecuteRunRegistry) {
|
go func(inner *schema.StreamReader[*filesystem.ExecuteResponse], command string, cancel context.CancelFunc, timeoutCleanup context.CancelFunc, tctx context.Context, conversationID string, reg mcp.EinoExecuteRunRegistry, toolReg mcp.ToolRunRegistry, execID string, toolCallID string, noOutputSec int) {
|
||||||
var innerCloseOnce sync.Once
|
var innerCloseOnce sync.Once
|
||||||
closeInner := func() {
|
closeInner := func() {
|
||||||
innerCloseOnce.Do(func() { inner.Close() })
|
innerCloseOnce.Do(func() { inner.Close() })
|
||||||
@@ -147,6 +161,9 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
|||||||
if reg != nil && conversationID != "" {
|
if reg != nil && conversationID != "" {
|
||||||
defer reg.UnregisterActiveEinoExecute(conversationID)
|
defer reg.UnregisterActiveEinoExecute(conversationID)
|
||||||
}
|
}
|
||||||
|
if toolReg != nil && conversationID != "" && execID != "" {
|
||||||
|
defer toolReg.UnregisterRunningTool(conversationID, execID)
|
||||||
|
}
|
||||||
|
|
||||||
// ctx 取消时关闭内层流,避免 amass 等长时间无换行输出时 Recv 永久阻塞。
|
// ctx 取消时关闭内层流,避免 amass 等长时间无换行输出时 Recv 永久阻塞。
|
||||||
stopWatch := make(chan struct{})
|
stopWatch := make(chan struct{})
|
||||||
@@ -165,43 +182,92 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
|||||||
exitCode := 0
|
exitCode := 0
|
||||||
hasExitCode := false
|
hasExitCode := false
|
||||||
|
|
||||||
|
idleWatch := security.NewShellInactivityWatch(noOutputSec)
|
||||||
|
if idleWatch != nil {
|
||||||
|
defer idleWatch.Stop()
|
||||||
|
}
|
||||||
|
|
||||||
|
type execRecvMsg struct {
|
||||||
|
resp *filesystem.ExecuteResponse
|
||||||
|
err error
|
||||||
|
}
|
||||||
|
recvCh := make(chan execRecvMsg, 1)
|
||||||
|
go func() {
|
||||||
|
for {
|
||||||
|
resp, rerr := inner.Recv()
|
||||||
|
recvCh <- execRecvMsg{resp: resp, err: rerr}
|
||||||
|
if rerr != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
|
||||||
|
fireInactivityTimeout := func() {
|
||||||
|
success = false
|
||||||
|
invokeErr = fmt.Errorf("shell inactivity timeout (%ds)", idleWatch.Sec)
|
||||||
|
msg := security.ShellNoOutputTimeoutMessage(idleWatch.Sec)
|
||||||
|
_ = outW.Send(&filesystem.ExecuteResponse{Output: msg}, nil)
|
||||||
|
sb.WriteString(msg)
|
||||||
|
if w.outputChunk != nil && toolCallID != "" {
|
||||||
|
w.outputChunk("execute", toolCallID, msg)
|
||||||
|
}
|
||||||
|
if cancel != nil {
|
||||||
|
cancel()
|
||||||
|
}
|
||||||
|
closeInner()
|
||||||
|
}
|
||||||
|
|
||||||
|
recvLoop:
|
||||||
for {
|
for {
|
||||||
resp, rerr := inner.Recv()
|
var idleCh <-chan struct{}
|
||||||
if errors.Is(rerr, io.EOF) {
|
if idleWatch != nil {
|
||||||
break
|
idleCh = idleWatch.Expired
|
||||||
}
|
}
|
||||||
if rerr != nil {
|
select {
|
||||||
success = false
|
case <-idleCh:
|
||||||
invokeErr = rerr
|
fireInactivityTimeout()
|
||||||
// 单次 execute 超时须与 MCP 工具一致:写入工具结果尾标、继续迭代,不得向 ADK 流注入硬错误。
|
break recvLoop
|
||||||
if einoExecuteRecvErrIsToolTimeout(rerr, tctx) {
|
case msg := <-recvCh:
|
||||||
invokeErr = context.DeadlineExceeded
|
rerr := msg.err
|
||||||
break
|
resp := msg.resp
|
||||||
|
if errors.Is(rerr, io.EOF) {
|
||||||
|
break recvLoop
|
||||||
}
|
}
|
||||||
if errors.Is(rerr, context.Canceled) || (tctx != nil && errors.Is(tctx.Err(), context.Canceled)) {
|
if rerr != nil {
|
||||||
invokeErr = context.Canceled
|
|
||||||
break
|
|
||||||
}
|
|
||||||
_ = outW.Send(nil, rerr)
|
|
||||||
break
|
|
||||||
}
|
|
||||||
if resp != nil {
|
|
||||||
if resp.ExitCode != nil {
|
|
||||||
hasExitCode = true
|
|
||||||
exitCode = *resp.ExitCode
|
|
||||||
}
|
|
||||||
var appended string
|
|
||||||
if resp.Output != "" {
|
|
||||||
sb.WriteString(resp.Output)
|
|
||||||
appended = resp.Output
|
|
||||||
}
|
|
||||||
if w.outputChunk != nil && strings.TrimSpace(appended) != "" {
|
|
||||||
w.outputChunk("execute", tid, appended)
|
|
||||||
}
|
|
||||||
if outW.Send(resp, nil) {
|
|
||||||
success = false
|
success = false
|
||||||
invokeErr = fmt.Errorf("execute stream closed by consumer")
|
invokeErr = rerr
|
||||||
break
|
if einoExecuteRecvErrIsToolTimeout(rerr, tctx) {
|
||||||
|
invokeErr = context.DeadlineExceeded
|
||||||
|
break recvLoop
|
||||||
|
}
|
||||||
|
if errors.Is(rerr, context.Canceled) || (tctx != nil && errors.Is(tctx.Err(), context.Canceled)) {
|
||||||
|
invokeErr = context.Canceled
|
||||||
|
break recvLoop
|
||||||
|
}
|
||||||
|
_ = outW.Send(nil, rerr)
|
||||||
|
break recvLoop
|
||||||
|
}
|
||||||
|
if resp != nil {
|
||||||
|
if resp.ExitCode != nil {
|
||||||
|
hasExitCode = true
|
||||||
|
exitCode = *resp.ExitCode
|
||||||
|
}
|
||||||
|
var appended string
|
||||||
|
if resp.Output != "" {
|
||||||
|
if idleWatch != nil {
|
||||||
|
idleWatch.Bump()
|
||||||
|
}
|
||||||
|
sb.WriteString(resp.Output)
|
||||||
|
appended = resp.Output
|
||||||
|
}
|
||||||
|
if w.outputChunk != nil && strings.TrimSpace(appended) != "" {
|
||||||
|
w.outputChunk("execute", toolCallID, appended)
|
||||||
|
}
|
||||||
|
if outW.Send(resp, nil) {
|
||||||
|
success = false
|
||||||
|
invokeErr = fmt.Errorf("execute stream closed by consumer")
|
||||||
|
break recvLoop
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -248,12 +314,14 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
|||||||
_ = outW.Send(&filesystem.ExecuteResponse{Output: text + "\n"}, nil)
|
_ = outW.Send(&filesystem.ExecuteResponse{Output: text + "\n"}, nil)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if w.recordMonitor != nil {
|
if w.finishMonitor != nil {
|
||||||
w.recordMonitor(tid, command, sb.String(), success, invokeErr)
|
w.finishMonitor(execID, toolCallID, command, sb.String(), success, invokeErr)
|
||||||
|
}
|
||||||
|
if w.invokeNotify != nil {
|
||||||
|
w.invokeNotify.Fire(toolCallID, "execute", agentTag, success, sb.String(), invokeErr)
|
||||||
}
|
}
|
||||||
w.invokeNotify.Fire(tid, "execute", agentTag, success, sb.String(), invokeErr)
|
|
||||||
outW.Close()
|
outW.Close()
|
||||||
}(sr, userCmd, execCancel, timeoutCancel, execCtx, convID, execReg)
|
}(sr, userCmd, execCancel, timeoutCancel, execCtx, convID, execReg, toolRunReg, monitorExecID, tid, w.shellNoOutputTimeoutSec)
|
||||||
|
|
||||||
return outR, nil
|
return outR, nil
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -19,9 +19,15 @@ type mockStreamingShell struct {
|
|||||||
immediateErr error
|
immediateErr error
|
||||||
recvErr error
|
recvErr error
|
||||||
output string
|
output string
|
||||||
|
called bool
|
||||||
|
lastCommand string
|
||||||
}
|
}
|
||||||
|
|
||||||
func (m *mockStreamingShell) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) {
|
func (m *mockStreamingShell) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) {
|
||||||
|
m.called = true
|
||||||
|
if input != nil {
|
||||||
|
m.lastCommand = input.Command
|
||||||
|
}
|
||||||
if m.immediateErr != nil {
|
if m.immediateErr != nil {
|
||||||
return nil, m.immediateErr
|
return nil, m.immediateErr
|
||||||
}
|
}
|
||||||
@@ -38,6 +44,135 @@ func (m *mockStreamingShell) ExecuteStreaming(ctx context.Context, input *filesy
|
|||||||
return outR, nil
|
return outR, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestEinoStreamingShellWrap_PreparesNonInteractiveCommand(t *testing.T) {
|
||||||
|
inner := &mockStreamingShell{output: "ok\n"}
|
||||||
|
wrap := &einoStreamingShellWrap{inner: inner}
|
||||||
|
sr, err := wrap.ExecuteStreaming(context.Background(), &filesystem.ExecuteRequest{Command: "echo ok"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ExecuteStreaming: %v", err)
|
||||||
|
}
|
||||||
|
defer sr.Close()
|
||||||
|
for {
|
||||||
|
_, rerr := sr.Recv()
|
||||||
|
if errors.Is(rerr, io.EOF) {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
if rerr != nil {
|
||||||
|
t.Fatalf("recv: %v", rerr)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !strings.Contains(inner.lastCommand, "exec </dev/null") {
|
||||||
|
t.Fatalf("missing stdin redirect in inner command: %q", inner.lastCommand)
|
||||||
|
}
|
||||||
|
if !strings.Contains(inner.lastCommand, "GIT_PAGER=cat") {
|
||||||
|
t.Fatalf("missing pager export in inner command: %q", inner.lastCommand)
|
||||||
|
}
|
||||||
|
if !strings.Contains(inner.lastCommand, "PYTHONUNBUFFERED=1") {
|
||||||
|
t.Fatalf("missing python unbuffer in inner command: %q", inner.lastCommand)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestEinoStreamingShellWrap_NoOutputTimeout(t *testing.T) {
|
||||||
|
inner := &mockStreamingShellHanging{}
|
||||||
|
notify := einomcp.NewToolInvokeNotifyHolder()
|
||||||
|
var fired string
|
||||||
|
notify.Set(func(toolCallID, toolName, einoAgent string, success bool, content string, invokeErr error) {
|
||||||
|
fired = content
|
||||||
|
})
|
||||||
|
wrap := &einoStreamingShellWrap{
|
||||||
|
inner: inner,
|
||||||
|
invokeNotify: notify,
|
||||||
|
shellNoOutputTimeoutSec: 1,
|
||||||
|
}
|
||||||
|
sr, err := wrap.ExecuteStreaming(context.Background(), &filesystem.ExecuteRequest{Command: "sudo whoami"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ExecuteStreaming: %v", err)
|
||||||
|
}
|
||||||
|
defer sr.Close()
|
||||||
|
var got strings.Builder
|
||||||
|
for {
|
||||||
|
resp, rerr := sr.Recv()
|
||||||
|
if errors.Is(rerr, io.EOF) {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
if rerr != nil {
|
||||||
|
t.Fatalf("recv: %v", rerr)
|
||||||
|
}
|
||||||
|
if resp != nil {
|
||||||
|
got.WriteString(resp.Output)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !inner.called {
|
||||||
|
t.Fatal("inner shell should run (no command blacklist)")
|
||||||
|
}
|
||||||
|
out := got.String()
|
||||||
|
if !strings.Contains(out, "没有新的输出") && !strings.Contains(out, "no new output") {
|
||||||
|
t.Fatalf("expected inactivity timeout message, got: %q notify=%q", out, fired)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type mockStreamingShellPartialThenHang struct {
|
||||||
|
called bool
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *mockStreamingShellPartialThenHang) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) {
|
||||||
|
m.called = true
|
||||||
|
outR, outW := schema.Pipe[*filesystem.ExecuteResponse](4)
|
||||||
|
go func() {
|
||||||
|
_ = outW.Send(&filesystem.ExecuteResponse{Output: "[sudo] password:\n"}, nil)
|
||||||
|
<-ctx.Done()
|
||||||
|
outW.Close()
|
||||||
|
}()
|
||||||
|
return outR, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestEinoStreamingShellWrap_InactivityAfterPartialOutput(t *testing.T) {
|
||||||
|
inner := &mockStreamingShellPartialThenHang{}
|
||||||
|
wrap := &einoStreamingShellWrap{
|
||||||
|
inner: inner,
|
||||||
|
shellNoOutputTimeoutSec: 1,
|
||||||
|
}
|
||||||
|
start := time.Now()
|
||||||
|
sr, err := wrap.ExecuteStreaming(context.Background(), &filesystem.ExecuteRequest{Command: "sudo whoami"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ExecuteStreaming: %v", err)
|
||||||
|
}
|
||||||
|
defer sr.Close()
|
||||||
|
var got strings.Builder
|
||||||
|
for {
|
||||||
|
resp, rerr := sr.Recv()
|
||||||
|
if errors.Is(rerr, io.EOF) {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
if rerr != nil {
|
||||||
|
t.Fatalf("recv: %v", rerr)
|
||||||
|
}
|
||||||
|
if resp != nil {
|
||||||
|
got.WriteString(resp.Output)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if time.Since(start) > 5*time.Second {
|
||||||
|
t.Fatalf("expected inactivity timeout ~1s, took %v", time.Since(start))
|
||||||
|
}
|
||||||
|
if !strings.Contains(got.String(), "没有新的输出") && !strings.Contains(got.String(), "no new output") {
|
||||||
|
t.Fatalf("expected inactivity message, got: %q", got.String())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type mockStreamingShellHanging struct {
|
||||||
|
called bool
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *mockStreamingShellHanging) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) {
|
||||||
|
m.called = true
|
||||||
|
outR, outW := schema.Pipe[*filesystem.ExecuteResponse](4)
|
||||||
|
go func() {
|
||||||
|
<-ctx.Done()
|
||||||
|
outW.Close()
|
||||||
|
}()
|
||||||
|
return outR, nil
|
||||||
|
}
|
||||||
|
|
||||||
func TestEinoExecuteRecvErrIsToolTimeout(t *testing.T) {
|
func TestEinoExecuteRecvErrIsToolTimeout(t *testing.T) {
|
||||||
tctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
|
tctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
|
||||||
defer cancel()
|
defer cancel()
|
||||||
|
|||||||
@@ -63,10 +63,43 @@ func toolCallArgsFromAccumulated(msgs []adk.Message, toolCallID, expectToolName
|
|||||||
return map[string]interface{}{}
|
return map[string]interface{}{}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// beginEinoADKFilesystemToolMonitor 在 Eino ADK filesystem 工具开始调用时写入 running 状态。
|
||||||
|
func beginEinoADKFilesystemToolMonitor(
|
||||||
|
ag *agent.Agent,
|
||||||
|
rec einomcp.ExecutionRecorder,
|
||||||
|
binder *MCPExecutionBinder,
|
||||||
|
toolCallID, toolName string,
|
||||||
|
) {
|
||||||
|
if ag == nil || rec == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
name := strings.TrimSpace(toolName)
|
||||||
|
if name == "" || strings.EqualFold(name, "execute") {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if !isBuiltinEinoADKFilesystemToolName(name) {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
tid := strings.TrimSpace(toolCallID)
|
||||||
|
if tid == "" {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
storedName := "eino_fs::" + strings.ToLower(name)
|
||||||
|
id := ag.BeginLocalToolExecution(storedName, map[string]interface{}{})
|
||||||
|
if id == "" {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
rec(id, tid)
|
||||||
|
if binder != nil {
|
||||||
|
binder.Bind(tid, id)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// recordEinoADKFilesystemToolMonitor 将 Eino ADK filesystem 中间件工具结果写入 MCP 监控(与 execute / MCP 桥芯片一致)。
|
// recordEinoADKFilesystemToolMonitor 将 Eino ADK filesystem 中间件工具结果写入 MCP 监控(与 execute / MCP 桥芯片一致)。
|
||||||
func recordEinoADKFilesystemToolMonitor(
|
func recordEinoADKFilesystemToolMonitor(
|
||||||
ag *agent.Agent,
|
ag *agent.Agent,
|
||||||
rec einomcp.ExecutionRecorder,
|
rec einomcp.ExecutionRecorder,
|
||||||
|
binder *MCPExecutionBinder,
|
||||||
toolName string,
|
toolName string,
|
||||||
toolCallID string,
|
toolCallID string,
|
||||||
msgs []adk.Message,
|
msgs []adk.Message,
|
||||||
@@ -94,8 +127,12 @@ func recordEinoADKFilesystemToolMonitor(
|
|||||||
invErr = errors.New(t)
|
invErr = errors.New(t)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
id := ag.RecordLocalToolExecution(storedName, args, resultText, invErr)
|
execID := ""
|
||||||
if id != "" {
|
if binder != nil {
|
||||||
|
execID = binder.ExecutionID(toolCallID)
|
||||||
|
}
|
||||||
|
id := ag.FinishLocalToolExecution(execID, storedName, args, resultText, invErr)
|
||||||
|
if id != "" && execID == "" {
|
||||||
rec(id, toolCallID)
|
rec(id, toolCallID)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -81,7 +81,7 @@ func RunEinoSingleChatModelAgent(
|
|||||||
}
|
}
|
||||||
|
|
||||||
toolInvokeNotify := einomcp.NewToolInvokeNotifyHolder()
|
toolInvokeNotify := einomcp.NewToolInvokeNotifyHolder()
|
||||||
einoExecMonitor := newEinoExecuteMonitorCallback(ag, recorder)
|
einoExecBegin, einoExecFinish := newEinoExecuteMonitorCallbacks(ag, recorder)
|
||||||
mainDefs := ag.ToolsForRole(roleTools)
|
mainDefs := ag.ToolsForRole(roleTools)
|
||||||
mainTools, err := einomcp.ToolsFromDefinitions(ag, holder, mainDefs, recorder, nil, toolInvokeNotify, einoSingleAgentName)
|
mainTools, err := einomcp.ToolsFromDefinitions(ag, holder, mainDefs, recorder, nil, toolInvokeNotify, einoSingleAgentName)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -136,7 +136,7 @@ func RunEinoSingleChatModelAgent(
|
|||||||
}
|
}
|
||||||
if einoSkillMW != nil {
|
if einoSkillMW != nil {
|
||||||
if einoFSTools && einoLoc != nil {
|
if einoFSTools && einoLoc != nil {
|
||||||
fsMw, fsErr := subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, einoSingleAgentName, einoExecMonitor, agentToolTimeoutMinutes(appCfg), nil)
|
fsMw, fsErr := subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, einoSingleAgentName, einoExecBegin, einoExecFinish, agentToolTimeoutMinutes(appCfg), agentShellNoOutputTimeoutSeconds(appCfg), nil)
|
||||||
if fsErr != nil {
|
if fsErr != nil {
|
||||||
return nil, fmt.Errorf("eino single filesystem 中间件: %w", fsErr)
|
return nil, fmt.Errorf("eino single filesystem 中间件: %w", fsErr)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -81,8 +81,10 @@ func subAgentFilesystemMiddleware(
|
|||||||
loc *localbk.Local,
|
loc *localbk.Local,
|
||||||
invokeNotify *einomcp.ToolInvokeNotifyHolder,
|
invokeNotify *einomcp.ToolInvokeNotifyHolder,
|
||||||
einoAgentName string,
|
einoAgentName string,
|
||||||
recordMonitor func(toolCallID, command, stdout string, success bool, invokeErr error),
|
beginMonitor func(toolCallID, command string) string,
|
||||||
|
finishMonitor func(executionID, toolCallID, command, stdout string, success bool, invokeErr error),
|
||||||
toolTimeoutMinutes int,
|
toolTimeoutMinutes int,
|
||||||
|
shellNoOutputTimeoutSec int,
|
||||||
outputChunk func(toolName, toolCallID, chunk string),
|
outputChunk func(toolName, toolCallID, chunk string),
|
||||||
) (adk.ChatModelAgentMiddleware, error) {
|
) (adk.ChatModelAgentMiddleware, error) {
|
||||||
if loc == nil {
|
if loc == nil {
|
||||||
@@ -91,12 +93,14 @@ func subAgentFilesystemMiddleware(
|
|||||||
return filesystem.New(ctx, &filesystem.MiddlewareConfig{
|
return filesystem.New(ctx, &filesystem.MiddlewareConfig{
|
||||||
Backend: loc,
|
Backend: loc,
|
||||||
StreamingShell: &einoStreamingShellWrap{
|
StreamingShell: &einoStreamingShellWrap{
|
||||||
inner: loc,
|
inner: loc,
|
||||||
invokeNotify: invokeNotify,
|
invokeNotify: invokeNotify,
|
||||||
einoAgentName: strings.TrimSpace(einoAgentName),
|
einoAgentName: strings.TrimSpace(einoAgentName),
|
||||||
outputChunk: outputChunk,
|
outputChunk: outputChunk,
|
||||||
recordMonitor: recordMonitor,
|
beginMonitor: beginMonitor,
|
||||||
toolTimeoutMinutes: toolTimeoutMinutes,
|
finishMonitor: finishMonitor,
|
||||||
|
toolTimeoutMinutes: toolTimeoutMinutes,
|
||||||
|
shellNoOutputTimeoutSec: shellNoOutputTimeoutSec,
|
||||||
},
|
},
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
@@ -108,3 +112,18 @@ func agentToolTimeoutMinutes(cfg *config.Config) int {
|
|||||||
}
|
}
|
||||||
return cfg.Agent.ToolTimeoutMinutes
|
return cfg.Agent.ToolTimeoutMinutes
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// agentShellNoOutputTimeoutSeconds:0=默认 300s(5 分钟);-1=关闭;>0=自定义秒数。
|
||||||
|
func agentShellNoOutputTimeoutSeconds(cfg *config.Config) int {
|
||||||
|
if cfg == nil {
|
||||||
|
return 300
|
||||||
|
}
|
||||||
|
v := cfg.Agent.ShellNoOutputTimeoutSeconds
|
||||||
|
if v < 0 {
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
if v == 0 {
|
||||||
|
return 300
|
||||||
|
}
|
||||||
|
return v
|
||||||
|
}
|
||||||
|
|||||||
@@ -120,7 +120,7 @@ func RunDeepAgent(
|
|||||||
mcpIDs = append(mcpIDs, id)
|
mcpIDs = append(mcpIDs, id)
|
||||||
mcpIDsMu.Unlock()
|
mcpIDsMu.Unlock()
|
||||||
}
|
}
|
||||||
einoExecMonitor := newEinoExecuteMonitorCallback(ag, recorder)
|
einoExecBegin, einoExecFinish := newEinoExecuteMonitorCallbacks(ag, recorder)
|
||||||
|
|
||||||
// 与单代理流式一致:在 response_start / response_delta 的 data 中带当前 mcpExecutionIds,供主聊天绑定复制与展示。
|
// 与单代理流式一致:在 response_start / response_delta 的 data 中带当前 mcpExecutionIds,供主聊天绑定复制与展示。
|
||||||
snapshotMCPIDs := func() []string {
|
snapshotMCPIDs := func() []string {
|
||||||
@@ -223,7 +223,7 @@ func RunDeepAgent(
|
|||||||
}
|
}
|
||||||
if einoSkillMW != nil {
|
if einoSkillMW != nil {
|
||||||
if einoFSTools && einoLoc != nil {
|
if einoFSTools && einoLoc != nil {
|
||||||
subFs, fsErr := subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, id, einoExecMonitor, agentToolTimeoutMinutes(appCfg), nil)
|
subFs, fsErr := subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, id, einoExecBegin, einoExecFinish, agentToolTimeoutMinutes(appCfg), agentShellNoOutputTimeoutSeconds(appCfg), nil)
|
||||||
if fsErr != nil {
|
if fsErr != nil {
|
||||||
return nil, fmt.Errorf("子代理 %q filesystem 中间件: %w", id, fsErr)
|
return nil, fmt.Errorf("子代理 %q filesystem 中间件: %w", id, fsErr)
|
||||||
}
|
}
|
||||||
@@ -358,12 +358,14 @@ func RunDeepAgent(
|
|||||||
if einoLoc != nil && einoFSTools {
|
if einoLoc != nil && einoFSTools {
|
||||||
deepBackend = einoLoc
|
deepBackend = einoLoc
|
||||||
deepShell = &einoStreamingShellWrap{
|
deepShell = &einoStreamingShellWrap{
|
||||||
inner: einoLoc,
|
inner: einoLoc,
|
||||||
invokeNotify: toolInvokeNotify,
|
invokeNotify: toolInvokeNotify,
|
||||||
einoAgentName: orchestratorName,
|
einoAgentName: orchestratorName,
|
||||||
outputChunk: nil,
|
outputChunk: nil,
|
||||||
recordMonitor: einoExecMonitor,
|
beginMonitor: einoExecBegin,
|
||||||
toolTimeoutMinutes: agentToolTimeoutMinutes(appCfg),
|
finishMonitor: einoExecFinish,
|
||||||
|
toolTimeoutMinutes: agentToolTimeoutMinutes(appCfg),
|
||||||
|
shellNoOutputTimeoutSec: agentShellNoOutputTimeoutSeconds(appCfg),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -428,7 +430,7 @@ func RunDeepAgent(
|
|||||||
// 构建 filesystem 中间件(与 Deep sub-agent 一致)
|
// 构建 filesystem 中间件(与 Deep sub-agent 一致)
|
||||||
var peFsMw adk.ChatModelAgentMiddleware
|
var peFsMw adk.ChatModelAgentMiddleware
|
||||||
if einoSkillMW != nil && einoFSTools && einoLoc != nil {
|
if einoSkillMW != nil && einoFSTools && einoLoc != nil {
|
||||||
peFsMw, err = subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, "executor", einoExecMonitor, agentToolTimeoutMinutes(appCfg), nil)
|
peFsMw, err = subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, "executor", einoExecBegin, einoExecFinish, agentToolTimeoutMinutes(appCfg), agentShellNoOutputTimeoutSeconds(appCfg), nil)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, fmt.Errorf("plan_execute filesystem 中间件: %w", err)
|
return nil, fmt.Errorf("plan_execute filesystem 中间件: %w", err)
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user