From 8e9772baa8899e7386486b9bd9c7c3bf77910a7a Mon Sep 17 00:00:00 2001 From: kite Date: Sat, 6 Jun 2026 13:22:32 +0800 Subject: [PATCH] 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 | 25 +++++++++++++++---------- internal/diff/relocation_test.go | 4 ++++ internal/llm/client.go | 22 +++++++++++++++------- 3 files changed, 34 insertions(+), 17 deletions(-) diff --git a/internal/agent/agent.go b/internal/agent/agent.go index db11a90..f2a28f1 100644 --- a/internal/agent/agent.go +++ b/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 } } diff --git a/internal/diff/relocation_test.go b/internal/diff/relocation_test.go index 8b3a81b..ce5ca5a 100644 --- a/internal/diff/relocation_test.go +++ b/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{ diff --git a/internal/llm/client.go b/internal/llm/client.go index 4948ffa..2217bc4 100644 --- a/internal/llm/client.go +++ b/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 }