nub/internal/ctxasm/assemble.go

// Package ctxasm baut den System-Block eines Requests auf: stabil vor
// volatil, mit einem Cache-Breakpoint am Ende (5.4). Alles Dynamische
// (Git-Status, Datum, Tool-Ausgaben) gehört in Messages, nie hierher.
package ctxasm

import (
	_ "embed"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"

	"nub/internal/llm"
	"nub/internal/tokens"
)

//go:embed default_system.md
var defaultSystemPrompt string

// Options steuert den Assemble-Lauf. RepoRoot/Cwd werden für die
// Umgebungs-Zeile und die walk_up-Dateisuche gebraucht.
type Options struct {
	RepoRoot    string
	Cwd         string
	Files       []string // konfigurierte Kontext-Dateinamen, z.B. AGENTS.md
	WalkUp      bool
	MaxTokens   int    // Default 20000, wenn <= 0
	SkillsIndex string // vorgerendert (internal/skill.RenderIndex), leer = kein Block
}

// FileOrigin protokolliert Herkunft und Kosten jeder geladenen Kontext-Datei
// (Grundlage für `nub context`).
type FileOrigin struct {
	Path      string // repo-relativ
	Source    string // absoluter Pfad, aus dem gelesen wurde
	Tokens    int
	Truncated bool
}

type Meta struct {
	Files       []FileOrigin
	Warnings    []string
	TotalTokens int
}

// Assemble baut die System-Blöcke in der Reihenfolge: Basis-System-Prompt,
// statischer Umgebungsblock, konfigurierte Kontext-Dateien. Der letzte Block
// bekommt CacheMark=true.
func Assemble(opts Options) ([]llm.Block, Meta, error) {
	budget := opts.MaxTokens
	if budget <= 0 {
		budget = 20000
	}

	var blocks []llm.Block
	var meta Meta

	systemText := systemPromptText(opts.RepoRoot)
	envText := environmentBlock(opts.RepoRoot)
	blocks = append(blocks, llm.Block{Kind: llm.KindText, Text: systemText})
	blocks = append(blocks, llm.Block{Kind: llm.KindText, Text: envText})

	// Das Budget gilt für die konfigurierten Kontext-Dateien, nicht für den
	// (kleinen, festen) System-Prompt/Env-Block.
	used := 0
	for _, name := range opts.Files {
		for _, path := range resolveFilePaths(opts.RepoRoot, opts.Cwd, name, opts.WalkUp) {
			data, err := os.ReadFile(path)
			if err != nil {
				continue
			}
			rel, err := filepath.Rel(opts.RepoRoot, path)
			if err != nil {
				rel = path
			}

			text := string(data)
			tok := tokens.Estimate(text)
			truncated := false

			if used+tok > budget {
				remaining := budget - used
				if remaining <= 0 {
					meta.Warnings = append(meta.Warnings, fmt.Sprintf("context budget (%d tokens) exceeded, skipping %s", budget, rel))
					continue
				}
				maxChars := remaining * 4
				if maxChars < len(text) {
					text = text[:maxChars] + "\n... [truncated: context budget exceeded] ..."
					tok = tokens.Estimate(text)
					truncated = true
					meta.Warnings = append(meta.Warnings, fmt.Sprintf("%s truncated: context budget (%d tokens) exceeded", rel, budget))
				}
			}

			blocks = append(blocks, llm.Block{Kind: llm.KindText, Text: fmt.Sprintf("# %s\n\n%s", rel, text)})
			used += tok
			meta.Files = append(meta.Files, FileOrigin{Path: rel, Source: path, Tokens: tok, Truncated: truncated})
		}
	}
	meta.TotalTokens = used

	if opts.SkillsIndex != "" {
		blocks = append(blocks, llm.Block{Kind: llm.KindText, Text: opts.SkillsIndex})
	}

	if len(blocks) > 0 {
		blocks[len(blocks)-1].CacheMark = true
	}

	warnStaleRepoMap(opts.RepoRoot, meta.Files, &meta.Warnings)

	return blocks, meta, nil
}

func systemPromptText(repoRoot string) string {
	if data, err := os.ReadFile(filepath.Join(repoRoot, "SYSTEM.md")); err == nil {
		return string(data)
	}
	return defaultSystemPrompt
}

func environmentBlock(repoRoot string) string {
	var b strings.Builder
	fmt.Fprintf(&b, "OS: %s\n", runtime.GOOS)
	shell := os.Getenv("SHELL")
	if shell == "" {
		shell = "unknown"
	}
	fmt.Fprintf(&b, "Shell: %s\n", shell)
	fmt.Fprintf(&b, "RepoRoot: %s\n", repoRoot)
	if branch, err := gitBranch(repoRoot); err == nil && branch != "" {
		fmt.Fprintf(&b, "Branch: %s\n", branch)
	}
	return b.String()
}

func gitBranch(repoRoot string) (string, error) {
	out, err := exec.Command("git", "-C", repoRoot, "rev-parse", "--abbrev-ref", "HEAD").Output()
	if err != nil {
		return "", err
	}
	return strings.TrimSpace(string(out)), nil
}

// resolveFilePaths sammelt alle Treffer für `name` zwischen Cwd und RepoRoot
// (walk_up) in der Reihenfolge RepoRoot -> Cwd (näher am Arbeitsverzeichnis
// gewinnt inhaltlich, weil später gelesen). Ohne walk_up nur RepoRoot.
func resolveFilePaths(repoRoot, cwd, name string, walkUp bool) []string {
	if !walkUp {
		p := filepath.Join(repoRoot, name)
		if fileExists(p) {
			return []string{p}
		}
		return nil
	}

	var dirs []string
	d := filepath.Clean(cwd)
	root := filepath.Clean(repoRoot)
	for {
		dirs = append(dirs, d)
		if d == root {
			break
		}
		parent := filepath.Dir(d)
		if parent == d {
			break // cwd liegt nicht unterhalb von repoRoot
		}
		d = parent
	}

	var out []string
	for i := len(dirs) - 1; i >= 0; i-- {
		p := filepath.Join(dirs[i], name)
		if fileExists(p) {
			out = append(out, p)
		}
	}
	return out
}

func fileExists(path string) bool {
	info, err := os.Stat(path)
	return err == nil && !info.IsDir()
}