package ai import ( "context" "fmt" "os" "os/exec" "path/filepath" "strings" "time" ) type CommandContextProvider struct { NameStr string Prefixes []string GatherCmd []string Label string } func (p *CommandContextProvider) Name() string { return p.NameStr } func (p *CommandContextProvider) Matches(buf string) bool { trimmed := strings.ToLower(strings.TrimSpace(buf)) for _, prefix := range p.Prefixes { if strings.HasPrefix(trimmed, prefix) { return true } } return false } func (p *CommandContextProvider) Gather(ctx context.Context) (string, error) { if len(p.GatherCmd) == 0 { return "", nil } ctxTimeout, cancel := context.WithTimeout(ctx, 1000*time.Millisecond) defer cancel() out, err := exec.CommandContext(ctxTimeout, p.GatherCmd[0], p.GatherCmd[1:]...).Output() if err != nil { return "", err } if s := strings.TrimSpace(string(out)); s != "" { // Cap gathered command output to 1000 characters to keep prompt concise and avoid blowing up token budget if len(s) > 1000 { s = s[:1000] + "\n... (truncated)" } return p.Label + ":\n" + s, nil } return "", nil } var allowedHelpCommands = map[string]bool{ "git": true, "docker": true, "kubectl": true, "npm": true, "yarn": true, "pnpm": true, "cargo": true, "go": true, "systemctl": true, "helm": true, "terraform": true, "aws": true, "gcloud": true, "az": true, "make": true, "bun": true, "pip": true, "python": true, "python3": true, "node": true, "deno": true, "tar": true, "curl": true, "wget": true, "ssh": true, "podman": true, "tofu": true, "ansible": true, "gh": true, "nix": true, } func isAllowedForHelp(cmdName string) bool { if strings.ContainsAny(cmdName, "/\\") { return false } return allowedHelpCommands[cmdName] } type universalProvider struct { cwd string buf string } func (p *universalProvider) Name() string { firstWord := "" if fields := strings.Fields(p.buf); len(fields) > 0 { firstWord = fields[0] } return "universal:" + p.cwd + ":" + firstWord } func (p *universalProvider) Matches(buf string) bool { return true } func (p *universalProvider) Gather(ctx context.Context) (string, error) { ctxTimeout, cancel := context.WithTimeout(ctx, 1200*time.Millisecond) defer cancel() var sb strings.Builder ExtractScriptsAndTargets(&sb, p.cwd, "") if entries, err := os.ReadDir(p.cwd); err == nil { var names []string for i, e := range entries { if i >= 30 { names = append(names, "...") break } name := e.Name() if e.IsDir() { name += "/" if !strings.HasPrefix(e.Name(), ".") && e.Name() != "node_modules" && i < 15 { ExtractScriptsAndTargets(&sb, filepath.Join(p.cwd, e.Name()), e.Name()) } } names = append(names, name) } if len(names) > 0 { fmt.Fprintf(&sb, "Files in Cwd: %s\n\n", strings.Join(names, ", ")) } } cmd := exec.CommandContext(ctxTimeout, "git", "-C", p.cwd, "rev-parse", "--is-inside-work-tree") if cmd.Run() == nil { statusOut, _ := exec.CommandContext(ctxTimeout, "git", "-C", p.cwd, "status", "-s").Output() statusStr := strings.TrimSpace(string(statusOut)) if len(statusStr) > 1000 { statusStr = statusStr[:1000] + "\n... (truncated)" } diffOut, _ := exec.CommandContext(ctxTimeout, "git", "-C", p.cwd, "diff", "--staged").Output() diffStr := strings.TrimSpace(string(diffOut)) if len(diffStr) > 1500 { diffStr = diffStr[:1500] + "\n... (truncated)" } logOut, _ := exec.CommandContext(ctxTimeout, "git", "-C", p.cwd, "log", "-n", "5", "--no-decorate", "--pretty=format:%s").Output() logStr := strings.TrimSpace(string(logOut)) sb.WriteString("Git Repository State:\n") if statusStr != "" { fmt.Fprintf(&sb, "Status:\n%s\n\n", statusStr) } if diffStr != "" { fmt.Fprintf(&sb, "Staged Diff:\n%s\n\n", diffStr) } if logStr != "" { fmt.Fprintf(&sb, "User's recent commit messages (MUST follow this exact style, formatting, language, and casing conventions):\n%s\n", logStr) } } if fields := strings.Fields(p.buf); len(fields) > 0 { cmdName := fields[0] if isAllowedForHelp(cmdName) { ctxHelp, cancel := context.WithTimeout(ctx, 500*time.Millisecond) defer cancel() helpOut, err := exec.CommandContext(ctxHelp, cmdName, "--help").CombinedOutput() if err == nil { helpStr := strings.TrimSpace(string(helpOut)) if len(helpStr) > 600 { helpStr = helpStr[:600] } if helpStr != "" { fmt.Fprintf(&sb, "\nCommand help (%s --help):\n%s\n", cmdName, helpStr) } } } } return strings.TrimSpace(sb.String()), nil }