iop/apps/edge/internal/service/node_command.go

package service

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

	toki "git.toki-labs.com/toki/proto-socket/go"

	runtime "iop/packages/go/execution"
	iop "iop/proto/gen/iop"
)

// BuildNodeCommandRequest builds a NodeCommandRequest for any supported type.
// idPrefix is used to namespace request_id (e.g. "status", "caps", "sessions").
func BuildNodeCommandRequest(cmdType iop.NodeCommandType, idPrefix, adapter, targetName, sessionID string, timeoutSec int) *iop.NodeCommandRequest {
	return buildNodeCommandRequest(cmdType, idPrefix, adapter, targetName, sessionID, timeoutSec)
}

func buildNodeCommandRequest(cmdType iop.NodeCommandType, idPrefix, adapter, targetName, sessionID string, timeoutSec int) *iop.NodeCommandRequest {
	if idPrefix == "" {
		idPrefix = "cmd"
	}
	return &iop.NodeCommandRequest{
		RequestId:  fmt.Sprintf("%s-%d", idPrefix, time.Now().UnixNano()),
		Type:       cmdType,
		Adapter:    adapter,
		Target:     targetName,
		SessionId:  NormalizeSessionID(sessionID),
		TimeoutSec: int32(normalizeTimeoutSec(timeoutSec)),
	}
}

func StatusWaitTimeout(req *iop.NodeCommandRequest) time.Duration {
	return time.Duration(normalizeTimeoutSec(int(req.GetTimeoutSec()))+5) * time.Second
}

// NodeCommandRequestSpec is the surface-neutral input for ops console node
// commands (capabilities, transport_status, ollama_api).
type NodeCommandRequestSpec struct {
	NodeRef    string
	Adapter    string
	Target     string
	SessionID  string
	TimeoutSec int
}

// NodeCommandView is the surface-neutral result for non-usage-status node
// commands. Result mirrors the proto map and is empty when the node returned
// no payload.
type NodeCommandView struct {
	NodeID            string
	NodeLabel         string
	Adapter           string
	Target            string
	SessionID         string
	Type              iop.NodeCommandType
	Result            map[string]string
	ProviderSnapshots []*iop.ProviderSnapshot
}

// Capabilities dispatches a CAPABILITIES node command and returns the result map.
func (s *Service) Capabilities(_ context.Context, req NodeCommandRequestSpec) (NodeCommandView, error) {
	return s.sendNodeCommand(req, iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES, "caps")
}

// TransportStatus dispatches a TRANSPORT_STATUS node command and returns the result map.
func (s *Service) TransportStatus(_ context.Context, req NodeCommandRequestSpec) (NodeCommandView, error) {
	return s.sendNodeCommand(req, iop.NodeCommandType_NODE_COMMAND_TYPE_TRANSPORT_STATUS, "transport")
}

type OllamaAPIRequest struct {
	NodeRef    string
	Adapter    string
	Target     string
	Method     string
	Path       string
	Body       string
	TimeoutSec int
}

type OllamaAPIView struct {
	StatusCode  int
	ContentType string
	Body        string
}

func (s *Service) OllamaAPI(_ context.Context, req OllamaAPIRequest) (OllamaAPIView, error) {
	entry, err := s.ResolveDispatchReady(req.NodeRef)
	if err != nil {
		return OllamaAPIView{}, err
	}
	commandReq := buildNodeCommandRequest(iop.NodeCommandType_NODE_COMMAND_TYPE_OLLAMA_API, "ollama", req.Adapter, req.Target, "", req.TimeoutSec)
	commandReq.Metadata = map[string]string{
		"ollama_method": req.Method,
		"ollama_path":   req.Path,
		"ollama_body":   req.Body,
	}
	resp, err := toki.SendRequestTyped[*iop.NodeCommandRequest, *iop.NodeCommandResponse](
		&entry.Client.Communicator,
		commandReq,
		StatusWaitTimeout(commandReq),
	)
	if err != nil {
		return OllamaAPIView{}, fmt.Errorf("transport error: %w", err)
	}
	if resp.GetError() != "" {
		return OllamaAPIView{}, fmt.Errorf("node reported error: %s", resp.GetError())
	}
	statusCode, _ := strconv.Atoi(resp.GetResult()["status_code"])
	if statusCode == 0 {
		statusCode = 200
	}
	return OllamaAPIView{
		StatusCode:  statusCode,
		ContentType: resp.GetResult()["content_type"],
		Body:        resp.GetResult()["body"],
	}, nil
}

func (s *Service) sendNodeCommand(req NodeCommandRequestSpec, cmdType iop.NodeCommandType, idPrefix string) (NodeCommandView, error) {
	entry, err := s.ResolveDispatchReady(req.NodeRef)
	if err != nil {
		return NodeCommandView{}, err
	}
	commandReq := buildNodeCommandRequest(cmdType, idPrefix, req.Adapter, req.Target, req.SessionID, req.TimeoutSec)
	resp, err := toki.SendRequestTyped[*iop.NodeCommandRequest, *iop.NodeCommandResponse](
		&entry.Client.Communicator,
		commandReq,
		StatusWaitTimeout(commandReq),
	)
	if err != nil {
		return NodeCommandView{}, fmt.Errorf("transport error: %w", err)
	}
	if resp.GetError() != "" {
		return NodeCommandView{}, fmt.Errorf("node reported error: %s", resp.GetError())
	}
	if cmdType == iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES && s.queue != nil {
		if evidence, ok := capabilitiesProbeEvidenceFromResponse(commandReq, resp); ok {
			s.queue.applyProviderProbeEvidence(
				entry.NodeID,
				entry.ConnectionGeneration,
				evidence.adapter,
				evidence.target,
				evidence.status,
				evidence.sequence,
				func() bool {
					return s.registry != nil && s.registry.IsCurrentOwnerGeneration(entry.NodeID, entry.ConnectionGeneration)
				},
			)
		}
	}
	return NodeCommandView{
		NodeID:            entry.NodeID,
		NodeLabel:         nodeLabel(entry),
		Adapter:           commandReq.GetAdapter(),
		Target:            commandReq.GetTarget(),
		SessionID:         commandReq.GetSessionId(),
		Type:              resp.GetType(),
		Result:            resp.GetResult(),
		ProviderSnapshots: resp.GetProviderSnapshots(),
	}, nil
}

type capabilitiesProbeEvidence struct {
	adapter  string
	target   string
	status   runtime.ProviderStatus
	sequence uint64
}

// capabilitiesProbeEvidenceFromResponse accepts only the stable, exact binding
// emitted by the Node CAPABILITIES probe. Older Nodes omit the sequence and are
// harmless no-ops. Empty/malformed identity, response-envelope mismatch, and
// non-baseline status values also fail closed.
func capabilitiesProbeEvidenceFromResponse(req *iop.NodeCommandRequest, resp *iop.NodeCommandResponse) (capabilitiesProbeEvidence, bool) {
	if req == nil || resp == nil || req.GetType() != iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES ||
		resp.GetType() != iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES {
		return capabilitiesProbeEvidence{}, false
	}
	adapter := strings.TrimSpace(req.GetAdapter())
	target := strings.TrimSpace(req.GetTarget())
	if adapter == "" || target == "" || resp.GetAdapter() != adapter || resp.GetTarget() != target {
		return capabilitiesProbeEvidence{}, false
	}
	result := resp.GetResult()
	if strings.TrimSpace(result["adapter_key"]) != adapter || strings.TrimSpace(result["target"]) != target {
		return capabilitiesProbeEvidence{}, false
	}
	sequence, err := strconv.ParseUint(result["health_observation_seq"], 10, 64)
	if err != nil || sequence == 0 {
		return capabilitiesProbeEvidence{}, false
	}
	status := runtime.ProviderStatus(strings.TrimSpace(result["provider_status"]))
	if normalized := runtime.NormalizeProviderStatus(status); normalized != status {
		return capabilitiesProbeEvidence{}, false
	}
	return capabilitiesProbeEvidence{adapter: adapter, target: target, status: status, sequence: sequence}, true
}