fix: resolve CI failures and add benchmark CLI features

- Fix Gemini streaming hang by detecting finishReason completion signal
- Fix sandbox TestNewExecutorWithAuto by disabling IPC in CI environment
- Fix lint issues in performance_benchmark_test.go (unused variables, fmt.Sprintf)
- Add adaptive iteration flags (-adaptive, -base-iters, -max-iters, -early-termination)
- Add recursion depth flag (-recursion-depth) to benchmark CLI
- Convert export_oolong.py to uv inline script (PEP 723)

Author: XiaoConstantine

cmd/benchmark/export_oolong.py

@@ -1,18 +1,19 @@
-#!/usr/bin/env python3
+#!/usr/bin/env -S uv run
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "datasets",
+# ]
+# ///
 """Export OOLONG dataset to JSON for Go benchmark."""
 
 import json
 import sys
-from pathlib import Path
+
+from datasets import load_dataset  # noqa: E402 - managed by uv
 
 
 def main():
-    try:
-        from datasets import load_dataset
-    except ImportError:
-        print("Error: datasets library not installed")
-        print("Install with: pip install datasets")
-        sys.exit(1)
 
     num_tasks = int(sys.argv[1]) if len(sys.argv) > 1 else 10
     output_file = sys.argv[2] if len(sys.argv) > 2 else "oolong_tasks.json"

cmd/benchmark/main.go

@@ -198,13 +198,18 @@ func runBaseline(ctx context.Context, task Task, client providers.Client) Benchm
 
 // RLMOptions holds configuration for RLM runs.
 type RLMOptions struct {
-	LogDir        string
-	Verbose       bool
-	Streaming     bool
-	Pool          *repl.REPLPool
-	Compression   bool
-	VerbatimIters int
-	Backend       string
+	LogDir         string
+	Verbose        bool
+	Streaming      bool
+	Pool           *repl.REPLPool
+	Compression    bool
+	VerbatimIters  int
+	Backend        string
+	Adaptive       bool
+	BaseIters      int
+	MaxIters       int
+	EarlyTerminate bool
+	RecursionDepth int
 }
 
 // ModelClient is an interface for clients that expose their model name.
@@ -240,7 +245,7 @@ func runRLM(ctx context.Context, task Task, client providers.Client, opts RLMOpt
 
 	// Build RLM options
 	rlmOpts := []rlm.Option{
-		rlm.WithMaxIterations(30),
+		rlm.WithMaxIterations(opts.MaxIters),
 		rlm.WithVerbose(opts.Verbose),
 		rlm.WithLogger(log),
 	}
@@ -257,6 +262,21 @@ func runRLM(ctx context.Context, task Task, client providers.Client, opts RLMOpt
 		rlmOpts = append(rlmOpts, rlm.WithHistoryCompression(opts.VerbatimIters, 500))
 	}
 
+	if opts.Adaptive {
+		rlmOpts = append(rlmOpts, rlm.WithAdaptiveIterationConfig(rlm.AdaptiveIterationConfig{
+			Enabled:                true,
+			BaseIterations:         opts.BaseIters,
+			MaxIterations:          opts.MaxIters,
+			ContextScaleFactor:     100000, // 100KB per additional iteration
+			EnableEarlyTermination: opts.EarlyTerminate,
+			ConfidenceThreshold:    1,
+		}))
+	}
+
+	if opts.RecursionDepth > 0 {
+		rlmOpts = append(rlmOpts, rlm.WithMaxRecursionDepth(opts.RecursionDepth))
+	}
+
 	r := rlm.New(client, client, rlmOpts...)
 
 	start := time.Now()
@@ -370,6 +390,11 @@ func main() {
 		verbatimIters       = flag.Int("verbatim-iters", 3, "Keep last N iterations verbatim (requires -compression)")
 		enablePrefixCaching = flag.Bool("prefix-caching", true, "Enable Anthropic prefix caching")
 		enableAsync         = flag.Bool("async", false, "Use async sub-LLM queries where possible")
+		enableAdaptive      = flag.Bool("adaptive", false, "Enable adaptive iteration strategy")
+		baseIters           = flag.Int("base-iters", 10, "Base iterations for adaptive mode")
+		maxIters            = flag.Int("max-iters", 30, "Maximum iterations (used in both adaptive and fixed modes)")
+		enableEarlyTerm     = flag.Bool("early-termination", true, "Enable early termination on confidence (requires -adaptive)")
+		recursionDepth      = flag.Int("recursion-depth", 1, "Max recursion depth for multi-depth RLM (0=disabled, 1=paper default)")
 	)
 	flag.Parse()
 
@@ -410,6 +435,7 @@ func main() {
 	fmt.Printf("Provider: %s\n", provider)
 	fmt.Printf("Model: %s\n", *model)
 	fmt.Printf("Tasks: %d\n", len(tasks))
+	fmt.Printf("Max Iterations: %d\n", *maxIters)
 	if *enableStreaming {
 		fmt.Println("Streaming: enabled")
 	}
@@ -419,6 +445,12 @@ func main() {
 	if *enableCompression {
 		fmt.Printf("History Compression: enabled (verbatim=%d)\n", *verbatimIters)
 	}
+	if *enableAdaptive {
+		fmt.Printf("Adaptive Iteration: enabled (base=%d, max=%d, early-term=%v)\n", *baseIters, *maxIters, *enableEarlyTerm)
+	}
+	if *recursionDepth > 0 {
+		fmt.Printf("Multi-Depth Recursion: enabled (max-depth=%d)\n", *recursionDepth)
+	}
 	if *enablePrefixCaching && provider == providers.Anthropic {
 		fmt.Println("Prefix Caching: enabled (Anthropic)")
 	} else if provider == providers.Anthropic {
@@ -454,13 +486,18 @@ func main() {
 
 	// Build RLM options
 	rlmOpts := RLMOptions{
-		LogDir:        *logDir,
-		Verbose:       *verbose,
-		Streaming:     *enableStreaming,
-		Pool:          pool,
-		Compression:   *enableCompression,
-		VerbatimIters: *verbatimIters,
-		Backend:       string(provider),
+		LogDir:         *logDir,
+		Verbose:        *verbose,
+		Streaming:      *enableStreaming,
+		Pool:           pool,
+		Compression:    *enableCompression,
+		VerbatimIters:  *verbatimIters,
+		Backend:        string(provider),
+		Adaptive:       *enableAdaptive,
+		BaseIters:      *baseIters,
+		MaxIters:       *maxIters,
+		EarlyTerminate: *enableEarlyTerm,
+		RecursionDepth: *recursionDepth,
 	}
 
 	var baselineResults []BenchmarkResult
@@ -554,6 +591,16 @@ func main() {
 		output := map[string]any{
 			"model":            *model,
 			"prefix_caching":   *enablePrefixCaching,
+			"streaming":        *enableStreaming,
+			"pooling":          *enablePooling,
+			"pool_size":        *poolSize,
+			"compression":      *enableCompression,
+			"verbatim_iters":   *verbatimIters,
+			"adaptive":         *enableAdaptive,
+			"base_iters":       *baseIters,
+			"max_iters":        *maxIters,
+			"early_termination": *enableEarlyTerm,
+			"recursion_depth":  *recursionDepth,
 			"baseline_results": baselineResults,
 			"rlm_results":      rlmResults,
 			"baseline_summary": baselineSummary,

pkg/benchmark/performance_benchmark_test.go

@@ -4,7 +4,6 @@ package benchmark
 
 import (
 	"context"
-	"fmt"
 	"os"
 	"strings"
 	"sync"
@@ -189,8 +188,8 @@ func TestStreamingVsNonStreaming(t *testing.T) {
 	avgFirstChunk := avgDuration(firstChunkTimes)
 
 	t.Logf("\n=== Streaming vs Non-Streaming ===")
-	t.Logf("Non-streaming avg: %v", avgNonStreaming)
-	t.Logf("Streaming avg:     %v", avgStreaming)
+	t.Logf("Non-streaming avg: %v (tokens: %v)", avgNonStreaming, nonStreamingTokens)
+	t.Logf("Streaming avg:     %v (tokens: %v)", avgStreaming, streamingTokens)
 	t.Logf("First chunk avg:   %v", avgFirstChunk)
 	if avgNonStreaming > 0 {
 		improvement := float64(avgNonStreaming-avgFirstChunk) / float64(avgNonStreaming) * 100
@@ -364,6 +363,7 @@ func TestGeminiImplicitCaching(t *testing.T) {
 	}
 
 	t.Logf("\n=== Gemini Implicit Caching ===")
+	t.Logf("Total tokens across runs: %v", totalTokens)
 	t.Logf("Cache read tokens across runs: %v", cacheReadTokens)
 	t.Logf("Timing across runs: %v", timings)
 
@@ -740,13 +740,13 @@ func TestIdentifyBottlenecks(t *testing.T) {
 	execTimes := make([]time.Duration, 5)
 	for i := 0; i < 5; i++ {
 		start := time.Now()
-		code := fmt.Sprintf(`
+		code := `
 			sum := 0
 			for j := 0; j < 1000; j++ {
 				sum += j
 			}
 			fmt.Println(sum)
-		`)
+		`
 		_, _ = r.Execute(context.Background(), code)
 		execTimes[i] = time.Since(start)
 	}

pkg/providers/gemini.go

@@ -209,13 +209,14 @@ func (c *GeminiClient) doRequest(ctx context.Context, reqBody geminiRequest) (st
 		return "", empty, fmt.Errorf("marshal request: %w", err)
 	}
 
-	url := fmt.Sprintf("%s/v1beta/models/%s:generateContent?key=%s", c.baseURL, c.model, c.apiKey)
+	url := fmt.Sprintf("%s/v1beta/models/%s:generateContent", c.baseURL, c.model)
 	req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(jsonBody))
 	if err != nil {
 		return "", empty, fmt.Errorf("create request: %w", err)
 	}
 
 	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("x-goog-api-key", c.apiKey)
 
 	resp, err := c.httpClient.Do(req)
 	if err != nil {
@@ -335,13 +336,14 @@ func (c *GeminiClient) doStreamRequest(ctx context.Context, reqBody geminiReques
 		return core.LLMResponse{}, fmt.Errorf("marshal request: %w", err)
 	}
 
-	url := fmt.Sprintf("%s/v1beta/models/%s:streamGenerateContent?alt=sse&key=%s", c.baseURL, c.model, c.apiKey)
+	url := fmt.Sprintf("%s/v1beta/models/%s:streamGenerateContent?alt=sse", c.baseURL, c.model)
 	req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(jsonBody))
 	if err != nil {
 		return core.LLMResponse{}, fmt.Errorf("create request: %w", err)
 	}
 
 	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("x-goog-api-key", c.apiKey)
 
 	resp, err := c.httpClient.Do(req)
 	if err != nil {
@@ -356,6 +358,7 @@ func (c *GeminiClient) doStreamRequest(ctx context.Context, reqBody geminiReques
 
 	var fullContent strings.Builder
 	var promptTokens, completionTokens, cachedContentTokens int
+	var streamCompleted bool
 
 	scanner := bufio.NewScanner(resp.Body)
 	buf := make([]byte, 0, 64*1024)
@@ -410,13 +413,30 @@ func (c *GeminiClient) doStreamRequest(ctx context.Context, reqBody geminiReques
 		if chunk.UsageMetadata.CachedContentTokenCount > 0 {
 			cachedContentTokens = chunk.UsageMetadata.CachedContentTokenCount
 		}
+
+		// Check for stream completion signal from Gemini
+		for _, candidate := range chunk.Candidates {
+			if candidate.FinishReason == "STOP" || candidate.FinishReason == "MAX_TOKENS" ||
+				candidate.FinishReason == "SAFETY" || candidate.FinishReason == "RECITATION" {
+				// Stream complete - signal handler and exit loop
+				streamCompleted = true
+				if handler != nil {
+					if err := handler("", true); err != nil {
+						return core.LLMResponse{}, fmt.Errorf("handler error: %w", err)
+					}
+				}
+				goto streamDone
+			}
+		}
 	}
 
+streamDone:
 	if err := scanner.Err(); err != nil {
 		return core.LLMResponse{}, fmt.Errorf("scanner error: %w", err)
 	}
 
-	if handler != nil {
+	// Only call handler if we didn't already signal completion via finishReason
+	if handler != nil && !streamCompleted {
 		if err := handler("", true); err != nil {
 			return core.LLMResponse{}, fmt.Errorf("handler error: %w", err)
 		}

pkg/providers/gemini_test.go

@@ -21,8 +21,8 @@ func TestGeminiClient_Complete_Success(t *testing.T) {
 		if !strings.Contains(r.URL.Path, "/v1beta/models/") {
 			t.Errorf("Expected /v1beta/models/ in path, got %s", r.URL.Path)
 		}
-		if !strings.Contains(r.URL.RawQuery, "key=test-key") {
-			t.Errorf("Expected key=test-key in query, got %s", r.URL.RawQuery)
+		if r.Header.Get("x-goog-api-key") != "test-key" {
+			t.Errorf("Expected x-goog-api-key header to be test-key, got %s", r.Header.Get("x-goog-api-key"))
 		}
 
 		var req geminiRequest

pkg/sandbox/container_test.go

@@ -318,6 +318,7 @@ func TestNewExecutorWithAuto(t *testing.T) {
 	client := NewMockLLMClient()
 	cfg := DefaultConfig()
 	cfg.Backend = BackendAuto
+	cfg.EnableIPC = false // Disable IPC for basic test to avoid network issues in CI
 
 	exec, err := New(client, cfg)
 	if err != nil {
@@ -349,7 +350,7 @@ func TestNewExecutorWithAuto(t *testing.T) {
 	}
 
 	if !strings.Contains(result.Stdout, "Hello from auto backend!") {
-		t.Errorf("Expected greeting in output, got: %s", result.Stdout)
+		t.Errorf("Expected greeting in output, got stdout: %q, stderr: %q", result.Stdout, result.Stderr)
 	}
 }