feat(agent): integrate session recording and token tracking for re-location task

Re-location LLM calls were invisible in session history and their token
  consumption was unaccounted for. Refactor the code_comment handling path
  to record re-location requests/responses into session JSONL, accumulate
  token usage into global counters, and display them in the viewer.

  Also fixes: async context cancellation risk (WithoutCancel), duplicate
  telemetry recording, missing timeout enforcement, missing {existing_code}
  placeholder in the re-location prompt, and renames Parse to ParseComments
  for clarity.

internal/agent/agent.go

@@ -842,22 +842,67 @@ func (a *Agent) executeToolCall(ctx context.Context, newPath string, call llm.To
 
 	startTime := time.Now()
 
-	// Async path for code_comment when worker pool is configured
-	// Mirrors Java: pendingCommentFutures.add(subtaskExecutor.submit(() -> getCodeComments(...)))
-	if t == tool.CodeComment && a.args.CommentWorkerPool != nil {
-		if rec != nil {
-			rec.AddToolResult(t.Name(), call.Function.Arguments, "(async)")
-		}
-		pool := a.args.CommentWorkerPool
+	// code_comment: parse → resolve line numbers → re-locate if needed → add to collector
+	if t == tool.CodeComment {
 		telemetry.PrintToolCallStarted(t.Name(), args)
-		pool.Submit(func() ([]model.LlmComment, error) {
-			_, _ = p.Execute(args) // provider parses args and adds to collector
-			telemetry.RecordToolCall(ctx, t.Name(), time.Since(startTime), true)
-			return []model.LlmComment{}, nil
-		})
-		// Return immediate success - actual comment processing continues off
-		// the critical path, exactly like Java's subtaskExecutor.submit for CODE_COMMENT.
-		telemetry.RecordToolCall(ctx, t.Name(), time.Since(startTime), true)
+
+		comments, errMsg := tool.ParseComments(args)
+		if errMsg != "" {
+			telemetry.RecordToolCall(ctx, t.Name(), time.Since(startTime), false)
+			return tool.Of(errMsg)
+		}
+
+		resolveAndCollect := func(rctx context.Context) {
+			for i := range comments {
+				cm := &comments[i]
+				d := a.findDiff(cm.Path)
+				if d != nil {
+					if !diff.ResolveComment(cm, d) && a.args.Template.ReLocationTask != nil {
+						rlStart := time.Now()
+						_, resp, msgs := diff.ReLocateComment(rctx, cm, d, a.args.LLMClient, a.args.Template.ReLocationTask, a.args.Model, a.args.Template.MaxTokens)
+						if msgs != nil {
+							fs := a.session.GetOrCreateFileSession(cm.Path)
+							rlRec := fs.AppendTaskRecord(session.ReLocationTask, msgs)
+							if resp != nil {
+								rlRec.SetResponse(resp, time.Since(rlStart))
+								if resp.Usage != nil {
+									atomic.AddInt64(&a.totalTokensUsed, int64(resp.Usage.TotalTokens))
+									atomic.AddInt64(&a.totalInputTokens, int64(resp.Usage.PromptTokens+resp.Usage.CacheReadTokens))
+									atomic.AddInt64(&a.totalOutputTokens, int64(resp.Usage.CompletionTokens+resp.Usage.CacheWriteTokens))
+								}
+							} else {
+								rlRec.SetError(fmt.Errorf("re-location LLM call failed"), time.Since(rlStart))
+							}
+						}
+					}
+				}
+				a.args.CommentCollector.Add(*cm)
+			}
+		}
+
+		if a.args.CommentWorkerPool != nil {
+			if rec != nil {
+				rec.AddToolResult(t.Name(), call.Function.Arguments, "(async)")
+			}
+			pool := a.args.CommentWorkerPool
+			asyncCtx := context.WithoutCancel(ctx)
+			toolName := t.Name()
+			pool.Submit(func() ([]model.LlmComment, error) {
+				resolveAndCollect(asyncCtx)
+				telemetry.PrintToolCallFinished(toolName, time.Since(startTime))
+				return []model.LlmComment{}, nil
+			})
+			telemetry.RecordToolCall(asyncCtx, toolName, time.Since(startTime), true)
+			return tool.Of(tool.CommentSucceed)
+		}
+
+		resolveAndCollect(ctx)
+		dur := time.Since(startTime)
+		telemetry.RecordToolCall(ctx, t.Name(), dur, true)
+		telemetry.PrintToolCallFinished(t.Name(), dur)
+		if rec != nil {
+			rec.AddToolResult(t.Name(), call.Function.Arguments, tool.CommentSucceed)
+		}
 		return tool.Of(tool.CommentSucceed)
 	}
 
@@ -879,6 +924,16 @@ func (a *Agent) executeToolCall(ctx context.Context, newPath string, call llm.To
 	return tool.Of(result)
 }
 
+// findDiff returns the Diff for the given file path, or nil if not found.
+func (a *Agent) findDiff(path string) *model.Diff {
+	for i := range a.diffs {
+		if a.diffs[i].NewPath == path || a.diffs[i].OldPath == path {
+			return &a.diffs[i]
+		}
+	}
+	return nil
+}
+
 // collectPendingComments waits for any async workers then returns aggregated comments from the collector.
 func (a *Agent) collectPendingComments() []model.LlmComment {
 	if a.args.CommentWorkerPool != nil {

internal/config/template/task_template.json

@@ -42,5 +42,18 @@
   "MAX_TOOL_REQUEST_TIMES": 20,
   "MAX_SUBTASK_EXECUTION_TIME_MINUTES": 5,
   "PLAN_MODE_LINE_THRESHOLD": 50,
-  "MAX_TOKENS": 58888
+  "MAX_TOKENS": 58888,
+  "RE_LOCATION_TASK": {
+    "messages": [
+      {
+        "role": "system",
+        "content": "You are a code location assistant. Given a unified diff and a review comment, your sole task is to extract the exact code snippet from the diff that the comment refers to. /no_think"
+      },
+      {
+        "role": "user",
+        "content": "Below is a unified diff and a review comment. Identify the minimal contiguous code range in the diff that the comment targets.\n\nRules:\n1. Copy the relevant lines VERBATIM from the diff — do not rewrite, reformat, or add anything.\n2. Strip leading diff markers (`+`, `-`, ` `) from each line before outputting.\n3. Include only the lines directly related to the issue — no surrounding context.\n4. If multiple disjoint locations apply, pick the single most relevant one.\n5. Output ONLY a fenced code block. No explanation, no commentary.\n\n**Diff:**\n```diff\n{diff}```\n\n**Original code snippet (failed to match):**\n```\n{existing_code}\n```\n\n**Review comment:**\n{suggestion_content}"
+      }
+    ],
+    "timeout": 60
+  }
 }

internal/config/template/template.go

@@ -18,6 +18,7 @@ type Template struct {
 	MaxToolRequestTimes   int              `json:"MAX_TOOL_REQUEST_TIMES"`
 	MaxSubtaskExecMinutes int              `json:"MAX_SUBTASK_EXECUTION_TIME_MINUTES"`
 	PlanModeLineThreshold int              `json:"PLAN_MODE_LINE_THRESHOLD"`
+	ReLocationTask        *LlmConversation `json:"RE_LOCATION_TASK,omitempty"`
 }
 
 //go:embed task_template.json

internal/diff/relocation.go

@@ -0,0 +1,91 @@
+package diff
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"time"
+
+	"github.com/open-code-review/open-code-review/internal/config/template"
+	"github.com/open-code-review/open-code-review/internal/llm"
+	"github.com/open-code-review/open-code-review/internal/model"
+	"github.com/open-code-review/open-code-review/internal/stdout"
+)
+
+// ReLocateComment calls the LLM to regenerate a precise existing_code snippet
+// when text-based matching fails, then retries ResolveComment with the new snippet.
+// Returns (success, response, requestMessages) so the caller can record session
+// history and track token usage. Response and messages are nil on early exits.
+func ReLocateComment(
+	ctx context.Context,
+	cm *model.LlmComment,
+	d *model.Diff,
+	client llm.LLMClient,
+	task *template.LlmConversation,
+	modelName string,
+	maxTokens int,
+) (bool, *llm.ChatResponse, []llm.Message) {
+	if task == nil || len(task.Messages) == 0 {
+		return false, nil, nil
+	}
+
+	if task.Timeout > 0 {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(ctx, time.Duration(task.Timeout)*time.Second)
+		defer cancel()
+	}
+
+	messages := make([]llm.Message, 0, len(task.Messages))
+	for _, m := range task.Messages {
+		content := m.Content
+		content = strings.ReplaceAll(content, "{diff}", d.Diff)
+		content = strings.ReplaceAll(content, "{existing_code}", cm.ExistingCode)
+		content = strings.ReplaceAll(content, "{suggestion_content}", cm.Content)
+		messages = append(messages, llm.NewTextMessage(m.Role, content))
+	}
+
+	resp, err := client.CompletionsWithCtx(ctx, llm.ChatRequest{
+		Model:     modelName,
+		Messages:  messages,
+		MaxTokens: maxTokens,
+	})
+	if err != nil {
+		fmt.Fprintf(stdout.Writer(), "[ocr] Re-location LLM call failed for %s: %v\n", cm.Path, err)
+		return false, nil, messages
+	}
+
+	code := extractCodeBlock(resp.Content())
+	if code == "" {
+		return false, resp, messages
+	}
+
+	original := cm.ExistingCode
+	cm.ExistingCode = code
+	if ResolveComment(cm, d) {
+		return true, resp, messages
+	}
+	cm.ExistingCode = original
+	return false, resp, messages
+}
+
+// extractCodeBlock extracts the content of the first fenced code block from text.
+// Returns empty string if no code block is found.
+func extractCodeBlock(text string) string {
+	text = strings.TrimSpace(text)
+	start := strings.Index(text, "```")
+	if start < 0 {
+		return ""
+	}
+	afterOpen := start + 3
+	// Skip optional language tag on the opening fence line.
+	if nl := strings.IndexByte(text[afterOpen:], '\n'); nl >= 0 {
+		afterOpen += nl + 1
+	} else {
+		return ""
+	}
+	end := strings.Index(text[afterOpen:], "```")
+	if end < 0 {
+		return ""
+	}
+	return strings.TrimSpace(text[afterOpen : afterOpen+end])
+}

internal/diff/relocation_test.go

@@ -0,0 +1,226 @@
+package diff
+
+import (
+	"context"
+	"errors"
+	"net/http"
+	"testing"
+
+	"github.com/open-code-review/open-code-review/internal/config/template"
+	"github.com/open-code-review/open-code-review/internal/llm"
+	"github.com/open-code-review/open-code-review/internal/model"
+)
+
+type mockLLMClient struct {
+	response *llm.ChatResponse
+	err      error
+}
+
+func (m *mockLLMClient) Completions(req llm.ChatRequest) (*llm.ChatResponse, error) {
+	return m.response, m.err
+}
+
+func (m *mockLLMClient) CompletionsWithCtx(_ context.Context, req llm.ChatRequest) (*llm.ChatResponse, error) {
+	return m.response, m.err
+}
+
+func (m *mockLLMClient) StreamCompletion(req llm.ChatRequest, cb func(chunk []byte) error) error {
+	return m.err
+}
+
+func newMockResponse(content string) *llm.ChatResponse {
+	return &llm.ChatResponse{
+		Choices: []llm.Choice{
+			{Message: llm.ResponseMessage{Role: "assistant", Content: &content}},
+		},
+		Headers: http.Header{},
+	}
+}
+
+func makeTask() *template.LlmConversation {
+	return &template.LlmConversation{
+		Timeout: 60,
+		Messages: []template.ChatMessage{
+			{Role: "system", Content: "you are a helper"},
+			{Role: "user", Content: "diff:\n{diff}\n\ncomment:\n{suggestion_content}"},
+		},
+	}
+}
+
+func makeDiff() *model.Diff {
+	return &model.Diff{
+		NewPath: "main.go",
+		Diff: `@@ -10,6 +10,8 @@
+ import "fmt"
+
+ func main() {
++    x := 1
++    y := 2
+     fmt.Println("hello")
+ }
+`,
+	}
+}
+
+func TestResolveComment_TextMatchSuccess(t *testing.T) {
+	cm := model.LlmComment{
+		Path:         "main.go",
+		Content:      "unused variable",
+		ExistingCode: "x := 1\ny := 2",
+	}
+	d := makeDiff()
+
+	ok := ResolveComment(&cm, d)
+	if !ok {
+		t.Fatal("expected ResolveComment to succeed")
+	}
+	if cm.StartLine == 0 || cm.EndLine == 0 {
+		t.Fatalf("expected non-zero lines, got %d-%d", cm.StartLine, cm.EndLine)
+	}
+}
+
+func TestResolveComment_AlreadyResolved(t *testing.T) {
+	cm := model.LlmComment{
+		Path:         "main.go",
+		Content:      "test",
+		ExistingCode: "whatever",
+		StartLine:    5,
+		EndLine:      10,
+	}
+	d := makeDiff()
+	ok := ResolveComment(&cm, d)
+	if !ok {
+		t.Fatal("expected true for already-resolved comment")
+	}
+	if cm.StartLine != 5 || cm.EndLine != 10 {
+		t.Fatal("should not change already-resolved lines")
+	}
+}
+
+func TestResolveComment_EmptyExistingCode(t *testing.T) {
+	cm := model.LlmComment{Path: "main.go", Content: "test"}
+	d := makeDiff()
+	ok := ResolveComment(&cm, d)
+	if ok {
+		t.Fatal("expected false for empty ExistingCode")
+	}
+}
+
+func TestReLocateComment_LLMReturnsValidCode(t *testing.T) {
+	cm := model.LlmComment{
+		Path:         "main.go",
+		Content:      "unused variable",
+		ExistingCode: "totally wrong code that won't match",
+	}
+	d := makeDiff()
+
+	client := &mockLLMClient{
+		response: newMockResponse("Here is the code:\n```go\nx := 1\ny := 2\n```\n"),
+	}
+
+	ok, resp, msgs := ReLocateComment(context.Background(), &cm, d, client, makeTask(), "test-model", 1000)
+	if !ok {
+		t.Fatal("expected re-location to succeed")
+	}
+	if resp == nil {
+		t.Fatal("expected non-nil response")
+	}
+	if len(msgs) == 0 {
+		t.Fatal("expected non-empty messages")
+	}
+	if cm.StartLine == 0 || cm.EndLine == 0 {
+		t.Fatalf("expected non-zero lines after re-location, got %d-%d", cm.StartLine, cm.EndLine)
+	}
+}
+
+func TestReLocateComment_LLMReturnsInvalidContent(t *testing.T) {
+	cm := model.LlmComment{
+		Path:         "main.go",
+		Content:      "unused variable",
+		ExistingCode: "totally wrong code",
+	}
+	d := makeDiff()
+
+	client := &mockLLMClient{
+		response: newMockResponse("I cannot find the code."),
+	}
+
+	ok, resp, msgs := ReLocateComment(context.Background(), &cm, d, client, makeTask(), "test-model", 1000)
+	if ok {
+		t.Fatal("expected re-location to fail for invalid LLM response")
+	}
+	if resp == nil {
+		t.Fatal("expected non-nil response even on failure")
+	}
+	if len(msgs) == 0 {
+		t.Fatal("expected non-empty messages")
+	}
+	if cm.StartLine != 0 || cm.EndLine != 0 {
+		t.Fatal("lines should remain 0-0")
+	}
+}
+
+func TestReLocateComment_LLMError(t *testing.T) {
+	cm := model.LlmComment{
+		Path:         "main.go",
+		Content:      "test",
+		ExistingCode: "bad code",
+	}
+	d := makeDiff()
+
+	client := &mockLLMClient{err: errors.New("network error")}
+
+	ok, resp, msgs := ReLocateComment(context.Background(), &cm, d, client, makeTask(), "test-model", 1000)
+	if ok {
+		t.Fatal("expected false on LLM error")
+	}
+	if resp != nil {
+		t.Fatal("expected nil response on error")
+	}
+	if len(msgs) == 0 {
+		t.Fatal("expected non-empty messages even on error")
+	}
+}
+
+func TestReLocateComment_NilTask(t *testing.T) {
+	cm := model.LlmComment{
+		Path:         "main.go",
+		Content:      "test",
+		ExistingCode: "bad code",
+	}
+	d := makeDiff()
+	client := &mockLLMClient{}
+
+	ok, resp, msgs := ReLocateComment(context.Background(), &cm, d, client, nil, "test-model", 1000)
+	if ok {
+		t.Fatal("expected false when task is nil")
+	}
+	if resp != nil {
+		t.Fatal("expected nil response when task is nil")
+	}
+	if msgs != nil {
+		t.Fatal("expected nil messages when task is nil")
+	}
+}
+
+func TestExtractCodeBlock(t *testing.T) {
+	tests := []struct {
+		name  string
+		input string
+		want  string
+	}{
+		{"with language tag", "```go\nfoo\nbar\n```", "foo\nbar"},
+		{"without language tag", "```\nfoo\n```", "foo"},
+		{"with surrounding text", "Here:\n```\ncode\n```\ndone", "code"},
+		{"no code block", "just text", ""},
+		{"empty block", "```\n```", ""},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := extractCodeBlock(tt.input)
+			if got != tt.want {
+				t.Errorf("extractCodeBlock() = %q, want %q", got, tt.want)
+			}
+		})
+	}
+}

internal/diff/resolver.go

@@ -54,6 +54,21 @@ func ResolveLineNumbers(comments []model.LlmComment, diffs []model.Diff) []model
 	return result
 }
 
+// ResolveComment attempts to resolve StartLine/EndLine for a single comment
+// by matching ExistingCode against the diff. Returns true on success.
+func ResolveComment(cm *model.LlmComment, d *model.Diff) bool {
+	if cm.StartLine > 0 || cm.EndLine > 0 {
+		return true
+	}
+	if cm.ExistingCode == "" {
+		return false
+	}
+	if resolveFromHunk(d, cm) {
+		return true
+	}
+	return resolveFromFileContent(d, cm)
+}
+
 // indexedLine pairs a normalized line with its absolute file line number.
 type indexedLine struct {
 	lineNum int

internal/session/history.go

@@ -18,6 +18,7 @@ const (
 	PlanTask              TaskType = "plan_task"
 	MainTask              TaskType = "main_task"
 	MemoryCompressionTask TaskType = "memory_compression_task"
+	ReLocationTask        TaskType = "re_location_task"
 )
 
 // SessionHistory is the top-level container for an entire CR run.

internal/tool/code_comment.go

@@ -19,8 +19,20 @@ func (p *CodeCommentProvider) Execute(args map[string]any) (string, error) {
 		return "Error: comment collector is not configured", nil
 	}
 
-	// Parse the "comments" array from the tool call arguments.
-	// LLM sometimes double-encodes the array as a JSON string.
+	comments, errMsg := ParseComments(args)
+	if errMsg != "" {
+		return errMsg, nil
+	}
+
+	for i := range comments {
+		p.Collector.Add(comments[i])
+	}
+	return CommentSucceed, nil
+}
+
+// ParseComments extracts LlmComment entries from tool call arguments without writing
+// to the Collector. Returns parsed comments and an error message (empty on success).
+func ParseComments(args map[string]any) ([]model.LlmComment, string) {
 	var rawComments []any
 	if arr, ok := args["comments"].([]any); ok && len(arr) > 0 {
 		rawComments = arr
@@ -29,9 +41,10 @@ func (p *CodeCommentProvider) Execute(args map[string]any) (string, error) {
 	}
 	if len(rawComments) == 0 {
 		raw, _ := json.Marshal(args)
-		return fmt.Sprintf("Error: 'comments' array is required. Got args: %s", string(raw)), nil
+		return nil, fmt.Sprintf("Error: 'comments' array is required. Got args: %s", string(raw))
 	}
 
+	var comments []model.LlmComment
 	for _, raw := range rawComments {
 		obj, ok := raw.(map[string]any)
 		if !ok {
@@ -60,7 +73,7 @@ func (p *CodeCommentProvider) Execute(args map[string]any) (string, error) {
 			continue
 		}
 
-		p.Collector.Add(cm)
+		comments = append(comments, cm)
 	}
-	return CommentSucceed, nil
+	return comments, ""
 }

internal/viewer/server.go

@@ -88,6 +88,8 @@ func parseTemplate(name string) (*template.Template, error) {
 				return "task-main"
 			case MemoryCompressionTask:
 				return "task-memory"
+			case ReLocationTask:
+				return "task-relocation"
 			default:
 				return "task-default"
 			}
@@ -96,7 +98,7 @@ func parseTemplate(name string) (*template.Template, error) {
 			Type  TaskType
 			Cards []*TaskCard
 		} {
-			order := []TaskType{PlanTask, MainTask, MemoryCompressionTask}
+			order := []TaskType{PlanTask, MainTask, ReLocationTask, MemoryCompressionTask}
 			var result []struct {
 				Type  TaskType
 				Cards []*TaskCard
@@ -110,7 +112,7 @@ func parseTemplate(name string) (*template.Template, error) {
 				}
 			}
 			for tt, cards := range tasks {
-				if tt != PlanTask && tt != MainTask && tt != MemoryCompressionTask {
+				if tt != PlanTask && tt != MainTask && tt != ReLocationTask && tt != MemoryCompressionTask {
 					result = append(result, struct {
 						Type  TaskType
 						Cards []*TaskCard

internal/viewer/static/style.css

@@ -228,6 +228,7 @@ h3 { margin: 1.5rem 0 0.75rem; color: #24292e; }
 .task-main .task-type-dot { background: #0366d6; }
 .task-plan .task-type-dot { background: #8250df; }
 .task-memory .task-type-dot { background: #8b949e; }
+.task-relocation .task-type-dot { background: #e36209; }
 .task-default .task-type-dot { background: #57606a; }
 
 /* Card */
@@ -244,6 +245,7 @@ h3 { margin: 1.5rem 0 0.75rem; color: #24292e; }
 .task-main .card { border-left: 3px solid #0366d6; }
 .task-plan .card { border-left: 3px solid #8250df; }
 .task-memory .card { border-left: 3px solid #8b949e; }
+.task-relocation .card { border-left: 3px solid #e36209; }
 
 .card-header {
     background: #f6f8fa;

internal/viewer/store.go

@@ -208,6 +208,7 @@ const (
 	PlanTask              TaskType = "plan_task"
 	MainTask              TaskType = "main_task"
 	MemoryCompressionTask TaskType = "memory_compression_task"
+	ReLocationTask        TaskType = "re_location_task"
 )
 
 // TaskCard links an LLM request with its response and tool calls.