/ internal / agent / statecache.go
statecache.go
  1  package agent
  2  
  3  import (
  4  	"encoding/json"
  5  	"path/filepath"
  6  	"sort"
  7  	"strconv"
  8  	"strings"
  9  )
 10  
 11  type StateDomain string
 12  
 13  const (
 14  	StateDomainBrowser    StateDomain = "browser"
 15  	StateDomainFilesystem StateDomain = "filesystem"
 16  	StateDomainProcess    StateDomain = "process"
 17  )
 18  
 19  type StateRef struct {
 20  	Domain StateDomain
 21  	Scope  string
 22  }
 23  
 24  type CallStateTraits struct {
 25  	Reads        []StateRef
 26  	Writes       []StateRef
 27  	UnknownWrite bool
 28  	Cacheable    bool
 29  }
 30  
 31  type stateVersionTracker struct {
 32  	versions map[string]int
 33  }
 34  
 35  func newStateVersionTracker() *stateVersionTracker {
 36  	return &stateVersionTracker{versions: make(map[string]int)}
 37  }
 38  
 39  func (t *stateVersionTracker) fingerprint(refs []StateRef) string {
 40  	if len(refs) == 0 {
 41  		return ""
 42  	}
 43  	seen := make(map[string]bool, len(refs))
 44  	parts := make([]string, 0, len(refs))
 45  	for _, ref := range refs {
 46  		key := stateRefKey(ref)
 47  		if key == "" || seen[key] {
 48  			continue
 49  		}
 50  		seen[key] = true
 51  		parts = append(parts, key+"="+strconv.Itoa(t.versions[key]))
 52  	}
 53  	sort.Strings(parts)
 54  	return strings.Join(parts, "|")
 55  }
 56  
 57  func (t *stateVersionTracker) bump(refs []StateRef) {
 58  	for _, ref := range refs {
 59  		key := stateRefKey(ref)
 60  		if key == "" {
 61  			continue
 62  		}
 63  		t.versions[key]++
 64  	}
 65  }
 66  
 67  func stateRefKey(ref StateRef) string {
 68  	scope := strings.TrimSpace(ref.Scope)
 69  	if scope == "" {
 70  		scope = "*"
 71  	}
 72  	return string(ref.Domain) + "\x00" + scope
 73  }
 74  
 75  func browserStateRef() StateRef {
 76  	return StateRef{Domain: StateDomainBrowser, Scope: "active"}
 77  }
 78  
 79  func filesystemStateRef(path string) StateRef {
 80  	path = strings.TrimSpace(path)
 81  	if path == "" {
 82  		return StateRef{}
 83  	}
 84  	return StateRef{Domain: StateDomainFilesystem, Scope: filepath.Clean(path)}
 85  }
 86  
 87  func processSessionStateRef() StateRef {
 88  	return StateRef{Domain: StateDomainProcess, Scope: "session"}
 89  }
 90  
 91  func resolveCallStateTraits(toolName, argsJSON string) CallStateTraits {
 92  	switch toolName {
 93  	case "browser_snapshot", "browser_take_screenshot", "browser_tabs":
 94  		return CallStateTraits{
 95  			Reads:     []StateRef{browserStateRef()},
 96  			Cacheable: true,
 97  		}
 98  	case "browser_navigate", "browser_click", "browser_type", "browser_press_key", "browser_drag", "browser_select_option":
 99  		return CallStateTraits{
100  			Writes: []StateRef{browserStateRef()},
101  		}
102  	case "file_read":
103  		if ref := filesystemStateRef(extractPathArg(argsJSON)); ref != (StateRef{}) {
104  			return CallStateTraits{
105  				Reads:     []StateRef{ref},
106  				Cacheable: true,
107  			}
108  		}
109  	case "file_write", "file_edit":
110  		if ref := filesystemStateRef(extractPathArg(argsJSON)); ref != (StateRef{}) {
111  			return CallStateTraits{
112  				Writes: []StateRef{ref},
113  			}
114  		}
115  	case "bash":
116  		return CallStateTraits{
117  			Writes:       []StateRef{processSessionStateRef()},
118  			UnknownWrite: true,
119  		}
120  	}
121  
122  	if strings.HasPrefix(toolName, "browser_") {
123  		return CallStateTraits{
124  			Writes: []StateRef{browserStateRef()},
125  		}
126  	}
127  
128  	return CallStateTraits{}
129  }
130  
131  func resolveFallbackReadStateTraits(tool Tool, argsJSON string) CallStateTraits {
132  	if tool == nil {
133  		return CallStateTraits{}
134  	}
135  	readOnly, ok := tool.(ReadOnlyChecker)
136  	if !ok || !readOnly.IsReadOnlyCall(argsJSON) {
137  		return CallStateTraits{}
138  	}
139  	return CallStateTraits{
140  		Reads:     []StateRef{processSessionStateRef()},
141  		Cacheable: true,
142  	}
143  }
144  
145  func buildStateAwareCacheKey(toolName string, args json.RawMessage, traits CallStateTraits, tracker *stateVersionTracker) string {
146  	if !traits.Cacheable || tracker == nil {
147  		return ""
148  	}
149  	base := toolName + "\x00" + normalizeJSON(args)
150  	fingerprint := tracker.fingerprint(traits.Reads)
151  	if fingerprint == "" {
152  		return ""
153  	}
154  	return base + "\x00" + fingerprint
155  }