kandev/apps/backend/internal/agentctl/client/agent.go

package client

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

	"github.com/gorilla/websocket"
	"github.com/kandev/kandev/internal/agentctl/tracing"
	"github.com/kandev/kandev/internal/agentctl/types"
	"github.com/kandev/kandev/internal/agentctl/types/streams"
	v1 "github.com/kandev/kandev/pkg/api/v1"
	ws "github.com/kandev/kandev/pkg/websocket"
	"go.uber.org/zap"
)

// Re-export stream types for convenience.
type (
	AgentEvent                = streams.AgentEvent
	PermissionNotification    = streams.PermissionNotification
	PermissionOption          = streams.PermissionOption
	PermissionRespondRequest  = streams.PermissionRespondRequest
	PermissionRespondResponse = streams.PermissionRespondResponse
)

// AgentInfo contains information about the connected agent.
type AgentInfo struct {
	Name    string `json:"name"`
	Version string `json:"version"`
}

// InitializeResponse from agentctl
type InitializeResponse struct {
	Success   bool       `json:"success"`
	AgentInfo *AgentInfo `json:"agent_info,omitempty"`
	Error     string     `json:"error,omitempty"`
}

// Initialize sends the ACP initialize request via the agent WebSocket stream.
func (c *Client) Initialize(ctx context.Context, clientName, clientVersion string) (*AgentInfo, error) {
	payload := struct {
		ClientName    string `json:"client_name"`
		ClientVersion string `json:"client_version"`
	}{
		ClientName:    clientName,
		ClientVersion: clientVersion,
	}

	resp, err := c.sendStreamRequest(ctx, "agent.initialize", payload)
	if err != nil {
		return nil, fmt.Errorf("initialize request failed: %w", err)
	}

	if resp.Type == ws.MessageTypeError {
		var errPayload ws.ErrorPayload
		if err := resp.ParsePayload(&errPayload); err != nil {
			return nil, fmt.Errorf("initialize failed: unable to parse error")
		}
		return nil, fmt.Errorf("initialize failed: %s", errPayload.Message)
	}

	var result InitializeResponse
	if err := resp.ParsePayload(&result); err != nil {
		return nil, fmt.Errorf("failed to parse initialize response: %w", err)
	}
	if !result.Success {
		return nil, fmt.Errorf("initialize failed: %s", result.Error)
	}
	return result.AgentInfo, nil
}

// NewSessionResponse from agentctl
type NewSessionResponse struct {
	Success   bool   `json:"success"`
	SessionID string `json:"session_id,omitempty"`
	Error     string `json:"error,omitempty"`
}

// createSessionRequest sends a session creation request and parses the response.
// Used by both NewSession and ResetSession which share the same payload/response format.
func (c *Client) createSessionRequest(ctx context.Context, action, cwd string, mcpServers []types.McpServer) (string, error) {
	payload := struct {
		Cwd        string            `json:"cwd"`
		McpServers []types.McpServer `json:"mcp_servers,omitempty"`
	}{Cwd: cwd, McpServers: mcpServers}

	resp, err := c.sendStreamRequest(ctx, action, payload)
	if err != nil {
		return "", fmt.Errorf("%s request failed: %w", action, err)
	}

	if resp.Type == ws.MessageTypeError {
		var errPayload ws.ErrorPayload
		if err := resp.ParsePayload(&errPayload); err != nil {
			return "", fmt.Errorf("%s failed: unable to parse error", action)
		}
		return "", fmt.Errorf("%s failed: %s", action, errPayload.Message)
	}

	var result NewSessionResponse
	if err := resp.ParsePayload(&result); err != nil {
		return "", fmt.Errorf("failed to parse %s response: %w", action, err)
	}
	if !result.Success {
		return "", fmt.Errorf("%s failed: %s", action, result.Error)
	}
	return result.SessionID, nil
}

// NewSession creates a new ACP session via the agent WebSocket stream.
func (c *Client) NewSession(ctx context.Context, cwd string, mcpServers []types.McpServer) (string, error) {
	return c.createSessionRequest(ctx, "agent.session.new", cwd, mcpServers)
}

// ResetSession creates a new session on the same connection, resetting context without
// restarting the subprocess. Returns the new session ID or an error if not supported.
func (c *Client) ResetSession(ctx context.Context, cwd string, mcpServers []types.McpServer) (string, error) {
	return c.createSessionRequest(ctx, "agent.session.reset", cwd, mcpServers)
}

// LoadSession resumes an existing ACP session via the agent WebSocket stream.
// mcpServers are forwarded to the agentctl handler so agents that receive MCP configs
// via the protocol (e.g. Auggie) can reconnect to MCP servers on the new instance.
func (c *Client) LoadSession(ctx context.Context, sessionID string, mcpServers []types.McpServer) error {
	payload := struct {
		SessionID  string            `json:"session_id"`
		McpServers []types.McpServer `json:"mcp_servers,omitempty"`
	}{SessionID: sessionID, McpServers: mcpServers}

	resp, err := c.sendStreamRequest(ctx, "agent.session.load", payload)
	if err != nil {
		return fmt.Errorf("load session request failed: %w", err)
	}

	if resp.Type == ws.MessageTypeError {
		var errPayload ws.ErrorPayload
		if err := resp.ParsePayload(&errPayload); err != nil {
			return fmt.Errorf("load session failed: unable to parse error")
		}
		return fmt.Errorf("load session failed: %s", errPayload.Message)
	}

	var result struct {
		Success bool   `json:"success"`
		Error   string `json:"error,omitempty"`
	}
	if err := resp.ParsePayload(&result); err != nil {
		return fmt.Errorf("failed to parse load session response: %w", err)
	}
	if !result.Success {
		return fmt.Errorf("load session failed: %s", result.Error)
	}
	return nil
}

// SetMode changes the agent's session mode via the agent WebSocket stream.
func (c *Client) SetMode(ctx context.Context, sessionID, modeID string) error {
	payload := struct {
		SessionID string `json:"session_id"`
		ModeID    string `json:"mode_id"`
	}{SessionID: sessionID, ModeID: modeID}

	resp, err := c.sendStreamRequest(ctx, "agent.session.set_mode", payload)
	if err != nil {
		return fmt.Errorf("set mode request failed: %w", err)
	}

	if resp.Type == ws.MessageTypeError {
		var errPayload ws.ErrorPayload
		if err := resp.ParsePayload(&errPayload); err != nil {
			return fmt.Errorf("set mode failed: unable to parse error")
		}
		return fmt.Errorf("set mode failed: %s", errPayload.Message)
	}

	return nil
}

// SetModel changes the agent's model via the agent WebSocket stream.
func (c *Client) SetModel(ctx context.Context, modelID string) error {
	payload := struct {
		ModelID string `json:"model_id"`
	}{ModelID: modelID}

	resp, err := c.sendStreamRequest(ctx, "agent.session.set_model", payload)
	if err != nil {
		return fmt.Errorf("set model request failed: %w", err)
	}

	if resp.Type == ws.MessageTypeError {
		var errPayload ws.ErrorPayload
		if err := resp.ParsePayload(&errPayload); err != nil {
			return fmt.Errorf("set model failed: unable to parse error")
		}
		return fmt.Errorf("set model failed: %s", errPayload.Message)
	}

	return nil
}

// SetConfigOption sets a session config option via the agent WebSocket stream.
func (c *Client) SetConfigOption(ctx context.Context, configID, value string) error {
	payload := struct {
		ConfigID string `json:"config_id"`
		Value    string `json:"value"`
	}{ConfigID: configID, Value: value}

	resp, err := c.sendStreamRequest(ctx, "agent.session.set_config_option", payload)
	if err != nil {
		return fmt.Errorf("set config option request failed: %w", err)
	}

	if resp.Type == ws.MessageTypeError {
		var errPayload ws.ErrorPayload
		if err := resp.ParsePayload(&errPayload); err != nil {
			return fmt.Errorf("set config option failed: unable to parse error")
		}
		return fmt.Errorf("set config option failed: %s", errPayload.Message)
	}

	return nil
}

// Authenticate triggers authentication for a given auth method via the agent WebSocket stream.
func (c *Client) Authenticate(ctx context.Context, methodID string) error {
	payload := struct {
		MethodID string `json:"method_id"`
	}{MethodID: methodID}

	resp, err := c.sendStreamRequest(ctx, "agent.session.authenticate", payload)
	if err != nil {
		return fmt.Errorf("authenticate request failed: %w", err)
	}

	if resp.Type == ws.MessageTypeError {
		var errPayload ws.ErrorPayload
		if err := resp.ParsePayload(&errPayload); err != nil {
			return fmt.Errorf("authenticate failed: unable to parse error")
		}
		return fmt.Errorf("authenticate failed: %s", errPayload.Message)
	}

	return nil
}

// Prompt sends a fire-and-forget prompt to the agent via the agent WebSocket stream.
// The server returns an accepted response immediately; completion is signaled via stream events.
// Attachments (images) are passed to the agent if provided.
func (c *Client) Prompt(ctx context.Context, text string, attachments []v1.MessageAttachment) error {
	payload := struct {
		Text        string                 `json:"text"`
		Attachments []v1.MessageAttachment `json:"attachments,omitempty"`
	}{Text: text, Attachments: attachments}

	resp, err := c.sendStreamRequest(ctx, "agent.prompt", payload)
	if err != nil {
		return fmt.Errorf("prompt request failed: %w", err)
	}

	if resp.Type == ws.MessageTypeError {
		var errPayload ws.ErrorPayload
		if err := resp.ParsePayload(&errPayload); err != nil {
			return fmt.Errorf("prompt failed: unable to parse error")
		}
		return fmt.Errorf("prompt failed: %s", errPayload.Message)
	}

	var result struct {
		Success bool   `json:"success"`
		Error   string `json:"error,omitempty"`
	}
	if err := resp.ParsePayload(&result); err != nil {
		return fmt.Errorf("failed to parse prompt response: %w", err)
	}
	if !result.Success {
		c.logger.Warn("prompt returned failure response", zap.String("error", result.Error))
		return fmt.Errorf("prompt failed: %s", result.Error)
	}
	return nil
}

// Note: AgentEvent, PermissionNotification, PermissionOption, and
// PermissionRespondRequest are re-exported from streams package at the top of this file.

// MCPHandler is the interface for handling MCP requests from agentctl.
type MCPHandler interface {
	// Dispatch handles an MCP request and returns a response.
	Dispatch(ctx context.Context, msg *ws.Message) (*ws.Message, error)
}

// StreamUpdates opens a WebSocket connection for streaming agent events.
// Events include message chunks, reasoning, tool calls, plan updates, and completion/error.
// If mcpHandler is provided, MCP requests from agentctl will be dispatched to it and responses sent back.
// If onDisconnect is provided, it is called when the WebSocket read goroutine exits (e.g., on error or close).
func (c *Client) StreamUpdates(ctx context.Context, handler func(AgentEvent), mcpHandler MCPHandler, onDisconnect func(err error)) error {
	wsURL := "ws" + c.baseURL[4:] + "/api/v1/agent/stream"

	conn, _, err := websocket.DefaultDialer.DialContext(ctx, wsURL, c.wsAuthHeaders())
	if err != nil {
		return fmt.Errorf("failed to connect to updates stream: %w", err)
	}

	c.mu.Lock()
	c.agentStreamConn = conn
	c.mu.Unlock()

	c.logger.Info("connected to updates stream", zap.String("url", wsURL))

	// writeMessage uses the shared stream write mutex for thread-safe writes
	writeMessage := func(data []byte) error {
		c.streamWriteMu.Lock()
		defer c.streamWriteMu.Unlock()
		return conn.WriteMessage(websocket.TextMessage, data)
	}

	go c.readUpdatesStream(ctx, conn, handler, mcpHandler, onDisconnect, writeMessage)

	return nil
}

// readUpdatesStream is the read loop for the agent updates WebSocket stream.
func (c *Client) readUpdatesStream(
	ctx context.Context,
	conn *websocket.Conn,
	handler func(AgentEvent),
	mcpHandler MCPHandler,
	onDisconnect func(err error),
	writeMessage func([]byte) error,
) {
	var lastErr error
	defer func() {
		// Clean up pending requests before signaling disconnect
		c.cleanupPendingRequests()

		c.mu.Lock()
		c.agentStreamConn = nil
		c.mu.Unlock()
		if err := conn.Close(); err != nil {
			c.logger.Debug("failed to close updates websocket", zap.Error(err))
		}
		// Signal disconnect to caller
		if onDisconnect != nil {
			onDisconnect(lastErr)
		}
	}()

	for {
		_, message, err := conn.ReadMessage()
		if err != nil {
			if websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) {
				c.logger.Info("updates stream closed normally")
				// Normal close — don't report as disconnect error
			} else {
				c.logger.Debug("updates stream error", zap.Error(err))
				lastErr = err
			}
			return
		}

		// Try to parse as ws.Message to check message type
		var wsMsg ws.Message
		if err := json.Unmarshal(message, &wsMsg); err == nil {
			// Check if this is a response/error to a pending request
			if (wsMsg.Type == ws.MessageTypeResponse || wsMsg.Type == ws.MessageTypeError) && c.resolvePendingRequest(&wsMsg) {
				continue
			}

			// This is an MCP request - dispatch it
			if wsMsg.Type == ws.MessageTypeRequest {
				if mcpHandler != nil {
					sessionID, pendingID := extractMCPRequestCorrelation(wsMsg.Payload)
					c.logger.Debug("received MCP request from agent stream",
						zap.String("request_id", wsMsg.ID),
						zap.String("action", wsMsg.Action),
						zap.String("session_id", sessionID),
						zap.String("pending_id", pendingID))
					go c.dispatchMCPRequest(ctx, wsMsg, mcpHandler, writeMessage)
				}
				continue
			}
		}

		// Not a ws.Message request/response, parse as agent event
		var event AgentEvent
		if err := json.Unmarshal(message, &event); err != nil {
			c.logger.Warn("failed to parse agent event", zap.Error(err))
			continue
		}

		tracing.TraceAgentEvent(ctx, event.Type, event.SessionID, c.executionID, message)
		handler(event)
	}
}

// dispatchMCPRequest handles an incoming MCP request message, dispatches it to the handler,
// and writes the response back over the stream.
func (c *Client) dispatchMCPRequest(ctx context.Context, msg ws.Message, mcpHandler MCPHandler, writeMessage func([]byte) error) {
	start := time.Now()
	sessionID, pendingID := extractMCPRequestCorrelation(msg.Payload)
	ctx, span := tracing.TraceMCPDispatch(ctx, msg.Action, msg.ID, c.executionID)
	defer span.End()

	c.logger.Debug("dispatching MCP request",
		zap.String("request_id", msg.ID),
		zap.String("action", msg.Action),
		zap.String("session_id", sessionID),
		zap.String("pending_id", pendingID))

	resp, err := mcpHandler.Dispatch(ctx, &msg)
	if err != nil {
		c.logger.Error("MCP dispatch error",
			zap.String("request_id", msg.ID),
			zap.String("action", msg.Action),
			zap.String("session_id", sessionID),
			zap.String("pending_id", pendingID),
			zap.Duration("duration", time.Since(start)),
			zap.Error(err))
		tracing.TraceMCPResponse(span, err)
		resp, _ = ws.NewError(msg.ID, msg.Action, ws.ErrorCodeInternalError, err.Error(), nil)
	} else {
		c.logger.Debug("MCP request dispatched",
			zap.String("request_id", msg.ID),
			zap.String("action", msg.Action),
			zap.String("session_id", sessionID),
			zap.String("pending_id", pendingID),
			zap.Bool("has_response", resp != nil),
			zap.Duration("duration", time.Since(start)))
		tracing.TraceMCPResponse(span, nil)
	}
	if resp != nil {
		data, err := json.Marshal(resp)
		if err != nil {
			c.logger.Error("failed to marshal MCP response",
				zap.String("request_id", msg.ID),
				zap.String("action", msg.Action),
				zap.String("session_id", sessionID),
				zap.String("pending_id", pendingID),
				zap.Error(err))
			return
		}
		writeStarted := time.Now()
		if err := writeMessage(data); err != nil {
			c.logger.Warn("failed to write MCP response",
				zap.String("request_id", msg.ID),
				zap.String("action", msg.Action),
				zap.String("session_id", sessionID),
				zap.String("pending_id", pendingID),
				zap.Int("response_bytes", len(data)),
				zap.Duration("write_duration", time.Since(writeStarted)),
				zap.Error(err))
			return
		}
		c.logger.Debug("wrote MCP response to agent stream",
			zap.String("request_id", msg.ID),
			zap.String("action", msg.Action),
			zap.String("session_id", sessionID),
			zap.String("pending_id", pendingID),
			zap.String("response_type", string(resp.Type)),
			zap.Int("response_bytes", len(data)),
			zap.Duration("write_duration", time.Since(writeStarted)),
			zap.Duration("total_duration", time.Since(start)))
	}
}

func extractMCPRequestCorrelation(payload json.RawMessage) (sessionID string, pendingID string) {
	if len(payload) == 0 {
		return "", ""
	}

	var payloadMap map[string]interface{}
	if err := json.Unmarshal(payload, &payloadMap); err != nil {
		return "", ""
	}

	if v, ok := payloadMap["session_id"].(string); ok {
		sessionID = v
	}
	if v, ok := payloadMap["pending_id"].(string); ok {
		pendingID = v
	}
	return sessionID, pendingID
}

// CloseUpdatesStream closes the agent events stream connection.
func (c *Client) CloseUpdatesStream() {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.agentStreamConn != nil {
		if err := c.agentStreamConn.Close(); err != nil {
			c.logger.Debug("failed to close agent events stream", zap.Error(err))
		}
		c.agentStreamConn = nil
	}
}

// CancelResponse is the response from the agentctl cancel endpoint.
type CancelResponse struct {
	Success bool   `json:"success"`
	Error   string `json:"error,omitempty"`
}

// Cancel interrupts the current agent turn via the agent WebSocket stream.
func (c *Client) Cancel(ctx context.Context) error {
	resp, err := c.sendStreamRequest(ctx, "agent.cancel", nil)
	if err != nil {
		return fmt.Errorf("cancel request failed: %w", err)
	}

	if resp.Type == ws.MessageTypeError {
		var errPayload ws.ErrorPayload
		if err := resp.ParsePayload(&errPayload); err != nil {
			return fmt.Errorf("cancel failed: unable to parse error")
		}
		return fmt.Errorf("cancel failed: %s", errPayload.Message)
	}

	var result CancelResponse
	if err := resp.ParsePayload(&result); err != nil {
		return fmt.Errorf("failed to parse cancel response: %w", err)
	}
	if !result.Success {
		return fmt.Errorf("cancel failed: %s", result.Error)
	}
	return nil
}

// GetAgentStderr returns recent stderr lines from the agent process via the agent WebSocket stream.
func (c *Client) GetAgentStderr(ctx context.Context) ([]string, error) {
	resp, err := c.sendStreamRequest(ctx, "agent.stderr", nil)
	if err != nil {
		return nil, fmt.Errorf("agent stderr request failed: %w", err)
	}

	if resp.Type == ws.MessageTypeError {
		var errPayload ws.ErrorPayload
		if err := resp.ParsePayload(&errPayload); err != nil {
			return nil, fmt.Errorf("agent stderr failed: unable to parse error")
		}
		return nil, fmt.Errorf("agent stderr failed: %s", errPayload.Message)
	}

	var result struct {
		Lines []string `json:"lines"`
	}
	if err := resp.ParsePayload(&result); err != nil {
		return nil, fmt.Errorf("failed to parse agent stderr response: %w", err)
	}
	return result.Lines, nil
}

// RespondToPermission sends a response to a permission request via the agent WebSocket stream.
func (c *Client) RespondToPermission(ctx context.Context, pendingID, optionID string, cancelled bool) error {
	payload := PermissionRespondRequest{
		PendingID: pendingID,
		OptionID:  optionID,
		Cancelled: cancelled,
	}

	resp, err := c.sendStreamRequest(ctx, "agent.permissions.respond", payload)
	if err != nil {
		return fmt.Errorf("permission response request failed: %w", err)
	}

	if resp.Type == ws.MessageTypeError {
		var errPayload ws.ErrorPayload
		if err := resp.ParsePayload(&errPayload); err != nil {
			return fmt.Errorf("permission response failed: unable to parse error")
		}
		return fmt.Errorf("permission response failed: %s", errPayload.Message)
	}

	var result PermissionRespondResponse
	if err := resp.ParsePayload(&result); err != nil {
		return fmt.Errorf("failed to parse permission response: %w", err)
	}
	if !result.Success {
		return fmt.Errorf("permission response failed: %s", result.Error)
	}
	return nil
}