update QoL and ai integration

ai_test.go

@@ -3,11 +3,15 @@ package main
 import (
 	"context"
 	"encoding/json"
+	"errors"
+	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"os"
 	"path/filepath"
+	"slices"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -26,6 +30,10 @@ func TestAIDefaultsAreDisabledAndBounded(t *testing.T) {
 	if config.AI.MaxRunBytes < config.AI.MaxRequestBytes || config.AI.MaxCalls < 1 {
 		t.Fatalf("unbounded defaults: %#v", config.AI)
 	}
+	if config.AI.MaxContextRounds != 2 || config.AI.MaxContextFiles != 8 ||
+		len(config.AI.SensitivePaths) == 0 {
+		t.Fatalf("focused-context defaults: %#v", config.AI)
+	}
 }
 
 func TestPrepareAIDiffFiltersAndRedacts(t *testing.T) {
@@ -160,6 +168,54 @@ func TestPullRequestDiffUsesAuthenticatedGHESRESTEndpoint(t *testing.T) {
 	}
 }
 
+func TestRepositoryTreeAndBlobUseAuthenticatedExactCommitEndpoints(t *testing.T) {
+	var requests []string
+	server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
+		if request.Header.Get("Authorization") != "Bearer token" {
+			t.Fatalf("authorization = %q", request.Header.Get("Authorization"))
+		}
+		requests = append(requests, request.URL.RequestURI())
+		switch request.URL.Path {
+		case "/api/v3/repos/owner/repo/git/commits/head-oid":
+			_, _ = writer.Write([]byte(`{"tree":{"sha":"tree-oid"}}`))
+		case "/api/v3/repos/owner/repo/git/trees/tree-oid":
+			_, _ = writer.Write([]byte(`{"truncated":false,"tree":[
+				{"path":"main.go","type":"blob","mode":"100644","sha":"blob-oid","size":13},
+				{"path":"nested","type":"tree","sha":"nested-tree"}
+			]}`))
+		case "/api/v3/repos/owner/repo/git/blobs/blob-oid":
+			_, _ = writer.Write([]byte(`{
+				"sha":"blob-oid","size":13,"encoding":"base64",
+				"content":"cGFja2FnZSBtYWluCg=="
+			}`))
+		default:
+			t.Fatalf("unexpected request: %s", request.URL.RequestURI())
+		}
+	}))
+	defer server.Close()
+	client := NewGitHubClient(server.URL+"/api/graphql", "token")
+	tree, err := client.RepositoryTree(context.Background(), "owner", "repo", "head-oid")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if tree.CommitOID != "head-oid" || tree.Truncated || len(tree.Entries) != 2 ||
+		tree.Entries[0].Path != "main.go" || tree.Entries[0].OID != "blob-oid" ||
+		tree.Entries[0].Mode != "100644" {
+		t.Fatalf("tree = %#v", tree)
+	}
+	content, err := client.RepositoryBlob(context.Background(), "owner", "repo", "blob-oid", 100)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if string(content) != "package main\n" {
+		t.Fatalf("blob = %q", content)
+	}
+	if len(requests) != 3 || requests[1] !=
+		"/api/v3/repos/owner/repo/git/trees/tree-oid?recursive=1" {
+		t.Fatalf("requests = %#v", requests)
+	}
+}
+
 func TestAIStoreIsPrivateAndMarksOldHeadOutdated(t *testing.T) {
 	store := NewAIStore(t.TempDir())
 	pr := PRDetails{PullRequest: PullRequest{Owner: "o", Repository: "r", Number: 7}, HeadOID: "head-1"}
@@ -225,6 +281,340 @@ func (s fakeAIDiffService) PullRequestDiff(context.Context, string, string, int)
 	return s.diff, nil
 }
 
+type fakeAIRepositoryService struct {
+	mu        sync.Mutex
+	diffCalls int
+	treeCalls []string
+	blobCalls []string
+	trees     map[string]AIRepositoryTree
+	blobs     map[string][]byte
+}
+
+func (s *fakeAIRepositoryService) PullRequestDiff(
+	context.Context, string, string, int,
+) (string, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.diffCalls++
+	return "", errors.New("focused discussion must not request the PR diff")
+}
+
+func (s *fakeAIRepositoryService) RepositoryTree(
+	_ context.Context, _, _, commitOID string,
+) (AIRepositoryTree, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.treeCalls = append(s.treeCalls, commitOID)
+	tree, ok := s.trees[commitOID]
+	if !ok {
+		return AIRepositoryTree{}, errors.New("tree not found")
+	}
+	return tree, nil
+}
+
+func (s *fakeAIRepositoryService) RepositoryBlob(
+	_ context.Context, _, _, oid string, maxBytes int,
+) ([]byte, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.blobCalls = append(s.blobCalls, oid)
+	content, ok := s.blobs[oid]
+	if !ok {
+		return nil, errors.New("blob not found")
+	}
+	if len(content) > maxBytes {
+		return nil, errors.New("blob too large")
+	}
+	return append([]byte(nil), content...), nil
+}
+
+type scriptedAIProvider struct {
+	responses []AIInferenceResponse
+	requests  []AIInferenceRequest
+}
+
+func (p *scriptedAIProvider) Name() string { return "scripted" }
+func (p *scriptedAIProvider) Status(context.Context) AIProviderStatus {
+	return AIProviderStatus{Ready: true, Model: "gpt-test"}
+}
+func (p *scriptedAIProvider) Generate(
+	_ context.Context, request AIInferenceRequest,
+) (AIInferenceResponse, error) {
+	p.requests = append(p.requests, request)
+	if len(p.responses) == 0 {
+		return AIInferenceResponse{}, errors.New("unexpected provider call")
+	}
+	response := p.responses[0]
+	p.responses = p.responses[1:]
+	return response, nil
+}
+
+func focusedAIRepository() *fakeAIRepositoryService {
+	return &fakeAIRepositoryService{
+		trees: map[string]AIRepositoryTree{
+			"head": {
+				CommitOID: "head",
+				Entries: []AIRepositoryEntry{
+					{Path: ".env", OID: "env", Type: "blob", Size: 10},
+					{Path: "go.sum", OID: "sum", Type: "blob", Size: 10},
+					{Path: "helper.go", OID: "helper", Type: "blob", Size: 24},
+					{Path: "main.go", OID: "main", Type: "blob", Size: 40},
+				},
+			},
+		},
+		blobs: map[string][]byte{
+			"main":   []byte("package main\n\nfunc selected() {}\n"),
+			"helper": []byte("package main\n\nfunc helper() {}\n"),
+		},
+	}
+}
+
+func focusedAIDetails() PRDetails {
+	return PRDetails{
+		PullRequest: PullRequest{
+			Owner: "owner", Repository: "repo", Number: 7, Title: "Focused change",
+		},
+		Body: "PR-BODY-MUST-NOT-BE-SENT", BaseRef: "main", BaseOID: "base",
+		HeadRef: "feature", HeadOID: "head",
+		Threads: []ReviewThread{
+			{
+				ID: "selected", Path: "main.go", Line: 3,
+				Comments: []ReviewComment{{
+					Author: "reviewer", Body: "SELECTED-THREAD-CONTEXT",
+					DiffHunk: "@@ -1 +1 @@\n-old\n+new", OriginalCommitOID: "old",
+				}},
+			},
+			{
+				ID: "unrelated", Path: "other.go",
+				Comments: []ReviewComment{{Author: "other", Body: "UNRELATED-THREAD-CONTEXT"}},
+			},
+		},
+	}
+}
+
+func TestFocusedAIPrepareUsesOnlySelectedThreadAndExactHeadContext(t *testing.T) {
+	repository := focusedAIRepository()
+	config := defaultAIConfig()
+	config.Enabled = true
+	controller := &AIController{
+		config: config, provider: &scriptedAIProvider{},
+		diffs: repository, repository: repository, store: NewAIStore(t.TempDir()),
+	}
+	preview, err := controller.Prepare(
+		context.Background(), focusedAIDetails(), "selected", "Please explain this.",
+	)
+	if err != nil {
+		t.Fatal(err)
+	}
+	prompt := preview.thread.basePrompt
+	for _, wanted := range []string{
+		"SELECTED-THREAD-CONTEXT", "Please explain this.", "package main",
+		"helper.go", "go.sum [unavailable: excluded]", "TARGET FILE main.go",
+	} {
+		if !strings.Contains(prompt, wanted) {
+			t.Fatalf("focused prompt is missing %q:\n%s", wanted, prompt)
+		}
+	}
+	for _, forbidden := range []string{
+		"UNRELATED-THREAD-CONTEXT", "PR-BODY-MUST-NOT-BE-SENT", ".env",
+	} {
+		if strings.Contains(prompt, forbidden) {
+			t.Fatalf("focused prompt leaked %q:\n%s", forbidden, prompt)
+		}
+	}
+	if repository.diffCalls != 0 || preview.Calls != 3 ||
+		preview.ContextRounds != 2 || preview.ContextFiles != 8 ||
+		preview.TreeHidden != 1 || preview.TreeUnavailable != 1 {
+		t.Fatalf("preview=%#v diffCalls=%d", preview, repository.diffCalls)
+	}
+}
+
+func TestFocusedAIDiscussionFetchesRequestedFilesAndStoresAnswer(t *testing.T) {
+	repository := focusedAIRepository()
+	provider := &scriptedAIProvider{responses: []AIInferenceResponse{
+		{
+			Model: "gpt-test",
+			Content: []byte(`{"action":"request_files","answer":"","requested_files":[
+				{"path":"helper.go","reason":"Need the helper"},
+				{"path":"../secret","reason":"Invalid"}
+			]}`),
+		},
+		{
+			Model: "gpt-test",
+			Content: []byte(`{
+				"action":"answer","answer":"The helper confirms the behavior.",
+				"requested_files":[]
+			}`),
+		},
+	}}
+	config := defaultAIConfig()
+	config.Enabled = true
+	store := NewAIStore(t.TempDir())
+	controller := &AIController{
+		config: config, provider: provider, diffs: repository,
+		repository: repository, store: store,
+	}
+	details := focusedAIDetails()
+	preview, err := controller.Prepare(
+		context.Background(), details, "selected", "Please explain this.",
+	)
+	if err != nil {
+		t.Fatal(err)
+	}
+	result, err := controller.Run(context.Background(), preview)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(provider.requests) != 2 ||
+		!strings.Contains(provider.requests[1].Prompt, "func helper()") ||
+		!strings.Contains(provider.requests[1].Prompt, "../secret: unavailable (invalid path)") {
+		t.Fatalf("provider requests = %#v", provider.requests)
+	}
+	if repository.diffCalls != 0 || result.Comments != 1 ||
+		result.Timing.Calls != 2 || result.Timing.Files != 1 {
+		t.Fatalf("result=%#v diffCalls=%d", result, repository.diffCalls)
+	}
+	thread := result.Details.Threads[0]
+	if len(thread.Comments) != 3 ||
+		thread.Comments[1].Origin != reviewOriginLocalAIUser ||
+		thread.Comments[2].Body != "The helper confirms the behavior." {
+		t.Fatalf("stored discussion = %#v", thread.Comments)
+	}
+}
+
+func TestFocusedAIPrepareFallsBackToReviewCommitForDeletedTarget(t *testing.T) {
+	repository := focusedAIRepository()
+	repository.trees["head"] = AIRepositoryTree{
+		CommitOID: "head",
+		Entries:   []AIRepositoryEntry{{Path: "helper.go", OID: "helper", Type: "blob", Size: 24}},
+	}
+	repository.trees["old"] = AIRepositoryTree{
+		CommitOID: "old",
+		Entries:   []AIRepositoryEntry{{Path: "main.go", OID: "old-main", Type: "blob", Size: 20}},
+	}
+	repository.blobs["old-main"] = []byte("package old\n")
+	config := defaultAIConfig()
+	config.Enabled = true
+	controller := &AIController{
+		config: config, provider: &scriptedAIProvider{},
+		diffs: repository, repository: repository, store: NewAIStore(t.TempDir()),
+	}
+	preview, err := controller.Prepare(
+		context.Background(), focusedAIDetails(), "selected", "What happened?",
+	)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if preview.InitialRevision != "old" ||
+		!strings.Contains(preview.thread.basePrompt, "package old") ||
+		!slices.Equal(repository.treeCalls, []string{"head", "old"}) {
+		t.Fatalf("preview=%#v treeCalls=%#v", preview, repository.treeCalls)
+	}
+}
+
+func TestFocusedAIPrepareRejectsSensitiveTarget(t *testing.T) {
+	repository := focusedAIRepository()
+	details := focusedAIDetails()
+	details.Threads[0].Path = ".env"
+	config := defaultAIConfig()
+	config.Enabled = true
+	controller := &AIController{
+		config: config, provider: &scriptedAIProvider{},
+		diffs: repository, repository: repository, store: NewAIStore(t.TempDir()),
+	}
+	_, err := controller.Prepare(context.Background(), details, "selected", "Explain.")
+	if err == nil || !strings.Contains(err.Error(), "sensitive path") ||
+		len(repository.treeCalls) != 0 {
+		t.Fatalf("err=%v treeCalls=%#v", err, repository.treeCalls)
+	}
+}
+
+func TestFocusedAIDiscussionEnforcesConfiguredFileAndRoundLimits(t *testing.T) {
+	repository := focusedAIRepository()
+	headTree := repository.trees["head"]
+	requests := make([]string, 0, 9)
+	for index := range 9 {
+		path := fmt.Sprintf("extra-%d.go", index)
+		oid := fmt.Sprintf("extra-%d", index)
+		requests = append(requests, fmt.Sprintf(
+			`{"path":%q,"reason":"Need context"}`, path,
+		))
+		headTree.Entries = append(
+			headTree.Entries,
+			AIRepositoryEntry{Path: path, OID: oid, Type: "blob", Size: 10},
+		)
+		repository.blobs[oid] = []byte("package x\n")
+	}
+	repository.trees["head"] = headTree
+	provider := &scriptedAIProvider{responses: []AIInferenceResponse{
+		{
+			Model: "gpt-test",
+			Content: []byte(fmt.Sprintf(
+				`{"action":"request_files","answer":"","requested_files":[%s]}`,
+				strings.Join(requests, ","),
+			)),
+		},
+		{Model: "gpt-test", Content: []byte(`{"answer":"Final bounded answer."}`)},
+	}}
+	config := defaultAIConfig()
+	config.Enabled = true
+	config.MaxContextRounds = 1
+	config.MaxContextFiles = 8
+	store := NewAIStore(t.TempDir())
+	controller := &AIController{
+		config: config, provider: provider, diffs: repository,
+		repository: repository, store: store,
+	}
+	preview, err := controller.Prepare(
+		context.Background(), focusedAIDetails(), "selected", "Investigate.",
+	)
+	if err != nil {
+		t.Fatal(err)
+	}
+	result, err := controller.Run(context.Background(), preview)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(provider.requests) != 2 || result.Timing.Calls != 2 ||
+		result.Timing.Files != 8 {
+		t.Fatalf("result=%#v requests=%d", result, len(provider.requests))
+	}
+	finalPrompt := provider.requests[1].Prompt
+	if !strings.Contains(finalPrompt, "extra-7.go") ||
+		!strings.Contains(finalPrompt, "extra-8.go: unavailable (file limit)") {
+		t.Fatalf("final prompt did not enforce file limit:\n%s", finalPrompt)
+	}
+	if !strings.Contains(string(provider.requests[1].Schema), `"required":["answer"]`) {
+		t.Fatalf("final call did not use answer-only schema: %s", provider.requests[1].Schema)
+	}
+}
+
+func TestAIConfirmationShowsAutomaticContextConsent(t *testing.T) {
+	config := defaultAIConfig()
+	config.Enabled = true
+	controller := &AIController{
+		config: config, provider: &scriptedAIProvider{},
+	}
+	app := NewApp(nil, "", "", false, 10, time.Minute)
+	app.width, app.height = 100, 35
+	app.ai, app.aiMode = controller, aiConfirm
+	app.aiPreview = AIPreview{
+		Files: 1, Bytes: 1200, Calls: 3, HeadOID: "head-oid", Model: "gpt-test",
+		Included: []string{"main.go"}, TreeEntries: 20, TreeUnavailable: 3,
+		TreeHidden: 2, TreeTruncated: true, ContextRounds: 2, ContextFiles: 8,
+		InitialRevision: "head-oid", thread: &aiThreadPrepared{},
+	}
+	plain := strings.Join(strings.Fields(ansi.Strip(app.viewAI())), " ")
+	for _, wanted := range []string{
+		"20 visible tree entries", "3 unavailable", "2 sensitive hidden",
+		"truncated", "2 automatic request round(s)", "and 8 additional",
+	} {
+		if !strings.Contains(plain, wanted) {
+			t.Fatalf("confirmation is missing %q:\n%s", wanted, plain)
+		}
+	}
+}
+
 func TestAIControllerAcceptsOnlyChangedLinesAndDeduplicates(t *testing.T) {
 	output := aiOutput{}
 	output.Findings = append(output.Findings, aiFinding{
@@ -475,6 +865,40 @@ func TestAIDiscussionStoresUserMessageBeforeProviderResponse(t *testing.T) {
 	}
 }
 
+func TestAIAnnotationsMergeIntoThreadTimelineByCreationTime(t *testing.T) {
+	rootTime := time.Date(2026, time.July, 29, 15, 0, 0, 0, time.Local)
+	aiTime := rootTime.Add(40 * time.Minute)
+	replyTime := rootTime.Add(58 * time.Minute)
+	pr := PRDetails{Threads: []ReviewThread{{
+		ID: "thread-1",
+		Comments: []ReviewComment{
+			{ID: "root", Body: "Root", CreatedAt: rootTime},
+			{ID: "remote-reply", Body: "Remote reply", CreatedAt: replyTime},
+		},
+	}}}
+	state := &aiStoredState{
+		Version: 1,
+		Annotations: map[string][]ReviewComment{
+			"thread-1": {{
+				ID: "local-ai-comment", Body: "Earlier AI discussion",
+				CreatedAt: aiTime, Origin: reviewOriginLocalAI,
+			}},
+		},
+	}
+
+	merged := state.Merge(pr)
+	comments := merged.Threads[0].Comments
+	if len(comments) != 3 ||
+		comments[0].ID != "root" ||
+		comments[1].ID != "local-ai-comment" ||
+		comments[2].ID != "remote-reply" {
+		t.Fatalf("merged timeline = %#v", comments)
+	}
+	if pr.Threads[0].Comments[1].ID != "remote-reply" {
+		t.Fatal("merge mutated the GitHub thread")
+	}
+}
+
 func TestAIDiscussionUsesViewerGitHubLogin(t *testing.T) {
 	state := &aiStoredState{
 		Version: 1, Annotations: make(map[string][]ReviewComment),