MuyueWorkspace/internal/agent/delegate.go

package agent

import (
	"context"
	"encoding/json"
	"fmt"
	"sync"
	"time"
)

type DelegateTaskParams struct {
	Task        string `json:"task" description:"Description of the sub-task to delegate"`
	Context     string `json:"context,omitempty" description:"Additional context for the sub-task"`
	Timeout     int    `json:"timeout,omitempty" description:"Timeout per sub-task in seconds (default 120, max 300)"`
	MaxParallel int    `json:"max_parallel,omitempty" description:"Maximum parallel sub-tasks (default 3, max 5)"`
}

type DelegateMultiParams struct {
	Tasks       []DelegateTaskParams `json:"tasks" description:"List of sub-tasks to execute in parallel"`
	MaxParallel int                  `json:"max_parallel,omitempty" description:"Maximum parallel sub-tasks (default 3, max 5)"`
}

type SubTaskResult struct {
	Task    string `json:"task"`
	Success bool   `json:"success"`
	Result  string `json:"result"`
	Error   string `json:"error,omitempty"`
}

type DelegateResponse struct {
	TotalTasks int             `json:"total_tasks"`
	Successful int             `json:"successful"`
	Failed     int             `json:"failed"`
	Results    []SubTaskResult `json:"results"`
	Duration   string          `json:"duration"`
}

func NewDelegateTool(registry *Registry) (*ToolDefinition, error) {
	return NewTool("delegate_task",
		"Delegate one or more tasks for parallel execution. Each sub-task runs in isolation with its own context. Returns aggregated results from all sub-tasks. Use for independent tasks that can run concurrently.",
		func(ctx context.Context, p DelegateTaskParams) (ToolResponse, error) {
			if p.Task == "" {
				return TextErrorResponse("task is required"), nil
			}

			timeout := time.Duration(p.Timeout) * time.Second
			if timeout == 0 {
				timeout = 120 * time.Second
			}
			if timeout > 300*time.Second {
				timeout = 300 * time.Second
			}

			result := executeSubTask(ctx, p.Task, p.Context, timeout, registry)
			resp := DelegateResponse{
				TotalTasks: 1,
				Successful: 0,
				Results:    []SubTaskResult{result},
				Duration:   "N/A",
			}
			if result.Success {
				resp.Successful = 1
			} else {
				resp.Failed = 1
			}

			data, _ := json.MarshalIndent(resp, "", "  ")
			return TextResponse(string(data)), nil
		})
}

func NewDelegateMultiTool(registry *Registry) (*ToolDefinition, error) {
	return NewTool("delegate_multi",
		"Execute multiple independent tasks in parallel using goroutines. Each task runs in its own isolated context. Returns aggregated results. Use for batch operations, parallel analysis, or concurrent file processing.",
		func(ctx context.Context, p DelegateMultiParams) (ToolResponse, error) {
			if len(p.Tasks) == 0 {
				return TextErrorResponse("tasks list is required"), nil
			}

			maxParallel := p.MaxParallel
			if maxParallel <= 0 {
				maxParallel = 3
			}
			if maxParallel > 5 {
				maxParallel = 5
			}

			if len(p.Tasks) > 10 {
				return TextErrorResponse("maximum 10 tasks per delegation"), nil
			}

			start := time.Now()
			results := executeParallelTasks(ctx, p.Tasks, maxParallel, registry)
			duration := time.Since(start)

			resp := DelegateResponse{
				TotalTasks: len(results),
				Results:    results,
				Duration:   duration.Round(time.Millisecond).String(),
			}

			for _, r := range results {
				if r.Success {
					resp.Successful++
				} else {
					resp.Failed++
				}
			}

			data, _ := json.MarshalIndent(resp, "", "  ")
			return TextResponse(string(data)), nil
		})
}

func executeSubTask(ctx context.Context, task, contextInfo string, timeout time.Duration, registry *Registry) SubTaskResult {
	taskCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	result := SubTaskResult{
		Task: truncateString(task, 100),
	}

	if contextInfo != "" {
		result.Task = fmt.Sprintf("%s (context: %s)", result.Task, truncateString(contextInfo, 50))
	}

	done := make(chan struct{})
	go func() {
		defer close(done)

		terminalTool, ok := registry.Get("terminal")
		if !ok {
			result.Error = "terminal tool not available"
			return
		}

		args, _ := json.Marshal(TerminalParams{
			Command: task,
			Timeout: int(timeout.Seconds()),
		})

		resp, err := terminalTool.Execute(taskCtx, ToolCall{
			ID:       fmt.Sprintf("delegate_%d", time.Now().UnixNano()),
			Name:     "terminal",
			Arguments: args,
		})
		if err != nil {
			result.Error = err.Error()
			return
		}

		result.Result = resp.Content
		result.Success = !resp.IsError
		if resp.IsError {
			result.Error = resp.Content
		}
	}()

	select {
	case <-done:
		return result
	case <-taskCtx.Done():
		result.Error = fmt.Sprintf("sub-task timed out after %v", timeout)
		return result
	}
}

func executeParallelTasks(ctx context.Context, tasks []DelegateTaskParams, maxParallel int, registry *Registry) []SubTaskResult {
	results := make([]SubTaskResult, len(tasks))

	sem := make(chan struct{}, maxParallel)
	var wg sync.WaitGroup

	for i, task := range tasks {
		wg.Add(1)
		go func(idx int, t DelegateTaskParams) {
			defer wg.Done()

			sem <- struct{}{}
			defer func() { <-sem }()

			timeout := time.Duration(t.Timeout) * time.Second
			if timeout == 0 {
				timeout = 120 * time.Second
			}
			if timeout > 300*time.Second {
				timeout = 300 * time.Second
			}

			results[idx] = executeSubTask(ctx, t.Task, t.Context, timeout, registry)
		}(i, task)
	}

	wg.Wait()
	return results
}

func truncateString(s string, maxLen int) string {
	if len(s) <= maxLen {
		return s
	}
	return s[:maxLen] + "..."
}