coder/coderd/x/chatd/subagent.go

package chatd

import (
	"context"
	"database/sql"
	"encoding/json"
	"sort"
	"strings"
	"time"

	"charm.land/fantasy"
	"github.com/google/uuid"
	"golang.org/x/xerrors"

	"github.com/coder/coder/v2/coderd/database"
	coderdpubsub "github.com/coder/coder/v2/coderd/pubsub"
	"github.com/coder/coder/v2/coderd/x/chatd/chatprompt"
	"github.com/coder/coder/v2/coderd/x/chatd/chatprovider"
	"github.com/coder/coder/v2/codersdk"
)

var ErrSubagentNotDescendant = xerrors.New("target chat is not a descendant of current chat")

const (
	subagentAwaitPollInterval  = 200 * time.Millisecond
	subagentAwaitFallbackPoll  = 5 * time.Second
	defaultSubagentWaitTimeout = 5 * time.Minute
)

// computerUseSubagentSystemPrompt is the system prompt prepended to
// every computer use subagent chat. It instructs the model on how to
// interact with the desktop environment via the computer tool.
const computerUseSubagentSystemPrompt = `You are a computer use agent with access to a desktop environment. You can see the screen, move the mouse, click, type, scroll, and drag.

Your primary tool is the "computer" tool which lets you interact with the desktop. After every action you take, you will receive a screenshot showing the current state of the screen. Use these screenshots to verify your actions and plan next steps.

Guidelines:
- Always start by taking a screenshot to see the current state of the desktop.
- Be precise with coordinates when clicking or typing.
- Wait for UI elements to load before interacting with them.
- If an action doesn't produce the expected result, try alternative approaches.
- Report what you accomplished when done.`

type spawnAgentArgs struct {
	Prompt string `json:"prompt"`
	Title  string `json:"title,omitempty"`
}

type spawnComputerUseAgentArgs struct {
	Prompt string `json:"prompt"`
	Title  string `json:"title,omitempty"`
}

type waitAgentArgs struct {
	ChatID         string `json:"chat_id"`
	TimeoutSeconds *int   `json:"timeout_seconds,omitempty"`
}

type messageAgentArgs struct {
	ChatID    string `json:"chat_id"`
	Message   string `json:"message"`
	Interrupt bool   `json:"interrupt,omitempty"`
}

type closeAgentArgs struct {
	ChatID string `json:"chat_id"`
}

// isAnthropicConfigured reports whether an Anthropic API key is
// available, either from static provider keys or from the database.
func (p *Server) isAnthropicConfigured(ctx context.Context) bool {
	if p.providerAPIKeys.APIKey("anthropic") != "" {
		return true
	}
	dbProviders, err := p.db.GetEnabledChatProviders(ctx)
	if err != nil {
		return false
	}
	for _, prov := range dbProviders {
		if chatprovider.NormalizeProvider(prov.Provider) == "anthropic" && strings.TrimSpace(prov.APIKey) != "" {
			return true
		}
	}
	return false
}

func (p *Server) isDesktopEnabled(ctx context.Context) bool {
	enabled, err := p.db.GetChatDesktopEnabled(ctx)
	if err != nil {
		return false
	}
	return enabled
}

func (p *Server) subagentTools(ctx context.Context, currentChat func() database.Chat) []fantasy.AgentTool {
	tools := []fantasy.AgentTool{
		fantasy.NewAgentTool(
			"spawn_agent",
			"Spawn a delegated child agent to work on a clearly scoped, "+
				"independent task in parallel. Use this when the task is "+
				"self-contained and would benefit from a separate agent "+
				"(e.g. fixing a specific bug, writing a single module, "+
				"running a migration). Do NOT use for simple or quick "+
				"operations you can handle directly with execute, "+
				"read_file, or write_file - for example, reading a group "+
				"of files and outputting them verbatim does not need a "+
				"subagent. Reserve subagents for tasks that require "+
				"intellectual work such as code analysis, writing new "+
				"code, or complex refactoring. Be careful when running "+
				"parallel subagents: if two subagents modify the same "+
				"files they will conflict with each other, so ensure "+
				"parallel subagent tasks are independent. "+
				"The child agent receives the same workspace tools but "+
				"cannot spawn its own subagents. After spawning, use "+
				"wait_agent to collect the result.",
			func(ctx context.Context, args spawnAgentArgs, _ fantasy.ToolCall) (fantasy.ToolResponse, error) {
				if currentChat == nil {
					return fantasy.NewTextErrorResponse("subagent callbacks are not configured"), nil
				}

				parent := currentChat()
				if parent.ParentChatID.Valid {
					return fantasy.NewTextErrorResponse("delegated chats cannot create child subagents"), nil
				}

				parent, err := p.db.GetChatByID(ctx, parent.ID)
				if err != nil {
					return fantasy.NewTextErrorResponse(err.Error()), nil
				}
				childChat, err := p.createChildSubagentChat(
					ctx,
					parent,
					args.Prompt,
					args.Title,
				)
				if err != nil {
					return fantasy.NewTextErrorResponse(err.Error()), nil
				}

				return toolJSONResponse(map[string]any{
					"chat_id": childChat.ID.String(),
					"title":   childChat.Title,
					"status":  string(childChat.Status),
				}), nil
			},
		),
		fantasy.NewAgentTool(
			"wait_agent",
			"Wait until a spawned child agent finishes its task. "+
				"Returns the agent's final response and status. "+
				"Call this after spawn_agent to collect the result "+
				"before continuing your own work.",
			func(ctx context.Context, args waitAgentArgs, _ fantasy.ToolCall) (fantasy.ToolResponse, error) {
				if currentChat == nil {
					return fantasy.NewTextErrorResponse("subagent callbacks are not configured"), nil
				}

				targetChatID, err := parseSubagentToolChatID(args.ChatID)
				if err != nil {
					return fantasy.NewTextErrorResponse(err.Error()), nil
				}

				timeout := defaultSubagentWaitTimeout
				if args.TimeoutSeconds != nil {
					timeout = time.Duration(*args.TimeoutSeconds) * time.Second
				}

				parent := currentChat()
				targetChat, report, err := p.awaitSubagentCompletion(
					ctx,
					parent.ID,
					targetChatID,
					timeout,
				)
				if err != nil {
					return fantasy.NewTextErrorResponse(err.Error()), nil
				}

				return toolJSONResponse(map[string]any{
					"chat_id": targetChatID.String(),
					"title":   targetChat.Title,
					"report":  report,
					"status":  string(targetChat.Status),
				}), nil
			},
		),
		fantasy.NewAgentTool(
			"message_agent",
			"Send a follow-up message to a previously spawned child "+
				"agent. Use this to provide additional instructions, "+
				"corrections, or context to a running or completed "+
				"agent. After sending, use wait_agent to collect the "+
				"updated response.",
			func(ctx context.Context, args messageAgentArgs, _ fantasy.ToolCall) (fantasy.ToolResponse, error) {
				if currentChat == nil {
					return fantasy.NewTextErrorResponse("subagent callbacks are not configured"), nil
				}

				targetChatID, err := parseSubagentToolChatID(args.ChatID)
				if err != nil {
					return fantasy.NewTextErrorResponse(err.Error()), nil
				}

				parent := currentChat()
				busyBehavior := SendMessageBusyBehaviorQueue
				if args.Interrupt {
					busyBehavior = SendMessageBusyBehaviorInterrupt
				}
				targetChat, err := p.sendSubagentMessage(
					ctx,
					parent.ID,
					targetChatID,
					args.Message,
					busyBehavior,
				)
				if err != nil {
					return fantasy.NewTextErrorResponse(err.Error()), nil
				}

				return toolJSONResponse(map[string]any{
					"chat_id":     targetChatID.String(),
					"title":       targetChat.Title,
					"status":      string(targetChat.Status),
					"interrupted": args.Interrupt,
				}), nil
			},
		),
		fantasy.NewAgentTool(
			"close_agent",
			"Immediately stop a spawned child agent. Use this to "+
				"cancel a subagent that is stuck, no longer needed, "+
				"or working on the wrong approach.",
			func(ctx context.Context, args closeAgentArgs, _ fantasy.ToolCall) (fantasy.ToolResponse, error) {
				if currentChat == nil {
					return fantasy.NewTextErrorResponse("subagent callbacks are not configured"), nil
				}

				targetChatID, err := parseSubagentToolChatID(args.ChatID)
				if err != nil {
					return fantasy.NewTextErrorResponse(err.Error()), nil
				}

				parent := currentChat()
				targetChat, err := p.closeSubagent(
					ctx,
					parent.ID,
					targetChatID,
				)
				if err != nil {
					return fantasy.NewTextErrorResponse(err.Error()), nil
				}

				return toolJSONResponse(map[string]any{
					"chat_id":    targetChatID.String(),
					"title":      targetChat.Title,
					"terminated": true,
					"status":     string(targetChat.Status),
				}), nil
			},
		),
	}

	// Only include the computer use tool when an Anthropic
	// provider is configured and desktop is enabled.
	if p.isAnthropicConfigured(ctx) && p.isDesktopEnabled(ctx) {
		tools = append(tools, fantasy.NewAgentTool(
			"spawn_computer_use_agent",
			"Spawn a dedicated computer use agent that can see the desktop "+
				"(take screenshots) and interact with it (mouse, keyboard, "+
				"scroll). The agent runs on a model optimized for computer "+
				"use and has the same workspace tools as a standard subagent "+
				"plus the native Anthropic computer tool. Use this for tasks "+
				"that require visual interaction with a desktop GUI (e.g. "+
				"browser automation, GUI testing, visual inspection). After "+
				"spawning, use wait_agent to collect the result.",
			func(ctx context.Context, args spawnComputerUseAgentArgs, _ fantasy.ToolCall) (fantasy.ToolResponse, error) {
				if currentChat == nil {
					return fantasy.NewTextErrorResponse("subagent callbacks are not configured"), nil
				}

				parent := currentChat()
				if parent.ParentChatID.Valid {
					return fantasy.NewTextErrorResponse("delegated chats cannot create child subagents"), nil
				}

				parent, err := p.db.GetChatByID(ctx, parent.ID)
				if err != nil {
					return fantasy.NewTextErrorResponse(err.Error()), nil
				}

				prompt := strings.TrimSpace(args.Prompt)
				if prompt == "" {
					return fantasy.NewTextErrorResponse("prompt is required"), nil
				}

				title := strings.TrimSpace(args.Title)
				if title == "" {
					title = subagentFallbackChatTitle(prompt)
				}

				rootChatID := parent.ID
				if parent.RootChatID.Valid {
					rootChatID = parent.RootChatID.UUID
				}
				if parent.LastModelConfigID == uuid.Nil {
					return fantasy.NewTextErrorResponse("parent chat model config id is required"), nil
				}

				// Create the child chat with Mode set to
				// computer_use. This signals runChat to use the
				// predefined computer use model and include the
				// computer tool.
				childChat, err := p.CreateChat(ctx, CreateOptions{
					OwnerID:     parent.OwnerID,
					WorkspaceID: parent.WorkspaceID,
					ParentChatID: uuid.NullUUID{
						UUID:  parent.ID,
						Valid: true,
					},
					RootChatID: uuid.NullUUID{
						UUID:  rootChatID,
						Valid: true,
					},
					ModelConfigID:      parent.LastModelConfigID,
					Title:              title,
					ChatMode:           database.NullChatMode{ChatMode: database.ChatModeComputerUse, Valid: true},
					SystemPrompt:       computerUseSubagentSystemPrompt + "\n\n" + prompt,
					InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText(prompt)},
				})
				if err != nil {
					return fantasy.NewTextErrorResponse(err.Error()), nil
				}

				return toolJSONResponse(map[string]any{
					"chat_id": childChat.ID.String(),
					"title":   childChat.Title,
					"status":  string(childChat.Status),
				}), nil
			},
		))
	}

	return tools
}

func parseSubagentToolChatID(raw string) (uuid.UUID, error) {
	chatID, err := uuid.Parse(strings.TrimSpace(raw))
	if err != nil {
		return uuid.Nil, xerrors.New("chat_id must be a valid UUID")
	}
	return chatID, nil
}

func (p *Server) createChildSubagentChat(
	ctx context.Context,
	parent database.Chat,
	prompt string,
	title string,
) (database.Chat, error) {
	if parent.ParentChatID.Valid {
		return database.Chat{}, xerrors.New("delegated chats cannot create child subagents")
	}

	prompt = strings.TrimSpace(prompt)
	if prompt == "" {
		return database.Chat{}, xerrors.New("prompt is required")
	}

	title = strings.TrimSpace(title)
	if title == "" {
		title = subagentFallbackChatTitle(prompt)
	}

	rootChatID := parent.ID
	if parent.RootChatID.Valid {
		rootChatID = parent.RootChatID.UUID
	}
	if parent.LastModelConfigID == uuid.Nil {
		return database.Chat{}, xerrors.New("parent chat model config id is required")
	}

	child, err := p.CreateChat(ctx, CreateOptions{
		OwnerID:     parent.OwnerID,
		WorkspaceID: parent.WorkspaceID,
		ParentChatID: uuid.NullUUID{
			UUID:  parent.ID,
			Valid: true,
		},
		RootChatID: uuid.NullUUID{
			UUID:  rootChatID,
			Valid: true,
		},
		ModelConfigID:      parent.LastModelConfigID,
		Title:              title,
		InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText(prompt)},
	})
	if err != nil {
		return database.Chat{}, xerrors.Errorf("create child chat: %w", err)
	}

	return child, nil
}

func (p *Server) sendSubagentMessage(
	ctx context.Context,
	parentChatID uuid.UUID,
	targetChatID uuid.UUID,
	message string,
	busyBehavior SendMessageBusyBehavior,
) (database.Chat, error) {
	message = strings.TrimSpace(message)
	if message == "" {
		return database.Chat{}, xerrors.New("message is required")
	}

	isDescendant, err := isSubagentDescendant(ctx, p.db, parentChatID, targetChatID)
	if err != nil {
		return database.Chat{}, err
	}
	if !isDescendant {
		return database.Chat{}, ErrSubagentNotDescendant
	}

	// Look up the target chat to get the owner for CreatedBy.
	targetChat, err := p.db.GetChatByID(ctx, targetChatID)
	if err != nil {
		return database.Chat{}, xerrors.Errorf("get target chat: %w", err)
	}

	sendResult, err := p.SendMessage(ctx, SendMessageOptions{
		ChatID:       targetChatID,
		CreatedBy:    targetChat.OwnerID,
		Content:      []codersdk.ChatMessagePart{codersdk.ChatMessageText(message)},
		BusyBehavior: busyBehavior,
	})
	if err != nil {
		return database.Chat{}, err
	}

	return sendResult.Chat, nil
}

func (p *Server) awaitSubagentCompletion(
	ctx context.Context,
	parentChatID uuid.UUID,
	targetChatID uuid.UUID,
	timeout time.Duration,
) (database.Chat, string, error) {
	isDescendant, err := isSubagentDescendant(ctx, p.db, parentChatID, targetChatID)
	if err != nil {
		return database.Chat{}, "", err
	}
	if !isDescendant {
		return database.Chat{}, "", ErrSubagentNotDescendant
	}

	// Check immediately before entering the poll loop.
	targetChat, report, done, checkErr := p.checkSubagentCompletion(ctx, targetChatID)
	if checkErr != nil {
		return database.Chat{}, "", checkErr
	}
	if done {
		return handleSubagentDone(targetChat, report)
	}

	if timeout <= 0 {
		timeout = defaultSubagentWaitTimeout
	}
	timer := p.clock.NewTimer(timeout, "chatd", "subagent_await")
	defer timer.Stop()

	// When pubsub is available, subscribe for fast status
	// notifications and use a less aggressive fallback poll.
	// Without pubsub (single-instance / in-memory) fall back
	// to the original 200ms polling.
	pollInterval := subagentAwaitPollInterval
	var notifyCh <-chan struct{}
	if p.pubsub != nil {
		pollInterval = subagentAwaitFallbackPoll
		ch := make(chan struct{}, 1)
		notifyCh = ch
		cancel, subErr := p.pubsub.SubscribeWithErr(
			coderdpubsub.ChatStreamNotifyChannel(targetChatID),
			func(_ context.Context, _ []byte, _ error) {
				// Non-blocking send so we never stall the
				// pubsub dispatch goroutine.
				select {
				case ch <- struct{}{}:
				default:
				}
			},
		)
		if subErr == nil {
			defer cancel()
		} else {
			// Subscription failed; fall back to fast polling.
			pollInterval = subagentAwaitPollInterval
			notifyCh = nil
		}
	}

	ticker := p.clock.NewTicker(pollInterval, "chatd", "subagent_poll")
	defer ticker.Stop()

	for {
		select {
		case <-notifyCh:
		case <-ticker.C:
		case <-timer.C:
			return database.Chat{}, "", xerrors.New("timed out waiting for delegated subagent completion")
		case <-ctx.Done():
			return database.Chat{}, "", ctx.Err()
		}

		targetChat, report, done, checkErr = p.checkSubagentCompletion(ctx, targetChatID)
		if checkErr != nil {
			return database.Chat{}, "", checkErr
		}
		if done {
			return handleSubagentDone(targetChat, report)
		}
	}
}

// handleSubagentDone translates a completed subagent check into the
// appropriate return value, surfacing error-status chats as errors.
func handleSubagentDone(
	chat database.Chat,
	report string,
) (database.Chat, string, error) {
	if chat.Status == database.ChatStatusError {
		reason := strings.TrimSpace(report)
		if reason == "" {
			reason = "agent reached error status"
		}
		return database.Chat{}, "", xerrors.New(reason)
	}
	return chat, report, nil
}

func (p *Server) closeSubagent(
	ctx context.Context,
	parentChatID uuid.UUID,
	targetChatID uuid.UUID,
) (database.Chat, error) {
	isDescendant, err := isSubagentDescendant(ctx, p.db, parentChatID, targetChatID)
	if err != nil {
		return database.Chat{}, err
	}
	if !isDescendant {
		return database.Chat{}, ErrSubagentNotDescendant
	}

	targetChat, err := p.db.GetChatByID(ctx, targetChatID)
	if err != nil {
		return database.Chat{}, xerrors.Errorf("get target chat: %w", err)
	}

	if targetChat.Status == database.ChatStatusWaiting {
		return targetChat, nil
	}

	updatedChat := p.InterruptChat(ctx, targetChat)
	if updatedChat.Status != database.ChatStatusWaiting {
		return database.Chat{}, xerrors.New("set target chat waiting")
	}
	return updatedChat, nil
}

func (p *Server) checkSubagentCompletion(
	ctx context.Context,
	chatID uuid.UUID,
) (database.Chat, string, bool, error) {
	chat, err := p.db.GetChatByID(ctx, chatID)
	if err != nil {
		return database.Chat{}, "", false, xerrors.Errorf("get chat: %w", err)
	}

	if chat.Status == database.ChatStatusPending || chat.Status == database.ChatStatusRunning {
		return database.Chat{}, "", false, nil
	}

	report, err := latestSubagentAssistantMessage(ctx, p.db, chatID)
	if err != nil {
		return database.Chat{}, "", false, err
	}

	return chat, report, true, nil
}

func latestSubagentAssistantMessage(
	ctx context.Context,
	store database.Store,
	chatID uuid.UUID,
) (string, error) {
	messages, err := store.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{
		ChatID:  chatID,
		AfterID: 0,
	})
	if err != nil {
		return "", xerrors.Errorf("get chat messages: %w", err)
	}

	sort.Slice(messages, func(i, j int) bool {
		if messages[i].CreatedAt.Equal(messages[j].CreatedAt) {
			return messages[i].ID < messages[j].ID
		}
		return messages[i].CreatedAt.Before(messages[j].CreatedAt)
	})

	for i := len(messages) - 1; i >= 0; i-- {
		message := messages[i]
		if message.Role != database.ChatMessageRoleAssistant ||
			message.Visibility == database.ChatMessageVisibilityModel {
			continue
		}

		content, parseErr := chatprompt.ParseContent(message)
		if parseErr != nil {
			continue
		}
		text := strings.TrimSpace(contentBlocksToText(content))
		if text == "" {
			continue
		}
		return text, nil
	}

	return "", nil
}

// isSubagentDescendant reports whether targetChatID is a descendant
// of ancestorChatID by walking up the parent chain from the target.
// This is O(depth) DB queries instead of O(nodes) BFS.
func isSubagentDescendant(
	ctx context.Context,
	store database.Store,
	ancestorChatID uuid.UUID,
	targetChatID uuid.UUID,
) (bool, error) {
	if ancestorChatID == targetChatID {
		return false, nil
	}

	currentID := targetChatID
	visited := map[uuid.UUID]struct{}{} // cycle protection
	for {
		if _, seen := visited[currentID]; seen {
			return false, nil
		}
		visited[currentID] = struct{}{}

		chat, err := store.GetChatByID(ctx, currentID)
		if err != nil {
			if xerrors.Is(err, sql.ErrNoRows) {
				return false, nil // chain broken; not a confirmed descendant
			}
			return false, xerrors.Errorf("get chat %s: %w", currentID, err)
		}
		if !chat.ParentChatID.Valid {
			return false, nil // reached root without finding ancestor
		}
		if chat.ParentChatID.UUID == ancestorChatID {
			return true, nil
		}
		currentID = chat.ParentChatID.UUID
	}
}

func subagentFallbackChatTitle(message string) string {
	const maxWords = 6
	const maxRunes = 80

	words := strings.Fields(message)
	if len(words) == 0 {
		return "New Chat"
	}

	truncated := false
	if len(words) > maxWords {
		words = words[:maxWords]
		truncated = true
	}

	title := strings.Join(words, " ")
	if truncated {
		title += "..."
	}

	return subagentTruncateRunes(title, maxRunes)
}

func subagentTruncateRunes(value string, maxRunes int) string {
	if maxRunes <= 0 {
		return ""
	}

	runes := []rune(value)
	if len(runes) <= maxRunes {
		return value
	}

	return string(runes[:maxRunes])
}

func toolJSONResponse(result map[string]any) fantasy.ToolResponse {
	data, err := json.Marshal(result)
	if err != nil {
		return fantasy.NewTextResponse("{}")
	}
	return fantasy.NewTextResponse(string(data))
}