iop/apps/node/internal/node/command_handler.go

package node

import (
	"context"
	"fmt"
	"sort"
	"strconv"
	"strings"
	"time"

	"go.uber.org/zap"

	"iop/apps/node/internal/runtime"
	"iop/apps/node/internal/transport"
	iop "iop/proto/gen/iop"
)

// OnCommandRequest handles an incoming node command from a transport Session.
func (n *Node) OnCommandRequest(ctx context.Context, sess *transport.Session, req *iop.NodeCommandRequest) (*iop.NodeCommandResponse, error) {
	n.logger.Info("command request",
		zap.String("request_id", req.GetRequestId()),
		zap.String("type", req.GetType().String()),
		zap.String("adapter", req.GetAdapter()),
		zap.String("target", req.GetTarget()),
	)

	cmdType, ok := protoCommandTypeToDomain(req.GetType())
	if !ok {
		return n.commandErrorResponse(req, fmt.Sprintf("node: unsupported command type %q", req.GetType().String())), nil
	}

	execCtx := ctx
	if req.GetTimeoutSec() > 0 {
		var cancel context.CancelFunc
		execCtx, cancel = context.WithTimeout(ctx, time.Duration(req.GetTimeoutSec())*time.Second)
		defer cancel()
	}

	switch cmdType {
	case runtime.CommandTypeCapabilities:
		return n.handleCapabilitiesCommand(execCtx, req), nil
	case runtime.CommandTypeTransportStatus:
		return n.handleTransportStatusCommand(sess, req), nil
	default:
		return n.dispatchAdapterCommand(execCtx, req), nil
	}
}

func (n *Node) handleCapabilitiesCommand(ctx context.Context, req *iop.NodeCommandRequest) *iop.NodeCommandResponse {
	adapter, err := n.router.LookupAdapter(req.GetAdapter())
	if err != nil {
		return n.commandErrorResponse(req, fmt.Sprintf("node: %s", err.Error()))
	}
	caps, err := adapter.Capabilities(ctx)
	if err != nil {
		return n.commandErrorResponse(req, err.Error())
	}

	targets := append([]string(nil), caps.Targets...)
	providerStatus := caps.ProviderStatus
	providerDetail := ""

	if prober, ok := adapter.(runtime.ProviderProber); ok {
		probeRes, err := prober.ProbeProvider(ctx, req.GetTarget())
		if err != nil {
			providerStatus = runtime.ProviderStatusUnavailable
			providerDetail = err.Error()
		} else {
			providerStatus = probeRes.Status
			providerDetail = probeRes.Detail
			if len(probeRes.Targets) > 0 {
				targets = append([]string(nil), probeRes.Targets...)
			}
		}
	}

	sort.Strings(targets)

	n.adapterGatesMu.Lock()
	gate, ok := n.adapterGates[req.GetAdapter()]
	n.adapterGatesMu.Unlock()
	inFlight := 0
	queued := 0
	if ok {
		inFlight = gate.activeCount()
	}

	result := map[string]string{
		"adapter":         caps.AdapterName,
		"instance_key":    caps.InstanceKey,
		"targets":         strings.Join(targets, ","),
		"max_concurrency": strconv.Itoa(caps.MaxConcurrency),
		"provider_status": string(runtime.NormalizeProviderStatus(providerStatus)),
		"capacity":        strconv.Itoa(caps.MaxConcurrency),
		"in_flight":       strconv.Itoa(inFlight),
		"queued":          strconv.Itoa(queued),
	}
	if providerDetail != "" {
		result["provider_detail"] = providerDetail
	}

	providerSnapshot := &iop.ProviderSnapshot{
		Adapter:  req.GetAdapter(),
		Status:   string(runtime.NormalizeProviderStatus(providerStatus)),
		Capacity: int32(caps.MaxConcurrency),
		InFlight: int32(inFlight),
		Queued:   int32(queued),
	}

	return &iop.NodeCommandResponse{
		RequestId:         req.GetRequestId(),
		Type:              req.GetType(),
		Adapter:           req.GetAdapter(),
		Target:            req.GetTarget(),
		SessionId:         req.GetSessionId(),
		Result:            result,
		ProviderSnapshots: []*iop.ProviderSnapshot{providerSnapshot},
	}
}

func sessionConnected(sess *transport.Session) (alive bool) {
	if sess == nil {
		return false
	}
	defer func() {
		if r := recover(); r != nil {
			alive = false
		}
	}()
	return sess.IsAlive()
}

func (n *Node) handleTransportStatusCommand(sess *transport.Session, req *iop.NodeCommandRequest) *iop.NodeCommandResponse {
	connected := "false"
	state := "disconnected"
	if sessionConnected(sess) {
		connected = "true"
		state = "connected"
	}
	result := map[string]string{
		"node_id":    n.nodeID,
		"connected":  connected,
		"state":      state,
		"adapter":    req.GetAdapter(),
		"target":     req.GetTarget(),
		"session_id": normalizeSessionID(req.GetSessionId()),
	}
	return &iop.NodeCommandResponse{
		RequestId: req.GetRequestId(),
		Type:      req.GetType(),
		Adapter:   req.GetAdapter(),
		Target:    req.GetTarget(),
		SessionId: normalizeSessionID(req.GetSessionId()),
		Result:    result,
	}
}

func (n *Node) dispatchAdapterCommand(ctx context.Context, req *iop.NodeCommandRequest) *iop.NodeCommandResponse {
	adapter, err := n.router.LookupAdapter(req.GetAdapter())
	if err != nil {
		return n.commandErrorResponse(req, fmt.Sprintf("node: %s", err.Error()))
	}
	handler, ok := adapter.(runtime.CommandHandler)
	if !ok {
		return n.commandErrorResponse(req, fmt.Sprintf("node: adapter %q does not support commands", req.GetAdapter()))
	}
	domainReq := toDomainCommandRequest(req)
	domainResp, err := handler.HandleCommand(ctx, domainReq)
	if err != nil {
		return n.commandErrorResponse(req, err.Error())
	}
	return toProtoCommandResponse(domainResp)
}

func (n *Node) commandErrorResponse(req *iop.NodeCommandRequest, msg string) *iop.NodeCommandResponse {
	return &iop.NodeCommandResponse{
		RequestId: req.GetRequestId(),
		Type:      req.GetType(),
		Adapter:   req.GetAdapter(),
		Target:    req.GetTarget(),
		SessionId: req.GetSessionId(),
		Error:     msg,
	}
}

// protoCommandTypeToDomain maps a proto NodeCommandType to its runtime
// CommandType. The second return is false for UNSPECIFIED or unknown values
// so callers can reject unsupported commands before dispatching.
func protoCommandTypeToDomain(t iop.NodeCommandType) (runtime.CommandType, bool) {
	switch t {
	case iop.NodeCommandType_NODE_COMMAND_TYPE_USAGE_STATUS:
		return runtime.CommandTypeUsageStatus, true
	case iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES:
		return runtime.CommandTypeCapabilities, true
	case iop.NodeCommandType_NODE_COMMAND_TYPE_SESSION_LIST:
		return runtime.CommandTypeSessionList, true
	case iop.NodeCommandType_NODE_COMMAND_TYPE_TRANSPORT_STATUS:
		return runtime.CommandTypeTransportStatus, true
	case iop.NodeCommandType_NODE_COMMAND_TYPE_OLLAMA_API:
		return runtime.CommandTypeOllamaAPI, true
	default:
		return "", false
	}
}

func domainCommandTypeToProto(t runtime.CommandType) iop.NodeCommandType {
	switch t {
	case runtime.CommandTypeUsageStatus:
		return iop.NodeCommandType_NODE_COMMAND_TYPE_USAGE_STATUS
	case runtime.CommandTypeCapabilities:
		return iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES
	case runtime.CommandTypeSessionList:
		return iop.NodeCommandType_NODE_COMMAND_TYPE_SESSION_LIST
	case runtime.CommandTypeTransportStatus:
		return iop.NodeCommandType_NODE_COMMAND_TYPE_TRANSPORT_STATUS
	case runtime.CommandTypeOllamaAPI:
		return iop.NodeCommandType_NODE_COMMAND_TYPE_OLLAMA_API
	default:
		return iop.NodeCommandType_NODE_COMMAND_TYPE_UNSPECIFIED
	}
}

func toDomainCommandRequest(req *iop.NodeCommandRequest) runtime.CommandRequest {
	cmdType, _ := protoCommandTypeToDomain(req.GetType())
	return runtime.CommandRequest{
		RequestID:  req.GetRequestId(),
		Type:       cmdType,
		Adapter:    req.GetAdapter(),
		Target:     req.GetTarget(),
		SessionID:  normalizeSessionID(req.GetSessionId()),
		TimeoutSec: int(req.GetTimeoutSec()),
		Metadata:   req.GetMetadata(),
	}
}

func toProtoCommandResponse(resp runtime.CommandResponse) *iop.NodeCommandResponse {
	out := &iop.NodeCommandResponse{
		RequestId: resp.RequestID,
		Type:      domainCommandTypeToProto(resp.Type),
		Adapter:   resp.Adapter,
		Target:    resp.Target,
		SessionId: resp.SessionID,
		Result:    resp.Result,
	}
	if resp.UsageStatus != nil {
		out.UsageStatus = &iop.AgentUsageStatus{
			RawOutput:       resp.UsageStatus.RawOutput,
			DailyLimit:      resp.UsageStatus.DailyLimit,
			DailyResetTime:  resp.UsageStatus.DailyResetTime,
			WeeklyLimit:     resp.UsageStatus.WeeklyLimit,
			WeeklyResetTime: resp.UsageStatus.WeeklyResetTime,
			Metadata:        resp.UsageStatus.Metadata,
		}
	}
	return out
}