llmgw/internal/web/connect_test.go

package web

import (
	"net/http"
	"net/http/httptest"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"
)

func TestEachToolGetsItsConfigurationWithTheKey(t *testing.T) {
	c := connectFor("continue", "https://gw.example", "qwen2.5-coder-7b", "sk-llmgw-abc", ModelFacts{})
	if !strings.Contains(c.Files[0].Body, "apiBase: https://gw.example/v1") || !strings.Contains(c.Files[0].Body, "apiKey: sk-llmgw-abc") ||
		c.Files[0].Path != "~/.continue/assistants/llmgw.yaml" {
		t.Errorf("continue = %+v", c.Files)
	}
	// The key goes in the environment, not the URL, where it would be kept in
	// shell history and proxy logs.
	if strings.Contains(strings.Split(c.Command, "|")[0], "sk-llmgw-abc") {
		t.Errorf("the key is in the URL: %s", c.Command)
	}
	x := connectFor("codex", "https://gw.example", "qwen2.5-coder-7b", "sk-llmgw-abc", ModelFacts{})
	if strings.Contains(x.Files[0].Body, "sk-llmgw-abc") || !strings.Contains(x.Files[0].Body, `env_key = "LLMGW_API_KEY"`) {
		t.Errorf("codex writes the key into its file: %s", x.Files[0].Body)
	}
	cl := connectFor("cline", "https://gw.example", "m", "sk-llmgw-abc", ModelFacts{})
	if len(cl.Fields) != 4 || cl.Fields[1].Value != "https://gw.example/v1" || cl.Command != "" {
		t.Errorf("cline = %+v", cl)
	}
}

// runScript fetches a setup script and runs it with HOME in a temporary
// directory, as a person's shell would.
func runScript(t *testing.T, tool string, env ...string) string {
	t.Helper()
	p := &AccountPages{}
	w := httptest.NewRecorder()
	r := httptest.NewRequest("GET", "http://gw.example/setup/"+tool+"?model=qwen2.5-coder-7b", nil)
	r.SetPathValue("tool", tool)
	p.setupScript(w, r)
	if w.Code != http.StatusOK {
		t.Fatalf("%d %s", w.Code, w.Body.String())
	}
	home := t.TempDir()
	cmd := exec.Command("sh")
	cmd.Stdin = strings.NewReader(w.Body.String())
	cmd.Env = append([]string{"HOME=" + home, "PATH=" + os.Getenv("PATH")}, env...)
	if out, err := cmd.CombinedOutput(); err != nil {
		t.Fatalf("%v: %s", err, out)
	}
	return home
}

func TestTheContinueScriptWritesItsOwnFileWithTheKey(t *testing.T) {
	home := runScript(t, "continue", "LLMGW_KEY=sk-llmgw-xyz")
	b, err := os.ReadFile(filepath.Join(home, ".continue", "assistants", "llmgw.yaml"))
	if err != nil {
		t.Fatal(err)
	}
	if !strings.Contains(string(b), "apiKey: sk-llmgw-xyz") || strings.Contains(string(b), "__LLMGW_KEY__") {
		t.Errorf("wrote %s", b)
	}
	if fi, _ := os.Stat(filepath.Join(home, ".continue", "assistants", "llmgw.yaml")); fi.Mode().Perm()&0o077 != 0 {
		t.Errorf("a file holding a key is readable by others: %v", fi.Mode())
	}
}

func TestTheCodexScriptKeepsTheRestOfTheConfig(t *testing.T) {
	// Run twice over a config with the person's own settings: theirs stay,
	// and ours appear once.
	p := &AccountPages{}
	w := httptest.NewRecorder()
	r := httptest.NewRequest("GET", "http://gw.example/setup/codex?model=qwen2.5-coder-7b", nil)
	r.SetPathValue("tool", "codex")
	p.setupScript(w, r)
	home := t.TempDir()
	if err := os.MkdirAll(filepath.Join(home, ".codex"), 0o755); err != nil {
		t.Fatal(err)
	}
	mine := "model = \"o3\"\n\n[mcp_servers.docs]\ncommand = \"docs-mcp\"\n"
	cfg := filepath.Join(home, ".codex", "config.toml")
	if err := os.WriteFile(cfg, []byte(mine), 0o644); err != nil {
		t.Fatal(err)
	}
	for range 2 {
		cmd := exec.Command("sh")
		cmd.Stdin = strings.NewReader(w.Body.String())
		cmd.Env = []string{"HOME=" + home, "PATH=" + os.Getenv("PATH")}
		if out, err := cmd.CombinedOutput(); err != nil {
			t.Fatalf("%v: %s", err, out)
		}
	}
	b, _ := os.ReadFile(cfg)
	got := string(b)
	if !strings.Contains(got, `model = "o3"`) || !strings.Contains(got, "[mcp_servers.docs]") {
		t.Errorf("the person's settings were lost:\n%s", got)
	}
	if strings.Count(got, "[model_providers.llmgw]") != 1 || strings.Count(got, "[profiles.llmgw]") != 1 {
		t.Errorf("ours should be there once:\n%s", got)
	}
	if backups, _ := filepath.Glob(cfg + ".bak.*"); len(backups) == 0 {
		t.Error("no copy was kept of the file as it was")
	}
}

func TestASetupScriptRefusesANameAShellWouldRead(t *testing.T) {
	p := &AccountPages{}
	w := httptest.NewRecorder()
	r := httptest.NewRequest("GET", "http://gw.example/setup/codex?model=$(rm%20-rf%20~)", nil)
	r.SetPathValue("tool", "codex")
	p.setupScript(w, r)
	if w.Code != http.StatusBadRequest {
		t.Errorf("%d", w.Code)
	}
}

func TestContinueIsToldTheModelsContextAndWhetherItCanUseTools(t *testing.T) {
	// Told a context far larger than the gateway allows, Continue sends
	// prompts that are refused; told a model uses tools when its calls come
	// back as text, its agent stalls.
	agent := continueYAML("https://gw.example", "coder-30b", "k", ModelFacts{ContextLength: 32768, Tools: true})
	for _, want := range []string{"tool_use", "contextLength: 32768", "maxTokens: 4096"} {
		if !strings.Contains(agent, want) {
			t.Errorf("%q missing:\n%s", want, agent)
		}
	}
	chat := continueYAML("https://gw.example", "coder", "k", ModelFacts{ContextLength: 4096})
	if strings.Contains(chat, "tool_use") || !strings.Contains(chat, "contextLength: 4096") ||
		!strings.Contains(chat, "maxTokens: 1024") || !strings.Contains(chat, "its agent cannot act") {
		t.Errorf("a model without tools:\n%s", chat)
	}
	if unknown := continueYAML("https://gw.example", "m", "k", ModelFacts{}); strings.Contains(unknown, "contextLength") {
		t.Errorf("a context nobody knows was written down:\n%s", unknown)
	}
}

func TestContinueNeverAsksForMoreThanTheGatewayAllows(t *testing.T) {
	// Asked for more than max_output, the gateway refuses the request.
	y := continueYAML("https://gw.example", "m", "k", ModelFacts{ContextLength: 32768, MaxOutput: 2048, Tools: true})
	if !strings.Contains(y, "maxTokens: 2048") {
		t.Errorf("not capped at max_output:\n%s", y)
	}
}

func TestVSCodeGetsTheExtensionAndALinkBack(t *testing.T) {
	c := connectFor("vscode", "https://gw.example", "m", "sk-llmgw-abc", ModelFacts{})
	if !strings.Contains(c.Command, "https://gw.example/setup/vscode/llmgw.vsix") || !strings.Contains(c.Command, "code --install-extension llmgw.vsix") {
		t.Errorf("command = %s", c.Command)
	}
	if strings.Contains(c.Command, "sk-llmgw-abc") {
		t.Errorf("the key is in the download command: %s", c.Command)
	}
	link := string(c.Link)
	if !strings.HasPrefix(link, "vscode://llmgw.llmgw/connect?") || !strings.Contains(link, "gateway=https%3A%2F%2Fgw.example") ||
		!strings.Contains(link, "key=sk-llmgw-abc") || strings.Contains(link, "state=") {
		t.Errorf("link = %s", link)
	}
	if len(c.Fields) != 2 || c.Fields[0].Value != "https://gw.example" {
		t.Errorf("fields = %+v", c.Fields)
	}
}

func TestOnlyAStateShapedLikeOneGoesBackInTheLink(t *testing.T) {
	if got := signInState("abcDEF0123456789_-xy"); got != "abcDEF0123456789_-xy" {
		t.Errorf("a good state was dropped: %q", got)
	}
	for _, bad := range []string{"", "short", "has space in it here", `"><script>`, strings.Repeat("a", 65)} {
		if got := signInState(bad); got != "" {
			t.Errorf("signInState(%q) = %q", bad, got)
		}
	}
	if l := string(vscodeLink("https://gw", "k", signInState("abcdefghijklmnop"))); !strings.Contains(l, "state=abcdefghijklmnop") {
		t.Errorf("link = %s", l)
	}
}

func TestTheExtensionIsServedForThisGateway(t *testing.T) {
	p := &AccountPages{version: "v0.15.0"}
	w := httptest.NewRecorder()
	p.vsix(w, httptest.NewRequest("GET", "http://gw.example/setup/vscode/llmgw.vsix", nil))
	if w.Code != http.StatusOK || w.Header().Get("Cache-Control") != "no-store" ||
		!strings.Contains(w.Header().Get("Content-Disposition"), "llmgw.vsix") {
		t.Fatalf("%d %v", w.Code, w.Header())
	}
	if !strings.HasPrefix(w.Body.String(), "PK") {
		t.Error("not a zip")
	}
}