llmgw/internal/web/codeblocks.go

package web

import (
	"fmt"
	"path"
	"strings"
)

// CodeFile is one fenced block from a reply, with the path it should be
// written to.
type CodeFile struct {
	Path string
	Body string
}

// maxCodeFiles and maxCodeBytes bound what one reply can turn into.
//
// The reply is already bounded by the token cap, so these are not the real
// defence; they are there so a change to that cap, or a model that emits
// something pathological, cannot turn a download into an archive nobody asked
// for.
const (
	maxCodeFiles = 64
	maxCodeBytes = 4 << 20
)

// ExtractCode pulls fenced code blocks out of a model reply and decides where
// each one belongs.
//
// The filename is whatever the model said, and models say it in several ways:
// on the fence itself (```go:main.go, ```main.go), or on the line above the
// fence as a heading or a comment. All of them are guesses, and a block whose
// name cannot be worked out still comes back — under a generated name, because
// dropping someone's code because it was unlabelled is worse than naming it
// badly.
func ExtractCode(reply string) []CodeFile {
	lines := strings.Split(reply, "\n")
	var out []CodeFile
	used := map[string]int{}

	for i := 0; i < len(lines); i++ {
		fence, info, ok := openingFence(lines[i])
		if !ok {
			continue
		}
		var body []string
		j := i + 1
		for ; j < len(lines); j++ {
			if isClosingFence(lines[j], fence) {
				break
			}
			body = append(body, lines[j])
		}

		lang, named := splitInfo(info)
		if named == "" {
			named = nameFromNearbyLine(lines, i)
		}
		p := safePath(named)
		if p == "" {
			p = fmt.Sprintf("snippet-%d%s", len(out)+1, extensionFor(lang))
		}
		// Two blocks claiming one path would silently lose the first. Models
		// do this when they show a file, then show it again with one line
		// changed.
		if n := used[p]; n > 0 {
			ext := path.Ext(p)
			p = fmt.Sprintf("%s-%d%s", strings.TrimSuffix(p, ext), n+1, ext)
		}
		used[p]++

		out = append(out, CodeFile{Path: p, Body: strings.Join(body, "\n") + "\n"})
		if len(out) >= maxCodeFiles {
			break
		}
		i = j
	}

	total := 0
	for i := range out {
		total += len(out[i].Body)
		if total > maxCodeBytes {
			return out[:i]
		}
	}
	return out
}

// openingFence recognises ``` and ~~~ fences of three characters or more,
// returning the fence itself so the matching close can be found.
func openingFence(line string) (fence, info string, ok bool) {
	t := strings.TrimSpace(line)
	for _, c := range []string{"`", "~"} {
		if !strings.HasPrefix(t, strings.Repeat(c, 3)) {
			continue
		}
		n := 0
		for n < len(t) && string(t[n]) == c {
			n++
		}
		info = strings.TrimSpace(t[n:])
		// An info string containing a backtick is not a fence at all; that is
		// inline code with backticks in it.
		if c == "`" && strings.Contains(info, "`") {
			return "", "", false
		}
		return t[:n], info, true
	}
	return "", "", false
}

func isClosingFence(line, fence string) bool {
	t := strings.TrimSpace(line)
	return strings.HasPrefix(t, fence) && strings.TrimRight(t, string(fence[0])) == ""
}

// splitInfo separates the language from a filename on the fence line.
//
// Handles ```go:path/main.go, ```go path/main.go, ```path/main.go and bare
// ```go. A token is treated as a name rather than a language when it looks
// like one: it has a dot or a slash in it.
func splitInfo(info string) (lang, name string) {
	if info == "" {
		return "", ""
	}
	if l, n, found := strings.Cut(info, ":"); found {
		return strings.TrimSpace(l), strings.TrimSpace(n)
	}
	fields := strings.Fields(info)
	if len(fields) == 0 {
		return "", ""
	}
	first := fields[0]
	// Strip the attribute forms some tools use: title="x", file=x.
	for _, f := range fields[1:] {
		for _, k := range []string{"title=", "file=", "filename=", "name="} {
			if v, ok := strings.CutPrefix(f, k); ok {
				return first, strings.Trim(v, `"'`)
			}
		}
	}
	if looksLikePath(first) {
		return "", first
	}
	if len(fields) > 1 && looksLikePath(fields[1]) {
		return first, fields[1]
	}
	return first, ""
}

func looksLikePath(s string) bool {
	if strings.Contains(s, "/") || strings.Contains(s, ".") {
		return true
	}
	// Files whose whole name is their type. ```Dockerfile is a language tag to
	// a syntax highlighter and a filename to everyone else, and the reader
	// wants a file called Dockerfile either way.
	switch s {
	case "Dockerfile", "Makefile", "Jenkinsfile", "Vagrantfile", "Procfile", "Gemfile", "Rakefile", "CMakeLists":
		return true
	}
	return false
}

// nameFromNearbyLine looks just above the fence for a filename.
//
// Models very often write the path as a short line of its own — as a heading,
// bold, in backticks, or as a comment — and then open the fence. Only the two
// lines immediately above are considered, because anything further away is
// prose that happens to mention a file.
func nameFromNearbyLine(lines []string, fence int) string {
	for k := fence - 1; k >= 0 && k >= fence-2; k-- {
		t := strings.TrimSpace(lines[k])
		if t == "" {
			continue
		}
		t = strings.Trim(t, "*_#:` ")
		t = strings.TrimPrefix(t, "//")
		t = strings.TrimPrefix(t, "#")
		for _, k := range []string{"File:", "file:", "FILE:", "Filename:", "filename:"} {
			t = strings.TrimSpace(strings.TrimPrefix(t, k))
		}
		t = strings.TrimSpace(t)
		// One token that looks like a path. A sentence mentioning main.go is
		// not a filename line.
		if fields := strings.Fields(t); len(fields) == 1 && looksLikePath(fields[0]) {
			return fields[0]
		}
		return ""
	}
	return ""
}

// safePath turns a model-supplied name into something safe to write inside an
// archive, or "" when it cannot.
//
// This is the part that matters. The name comes from a language model and goes
// into a zip that someone will extract, so "../../.ssh/authorized_keys" and
// "/etc/passwd" and "C:\\Windows\\..." are all things to expect rather than
// things to hope against. Everything is forced to a relative path under the
// archive root, and anything that still escapes is rejected outright rather
// than repaired into something surprising.
func safePath(name string) string {
	name = strings.TrimSpace(name)
	if name == "" {
		return ""
	}
	// Windows separators and drive letters, before anything else looks at it.
	name = strings.ReplaceAll(name, `\`, "/")
	if len(name) > 1 && name[1] == ':' {
		name = name[2:]
	}
	name = strings.TrimPrefix(name, "/")

	cleaned := path.Clean(name)
	if cleaned == "." || cleaned == "/" {
		return ""
	}
	if strings.HasPrefix(cleaned, "../") || cleaned == ".." || strings.HasPrefix(cleaned, "/") {
		return ""
	}
	for _, seg := range strings.Split(cleaned, "/") {
		if seg == "" || seg == "." || seg == ".." {
			return ""
		}
		// A NUL or a newline in a path is not a filename anyone meant.
		if strings.ContainsAny(seg, "\x00\n\r") {
			return ""
		}
		if len(seg) > 128 {
			return ""
		}
	}
	if len(cleaned) > 512 || strings.Count(cleaned, "/") > 16 {
		return ""
	}
	return cleaned
}

// extensionFor gives an unnamed block a plausible extension from its language.
func extensionFor(lang string) string {
	switch strings.ToLower(lang) {
	case "go":
		return ".go"
	case "python", "py":
		return ".py"
	case "javascript", "js":
		return ".js"
	case "typescript", "ts":
		return ".ts"
	case "rust", "rs":
		return ".rs"
	case "c":
		return ".c"
	case "cpp", "c++":
		return ".cpp"
	case "java":
		return ".java"
	case "sh", "bash", "shell":
		return ".sh"
	case "sql":
		return ".sql"
	case "yaml", "yml":
		return ".yaml"
	case "json":
		return ".json"
	case "html":
		return ".html"
	case "css":
		return ".css"
	case "":
		return ".txt"
	default:
		return "." + strings.ToLower(lang)
	}
}