diff --git a/internal/mcp/execution_control_tools.go b/internal/mcp/execution_control_tools.go index f8b56783..7d74af25 100644 --- a/internal/mcp/execution_control_tools.go +++ b/internal/mcp/execution_control_tools.go @@ -15,6 +15,8 @@ import ( const ( defaultExecutionWaitTimeout = 60 * time.Second maxExecutionWaitTimeout = 10 * time.Minute + defaultPartialPreviewBytes = 4096 + maxPartialPreviewBytes = 64 * 1024 ) // RegisterExecutionControlTools exposes execution handle operations to Eino as @@ -32,7 +34,9 @@ func RegisterExecutionControlTools(server *Server, external *ExternalMCPManager) InputSchema: map[string]interface{}{ "type": "object", "properties": map[string]interface{}{ - "execution_id": map[string]interface{}{"type": "string", "description": "工具执行 ID"}, + "execution_id": map[string]interface{}{"type": "string", "description": "工具执行 ID"}, + "include_partial_output": map[string]interface{}{"type": "boolean", "description": "是否返回运行中已产生输出的尾部预览，默认 true"}, + "partial_output_max_bytes": map[string]interface{}{"type": "number", "description": "partial_output 最多返回字节数，默认 4096，最大 65536"}, }, "required": []string{"execution_id"}, }, @@ -45,7 +49,7 @@ func RegisterExecutionControlTools(server *Server, external *ExternalMCPManager) if exec == nil { return textToolResult("未找到该 execution_id: "+id, true), nil } - return textToolResult(formatExecutionForModel(exec), false), nil + return textToolResult(formatExecutionForModel(exec, executionFormatOptionsFromArgs(args)), false), nil }) server.RegisterTool(Tool{ @@ -55,8 +59,10 @@ func RegisterExecutionControlTools(server *Server, external *ExternalMCPManager) InputSchema: map[string]interface{}{ "type": "object", "properties": map[string]interface{}{ - "execution_id": map[string]interface{}{"type": "string", "description": "工具执行 ID"}, - "timeout_seconds": map[string]interface{}{"type": "number", "description": "本次最多等待秒数，默认 60，最大 600"}, + "execution_id": map[string]interface{}{"type": "string", "description": "工具执行 ID"}, + "timeout_seconds": map[string]interface{}{"type": "number", "description": "本次最多等待秒数，默认 60，最大 600"}, + "include_partial_output": map[string]interface{}{"type": "boolean", "description": "是否返回运行中已产生输出的尾部预览，默认 true"}, + "partial_output_max_bytes": map[string]interface{}{"type": "number", "description": "partial_output 最多返回字节数，默认 4096，最大 65536"}, }, "required": []string{"execution_id"}, }, @@ -73,7 +79,7 @@ func RegisterExecutionControlTools(server *Server, external *ExternalMCPManager) if snap == nil || snap.Execution == nil { return textToolResult("未找到该 execution_id: "+id, true), nil } - body := formatExecutionForModel(snap.Execution) + body := formatExecutionForModel(snap.Execution, executionFormatOptionsFromArgs(args)) if errors.Is(err, ErrExecutionWaitTimeout) { body += "\n\n本次等待已到达 timeout_seconds，上述 execution 仍未完成。可继续等待、取消，或采用其他步骤。" } @@ -144,7 +150,25 @@ func lookupToolExecution(server *Server, external *ExternalMCPManager, id string return nil } -func formatExecutionForModel(exec *ToolExecution) string { +type executionFormatOptions struct { + includePartialOutput bool + partialMaxBytes int +} + +func executionFormatOptionsFromArgs(args map[string]interface{}) executionFormatOptions { + includePartial := true + if raw, ok := args["include_partial_output"]; ok { + if b, ok := raw.(bool); ok { + includePartial = b + } else if s := strings.TrimSpace(fmt.Sprint(raw)); s != "" { + includePartial = strings.EqualFold(s, "true") || s == "1" || strings.EqualFold(s, "yes") + } + } + maxBytes := intArg(args, "partial_output_max_bytes", defaultPartialPreviewBytes, maxPartialPreviewBytes) + return executionFormatOptions{includePartialOutput: includePartial, partialMaxBytes: maxBytes} +} + +func formatExecutionForModel(exec *ToolExecution, opts executionFormatOptions) string { if exec == nil { return "execution: null" } @@ -167,6 +191,15 @@ func formatExecutionForModel(exec *ToolExecution) string { payload["result"] = ToolResultPlainText(exec.Result) payload["is_error"] = exec.Result.IsError } + if opts.includePartialOutput && exec.PartialOutput != "" { + partial := tailStringBytes(exec.PartialOutput, opts.partialMaxBytes) + payload["partial_output"] = partial + payload["partial_output_bytes"] = exec.PartialOutputBytes + payload["partial_output_truncated"] = exec.PartialOutputTruncated || len([]byte(partial)) < len([]byte(exec.PartialOutput)) + if exec.PartialOutputUpdatedAt != nil { + payload["partial_output_updated_at"] = exec.PartialOutputUpdatedAt.Format(time.RFC3339) + } + } b, err := json.MarshalIndent(payload, "", " ") if err != nil { return fmt.Sprintf("execution_id: %s\nstatus: %s\nerror: %s", exec.ID, exec.Status, exec.Error) @@ -174,6 +207,17 @@ func formatExecutionForModel(exec *ToolExecution) string { return string(b) } +func tailStringBytes(s string, maxBytes int) string { + if maxBytes <= 0 { + maxBytes = defaultPartialPreviewBytes + } + b := []byte(s) + if len(b) <= maxBytes { + return s + } + return string(b[len(b)-maxBytes:]) +} + func textToolResult(text string, isErr bool) *ToolResult { return &ToolResult{Content: []Content{{Type: "text", Text: text}}, IsError: isErr} } @@ -222,3 +266,31 @@ func durationSecondsArg(args map[string]interface{}, key string, def, max time.D } return d } + +func intArg(args map[string]interface{}, key string, def, max int) int { + if args == nil { + return def + } + var n int + switch v := args[key].(type) { + case int: + n = v + case int64: + n = int(v) + case float64: + n = int(v) + case json.Number: + i, _ := v.Int64() + n = int(i) + case string: + i, _ := strconv.Atoi(strings.TrimSpace(v)) + n = i + } + if n <= 0 { + return def + } + if max > 0 && n > max { + return max + } + return n +} diff --git a/internal/mcp/execution_service.go b/internal/mcp/execution_service.go index 61a1c325..1e2ad6d5 100644 --- a/internal/mcp/execution_service.go +++ b/internal/mcp/execution_service.go @@ -348,6 +348,22 @@ func (s *ExecutionService) Get(executionID string) (*ExecutionSnapshot, error) { return s.getPersistedSnapshot(executionID) } +func (s *ExecutionService) AppendPartialOutput(executionID, chunk string) bool { + id := strings.TrimSpace(executionID) + if s == nil || id == "" || chunk == "" { + return false + } + now := time.Now() + s.mu.Lock() + defer s.mu.Unlock() + entry := s.entries[id] + if entry == nil || entry.exec == nil { + return false + } + appendPartialOutput(entry.exec, chunk, defaultPartialOutputMaxBytes, now) + return true +} + func (s *ExecutionService) Cancel(executionID, note string) bool { id := strings.TrimSpace(executionID) if id == "" || s == nil { @@ -582,5 +598,28 @@ func cloneToolExecution(in *ToolExecution) *ToolExecution { t := *in.EndTime out.EndTime = &t } + if in.PartialOutputUpdatedAt != nil { + t := *in.PartialOutputUpdatedAt + out.PartialOutputUpdatedAt = &t + } return &out } + +func appendPartialOutput(exec *ToolExecution, chunk string, maxBytes int, updatedAt time.Time) { + if exec == nil || chunk == "" { + return + } + if maxBytes <= 0 { + maxBytes = defaultPartialOutputMaxBytes + } + exec.PartialOutputBytes += int64(len([]byte(chunk))) + combined := exec.PartialOutput + chunk + if len([]byte(combined)) > maxBytes { + b := []byte(combined) + combined = string(b[len(b)-maxBytes:]) + exec.PartialOutputTruncated = true + } + exec.PartialOutput = combined + t := updatedAt + exec.PartialOutputUpdatedAt = &t +} diff --git a/internal/mcp/server.go b/internal/mcp/server.go index 2bf2c8ec..0ed8abb3 100644 --- a/internal/mcp/server.go +++ b/internal/mcp/server.go @@ -59,6 +59,8 @@ type Server struct { spillRootDir string } +const defaultPartialOutputMaxBytes = 64 * 1024 + // SetToolAuthorizer installs the common policy decision point for every // user-attributed tool call, whether it originates from HTTP or an Agent. func (s *Server) SetToolAuthorizer(authorizer func(context.Context, string, map[string]interface{}) error) { @@ -1123,6 +1125,25 @@ func (s *Server) FinishToolExecution(ctx context.Context, executionID, toolName return id } +// AppendToolExecutionPartialOutput records a bounded tail preview for a running local execution. +// The final Result remains authoritative and is written only when the tool finishes. +func (s *Server) AppendToolExecutionPartialOutput(executionID, chunk string) { + if s == nil || strings.TrimSpace(executionID) == "" || chunk == "" { + return + } + id := strings.TrimSpace(executionID) + if s.executionService != nil && s.executionService.AppendPartialOutput(id, chunk) { + return + } + now := time.Now() + s.mu.Lock() + exec := s.executions[id] + if exec != nil { + appendPartialOutput(exec, chunk, defaultPartialOutputMaxBytes, now) + } + s.mu.Unlock() +} + // RecordCompletedToolInvocation 将已在其它路径完成的工具调用写入监控存储（格式与 CallTool 结束后一致）， // 用于 Eino ADK filesystem execute 等未经过 CallTool 的场景；返回 executionId 供助手消息 mcpExecutionIds 关联。 func (s *Server) RecordCompletedToolInvocation(ctx context.Context, toolName string, args map[string]interface{}, resultText string, invokeErr error) string { @@ -1206,6 +1227,24 @@ func (s *Server) unregisterRunningCancel(id string) { s.runningCancelsMu.Unlock() } +// RegisterToolExecutionCancel lets non-ExecutionService tool paths, such as Eino +// filesystem execute, participate in cancel_tool_execution by execution_id. +func (s *Server) RegisterToolExecutionCancel(id string, cancel context.CancelFunc) { + id = strings.TrimSpace(id) + if s == nil || id == "" || cancel == nil { + return + } + s.registerRunningCancel(id, cancel) +} + +func (s *Server) UnregisterToolExecutionCancel(id string) { + id = strings.TrimSpace(id) + if s == nil || id == "" { + return + } + s.unregisterRunningCancel(id) +} + func (s *Server) readAbortUserNote(id string) string { s.runningCancelsMu.Lock() defer s.runningCancelsMu.Unlock() diff --git a/internal/mcp/server_authorization_test.go b/internal/mcp/server_authorization_test.go index ecfcf3cf..90f0d1a9 100644 --- a/internal/mcp/server_authorization_test.go +++ b/internal/mcp/server_authorization_test.go @@ -180,3 +180,41 @@ func TestWaitToolExecutionTimeoutIsObservationNotFailure(t *testing.T) { } close(release) } + +func TestGetToolExecutionIncludesBoundedPartialOutput(t *testing.T) { + server := NewServer(zap.NewNop()) + RegisterExecutionControlTools(server, nil) + + executionID := server.BeginToolExecution(context.Background(), "execute", map[string]interface{}{"command": "demo"}) + if executionID == "" { + t.Fatal("missing execution id") + } + server.AppendToolExecutionPartialOutput(executionID, "first\n") + server.AppendToolExecutionPartialOutput(executionID, strings.Repeat("x", 32)) + + result, _, err := server.CallTool(context.Background(), "get_tool_execution", map[string]interface{}{ + "execution_id": executionID, + "partial_output_max_bytes": 8, + }) + if err != nil { + t.Fatalf("get_tool_execution: %v", err) + } + body := ToolResultPlainText(result) + if !strings.Contains(body, `"partial_output": "xxxxxxxx"`) { + t.Fatalf("missing bounded partial output: %s", body) + } + if !strings.Contains(body, `"partial_output_bytes": 38`) { + t.Fatalf("missing partial byte count: %s", body) + } + + result, _, err = server.CallTool(context.Background(), "get_tool_execution", map[string]interface{}{ + "execution_id": executionID, + "include_partial_output": false, + }) + if err != nil { + t.Fatalf("get_tool_execution without partial: %v", err) + } + if body := ToolResultPlainText(result); strings.Contains(body, "partial_output") { + t.Fatalf("partial output should be omitted: %s", body) + } +} diff --git a/internal/mcp/types.go b/internal/mcp/types.go index d8217286..6922e047 100644 --- a/internal/mcp/types.go +++ b/internal/mcp/types.go @@ -199,6 +199,12 @@ type ToolExecution struct { StartTime time.Time `json:"startTime"` EndTime *time.Time `json:"endTime,omitempty"` Duration time.Duration `json:"duration,omitempty"` + // PartialOutput is a bounded tail preview of output produced by a running tool. + // It is intentionally separate from Result, which remains the final canonical tool result. + PartialOutput string `json:"partialOutput,omitempty"` + PartialOutputBytes int64 `json:"partialOutputBytes,omitempty"` + PartialOutputTruncated bool `json:"partialOutputTruncated,omitempty"` + PartialOutputUpdatedAt *time.Time `json:"partialOutputUpdatedAt,omitempty"` // ConversationID 仅 API 展示用（进行中的 Agent 任务），不写入 tool_executions 表。 ConversationID string `json:"conversationId,omitempty"` OwnerUserID string `json:"-"` diff --git a/internal/security/executor.go b/internal/security/executor.go index 04af6100..a1f6ae97 100644 --- a/internal/security/executor.go +++ b/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 重试") diff --git a/internal/security/executor_test.go b/internal/security/executor_test.go index 1c40dfed..4b62889a 100644 --- a/internal/security/executor_test.go +++ b/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{}{