iop/apps/edge/internal/transport/server_test.go

package transport

import (
	"testing"

	toki "git.toki-labs.com/toki/proto-socket/go"
	"google.golang.org/protobuf/proto"

	edgenode "iop/apps/edge/internal/node"
	"iop/packages/config"
	iop "iop/proto/gen/iop"
)

func TestEdgeParserMap_NodeCommandResponse(t *testing.T) {
	parsers := edgeParserMap()
	original := &iop.NodeCommandResponse{
		Type:      iop.NodeCommandType_NODE_COMMAND_TYPE_USAGE_STATUS,
		Adapter:   "cli",
		Target:    "codex",
		SessionId: "default",
		UsageStatus: &iop.AgentUsageStatus{
			RawOutput: "test",
		},
	}
	payload, err := proto.Marshal(original)
	if err != nil {
		t.Fatalf("marshal: %v", err)
	}

	key := toki.TypeNameOf(original)
	parser, ok := parsers[key]
	if !ok {
		t.Fatalf("parser not found for key: %s", key)
	}

	parsed, err := parser(payload)
	if err != nil {
		t.Fatalf("parse: %v", err)
	}
	got := parsed.(*iop.NodeCommandResponse)
	if got.GetType() != original.GetType() ||
		got.GetAdapter() != original.GetAdapter() ||
		got.GetTarget() != original.GetTarget() ||
		got.GetSessionId() != original.GetSessionId() ||
		got.GetUsageStatus().GetRawOutput() != original.GetUsageStatus().GetRawOutput() {
		t.Fatalf("unexpected node command response: %+v", got)
	}
}

func TestEdgeParserMap_NodeCommandResponse_NewTypes(t *testing.T) {
	parsers := edgeParserMap()
	cases := []struct {
		cmdType iop.NodeCommandType
		result  map[string]string
	}{
		{iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES, map[string]string{"targets": "codex,gemini"}},
		{iop.NodeCommandType_NODE_COMMAND_TYPE_SESSION_LIST, map[string]string{"count": "2"}},
		{iop.NodeCommandType_NODE_COMMAND_TYPE_TRANSPORT_STATUS, map[string]string{"connected": "true"}},
	}
	for _, tc := range cases {
		original := &iop.NodeCommandResponse{
			RequestId: "req-1",
			Type:      tc.cmdType,
			Adapter:   "cli",
			Target:    "codex",
			SessionId: "default",
			Result:    tc.result,
		}
		payload, err := proto.Marshal(original)
		if err != nil {
			t.Fatalf("marshal %v: %v", tc.cmdType, err)
		}
		parser, ok := parsers[toki.TypeNameOf(original)]
		if !ok {
			t.Fatalf("parser not found for NodeCommandResponse")
		}
		parsed, err := parser(payload)
		if err != nil {
			t.Fatalf("parse %v: %v", tc.cmdType, err)
		}
		got := parsed.(*iop.NodeCommandResponse)
		if got.GetType() != tc.cmdType {
			t.Fatalf("type: got %v want %v", got.GetType(), tc.cmdType)
		}
		for k, v := range tc.result {
			if got.GetResult()[k] != v {
				t.Fatalf("result[%q]: got %q want %q", k, got.GetResult()[k], v)
			}
		}
	}
}

func TestEdgeParserMap_EdgeNodeEvent(t *testing.T) {
	parsers := edgeParserMap()
	original := &iop.EdgeNodeEvent{
		Type:   "node.disconnected",
		Source: "edge",
		NodeId: "node-1",
		Alias:  "local-node",
		Reason: "transport_closed",
	}
	payload, err := proto.Marshal(original)
	if err != nil {
		t.Fatalf("marshal: %v", err)
	}

	key := toki.TypeNameOf(original)
	parser, ok := parsers[key]
	if !ok {
		t.Fatalf("parser not found for key: %s", key)
	}

	parsed, err := parser(payload)
	if err != nil {
		t.Fatalf("parse: %v", err)
	}
	got := parsed.(*iop.EdgeNodeEvent)
	if got.GetType() != original.GetType() ||
		got.GetSource() != original.GetSource() ||
		got.GetNodeId() != original.GetNodeId() ||
		got.GetAlias() != original.GetAlias() ||
		got.GetReason() != original.GetReason() {
		t.Fatalf("unexpected edge node event: %+v", got)
	}
}

func TestServerEnrichesRunEventNodeAlias(t *testing.T) {
	registry := edgenode.NewRegistry()
	registry.Register(&edgenode.NodeEntry{NodeID: "node-1", Alias: "alias-1"})
	s := &Server{registry: registry}

	event := &iop.RunEvent{RunId: "run-1", Type: "start", NodeId: "node-1"}
	s.enrichRunEvent(event)
	if event.GetNodeAlias() != "alias-1" {
		t.Fatalf("expected alias-1, got %q", event.GetNodeAlias())
	}

	preset := &iop.RunEvent{RunId: "run-2", Type: "start", NodeId: "node-1", NodeAlias: "preset"}
	s.enrichRunEvent(preset)
	if preset.GetNodeAlias() != "preset" {
		t.Fatalf("expected preset to be preserved, got %q", preset.GetNodeAlias())
	}

	unknown := &iop.RunEvent{RunId: "run-3", Type: "start", NodeId: "node-x"}
	s.enrichRunEvent(unknown)
	if unknown.GetNodeAlias() != "" {
		t.Fatalf("expected empty alias for unknown node, got %q", unknown.GetNodeAlias())
	}
}

func findCLIAdapter(payload *iop.NodeConfigPayload) *iop.AdapterConfig {
	for _, a := range payload.GetAdapters() {
		if a.GetType() == "cli" {
			return a
		}
	}
	return nil
}

func TestBuildConfigPayload_CLIOneof(t *testing.T) {
	rec := &edgenode.NodeRecord{
		ID:    "node-local",
		Alias: "local",
		Adapters: config.AdaptersConf{
			CLI: config.CLIConf{
				Enabled: true,
				Profiles: map[string]config.CLIProfileConf{
					"codex": {
						Command:               "codex",
						Args:                  []string{"--foo", "--bar"},
						Env:                   []string{"KEY=val"},
						Persistent:            true,
						Terminal:              true,
						ResponseIdleTimeoutMS: 1500,
						StartupIdleTimeoutMS:  300,
						OutputFormat:          "codex-json",
					},
				},
			},
		},
	}

	payload, err := edgenode.BuildConfigPayload(rec)
	if err != nil {
		t.Fatalf("build config payload: %v", err)
	}

	cli := findCLIAdapter(payload)
	if cli == nil {
		t.Fatal("expected cli adapter")
	}
	if cli.GetSettings() != nil {
		t.Fatalf("expected settings to be nil for cli adapter, got %v", cli.GetSettings())
	}
	cliCfg := cli.GetCli()
	if cliCfg == nil {
		t.Fatal("expected typed cli config")
	}
	prof, ok := cliCfg.GetProfiles()["codex"]
	if !ok {
		t.Fatal("expected codex profile")
	}
	if prof.GetCommand() != "codex" {
		t.Fatalf("command: %q", prof.GetCommand())
	}
	if got := prof.GetArgs(); len(got) != 2 || got[0] != "--foo" || got[1] != "--bar" {
		t.Fatalf("args: %#v", got)
	}
	if got := prof.GetEnv(); len(got) != 1 || got[0] != "KEY=val" {
		t.Fatalf("env: %#v", got)
	}
	if !prof.GetPersistent() {
		t.Fatal("persistent")
	}
	if !prof.GetTerminal() {
		t.Fatal("terminal")
	}
	if prof.GetResponseIdleTimeoutMs() != 1500 {
		t.Fatalf("response_idle_timeout_ms: %d", prof.GetResponseIdleTimeoutMs())
	}
	if prof.GetStartupIdleTimeoutMs() != 300 {
		t.Fatalf("startup_idle_timeout_ms: %d", prof.GetStartupIdleTimeoutMs())
	}
	if prof.GetOutputFormat() != "codex-json" {
		t.Fatalf("output_format: %q", prof.GetOutputFormat())
	}
}

func TestBuildConfigPayload_CLICompletionMarker(t *testing.T) {
	rec := &edgenode.NodeRecord{
		ID:    "node-local",
		Alias: "local",
		Adapters: config.AdaptersConf{
			CLI: config.CLIConf{
				Enabled: true,
				Profiles: map[string]config.CLIProfileConf{
					"codex": {
						Command:    "codex",
						Persistent: true,
						CompletionMarker: config.CompletionMarkerConf{
							Line:  "[[DONE]]",
							Regex: `^\[\[DONE\]\]$`,
						},
					},
				},
			},
		},
	}

	payload, err := edgenode.BuildConfigPayload(rec)
	if err != nil {
		t.Fatalf("build config payload: %v", err)
	}
	cli := findCLIAdapter(payload)
	if cli == nil {
		t.Fatal("expected cli adapter")
	}
	prof := cli.GetCli().GetProfiles()["codex"]
	marker := prof.GetCompletionMarker()
	if marker == nil {
		t.Fatal("expected completion_marker")
	}
	if marker.GetLine() != "[[DONE]]" {
		t.Fatalf("line: %q", marker.GetLine())
	}
	if marker.GetRegex() != `^\[\[DONE\]\]$` {
		t.Fatalf("regex: %q", marker.GetRegex())
	}
}

func TestBuildConfigPayload_OllamaVllmOneof(t *testing.T) {
	rec := &edgenode.NodeRecord{
		Adapters: config.AdaptersConf{
			Ollama: config.OllamaConf{Enabled: true, BaseURL: "http://localhost:11434"},
			Vllm:   config.VllmConf{Enabled: true, Endpoint: "http://localhost:8000"},
		},
	}
	payload, err := edgenode.BuildConfigPayload(rec)
	if err != nil {
		t.Fatalf("build: %v", err)
	}
	var ollama, vllm *iop.AdapterConfig
	for _, a := range payload.GetAdapters() {
		switch a.GetType() {
		case "ollama":
			ollama = a
		case "vllm":
			vllm = a
		}
	}
	if ollama == nil || ollama.GetOllama().GetBaseUrl() != "http://localhost:11434" {
		t.Fatalf("ollama: %+v", ollama)
	}
	if vllm == nil || vllm.GetVllm().GetEndpoint() != "http://localhost:8000" {
		t.Fatalf("vllm: %+v", vllm)
	}
}