zzet--gortex/internal/mcp/response_buffer.go

package mcp

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/mark3labs/mcp-go/mcp"
)

// defaultResponseBufferCap is how many recent tool responses a session
// keeps available for post-filter re-cutting.
const defaultResponseBufferCap = 8

// minCapturedResponseBytes is the capture floor: responses smaller
// than this are cheap to re-fetch, so they never displace a ring slot.
const minCapturedResponseBytes = 1024

// postFilterTools never have their own responses captured — re-cutting
// a re-cut would only churn the ring.
var postFilterTools = map[string]bool{
	"ctx_stats":    true,
	"ctx_peek":     true,
	"ctx_slice":    true,
	"ctx_grep":     true,
	"grep_results": true,
	"head_results": true,
}

// bufferedResponse is one captured tool response held for re-cutting.
type bufferedResponse struct {
	ID         string
	Tool       string
	Text       string
	CapturedAt time.Time
}

// responseBuffer is a per-session ring of recent large tool responses.
// It backs the post-filter tools (ctx_grep, ctx_slice, …) so an agent
// can re-cut a prior result without re-issuing the original query.
type responseBuffer struct {
	mu      sync.Mutex
	entries []bufferedResponse
	seq     int
}

// capture stores a response and returns its handle ID. The oldest
// entry is evicted, and its backing memory released, once the ring is
// full.
func (b *responseBuffer) capture(tool, text string) string {
	b.mu.Lock()
	defer b.mu.Unlock()
	b.seq++
	id := fmt.Sprintf("resp_%d", b.seq)
	b.entries = append(b.entries, bufferedResponse{
		ID:         id,
		Tool:       tool,
		Text:       text,
		CapturedAt: time.Now(),
	})
	if len(b.entries) > defaultResponseBufferCap {
		trimmed := make([]bufferedResponse, defaultResponseBufferCap)
		copy(trimmed, b.entries[len(b.entries)-defaultResponseBufferCap:])
		b.entries = trimmed
	}
	return id
}

// get resolves a handle. An empty or "latest" id returns the most
// recent capture; the bool is false when the buffer is empty or the id
// is unknown.
func (b *responseBuffer) get(id string) (bufferedResponse, bool) {
	b.mu.Lock()
	defer b.mu.Unlock()
	if len(b.entries) == 0 {
		return bufferedResponse{}, false
	}
	if id == "" || id == "latest" {
		return b.entries[len(b.entries)-1], true
	}
	for i := len(b.entries) - 1; i >= 0; i-- {
		if b.entries[i].ID == id {
			return b.entries[i], true
		}
	}
	return bufferedResponse{}, false
}

// list returns every buffered response, most recent first.
func (b *responseBuffer) list() []bufferedResponse {
	b.mu.Lock()
	defer b.mu.Unlock()
	out := make([]bufferedResponse, 0, len(b.entries))
	for i := len(b.entries) - 1; i >= 0; i-- {
		out = append(out, b.entries[i])
	}
	return out
}

// responseBufferFor returns the calling session's response buffer,
// allocating it on first use.
func (s *Server) responseBufferFor(ctx context.Context) *responseBuffer {
	sess := s.sessionFor(ctx)
	sess.mu.Lock()
	defer sess.mu.Unlock()
	if sess.responses == nil {
		sess.responses = &responseBuffer{}
	}
	return sess.responses
}

// captureResponse stores a large, successful, non-post-filter tool
// response in the session ring so it can be re-cut later. Called from
// wrapToolHandler for every tool call.
func (s *Server) captureResponse(ctx context.Context, tool string, res *mcp.CallToolResult) {
	if res == nil || res.IsError || postFilterTools[tool] {
		return
	}
	text := normalizeForBuffer(toolResultText(res))
	if len(text) < minCapturedResponseBytes {
		return
	}
	s.responseBufferFor(ctx).capture(tool, text)
}

// normalizeForBuffer renders a compact-JSON response as indented text
// so the post-filter tools' line-based slicing and grep land on
// meaningful lines. Non-JSON text (GCX, TOON) is already
// line-structured and passes through unchanged.
func normalizeForBuffer(text string) string {
	t := strings.TrimSpace(text)
	if t == "" || (t[0] != '{' && t[0] != '[') {
		return text
	}
	var buf bytes.Buffer
	if err := json.Indent(&buf, []byte(text), "", "  "); err != nil {
		return text
	}
	return buf.String()
}

// toolResultText extracts the concatenated text content of a tool
// result. The common single-block case is returned without a copy.
func toolResultText(res *mcp.CallToolResult) string {
	if res == nil || len(res.Content) == 0 {
		return ""
	}
	if len(res.Content) == 1 {
		if tc, ok := res.Content[0].(mcp.TextContent); ok {
			return tc.Text
		}
		return ""
	}
	var b strings.Builder
	for _, c := range res.Content {
		if tc, ok := c.(mcp.TextContent); ok {
			b.WriteString(tc.Text)
		}
	}
	return b.String()
}