iop/apps/edge/internal/opsconsole/status.go

package opsconsole

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

	edgeservice "iop/apps/edge/internal/service"
	iop "iop/proto/gen/iop"
)

func BuildNodeCommandRequest(adapter, targetName, sessionID string, timeoutSec int) (*iop.NodeCommandRequest, string) {
	req := edgeservice.BuildUsageStatusRequest(adapter, targetName, sessionID, timeoutSec)
	return req, req.GetRequestId()
}

func StatusWaitTimeout(req *iop.NodeCommandRequest) time.Duration {
	return edgeservice.StatusWaitTimeout(req)
}

func SendStatus(ctx context.Context, edgeSvc *edgeservice.Service, out io.Writer, target *TargetState) error {
	snap, err := edgeSvc.ResolveNodeSnapshot(target.NodeRef)
	if err != nil {
		return err
	}
	sessionID := edgeservice.NormalizeSessionID(target.SessionID)
	fmt.Fprintf(out, "[edge] sent command=status node=%s adapter=%s target=%s session=%s\n", snap.Label, target.Adapter, target.Target, sessionID)

	result, err := edgeSvc.UsageStatus(ctx, edgeservice.UsageStatusRequest{
		NodeRef:    target.NodeRef,
		Adapter:    target.Adapter,
		Target:     target.Target,
		SessionID:  target.SessionID,
		TimeoutSec: target.TimeoutSec,
	})
	if err != nil {
		return err
	}

	FormatUsageStatus(out, result.NodeLabel, result.Target, result.SessionID, result.UsageStatus)
	return nil
}

// SendCapabilities dispatches a CAPABILITIES command and renders the result.
func SendCapabilities(ctx context.Context, edgeSvc *edgeservice.Service, out io.Writer, target *TargetState) error {
	return sendNodeCommandView(ctx, out, target, "capabilities",
		edgeSvc.Capabilities,
	)
}

// SendSessionList dispatches a SESSION_LIST command and renders the result.
func SendSessionList(ctx context.Context, edgeSvc *edgeservice.Service, out io.Writer, target *TargetState) error {
	return sendNodeCommandView(ctx, out, target, "sessions",
		edgeSvc.SessionList,
	)
}

// SendTransportStatus dispatches a TRANSPORT_STATUS command and renders the result.
func SendTransportStatus(ctx context.Context, edgeSvc *edgeservice.Service, out io.Writer, target *TargetState) error {
	return sendNodeCommandView(ctx, out, target, "transport",
		edgeSvc.TransportStatus,
	)
}

type nodeCommandFunc func(context.Context, edgeservice.NodeCommandRequestSpec) (edgeservice.NodeCommandView, error)

func sendNodeCommandView(ctx context.Context, out io.Writer, target *TargetState, label string, send nodeCommandFunc) error {
	view, err := send(ctx, edgeservice.NodeCommandRequestSpec{
		NodeRef:    target.NodeRef,
		Adapter:    target.Adapter,
		Target:     target.Target,
		SessionID:  target.SessionID,
		TimeoutSec: target.TimeoutSec,
	})
	if err != nil {
		return err
	}
	FormatNodeCommandView(out, label, view)
	return nil
}

// FormatNodeCommandView renders a node command result map in a stable,
// sorted order so output is predictable for tests and humans.
func FormatNodeCommandView(out io.Writer, label string, view edgeservice.NodeCommandView) {
	fmt.Fprintf(out, "[node-%s-%s] target=%s session=%s\n", view.NodeLabel, label, view.Target, view.SessionID)
	if len(view.Result) == 0 {
		fmt.Fprintf(out, "no %s payload returned\n", label)
		return
	}
	keys := make([]string, 0, len(view.Result))
	for k := range view.Result {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	for _, k := range keys {
		fmt.Fprintf(out, "  %s = %s\n", k, view.Result[k])
	}
}

func FormatUsageStatus(out io.Writer, nodeAlias, targetName, sessionID string, status *iop.AgentUsageStatus) {
	fmt.Fprintf(out, "[node-%s-status] target=%s session=%s\n", nodeAlias, targetName, sessionID)

	if status == nil {
		fmt.Fprintln(out, "no usage status provided")
		return
	}

	hasParsedLimits := false
	dailyLabel := usageStatusLabel(status.GetMetadata(), "daily_label", "Daily limit")
	if status.GetDailyLimit() != "" {
		fmt.Fprintf(out, "%s: %s remaining%s\n", dailyLabel, status.GetDailyLimit(), resetSuffix(status.GetDailyResetTime()))
		hasParsedLimits = true
	}
	weeklyLabel := usageStatusLabel(status.GetMetadata(), "weekly_label", "Weekly limit")
	if status.GetWeeklyLimit() != "" {
		fmt.Fprintf(out, "%s: %s remaining%s\n", weeklyLabel, status.GetWeeklyLimit(), resetSuffix(status.GetWeeklyResetTime()))
		hasParsedLimits = true
	}
	metadata := status.GetMetadata()
	if used := usageStatusMetadata(metadata, "used_percent"); used != "" {
		fmt.Fprintf(out, "Used: %s\n", used)
	}
	if limit := usageStatusMetadata(metadata, "usage_limit"); limit != "" {
		fmt.Fprintf(out, "Usage limit: %s\n", limit)
	}
	if formatGeminiModelUsage(out, metadata) {
		hasParsedLimits = true
	}

	if !hasParsedLimits {
		if status.GetRawOutput() == "" {
			fmt.Fprintln(out, "raw output did not include parsed limits and was empty")
		} else {
			fmt.Fprintln(out, "raw output did not include parsed limits:")
			lines := strings.Split(status.GetRawOutput(), "\n")
			for i, line := range lines {
				if i >= 5 {
					fmt.Fprintln(out, "...")
					break
				}
				fmt.Fprintln(out, line)
			}
		}
	}
}

func usageStatusLabel(metadata map[string]string, key, fallback string) string {
	if metadata != nil {
		if label, ok := metadata[key]; ok {
			return label
		}
	}
	return fallback
}

func usageStatusMetadata(metadata map[string]string, key string) string {
	if metadata == nil {
		return ""
	}
	return metadata[key]
}

func formatGeminiModelUsage(out io.Writer, metadata map[string]string) bool {
	countRaw := usageStatusMetadata(metadata, "model_usage_count")
	if countRaw == "" {
		return false
	}
	count, err := strconv.Atoi(countRaw)
	if err != nil || count <= 0 {
		return false
	}

	printed := false
	for i := 0; i < count; i++ {
		prefix := fmt.Sprintf("model_usage_%d", i)
		name := usageStatusMetadata(metadata, prefix+"_name")
		if name == "" {
			continue
		}
		used := usageStatusMetadata(metadata, prefix+"_used_percent")
		reset := usageStatusMetadata(metadata, prefix+"_reset")
		if !printed {
			fmt.Fprintln(out, "Model usage:")
			printed = true
		}
		switch {
		case used != "" && reset != "":
			fmt.Fprintf(out, "  %s: %s used (resets %s)\n", name, used, reset)
		case used != "":
			fmt.Fprintf(out, "  %s: %s used\n", name, used)
		case reset != "":
			fmt.Fprintf(out, "  %s: resets %s\n", name, reset)
		default:
			fmt.Fprintf(out, "  %s\n", name)
		}
	}
	return printed
}

func resetSuffix(resetTime string) string {
	if resetTime == "" {
		return ""
	}
	return " (resets " + resetTime + ")"
}