feat(agent): add bounded cache, resilient KV, and agent loop hardening
- pkg/agent/bounded_cache.go: concurrent bounded cache with eviction for session→conversationID mappings to prevent unbounded growth - pkg/agent/resilient_kv.go: ResilientKV wraps KVDelegate with circuit breaker; degrades gracefully on persistent DB failures - agent_run, loop, summarizer: integrate bounded cache and resilient KV - command_handler, context, offloading_tool_runtime, securebus_runtime, toolloop: supporting updates - loop_test: expanded tests
This commit is contained in:
ZanzyTHEbar
parent
88667b29ec
commit
7f1400f6c0
11 changed files with 843 additions and 142 deletions
|
|
@ -21,6 +21,7 @@ import (
|
|||
"github.com/ZanzyTHEbar/dragonscale/pkg/messages"
|
||||
"github.com/ZanzyTHEbar/dragonscale/pkg/tools"
|
||||
"github.com/ZanzyTHEbar/dragonscale/pkg/utils"
|
||||
"golang.org/x/sync/errgroup"
|
||||
)
|
||||
|
||||
type assembledContext struct {
|
||||
|
|
@ -41,12 +42,12 @@ func (al *AgentLoop) prepareRuntimeState(ctx context.Context, sessionKey string)
|
|||
|
||||
var conversationID ids.UUID
|
||||
if cached, ok := al.conversationIDs.Load(sessionKey); ok {
|
||||
conversationID = cached.(ids.UUID)
|
||||
conversationID = cached
|
||||
} else {
|
||||
al.conversationMu.Lock()
|
||||
defer al.conversationMu.Unlock()
|
||||
if cached, ok := al.conversationIDs.Load(sessionKey); ok {
|
||||
conversationID = cached.(ids.UUID)
|
||||
conversationID = cached
|
||||
} else {
|
||||
conversationID = ids.New()
|
||||
title := sessionKey
|
||||
|
|
@ -130,18 +131,87 @@ func (al *AgentLoop) recordChannelState(ctx context.Context, opts processOptions
|
|||
return al.RecordLastChannel(ctx, channelKey)
|
||||
}
|
||||
|
||||
type ctxBlockCacheEntry struct {
|
||||
focusBlock string
|
||||
knowledge string
|
||||
cachedAt time.Time
|
||||
}
|
||||
|
||||
const ctxBlockCacheTTL = 2 * time.Minute
|
||||
|
||||
func (al *AgentLoop) refreshContextBlocks(ctx context.Context, opts processOptions) {
|
||||
al.updateToolContexts(opts.Channel, opts.ChatID)
|
||||
|
||||
block := al.obsManager.LoadBlock(ctx, opts.SessionKey)
|
||||
al.contextBuilder.SetObservationBlock(block)
|
||||
var (
|
||||
obsBlock string
|
||||
kb string
|
||||
focusBlock string
|
||||
)
|
||||
|
||||
kb := tools.LoadKnowledgeBlock(ctx, al.memDelegate, opts.SessionKey)
|
||||
al.contextBuilder.SetKnowledgeBlock(kb)
|
||||
// Check if focus/knowledge can be served from cache.
|
||||
// Invalidated by: (a) start_focus / complete_focus OnChange callbacks,
|
||||
// (b) TTL expiry — catches external KV modifications outside focus tools.
|
||||
_, dirty := al.focusDirty.LoadAndDelete(opts.SessionKey)
|
||||
|
||||
if al.identitySync != nil {
|
||||
_ = al.identitySync.CheckAndSync(ctx)
|
||||
useCached := false
|
||||
if !dirty {
|
||||
if cached, ok := al.ctxBlockCache.Load(opts.SessionKey); ok {
|
||||
entry := cached.(ctxBlockCacheEntry)
|
||||
if time.Since(entry.cachedAt) < ctxBlockCacheTTL {
|
||||
useCached = true
|
||||
kb = entry.knowledge
|
||||
focusBlock = entry.focusBlock
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
g, gCtx := errgroup.WithContext(ctx)
|
||||
|
||||
g.Go(func() error {
|
||||
obsBlock = al.obsManager.LoadBlock(gCtx, opts.SessionKey)
|
||||
return nil
|
||||
})
|
||||
if !useCached {
|
||||
g.Go(func() error {
|
||||
kb = tools.LoadKnowledgeBlock(gCtx, al.memDelegate, opts.SessionKey)
|
||||
return nil
|
||||
})
|
||||
g.Go(func() error {
|
||||
if fs, ok := tools.LoadFocusState(gCtx, al.memDelegate, opts.SessionKey); ok {
|
||||
focusBlock = fs.FormatBlock()
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
g.Go(func() error {
|
||||
if al.identitySync != nil {
|
||||
_ = al.identitySync.CheckAndSync(gCtx)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
|
||||
_ = g.Wait()
|
||||
|
||||
al.ctxBlockCache.Store(opts.SessionKey, ctxBlockCacheEntry{
|
||||
focusBlock: focusBlock,
|
||||
knowledge: kb,
|
||||
cachedAt: time.Now(),
|
||||
})
|
||||
|
||||
// Prune stale entries from other sessions to prevent unbounded growth.
|
||||
al.ctxBlockCache.Range(func(key, value any) bool {
|
||||
if entry, ok := value.(ctxBlockCacheEntry); ok {
|
||||
if time.Since(entry.cachedAt) > 2*ctxBlockCacheTTL {
|
||||
al.ctxBlockCache.Delete(key)
|
||||
al.focusDirty.Delete(key) // clean up orphaned dirty flags
|
||||
}
|
||||
}
|
||||
return true
|
||||
})
|
||||
|
||||
al.contextBuilder.SetObservationBlock(obsBlock)
|
||||
al.contextBuilder.SetKnowledgeBlock(kb)
|
||||
al.contextBuilder.SetFocusBlock(focusBlock)
|
||||
}
|
||||
|
||||
func (al *AgentLoop) loadSessionState(ctx context.Context, opts processOptions) ([]messages.Message, string) {
|
||||
|
|
@ -248,6 +318,7 @@ func (al *AgentLoop) createFantasyAgent(ctx context.Context, opts processOptions
|
|||
Queries: al.queries,
|
||||
ConversationID: conversationID,
|
||||
RunID: runID,
|
||||
ThresholdChars: al.offloadThresholdChars,
|
||||
}
|
||||
toolRuntime := SecureBusToolRuntime{
|
||||
Base: baseRuntime,
|
||||
|
|
@ -273,13 +344,18 @@ func (al *AgentLoop) createFantasyAgent(ctx context.Context, opts processOptions
|
|||
// postProcess handles the common finalization after Generate or Stream:
|
||||
// extract final text, save session, summarize, observe, optionally send response.
|
||||
func (al *AgentLoop) postProcess(ctx context.Context, opts processOptions, finalContent string, stepCount int) string {
|
||||
al.sessions.Save(opts.SessionKey)
|
||||
// Snapshot BEFORE summarization can truncate history, preventing
|
||||
// observation manager from seeing an incomplete view.
|
||||
tail := al.sessionsToMessagePairs(opts.SessionKey)
|
||||
|
||||
// Defer disk/DB persistence off the response path; in-memory state
|
||||
// is already consistent for observation/summarization reads.
|
||||
go al.sessions.Save(opts.SessionKey)
|
||||
|
||||
if opts.EnableSummary {
|
||||
al.maybeSummarize(ctx, opts.SessionKey, opts.Channel, opts.ChatID)
|
||||
go al.maybeSummarize(context.WithoutCancel(ctx), opts.SessionKey, opts.Channel, opts.ChatID)
|
||||
}
|
||||
|
||||
tail := al.sessionsToMessagePairs(opts.SessionKey)
|
||||
al.obsManager.MaybeObserveAsync(ctx, opts.SessionKey, tail)
|
||||
|
||||
if opts.SendResponse {
|
||||
|
|
@ -498,7 +574,7 @@ func (al *AgentLoop) runStreaming(ctx context.Context, opts processOptions, ac a
|
|||
// runLLMIteration — DELETED. Replaced by Fantasy's internal agent loop.
|
||||
|
||||
// auditStep logs tool calls from a Fantasy step result to the audit log.
|
||||
func (al *AgentLoop) auditStep(ctx context.Context, step fantasy.StepResult, sessionKey string) {
|
||||
func (al *AgentLoop) auditStep(_ context.Context, step fantasy.StepResult, sessionKey string) {
|
||||
toolCalls := step.Content.ToolCalls()
|
||||
if len(toolCalls) == 0 {
|
||||
return
|
||||
|
|
@ -513,12 +589,12 @@ func (al *AgentLoop) auditStep(ctx context.Context, step fantasy.StepResult, ses
|
|||
Target: tc.ToolName,
|
||||
Input: tc.Input,
|
||||
}
|
||||
aCtx, cancel := context.WithTimeout(ctx, time.Second)
|
||||
if err := al.memDelegate.InsertAuditEntry(aCtx, entry); err != nil {
|
||||
logger.WarnCF("agent", "Failed to log audit entry",
|
||||
map[string]interface{}{"tool": tc.ToolName, "error": err.Error()})
|
||||
select {
|
||||
case al.auditChan <- entry:
|
||||
default:
|
||||
logger.WarnCF("agent", "Audit channel full, dropping entry",
|
||||
map[string]interface{}{"tool": tc.ToolName})
|
||||
}
|
||||
cancel()
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
61
pkg/agent/bounded_cache.go
Normal file
61
pkg/agent/bounded_cache.go
Normal file
|
|
@ -0,0 +1,61 @@
|
|||
package agent
|
||||
|
||||
import "sync"
|
||||
|
||||
// boundedCache is a concurrent-safe, bounded key-value cache that evicts
|
||||
// the oldest half of entries when capacity is reached. Sufficient for
|
||||
// caching session→conversationID mappings where exact LRU ordering is
|
||||
// not critical but unbounded growth must be prevented.
|
||||
type boundedCache[K comparable, V any] struct {
|
||||
mu sync.RWMutex
|
||||
items map[K]V
|
||||
order []K
|
||||
capacity int
|
||||
}
|
||||
|
||||
func newBoundedCache[K comparable, V any](capacity int) *boundedCache[K, V] {
|
||||
if capacity <= 0 {
|
||||
capacity = 1024
|
||||
}
|
||||
return &boundedCache[K, V]{
|
||||
items: make(map[K]V, capacity),
|
||||
order: make([]K, 0, capacity),
|
||||
capacity: capacity,
|
||||
}
|
||||
}
|
||||
|
||||
func (c *boundedCache[K, V]) Load(key K) (V, bool) {
|
||||
c.mu.RLock()
|
||||
defer c.mu.RUnlock()
|
||||
v, ok := c.items[key]
|
||||
return v, ok
|
||||
}
|
||||
|
||||
func (c *boundedCache[K, V]) Store(key K, value V) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
if _, exists := c.items[key]; !exists {
|
||||
if len(c.items) >= c.capacity {
|
||||
c.evictHalfLocked()
|
||||
}
|
||||
c.order = append(c.order, key)
|
||||
}
|
||||
c.items[key] = value
|
||||
}
|
||||
|
||||
func (c *boundedCache[K, V]) evictHalfLocked() {
|
||||
evictCount := len(c.order) / 2
|
||||
if evictCount == 0 {
|
||||
evictCount = 1
|
||||
}
|
||||
for _, k := range c.order[:evictCount] {
|
||||
delete(c.items, k)
|
||||
}
|
||||
c.order = append(c.order[:0], c.order[evictCount:]...)
|
||||
}
|
||||
|
||||
func (c *boundedCache[K, V]) Len() int {
|
||||
c.mu.RLock()
|
||||
defer c.mu.RUnlock()
|
||||
return len(c.items)
|
||||
}
|
||||
|
|
@ -157,8 +157,7 @@ func (al *AgentLoop) handleSwitchChannel(ctx context.Context, target string) str
|
|||
if channel == "cli" {
|
||||
al.outputOverride.Store(outputTarget{})
|
||||
if al.state != nil {
|
||||
_ = al.state.SetLastChannel(ctx, "cli")
|
||||
_ = al.state.SetLastChatID(ctx, "")
|
||||
_ = al.state.SetChannelAndChatID(ctx, "cli", "")
|
||||
}
|
||||
return "Cleared output channel override to CLI defaults"
|
||||
}
|
||||
|
|
@ -185,11 +184,8 @@ func (al *AgentLoop) handleSwitchChannel(ctx context.Context, target string) str
|
|||
ChatID: chatID,
|
||||
})
|
||||
if al.state != nil {
|
||||
if err := al.state.SetLastChannel(ctx, channel); err != nil {
|
||||
return fmt.Sprintf("Output redirection set to %s:%s, but failed to persist channel: %v", channel, chatID, err)
|
||||
}
|
||||
if err := al.state.SetLastChatID(ctx, chatID); err != nil {
|
||||
return fmt.Sprintf("Output redirection set to %s:%s, but failed to persist chat id: %v", channel, chatID, err)
|
||||
if err := al.state.SetChannelAndChatID(ctx, channel, chatID); err != nil {
|
||||
return fmt.Sprintf("Output redirection set to %s:%s, but failed to persist state: %v", channel, chatID, err)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ import (
|
|||
"path/filepath"
|
||||
"runtime"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/ZanzyTHEbar/dragonscale/pkg"
|
||||
|
|
@ -26,11 +27,21 @@ type ContextBuilder struct {
|
|||
delegate memory.MemoryDelegate // Direct delegate for document loading (may be nil)
|
||||
tools *tools.ToolRegistry // Direct reference to tool registry
|
||||
observationBlock string // Pre-rendered observation block for prompt injection
|
||||
focusBlock string // Pre-rendered active focus block
|
||||
knowledgeBlock string // Pre-rendered knowledge block from Focus completions
|
||||
dagBlock string // Pre-rendered DAG compressed history
|
||||
contextWindow int // Max tokens for context window (0 = no limit)
|
||||
|
||||
cacheMu sync.Mutex
|
||||
skillsCache string
|
||||
skillsCacheAt time.Time
|
||||
skillsDirsMtime time.Time // last known mtime of skills directories
|
||||
bootstrapCache string
|
||||
bootstrapAt time.Time
|
||||
}
|
||||
|
||||
const promptCacheTTL = 30 * time.Second
|
||||
|
||||
func NewContextBuilder(workspace string) *ContextBuilder {
|
||||
// Primary skills dir: XDG data dir (installed skills).
|
||||
// Falls back to workspace/skills for legacy setups.
|
||||
|
|
@ -78,6 +89,11 @@ func (cb *ContextBuilder) SetObservationBlock(block string) {
|
|||
cb.observationBlock = block
|
||||
}
|
||||
|
||||
// SetFocusBlock sets the pre-rendered active focus block for prompt injection.
|
||||
func (cb *ContextBuilder) SetFocusBlock(block string) {
|
||||
cb.focusBlock = block
|
||||
}
|
||||
|
||||
// SetKnowledgeBlock sets the pre-rendered knowledge block from completed Focus sessions.
|
||||
func (cb *ContextBuilder) SetKnowledgeBlock(block string) {
|
||||
cb.knowledgeBlock = block
|
||||
|
|
@ -155,27 +171,28 @@ func (cb *ContextBuilder) buildToolsSection() string {
|
|||
return sb.String()
|
||||
}
|
||||
|
||||
type contextSection struct {
|
||||
name string
|
||||
content string
|
||||
priority int // lower = higher priority (kept first when trimming)
|
||||
}
|
||||
|
||||
func (cb *ContextBuilder) BuildSystemPrompt() string {
|
||||
type section struct {
|
||||
name string
|
||||
content string
|
||||
priority int // lower = higher priority (kept first when trimming)
|
||||
}
|
||||
|
||||
// Collect sections in priority order
|
||||
sections := []section{}
|
||||
sections := []contextSection{}
|
||||
|
||||
// P0: Core identity (always included)
|
||||
sections = append(sections, section{"identity", cb.getIdentity(), 0})
|
||||
sections = append(sections, contextSection{"identity", cb.getIdentity(), 0})
|
||||
|
||||
// P1: Bootstrap files (user identity)
|
||||
if bc := cb.LoadBootstrapFiles(); bc != "" {
|
||||
sections = append(sections, section{"bootstrap", bc, 1})
|
||||
// P1: Bootstrap files (user identity) — cached with TTL
|
||||
if bc := cb.cachedBootstrapFiles(); bc != "" {
|
||||
sections = append(sections, contextSection{"bootstrap", bc, 1})
|
||||
}
|
||||
|
||||
// P2: Skills index (lightweight Level 1 metadata)
|
||||
if summary := cb.skillsLoader.BuildSkillsSummary(); summary != "" {
|
||||
sections = append(sections, section{"skills", fmt.Sprintf(`# Skills
|
||||
// P2: Skills index (lightweight Level 1 metadata) — cached with TTL
|
||||
if summary := cb.cachedSkillsSummary(); summary != "" {
|
||||
sections = append(sections, contextSection{"skills", fmt.Sprintf(`# Skills
|
||||
|
||||
The following skills extend your capabilities. To use a skill:
|
||||
1. Use **skill_search** to find relevant skills by keyword
|
||||
|
|
@ -190,27 +207,34 @@ Do NOT assume skill content — always load before applying.
|
|||
// P3: Working context (hot tier — highly dynamic, high value)
|
||||
if cb.memoryStore != nil {
|
||||
if wc := cb.buildWorkingContextSection(); wc != "" {
|
||||
sections = append(sections, section{"working_context", wc, 3})
|
||||
sections = append(sections, contextSection{"working_context", wc, 3})
|
||||
}
|
||||
}
|
||||
|
||||
// P4: Observation block
|
||||
if cb.observationBlock != "" {
|
||||
sections = append(sections, section{"observations", "# Observations\n\n" + cb.observationBlock, 4})
|
||||
sections = append(sections, contextSection{"observations", "# Observations\n\n" + cb.observationBlock, 4})
|
||||
}
|
||||
|
||||
// P5: Knowledge block
|
||||
// P5: Active focus block (current investigation context)
|
||||
if cb.focusBlock != "" {
|
||||
sections = append(sections, contextSection{"focus", cb.focusBlock, 5})
|
||||
}
|
||||
|
||||
// P6: Knowledge block (historical completions)
|
||||
if cb.knowledgeBlock != "" {
|
||||
sections = append(sections, section{"knowledge", cb.knowledgeBlock, 5})
|
||||
sections = append(sections, contextSection{"knowledge", cb.knowledgeBlock, 6})
|
||||
}
|
||||
|
||||
// P6: DAG compressed history (lowest priority — can be reconstructed)
|
||||
// P7: DAG compressed history (lowest priority — can be reconstructed)
|
||||
if cb.dagBlock != "" {
|
||||
sections = append(sections, section{"dag", "# Conversation History (Compressed)\n\n" + cb.dagBlock, 6})
|
||||
sections = append(sections, contextSection{"dag", "# Conversation History (Compressed)\n\n" + cb.dagBlock, 7})
|
||||
}
|
||||
|
||||
// Token budget enforcement: if we exceed ~40% of context window for the
|
||||
// system prompt, trim lowest-priority sections first.
|
||||
// Proportional budget enforcement: each section gets a share of the
|
||||
// token budget proportional to its priority weight. Surplus from small
|
||||
// sections redistributes to higher-priority ones. Sections that still
|
||||
// exceed their allocation are truncated rather than dropped entirely.
|
||||
budgetTokens := cb.tokenBudgetTokens()
|
||||
totalTokens := 0
|
||||
sectionTokens := make([]int, len(sections))
|
||||
|
|
@ -220,19 +244,7 @@ Do NOT assume skill content — always load before applying.
|
|||
}
|
||||
|
||||
if budgetTokens > 0 && totalTokens > budgetTokens {
|
||||
logger.WarnCF("context", "System prompt exceeds token budget, trimming low-priority sections",
|
||||
map[string]interface{}{
|
||||
"total_tokens": totalTokens,
|
||||
"budget_tokens": budgetTokens,
|
||||
"sections": len(sections),
|
||||
})
|
||||
// Trim from lowest priority (highest number) first
|
||||
for i := len(sections) - 1; i >= 0 && totalTokens > budgetTokens; i-- {
|
||||
if sections[i].priority >= 5 { // only trim P5+ (knowledge, dag)
|
||||
totalTokens -= sectionTokens[i]
|
||||
sections[i].content = ""
|
||||
}
|
||||
}
|
||||
sections, sectionTokens = cb.applyProportionalBudget(sections, sectionTokens, budgetTokens)
|
||||
}
|
||||
|
||||
parts := make([]string, 0, len(sections))
|
||||
|
|
@ -262,10 +274,136 @@ func (cb *ContextBuilder) tokenBudgetTokens() int {
|
|||
if cb.contextWindow <= 0 {
|
||||
return 0
|
||||
}
|
||||
// Reserve ~40% of context window for system prompt
|
||||
return int(float64(cb.contextWindow) * 0.4)
|
||||
}
|
||||
|
||||
// priorityWeight maps section priority to a relative budget weight.
|
||||
// Higher priority (lower number) gets more weight.
|
||||
var priorityWeight = [8]float64{
|
||||
0: 0.25, // identity
|
||||
1: 0.18, // bootstrap
|
||||
2: 0.12, // skills
|
||||
3: 0.15, // working context
|
||||
4: 0.10, // observations
|
||||
5: 0.08, // focus
|
||||
6: 0.06, // knowledge
|
||||
7: 0.06, // DAG
|
||||
}
|
||||
|
||||
// applyProportionalBudget distributes tokens among sections using priority
|
||||
// weights. Sections that fit within their share keep full content. Surplus
|
||||
// redistributes to higher-priority sections. Oversized sections are truncated
|
||||
// to their share rather than dropped.
|
||||
func (cb *ContextBuilder) applyProportionalBudget(sections []contextSection, sectionTokens []int, budget int) ([]contextSection, []int) {
|
||||
n := len(sections)
|
||||
allocations := make([]int, n)
|
||||
totalWeight := 0.0
|
||||
for _, s := range sections {
|
||||
p := s.priority
|
||||
if p >= len(priorityWeight) {
|
||||
p = len(priorityWeight) - 1
|
||||
}
|
||||
totalWeight += priorityWeight[p]
|
||||
}
|
||||
|
||||
// First pass: proportional allocation
|
||||
for i, s := range sections {
|
||||
p := s.priority
|
||||
if p >= len(priorityWeight) {
|
||||
p = len(priorityWeight) - 1
|
||||
}
|
||||
allocations[i] = int(float64(budget) * priorityWeight[p] / totalWeight)
|
||||
}
|
||||
|
||||
// Second pass: redistribute surplus from sections that fit within allocation
|
||||
surplus := 0
|
||||
deficitIndices := []int{}
|
||||
for i := range sections {
|
||||
if sectionTokens[i] <= allocations[i] {
|
||||
surplus += allocations[i] - sectionTokens[i]
|
||||
allocations[i] = sectionTokens[i]
|
||||
} else {
|
||||
deficitIndices = append(deficitIndices, i)
|
||||
}
|
||||
}
|
||||
if surplus > 0 && len(deficitIndices) > 0 {
|
||||
share := surplus / len(deficitIndices)
|
||||
for _, idx := range deficitIndices {
|
||||
allocations[idx] += share
|
||||
}
|
||||
}
|
||||
|
||||
// Third pass: truncate oversized sections
|
||||
for i := range sections {
|
||||
if sectionTokens[i] > allocations[i] && allocations[i] > 0 {
|
||||
sections[i].content = truncateToTokenBudget(sections[i].content, allocations[i])
|
||||
sectionTokens[i] = allocations[i]
|
||||
}
|
||||
}
|
||||
|
||||
logger.WarnCF("context", "Proportional budget applied",
|
||||
map[string]interface{}{
|
||||
"budget": budget,
|
||||
"sections": n,
|
||||
"surplus": surplus,
|
||||
"deficits": len(deficitIndices),
|
||||
})
|
||||
|
||||
return sections, sectionTokens
|
||||
}
|
||||
|
||||
// truncateToTokenBudget trims content to fit within a token budget,
|
||||
// cutting at paragraph boundaries when possible.
|
||||
func truncateToTokenBudget(content string, budget int) string {
|
||||
if observation.EstimateTokens(content) <= budget {
|
||||
return content
|
||||
}
|
||||
|
||||
// Rough char estimate: 1 token ~= 4 chars
|
||||
maxChars := budget * 4
|
||||
if maxChars >= len(content) {
|
||||
return content
|
||||
}
|
||||
|
||||
truncated := content[:maxChars]
|
||||
// Try to cut at the last paragraph boundary
|
||||
if idx := strings.LastIndex(truncated, "\n\n"); idx > maxChars/2 {
|
||||
truncated = truncated[:idx]
|
||||
}
|
||||
return truncated + "\n\n[...section truncated to fit context budget]"
|
||||
}
|
||||
|
||||
func (cb *ContextBuilder) cachedSkillsSummary() string {
|
||||
cb.cacheMu.Lock()
|
||||
defer cb.cacheMu.Unlock()
|
||||
|
||||
if cb.skillsCache != "" {
|
||||
currentMtime := cb.skillsLoader.DirsMtime()
|
||||
if currentMtime.Equal(cb.skillsDirsMtime) && time.Since(cb.skillsCacheAt) < promptCacheTTL {
|
||||
return cb.skillsCache
|
||||
}
|
||||
cb.skillsDirsMtime = currentMtime
|
||||
}
|
||||
|
||||
cb.skillsCache = cb.skillsLoader.BuildSkillsSummary()
|
||||
cb.skillsCacheAt = time.Now()
|
||||
if cb.skillsDirsMtime.IsZero() {
|
||||
cb.skillsDirsMtime = cb.skillsLoader.DirsMtime()
|
||||
}
|
||||
return cb.skillsCache
|
||||
}
|
||||
|
||||
func (cb *ContextBuilder) cachedBootstrapFiles() string {
|
||||
cb.cacheMu.Lock()
|
||||
defer cb.cacheMu.Unlock()
|
||||
if cb.bootstrapCache != "" && time.Since(cb.bootstrapAt) < promptCacheTTL {
|
||||
return cb.bootstrapCache
|
||||
}
|
||||
cb.bootstrapCache = cb.LoadBootstrapFiles()
|
||||
cb.bootstrapAt = time.Now()
|
||||
return cb.bootstrapCache
|
||||
}
|
||||
|
||||
func (cb *ContextBuilder) LoadBootstrapFiles() string {
|
||||
if cb.delegate == nil {
|
||||
return ""
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ import (
|
|||
"github.com/ZanzyTHEbar/dragonscale/pkg/channels"
|
||||
"github.com/ZanzyTHEbar/dragonscale/pkg/config"
|
||||
"github.com/ZanzyTHEbar/dragonscale/pkg/constants"
|
||||
"github.com/ZanzyTHEbar/dragonscale/pkg/ids"
|
||||
"github.com/ZanzyTHEbar/dragonscale/pkg/logger"
|
||||
"github.com/ZanzyTHEbar/dragonscale/pkg/memory"
|
||||
"github.com/ZanzyTHEbar/dragonscale/pkg/memory/dag"
|
||||
|
|
@ -37,35 +38,41 @@ import (
|
|||
)
|
||||
|
||||
type AgentLoop struct {
|
||||
bus *bus.MessageBus
|
||||
languageModel fantasy.LanguageModel
|
||||
workspace string
|
||||
model string
|
||||
contextWindow int // Maximum context window size in tokens
|
||||
maxIterations int
|
||||
sessions *session.SessionManager
|
||||
state *state.Manager
|
||||
contextBuilder *ContextBuilder
|
||||
tools *tools.ToolRegistry
|
||||
memoryStore *memstore.MemoryStore // 3-tier MemGPT memory (always initialized)
|
||||
memDelegate memory.MemoryDelegate // DB delegate (always initialized)
|
||||
obsManager *observation.Manager // Observational memory (always initialized)
|
||||
secureBus *securebus.Bus // ITR SecureBus (always initialized)
|
||||
queries *memsqlc.Queries // SQL query surface for runtime persistence
|
||||
kvDelegate KVDelegate // KV adapter for offloaded tool results
|
||||
stateStore *StateStore // Agent run state persistence
|
||||
conversationIDs sync.Map // Owner: agent_run.go — wrote by prepareRuntimeState, read in prepareRuntimeState/load path
|
||||
conversationMu sync.Mutex // serializes conversation creation path
|
||||
identitySync *dragonsync.IdentitySync // File→DB sync for identity docs (nil if memory disabled)
|
||||
activeSessionKey atomic.Value // Owner: agent_run.go — written in runAgentLoop, read by router/toolloop for context routing
|
||||
running atomic.Bool // Owner: loop.go — lifecycle gate controlled by Run/Stop only
|
||||
summarizing sync.Map // Owner: summarizer.go — intended for async summarization lockout, currently gated by TODO path
|
||||
summarizeFailures sync.Map // Owner: summarizer.go — write/read in forceCompression + summarizeSession error paths
|
||||
cfg *config.Config // Stored for subagent factory access
|
||||
channelManager *channels.Manager
|
||||
commandRegistry []SlashCommand
|
||||
outputOverride atomic.Value // Owner: command_handler.go — CLI output redirection target for internal messages
|
||||
toolResultSearch fantasy.AgentTool
|
||||
bus *bus.MessageBus
|
||||
languageModel fantasy.LanguageModel
|
||||
workspace string
|
||||
model string
|
||||
contextWindow int // Maximum context window size in tokens
|
||||
maxIterations int
|
||||
sessions *session.SessionManager
|
||||
state *state.Manager
|
||||
contextBuilder *ContextBuilder
|
||||
tools *tools.ToolRegistry
|
||||
memoryStore *memstore.MemoryStore // 3-tier MemGPT memory (always initialized)
|
||||
memDelegate memory.MemoryDelegate // DB delegate (always initialized)
|
||||
obsManager *observation.Manager // Observational memory (always initialized)
|
||||
secureBus *securebus.Bus // ITR SecureBus (always initialized)
|
||||
queries *memsqlc.Queries // SQL query surface for runtime persistence
|
||||
kvDelegate KVDelegate // KV adapter for offloaded tool results
|
||||
stateStore *StateStore // Agent run state persistence
|
||||
offloadThresholdChars int // Char threshold for tool result offloading (derived from token config)
|
||||
conversationIDs *boundedCache[string, ids.UUID] // Owner: agent_run.go — wrote by prepareRuntimeState, read in prepareRuntimeState/load path
|
||||
conversationMu sync.Mutex // serializes conversation creation path
|
||||
identitySync *dragonsync.IdentitySync // File→DB sync for identity docs (nil if memory disabled)
|
||||
activeSessionKey atomic.Value // Owner: agent_run.go — written in runAgentLoop, read by router/toolloop for context routing
|
||||
running atomic.Bool // Owner: loop.go — lifecycle gate controlled by Run/Stop only
|
||||
summarizing sync.Map // Owner: summarizer.go — intended for async summarization lockout, currently gated by TODO path
|
||||
summarizeFailures sync.Map // Owner: summarizer.go — write/read in forceCompression + summarizeSession error paths
|
||||
dagCache sync.Map // Owner: summarizer.go — sessionKey → dagCacheEntry, skips recompression when compressible count unchanged
|
||||
auditChan chan *memory.AuditEntry // Buffered channel for async audit logging; drained by background worker
|
||||
auditDone chan struct{} // Closed when audit worker exits
|
||||
focusDirty sync.Map // sessionKey → struct{}: set by focus tool callbacks, cleared after context reload
|
||||
ctxBlockCache sync.Map // sessionKey → ctxBlockCacheEntry: cached focus + knowledge blocks
|
||||
cfg *config.Config // Stored for subagent factory access
|
||||
channelManager *channels.Manager
|
||||
commandRegistry []SlashCommand
|
||||
outputOverride atomic.Value // Owner: command_handler.go — CLI output redirection target for internal messages
|
||||
toolResultSearch fantasy.AgentTool
|
||||
}
|
||||
|
||||
type outputTarget struct {
|
||||
|
|
@ -143,7 +150,7 @@ func NewAgentLoop(ctx context.Context, cfg *config.Config, msgBus *bus.MessageBu
|
|||
del.Close()
|
||||
return nil, fmt.Errorf("memory delegate queries are not initialized")
|
||||
}
|
||||
kv := NewDelegateKV(memDelegate, pkg.NAME)
|
||||
kv := NewResilientKV(NewDelegateKV(memDelegate, pkg.NAME))
|
||||
stateStore := NewStateStore(queries)
|
||||
|
||||
offloadThreshold := cfg.Memory.OffloadThresholdTokens
|
||||
|
|
@ -277,13 +284,6 @@ func NewAgentLoop(ctx context.Context, cfg *config.Config, msgBus *bus.MessageBu
|
|||
toolsRegistry.MarkGateway(name)
|
||||
}
|
||||
|
||||
// Wire skills loader into tool_search for unified discovery
|
||||
if ts, ok := toolsRegistry.Get("tool_search"); ok {
|
||||
if tst, ok := ts.(*tools.ToolSearchTool); ok {
|
||||
tst.SetSkillsLoader(contextBuilder.SkillsLoader())
|
||||
}
|
||||
}
|
||||
|
||||
// Observation manager
|
||||
callModelFn := func(ctx context.Context, prompt string) (string, error) {
|
||||
temp := 0.3
|
||||
|
|
@ -302,30 +302,39 @@ func NewAgentLoop(ctx context.Context, cfg *config.Config, msgBus *bus.MessageBu
|
|||
}
|
||||
obsManager := observation.NewManager(memDelegate, pkg.NAME, callModelFn, observation.DefaultManagerConfig())
|
||||
|
||||
auditCh := make(chan *memory.AuditEntry, 256)
|
||||
auditDone := make(chan struct{})
|
||||
|
||||
al := &AgentLoop{
|
||||
bus: msgBus,
|
||||
languageModel: model,
|
||||
workspace: workspace,
|
||||
model: cfg.Agents.Defaults.Model,
|
||||
contextWindow: cfg.Agents.Defaults.MaxTokens,
|
||||
maxIterations: cfg.Agents.Defaults.MaxToolIterations,
|
||||
sessions: sessionsManager,
|
||||
state: stateManager,
|
||||
contextBuilder: contextBuilder,
|
||||
tools: toolsRegistry,
|
||||
memoryStore: ms,
|
||||
memDelegate: memDelegate,
|
||||
obsManager: obsManager,
|
||||
queries: queries,
|
||||
kvDelegate: kv,
|
||||
stateStore: stateStore,
|
||||
toolResultSearch: NewToolResultSearchTool(queries, kv),
|
||||
identitySync: idSync,
|
||||
summarizing: sync.Map{},
|
||||
commandRegistry: defaultSlashCommands(),
|
||||
cfg: cfg,
|
||||
bus: msgBus,
|
||||
languageModel: model,
|
||||
workspace: workspace,
|
||||
model: cfg.Agents.Defaults.Model,
|
||||
contextWindow: cfg.Agents.Defaults.MaxTokens,
|
||||
maxIterations: cfg.Agents.Defaults.MaxToolIterations,
|
||||
sessions: sessionsManager,
|
||||
state: stateManager,
|
||||
contextBuilder: contextBuilder,
|
||||
tools: toolsRegistry,
|
||||
memoryStore: ms,
|
||||
memDelegate: memDelegate,
|
||||
obsManager: obsManager,
|
||||
queries: queries,
|
||||
kvDelegate: kv,
|
||||
stateStore: stateStore,
|
||||
offloadThresholdChars: offloadThreshold * 4,
|
||||
toolResultSearch: NewToolResultSearchTool(queries, kv),
|
||||
conversationIDs: newBoundedCache[string, ids.UUID](1024),
|
||||
identitySync: idSync,
|
||||
summarizing: sync.Map{},
|
||||
auditChan: auditCh,
|
||||
auditDone: auditDone,
|
||||
commandRegistry: defaultSlashCommands(),
|
||||
cfg: cfg,
|
||||
}
|
||||
|
||||
go al.auditWorker(ctx, auditCh, auditDone)
|
||||
|
||||
for _, apply := range opts {
|
||||
if apply != nil {
|
||||
apply(al)
|
||||
|
|
@ -339,8 +348,25 @@ func NewAgentLoop(ctx context.Context, cfg *config.Config, msgBus *bus.MessageBu
|
|||
}
|
||||
return ""
|
||||
}
|
||||
toolsRegistry.Register(tools.NewStartFocusTool(memDelegate, sessionsManager, sessionKeyFn))
|
||||
toolsRegistry.Register(tools.NewCompleteFocusTool(memDelegate, sessionsManager, sessionKeyFn))
|
||||
focusInvalidate := func() {
|
||||
if sk := sessionKeyFn(); sk != "" {
|
||||
al.focusDirty.Store(sk, struct{}{})
|
||||
}
|
||||
}
|
||||
startFocus := tools.NewStartFocusTool(memDelegate, sessionsManager, sessionKeyFn)
|
||||
startFocus.OnChange = focusInvalidate
|
||||
toolsRegistry.Register(startFocus)
|
||||
completeFocus := tools.NewCompleteFocusTool(memDelegate, sessionsManager, sessionKeyFn)
|
||||
completeFocus.OnChange = focusInvalidate
|
||||
toolsRegistry.Register(completeFocus)
|
||||
toolsRegistry.Register(tools.NewFocusHistoryTool(memDelegate, sessionKeyFn))
|
||||
// Wire context into tool_search for focus-aware bias and unified discovery.
|
||||
if ts, ok := toolsRegistry.Get("tool_search"); ok {
|
||||
if tst, ok := ts.(*tools.ToolSearchTool); ok {
|
||||
tst.SetSkillsLoader(contextBuilder.SkillsLoader())
|
||||
tst.SetFocusContext(memDelegate, sessionKeyFn)
|
||||
}
|
||||
}
|
||||
|
||||
// DAG tools: dag_expand, dag_describe, dag_grep (require delegate-backed session)
|
||||
dagDeps := tools.DAGToolDeps{
|
||||
|
|
@ -418,6 +444,9 @@ func (al *AgentLoop) Run(ctx context.Context) error {
|
|||
|
||||
func (al *AgentLoop) Stop() {
|
||||
al.running.Store(false)
|
||||
close(al.auditChan)
|
||||
<-al.auditDone // wait for audit worker to drain
|
||||
al.sessions.Close()
|
||||
if al.identitySync != nil {
|
||||
al.identitySync.Close()
|
||||
}
|
||||
|
|
@ -430,6 +459,50 @@ func (al *AgentLoop) Stop() {
|
|||
}
|
||||
}
|
||||
|
||||
// auditWorker drains the audit channel, accumulating entries and flushing
|
||||
// either when the batch reaches 32 entries or after 100ms of inactivity.
|
||||
func (al *AgentLoop) auditWorker(_ context.Context, ch <-chan *memory.AuditEntry, done chan<- struct{}) {
|
||||
defer close(done)
|
||||
|
||||
const maxBatch = 32
|
||||
const flushInterval = 100 * time.Millisecond
|
||||
|
||||
buf := make([]*memory.AuditEntry, 0, maxBatch)
|
||||
ticker := time.NewTicker(flushInterval)
|
||||
defer ticker.Stop()
|
||||
|
||||
// Use context.Background so flushes succeed even during shutdown
|
||||
// when the parent context may already be cancelled.
|
||||
flush := func() {
|
||||
if len(buf) == 0 {
|
||||
return
|
||||
}
|
||||
fCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
if err := al.memDelegate.InsertAuditEntryBatch(fCtx, buf); err != nil {
|
||||
logger.WarnCF("agent", "Async audit batch insert failed",
|
||||
map[string]interface{}{"count": len(buf), "error": err.Error()})
|
||||
}
|
||||
cancel()
|
||||
buf = buf[:0]
|
||||
}
|
||||
|
||||
for {
|
||||
select {
|
||||
case entry, ok := <-ch:
|
||||
if !ok {
|
||||
flush()
|
||||
return
|
||||
}
|
||||
buf = append(buf, entry)
|
||||
if len(buf) >= maxBatch {
|
||||
flush()
|
||||
}
|
||||
case <-ticker.C:
|
||||
flush()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (al *AgentLoop) RegisterTool(tool tools.Tool) {
|
||||
al.tools.Register(tool)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -433,6 +433,53 @@ func TestToolContext_Updates(t *testing.T) {
|
|||
var _ tools.ContextualTool = ctxTool
|
||||
}
|
||||
|
||||
func TestStartupInfo_IncludesFocusTools(t *testing.T) {
|
||||
t.Parallel()
|
||||
tmpDir, err := os.MkdirTemp("", "agent-test-*")
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create temp dir: %v", err)
|
||||
}
|
||||
defer os.RemoveAll(tmpDir)
|
||||
|
||||
cfg := &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
Workspace: tmpDir,
|
||||
Model: "test-model",
|
||||
MaxTokens: 4096,
|
||||
MaxToolIterations: 10,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
msgBus := bus.NewMessageBus()
|
||||
model := newMockLanguageModel("")
|
||||
al := mustNewAgentLoop(t, cfg, msgBus, model)
|
||||
|
||||
info := al.GetStartupInfo()
|
||||
toolsInfo := info["tools"].(map[string]interface{})
|
||||
toolsList := toolsInfo["names"].([]string)
|
||||
|
||||
requiredTools := map[string]bool{
|
||||
"start_focus": false,
|
||||
"complete_focus": false,
|
||||
"focus_history": false,
|
||||
"tool_search": false,
|
||||
"tool_call": false,
|
||||
}
|
||||
for _, name := range toolsList {
|
||||
if _, ok := requiredTools[name]; ok {
|
||||
requiredTools[name] = true
|
||||
}
|
||||
}
|
||||
|
||||
for name, ok := range requiredTools {
|
||||
if !ok {
|
||||
t.Errorf("expected startup tool list to include %q", name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestToolRegistry_GetDefinitions verifies tool definitions can be retrieved
|
||||
func TestToolRegistry_GetDefinitions(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
|
@ -954,3 +1001,77 @@ func TestPersistOversizedRecoveryRefs_CreatesRecoverableReferences(t *testing.T)
|
|||
t.Fatalf("expected recovered content in output, got: %s", res.ForLLM)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRefreshContextBlocks_LoadsActiveFocusStateForPrompt(t *testing.T) {
|
||||
t.Parallel()
|
||||
tmpDir, err := os.MkdirTemp("", "agent-focus-context-test-*")
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create temp dir: %v", err)
|
||||
}
|
||||
defer os.RemoveAll(tmpDir)
|
||||
|
||||
cfg := &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
Workspace: tmpDir,
|
||||
Model: "test-model",
|
||||
MaxTokens: 4096,
|
||||
MaxToolIterations: 10,
|
||||
},
|
||||
},
|
||||
}
|
||||
msgBus := bus.NewMessageBus()
|
||||
model := newMockLanguageModel("ok")
|
||||
al := mustNewAgentLoop(t, cfg, msgBus, model)
|
||||
|
||||
sessionKey := "refresh-focus-session"
|
||||
startTool := tools.NewStartFocusTool(al.memDelegate, al.sessions, func() string { return sessionKey })
|
||||
result := startTool.Execute(context.Background(), map[string]interface{}{
|
||||
"topic": "investigate timeout issue",
|
||||
"goal": "reduce API latency",
|
||||
"steps": []interface{}{"collect traces", "analyze retries"},
|
||||
})
|
||||
if result.IsError {
|
||||
t.Fatalf("expected focus start to succeed, got: %s", result.ForLLM)
|
||||
}
|
||||
|
||||
al.refreshContextBlocks(context.Background(), processOptions{SessionKey: sessionKey})
|
||||
|
||||
prompt := al.contextBuilder.BuildSystemPrompt()
|
||||
if !strings.Contains(prompt, "# Focus") {
|
||||
t.Fatalf("expected focus section in prompt, got: %s", prompt)
|
||||
}
|
||||
if !strings.Contains(prompt, "## reduce API latency") {
|
||||
t.Fatalf("expected focus goal in prompt, got: %s", prompt)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRefreshContextBlocks_ClearsFocusBlockWhenStateMissing(t *testing.T) {
|
||||
t.Parallel()
|
||||
tmpDir, err := os.MkdirTemp("", "agent-focus-context-miss-test-*")
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create temp dir: %v", err)
|
||||
}
|
||||
defer os.RemoveAll(tmpDir)
|
||||
|
||||
cfg := &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
Workspace: tmpDir,
|
||||
Model: "test-model",
|
||||
MaxTokens: 4096,
|
||||
MaxToolIterations: 10,
|
||||
},
|
||||
},
|
||||
}
|
||||
msgBus := bus.NewMessageBus()
|
||||
model := newMockLanguageModel("ok")
|
||||
al := mustNewAgentLoop(t, cfg, msgBus, model)
|
||||
|
||||
sessionKey := "missing-focus-session"
|
||||
al.refreshContextBlocks(context.Background(), processOptions{SessionKey: sessionKey})
|
||||
prompt := al.contextBuilder.BuildSystemPrompt()
|
||||
if strings.Contains(prompt, "# Focus") {
|
||||
t.Fatalf("did not expect focus section when focus state is missing, got: %s", prompt)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -3,10 +3,11 @@ package agent
|
|||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
jsonv2 "github.com/go-json-experiment/json"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
jsonv2 "github.com/go-json-experiment/json"
|
||||
|
||||
"charm.land/fantasy"
|
||||
"github.com/ZanzyTHEbar/dragonscale/pkg/dserrors"
|
||||
"github.com/ZanzyTHEbar/dragonscale/pkg/ids"
|
||||
|
|
|
|||
123
pkg/agent/resilient_kv.go
Normal file
123
pkg/agent/resilient_kv.go
Normal file
|
|
@ -0,0 +1,123 @@
|
|||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/ZanzyTHEbar/dragonscale/pkg/logger"
|
||||
)
|
||||
|
||||
// ResilientKV wraps a KVDelegate with a circuit breaker that degrades gracefully
|
||||
// on persistent failures. When the circuit is open, Put operations are silently
|
||||
// dropped (with a warning log), while Get and Scan return empty results. This
|
||||
// prevents transient DB hiccups from killing entire agent turns.
|
||||
type ResilientKV struct {
|
||||
inner KVDelegate
|
||||
|
||||
mu sync.Mutex
|
||||
failureCount int
|
||||
lastFailureTime time.Time
|
||||
circuitOpen bool
|
||||
|
||||
failureThreshold int
|
||||
resetTimeout time.Duration
|
||||
}
|
||||
|
||||
// ResilientKVOption configures a ResilientKV.
|
||||
type ResilientKVOption func(*ResilientKV)
|
||||
|
||||
func WithFailureThreshold(n int) ResilientKVOption {
|
||||
return func(r *ResilientKV) { r.failureThreshold = n }
|
||||
}
|
||||
|
||||
func WithResetTimeout(d time.Duration) ResilientKVOption {
|
||||
return func(r *ResilientKV) { r.resetTimeout = d }
|
||||
}
|
||||
|
||||
func NewResilientKV(inner KVDelegate, opts ...ResilientKVOption) *ResilientKV {
|
||||
r := &ResilientKV{
|
||||
inner: inner,
|
||||
failureThreshold: 5,
|
||||
resetTimeout: 30 * time.Second,
|
||||
}
|
||||
for _, o := range opts {
|
||||
o(r)
|
||||
}
|
||||
return r
|
||||
}
|
||||
|
||||
func (r *ResilientKV) isOpen() bool {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
|
||||
if !r.circuitOpen {
|
||||
return false
|
||||
}
|
||||
if time.Since(r.lastFailureTime) > r.resetTimeout {
|
||||
r.circuitOpen = false
|
||||
r.failureCount = 0
|
||||
logger.InfoCF("kv", "Circuit breaker reset, attempting KV operations again", nil)
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (r *ResilientKV) recordSuccess() {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.failureCount = 0
|
||||
}
|
||||
|
||||
func (r *ResilientKV) recordFailure() {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.failureCount++
|
||||
r.lastFailureTime = time.Now()
|
||||
if r.failureCount >= r.failureThreshold {
|
||||
r.circuitOpen = true
|
||||
logger.WarnCF("kv", "Circuit breaker opened after repeated failures",
|
||||
map[string]interface{}{"failures": r.failureCount, "threshold": r.failureThreshold})
|
||||
}
|
||||
}
|
||||
|
||||
func (r *ResilientKV) Put(ctx context.Context, key string, value []byte) error {
|
||||
if r.isOpen() {
|
||||
logger.WarnCF("kv", "Circuit open, dropping KV Put",
|
||||
map[string]interface{}{"key": key})
|
||||
return nil
|
||||
}
|
||||
err := r.inner.Put(ctx, key, value)
|
||||
if err != nil {
|
||||
r.recordFailure()
|
||||
return err
|
||||
}
|
||||
r.recordSuccess()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *ResilientKV) Get(ctx context.Context, key string) ([]byte, error) {
|
||||
if r.isOpen() {
|
||||
return nil, nil
|
||||
}
|
||||
data, err := r.inner.Get(ctx, key)
|
||||
if err != nil {
|
||||
r.recordFailure()
|
||||
return nil, err
|
||||
}
|
||||
r.recordSuccess()
|
||||
return data, nil
|
||||
}
|
||||
|
||||
func (r *ResilientKV) Scan(ctx context.Context, prefix string) ([]string, error) {
|
||||
if r.isOpen() {
|
||||
return nil, nil
|
||||
}
|
||||
keys, err := r.inner.Scan(ctx, prefix)
|
||||
if err != nil {
|
||||
r.recordFailure()
|
||||
return nil, err
|
||||
}
|
||||
r.recordSuccess()
|
||||
return keys, nil
|
||||
}
|
||||
|
|
@ -4,6 +4,7 @@ import (
|
|||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
fantasy "charm.land/fantasy"
|
||||
"github.com/ZanzyTHEbar/dragonscale/pkg/ids"
|
||||
|
|
@ -55,27 +56,35 @@ func (r SecureBusToolRuntime) Execute(
|
|||
|
||||
results := make([]fantasy.ToolResultContent, 0, len(toolCalls))
|
||||
|
||||
type deferredState struct {
|
||||
step int
|
||||
state string
|
||||
snapshot map[string]any
|
||||
}
|
||||
var pendingStates []deferredState
|
||||
|
||||
for i, tc := range toolCalls {
|
||||
step := r.StepIndex + i
|
||||
r.recordRunState(ctx, step, "tool_call", map[string]any{
|
||||
pendingStates = append(pendingStates, deferredState{step, "tool_call", map[string]any{
|
||||
"tool_name": tc.ToolName,
|
||||
})
|
||||
}})
|
||||
|
||||
reqID := ids.New().String()
|
||||
req := itr.NewToolExecRequest(reqID, r.SessionKey, tc.ToolCallID, tc.ToolName, tc.Input)
|
||||
busResp := r.Bus.Execute(ctx, req)
|
||||
|
||||
if busResp.IsError {
|
||||
// Policy violation or secret resolution failure.
|
||||
sanitized := sanitizePolicyError(busResp.Result)
|
||||
tr := fantasy.ToolResultContent{
|
||||
ToolCallID: tc.ToolCallID,
|
||||
ToolName: tc.ToolName,
|
||||
Result: fantasy.ToolResultOutputContentError{Error: errors.New(busResp.Result)},
|
||||
Result: fantasy.ToolResultOutputContentError{Error: errors.New(sanitized)},
|
||||
}
|
||||
r.recordRunState(ctx, step, "tool_call_error", map[string]any{
|
||||
"tool_name": tc.ToolName,
|
||||
"error": busResp.Result,
|
||||
})
|
||||
pendingStates = append(pendingStates, deferredState{step, "tool_call_error", map[string]any{
|
||||
"tool_name": tc.ToolName,
|
||||
"error": busResp.Result,
|
||||
"error_safe": sanitized,
|
||||
}})
|
||||
results = append(results, tr)
|
||||
if onResult != nil {
|
||||
if err := onResult(tr); err != nil {
|
||||
|
|
@ -96,9 +105,9 @@ func (r SecureBusToolRuntime) Execute(
|
|||
br = overrideResultText(br, busResp.Result)
|
||||
}
|
||||
results = append(results, br)
|
||||
r.recordRunState(ctx, step, "tool_result", map[string]any{
|
||||
pendingStates = append(pendingStates, deferredState{step, "tool_result", map[string]any{
|
||||
"tool_name": tc.ToolName,
|
||||
})
|
||||
}})
|
||||
if onResult != nil {
|
||||
if err := onResult(br); err != nil {
|
||||
return results, err
|
||||
|
|
@ -107,6 +116,11 @@ func (r SecureBusToolRuntime) Execute(
|
|||
}
|
||||
}
|
||||
|
||||
// Flush all buffered state writes in one pass
|
||||
for _, ps := range pendingStates {
|
||||
r.recordRunState(ctx, ps.step, ps.state, ps.snapshot)
|
||||
}
|
||||
|
||||
return results, nil
|
||||
}
|
||||
|
||||
|
|
@ -117,6 +131,27 @@ func (r SecureBusToolRuntime) recordRunState(ctx context.Context, stepIndex int,
|
|||
_, _ = r.StateStore.AddRunState(ctx, r.RunID, stepIndex, fantasy.ReActState(state), snapshot)
|
||||
}
|
||||
|
||||
// sanitizePolicyError strips internal details from policy/bus errors before
|
||||
// they reach the LLM. The full error is preserved in audit state only.
|
||||
func sanitizePolicyError(raw string) string {
|
||||
switch {
|
||||
case strings.Contains(raw, "recursion depth"):
|
||||
return "policy violation: recursion limit exceeded"
|
||||
case strings.Contains(raw, "network access denied"):
|
||||
return "policy violation: network access denied"
|
||||
case strings.Contains(raw, "filesystem access denied"):
|
||||
return "policy violation: filesystem access denied"
|
||||
case strings.Contains(raw, "secret injection failed"):
|
||||
return "policy violation: unable to resolve required secrets"
|
||||
case strings.Contains(raw, "invalid args JSON"):
|
||||
return "policy violation: invalid tool arguments"
|
||||
case strings.Contains(raw, "policy violation"):
|
||||
return "policy violation: access denied"
|
||||
default:
|
||||
return "tool execution denied"
|
||||
}
|
||||
}
|
||||
|
||||
// overrideResultText replaces the text output of a ToolResultContent with
|
||||
// the redacted version produced by the SecureBus.
|
||||
func overrideResultText(tr fantasy.ToolResultContent, text string) fantasy.ToolResultContent {
|
||||
|
|
|
|||
|
|
@ -322,7 +322,13 @@ func (al *AgentLoop) summarizeSession(parentCtx context.Context, sessionKey stri
|
|||
if omitted && finalSummary != "" {
|
||||
recoveryRefs, err := al.persistOversizedRecoveryRefs(ctx, sessionKey, omittedMessages)
|
||||
if err != nil {
|
||||
logger.WarnCF("agent", "Failed to persist DAG recovery references for oversized messages",
|
||||
// Retry once with a fresh timeout
|
||||
retryCtx, retryCancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
recoveryRefs, err = al.persistOversizedRecoveryRefs(retryCtx, sessionKey, omittedMessages)
|
||||
retryCancel()
|
||||
}
|
||||
if err != nil {
|
||||
logger.ErrorCF("agent", "Failed to persist DAG recovery references after retry",
|
||||
map[string]interface{}{
|
||||
"session_key": sessionKey,
|
||||
"error": err.Error(),
|
||||
|
|
@ -337,6 +343,12 @@ func (al *AgentLoop) summarizeSession(parentCtx context.Context, sessionKey stri
|
|||
}
|
||||
}
|
||||
|
||||
// Quality gate: reject summaries that are too short to be useful.
|
||||
// A valid summary of even a short conversation should be 20+ chars.
|
||||
if len(strings.TrimSpace(finalSummary)) < 20 {
|
||||
finalSummary = ""
|
||||
}
|
||||
|
||||
if finalSummary != "" {
|
||||
al.sessions.SetSummary(sessionKey, finalSummary)
|
||||
al.sessions.TruncateHistory(sessionKey, keepLast)
|
||||
|
|
@ -411,10 +423,24 @@ func (al *AgentLoop) sessionsToMessagePairs(sessionKey string) []observation.Mes
|
|||
return pairs
|
||||
}
|
||||
|
||||
// dagCacheEntry holds the cached DAG compression output for a session,
|
||||
// enabling skip of recompression and repersistence when the compressible
|
||||
// portion hasn't grown since the last call.
|
||||
type dagCacheEntry struct {
|
||||
msgCount int
|
||||
rendered string
|
||||
dag *dag.DAG
|
||||
persistFailed bool
|
||||
}
|
||||
|
||||
// applyDAGCompression compresses old history into a DAG summary block and
|
||||
// returns only the tail messages that should be passed as raw conversation.
|
||||
// The compressed portion is injected into the system prompt via contextBuilder.
|
||||
// When memDelegate implements dag.DAGPersister, the DAG is persisted for dag_expand/describe/grep.
|
||||
//
|
||||
// Incremental optimization: if the compressible message count matches the
|
||||
// cached count, the previous DAG and rendered block are reused without
|
||||
// recompression or repersistence.
|
||||
func (al *AgentLoop) applyDAGCompression(ctx context.Context, sessionKey string, history []messages.Message) []messages.Message {
|
||||
const minHistoryForDAG = 16
|
||||
|
||||
|
|
@ -445,6 +471,19 @@ func (al *AgentLoop) applyDAGCompression(ctx context.Context, sessionKey string,
|
|||
return history
|
||||
}
|
||||
|
||||
// Check DAG cache: skip recompression if compressible count hasn't changed
|
||||
if cached, ok := al.dagCache.Load(sessionKey); ok {
|
||||
entry := cached.(dagCacheEntry)
|
||||
if entry.msgCount == len(compressible) {
|
||||
al.contextBuilder.SetDAGBlock(entry.rendered)
|
||||
// Retry failed persistence from previous turn
|
||||
if entry.persistFailed {
|
||||
al.retryDAGPersist(ctx, sessionKey, entry)
|
||||
}
|
||||
return tail
|
||||
}
|
||||
}
|
||||
|
||||
dagMsgs := make([]dag.Message, len(compressible))
|
||||
for i, m := range compressible {
|
||||
dagMsgs[i] = dag.Message{Role: m.Role, Content: m.Content}
|
||||
|
|
@ -456,17 +495,23 @@ func (al *AgentLoop) applyDAGCompression(ctx context.Context, sessionKey string,
|
|||
rendered := dag.RenderDAGForBudget(d, budget.DAGSummaries)
|
||||
al.contextBuilder.SetDAGBlock(rendered)
|
||||
|
||||
// Cache the result
|
||||
al.dagCache.Store(sessionKey, dagCacheEntry{
|
||||
msgCount: len(compressible),
|
||||
rendered: rendered,
|
||||
dag: d,
|
||||
})
|
||||
|
||||
// Persist DAG for dag_expand, dag_describe, dag_grep (additive; in-memory behavior unchanged)
|
||||
if dp, ok := al.memDelegate.(dag.DAGPersister); ok {
|
||||
if err := dp.PersistDAG(ctx, pkg.NAME, sessionKey, &dag.PersistSnapshot{
|
||||
FromMsgIdx: 0,
|
||||
ToMsgIdx: len(compressible),
|
||||
MsgCount: len(compressible),
|
||||
DAG: d,
|
||||
}); err != nil {
|
||||
logger.WarnCF("agent", "DAG persist failed (non-fatal)",
|
||||
map[string]interface{}{"error": err.Error(), "session_key": sessionKey})
|
||||
}
|
||||
persistOK := al.tryDAGPersist(ctx, sessionKey, d, len(compressible))
|
||||
if !persistOK {
|
||||
// Mark for retry on next cache hit
|
||||
al.dagCache.Store(sessionKey, dagCacheEntry{
|
||||
msgCount: len(compressible),
|
||||
rendered: rendered,
|
||||
dag: d,
|
||||
persistFailed: true,
|
||||
})
|
||||
}
|
||||
|
||||
logger.DebugCF("agent", "DAG compression applied",
|
||||
|
|
@ -475,11 +520,41 @@ func (al *AgentLoop) applyDAGCompression(ctx context.Context, sessionKey string,
|
|||
"compressed_msgs": len(compressible),
|
||||
"tail_msgs": len(tail),
|
||||
"dag_nodes": len(d.Nodes),
|
||||
"cache_hit": false,
|
||||
})
|
||||
|
||||
return tail
|
||||
}
|
||||
|
||||
func (al *AgentLoop) tryDAGPersist(ctx context.Context, sessionKey string, d *dag.DAG, msgCount int) bool {
|
||||
dp, ok := al.memDelegate.(dag.DAGPersister)
|
||||
if !ok {
|
||||
return true
|
||||
}
|
||||
if err := dp.PersistDAG(ctx, pkg.NAME, sessionKey, &dag.PersistSnapshot{
|
||||
FromMsgIdx: 0,
|
||||
ToMsgIdx: msgCount,
|
||||
MsgCount: msgCount,
|
||||
DAG: d,
|
||||
}); err != nil {
|
||||
logger.WarnCF("agent", "DAG persist failed (will retry next turn)",
|
||||
map[string]interface{}{"error": err.Error(), "session_key": sessionKey})
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (al *AgentLoop) retryDAGPersist(ctx context.Context, sessionKey string, entry dagCacheEntry) {
|
||||
if al.tryDAGPersist(ctx, sessionKey, entry.dag, entry.msgCount) {
|
||||
al.dagCache.Store(sessionKey, dagCacheEntry{
|
||||
msgCount: entry.msgCount,
|
||||
rendered: entry.rendered,
|
||||
dag: entry.dag,
|
||||
persistFailed: false,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func (al *AgentLoop) estimateTokens(msgs []messages.Message) int {
|
||||
pairs := make([]observation.MessagePair, 0, len(msgs))
|
||||
for _, m := range msgs {
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ import (
|
|||
fantasy "charm.land/fantasy"
|
||||
"github.com/ZanzyTHEbar/dragonscale/pkg"
|
||||
dragonfantasy "github.com/ZanzyTHEbar/dragonscale/pkg/fantasy"
|
||||
"github.com/ZanzyTHEbar/dragonscale/pkg/ids"
|
||||
"github.com/ZanzyTHEbar/dragonscale/pkg/logger"
|
||||
memstore "github.com/ZanzyTHEbar/dragonscale/pkg/memory/store"
|
||||
"github.com/ZanzyTHEbar/dragonscale/pkg/tools"
|
||||
|
|
@ -38,7 +39,7 @@ func MakeUnifiedRunLoopFunc(al *AgentLoop) tools.RunLoopFunc {
|
|||
if strings.TrimSpace(baseSession) == "" {
|
||||
baseSession = "subagent:default"
|
||||
}
|
||||
sessionKey := baseSession + "::subagent"
|
||||
sessionKey := fmt.Sprintf("%s::subagent::%s", baseSession, ids.New().String()[:8])
|
||||
|
||||
conversationID, runID, err := al.prepareRuntimeState(ctx, sessionKey)
|
||||
if err != nil {
|
||||
|
|
@ -51,6 +52,7 @@ func MakeUnifiedRunLoopFunc(al *AgentLoop) tools.RunLoopFunc {
|
|||
Queries: al.queries,
|
||||
ConversationID: conversationID,
|
||||
RunID: runID,
|
||||
ThresholdChars: al.offloadThresholdChars,
|
||||
}
|
||||
toolRuntime := SecureBusToolRuntime{
|
||||
Base: baseRuntime,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue