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.ResolveNode(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.ResolveNode(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.ResolveNode(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 }