1  package agent
  2  
  3  import (
  4  	"crypto/sha256"
  5  	"encoding/hex"
  6  	"fmt"
  7  	"regexp"
  8  	"strings"
  9  )
 10  
 11  type ResultProfile string
 12  
 13  const (
 14  	ResultProfileDefault ResultProfile = "default"
 15  	ResultProfileTree    ResultProfile = "tree"
 16  )
 17  
 18  type ShapedResult struct {
 19  	Text      string
 20  	Signature string
 21  }
 22  
 23  var treeRefPattern = regexp.MustCompile(`\bref=e\d+\b`)
 24  var treeWhitespacePattern = regexp.MustCompile(`[ \t]+`)
 25  
 26  func resultProfileForTool(toolName string) ResultProfile {
 27  	switch toolName {
 28  	case "browser_snapshot":
 29  		return ResultProfileTree
 30  	default:
 31  		return ResultProfileDefault
 32  	}
 33  }
 34  
 35  func shapeContextResult(toolName, content string, previous *ShapedResult) ShapedResult {
 36  	switch resultProfileForTool(toolName) {
 37  	case ResultProfileTree:
 38  		return shapeTreeResult(content, previous)
 39  	default:
 40  		return ShapedResult{Text: content}
 41  	}
 42  }
 43  
 44  func shapeContextKey(toolName string, traits CallStateTraits, tracker *stateVersionTracker) string {
 45  	if tracker == nil || len(traits.Reads) == 0 {
 46  		return toolName
 47  	}
 48  	if fingerprint := tracker.fingerprint(traits.Reads); fingerprint != "" {
 49  		return toolName + "\x00" + fingerprint
 50  	}
 51  	return toolName
 52  }
 53  
 54  func shapeTreeResult(content string, previous *ShapedResult) ShapedResult {
 55  	normalizedLines := normalizeTreeLines(content)
 56  	if len(normalizedLines) == 0 {
 57  		return ShapedResult{Text: content}
 58  	}
 59  
 60  	signature := shortStableHash(strings.Join(normalizedLines, "\n"))
 61  	if len([]rune(content)) < 2000 && len(normalizedLines) < 30 {
 62  		return ShapedResult{Text: content, Signature: signature}
 63  	}
 64  
 65  	refCount := 0
 66  	for _, line := range normalizedLines {
 67  		if strings.Contains(line, "ref=*") {
 68  			refCount++
 69  		}
 70  	}
 71  
 72  	if previous != nil && previous.Signature == signature {
 73  		return ShapedResult{
 74  			Text:      fmt.Sprintf("[tree snapshot unchanged since last read; signature %s; %d lines; %d refs]", signature, len(normalizedLines), refCount),
 75  			Signature: signature,
 76  		}
 77  	}
 78  
 79  	header := fmt.Sprintf("[tree snapshot summary; signature %s; %d lines; %d refs]", signature, len(normalizedLines), refCount)
 80  	if previous != nil && previous.Signature != "" {
 81  		header = fmt.Sprintf("[tree snapshot changed since last read; signature %s; previous %s; %d lines; %d refs]", signature, previous.Signature, len(normalizedLines), refCount)
 82  	}
 83  
 84  	excerpt := buildTreeExcerpt(normalizedLines, 18, 1600)
 85  	if excerpt == "" {
 86  		return ShapedResult{Text: header, Signature: signature}
 87  	}
 88  	return ShapedResult{
 89  		Text:      header + "\n" + excerpt,
 90  		Signature: signature,
 91  	}
 92  }
 93  
 94  func normalizeTreeLines(content string) []string {
 95  	lines := strings.Split(content, "\n")
 96  	normalized := make([]string, 0, len(lines))
 97  	for _, line := range lines {
 98  		line = treeRefPattern.ReplaceAllString(line, "ref=*")
 99  		line = treeWhitespacePattern.ReplaceAllString(line, " ")
100  		line = strings.TrimSpace(line)
101  		if line == "" {
102  			continue
103  		}
104  		normalized = append(normalized, line)
105  	}
106  	return normalized
107  }
108  
109  func buildTreeExcerpt(lines []string, maxLines, maxChars int) string {
110  	if len(lines) == 0 || maxLines <= 0 || maxChars <= 0 {
111  		return ""
112  	}
113  
114  	var excerpt []string
115  	chars := 0
116  	for _, line := range lines {
117  		if len(excerpt) >= maxLines {
118  			break
119  		}
120  		if len([]rune(line)) > 140 {
121  			runes := []rune(line)
122  			line = string(runes[:140]) + "..."
123  		}
124  		nextChars := chars + len([]rune(line))
125  		if len(excerpt) > 0 {
126  			nextChars++
127  		}
128  		if nextChars > maxChars {
129  			break
130  		}
131  		excerpt = append(excerpt, line)
132  		chars = nextChars
133  	}
134  	return strings.Join(excerpt, "\n")
135  }
136  
137  func shortStableHash(content string) string {
138  	sum := sha256.Sum256([]byte(content))
139  	return hex.EncodeToString(sum[:])[:12]
140  }