recursive-llm-ts 4.7.0 → 4.9.0

package/go/rlm/context_overflow_test.go

@@ -153,7 +153,7 @@ func TestGetResponseTokenBudget(t *testing.T) {
 	config := DefaultContextOverflowConfig()
 	reducer := newContextReducer(rlm, config, obs)
 
-	budget := reducer.getResponseTokenBudget(32768)
+	budget := reducer.getResponseTokenBudget()
 	if budget != 10000 {
 		t.Errorf("expected response token budget 10000, got %d", budget)
 	}
@@ -169,7 +169,7 @@ func TestGetResponseTokenBudget_MaxCompletionTokens(t *testing.T) {
 	config := DefaultContextOverflowConfig()
 	reducer := newContextReducer(rlm, config, obs)
 
-	budget := reducer.getResponseTokenBudget(32768)
+	budget := reducer.getResponseTokenBudget()
 	if budget != 5000 {
 		t.Errorf("expected response token budget 5000, got %d", budget)
 	}
@@ -185,7 +185,7 @@ func TestGetResponseTokenBudget_NoMaxTokens(t *testing.T) {
 	config := DefaultContextOverflowConfig()
 	reducer := newContextReducer(rlm, config, obs)
 
-	budget := reducer.getResponseTokenBudget(32768)
+	budget := reducer.getResponseTokenBudget()
 	if budget != 0 {
 		t.Errorf("expected response token budget 0, got %d", budget)
 	}
@@ -899,3 +899,373 @@ func TestReduceForCompletion_DispatchesTextRank(t *testing.T) {
 		t.Errorf("expected reduced context for textrank strategy")
 	}
 }
+
+// ─── Model Token Limits Tests ────────────────────────────────────────────────
+
+func TestLookupModelTokenLimit_ExactMatch(t *testing.T) {
+	tests := []struct {
+		model    string
+		expected int
+	}{
+		{"gpt-4o", 128000},
+		{"gpt-4o-mini", 128000},
+		{"gpt-4", 8192},
+		{"gpt-4-32k", 32768},
+		{"gpt-3.5-turbo", 16385},
+		{"claude-3-opus", 200000},
+		{"claude-sonnet-4", 200000},
+		{"mistral-7b", 32768},
+	}
+
+	for _, tt := range tests {
+		limit := LookupModelTokenLimit(tt.model)
+		if limit != tt.expected {
+			t.Errorf("LookupModelTokenLimit(%q) = %d, expected %d", tt.model, limit, tt.expected)
+		}
+	}
+}
+
+func TestLookupModelTokenLimit_PrefixMatch(t *testing.T) {
+	// Versioned model names should match by prefix
+	tests := []struct {
+		model    string
+		expected int
+	}{
+		{"gpt-4o-mini-2024-07-18", 128000},
+		{"gpt-4o-2024-05-13", 128000},
+		{"claude-3-opus-20240229", 200000},
+		{"mistral-7b-instruct-v0.2", 32768},
+	}
+
+	for _, tt := range tests {
+		limit := LookupModelTokenLimit(tt.model)
+		if limit != tt.expected {
+			t.Errorf("LookupModelTokenLimit(%q) = %d, expected %d", tt.model, limit, tt.expected)
+		}
+	}
+}
+
+func TestLookupModelTokenLimit_Unknown(t *testing.T) {
+	limit := LookupModelTokenLimit("completely-unknown-model-xyz")
+	if limit != 0 {
+		t.Errorf("expected 0 for unknown model, got %d", limit)
+	}
+}
+
+func TestLookupModelTokenLimit_CaseInsensitive(t *testing.T) {
+	limit := LookupModelTokenLimit("GPT-4O-MINI")
+	if limit != 128000 {
+		t.Errorf("expected 128000 for case-insensitive match, got %d", limit)
+	}
+}
+
+func TestGetModelTokenLimit_ConfigOverride(t *testing.T) {
+	engine := New("gpt-4o-mini", Config{
+		APIKey: "test",
+		ContextOverflow: &ContextOverflowConfig{
+			Enabled:        true,
+			MaxModelTokens: 16384,
+		},
+	})
+
+	limit := engine.getModelTokenLimit()
+	if limit != 16384 {
+		t.Errorf("expected config override 16384, got %d", limit)
+	}
+}
+
+func TestGetModelTokenLimit_ModelLookup(t *testing.T) {
+	engine := New("gpt-4o-mini", Config{
+		APIKey: "test",
+	})
+
+	limit := engine.getModelTokenLimit()
+	if limit != 128000 {
+		t.Errorf("expected model lookup 128000, got %d", limit)
+	}
+}
+
+func TestGetModelTokenLimit_UnknownModel(t *testing.T) {
+	engine := New("custom-local-model", Config{
+		APIKey: "test",
+	})
+
+	limit := engine.getModelTokenLimit()
+	if limit != 0 {
+		t.Errorf("expected 0 for unknown model, got %d", limit)
+	}
+}
+
+// ─── Pre-emptive Overflow Tests ──────────────────────────────────────────────
+
+func TestPreemptiveReduceContext_SmallContext(t *testing.T) {
+	engine := New("gpt-4o-mini", Config{
+		APIKey: "test",
+	})
+
+	// Small context should pass through unchanged
+	context := "This is a small context that easily fits."
+	reduced, wasReduced, err := engine.PreemptiveReduceContext("What is this?", context, 500)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if wasReduced {
+		t.Error("expected no reduction for small context")
+	}
+	if reduced != context {
+		t.Error("expected context to be unchanged")
+	}
+}
+
+func TestPreemptiveReduceContext_LargeContext(t *testing.T) {
+	engine := New("gpt-4o-mini", Config{
+		APIKey: "test",
+		ContextOverflow: &ContextOverflowConfig{
+			Enabled:        true,
+			MaxModelTokens: 1000, // Very small limit to force overflow
+			Strategy:       "truncate",
+			SafetyMargin:   0.15,
+		},
+	})
+
+	// Create large context that exceeds the 1000 token limit
+	context := strings.Repeat("The revenue for Q4 was $4.2 billion, representing 23% year-over-year growth. ", 100)
+
+	reduced, wasReduced, err := engine.PreemptiveReduceContext("Summarize revenue", context, 300)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if !wasReduced {
+		t.Error("expected context to be reduced")
+	}
+	if len(reduced) >= len(context) {
+		t.Errorf("expected reduced context to be shorter: %d >= %d", len(reduced), len(context))
+	}
+}
+
+func TestPreemptiveReduceContext_DisabledOverflow(t *testing.T) {
+	engine := New("gpt-4o-mini", Config{
+		APIKey: "test",
+		ContextOverflow: &ContextOverflowConfig{
+			Enabled: false,
+		},
+	})
+
+	context := strings.Repeat("Large content. ", 10000)
+	reduced, wasReduced, err := engine.PreemptiveReduceContext("query", context, 500)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if wasReduced {
+		t.Error("expected no reduction when overflow is disabled")
+	}
+	if reduced != context {
+		t.Error("expected context unchanged when overflow is disabled")
+	}
+}
+
+func TestPreemptiveReduceContext_UnknownModel(t *testing.T) {
+	engine := New("custom-local-model", Config{
+		APIKey: "test",
+	})
+
+	// Unknown model with no config override → no pre-emptive check
+	context := strings.Repeat("Large content. ", 10000)
+	reduced, wasReduced, err := engine.PreemptiveReduceContext("query", context, 500)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if wasReduced {
+		t.Error("expected no reduction for unknown model with no config limit")
+	}
+	if reduced != context {
+		t.Error("expected context unchanged")
+	}
+}
+
+func TestPreemptiveReduceContext_AccountsForResponseBudget(t *testing.T) {
+	// With a high max_tokens, even moderate context should trigger reduction
+	engine := New("gpt-4o-mini", Config{
+		APIKey: "test",
+		ContextOverflow: &ContextOverflowConfig{
+			Enabled:        true,
+			MaxModelTokens: 2000,
+			Strategy:       "truncate",
+			SafetyMargin:   0.15,
+		},
+		ExtraParams: map[string]interface{}{
+			"max_tokens": float64(1000), // Large response budget
+		},
+	})
+
+	// Context of ~500 tokens + max_tokens 1000 + overhead = exceeds 2000
+	context := strings.Repeat("Revenue data: the company earned $4.2B in Q4 fiscal year. ", 30)
+
+	reduced, wasReduced, err := engine.PreemptiveReduceContext("Summarize", context, 300)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if !wasReduced {
+		t.Error("expected reduction when response budget + context exceeds limit")
+	}
+	if len(reduced) >= len(context) {
+		t.Errorf("expected reduced context: %d >= %d", len(reduced), len(context))
+	}
+}
+
+func TestPreemptiveReduceContext_TFIDFStrategy(t *testing.T) {
+	engine := New("gpt-4o-mini", Config{
+		APIKey: "test",
+		ContextOverflow: &ContextOverflowConfig{
+			Enabled:        true,
+			MaxModelTokens: 500,
+			Strategy:       "tfidf",
+			SafetyMargin:   0.15,
+		},
+	})
+
+	context := strings.Repeat("Machine learning models process large datasets effectively. ", 100)
+
+	reduced, wasReduced, err := engine.PreemptiveReduceContext("Tell me about ML", context, 200)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if !wasReduced {
+		t.Error("expected reduction with tfidf strategy")
+	}
+	if len(reduced) >= len(context) {
+		t.Errorf("expected shorter context: %d >= %d", len(reduced), len(context))
+	}
+}
+
+func TestPreemptiveReduceContext_TextRankStrategy(t *testing.T) {
+	engine := New("gpt-4o-mini", Config{
+		APIKey: "test",
+		ContextOverflow: &ContextOverflowConfig{
+			Enabled:        true,
+			MaxModelTokens: 500,
+			Strategy:       "textrank",
+			SafetyMargin:   0.15,
+		},
+	})
+
+	context := strings.Repeat("Neural networks are powerful computation models. ", 100)
+
+	reduced, wasReduced, err := engine.PreemptiveReduceContext("Explain neural nets", context, 200)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if !wasReduced {
+		t.Error("expected reduction with textrank strategy")
+	}
+	if len(reduced) >= len(context) {
+		t.Errorf("expected shorter context: %d >= %d", len(reduced), len(context))
+	}
+}
+
+func TestGetResponseTokenBudget_RLMMethod(t *testing.T) {
+	engine := &RLM{
+		extraParams: map[string]interface{}{
+			"max_tokens": float64(5000),
+		},
+	}
+	budget := engine.getResponseTokenBudget()
+	if budget != 5000 {
+		t.Errorf("expected 5000, got %d", budget)
+	}
+}
+
+func TestGetResponseTokenBudget_MaxCompletionTokensPreferred(t *testing.T) {
+	engine := &RLM{
+		extraParams: map[string]interface{}{
+			"max_tokens":            float64(5000),
+			"max_completion_tokens": float64(8000),
+		},
+	}
+	budget := engine.getResponseTokenBudget()
+	if budget != 8000 {
+		t.Errorf("expected max_completion_tokens=8000 preferred, got %d", budget)
+	}
+}
+
+func TestGetResponseTokenBudget_NoParams(t *testing.T) {
+	engine := &RLM{
+		extraParams: map[string]interface{}{
+			"temperature": 0.7,
+		},
+	}
+	budget := engine.getResponseTokenBudget()
+	if budget != 0 {
+		t.Errorf("expected 0 when no max_tokens set, got %d", budget)
+	}
+}
+
+// ─── Message Pruning Tests ───────────────────────────────────────────────────
+
+func TestPruneMessages_SmallHistory(t *testing.T) {
+	messages := []Message{
+		{Role: "system", Content: "You are helpful."},
+		{Role: "user", Content: "Hello"},
+		{Role: "assistant", Content: "Hi there!"},
+	}
+
+	result := pruneMessages(messages, 100)
+	if len(result) != 3 {
+		t.Errorf("expected 3 messages (no pruning needed), got %d", len(result))
+	}
+}
+
+func TestPruneMessages_PreservesSystemAndLast(t *testing.T) {
+	messages := []Message{
+		{Role: "system", Content: "System prompt"},
+		{Role: "user", Content: "First question"},
+		{Role: "assistant", Content: "First answer"},
+		{Role: "user", Content: "Second question"},
+		{Role: "assistant", Content: "Second answer"},
+		{Role: "user", Content: strings.Repeat("Third question with lots of context. ", 100)},
+		{Role: "assistant", Content: "Third answer"},
+	}
+
+	result := pruneMessages(messages, 50) // Very tight budget
+
+	// Should always keep system prompt (first) and last 2 messages
+	if len(result) < 3 {
+		t.Errorf("expected at least 3 messages, got %d", len(result))
+	}
+	if result[0].Role != "system" {
+		t.Error("first message should be system prompt")
+	}
+	if result[len(result)-1].Content != "Third answer" {
+		t.Error("last message should be the most recent")
+	}
+	if result[len(result)-2].Role != "user" {
+		t.Error("second-to-last should be the most recent user message")
+	}
+}
+
+func TestPruneMessages_KeepsRecentMiddleMessages(t *testing.T) {
+	messages := []Message{
+		{Role: "system", Content: "Short."},
+		{Role: "user", Content: "Q1"},
+		{Role: "assistant", Content: "A1"},
+		{Role: "user", Content: "Q2"},
+		{Role: "assistant", Content: "A2"},
+		{Role: "user", Content: "Q3"},
+		{Role: "assistant", Content: "A3"},
+	}
+
+	// Budget large enough for all
+	result := pruneMessages(messages, 10000)
+	if len(result) != 7 {
+		t.Errorf("expected all 7 messages with large budget, got %d", len(result))
+	}
+}
+
+// ─── Structured Completion Pre-emptive Integration Tests ─────────────────────
+
+func TestStructuredPromptOverhead_Constant(t *testing.T) {
+	// Verify the constant is reasonable (300-500 tokens for structured prompt instructions)
+	if structuredPromptOverhead < 200 || structuredPromptOverhead > 600 {
+		t.Errorf("structuredPromptOverhead=%d seems out of range (expected 200-600)", structuredPromptOverhead)
+	}
+}

package/go/rlm/doc.go

@@ -8,13 +8,13 @@
 //
 // To use this package in your Go project:
 //
-//	go get github.com/jbeck018/recursive-llm-ts/go
+//	go get github.com/howlerops/recursive-llm-ts/go
 //
 // # Basic Usage
 //
 // Create an RLM engine and execute a completion:
 //
-//	import "github.com/jbeck018/recursive-llm-ts/go/rlm"
+//	import "github.com/howlerops/recursive-llm-ts/go/rlm"
 //
 //	config := rlm.Config{
 //	    MaxDepth:      5,

package/go/rlm/meta_agent.go

@@ -74,7 +74,15 @@ func (ma *MetaAgent) OptimizeQuery(query string, context string) (string, error)
 		return query, nil
 	}
 
-	optimized := strings.TrimSpace(result)
+	// Track meta-agent token usage in the parent RLM's stats
+	if result.Usage != nil {
+		ma.rlm.stats.PromptTokens += result.Usage.PromptTokens
+		ma.rlm.stats.CompletionTokens += result.Usage.CompletionTokens
+		ma.rlm.stats.TotalTokens += result.Usage.TotalTokens
+	}
+	ma.rlm.stats.LlmCalls++
+
+	optimized := strings.TrimSpace(result.Content)
 	ma.obs.Debug("meta_agent", "Optimized query: %s", truncateStr(optimized, 200))
 	ma.obs.Event("meta_agent.query_optimized", map[string]string{
 		"original_length":  fmt.Sprintf("%d", len(query)),
@@ -136,7 +144,15 @@ func (ma *MetaAgent) OptimizeForStructured(query string, context string, schema
 		return query, nil
 	}
 
-	optimized := strings.TrimSpace(result)
+	// Track meta-agent token usage in the parent RLM's stats
+	if result.Usage != nil {
+		ma.rlm.stats.PromptTokens += result.Usage.PromptTokens
+		ma.rlm.stats.CompletionTokens += result.Usage.CompletionTokens
+		ma.rlm.stats.TotalTokens += result.Usage.TotalTokens
+	}
+	ma.rlm.stats.LlmCalls++
+
+	optimized := strings.TrimSpace(result.Content)
 	ma.obs.Debug("meta_agent", "Optimized structured query: %s", truncateStr(optimized, 200))
 	ma.obs.Event("meta_agent.structured_query_optimized", map[string]string{
 		"original_length":  fmt.Sprintf("%d", len(query)),

package/go/rlm/observability.go

@@ -445,6 +445,12 @@ func FormatStatsWithObservability(stats RLMStats, obs *Observer) map[string]inte
 		result["parsing_retries"] = stats.ParsingRetries
 	}
 
+	if stats.TotalTokens > 0 {
+		result["total_tokens"] = stats.TotalTokens
+		result["prompt_tokens"] = stats.PromptTokens
+		result["completion_tokens"] = stats.CompletionTokens
+	}
+
 	if obs != nil && obs.config.Debug {
 		events := obs.GetEvents()
 		if len(events) > 0 {

package/go/rlm/openai.go

@@ -34,6 +34,20 @@ type chatResponse struct {
 	Error *struct {
 		Message string `json:"message"`
 	} `json:"error"`
+	Usage *TokenUsage `json:"usage,omitempty"`
+}
+
+// TokenUsage represents token consumption from an LLM API response.
+type TokenUsage struct {
+	PromptTokens     int `json:"prompt_tokens"`
+	CompletionTokens int `json:"completion_tokens"`
+	TotalTokens      int `json:"total_tokens"`
+}
+
+// ChatCompletionResult holds the content and token usage from an LLM call.
+type ChatCompletionResult struct {
+	Content string
+	Usage   *TokenUsage
 }
 
 var (
@@ -48,7 +62,7 @@ var (
 	}
 )
 
-func CallChatCompletion(request ChatRequest) (string, error) {
+func CallChatCompletion(request ChatRequest) (ChatCompletionResult, error) {
 	endpoint := buildEndpoint(request.APIBase)
 	payload := map[string]interface{}{
 		"model":    request.Model,
@@ -61,7 +75,7 @@ func CallChatCompletion(request ChatRequest) (string, error) {
 
 	body, err := json.Marshal(payload)
 	if err != nil {
-		return "", err
+		return ChatCompletionResult{}, err
 	}
 
 	// Use shared client with connection pooling
@@ -76,7 +90,7 @@ func CallChatCompletion(request ChatRequest) (string, error) {
 
 	req, err := http.NewRequest(http.MethodPost, endpoint, bytes.NewReader(body))
 	if err != nil {
-		return "", err
+		return ChatCompletionResult{}, err
 	}
 	req.Header.Set("Content-Type", "application/json")
 	if request.APIKey != "" {
@@ -85,7 +99,7 @@ func CallChatCompletion(request ChatRequest) (string, error) {
 
 	resp, err := client.Do(req)
 	if err != nil {
-		return "", err
+		return ChatCompletionResult{}, err
 	}
 	defer func() {
 		_ = resp.Body.Close()
@@ -93,27 +107,30 @@ func CallChatCompletion(request ChatRequest) (string, error) {
 
 	responseBody, err := io.ReadAll(resp.Body)
 	if err != nil {
-		return "", err
+		return ChatCompletionResult{}, err
 	}
 
 	if resp.StatusCode >= http.StatusBadRequest {
-		return "", NewAPIError(resp.StatusCode, strings.TrimSpace(string(responseBody)))
+		return ChatCompletionResult{}, NewAPIError(resp.StatusCode, strings.TrimSpace(string(responseBody)))
 	}
 
 	var parsed chatResponse
 	if err := json.Unmarshal(responseBody, &parsed); err != nil {
-		return "", err
+		return ChatCompletionResult{}, err
 	}
 
 	if parsed.Error != nil && parsed.Error.Message != "" {
-		return "", errors.New(parsed.Error.Message)
+		return ChatCompletionResult{}, errors.New(parsed.Error.Message)
 	}
 
 	if len(parsed.Choices) == 0 {
-		return "", errors.New("no choices returned by LLM")
+		return ChatCompletionResult{}, errors.New("no choices returned by LLM")
 	}
 
-	return parsed.Choices[0].Message.Content, nil
+	return ChatCompletionResult{
+		Content: parsed.Choices[0].Message.Content,
+		Usage:   parsed.Usage,
+	}, nil
 }
 
 func buildEndpoint(apiBase string) string {

package/go/rlm/rlm.go

@@ -109,8 +109,41 @@ func (r *RLM) Completion(query string, context string) (string, RLMStats, error)
 		r.stats.Iterations = iteration + 1
 		r.observer.Debug("rlm", "Iteration %d/%d at depth %d", iteration+1, r.maxIterations, r.currentDepth)
 
+		// Pre-emptive message overflow check: prune older messages if history is growing too large.
+		// Regular completion stores context in the REPL env (not messages), but the iterative
+		// loop appends assistant+user messages each iteration which can accumulate.
+		if modelLimit := r.getModelTokenLimit(); modelLimit > 0 && len(messages) > 4 {
+			msgTokens := EstimateMessagesTokens(messages)
+			responseTokens := r.getResponseTokenBudget()
+			safetyMargin := 0.15
+			if r.contextOverflow != nil && r.contextOverflow.SafetyMargin > 0 {
+				safetyMargin = r.contextOverflow.SafetyMargin
+			}
+			available := modelLimit - responseTokens - int(float64(modelLimit)*safetyMargin)
+			if msgTokens > available {
+				r.observer.Debug("rlm", "Message history overflow: %d tokens > %d available, pruning middle messages", msgTokens, available)
+				messages = pruneMessages(messages, available)
+			}
+		}
+
 		response, err := r.callLLM(messages)
 		if err != nil {
+			// Check for context overflow and attempt recovery
+			if r.contextOverflow != nil && r.contextOverflow.Enabled {
+				if _, isOverflow := IsContextOverflow(err); isOverflow && len(messages) > 4 {
+					r.observer.Debug("rlm", "Context overflow on iteration %d, pruning messages and retrying", iteration+1)
+					modelLimit := r.getModelTokenLimit()
+					if modelLimit == 0 {
+						modelLimit = 32768 // Reasonable default
+					}
+					responseTokens := r.getResponseTokenBudget()
+					available := modelLimit - responseTokens - int(float64(modelLimit)*0.15)
+					messages = pruneMessages(messages, available)
+					// Retry this iteration
+					iteration--
+					continue
+				}
+			}
 			r.observer.Error("rlm", "LLM call failed on iteration %d: %v", iteration+1, err)
 			return "", r.stats, err
 		}
@@ -165,14 +198,22 @@ func (r *RLM) callLLM(messages []Message) (string, error) {
 	result, err := CallChatCompletion(request)
 	duration := time.Since(start)
 
-	r.observer.LLMCall(defaultModel, len(messages), 0, duration, err)
+	tokensUsed := 0
+	if result.Usage != nil {
+		r.stats.PromptTokens += result.Usage.PromptTokens
+		r.stats.CompletionTokens += result.Usage.CompletionTokens
+		r.stats.TotalTokens += result.Usage.TotalTokens
+		tokensUsed = result.Usage.TotalTokens
+	}
+
+	r.observer.LLMCall(defaultModel, len(messages), tokensUsed, duration, err)
 
 	if err != nil {
 		return "", err
 	}
 
-	r.observer.Debug("llm", "Response received (%d chars) in %s", len(result), duration)
-	return result, nil
+	r.observer.Debug("llm", "Response received (%d chars, %d tokens) in %s", len(result.Content), tokensUsed, duration)
+	return result.Content, nil
 }
 
 func (r *RLM) buildREPLEnv(query string, context string) map[string]interface{} {
@@ -214,6 +255,48 @@ func (r *RLM) buildREPLEnv(query string, context string) map[string]interface{}
 	return env
 }
 
+// pruneMessages removes older middle messages to fit within a token budget.
+// Preserves the first message (system prompt) and the last 2 messages (most recent exchange).
+func pruneMessages(messages []Message, targetTokens int) []Message {
+	if len(messages) <= 3 {
+		return messages
+	}
+
+	// Always keep: system prompt (first), last 2 messages (most recent exchange)
+	system := messages[0]
+	lastN := messages[len(messages)-2:]
+
+	// Start with the preserved messages
+	result := []Message{system}
+	currentTokens := EstimateMessagesTokens(append(result, lastN...))
+
+	if currentTokens >= targetTokens {
+		// Even the minimum set exceeds the budget; return it anyway
+		return append(result, lastN...)
+	}
+
+	// Add middle messages from most recent to oldest until budget is exceeded
+	middle := messages[1 : len(messages)-2]
+	for i := len(middle) - 1; i >= 0; i-- {
+		msgTokens := 4 + EstimateTokens(middle[i].Content)
+		if currentTokens+msgTokens > targetTokens {
+			break
+		}
+		result = append(result, middle[i])
+		currentTokens += msgTokens
+	}
+
+	// Reverse the added middle messages (they were added newest-first)
+	if len(result) > 1 {
+		added := result[1:]
+		for i, j := 0, len(added)-1; i < j; i, j = i+1, j-1 {
+			added[i], added[j] = added[j], added[i]
+		}
+	}
+
+	return append(result, lastN...)
+}
+
 // GetObserver returns the observer for external access to events/traces.
 func (r *RLM) GetObserver() *Observer {
 	return r.observer