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

package transport

import (
	"testing"

	toki "git.toki-labs.com/toki/common-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 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)
	}
}