fix(llm): add context support to StreamCompletion and memory compression to prevent infinite timeouts

StreamCompletion used http.NewRequest without context and retry with
context.Background(), making streaming calls uncancellable. Memory
compression called Completions() without context, leaving async
goroutines unable to be cancelled or timed out.

- Add StreamCompletionWithCtx to LLMClient interface and both
  OpenAI/Anthropic implementations
- Propagate context through runCompression, addNextMessage, and
  triggerAsyncCompression call chain
- Add 5-minute timeout cap on async compression goroutines
- Cancel in-flight compression HTTP requests when job is superseded

internal/agent/agent.go

@@ -153,6 +153,7 @@ type partitionResult struct {
 type compressionJob struct {
 	done    chan struct{}
 	rebuilt []llm.Message
+	cancel  context.CancelFunc
 }
 
 // Agent orchestrates the AI-powered code review.
@@ -848,7 +849,7 @@ func (a *Agent) performLlmCodeReview(ctx context.Context, messages []llm.Message
 			consecutiveEmptyRounds = 0
 		}
 
-		succeed := a.addNextMessage(content, calls, results, &messages, newPath)
+		succeed := a.addNextMessage(ctx, content, calls, results, &messages, newPath)
 		if !succeed {
 			fmt.Fprintf(stdout.Writer(), "[ocr] Context compression exceeded threshold for %s, stopping.\n", newPath)
 			break
@@ -999,7 +1000,7 @@ func (a *Agent) collectPendingComments() []model.LlmComment {
 // Implements dual-threshold compression:
 //   - 60% of MaxTokens: trigger async background compression (non-blocking)
 //   - 80% of MaxTokens: perform synchronous compression immediately
-func (a *Agent) addNextMessage(assistantContent string, toolCalls []llm.ToolCall, results []tool.ToolCallResult, messages *[]llm.Message, filePath string) bool {
+func (a *Agent) addNextMessage(ctx context.Context, assistantContent string, toolCalls []llm.ToolCall, results []tool.ToolCallResult, messages *[]llm.Message, filePath string) bool {
 	maxAllowed := a.args.Template.MaxTokens
 	softLimit := int(float64(maxAllowed) * tokenSoftThreshold)
 	warnLimit := int(float64(maxAllowed) * tokenWarningThreshold)
@@ -1012,13 +1013,13 @@ func (a *Agent) addNextMessage(assistantContent string, toolCalls []llm.ToolCall
 	// Hard threshold: synchronous compression.
 	if tokenCount > warnLimit {
 		a.cancelPendingCompression()
-		*messages, _ = a.runCompression(*messages, filePath)
+		*messages, _ = a.runCompression(ctx, *messages, filePath)
 		tokenCount = countMessagesTokens(*messages)
 	}
 
 	// Soft threshold: async compression.
 	if tokenCount > softLimit && a.pendingJob == nil {
-		a.triggerAsyncCompression(*messages, filePath)
+		a.triggerAsyncCompression(ctx, *messages, filePath)
 	}
 
 	// Add assistant message with tool_calls when present.
@@ -1037,7 +1038,7 @@ func (a *Agent) addNextMessage(assistantContent string, toolCalls []llm.ToolCall
 	finalCount := countMessagesTokens(*messages)
 	if finalCount > warnLimit {
 		a.cancelPendingCompression()
-		*messages, _ = a.runCompression(*messages, filePath)
+		*messages, _ = a.runCompression(ctx, *messages, filePath)
 	}
 
 	return countMessagesTokens(*messages) < warnLimit
@@ -1154,7 +1155,7 @@ func stripMarkdownFences(s string) string {
 // runCompression performs three-zone memory compression on the given messages.
 // It summarizes the compress zone while preserving the active zone intact.
 // Returns the rebuilt messages slice: [frozen] + [compressed_summary] + [active].
-func (a *Agent) runCompression(msgs []llm.Message, filePath string) ([]llm.Message, error) {
+func (a *Agent) runCompression(ctx context.Context, msgs []llm.Message, filePath string) ([]llm.Message, error) {
 	if len(a.args.Template.MemoryCompressionTask.Messages) == 0 || len(msgs) <= 2 {
 		return msgs[:min(len(msgs), 2)], nil
 	}
@@ -1173,7 +1174,7 @@ func (a *Agent) runCompression(msgs []llm.Message, filePath string) ([]llm.Messa
 	}
 
 	startTime := time.Now()
-	resp, err := a.args.LLMClient.Completions(llm.ChatRequest{
+	resp, err := a.args.LLMClient.CompletionsWithCtx(ctx, llm.ChatRequest{
 		Model:     a.args.Model,
 		Messages:  compressionMsgs,
 		MaxTokens: a.args.Template.MaxTokens,
@@ -1214,16 +1215,19 @@ func (a *Agent) runCompression(msgs []llm.Message, filePath string) ([]llm.Messa
 }
 
 // triggerAsyncCompression kicks off a background compression job.
-func (a *Agent) triggerAsyncCompression(messages []llm.Message, filePath string) {
+func (a *Agent) triggerAsyncCompression(ctx context.Context, messages []llm.Message, filePath string) {
 	msgSnapshot := copyMessages(messages)
 
-	job := &compressionJob{done: make(chan struct{})}
+	asyncCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Minute)
+
+	job := &compressionJob{done: make(chan struct{}), cancel: cancel}
 	a.compressionMu.Lock()
 	a.pendingJob = job
 	a.compressionMu.Unlock()
 
 	go func() {
-		rebuilt, _ := a.runCompression(msgSnapshot, filePath)
+		defer cancel()
+		rebuilt, _ := a.runCompression(asyncCtx, msgSnapshot, filePath)
 
 		a.compressionMu.Lock()
 		defer a.compressionMu.Unlock()
@@ -1267,6 +1271,7 @@ func (a *Agent) cancelPendingCompression() {
 	defer a.compressionMu.Unlock()
 
 	if a.pendingJob != nil {
+		a.pendingJob.cancel()
 		a.pendingJob = nil
 	}
 }

internal/diff/relocation_test.go

@@ -28,6 +28,10 @@ func (m *mockLLMClient) StreamCompletion(req llm.ChatRequest, cb func(chunk []by
 	return m.err
 }
 
+func (m *mockLLMClient) StreamCompletionWithCtx(_ context.Context, req llm.ChatRequest, cb func(chunk []byte) error) error {
+	return m.err
+}
+
 func newMockResponse(content string) *llm.ChatResponse {
 	return &llm.ChatResponse{
 		Choices: []llm.Choice{

internal/llm/client.go

@@ -37,6 +37,7 @@ type LLMClient interface {
 	Completions(req ChatRequest) (*ChatResponse, error)
 	CompletionsWithCtx(ctx context.Context, req ChatRequest) (*ChatResponse, error)
 	StreamCompletion(req ChatRequest, cb func(chunk []byte) error) error
+	StreamCompletionWithCtx(ctx context.Context, req ChatRequest, cb func(chunk []byte) error) error
 }
 
 // --- Shared data types ---
@@ -356,6 +357,11 @@ func (c *OpenAIClient) GeneralRequestWithCtx(ctx context.Context, messages []Mes
 
 // StreamCompletion initiates a streaming chat completion. The callback is invoked per chunk.
 func (c *OpenAIClient) StreamCompletion(req ChatRequest, cb func(chunk []byte) error) error {
+	return c.StreamCompletionWithCtx(context.Background(), req, cb)
+}
+
+// StreamCompletionWithCtx initiates a streaming chat completion with context support for cancellation and timeout.
+func (c *OpenAIClient) StreamCompletionWithCtx(ctx context.Context, req ChatRequest, cb func(chunk []byte) error) error {
 	req.Stream = true
 
 	model := req.Model
@@ -363,7 +369,7 @@ func (c *OpenAIClient) StreamCompletion(req ChatRequest, cb func(chunk []byte) e
 		model = c.cfg.Model
 	}
 
-	return c.withRetry(func() error {
+	return c.withRetryCtx(ctx, func() error {
 		body := make(map[string]any)
 		b, _ := json.Marshal(req)
 		json.Unmarshal(b, &body)
@@ -373,7 +379,7 @@ func (c *OpenAIClient) StreamCompletion(req ChatRequest, cb func(chunk []byte) e
 		}
 
 		payload, _ := json.Marshal(body)
-		httpReq, err := http.NewRequest(http.MethodPost, c.cfg.URL, bytes.NewReader(payload))
+		httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, c.cfg.URL, bytes.NewReader(payload))
 		if err != nil {
 			return fmt.Errorf("create request: %w", err)
 		}
@@ -528,6 +534,11 @@ func (c *AnthropicClient) CompletionsWithCtx(ctx context.Context, req ChatReques
 // StreamCompletion initiates a streaming chat completion using SSE. The callback
 // is invoked per chunk with raw JSON data stripped of the "data: " prefix.
 func (c *AnthropicClient) StreamCompletion(req ChatRequest, cb func(chunk []byte) error) error {
+	return c.StreamCompletionWithCtx(context.Background(), req, cb)
+}
+
+// StreamCompletionWithCtx initiates a streaming chat completion with context support for cancellation and timeout.
+func (c *AnthropicClient) StreamCompletionWithCtx(ctx context.Context, req ChatRequest, cb func(chunk []byte) error) error {
 	req.Stream = true
 
 	model := req.Model
@@ -535,7 +546,7 @@ func (c *AnthropicClient) StreamCompletion(req ChatRequest, cb func(chunk []byte
 		model = c.cfg.Model
 	}
 
-	return c.withRetry(func() error {
+	return c.withRetryCtx(ctx, func() error {
 		body := c.buildRequestBody(model, req)
 		body.Stream = true
 
@@ -544,7 +555,7 @@ func (c *AnthropicClient) StreamCompletion(req ChatRequest, cb func(chunk []byte
 			return fmt.Errorf("marshal request body: %w", err)
 		}
 
-		httpReq, err := http.NewRequest(http.MethodPost, c.cfg.URL, bytes.NewReader(payload))
+		httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, c.cfg.URL, bytes.NewReader(payload))
 		if err != nil {
 			return fmt.Errorf("create request: %w", err)
 		}
@@ -574,13 +585,11 @@ func (c *AnthropicClient) StreamCompletion(req ChatRequest, cb func(chunk []byte
 		for scanner.Scan() {
 			line := scanner.Text()
 
-			// Capture event type line: "event: message_delta"
 			if strings.HasPrefix(line, "event: ") {
 				eventType = strings.TrimPrefix(line, "event: ")
 				continue
 			}
 
-			// Skip empty lines and non-data lines
 			if !strings.HasPrefix(line, "data: ") {
 				continue
 			}
@@ -590,7 +599,6 @@ func (c *AnthropicClient) StreamCompletion(req ChatRequest, cb func(chunk []byte
 				continue
 			}
 
-			// message_stop signals end of stream
 			if eventType == "message_stop" {
 				break
 			}