diff --git a/internal/mcp/execution_service.go b/internal/mcp/execution_service.go index 8de9c376..61a1c325 100644 --- a/internal/mcp/execution_service.go +++ b/internal/mcp/execution_service.go @@ -76,6 +76,7 @@ type ExecutionService struct { abortUserNotes map[string]string maxInMemory int resultMaxBytes int + spillRootDir string } func NewExecutionService(storage MonitorStorage, logger *zap.Logger) *ExecutionService { @@ -101,6 +102,17 @@ func (s *ExecutionService) ConfigureToolResultMaxBytes(maxBytes int) { s.resultMaxBytes = maxBytes } +// ConfigureToolResultSpillRoot sets the reduction-compatible root used when +// oversized tool results are spilled to local files (empty → tmp/reduction). +func (s *ExecutionService) ConfigureToolResultSpillRoot(rootDir string) { + if s == nil { + return + } + s.mu.Lock() + defer s.mu.Unlock() + s.spillRootDir = strings.TrimSpace(rootDir) +} + func (s *ExecutionService) Submit(ctx context.Context, req ExecutionRequest) (*ExecutionHandle, error) { if s == nil { return nil, fmt.Errorf("execution service is nil") @@ -163,6 +175,10 @@ func (s *ExecutionService) Submit(ctx context.Context, req ExecutionRequest) (*E func (s *ExecutionService) runWorker(ctx context.Context, entry *executionEntry, onDone ExecutionDoneFunc) { id := entry.exec.ID + ctx = WithMCPExecutionID(ctx, id) + if conv := strings.TrimSpace(entry.exec.ConversationID); conv != "" { + ctx = WithMCPConversationID(ctx, conv) + } var release func() defer func() { if release != nil { @@ -212,7 +228,20 @@ func (s *ExecutionService) finishEntry(ctx context.Context, entry *executionEntr now := time.Now() s.mu.Lock() - result = NormalizeToolResultForStorage(result, s.resultMaxBytes) + spill := ToolResultSpillConfig{ + RootDir: s.spillRootDir, + ConversationID: entry.exec.ConversationID, + ExecutionID: id, + } + if ctx != nil { + if pid := MCPProjectIDFromContext(ctx); pid != "" { + spill.ProjectID = pid + } + if conv := MCPConversationIDFromContext(ctx); conv != "" { + spill.ConversationID = conv + } + } + result = NormalizeToolResultForStorageWithSpill(result, s.resultMaxBytes, spill) entry.result = result entry.err = err entry.exec.EndTime = &now diff --git a/internal/mcp/external_manager.go b/internal/mcp/external_manager.go index ed6516ff..ca9913fd 100644 --- a/internal/mcp/external_manager.go +++ b/internal/mcp/external_manager.go @@ -77,6 +77,7 @@ type ExternalMCPManager struct { executionService *ExecutionService toolWaitTimeout time.Duration toolResultMaxBytes int + spillRootDir string resilience ExternalMCPResilienceConfig serverRuntimes map[string]*externalMCPServerRuntime globalSemaphore chan struct{} @@ -143,6 +144,20 @@ func (m *ExternalMCPManager) ConfigureToolResultMaxBytes(maxBytes int) { } } +// ConfigureToolResultSpillRoot sets the local directory root used when oversized +// tool results are spilled (aligned with reduction_root_dir; empty → tmp/reduction). +func (m *ExternalMCPManager) ConfigureToolResultSpillRoot(rootDir string) { + if m == nil { + return + } + m.mu.Lock() + m.spillRootDir = strings.TrimSpace(rootDir) + m.mu.Unlock() + if m.executionService != nil { + m.executionService.ConfigureToolResultSpillRoot(rootDir) + } +} + // ConfigureToolWaitTimeoutSeconds controls how long an agent-facing tool call // waits for an external MCP execution before returning an execution_id that can // be polled with wait_tool_execution. seconds<=0 waits until completion. @@ -704,6 +719,7 @@ func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args if m.executionService == nil { m.executionService = NewExecutionService(m.storage, m.logger) m.executionService.ConfigureToolResultMaxBytes(m.toolResultMaxBytes) + m.executionService.ConfigureToolResultSpillRoot(m.spillRootDir) } var ownerUserID string if principal, ok := authctx.PrincipalFromContext(ctx); ok { diff --git a/internal/mcp/run_context.go b/internal/mcp/run_context.go index d858c2ce..7612032e 100644 --- a/internal/mcp/run_context.go +++ b/internal/mcp/run_context.go @@ -22,6 +22,8 @@ type EinoExecuteRunRegistry interface { type toolRunRegistryCtxKey struct{} type einoExecuteRunRegistryCtxKey struct{} type mcpConversationIDCtxKey struct{} +type mcpExecutionIDCtxKey struct{} +type mcpProjectIDCtxKey struct{} // WithToolRunRegistry 将登记器注入 ctx(Eino / 原生 Agent 任务 ctx)。 func WithToolRunRegistry(ctx context.Context, reg ToolRunRegistry) context.Context { @@ -78,6 +80,48 @@ func MCPConversationIDFromContext(ctx context.Context) string { return v } +// WithMCPExecutionID 将当前工具 executionId 注入 ctx,供超长输出落盘文件名对齐。 +func WithMCPExecutionID(ctx context.Context, executionID string) context.Context { + if ctx == nil { + return nil + } + id := strings.TrimSpace(executionID) + if id == "" { + return ctx + } + return context.WithValue(ctx, mcpExecutionIDCtxKey{}, id) +} + +// MCPExecutionIDFromContext 读取当前工具 executionId。 +func MCPExecutionIDFromContext(ctx context.Context) string { + if ctx == nil { + return "" + } + v, _ := ctx.Value(mcpExecutionIDCtxKey{}).(string) + return v +} + +// WithMCPProjectID 将项目 ID 注入 ctx,供 reduction/trunc 落盘路径与项目隔离对齐。 +func WithMCPProjectID(ctx context.Context, projectID string) context.Context { + if ctx == nil { + return nil + } + id := strings.TrimSpace(projectID) + if id == "" { + return ctx + } + return context.WithValue(ctx, mcpProjectIDCtxKey{}, id) +} + +// MCPProjectIDFromContext 读取项目 ID。 +func MCPProjectIDFromContext(ctx context.Context) string { + if ctx == nil { + return "" + } + v, _ := ctx.Value(mcpProjectIDCtxKey{}).(string) + return v +} + func notifyToolRunBegin(ctx context.Context, executionID string) { reg := ToolRunRegistryFromContext(ctx) if reg == nil { diff --git a/internal/mcp/server.go b/internal/mcp/server.go index 9251c667..d2838db5 100644 --- a/internal/mcp/server.go +++ b/internal/mcp/server.go @@ -56,6 +56,7 @@ type Server struct { executionService *ExecutionService toolWaitTimeout time.Duration toolResultMaxBytes int + spillRootDir string } // SetToolAuthorizer installs the common policy decision point for every @@ -128,6 +129,20 @@ func (s *Server) ConfigureToolResultMaxBytes(maxBytes int) { } } +// ConfigureToolResultSpillRoot sets the local directory root used when oversized +// tool results are spilled (aligned with reduction_root_dir; empty → tmp/reduction). +func (s *Server) ConfigureToolResultSpillRoot(rootDir string) { + if s == nil { + return + } + s.mu.Lock() + s.spillRootDir = strings.TrimSpace(rootDir) + s.mu.Unlock() + if s.executionService != nil { + s.executionService.ConfigureToolResultSpillRoot(rootDir) + } +} + // ConfigureHTTPToolCallTimeoutFromAgentMinutes 将 agent.tool_timeout_minutes 同步到经 HTTP POST /api/mcp 触发的 tools/call。 // minutes<=0 表示不设置硬性截止时间(与配置「0 不限制」一致);minutes>0 为该次调用的最长等待时间。 // 未调用前对 tools/call 使用默认 30 分钟(与历史硬编码一致)。 @@ -900,6 +915,7 @@ func (s *Server) CallTool(ctx context.Context, toolName string, args map[string] if s.executionService == nil { s.executionService = NewExecutionService(s.storage, s.logger) s.executionService.ConfigureToolResultMaxBytes(s.toolResultMaxBytes) + s.executionService.ConfigureToolResultSpillRoot(s.spillRootDir) } var ownerUserID string if principal, ok := authctx.PrincipalFromContext(ctx); ok { @@ -1036,6 +1052,7 @@ func (s *Server) FinishToolExecution(ctx context.Context, executionID, toolName s.mu.Lock() maxBytes := s.toolResultMaxBytes + spillRoot := s.spillRootDir exec, inMem := s.executions[id] if !inMem || exec == nil { exec = &ToolExecution{ @@ -1063,13 +1080,19 @@ func (s *Server) FinishToolExecution(ctx context.Context, executionID, toolName } exec.Duration = now.Sub(exec.StartTime) + spill := ToolResultSpillConfig{ + RootDir: spillRoot, + ProjectID: MCPProjectIDFromContext(ctx), + ConversationID: exec.ConversationID, + ExecutionID: id, + } if failed { st, msg := executionStatusAndMessage(invokeErr) exec.Status = st exec.Error = msg if strings.TrimSpace(resultText) != "" { finalResult = &ToolResult{Content: []Content{{Type: "text", Text: resultText}}} - finalResult = NormalizeToolResultForStorage(finalResult, maxBytes) + finalResult = NormalizeToolResultForStorageWithSpill(finalResult, maxBytes, spill) exec.Result = finalResult } } else { @@ -1079,7 +1102,7 @@ func (s *Server) FinishToolExecution(ctx context.Context, executionID, toolName text = "(无输出)" } finalResult = &ToolResult{Content: []Content{{Type: "text", Text: text}}} - finalResult = NormalizeToolResultForStorage(finalResult, maxBytes) + finalResult = NormalizeToolResultForStorageWithSpill(finalResult, maxBytes, spill) exec.Result = finalResult } s.mu.Unlock() @@ -1116,9 +1139,16 @@ func (s *Server) UpdateToolExecutionResult(executionID string, result *ToolResul return nil } s.mu.Lock() - result = NormalizeToolResultForStorage(result, s.toolResultMaxBytes) + spill := ToolResultSpillConfig{ + RootDir: s.spillRootDir, + ExecutionID: executionID, + } if exec, ok := s.executions[executionID]; ok && exec != nil { + spill.ConversationID = exec.ConversationID + result = NormalizeToolResultForStorageWithSpill(result, s.toolResultMaxBytes, spill) exec.Result = result + } else { + result = NormalizeToolResultForStorageWithSpill(result, s.toolResultMaxBytes, spill) } s.mu.Unlock() if s.storage != nil { diff --git a/internal/mcp/tool_result_guard.go b/internal/mcp/tool_result_guard.go index 70fdbb9b..851db4df 100644 --- a/internal/mcp/tool_result_guard.go +++ b/internal/mcp/tool_result_guard.go @@ -1,12 +1,29 @@ package mcp -import "fmt" +import "cyberstrike-ai/internal/tooloutput" const DefaultToolResultMaxBytes = 12000 +// ToolResultSpillConfig controls where oversized tool results are written on disk +// before the in-memory/DB/agent-facing payload is truncated. +type ToolResultSpillConfig struct { + RootDir string + ProjectID string + ConversationID string + ExecutionID string +} + // NormalizeToolResultForStorage returns the canonical result used by both the -// agent-facing response and monitor persistence. +// agent-facing response and monitor persistence. When maxBytes is exceeded the +// full text is spilled under the reduction cache tree and replaced with a +// notice that includes the file path. func NormalizeToolResultForStorage(result *ToolResult, maxBytes int) *ToolResult { + return NormalizeToolResultForStorageWithSpill(result, maxBytes, ToolResultSpillConfig{}) +} + +// NormalizeToolResultForStorageWithSpill is NormalizeToolResultForStorage with +// an explicit spill location (conversation/execution scoped). +func NormalizeToolResultForStorageWithSpill(result *ToolResult, maxBytes int, spill ToolResultSpillConfig) *ToolResult { if result == nil { return nil } @@ -25,48 +42,14 @@ func NormalizeToolResultForStorage(result *ToolResult, maxBytes int) *ToolResult return out } - remaining := maxBytes - truncated := false - for i := range out.Content { - if out.Content[i].Type != "text" { - continue - } - if remaining <= 0 { - out.Content[i].Text = "" - truncated = true - continue - } - text := out.Content[i].Text - if len(text) <= remaining { - remaining -= len(text) - continue - } - out.Content[i].Text = truncateUTF8Bytes(text, remaining) - remaining = 0 - truncated = true - } - - if truncated { - marker := fmt.Sprintf("\n\n...[tool output truncated: original %d bytes, kept %d bytes]...", total, maxBytes) - textBudget := maxBytes - len(marker) - if textBudget < 0 { - marker = truncateUTF8Bytes(marker, maxBytes) - textBudget = 0 - } - for i := range out.Content { - if out.Content[i].Type == "text" { - out.Content[i].Text = truncateUTF8Bytes(out.Content[i].Text, textBudget) + marker - remaining = 0 - for j := range out.Content { - if j != i && out.Content[j].Type == "text" { - out.Content[j].Text = "" - } - } - return out - } - } - out.Content = append(out.Content, Content{Type: "text", Text: marker}) - } + full := ToolResultPlainText(out) + bound := tooloutput.BoundWithSpill(full, maxBytes, tooloutput.SpillOpts{ + RootDir: spill.RootDir, + ProjectID: spill.ProjectID, + ConversationID: spill.ConversationID, + ExecutionID: spill.ExecutionID, + }) + out.Content = []Content{{Type: "text", Text: bound}} return out } @@ -80,20 +63,3 @@ func cloneToolResult(in *ToolResult) *ToolResult { } return &out } - -func truncateUTF8Bytes(s string, maxBytes int) string { - if maxBytes <= 0 { - return "" - } - if len(s) <= maxBytes { - return s - } - cut := maxBytes - for cut > 0 && (s[cut]&0xC0) == 0x80 { - cut-- - } - if cut <= 0 { - return "" - } - return s[:cut] -} diff --git a/internal/mcp/tool_result_guard_test.go b/internal/mcp/tool_result_guard_test.go index 4a0ff55c..fd7c8d55 100644 --- a/internal/mcp/tool_result_guard_test.go +++ b/internal/mcp/tool_result_guard_test.go @@ -2,6 +2,8 @@ package mcp import ( "context" + "os" + "path/filepath" "strings" "testing" "time" @@ -63,12 +65,15 @@ func TestServerCallToolStoresAndReturnsSameGuardedResult(t *testing.T) { storage := newInMemoryMonitorStorage() server := NewServerWithStorage(zap.NewNop(), storage) server.ConfigureToolWaitTimeoutSeconds(0) - server.ConfigureToolResultMaxBytes(50) + server.ConfigureToolResultMaxBytes(400) + spillRoot := t.TempDir() + server.ConfigureToolResultSpillRoot(spillRoot) server.RegisterTool(Tool{Name: "big", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) { - return &ToolResult{Content: []Content{{Type: "text", Text: strings.Repeat("x", 100)}}}, nil + return &ToolResult{Content: []Content{{Type: "text", Text: strings.Repeat("x", 800)}}}, nil }) - result, executionID, err := server.CallTool(context.Background(), "big", nil) + ctx := WithMCPConversationID(context.Background(), "conv-spill") + result, executionID, err := server.CallTool(ctx, "big", nil) if err != nil { t.Fatalf("CallTool: %v", err) } @@ -76,13 +81,29 @@ func TestServerCallToolStoresAndReturnsSameGuardedResult(t *testing.T) { t.Fatal("missing execution id") } returned := ToolResultPlainText(result) - if !strings.Contains(returned, "tool output truncated") || strings.Contains(returned, strings.Repeat("x", 100)) { - t.Fatalf("returned result was not guarded: %q", returned) + if !strings.Contains(returned, "") || !strings.Contains(returned, "Full output saved to:") { + t.Fatalf("returned result was not spilled: %q", returned) } - if len(returned) > 50 { + if len(returned) > 400 { t.Fatalf("returned result exceeded hard limit: len=%d text=%q", len(returned), returned) } + spillPath := filepath.Join(spillRoot, "conversations", "conv-spill", "trunc", executionID) + abs, err := filepath.Abs(spillPath) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(returned, abs) { + t.Fatalf("missing spill path %q in %q", abs, returned) + } + body, err := os.ReadFile(abs) + if err != nil { + t.Fatalf("read spill file: %v", err) + } + if string(body) != strings.Repeat("x", 800) { + t.Fatalf("spill body mismatch: len=%d", len(body)) + } + inMem, ok := server.GetExecution(executionID) if !ok || inMem == nil || inMem.Result == nil { t.Fatalf("missing in-memory execution: %#v", inMem) @@ -101,11 +122,14 @@ func TestServerCallToolStoresAndReturnsSameGuardedResult(t *testing.T) { func TestExecutionServiceStoresGuardedResult(t *testing.T) { service := NewExecutionService(nil, zap.NewNop()) - service.ConfigureToolResultMaxBytes(80) + service.ConfigureToolResultMaxBytes(400) + spillRoot := t.TempDir() + service.ConfigureToolResultSpillRoot(spillRoot) handle, err := service.Submit(context.Background(), ExecutionRequest{ - ToolName: "big", + ToolName: "big", + ConversationID: "svc-conv", Run: func(context.Context) (*ToolResult, error) { - return &ToolResult{Content: []Content{{Type: "text", Text: strings.Repeat("a", 200)}}}, nil + return &ToolResult{Content: []Content{{Type: "text", Text: strings.Repeat("a", 800)}}}, nil }, }) if err != nil { @@ -116,10 +140,19 @@ func TestExecutionServiceStoresGuardedResult(t *testing.T) { t.Fatalf("Wait: %v", err) } got := ToolResultPlainText(snap.Execution.Result) - if !strings.Contains(got, "tool output truncated") || strings.Contains(got, strings.Repeat("a", 64)) { - t.Fatalf("service result was not guarded: %q", got) + if !strings.Contains(got, "") { + t.Fatalf("service result was not spilled: %q", got) } - if len(got) > 80 { + if len(got) > 400 { t.Fatalf("service result exceeded hard limit: len=%d text=%q", len(got), got) } + path := filepath.Join(spillRoot, "conversations", "svc-conv", "trunc", handle.ID) + abs, _ := filepath.Abs(path) + body, err := os.ReadFile(abs) + if err != nil { + t.Fatalf("read spill: %v", err) + } + if string(body) != strings.Repeat("a", 800) { + t.Fatalf("unexpected spill body len=%d", len(body)) + } } diff --git a/internal/security/executor.go b/internal/security/executor.go index 0dafc55d..04af6100 100644 --- a/internal/security/executor.go +++ b/internal/security/executor.go @@ -16,8 +16,10 @@ import ( "cyberstrike-ai/internal/config" "cyberstrike-ai/internal/mcp" + "cyberstrike-ai/internal/tooloutput" "github.com/creack/pty" + "github.com/google/uuid" "go.uber.org/zap" ) @@ -38,6 +40,7 @@ type Executor struct { logger *zap.Logger shellNoOutputTimeoutSec int // execute/exec 无新输出空闲秒数;0=默认 300;-1=关闭(见 SetShellNoOutputTimeoutSeconds) toolOutputMaxBytes int + spillRootDir string } // NewExecutor 创建新的执行器 @@ -60,11 +63,34 @@ func (e *Executor) SetShellNoOutputTimeoutSeconds(sec int) { // SetToolOutputMaxBytes limits stdout/stderr retained and streamed by exec-like // tools. It should stay aligned with MCP result normalization so every channel -// sees the same bounded payload. +// sees the same bounded payload. Oversized full output is spilled to disk first. func (e *Executor) SetToolOutputMaxBytes(maxBytes int) { e.toolOutputMaxBytes = maxBytes } +// SetToolOutputSpillRoot sets the reduction-compatible root for spilling full +// exec stdout/stderr when the in-memory bound is exceeded (empty → tmp/reduction). +func (e *Executor) SetToolOutputSpillRoot(rootDir string) { + e.spillRootDir = strings.TrimSpace(rootDir) +} + +func (e *Executor) spillOptsFromContext(ctx context.Context) tooloutput.SpillOpts { + root := "" + if e != nil { + root = e.spillRootDir + } + opts := tooloutput.SpillOpts{RootDir: root} + if ctx != nil { + opts.ConversationID = mcp.MCPConversationIDFromContext(ctx) + opts.ProjectID = mcp.MCPProjectIDFromContext(ctx) + opts.ExecutionID = mcp.MCPExecutionIDFromContext(ctx) + } + if opts.ExecutionID == "" { + opts.ExecutionID = uuid.NewString() + } + return opts +} + // buildToolIndex 构建工具索引,将 O(n) 查找优化为 O(1) func (e *Executor) buildToolIndex() { e.toolIndex = make(map[string]*config.ToolConfig) @@ -157,9 +183,10 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st var output string var err error + spill := e.spillOptsFromContext(ctx) // 如果上层提供了 stdout/stderr 增量回调,则边执行边读取并回调。 if cb, ok := ctx.Value(ToolOutputCallbackCtxKey).(ToolOutputCallback); ok && cb != nil { - output, err = streamCommandOutput(ctx, cmd, cb, ResolveShellNoOutputTimeoutSeconds(e.shellNoOutputTimeoutSec), e.toolOutputMaxBytes) + output, err = streamCommandOutput(ctx, cmd, cb, ResolveShellNoOutputTimeoutSeconds(e.shellNoOutputTimeoutSec), e.toolOutputMaxBytes, spill) if err != nil && shouldRetryWithPTY(output) { e.logger.Info("检测到工具需要 TTY,使用 PTY 重试", zap.String("tool", toolName), @@ -167,11 +194,11 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st cmd2 := exec.CommandContext(ctx, toolConfig.Command, cmdArgs...) applyDefaultTerminalEnv(cmd2) _ = prepareShellCmdSession(cmd2) - output, err = runCommandWithPTY(ctx, cmd2, cb, e.toolOutputMaxBytes) + output, err = runCommandWithPTY(ctx, cmd2, cb, e.toolOutputMaxBytes, spill) } } else { // 非流式:内存缓冲 + ctx 取消杀进程组;行为对齐原 CombinedOutput,避免双流管道 fan-in 死锁。 - output, err = combinedOutputCancellableWithLimit(ctx, cmd, e.toolOutputMaxBytes) + output, err = combinedOutputCancellableWithLimit(ctx, cmd, e.toolOutputMaxBytes, spill) if err != nil && shouldRetryWithPTY(output) { e.logger.Info("检测到工具需要 TTY,使用 PTY 重试", zap.String("tool", toolName), @@ -179,7 +206,7 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st cmd2 := exec.CommandContext(ctx, toolConfig.Command, cmdArgs...) applyDefaultTerminalEnv(cmd2) _ = prepareShellCmdSession(cmd2) - output, err = runCommandWithPTY(ctx, cmd2, nil, e.toolOutputMaxBytes) + output, err = runCommandWithPTY(ctx, cmd2, nil, e.toolOutputMaxBytes, spill) } } if err != nil { @@ -920,9 +947,10 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int // 非后台命令:等待输出 var output string var err error + spill := e.spillOptsFromContext(ctx) // 若上层提供工具输出增量回调,则边执行边流式读取。 if cb, ok := ctx.Value(ToolOutputCallbackCtxKey).(ToolOutputCallback); ok && cb != nil { - output, err = streamCommandOutput(ctx, cmd, cb, ResolveShellNoOutputTimeoutSeconds(e.shellNoOutputTimeoutSec), e.toolOutputMaxBytes) + output, err = streamCommandOutput(ctx, cmd, cb, ResolveShellNoOutputTimeoutSeconds(e.shellNoOutputTimeoutSec), e.toolOutputMaxBytes, spill) if err != nil && shouldRetryWithPTY(output) { e.logger.Info("检测到系统命令需要 TTY,使用 PTY 重试") cmd2 := exec.CommandContext(ctx, shell, "-c", command) @@ -930,10 +958,10 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int cmd2.Dir = workDir } ConfigureShellCmdForAgentExecute(cmd2) - output, err = runCommandWithPTY(ctx, cmd2, cb, e.toolOutputMaxBytes) + output, err = runCommandWithPTY(ctx, cmd2, cb, e.toolOutputMaxBytes, spill) } } else { - output, err = combinedOutputCancellableWithLimit(ctx, cmd, e.toolOutputMaxBytes) + output, err = combinedOutputCancellableWithLimit(ctx, cmd, e.toolOutputMaxBytes, spill) if err != nil && shouldRetryWithPTY(output) { e.logger.Info("检测到系统命令需要 TTY,使用 PTY 重试") cmd2 := exec.CommandContext(ctx, shell, "-c", command) @@ -941,7 +969,7 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int cmd2.Dir = workDir } ConfigureShellCmdForAgentExecute(cmd2) - output, err = runCommandWithPTY(ctx, cmd2, nil, e.toolOutputMaxBytes) + output, err = runCommandWithPTY(ctx, cmd2, nil, e.toolOutputMaxBytes, spill) } } if err != nil { @@ -982,12 +1010,17 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int // 非流式路径不使用双流管道 fan-in,避免 stderr 撑满管道缓冲区时与 stdout 互相阻塞导致死锁。 // 无输出空闲检测由上层 agent.tool_timeout_minutes 兜底,不改变原 CombinedOutput 语义。 func combinedOutputCancellable(ctx context.Context, cmd *exec.Cmd) (string, error) { - return combinedOutputCancellableWithLimit(ctx, cmd, 0) + return combinedOutputCancellableWithLimit(ctx, cmd, 0, tooloutput.SpillOpts{}) } -func combinedOutputCancellableWithLimit(ctx context.Context, cmd *exec.Cmd, maxBytes int) (string, error) { - stdoutBuf := newBoundedOutputCollector(maxBytes) - stderrBuf := newBoundedOutputCollector(maxBytes) +func combinedOutputCancellableWithLimit(ctx context.Context, cmd *exec.Cmd, maxBytes int, spill tooloutput.SpillOpts) (string, error) { + var tee *tooloutput.Tee + if maxBytes > 0 { + tee = tooloutput.NewTee(spill) + defer func() { _ = tee.Close() }() + } + stdoutBuf := newBoundedOutputCollector(maxBytes, tee) + stderrBuf := newBoundedOutputCollector(maxBytes, tee) cmd.Stdout = stdoutBuf cmd.Stderr = stderrBuf @@ -1016,9 +1049,9 @@ func combinedOutputCancellableWithLimit(ctx context.Context, cmd *exec.Cmd, maxB case waitErr = <-done: case <-ctx.Done(): waitErr = <-done - return limitOutputString(joinCommandOutput(stdoutBuf.String(), stderrBuf.String()), maxBytes), ctx.Err() + return finalizeJoinedBoundedOutputs(stdoutBuf, stderrBuf, maxBytes, tee), ctx.Err() } - return limitOutputString(joinCommandOutput(stdoutBuf.String(), stderrBuf.String()), maxBytes), waitErr + return finalizeJoinedBoundedOutputs(stdoutBuf, stderrBuf, maxBytes, tee), waitErr } func joinCommandOutput(stdout, stderr string) string { @@ -1036,10 +1069,11 @@ type boundedOutputCollector struct { maxBytes int seenBytes int truncated bool + tee *tooloutput.Tee } -func newBoundedOutputCollector(maxBytes int) *boundedOutputCollector { - return &boundedOutputCollector{maxBytes: maxBytes} +func newBoundedOutputCollector(maxBytes int, tee *tooloutput.Tee) *boundedOutputCollector { + return &boundedOutputCollector{maxBytes: maxBytes, tee: tee} } func (b *boundedOutputCollector) Write(p []byte) (int, error) { @@ -1051,27 +1085,20 @@ func (b *boundedOutputCollector) WriteStringLimited(s string) string { if b == nil { return "" } + if b.tee != nil { + _, _ = b.tee.Write([]byte(s)) + } if b.maxBytes <= 0 { b.seenBytes += len(s) b.builder.WriteString(s) return s } b.seenBytes += len(s) - marker := b.truncationMarker() - contentLimit := b.maxBytes - len(marker) - if contentLimit < 0 { - marker = truncateStringBytes(marker, b.maxBytes) - contentLimit = 0 - } - if b.builder.Len() >= contentLimit { - if b.truncated { - return "" - } + if b.builder.Len() >= b.maxBytes { b.truncated = true - b.builder.WriteString(marker) - return marker + return "" } - remaining := contentLimit - b.builder.Len() + remaining := b.maxBytes - b.builder.Len() if len(s) <= remaining { b.builder.WriteString(s) return s @@ -1079,8 +1106,7 @@ func (b *boundedOutputCollector) WriteStringLimited(s string) string { kept := truncateStringBytes(s, remaining) b.builder.WriteString(kept) b.truncated = true - b.builder.WriteString(marker) - return kept + marker + return kept } func (b *boundedOutputCollector) String() string { @@ -1090,13 +1116,90 @@ func (b *boundedOutputCollector) String() string { return b.builder.String() } -func (b *boundedOutputCollector) truncationMarker() string { - return fmt.Sprintf("\n\n...[tool output limit reached: kept %d bytes; further output suppressed]...", b.maxBytes) +func finalizeJoinedBoundedOutputs(stdout, stderr *boundedOutputCollector, maxBytes int, tee *tooloutput.Tee) string { + if tee != nil { + _ = tee.Close() + } + truncated := (stdout != nil && stdout.truncated) || (stderr != nil && stderr.truncated) + seen := 0 + if stdout != nil { + seen += stdout.seenBytes + } + if stderr != nil { + seen += stderr.seenBytes + } + joined := joinCommandOutput( + func() string { + if stdout == nil { + return "" + } + return stdout.String() + }(), + func() string { + if stderr == nil { + return "" + } + return stderr.String() + }(), + ) + if maxBytes > 0 && !truncated && len(joined) > maxBytes { + truncated = true + seen = len(joined) + } + path := "" + if tee != nil { + path = tee.Path() + } + if truncated && maxBytes > 0 { + if path != "" { + return tooloutput.FormatPersistedFromFile(path, seen, maxBytes) + } + if len(joined) > maxBytes { + return truncateStringBytes(joined, maxBytes) + } + return joined + } + if path != "" { + _ = os.Remove(path) + } + if maxBytes > 0 && len(joined) > maxBytes { + return truncateStringBytes(joined, maxBytes) + } + return joined } -func limitOutputString(s string, maxBytes int) string { - collector := newBoundedOutputCollector(maxBytes) - return collector.WriteStringLimited(s) +func finalizeBoundedOutput(collector *boundedOutputCollector, maxBytes int, tee *tooloutput.Tee) string { + if tee != nil { + _ = tee.Close() + } + if collector == nil { + return "" + } + path := "" + if tee != nil { + path = tee.Path() + } + if collector.truncated && maxBytes > 0 { + if path != "" { + return tooloutput.FormatPersistedFromFile(path, collector.seenBytes, maxBytes) + } + return truncateStringBytes(collector.String(), maxBytes) + } + if path != "" { + _ = os.Remove(path) + } + out := collector.String() + if maxBytes > 0 && len(out) > maxBytes { + return tooloutput.BoundWithSpill(out, maxBytes, tooloutput.SpillOpts{}) + } + return out +} + +func limitOutputString(s string, maxBytes int, spill tooloutput.SpillOpts) string { + if maxBytes <= 0 || len(s) <= maxBytes { + return s + } + return tooloutput.BoundWithSpill(s, maxBytes, spill) } func truncateStringBytes(s string, maxBytes int) string { @@ -1118,7 +1221,7 @@ func truncateStringBytes(s string, maxBytes int) string { // streamCommandOutput 以“边读边回调”的方式读取命令 stdout/stderr。 // 使用定长块读取,避免按行读取在无换行输出时永久阻塞;ctx 取消时终止进程树。 -func streamCommandOutput(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback, noOutputSec int, maxBytes int) (string, error) { +func streamCommandOutput(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback, noOutputSec int, maxBytes int, spill tooloutput.SpillOpts) (string, error) { stdoutPipe, err := cmd.StdoutPipe() if err != nil { return "", err @@ -1170,7 +1273,12 @@ func streamCommandOutput(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallba close(chunks) }() - outBuilder := newBoundedOutputCollector(maxBytes) + tee := (*tooloutput.Tee)(nil) + if maxBytes > 0 { + tee = tooloutput.NewTee(spill) + defer func() { _ = tee.Close() }() + } + outBuilder := newBoundedOutputCollector(maxBytes, tee) var deltaBuilder strings.Builder lastFlush := time.Now() @@ -1214,7 +1322,7 @@ chunksLoop: return outBuilder.String(), ctx.Err() case <-idleCh: fireInactivity() - return outBuilder.String(), fmt.Errorf("shell inactivity timeout (%ds)", idleWatch.Sec) + return finalizeBoundedOutput(outBuilder, maxBytes, tee), fmt.Errorf("shell inactivity timeout (%ds)", idleWatch.Sec) case chunk, ok := <-chunks: if !ok { break chunksLoop @@ -1233,7 +1341,7 @@ chunksLoop: // 等待命令结束,返回最终退出状态 waitErr := session.Wait() - return outBuilder.String(), waitErr + return finalizeBoundedOutput(outBuilder, maxBytes, tee), waitErr } // applyDefaultTerminalEnv 为外部工具补齐常见的终端环境变量。 @@ -1286,14 +1394,14 @@ func shouldRetryWithPTY(output string) bool { // runCommandWithPTY 为子进程分配 PTY,适配需要交互式终端的工具(如 autorecon)。 // 若 cb != nil,将持续回调增量输出(用于 SSE)。 -func runCommandWithPTY(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback, maxBytes int) (string, error) { +func runCommandWithPTY(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback, maxBytes int, spill tooloutput.SpillOpts) (string, error) { if runtime.GOOS == "windows" { // PTY 方案为类 Unix;Windows 走原逻辑 if cb != nil { - return streamCommandOutput(ctx, cmd, cb, 0, maxBytes) + return streamCommandOutput(ctx, cmd, cb, 0, maxBytes, spill) } _ = prepareShellCmdSession(cmd) - return combinedOutputCancellableWithLimit(ctx, cmd, maxBytes) + return combinedOutputCancellableWithLimit(ctx, cmd, maxBytes, spill) } _ = prepareShellCmdSession(cmd) @@ -1320,7 +1428,12 @@ func runCommandWithPTY(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback }() defer close(done) - outBuilder := newBoundedOutputCollector(maxBytes) + tee := (*tooloutput.Tee)(nil) + if maxBytes > 0 { + tee = tooloutput.NewTee(spill) + defer func() { _ = tee.Close() }() + } + outBuilder := newBoundedOutputCollector(maxBytes, tee) var deltaBuilder strings.Builder lastFlush := time.Now() flush := func() { @@ -1355,7 +1468,7 @@ func runCommandWithPTY(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback flush() waitErr := cmd.Wait() - return outBuilder.String(), waitErr + return finalizeBoundedOutput(outBuilder, maxBytes, tee), waitErr } // executeInternalTool 执行内部工具(不执行外部命令) diff --git a/internal/security/executor_test.go b/internal/security/executor_test.go index a2be9c7b..1c40dfed 100644 --- a/internal/security/executor_test.go +++ b/internal/security/executor_test.go @@ -96,8 +96,11 @@ func TestExecuteSystemCommand_FailureFormat(t *testing.T) { func TestExecuteSystemCommand_OutputIsSourceLimited(t *testing.T) { executor, _ := setupTestExecutor(t) - executor.SetToolOutputMaxBytes(64) - res, err := executor.executeSystemCommand(context.Background(), map[string]interface{}{ + spillRoot := t.TempDir() + executor.SetToolOutputMaxBytes(200) + executor.SetToolOutputSpillRoot(spillRoot) + ctx := mcp.WithMCPConversationID(context.Background(), "exec-spill") + res, err := executor.executeSystemCommand(ctx, map[string]interface{}{ "command": "i=0; while [ $i -lt 2000 ]; do printf 0123456789; i=$((i+1)); done", "shell": "sh", }) @@ -108,10 +111,10 @@ func TestExecuteSystemCommand_OutputIsSourceLimited(t *testing.T) { t.Fatalf("expected success, got %+v", res) } text := res.Content[0].Text - if !strings.Contains(text, "tool output limit reached") { - t.Fatalf("missing output limit marker: %q", text) + if !strings.Contains(text, "") || !strings.Contains(text, "Full output saved to:") { + t.Fatalf("missing persisted-output notice: %q", text) } - if len(text) > 64 { + if len(text) > 200 { t.Fatalf("output exceeded hard limit: len=%d text=%q", len(text), text) } if strings.Contains(text, strings.Repeat("0123456789", 20)) { @@ -121,11 +124,14 @@ func TestExecuteSystemCommand_OutputIsSourceLimited(t *testing.T) { func TestExecuteSystemCommand_StreamingOutputIsSourceLimited(t *testing.T) { executor, _ := setupTestExecutor(t) - executor.SetToolOutputMaxBytes(64) + spillRoot := t.TempDir() + executor.SetToolOutputMaxBytes(200) + executor.SetToolOutputSpillRoot(spillRoot) var streamed strings.Builder ctx := context.WithValue(context.Background(), ToolOutputCallbackCtxKey, ToolOutputCallback(func(chunk string) { streamed.WriteString(chunk) })) + ctx = mcp.WithMCPConversationID(ctx, "exec-stream-spill") res, err := executor.executeSystemCommand(ctx, map[string]interface{}{ "command": "i=0; while [ $i -lt 2000 ]; do printf abcdefghij; i=$((i+1)); done", "shell": "sh", @@ -134,17 +140,21 @@ func TestExecuteSystemCommand_StreamingOutputIsSourceLimited(t *testing.T) { t.Fatalf("executeSystemCommand: %v", err) } text := res.Content[0].Text - if text != streamed.String() { - t.Fatalf("streamed output and returned output diverged\nstream=%q\nret=%q", streamed.String(), text) + if !strings.Contains(text, "") { + t.Fatalf("missing persisted-output notice: %q", text) } - if !strings.Contains(text, "tool output limit reached") { - t.Fatalf("missing output limit marker: %q", text) + if len(text) > 200 { + t.Fatalf("returned output exceeded hard limit: len=%d text=%q", len(text), text) } - if len(text) > 64 { - t.Fatalf("streamed output exceeded hard limit: len=%d text=%q", len(text), text) + // SSE only streams the bounded prefix; final agent-facing body is the spill notice. + if len(streamed.String()) > 200 { + t.Fatalf("streamed prefix exceeded hard limit: len=%d", len(streamed.String())) } - if strings.Contains(text, strings.Repeat("abcdefghij", 20)) { - t.Fatalf("streamed output kept too much data: len=%d", len(text)) + if streamed.Len() == 0 { + t.Fatal("expected some streamed prefix before truncation") + } + if strings.Contains(text, strings.Repeat("abcdefghij", 50)) { + t.Fatalf("returned output kept too much raw data: len=%d", len(text)) } }