llmgw/internal/web/connect.go

package web

import (
	"fmt"
	"html/template"
	"net/http"
	"net/url"
	"regexp"
	"strings"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
	"gitea.mynodes.xyz/ilker/llmgw/internal/vscodeext"
)

// Connecting an editor: a key made for it, and the configuration it needs
// with that key, this gateway's address and a served model already in it.
//
// A page cannot write files on somebody's machine, so what it can do is hand
// over the file, ready — to copy, to download, or to have a one-line command
// write. The command's script is served without authentication and holds no
// secret: the key travels in an environment variable, never in a URL, where
// it would land in shell history and proxy logs.

// editorTools are the tools the page knows how to configure.
var editorTools = []struct{ Key, Name string }{
	{"vscode", "VS Code"}, {"continue", "Continue"}, {"codex", "Codex CLI"}, {"cline", "Cline"},
}

// connectView is what the page shows after a key was made for a tool.
type connectView struct {
	Tool, ToolName, Model string
	Key                   string
	// Files are what the tool reads, with the key in them.
	Files []connectFile
	// Fields are for a tool that keeps its settings where no file reaches:
	// Cline's live in VS Code's secret storage, set through its own screen.
	Fields []connectField
	// Command writes the files; Run is how to start the tool on the gateway.
	Command, Run string
	Notes        []string
	// Link hands the key to a tool that takes it through a link of its own:
	// llmgw's VS Code extension, at vscode://llmgw.llmgw/connect.
	Link template.URL
}

type connectFile struct{ Path, Name, Body string }
type connectField struct{ Label, Value string }

// safeModel is what a model name may be to go into a script and a file
// unquoted: the names this gateway serves, and nothing a shell reads twice.
var safeModel = regexp.MustCompile(`^[A-Za-z0-9._:/@+-]{1,200}$`)

// ModelFacts is what an editor's configuration needs to know of a model: how
// much one request may hold, and whether its tool calls come back as calls.
type ModelFacts struct {
	ContextLength int
	// MaxOutput is the longest answer the gateway allows; 0 is no limit.
	MaxOutput int
	Tools     bool
}

// continueYAML is a Continue assistant pointing at the gateway: a file of its
// own under ~/.continue/assistants, so the person's config.yaml is not touched.
//
// What it says of the model is what the catalog says. Continue assumed a
// context far larger than a CPU replica's, and every request over it was
// refused; and told a model uses tools when its calls come back as text,
// Continue's agent stalls on them rather than acting.
func continueYAML(base, model, key string, f ModelFacts) string {
	var b strings.Builder
	fmt.Fprintf(&b, `# Written by llmgw for Continue: the gateway at %[1]s.
name: llmgw
version: 1.0.0
schema: v1
models:
  - name: %[2]s (llmgw)
    provider: openai
    model: %[2]s
    apiBase: %[1]s/v1
    apiKey: %[3]s
    roles:
      - chat
      - edit
      - apply
`, base, model, key)
	if f.Tools {
		b.WriteString("    capabilities:\n      - tool_use\n")
	}
	if f.ContextLength > 0 {
		fmt.Fprintf(&b, `    # What the gateway lets one request hold; Continue trims to fit.
    defaultCompletionOptions:
      contextLength: %d
      maxTokens: %d
`, f.ContextLength, maxTokensFor(f))
	}
	if !f.Tools {
		b.WriteString("# This model's catalog entry does not say it makes tool calls, so Continue\n" +
			"# chats and edits with it but its agent cannot act: choose one that does.\n")
	}
	b.WriteString(`# For completion as you type, add a small, fast model here with the role
# autocomplete; a large one makes every keystroke wait for it.
`)
	return b.String()
}

// codexTOML is the provider and profile Codex needs, and nothing else of its
// configuration: codex --profile llmgw uses the gateway, plain codex what it
// used before. The key is read from LLMGW_API_KEY rather than written down.
func codexTOML(base, model string) string {
	return fmt.Sprintf(`# Written by llmgw for Codex: the gateway at %[1]s.
[model_providers.llmgw]
name = "llmgw"
base_url = "%[1]s/v1"
env_key = "LLMGW_API_KEY"
wire_api = "responses"

[profiles.llmgw]
model_provider = "llmgw"
model = "%[2]s"
`, base, model)
}

// connectFor builds what the page shows for a tool, in English.
func connectFor(tool, base, model, key string, f ModelFacts) connectView {
	return connectIn(english, tool, base, model, key, f)
}

// connectIn builds what the page shows for a tool, its sentences in l. The
// files, commands and paths are the same in every language.
func connectIn(l *Lang, tool, base, model, key string, f ModelFacts) connectView {
	v := connectView{Tool: tool, Model: model, Key: key}
	for _, t := range editorTools {
		if t.Key == tool {
			v.ToolName = t.Name
		}
	}
	script := base + "/setup/" + tool + "?model=" + model
	switch tool {
	case "continue":
		v.Files = []connectFile{{Path: "~/.continue/assistants/llmgw.yaml", Name: "llmgw.yaml",
			Body: continueYAML(base, model, key, f)}}
		v.Command = fmt.Sprintf("curl -fsSL '%s' | LLMGW_KEY='%s' sh", script, key)
		v.Run = l.T("connect.run.continue")
	case "codex":
		v.Files = []connectFile{{Path: "~/.codex/config.toml", Name: "config.toml", Body: codexTOML(base, model)}}
		v.Command = fmt.Sprintf("curl -fsSL '%s' | sh", script)
		v.Run = fmt.Sprintf("LLMGW_API_KEY='%s' codex --profile llmgw", key)
		v.Notes = []string{l.T("connect.note.codex_blocks"), l.T("connect.note.codex_env")}
	case "vscode":
		download := base + "/setup/vscode/" + vscodeext.FileName
		v.Command = fmt.Sprintf("curl -fsSLo %[1]s '%[2]s' && code --install-extension %[1]s --force", vscodeext.FileName, download)
		v.Link = vscodeLink(base, key, "")
		v.Fields = []connectField{
			{l.T("connect.field.gateway"), base},
			{l.T("connect.field.key"), key},
		}
		v.Run = l.T("connect.run.vscode")
		v.Notes = []string{l.T("connect.note.vscode_agent")}
	case "cline":
		// The values are what Cline's own screen offers and takes, which is
		// English whatever language this page is in.
		v.Fields = []connectField{
			{l.T("connect.field.provider"), "OpenAI Compatible"},
			{l.T("connect.field.base_url"), base + "/v1"},
			{l.T("connect.field.key"), key},
			{l.T("connect.field.model"), model},
		}
		v.Run = l.T("connect.run.cline")
	}
	return v
}

// connect makes a key for an editor and shows its configuration.
func (p *AccountPages) connect(w http.ResponseWriter, r *http.Request) {
	noStore(w)
	sess, user, ok := p.session(r)
	if !ok {
		http.Redirect(w, r, "/login", http.StatusSeeOther)
		return
	}
	if !p.accounts.CheckCSRF(sess, r.FormValue("csrf")) {
		badForm(w, r)
		return
	}
	l := langOf(r)
	tool := r.FormValue("tool")
	known := false
	for _, t := range editorTools {
		known = known || t.Key == tool
	}
	model := strings.TrimSpace(r.FormValue("model"))
	if !known || !safeModel.MatchString(model) {
		p.renderProfile(w, r, sess, user, "", l.T("connect.err.choose"))
		return
	}
	project := r.FormValue("context")
	if project == personalValue {
		if p.personalOff {
			p.renderProfile(w, r, sess, user, "", l.T("profile.err.personal_off"))
			return
		}
		project = ""
	}
	name := tool + " " + time.Now().Format("2006-01-02")
	secret, k, err := store.NewKey(randomKeyID(), name, project)
	if err != nil {
		p.renderProfile(w, r, sess, user, "", l.T("profile.err.try_again"))
		return
	}
	k.CreatedAt = time.Now()
	if err := p.keys.CreateUserKey(r.Context(), user.ID, k); err != nil {
		msg := l.T("profile.err.try_again")
		if err == store.ErrNotMember {
			msg = l.T("profile.err.not_member")
		}
		p.renderProfile(w, r, sess, user, "", msg)
		return
	}
	p.log.Info("a user connected an editor", "user", user.ID, "tool", tool, "key", k.ID)
	c := connectIn(l, tool, baseURL(r), model, secret, p.facts(model))
	if tool == "vscode" {
		// The extension's sign-in sent a state of its own; the link carries
		// it back so the extension knows this is the page it opened.
		c.Link = vscodeLink(baseURL(r), secret, signInState(r.FormValue("state")))
	}
	p.render(w, r, sess, user, profileData{Connect: &c})
}

// vscodeLink is the link that hands VS Code's llmgw extension a gateway and
// a key. It is opened by the person's own browser and handled by their own
// editor; the key is on the page beside it already.
func vscodeLink(base, key, state string) template.URL {
	q := url.Values{"gateway": {base}, "key": {key}}
	if state != "" {
		q.Set("state", state)
	}
	return template.URL("vscode://" + vscodeext.ID + "/connect?" + q.Encode())
}

// stateShape is what the extension's sign-in state looks like: random,
// base64url. Anything else is dropped rather than carried into a link.
var stateShape = regexp.MustCompile(`^[A-Za-z0-9_-]{16,64}$`)

func signInState(s string) string {
	if stateShape.MatchString(s) {
		return s
	}
	return ""
}

// vsix serves the VS Code extension, made for this gateway: its address is
// the extension's setting to start with. Served without a session, like the
// setup scripts — it holds no secret, only the address the person typed.
func (p *AccountPages) vsix(w http.ResponseWriter, r *http.Request) {
	b, err := vscodeext.Build(p.version, baseURL(r))
	if err != nil {
		p.log.Error("could not make the VS Code extension", "err", err)
		http.Error(w, "could not make the extension", http.StatusInternalServerError)
		return
	}
	// Never from a cache, for the same reason as the scripts: it names the
	// address the request came to.
	w.Header().Set("Cache-Control", "no-store")
	w.Header().Set("Content-Type", "application/vsix")
	w.Header().Set("Content-Disposition", `attachment; filename="`+vscodeext.FileName+`"`)
	_, _ = w.Write(b)
}

// setupScript is the script the one-line command runs: it writes the tool's
// configuration with the key from the environment, keeping a copy of what
// was there. Served without a session — it holds nothing a stranger could
// use: this gateway's address, which they already typed, and a model name.
func (p *AccountPages) setupScript(w http.ResponseWriter, r *http.Request) {
	// Never from a cache: the script names this gateway's address, and one
	// cached for another request's Host would send a client's key elsewhere.
	w.Header().Set("Cache-Control", "no-store")
	tool := r.PathValue("tool")
	model := r.URL.Query().Get("model")
	if !safeModel.MatchString(model) {
		http.Error(w, "which model? add ?model=<name>", http.StatusBadRequest)
		return
	}
	base := baseURL(r)
	var body string
	switch tool {
	case "continue":
		body = fmt.Sprintf(`#!/bin/sh
# Connects Continue to the llmgw gateway at %[1]s.
set -eu
: "${LLMGW_KEY:?set LLMGW_KEY to the key the gateway gave you}"
dir="$HOME/.continue/assistants"
file="$dir/llmgw.yaml"
mkdir -p "$dir"
[ -f "$file" ] && cp "$file" "$file.bak.$(date +%%Y%%m%%d%%H%%M%%S)"
umask 077
cat > "$file" <<'LLMGW'
%[2]sLLMGW
# The key goes in last, from the environment: never through a URL.
tmp="$file.tmp"
sed "s|__LLMGW_KEY__|$LLMGW_KEY|" "$file" > "$tmp" && mv "$tmp" "$file"
echo "Wrote $file. In Continue, choose the llmgw assistant."
`, base, continueYAML(base, model, "__LLMGW_KEY__", p.facts(model)))
	case "codex":
		body = fmt.Sprintf(`#!/bin/sh
# Connects Codex to the llmgw gateway at %[1]s.
set -eu
dir="$HOME/.codex"
file="$dir/config.toml"
mkdir -p "$dir"
touch "$file"
cp "$file" "$file.bak.$(date +%%Y%%m%%d%%H%%M%%S)"
# Out with the blocks this wrote before, in with the new: every other line
# of the file stays as it was.
awk '
  /^\[/ { skip = ($0 == "[model_providers.llmgw]" || $0 == "[profiles.llmgw]" || $0 ~ /^# Written by llmgw/) }
  /^# Written by llmgw/ { skip = 1; next }
  !skip { print }
' "$file" > "$file.tmp"
cat >> "$file.tmp" <<'LLMGW'

%[2]sLLMGW
mv "$file.tmp" "$file"
echo "Added the llmgw provider to $file."
echo "Run: LLMGW_API_KEY=<your key> codex --profile llmgw"
`, base, codexTOML(base, model))
	default:
		http.NotFound(w, r)
		return
	}
	w.Header().Set("Content-Type", "text/x-shellscript; charset=utf-8")
	w.Header().Set("Cache-Control", "no-store")
	_, _ = w.Write([]byte(body))
}

// facts is what the catalog says of a model, when there is one to ask.
func (p *AccountPages) facts(model string) ModelFacts {
	if p.modelFacts == nil {
		return ModelFacts{}
	}
	return p.modelFacts(model)
}

// maxTokensFor is the answer length Continue asks for: a quarter of the
// context, at most 4096, and never over what the gateway allows — asked for
// more, it refuses the request.
func maxTokensFor(f ModelFacts) int {
	n := min(4096, max(256, f.ContextLength/4))
	if f.MaxOutput > 0 {
		n = min(n, f.MaxOutput)
	}
	return n
}