diple/ai_diff.go

package main

import (
	"bufio"
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"path"
	"path/filepath"
	"strconv"
	"strings"
)

func (c *GitHubClient) PullRequestDiff(ctx context.Context, owner, repo string, number int) (string, error) {
	requestURL := c.restBaseURL() + "/repos/" + url.PathEscape(owner) + "/" + url.PathEscape(repo) +
		"/pulls/" + strconv.Itoa(number)
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, requestURL, nil)
	if err != nil {
		return "", err
	}
	req.Header.Set("Authorization", "Bearer "+c.token)
	req.Header.Set("Accept", "application/vnd.github.v3.diff")
	req.Header.Set("User-Agent", "diple")
	response, err := c.http.Do(req)
	if err != nil {
		return "", err
	}
	defer response.Body.Close()
	if response.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(io.LimitReader(response.Body, 4096))
		return "", fmt.Errorf("GitHub returned %s: %s", response.Status, strings.TrimSpace(string(body)))
	}
	body, err := io.ReadAll(io.LimitReader(response.Body, (32<<20)+1))
	if err != nil {
		return "", err
	}
	if len(body) > 32<<20 {
		return "", fmt.Errorf("PR diff exceeds the 32 MiB safety limit")
	}
	return string(body), nil
}

func (c *CachedGitHubService) PullRequestDiff(ctx context.Context, owner, repo string, number int) (string, error) {
	service, ok := c.remote.(AIDiffService)
	if !ok {
		return "", fmt.Errorf("configured GitHub service cannot load pull request diffs")
	}
	return service.PullRequestDiff(ctx, owner, repo, number)
}

func (c *GitHubClient) RepositoryTree(
	ctx context.Context, owner, repo, commitOID string,
) (AIRepositoryTree, error) {
	baseURL := c.restBaseURL() + "/repos/" + url.PathEscape(owner) + "/" + url.PathEscape(repo)
	var commit struct {
		Tree struct {
			SHA string `json:"sha"`
		} `json:"tree"`
	}
	if err := c.getAIRepositoryJSON(
		ctx, baseURL+"/git/commits/"+url.PathEscape(commitOID), 1<<20, &commit,
	); err != nil {
		return AIRepositoryTree{}, fmt.Errorf("load repository commit: %w", err)
	}
	if commit.Tree.SHA == "" {
		return AIRepositoryTree{}, fmt.Errorf("load repository commit: GitHub returned no tree OID")
	}
	var tree struct {
		Truncated bool `json:"truncated"`
		Tree      []struct {
			Path string `json:"path"`
			Mode string `json:"mode"`
			Type string `json:"type"`
			SHA  string `json:"sha"`
			Size int64  `json:"size"`
		} `json:"tree"`
	}
	if err := c.getAIRepositoryJSON(
		ctx, baseURL+"/git/trees/"+url.PathEscape(commit.Tree.SHA)+"?recursive=1",
		8<<20, &tree,
	); err != nil {
		return AIRepositoryTree{}, fmt.Errorf("load repository tree: %w", err)
	}
	result := AIRepositoryTree{
		CommitOID: commitOID,
		Entries:   make([]AIRepositoryEntry, 0, len(tree.Tree)),
		Truncated: tree.Truncated,
	}
	for _, entry := range tree.Tree {
		clean := filepath.ToSlash(filepath.Clean(entry.Path))
		if clean == "." || clean == "" || filepath.IsAbs(clean) ||
			strings.HasPrefix(clean, "../") {
			continue
		}
		result.Entries = append(result.Entries, AIRepositoryEntry{
			Path: clean, OID: entry.SHA, Type: entry.Type, Mode: entry.Mode, Size: entry.Size,
		})
	}
	return result, nil
}

func (c *CachedGitHubService) RepositoryTree(
	ctx context.Context, owner, repo, commitOID string,
) (AIRepositoryTree, error) {
	service, ok := c.remote.(AIRepositoryService)
	if !ok {
		return AIRepositoryTree{}, fmt.Errorf("configured GitHub service cannot load repository trees")
	}
	return service.RepositoryTree(ctx, owner, repo, commitOID)
}

func (c *GitHubClient) RepositoryBlob(
	ctx context.Context, owner, repo, blobOID string, maxBytes int,
) ([]byte, error) {
	baseURL := c.restBaseURL() + "/repos/" + url.PathEscape(owner) + "/" + url.PathEscape(repo)
	var blob struct {
		Content  string `json:"content"`
		Encoding string `json:"encoding"`
		Size     int    `json:"size"`
		SHA      string `json:"sha"`
	}
	responseLimit := int64(max(16_384, maxBytes*2+8_192))
	if err := c.getAIRepositoryJSON(
		ctx, baseURL+"/git/blobs/"+url.PathEscape(blobOID), responseLimit, &blob,
	); err != nil {
		return nil, fmt.Errorf("load repository blob: %w", err)
	}
	if blob.SHA != "" && blob.SHA != blobOID {
		return nil, fmt.Errorf("load repository blob: GitHub returned an unexpected blob OID")
	}
	if blob.Size > maxBytes {
		return nil, fmt.Errorf("repository file exceeds the %d-byte AI limit", maxBytes)
	}
	if blob.Encoding != "base64" {
		return nil, fmt.Errorf("load repository blob: unsupported encoding %q", blob.Encoding)
	}
	content, err := base64.StdEncoding.DecodeString(strings.Map(func(r rune) rune {
		if r == '\r' || r == '\n' || r == ' ' || r == '\t' {
			return -1
		}
		return r
	}, blob.Content))
	if err != nil {
		return nil, fmt.Errorf("decode repository blob: %w", err)
	}
	if len(content) > maxBytes {
		return nil, fmt.Errorf("repository file exceeds the %d-byte AI limit", maxBytes)
	}
	return content, nil
}

func (c *CachedGitHubService) RepositoryBlob(
	ctx context.Context, owner, repo, blobOID string, maxBytes int,
) ([]byte, error) {
	service, ok := c.remote.(AIRepositoryService)
	if !ok {
		return nil, fmt.Errorf("configured GitHub service cannot load repository blobs")
	}
	return service.RepositoryBlob(ctx, owner, repo, blobOID, maxBytes)
}

func (c *GitHubClient) getAIRepositoryJSON(
	ctx context.Context, requestURL string, limit int64, output any,
) error {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, requestURL, nil)
	if err != nil {
		return err
	}
	req.Header.Set("Authorization", "Bearer "+c.token)
	req.Header.Set("Accept", "application/vnd.github+json")
	req.Header.Set("User-Agent", "diple")
	response, err := c.http.Do(req)
	if err != nil {
		return err
	}
	defer response.Body.Close()
	if response.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(io.LimitReader(response.Body, 4096))
		return fmt.Errorf("GitHub returned %s: %s", response.Status, strings.TrimSpace(string(body)))
	}
	data, err := io.ReadAll(io.LimitReader(response.Body, limit+1))
	if err != nil {
		return err
	}
	if int64(len(data)) > limit {
		return fmt.Errorf("GitHub response exceeds the %d-byte safety limit", limit)
	}
	if err := json.Unmarshal(data, output); err != nil {
		return fmt.Errorf("decode GitHub response: %w", err)
	}
	return nil
}

func prepareAIDiff(raw string, config AIConfig) ([]aiDiffFile, []string, int) {
	sections := strings.Split(raw, "\ndiff --git ")
	var files []aiDiffFile
	var excluded []string
	redactions := 0
	for sectionIndex, section := range sections {
		if sectionIndex > 0 {
			section = "diff --git " + section
		}
		file, ok := parseAIDiffFile(section)
		if !ok {
			continue
		}
		reason := aiExcludedReason(file, config)
		if reason != "" {
			excluded = append(excluded, file.Path+" ("+reason+")")
			continue
		}
		redacted, count := redactAISecrets(file.Text)
		file.Text = sanitizeAIControls(redacted)
		redactions += count
		files = append(files, file)
	}
	return files, excluded, redactions
}

func parseAIDiffFile(section string) (aiDiffFile, bool) {
	scanner := bufio.NewScanner(strings.NewReader(section))
	scanner.Buffer(make([]byte, 4096), 2<<20)
	file := aiDiffFile{
		ChangedLines: make(map[int]bool), DeletedLines: make(map[int]bool), Text: section,
	}
	oldPath := ""
	oldLine := 0
	newLine := 0
	inHunk := false
	for scanner.Scan() {
		line := scanner.Text()
		if strings.HasPrefix(line, "--- ") && !inHunk {
			candidate := strings.TrimPrefix(line, "--- ")
			if strings.HasPrefix(candidate, `"`) {
				if unquoted, err := strconv.Unquote(candidate); err == nil {
					candidate = unquoted
				}
			}
			if strings.HasPrefix(candidate, "a/") {
				oldPath = strings.TrimPrefix(candidate, "a/")
			}
		}
		if strings.HasPrefix(line, "+++ ") && !inHunk {
			candidate := strings.TrimPrefix(line, "+++ ")
			if strings.HasPrefix(candidate, `"`) {
				if unquoted, err := strconv.Unquote(candidate); err == nil {
					candidate = unquoted
				}
			}
			if strings.HasPrefix(candidate, "b/") {
				file.Path = strings.TrimPrefix(candidate, "b/")
			}
		}
		if strings.HasPrefix(line, "@@ ") {
			parts := strings.Fields(line)
			if len(parts) >= 3 {
				oldRange := strings.TrimPrefix(parts[1], "-")
				oldStart := strings.SplitN(oldRange, ",", 2)[0]
				oldLine, _ = strconv.Atoi(oldStart)
				rangePart := strings.TrimPrefix(parts[2], "+")
				start := strings.SplitN(rangePart, ",", 2)[0]
				newLine, _ = strconv.Atoi(start)
				inHunk = true
			}
			continue
		}
		if !inHunk || line == "" {
			continue
		}
		switch line[0] {
		case '+':
			file.ChangedLines[newLine] = true
			newLine++
		case '-':
			file.DeletedLines[oldLine] = true
			oldLine++
		default:
			oldLine++
			newLine++
		}
	}
	if file.Path == "" {
		file.Path = oldPath
	}
	if file.Path == "" || file.Path == "/dev/null" {
		return aiDiffFile{}, false
	}
	file.Path = filepath.ToSlash(filepath.Clean(file.Path))
	if strings.HasPrefix(file.Path, "../") || filepath.IsAbs(file.Path) {
		return aiDiffFile{}, false
	}
	return file, len(file.ChangedLines) > 0 || len(file.DeletedLines) > 0
}

func aiExcludedReason(file aiDiffFile, config AIConfig) string {
	if len(file.Text) > config.MaxFileBytes {
		return "oversized"
	}
	if strings.Contains(file.Text, "GIT binary patch") ||
		strings.Contains(file.Text, "Binary files ") || strings.IndexByte(file.Text, 0) >= 0 {
		return "binary"
	}
	if aiPathMatches(file.Path, config.SensitivePaths) {
		return "sensitive"
	}
	if aiPathMatches(file.Path, config.Exclude) {
		return "excluded"
	}
	return ""
}

func aiPathMatches(filePath string, patterns []string) bool {
	lower := strings.ToLower(filepath.ToSlash(filePath))
	for _, pattern := range patterns {
		pattern = filepath.ToSlash(pattern)
		if strings.HasSuffix(pattern, "/") {
			directory := strings.ToLower(pattern)
			if strings.HasPrefix(lower, directory) || strings.Contains(lower, "/"+directory) {
				return true
			}
		}
		if matched, _ := path.Match(strings.ToLower(pattern), path.Base(lower)); matched {
			return true
		}
		if matched, _ := path.Match(strings.ToLower(pattern), lower); matched {
			return true
		}
	}
	return false
}