state: expose estimateTokens/stripCold on interface; remove type assertion

This commit is contained in:
Levi Neely 2026-05-17 11:06:45 +02:00
parent 62774e60a6
commit 837c95c483
4 changed files with 20 additions and 4 deletions

View File

@ -75,10 +75,8 @@ func run(ctx context.Context, cfg agentConfig, state state) error {
for {
// Proactive context gate: strip cold material before calling the backend.
if s, ok := state.(*Session); ok {
if budget := contextBudget(ctx, cfg.Backend); budget > 0 && s.estimateTokens() > budget {
s.stripColdResults()
}
if budget := contextBudget(ctx, cfg.Backend); budget > 0 && state.estimateTokens() > budget {
state.stripCold()
}
sc := StepContext{

View File

@ -69,3 +69,13 @@ func (s *simpleState) update(msg backend.Message, results []toolResult) {
})
}
}
func (s *simpleState) estimateTokens() int {
chars := 0
for _, m := range s.msgs {
chars += len(m.Content)
}
return chars / 4
}
func (s *simpleState) stripCold() {}

View File

@ -531,6 +531,9 @@ func (s *Session) stripColdResults() {
}
}
// stripCold implements the state interface.
func (s *Session) stripCold() { s.stripColdResults() }
// contextDebug returns a multi-line breakdown of the history.
func (s *Session) contextDebug() string {
var sb strings.Builder

View File

@ -40,4 +40,9 @@ type state interface {
// Called once per loop iteration after Act completes.
update(assistant backend.Message, results []toolResult)
// estimateTokens returns a rough token count of the current history.
estimateTokens() int
// stripCold collapses large tool-result messages in-place to reduce context size.
stripCold()
}