subagent.go

  1package claudetool
  2
  3import (
  4	"context"
  5	"encoding/json"
  6	"fmt"
  7	"strings"
  8	"time"
  9
 10	"shelley.exe.dev/llm"
 11)
 12
 13// SubagentRunner is the interface for running a subagent conversation.
 14// This is implemented by the server package to avoid import cycles.
 15type SubagentRunner interface {
 16	// RunSubagent runs a subagent conversation and returns the last response.
 17	// If wait is false, it starts processing in background and returns immediately.
 18	// timeout is the maximum time to wait for a response.
 19	RunSubagent(ctx context.Context, conversationID, prompt string, wait bool, timeout time.Duration) (string, error)
 20}
 21
 22// SubagentDB is the database interface for subagent operations.
 23// This is implemented by the db package.
 24type SubagentDB interface {
 25	// GetOrCreateSubagentConversation retrieves or creates a subagent conversation.
 26	// Returns the conversation ID and the actual slug used (may differ from requested
 27	// slug if a numeric suffix was added for uniqueness).
 28	GetOrCreateSubagentConversation(ctx context.Context, slug, parentID, cwd string) (conversationID, actualSlug string, err error)
 29}
 30
 31// SubagentTool provides the ability to spawn and interact with subagent conversations.
 32type SubagentTool struct {
 33	DB                   SubagentDB
 34	ParentConversationID string
 35	WorkingDir           *MutableWorkingDir
 36	Runner               SubagentRunner
 37}
 38
 39const (
 40	subagentName        = "subagent"
 41	subagentDescription = `Spawn or interact with a subagent conversation.
 42
 43Subagents are independent conversations that can work on subtasks in parallel.
 44Use subagents for:
 45- Long-running tasks that you want to delegate
 46- Token-intensive tasks that produce lots of output, little of which is needed
 47- Parallel exploration of different approaches
 48- Breaking down complex problems into independent pieces
 49
 50Each subagent has its own slug identifier within this conversation.
 51You can send messages to existing subagents by using the same slug.
 52The tool returns the subagent's last response, or a status if the timeout is reached.
 53`
 54	subagentInputSchema = `
 55{
 56  "type": "object",
 57  "required": ["slug", "prompt"],
 58  "properties": {
 59    "slug": {
 60      "type": "string",
 61      "description": "A short identifier for this subagent (e.g., 'research-api', 'test-runner')"
 62    },
 63    "prompt": {
 64      "type": "string",
 65      "description": "The message to send to the subagent"
 66    },
 67    "timeout_seconds": {
 68      "type": "integer",
 69      "description": "How long to wait for a response (default: 60, max: 300)"
 70    },
 71    "wait": {
 72      "type": "boolean",
 73      "description": "Whether to wait for completion (default: true). If false, returns immediately."
 74    }
 75  }
 76}
 77`
 78)
 79
 80type subagentInput struct {
 81	Slug           string `json:"slug"`
 82	Prompt         string `json:"prompt"`
 83	TimeoutSeconds int    `json:"timeout_seconds,omitempty"`
 84	Wait           *bool  `json:"wait,omitempty"`
 85}
 86
 87// Tool returns an llm.Tool for the subagent functionality.
 88func (s *SubagentTool) Tool() *llm.Tool {
 89	return &llm.Tool{
 90		Name:        subagentName,
 91		Description: strings.TrimSpace(subagentDescription),
 92		InputSchema: llm.MustSchema(subagentInputSchema),
 93		Run:         s.Run,
 94	}
 95}
 96
 97func (s *SubagentTool) Run(ctx context.Context, m json.RawMessage) llm.ToolOut {
 98	var req subagentInput
 99	if err := json.Unmarshal(m, &req); err != nil {
100		return llm.ErrorfToolOut("failed to parse subagent input: %w", err)
101	}
102
103	// Validate slug
104	if req.Slug == "" {
105		return llm.ErrorfToolOut("slug is required")
106	}
107	req.Slug = sanitizeSlug(req.Slug)
108	if req.Slug == "" {
109		return llm.ErrorfToolOut("slug must contain alphanumeric characters")
110	}
111
112	if req.Prompt == "" {
113		return llm.ErrorfToolOut("prompt is required")
114	}
115
116	// Set defaults
117	timeout := 60 * time.Second
118	if req.TimeoutSeconds > 0 {
119		if req.TimeoutSeconds > 300 {
120			req.TimeoutSeconds = 300
121		}
122		timeout = time.Duration(req.TimeoutSeconds) * time.Second
123	}
124
125	wait := true
126	if req.Wait != nil {
127		wait = *req.Wait
128	}
129
130	// Get or create the subagent conversation
131	conversationID, actualSlug, err := s.DB.GetOrCreateSubagentConversation(ctx, req.Slug, s.ParentConversationID, s.WorkingDir.Get())
132	if err != nil {
133		return llm.ErrorfToolOut("failed to get/create subagent conversation: %w", err)
134	}
135
136	// Use the runner to execute the subagent
137	response, err := s.Runner.RunSubagent(ctx, conversationID, req.Prompt, wait, timeout)
138	if err != nil {
139		return llm.ErrorfToolOut("subagent error: %w", err)
140	}
141
142	// Include actual slug in response if it differs from requested
143	slugNote := ""
144	if actualSlug != req.Slug {
145		slugNote = fmt.Sprintf(" (Note: slug was changed to '%s' for uniqueness. Use '%s' for future messages to this subagent.)", actualSlug, actualSlug)
146	}
147
148	return llm.ToolOut{
149		LLMContent: llm.TextContent(fmt.Sprintf("Subagent '%s' response:%s\n%s", actualSlug, slugNote, response)),
150		Display: SubagentDisplayData{
151			Slug:           actualSlug,
152			ConversationID: conversationID,
153		},
154	}
155}
156
157// SubagentDisplayData is the display data sent to the UI for subagent tool results.
158type SubagentDisplayData struct {
159	Slug           string `json:"slug"`
160	ConversationID string `json:"conversation_id"`
161}
162
163func sanitizeSlug(slug string) string {
164	// Lowercase, keep alphanumeric and hyphens
165	var result strings.Builder
166	for _, r := range strings.ToLower(slug) {
167		if (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') || r == '-' {
168			result.WriteRune(r)
169		} else if r == ' ' || r == '_' {
170			result.WriteRune('-')
171		}
172	}
173	// Remove consecutive hyphens and trim
174	s := result.String()
175	for strings.Contains(s, "--") {
176		s = strings.ReplaceAll(s, "--", "-")
177	}
178	return strings.Trim(s, "-")
179}