kandev/apps/backend/internal/agent/runtime/agentctl/control.go

// Package agentctl provides a client for communicating with agentctl.
// This file contains the ControlClient for the agentctl control server API.
package client

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"sync"
	"time"

	"github.com/kandev/kandev/internal/common/logger"
	"go.uber.org/zap"
)

// ControlClient is a client for the agentctl control server API.
// It manages creation and deletion of agent instances.
// Used by both Docker and standalone runtimes.
type ControlClient struct {
	mu         sync.Mutex
	baseURL    string
	httpClient *http.Client
	logger     *logger.Logger
	authToken  string
}

// McpServerConfig holds configuration for an MCP server.
type McpServerConfig struct {
	Name    string            `json:"name"`
	URL     string            `json:"url,omitempty"`
	Type    string            `json:"type,omitempty"`
	Command string            `json:"command,omitempty"`
	Args    []string          `json:"args,omitempty"`
	Env     map[string]string `json:"env,omitempty"`
	Headers map[string]string `json:"headers,omitempty"`
}

// CreateInstanceRequest contains the parameters for creating a new agent instance.
type CreateInstanceRequest struct {
	ID                 string            `json:"id,omitempty"`
	WorkspacePath      string            `json:"workspace_path"`
	AgentCommand       string            `json:"agent_command,omitempty"`
	Protocol           string            `json:"protocol,omitempty"`       // Protocol adapter to use (acp, rest, mcp, codex)
	AgentType          string            `json:"agent_type,omitempty"`     // Agent type ID for debug file naming (e.g., "codex", "auggie")
	WorkspaceFlag      string            `json:"workspace_flag,omitempty"` // CLI flag for workspace path (e.g., "--workspace-root")
	Env                map[string]string `json:"env,omitempty"`
	AutoStart          bool              `json:"auto_start,omitempty"`
	McpServers         []McpServerConfig `json:"mcp_servers,omitempty"`
	SessionID          string            `json:"session_id,omitempty"`           // Task session ID for MCP tool calls
	TaskID             string            `json:"task_id,omitempty"`              // Task ID for MCP plan tool calls (server-side injection)
	DisableAskQuestion bool              `json:"disable_ask_question,omitempty"` // Disable ask_user_question MCP tool (TUI agents)
	AssumeMcpSse       bool              `json:"assume_mcp_sse,omitempty"`       // Assume agent supports SSE MCP servers
	AssumeMcpHttp      bool              `json:"assume_mcp_http,omitempty"`      // Assume agent supports HTTP MCP servers
	McpMode            string            `json:"mcp_mode,omitempty"`             // MCP tool mode: "task" (default) or "config"
	// RequiresProcessKill forces agentctl to kill the agent's process group
	// (not just close stdin) on shutdown. Required for agents whose runtime
	// keeps child processes (e.g. MCP servers) alive when stdin closes —
	// notably opencode acp.
	RequiresProcessKill bool `json:"requires_process_kill,omitempty"`

	// BaseBranches maps RepositoryName → base branch ref for the task's
	// per-repo diff stats. The empty key "" applies to the root /
	// single-repo tracker. Empty map disables the override and falls back
	// to the hardcoded origin/main → master priority list inside agentctl.
	BaseBranches map[string]string `json:"base_branches,omitempty"`
}

// CreateInstanceResponse contains the result of creating a new agent instance.
type CreateInstanceResponse struct {
	ID   string `json:"id"`
	Port int    `json:"port"`
}

// InstanceInfo contains information about an agent instance.
type InstanceInfo struct {
	ID            string            `json:"id"`
	Port          int               `json:"port"`
	Status        string            `json:"status"`
	WorkspacePath string            `json:"workspace_path"`
	AgentCommand  string            `json:"agent_command"`
	Env           map[string]string `json:"env,omitempty"`
	CreatedAt     time.Time         `json:"created_at"`
}

// ControlClientOption configures optional ControlClient settings.
type ControlClientOption func(*ControlClient)

// WithControlAuthToken sets the Bearer token for authenticating control requests.
func WithControlAuthToken(token string) ControlClientOption {
	return func(c *ControlClient) {
		c.authToken = token
	}
}

// NewControlClient creates a new ControlClient for the agentctl control server.
func NewControlClient(host string, port int, log *logger.Logger, opts ...ControlClientOption) *ControlClient {
	c := &ControlClient{
		baseURL: fmt.Sprintf("http://%s:%d", host, port),
		httpClient: &http.Client{
			Timeout: 30 * time.Second,
		},
		logger: log.WithFields(zap.String("component", "agentctl-control")),
	}
	for _, opt := range opts {
		opt(c)
	}
	if c.authToken != "" {
		c.httpClient.Transport = &authTransport{token: c.authToken}
	}
	return c
}

// AuthToken returns the current auth token. Used to propagate the token
// from the ControlClient to per-instance Clients after a handshake.
func (c *ControlClient) AuthToken() string {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.authToken
}

// SetAuthToken sets the Bearer token for future requests.
// Used after a successful Handshake to authenticate subsequent calls.
func (c *ControlClient) SetAuthToken(token string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.authToken = token
	c.httpClient.Transport = &authTransport{token: token}
}

// Handshake performs the bootstrap handshake with agentctl.
// It sends the one-time nonce and receives the self-generated auth token.
// On success, the token is automatically set for future requests.
func (c *ControlClient) Handshake(ctx context.Context, nonce string) (string, error) {
	body, err := json.Marshal(map[string]string{"nonce": nonce})
	if err != nil {
		return "", fmt.Errorf("failed to marshal handshake request: %w", err)
	}

	req, err := http.NewRequestWithContext(ctx, "POST", c.baseURL+"/auth/handshake", bytes.NewReader(body))
	if err != nil {
		return "", err
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return "", fmt.Errorf("handshake failed: %w", err)
	}
	defer func() { _ = resp.Body.Close() }()

	if resp.StatusCode != http.StatusOK {
		var errResp struct {
			Error string `json:"error"`
		}
		if decErr := json.NewDecoder(resp.Body).Decode(&errResp); decErr == nil && errResp.Error != "" {
			return "", fmt.Errorf("handshake rejected: %s (status %d)", errResp.Error, resp.StatusCode)
		}
		return "", fmt.Errorf("handshake failed: status %d", resp.StatusCode)
	}

	var result struct {
		Token string `json:"token"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return "", fmt.Errorf("failed to decode handshake response: %w", err)
	}
	if result.Token == "" {
		return "", fmt.Errorf("handshake returned empty token")
	}

	// Automatically set the token for future requests
	c.SetAuthToken(result.Token)
	c.logger.Info("bootstrap handshake completed")

	return result.Token, nil
}

// Health checks if the agentctl control server is healthy.
func (c *ControlClient) Health(ctx context.Context) error {
	req, err := http.NewRequestWithContext(ctx, "GET", c.baseURL+"/health", nil)
	if err != nil {
		return err
	}

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return err
	}
	defer func() {
		_ = resp.Body.Close()
	}()

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("health check failed: %d", resp.StatusCode)
	}
	return nil
}

// CreateInstance creates a new agent instance.
func (c *ControlClient) CreateInstance(ctx context.Context, req *CreateInstanceRequest) (*CreateInstanceResponse, error) {
	body, err := json.Marshal(req)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal request: %w", err)
	}

	httpReq, err := http.NewRequestWithContext(ctx, "POST", c.baseURL+"/api/v1/instances", bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	httpReq.Header.Set("Content-Type", "application/json")

	resp, err := c.httpClient.Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("failed to create instance: %w", err)
	}
	defer func() { _ = resp.Body.Close() }()

	if resp.StatusCode != http.StatusCreated && resp.StatusCode != http.StatusOK {
		var errResp struct {
			Error string `json:"error"`
		}
		if err := json.NewDecoder(resp.Body).Decode(&errResp); err != nil {
			return nil, fmt.Errorf("failed to decode error response: %w", err)
		}
		return nil, fmt.Errorf("failed to create instance: %s (status %d)", errResp.Error, resp.StatusCode)
	}

	var result CreateInstanceResponse
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return nil, fmt.Errorf("failed to decode response: %w", err)
	}

	c.logger.Info("created agent instance",
		zap.String("instance_id", result.ID),
		zap.Int("port", result.Port))

	return &result, nil
}

// DeleteInstance stops and removes an agent instance.
func (c *ControlClient) DeleteInstance(ctx context.Context, instanceID string) error {
	req, err := http.NewRequestWithContext(ctx, "DELETE", c.baseURL+"/api/v1/instances/"+instanceID, nil)
	if err != nil {
		return err
	}

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return fmt.Errorf("failed to delete instance: %w", err)
	}
	defer func() {
		_ = resp.Body.Close()
	}()

	if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusNoContent {
		var errResp struct {
			Error string `json:"error"`
		}
		if err := json.NewDecoder(resp.Body).Decode(&errResp); err != nil {
			return fmt.Errorf("failed to decode error response: %w", err)
		}
		return fmt.Errorf("failed to delete instance: %s (status %d)", errResp.Error, resp.StatusCode)
	}

	c.logger.Info("deleted agent instance", zap.String("instance_id", instanceID))
	return nil
}

// GetInstance gets information about a specific instance.
func (c *ControlClient) GetInstance(ctx context.Context, instanceID string) (*InstanceInfo, error) {
	req, err := http.NewRequestWithContext(ctx, "GET", c.baseURL+"/api/v1/instances/"+instanceID, nil)
	if err != nil {
		return nil, err
	}

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("failed to get instance: %w", err)
	}
	defer func() { _ = resp.Body.Close() }()

	if resp.StatusCode == http.StatusNotFound {
		return nil, fmt.Errorf("instance %q not found", instanceID)
	}
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("failed to get instance: status %d", resp.StatusCode)
	}

	var info InstanceInfo
	if err := json.NewDecoder(resp.Body).Decode(&info); err != nil {
		return nil, fmt.Errorf("failed to decode response: %w", err)
	}
	return &info, nil
}

// ListInstances lists all running agent instances.
func (c *ControlClient) ListInstances(ctx context.Context) ([]*InstanceInfo, error) {
	req, err := http.NewRequestWithContext(ctx, "GET", c.baseURL+"/api/v1/instances", nil)
	if err != nil {
		return nil, err
	}

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("failed to list instances: %w", err)
	}
	defer func() { _ = resp.Body.Close() }()

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("failed to list instances: status %d", resp.StatusCode)
	}

	var result struct {
		Instances []*InstanceInfo `json:"instances"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return nil, fmt.Errorf("failed to decode response: %w", err)
	}
	return result.Instances, nil
}