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 }