From a4325e2abe26f6beb388e0ba9f4ae596c9923b86 Mon Sep 17 00:00:00 2001 From: viettranx Date: Fri, 27 Feb 2026 12:11:25 +0700 Subject: [PATCH] feat: Improve agent context management with new compaction settings, session-level context window tracking, and calibrated token estimation for adaptive throttling. --- cmd/gateway.go | 11 ++++- internal/agent/loop.go | 15 ++++++ internal/agent/loop_history.go | 31 ++++++++++--- internal/agent/loop_tracing.go | 24 ++++++++++ internal/config/config.go | 2 + internal/sessions/manager.go | 46 +++++++++++++++++++ internal/store/file/sessions.go | 19 ++++++++ internal/store/pg/sessions.go | 35 ++++++++++++++ internal/store/session_store.go | 9 ++++ .../config-sections/compaction-section.tsx | 31 ++++++++++--- .../config-sections/config-section.tsx | 27 +++++++++++ .../context-pruning-section.tsx | 29 +++++------- .../config-sections/memory-section.tsx | 17 ++++--- .../config-sections/quality-gates-section.tsx | 15 +++--- .../config-sections/sandbox-section.tsx | 19 ++++---- .../config-sections/subagents-section.tsx | 16 +++---- .../config-sections/tool-policy-section.tsx | 13 ++---- ui/web/src/types/agent.ts | 2 + 18 files changed, 284 insertions(+), 77 deletions(-) diff --git a/cmd/gateway.go b/cmd/gateway.go index b53a94b4..ffe92dfb 100644 --- a/cmd/gateway.go +++ b/cmd/gateway.go @@ -822,10 +822,17 @@ func runGateway() { // Adaptive throttle: reduce per-session concurrency when nearing the summary threshold. // This prevents concurrent runs from racing with summarization. + // Uses calibrated token estimation (actual prompt tokens from last LLM call) + // and the agent's real context window (cached on session by the Loop). sched.SetTokenEstimateFunc(func(sessionKey string) (int, int) { history := sessStore.GetHistory(sessionKey) - tokens := agent.EstimateTokens(history) - return tokens, 200000 // default context window + lastPT, lastMC := sessStore.GetLastPromptTokens(sessionKey) + tokens := agent.EstimateTokensWithCalibration(history, lastPT, lastMC) + cw := sessStore.GetContextWindow(sessionKey) + if cw <= 0 { + cw = 200000 // fallback for sessions not yet processed + } + return tokens, cw }) // Subscribe to agent events for channel streaming/reaction forwarding. diff --git a/internal/agent/loop.go b/internal/agent/loop.go index 8ca0b392..1f276ff0 100644 --- a/internal/agent/loop.go +++ b/internal/agent/loop.go @@ -451,6 +451,12 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error) ) } + // 0. Cache agent's context window on the session (first run only). + // Enables scheduler's adaptive throttle to use the real value instead of hardcoded 200K. + if l.sessions.GetContextWindow(req.SessionKey) <= 0 { + l.sessions.SetContextWindow(req.SessionKey, l.contextWindow) + } + // 1. Build messages from session history history := l.sessions.GetHistory(req.SessionKey) summary := l.sessions.GetSummary(req.SessionKey) @@ -757,6 +763,15 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error) // Write session metadata (matching TS session entry updates) l.sessions.UpdateMetadata(req.SessionKey, l.model, l.provider.Name(), req.Channel) l.sessions.AccumulateTokens(req.SessionKey, int64(totalUsage.PromptTokens), int64(totalUsage.CompletionTokens)) + + // Calibrate token estimation: store actual prompt tokens + message count. + // Next time EstimateTokensWithCalibration() is called, it uses this as a base + // instead of the chars/3 heuristic (more accurate for multilingual content). + if totalUsage.PromptTokens > 0 { + msgCount := len(history) + len(pendingMsgs) + l.sessions.SetLastPromptTokens(req.SessionKey, totalUsage.PromptTokens, msgCount) + } + l.sessions.Save(req.SessionKey) // Bootstrap auto-cleanup: after enough conversation turns, remove BOOTSTRAP.md diff --git a/internal/agent/loop_history.go b/internal/agent/loop_history.go index 5dc0aefe..e85d47cf 100644 --- a/internal/agent/loop_history.go +++ b/internal/agent/loop_history.go @@ -263,10 +263,23 @@ func sanitizeHistory(msgs []providers.Message) []providers.Message { func (l *Loop) maybeSummarize(ctx context.Context, sessionKey string) { history := l.sessions.GetHistory(sessionKey) - tokenEstimate := EstimateTokens(history) - threshold := l.contextWindow * 75 / 100 - if len(history) <= 50 && tokenEstimate <= threshold { + // Use calibrated token estimation when available. + lastPT, lastMC := l.sessions.GetLastPromptTokens(sessionKey) + tokenEstimate := EstimateTokensWithCalibration(history, lastPT, lastMC) + + // Resolve compaction thresholds from config with sensible defaults. + historyShare := 0.75 + if l.compactionCfg != nil && l.compactionCfg.MaxHistoryShare > 0 { + historyShare = l.compactionCfg.MaxHistoryShare + } + minMessages := 50 + if l.compactionCfg != nil && l.compactionCfg.MinMessages > 0 { + minMessages = l.compactionCfg.MinMessages + } + + threshold := int(float64(l.contextWindow) * historyShare) + if len(history) <= minMessages && tokenEstimate <= threshold { return } @@ -287,6 +300,12 @@ func (l *Loop) maybeSummarize(ctx context.Context, sessionKey string) { l.runMemoryFlush(ctx, sessionKey, flushSettings) } + // Resolve keepLast before spawning goroutine (reads config under caller's scope). + keepLast := 4 + if l.compactionCfg != nil && l.compactionCfg.KeepLastMessages > 0 { + keepLast = l.compactionCfg.KeepLastMessages + } + // Summarize in background (holds the per-session lock until done) go func() { defer sessionMu.Unlock() @@ -294,7 +313,7 @@ func (l *Loop) maybeSummarize(ctx context.Context, sessionKey string) { // Re-check: history may have been truncated by a concurrent summarize // that finished between our threshold check and acquiring the lock. history := l.sessions.GetHistory(sessionKey) - if len(history) <= 4 { + if len(history) <= keepLast { return } @@ -302,7 +321,7 @@ func (l *Loop) maybeSummarize(ctx context.Context, sessionKey string) { defer cancel() summary := l.sessions.GetSummary(sessionKey) - toSummarize := history[:len(history)-4] + toSummarize := history[:len(history)-keepLast] var sb string for _, m := range toSummarize { @@ -330,7 +349,7 @@ func (l *Loop) maybeSummarize(ctx context.Context, sessionKey string) { } l.sessions.SetSummary(sessionKey, SanitizeAssistantContent(resp.Content)) - l.sessions.TruncateHistory(sessionKey, 4) + l.sessions.TruncateHistory(sessionKey, keepLast) l.sessions.IncrementCompaction(sessionKey) l.sessions.Save(sessionKey) }() diff --git a/internal/agent/loop_tracing.go b/internal/agent/loop_tracing.go index c8ce9e87..8721360f 100644 --- a/internal/agent/loop_tracing.go +++ b/internal/agent/loop_tracing.go @@ -246,3 +246,27 @@ func EstimateTokens(messages []providers.Message) int { } return total } + +// EstimateTokensWithCalibration uses actual prompt tokens from the last LLM +// response as a calibration base, then estimates only new messages on top. +// Falls back to EstimateTokens() when no calibration data is available. +func EstimateTokensWithCalibration(messages []providers.Message, lastPromptTokens, lastMsgCount int) int { + if lastPromptTokens <= 0 || lastMsgCount <= 0 { + return EstimateTokens(messages) + } + + currentCount := len(messages) + newMsgs := currentCount - lastMsgCount + if newMsgs <= 0 { + // No new messages since last calibration (or history was truncated). + // Use calibration value as-is; it's the best estimate we have. + return lastPromptTokens + } + + // Estimate only the new messages with the heuristic and add to base. + delta := 0 + for _, m := range messages[lastMsgCount:] { + delta += utf8.RuneCountInString(m.Content) / 3 + } + return lastPromptTokens + delta +} diff --git a/internal/config/config.go b/internal/config/config.go index 45c24680..b1def5ea 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -137,6 +137,8 @@ type AgentDefaults struct { type CompactionConfig struct { ReserveTokensFloor int `json:"reserveTokensFloor,omitempty"` // min reserve tokens (default 20000) MaxHistoryShare float64 `json:"maxHistoryShare,omitempty"` // max share of context for history (default 0.75) + MinMessages int `json:"minMessages,omitempty"` // min messages before compaction triggers (default 50) + KeepLastMessages int `json:"keepLastMessages,omitempty"` // messages to keep after compaction (default 4) MemoryFlush *MemoryFlushConfig `json:"memoryFlush,omitempty"` // pre-compaction flush } diff --git a/internal/sessions/manager.go b/internal/sessions/manager.go index c75c527b..8b447f95 100644 --- a/internal/sessions/manager.go +++ b/internal/sessions/manager.go @@ -32,6 +32,10 @@ type Session struct { Label string `json:"label,omitempty"` SpawnedBy string `json:"spawnedBy,omitempty"` SpawnDepth int `json:"spawnDepth,omitempty"` + + ContextWindow int `json:"contextWindow,omitempty"` + LastPromptTokens int `json:"lastPromptTokens,omitempty"` + LastMessageCount int `json:"lastMessageCount,omitempty"` } // Manager handles session lifecycle, persistence, and lookup. @@ -217,6 +221,45 @@ func (m *Manager) SetSpawnInfo(key, spawnedBy string, depth int) { } } +// SetContextWindow caches the agent's context window on the session. +func (m *Manager) SetContextWindow(key string, cw int) { + m.mu.Lock() + defer m.mu.Unlock() + if s, ok := m.sessions[key]; ok { + s.ContextWindow = cw + } +} + +// GetContextWindow returns the cached context window for a session (0 if unset). +func (m *Manager) GetContextWindow(key string) int { + m.mu.RLock() + defer m.mu.RUnlock() + if s, ok := m.sessions[key]; ok { + return s.ContextWindow + } + return 0 +} + +// SetLastPromptTokens records actual prompt tokens from the last LLM response. +func (m *Manager) SetLastPromptTokens(key string, tokens, msgCount int) { + m.mu.Lock() + defer m.mu.Unlock() + if s, ok := m.sessions[key]; ok { + s.LastPromptTokens = tokens + s.LastMessageCount = msgCount + } +} + +// GetLastPromptTokens returns the last known prompt tokens and message count. +func (m *Manager) GetLastPromptTokens(key string) (int, int) { + m.mu.RLock() + defer m.mu.RUnlock() + if s, ok := m.sessions[key]; ok { + return s.LastPromptTokens, s.LastMessageCount + } + return 0, 0 +} + // TruncateHistory keeps only the last N messages. func (m *Manager) TruncateHistory(key string, keepLast int) { m.mu.Lock() @@ -364,6 +407,9 @@ func (m *Manager) Save(key string) error { Label: s.Label, SpawnedBy: s.SpawnedBy, SpawnDepth: s.SpawnDepth, + ContextWindow: s.ContextWindow, + LastPromptTokens: s.LastPromptTokens, + LastMessageCount: s.LastMessageCount, } if len(s.Messages) > 0 { snapshot.Messages = make([]providers.Message, len(s.Messages)) diff --git a/internal/store/file/sessions.go b/internal/store/file/sessions.go index 106e9776..01747315 100644 --- a/internal/store/file/sessions.go +++ b/internal/store/file/sessions.go @@ -74,6 +74,22 @@ func (f *FileSessionStore) SetSpawnInfo(key, spawnedBy string, depth int) { f.mgr.SetSpawnInfo(key, spawnedBy, depth) } +func (f *FileSessionStore) SetContextWindow(key string, cw int) { + f.mgr.SetContextWindow(key, cw) +} + +func (f *FileSessionStore) GetContextWindow(key string) int { + return f.mgr.GetContextWindow(key) +} + +func (f *FileSessionStore) SetLastPromptTokens(key string, tokens, msgCount int) { + f.mgr.SetLastPromptTokens(key, tokens, msgCount) +} + +func (f *FileSessionStore) GetLastPromptTokens(key string) (int, int) { + return f.mgr.GetLastPromptTokens(key) +} + func (f *FileSessionStore) TruncateHistory(key string, keepLast int) { f.mgr.TruncateHistory(key, keepLast) } @@ -154,5 +170,8 @@ func sessionToData(s *sessions.Session) *store.SessionData { Label: s.Label, SpawnedBy: s.SpawnedBy, SpawnDepth: s.SpawnDepth, + ContextWindow: s.ContextWindow, + LastPromptTokens: s.LastPromptTokens, + LastMessageCount: s.LastMessageCount, } } diff --git a/internal/store/pg/sessions.go b/internal/store/pg/sessions.go index 01089561..40f90595 100644 --- a/internal/store/pg/sessions.go +++ b/internal/store/pg/sessions.go @@ -210,6 +210,41 @@ func (s *PGSessionStore) SetSpawnInfo(key, spawnedBy string, depth int) { } } +func (s *PGSessionStore) SetContextWindow(key string, cw int) { + s.mu.Lock() + defer s.mu.Unlock() + if data, ok := s.cache[key]; ok { + data.ContextWindow = cw + } +} + +func (s *PGSessionStore) GetContextWindow(key string) int { + s.mu.RLock() + defer s.mu.RUnlock() + if data, ok := s.cache[key]; ok { + return data.ContextWindow + } + return 0 +} + +func (s *PGSessionStore) SetLastPromptTokens(key string, tokens, msgCount int) { + s.mu.Lock() + defer s.mu.Unlock() + if data, ok := s.cache[key]; ok { + data.LastPromptTokens = tokens + data.LastMessageCount = msgCount + } +} + +func (s *PGSessionStore) GetLastPromptTokens(key string) (int, int) { + s.mu.RLock() + defer s.mu.RUnlock() + if data, ok := s.cache[key]; ok { + return data.LastPromptTokens, data.LastMessageCount + } + return 0, 0 +} + func (s *PGSessionStore) TruncateHistory(key string, keepLast int) { s.mu.Lock() defer s.mu.Unlock() diff --git a/internal/store/session_store.go b/internal/store/session_store.go index 6612d70a..1bb88a74 100644 --- a/internal/store/session_store.go +++ b/internal/store/session_store.go @@ -29,6 +29,11 @@ type SessionData struct { Label string `json:"label,omitempty"` SpawnedBy string `json:"spawnedBy,omitempty"` SpawnDepth int `json:"spawnDepth,omitempty"` + + // Adaptive throttle: cached per-session so scheduler reads without DB lookup. + ContextWindow int `json:"contextWindow,omitempty"` // agent's context window (set on first run) + LastPromptTokens int `json:"lastPromptTokens,omitempty"` // actual prompt tokens from last LLM response + LastMessageCount int `json:"lastMessageCount,omitempty"` // message count at time of last LLM call } // SessionInfo is lightweight session metadata for listing. @@ -68,6 +73,10 @@ type SessionStore interface { GetMemoryFlushCompactionCount(key string) int SetMemoryFlushDone(key string) SetSpawnInfo(key, spawnedBy string, depth int) + SetContextWindow(key string, cw int) + GetContextWindow(key string) int + SetLastPromptTokens(key string, tokens, msgCount int) + GetLastPromptTokens(key string) (tokens, msgCount int) TruncateHistory(key string, keepLast int) Reset(key string) Delete(key string) error diff --git a/ui/web/src/pages/agents/agent-detail/config-sections/compaction-section.tsx b/ui/web/src/pages/agents/agent-detail/config-sections/compaction-section.tsx index d9f46f30..eb8b9fca 100644 --- a/ui/web/src/pages/agents/agent-detail/config-sections/compaction-section.tsx +++ b/ui/web/src/pages/agents/agent-detail/config-sections/compaction-section.tsx @@ -1,8 +1,7 @@ import { Input } from "@/components/ui/input"; -import { Label } from "@/components/ui/label"; import { Switch } from "@/components/ui/switch"; import type { CompactionConfig } from "@/types/agent"; -import { ConfigSection, numOrUndef } from "./config-section"; +import { ConfigSection, InfoLabel, numOrUndef } from "./config-section"; interface CompactionSectionProps { enabled: boolean; @@ -21,7 +20,7 @@ export function CompactionSection({ enabled, value, onToggle, onChange }: Compac >
- + Reserve Tokens Floor
- + Max History Share (0-1)
+
+
+ Min Messages + onChange({ ...value, minMessages: numOrUndef(e.target.value) })} + /> +
+
+ Keep Last Messages + onChange({ ...value, keepLastMessages: numOrUndef(e.target.value) })} + /> +
+
- + Memory Flush
{value.memoryFlush?.enabled && (
- + Soft Threshold (tokens) + + + + + + + + {tip} + + + +
+ ); +} interface ConfigSectionProps { title: string; diff --git a/ui/web/src/pages/agents/agent-detail/config-sections/context-pruning-section.tsx b/ui/web/src/pages/agents/agent-detail/config-sections/context-pruning-section.tsx index e67cc7df..a74be7f2 100644 --- a/ui/web/src/pages/agents/agent-detail/config-sections/context-pruning-section.tsx +++ b/ui/web/src/pages/agents/agent-detail/config-sections/context-pruning-section.tsx @@ -1,5 +1,4 @@ import { Input } from "@/components/ui/input"; -import { Label } from "@/components/ui/label"; import { Switch } from "@/components/ui/switch"; import { Select, @@ -9,7 +8,7 @@ import { SelectValue, } from "@/components/ui/select"; import type { ContextPruningConfig } from "@/types/agent"; -import { ConfigSection, numOrUndef } from "./config-section"; +import { ConfigSection, InfoLabel, numOrUndef } from "./config-section"; interface ContextPruningSectionProps { enabled: boolean; @@ -28,7 +27,7 @@ export function ContextPruningSection({ enabled, value, onToggle, onChange }: Co >
- + Mode
- + Keep Last Assistants
- + Soft Trim Ratio (0-1)
- + Hard Clear Ratio (0-1)
- + Min Prunable Tool Chars Soft Trim
- + Max Chars
- + Head Chars
- + Tail Chars
-

- Tool results longer than Max Chars get trimmed, keeping Head + Tail chars. -

{/* Hard Clear */} @@ -141,10 +137,10 @@ export function ContextPruningSection({ enabled, value, onToggle, onChange }: Co onChange({ ...value, hardClear: { ...value.hardClear, enabled: v } }) } /> - + Enabled
- + Placeholder Text -

- Replacement text when old tool results are cleared entirely. -

diff --git a/ui/web/src/pages/agents/agent-detail/config-sections/memory-section.tsx b/ui/web/src/pages/agents/agent-detail/config-sections/memory-section.tsx index f1749d88..3baa1856 100644 --- a/ui/web/src/pages/agents/agent-detail/config-sections/memory-section.tsx +++ b/ui/web/src/pages/agents/agent-detail/config-sections/memory-section.tsx @@ -1,8 +1,7 @@ import { Input } from "@/components/ui/input"; -import { Label } from "@/components/ui/label"; import { Switch } from "@/components/ui/switch"; import type { MemoryConfig } from "@/types/agent"; -import { ConfigSection, numOrUndef } from "./config-section"; +import { ConfigSection, InfoLabel, numOrUndef } from "./config-section"; interface MemorySectionProps { enabled: boolean; @@ -24,11 +23,11 @@ export function MemorySection({ enabled, value, onToggle, onChange }: MemorySect checked={value.enabled ?? true} onCheckedChange={(v) => onChange({ ...value, enabled: v })} /> - + Enabled
- + Embedding Provider
- + Embedding Model
- + Max Results
- + Max Chunk Length
- + Vector Weight
- + Text Weight
- + Type ) : ( <> - + Command updateGate(i, { block_on_failure: v })} /> - + Block on Failure
{gate.block_on_failure && (
- + Max Retries
- + Timeout (seconds)
- + Mode
- + Workspace Access
- + Scope
- + Timeout (sec)
- + Memory (MB)
- + CPUs onChange({ ...value, network_enabled: v })} /> - + Network Enabled
); diff --git a/ui/web/src/pages/agents/agent-detail/config-sections/subagents-section.tsx b/ui/web/src/pages/agents/agent-detail/config-sections/subagents-section.tsx index 48cde350..7a6d7df9 100644 --- a/ui/web/src/pages/agents/agent-detail/config-sections/subagents-section.tsx +++ b/ui/web/src/pages/agents/agent-detail/config-sections/subagents-section.tsx @@ -1,5 +1,4 @@ import { Input } from "@/components/ui/input"; -import { Label } from "@/components/ui/label"; import { Select, SelectContent, @@ -8,7 +7,7 @@ import { SelectValue, } from "@/components/ui/select"; import type { SubagentsConfig } from "@/types/agent"; -import { ConfigSection, numOrUndef } from "./config-section"; +import { ConfigSection, InfoLabel, numOrUndef } from "./config-section"; interface SubagentsSectionProps { enabled: boolean; @@ -27,7 +26,7 @@ export function SubagentsSection({ enabled, value, onToggle, onChange }: Subagen >
- + Max Concurrent
- + Max Spawn Depth
- + Archive After (minutes)
- + Model Override onChange({ ...value, model: e.target.value || undefined })} /> -

- LLM model for sub-agents. Leave empty to inherit the parent agent's model. -

); diff --git a/ui/web/src/pages/agents/agent-detail/config-sections/tool-policy-section.tsx b/ui/web/src/pages/agents/agent-detail/config-sections/tool-policy-section.tsx index 557ec5ed..e3d8b321 100644 --- a/ui/web/src/pages/agents/agent-detail/config-sections/tool-policy-section.tsx +++ b/ui/web/src/pages/agents/agent-detail/config-sections/tool-policy-section.tsx @@ -1,5 +1,4 @@ import { Input } from "@/components/ui/input"; -import { Label } from "@/components/ui/label"; import { Select, SelectContent, @@ -8,7 +7,7 @@ import { SelectValue, } from "@/components/ui/select"; import type { ToolPolicyConfig } from "@/types/agent"; -import { ConfigSection, arrayToTags, tagsToArray } from "./config-section"; +import { ConfigSection, InfoLabel, arrayToTags, tagsToArray } from "./config-section"; interface ToolPolicySectionProps { enabled: boolean; @@ -26,7 +25,7 @@ export function ToolPolicySection({ enabled, value, onToggle, onChange }: ToolPo onToggle={onToggle} >
- + Profile
- + Allow onChange({ ...value, allow: tagsToArray(e.target.value) })} /> -

Comma-separated list of allowed tools.

- + Deny
- + Also Allow onChange({ ...value, alsoAllow: tagsToArray(e.target.value) })} /> -

Additional tools to add on top of profile defaults.

); diff --git a/ui/web/src/types/agent.ts b/ui/web/src/types/agent.ts index 4fa21e99..75a0373d 100644 --- a/ui/web/src/types/agent.ts +++ b/ui/web/src/types/agent.ts @@ -20,6 +20,8 @@ export interface SubagentsConfig { export interface CompactionConfig { reserveTokensFloor?: number; maxHistoryShare?: number; + minMessages?: number; + keepLastMessages?: number; memoryFlush?: { enabled?: boolean; softThresholdTokens?: number;