Add files via upload

internal/security/executor.go

@@ -74,6 +74,21 @@ func (e *Executor) SetToolOutputSpillRoot(rootDir string) {
 	e.spillRootDir = strings.TrimSpace(rootDir)
 }
 
+func (e *Executor) wrapToolOutputCallback(ctx context.Context, cb ToolOutputCallback) ToolOutputCallback {
+	executionID := mcp.MCPExecutionIDFromContext(ctx)
+	if e == nil || e.mcpServer == nil || strings.TrimSpace(executionID) == "" {
+		return cb
+	}
+	return func(chunk string) {
+		if chunk != "" {
+			e.mcpServer.AppendToolExecutionPartialOutput(executionID, chunk)
+		}
+		if cb != nil {
+			cb(chunk)
+		}
+	}
+}
+
 func (e *Executor) spillOptsFromContext(ctx context.Context) tooloutput.SpillOpts {
 	root := ""
 	if e != nil {
@@ -184,8 +199,9 @@ 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 {
+	// 如果上层提供了 stdout/stderr 增量回调，或当前处于 MCP execution 中，则边执行边读取并回调。
+	if cb, ok := ctx.Value(ToolOutputCallbackCtxKey).(ToolOutputCallback); (ok && cb != nil) || mcp.MCPExecutionIDFromContext(ctx) != "" {
+		cb = e.wrapToolOutputCallback(ctx, cb)
 		output, err = streamCommandOutput(ctx, cmd, cb, ResolveShellNoOutputTimeoutSeconds(e.shellNoOutputTimeoutSec), e.toolOutputMaxBytes, spill)
 		if err != nil && shouldRetryWithPTY(output) {
 			e.logger.Info("检测到工具需要 TTY，使用 PTY 重试",
@@ -948,8 +964,9 @@ 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 {
+	// 若上层提供工具输出增量回调，或当前处于 MCP execution 中，则边执行边流式读取。
+	if cb, ok := ctx.Value(ToolOutputCallbackCtxKey).(ToolOutputCallback); (ok && cb != nil) || mcp.MCPExecutionIDFromContext(ctx) != "" {
+		cb = e.wrapToolOutputCallback(ctx, cb)
 		output, err = streamCommandOutput(ctx, cmd, cb, ResolveShellNoOutputTimeoutSeconds(e.shellNoOutputTimeoutSec), e.toolOutputMaxBytes, spill)
 		if err != nil && shouldRetryWithPTY(output) {
 			e.logger.Info("检测到系统命令需要 TTY，使用 PTY 重试")

internal/security/executor_test.go

@@ -73,6 +73,43 @@ func TestExecuteSystemCommand_BackgroundDoesNotBlockOnChildStdout(t *testing.T)
 	}
 }
 
+func TestExecToolSoftWaitExposesPartialOutput(t *testing.T) {
+	executor, server := setupTestExecutor(t)
+	server.ConfigureToolWaitTimeoutSeconds(1)
+	mcp.RegisterExecutionControlTools(server, nil)
+	server.RegisterTool(mcp.Tool{Name: "exec", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
+		return executor.ExecuteTool(ctx, "exec", args)
+	})
+
+	result, executionID, err := server.CallTool(context.Background(), "exec", map[string]interface{}{
+		"command": "for i in 1 2 3 4; do echo partial-$i; sleep 0.3; done; sleep 5",
+		"shell":   "sh",
+	})
+	if err != nil {
+		t.Fatalf("CallTool exec: %v", err)
+	}
+	if executionID == "" || result == nil || !result.IsError {
+		t.Fatalf("expected soft wait timeout, id=%q result=%#v", executionID, result)
+	}
+
+	status, _, err := server.CallTool(context.Background(), "get_tool_execution", map[string]interface{}{
+		"execution_id":             executionID,
+		"include_partial_output":   true,
+		"partial_output_max_bytes": 4096,
+	})
+	if err != nil {
+		t.Fatalf("get_tool_execution: %v", err)
+	}
+	body := mcp.ToolResultPlainText(status)
+	if !strings.Contains(body, `"status": "running"`) {
+		t.Fatalf("expected running execution, got: %s", body)
+	}
+	if !strings.Contains(body, "partial-") || !strings.Contains(body, "partial_output") {
+		t.Fatalf("expected partial output in execution status, got: %s", body)
+	}
+	server.CancelToolExecution(executionID)
+}
+
 func TestExecuteSystemCommand_FailureFormat(t *testing.T) {
 	executor, _ := setupTestExecutor(t)
 	res, err := executor.executeSystemCommand(context.Background(), map[string]interface{}{