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

package service

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

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

	iop "iop/proto/gen/iop"
)

type UsageStatusRequest struct {
	NodeRef    string
	Adapter    string
	Target     string
	SessionID  string
	TimeoutSec int
}

// UsageStatusView is the surface-neutral DTO for the usage-status command.
// proto types are retained for the embedded AgentUsageStatus payload.
type UsageStatusView struct {
	NodeID      string
	NodeLabel   string
	Adapter     string
	Target      string
	SessionID   string
	UsageStatus *iop.AgentUsageStatus
}

// UsageStatusResult is retained as an alias for backward compatibility.
type UsageStatusResult = UsageStatusView

func (s *Service) UsageStatus(_ context.Context, req UsageStatusRequest) (UsageStatusResult, error) {
	entry, err := s.ResolveDispatchReady(req.NodeRef)
	if err != nil {
		return UsageStatusResult{}, err
	}

	commandReq := BuildUsageStatusRequest(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 UsageStatusResult{}, fmt.Errorf("transport error: %w", err)
	}
	if resp.GetError() != "" {
		return UsageStatusResult{}, fmt.Errorf("node reported error: %s", resp.GetError())
	}

	return UsageStatusResult{
		NodeID:      entry.NodeID,
		NodeLabel:   nodeLabel(entry),
		Adapter:     commandReq.GetAdapter(),
		Target:      commandReq.GetTarget(),
		SessionID:   commandReq.GetSessionId(),
		UsageStatus: resp.GetUsageStatus(),
	}, nil
}

func BuildUsageStatusRequest(adapter, targetName, sessionID string, timeoutSec int) *iop.NodeCommandRequest {
	return buildNodeCommandRequest(iop.NodeCommandType_NODE_COMMAND_TYPE_USAGE_STATUS, "status", adapter, targetName, sessionID, timeoutSec)
}

// 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, session_list, transport_status). UsageStatus keeps
// its own typed request shape.
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")
}

// SessionList dispatches a SESSION_LIST node command and returns the result map.
func (s *Service) SessionList(_ context.Context, req NodeCommandRequestSpec) (NodeCommandView, error) {
	return s.sendNodeCommand(req, iop.NodeCommandType_NODE_COMMAND_TYPE_SESSION_LIST, "sessions")
}

// 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())
	}
	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
}