feat(ai): add knowledge accumulation and enhance agentic loop

- Introduce KnowledgeAccumulator to persist facts across turns
- Enhance AgenticLoop to support knowledge injection and final text summaries
- Update chat service to wire up knowledge components
- Frontend updates to support enhanced chat capabilities
This commit is contained in:
rcourtman
2026-01-31 16:22:50 +00:00
parent 95a0d7a6bd
commit c5717d1a10
13 changed files with 1669 additions and 32 deletions
@@ -96,6 +96,18 @@ export function useChat(options: UseChatOptions = {}) {
return '';
};
const extractTokens = (data: unknown): { input: number; output: number } | null => {
if (!data || typeof data !== 'object') return null;
const record = data as Record<string, unknown>;
const input = Number(record.input_tokens ?? record.inputTokens ?? record.input);
const output = Number(record.output_tokens ?? record.outputTokens ?? record.output);
if (!Number.isFinite(input) && !Number.isFinite(output)) return null;
return {
input: Number.isFinite(input) && input > 0 ? input : 0,
output: Number.isFinite(output) && output > 0 ? output : 0,
};
};
const processEvent = (
assistantId: string,
event: StreamEvent
@@ -281,6 +293,10 @@ export function useChat(options: UseChatOptions = {}) {
}
case 'done': {
const tokens = extractTokens(event.data);
if (tokens && (tokens.input > 0 || tokens.output > 0)) {
return { ...msg, isStreaming: false, pendingTools: [], tokens };
}
return { ...msg, isStreaming: false, pendingTools: [] };
}
@@ -104,6 +104,7 @@ export const AIChat: Component<AIChatProps> = (props) => {
// Chat hook
const chat = useChat({ model: initialModelSelections[DEFAULT_SESSION_KEY] || '' });
const defaultModelLabel = createMemo(() => {
const fallback = defaultModel().trim();
if (!fallback) return '';
@@ -131,3 +131,9 @@ export interface StreamCompleteData {
output_tokens: number;
tool_calls?: ToolExecution[];
}
export interface StreamDoneData {
session_id?: string;
input_tokens?: number;
output_tokens?: number;
}
+220 -15
View File
@@ -40,6 +40,9 @@ type AgenticLoop struct {
// Per-session FSMs for workflow enforcement (set before each execution)
sessionFSM *SessionFSM
// Knowledge accumulator for fact extraction across turns
knowledgeAccumulator *KnowledgeAccumulator
// Budget checker called after each turn to enforce token spending limits
budgetChecker func() error
}
@@ -78,8 +81,10 @@ func (a *AgenticLoop) ExecuteWithTools(ctx context.Context, sessionID string, me
}
func (a *AgenticLoop) executeWithTools(ctx context.Context, sessionID string, messages []Message, tools []providers.Tool, callback StreamCallback) ([]Message, error) {
// Track this session for potential abort
// Snapshot maxTurns under the lock — callers may override via SetMaxTurns
// before calling ExecuteWithTools, and this avoids races with concurrent sessions.
a.mu.Lock()
maxTurns := a.maxTurns
a.aborted[sessionID] = false
a.mu.Unlock()
defer func() {
@@ -113,11 +118,11 @@ func (a *AgenticLoop) executeWithTools(ctx context.Context, sessionID string, me
// Track where each turn's messages begin in providerMessages for compaction.
// We keep the last N turns' tool results in full; older ones get compacted.
const compactionKeepTurns = 3 // Keep last 3 turns' tool results in full
const compactionMinChars = 500 // Only compact results longer than this
const compactionKeepTurns = 2 // Keep last 2 turns' tool results in full (KA preserves key facts)
const compactionMinChars = 300 // Only compact results longer than this
currentTurnStartIndex := len(providerMessages) // Initial messages are never compacted
for turn < a.maxTurns {
for turn < maxTurns {
// === CONTEXT COMPACTION: Compact old tool results to prevent context blowout ===
if turn > 0 {
compactOldToolResults(providerMessages, currentTurnStartIndex, compactionKeepTurns, compactionMinChars)
@@ -168,11 +173,23 @@ func (a *AgenticLoop) executeWithTools(ctx context.Context, sessionID string, me
// 2. It's the first turn
// 3. The user's message indicates they need live data or an action
// This prevents forcing tool calls on conceptual questions like "What is TCP?"
if writeCompletedLastTurn {
forcedTextOnly := false
if turn >= maxTurns-1 {
// Last turn before hitting the limit — force a text-only response so
// the model summarizes its findings instead of silently stopping.
req.ToolChoice = &providers.ToolChoice{Type: providers.ToolChoiceNone}
forcedTextOnly = true
log.Warn().
Int("turn", turn).
Int("max_turns", maxTurns).
Str("session_id", sessionID).
Msg("[AgenticLoop] Approaching max turns — forcing text-only response for summary")
} else if writeCompletedLastTurn {
// A write action completed successfully on the previous turn.
// Force text-only response so the model summarizes the result instead of
// making more tool calls (which often return stale cached data and cause loops).
req.ToolChoice = &providers.ToolChoice{Type: providers.ToolChoiceNone}
forcedTextOnly = true
writeCompletedLastTurn = false
log.Debug().
Str("session_id", sessionID).
@@ -182,6 +199,7 @@ func (a *AgenticLoop) executeWithTools(ctx context.Context, sessionID string, me
// The model already has the data it gathered — force it to produce a text
// response instead of continuing to call tools that will just be blocked again.
req.ToolChoice = &providers.ToolChoice{Type: providers.ToolChoiceNone}
forcedTextOnly = true
toolBlockedLastTurn = false
log.Debug().
Str("session_id", sessionID).
@@ -248,7 +266,11 @@ func (a *AgenticLoop) executeWithTools(ctx context.Context, sessionID string, me
// Format input for frontend display
// For control tools, show a human-readable summary instead of raw JSON to avoid "hallucination" look
inputStr := "{}"
rawInput := ""
if data.Input != nil {
if inputBytes, err := json.Marshal(data.Input); err == nil {
rawInput = string(inputBytes)
}
// Special handling for command execution tools to avoid showing raw JSON
if data.Name == "pulse_control" || data.Name == "pulse_run_command" || data.Name == "control" {
if cmd, ok := data.Input["command"].(string); ok {
@@ -277,9 +299,10 @@ func (a *AgenticLoop) executeWithTools(ctx context.Context, sessionID string, me
}
}
jsonData, _ := json.Marshal(ToolStartData{
ID: data.ID,
Name: data.Name,
Input: inputStr,
ID: data.ID,
Name: data.Name,
Input: inputStr,
RawInput: rawInput,
})
callback(StreamEvent{Type: "tool_start", Data: jsonData})
}
@@ -323,6 +346,18 @@ func (a *AgenticLoop) executeWithTools(ctx context.Context, sessionID string, me
return resultMessages, fmt.Errorf("provider error: %w", err)
}
// Guard: if we forced text-only but the model still returned tool calls
// (some providers like Gemini can hallucinate function calls from conversation
// history even when tools are not offered in the request), strip them so the
// model's text content is treated as the final response.
if forcedTextOnly && len(toolCalls) > 0 {
log.Warn().
Str("session_id", sessionID).
Int("stripped_tool_calls", len(toolCalls)).
Msg("[AgenticLoop] Model returned tool calls despite ToolChoiceNone — stripping them")
toolCalls = nil
}
// Check mid-run budget after each turn completes
if a.budgetChecker != nil {
if budgetErr := a.budgetChecker(); budgetErr != nil {
@@ -475,6 +510,7 @@ func (a *AgenticLoop) executeWithTools(ctx context.Context, sessionID string, me
}
log.Debug().Msg("Agentic loop complete - no tool calls")
resultMessages = a.ensureFinalTextResponse(ctx, sessionID, resultMessages, providerMessages, callback)
return resultMessages, nil
}
@@ -611,6 +647,58 @@ func (a *AgenticLoop) executeWithTools(ctx context.Context, sessionID string, me
continue
}
// === KNOWLEDGE GATE: Return cached facts for redundant tool calls ===
if a.knowledgeAccumulator != nil {
if keys := PredictFactKeys(tc.Name, tc.Input); len(keys) > 0 {
var cachedParts []string
for _, key := range keys {
if value, found := a.knowledgeAccumulator.Lookup(key); found {
cachedParts = append(cachedParts, value)
}
}
if len(cachedParts) > 0 {
cachedResult := fmt.Sprintf("Already known (from earlier investigation, matched keys: %s): %s. If you need fresh data, use a different query or approach.", strings.Join(keys, ","), strings.Join(cachedParts, "; "))
log.Info().
Str("tool", tc.Name).
Str("session_id", sessionID).
Strs("matched_keys", keys).
Int("cached_parts", len(cachedParts)).
Msg("[AgenticLoop] Knowledge gate: returning cached fact instead of re-executing tool")
jsonData, _ := json.Marshal(ToolEndData{
ID: tc.ID,
Name: tc.Name,
Input: "",
Output: cachedResult,
Success: true,
})
callback(StreamEvent{Type: "tool_end", Data: jsonData})
toolResultMsg := Message{
ID: uuid.New().String(),
Role: "user",
Timestamp: time.Now(),
ToolResult: &ToolResult{
ToolUseID: tc.ID,
Content: cachedResult,
IsError: false,
},
}
resultMessages = append(resultMessages, toolResultMsg)
providerMessages = append(providerMessages, providers.Message{
Role: "user",
ToolResult: &providers.ToolResult{
ToolUseID: tc.ID,
Content: cachedResult,
IsError: false,
},
})
continue
}
}
}
// Execute the tool
result, err := a.executor.ExecuteTool(ctx, tc.Name, tc.Input)
@@ -628,6 +716,20 @@ func (a *AgenticLoop) executeWithTools(ctx context.Context, sessionID string, me
log.Debug().
Str("tool", tc.Name).
Msg("[AgenticLoop] Tool succeeded - toolsSucceededThisEpisode set to true")
// Extract and accumulate knowledge facts
if a.knowledgeAccumulator != nil {
a.knowledgeAccumulator.SetTurn(turn)
facts := ExtractFacts(tc.Name, tc.Input, resultText)
for _, f := range facts {
a.knowledgeAccumulator.AddFact(f.Category, f.Key, f.Value)
log.Debug().
Str("tool", tc.Name).
Str("fact_key", f.Key).
Int("value_len", len(f.Value)).
Msg("[AgenticLoop] Stored knowledge fact")
}
}
}
}
@@ -786,7 +888,11 @@ func (a *AgenticLoop) executeWithTools(ctx context.Context, sessionID string, me
// Send tool_end event
// Convert input to JSON string for frontend display
inputStr := ""
rawInput := ""
if tc.Input != nil {
if inputBytes, err := json.Marshal(tc.Input); err == nil {
rawInput = string(inputBytes)
}
// Special handling for command execution tools to avoid showing raw JSON
if tc.Name == "pulse_control" || tc.Name == "pulse_run_command" || tc.Name == "control" {
if cmd, ok := tc.Input["command"].(string); ok {
@@ -814,11 +920,12 @@ func (a *AgenticLoop) executeWithTools(ctx context.Context, sessionID string, me
}
}
jsonData, _ := json.Marshal(ToolEndData{
ID: tc.ID,
Name: tc.Name,
Input: inputStr,
Output: resultText,
Success: !isError,
ID: tc.ID,
Name: tc.Name,
Input: inputStr,
RawInput: rawInput,
Output: resultText,
Success: !isError,
})
callback(StreamEvent{Type: "tool_end", Data: jsonData})
@@ -951,10 +1058,93 @@ func (a *AgenticLoop) executeWithTools(ctx context.Context, sessionID string, me
turn++
}
log.Warn().Int("max_turns", a.maxTurns).Msg("Agentic loop hit max turns limit")
log.Warn().Int("max_turns", maxTurns).Str("session_id", sessionID).Msg("Agentic loop hit max turns limit")
resultMessages = a.ensureFinalTextResponse(ctx, sessionID, resultMessages, providerMessages, callback)
return resultMessages, nil
}
// ensureFinalTextResponse checks if the result messages contain any assistant text.
// If not, it makes one last text-only LLM call to force the model to summarize its findings.
// This prevents the loop from exiting silently after making tool calls without answering.
func (a *AgenticLoop) ensureFinalTextResponse(
ctx context.Context,
sessionID string,
resultMessages []Message,
providerMessages []providers.Message,
callback StreamCallback,
) []Message {
// Check if any assistant message has text content
for i := len(resultMessages) - 1; i >= 0; i-- {
if resultMessages[i].Role == "assistant" && strings.TrimSpace(resultMessages[i].Content) != "" {
return resultMessages // Already has text — nothing to do
}
}
// No text content from the model. Make a final text-only call.
log.Warn().Str("session_id", sessionID).Msg("[AgenticLoop] No text content produced — making final summary call")
// Build clean message history for the summary call:
// 1. Strip any trailing empty assistant messages (the model already failed to produce
// text with these, so including them would just get the same empty result).
// 2. Append a user-role nudge to give the model a clear instruction.
cleanMessages := make([]providers.Message, len(providerMessages))
copy(cleanMessages, providerMessages)
for len(cleanMessages) > 0 {
last := cleanMessages[len(cleanMessages)-1]
if last.Role == "assistant" && strings.TrimSpace(last.Content) == "" && len(last.ToolCalls) == 0 {
cleanMessages = cleanMessages[:len(cleanMessages)-1]
} else {
break
}
}
cleanMessages = append(cleanMessages, providers.Message{
Role: "user",
Content: "You've gathered information above using tools. Now provide your analysis and answer to the original question. Summarize your findings concisely.",
})
summaryReq := providers.ChatRequest{
Messages: cleanMessages,
System: a.getSystemPrompt(),
ToolChoice: &providers.ToolChoice{Type: providers.ToolChoiceNone},
// No Tools field — completely omit tools to prevent hallucinated function calls
}
var summaryBuilder strings.Builder
summaryErr := a.provider.ChatStream(ctx, summaryReq, func(event providers.StreamEvent) {
switch event.Type {
case "content":
if data, ok := event.Data.(providers.ContentEvent); ok {
summaryBuilder.WriteString(data.Text)
jsonData, _ := json.Marshal(ContentData{Text: data.Text})
callback(StreamEvent{Type: "content", Data: jsonData})
}
case "done":
if data, ok := event.Data.(providers.DoneEvent); ok {
a.totalInputTokens += data.InputTokens
a.totalOutputTokens += data.OutputTokens
}
}
})
if summaryErr == nil && summaryBuilder.Len() > 0 {
summaryMsg := Message{
ID: uuid.New().String(),
Role: "assistant",
Content: summaryBuilder.String(),
Timestamp: time.Now(),
}
resultMessages = append(resultMessages, summaryMsg)
log.Info().Str("session_id", sessionID).Int("summary_len", summaryBuilder.Len()).Msg("[AgenticLoop] Final summary produced")
} else if summaryErr != nil {
log.Error().Err(summaryErr).Str("session_id", sessionID).Msg("[AgenticLoop] Final summary call failed")
} else {
log.Warn().Str("session_id", sessionID).Msg("[AgenticLoop] Final summary call returned empty content")
}
return resultMessages
}
// Abort aborts an ongoing session
func (a *AgenticLoop) Abort(sessionID string) {
a.mu.Lock()
@@ -978,6 +1168,14 @@ func (a *AgenticLoop) SetSessionFSM(fsm *SessionFSM) {
a.mu.Unlock()
}
// SetKnowledgeAccumulator sets the knowledge accumulator for fact extraction.
// This must be called before Execute to enable knowledge accumulation.
func (a *AgenticLoop) SetKnowledgeAccumulator(ka *KnowledgeAccumulator) {
a.mu.Lock()
a.knowledgeAccumulator = ka
a.mu.Unlock()
}
// SetMaxTurns overrides the maximum number of agentic turns for this loop.
func (a *AgenticLoop) SetMaxTurns(n int) {
a.mu.Lock()
@@ -1475,7 +1673,14 @@ confirmation prompt - you don't need to ask "Would you like me to...?" Just exec
needed and the system will prompt the user to approve if required.`
}
return a.baseSystemPrompt + modeContext
prompt := a.baseSystemPrompt + modeContext
// Append accumulated knowledge facts to system prompt
if ka := a.knowledgeAccumulator; ka != nil && ka.Len() > 0 {
prompt += "\n\n" + ka.Render()
}
return prompt
}
// AnswerQuestion provides an answer to a pending question
+78
View File
@@ -11,6 +11,7 @@ import (
"github.com/rcourtman/pulse-go-rewrite/internal/ai/approval"
"github.com/rcourtman/pulse-go-rewrite/internal/ai/providers"
"github.com/rcourtman/pulse-go-rewrite/internal/ai/tools"
"github.com/rcourtman/pulse-go-rewrite/internal/models"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
@@ -164,6 +165,83 @@ func TestAgenticLoop(t *testing.T) {
// In a real execution, loop would check aborted map
assert.True(t, loop.aborted["abort-me"])
})
t.Run("KnowledgeAccumulatorInjectedInSystemPrompt", func(t *testing.T) {
state := models.StateSnapshot{
VMs: []models.VM{{
ID: "vm-100",
VMID: 100,
Name: "vm-one",
Node: "node1",
Status: "running",
CPU: 0.5,
CPUs: 4,
Memory: models.Memory{
Usage: 0.42,
Used: 4 * 1024 * 1024 * 1024,
Total: 8 * 1024 * 1024 * 1024,
},
}},
}
execWithState := tools.NewPulseToolExecutor(tools.ExecutorConfig{
StateProvider: &mockStateProvider{state: state},
})
mockProvider := &MockProvider{}
loop := NewAgenticLoop(mockProvider, execWithState, "You are a helper")
ka := NewKnowledgeAccumulator()
loop.SetKnowledgeAccumulator(ka)
ctx := context.Background()
sessionID := "ka-session"
messages := []Message{{Role: "user", Content: "Get VM status"}}
mockProvider.On("ChatStream", mock.Anything, mock.MatchedBy(func(req providers.ChatRequest) bool {
return len(req.Messages) == 1
}), mock.Anything).Return(nil).Run(func(args mock.Arguments) {
callback := args.Get(2).(providers.StreamCallback)
callback(providers.StreamEvent{
Type: "tool_start",
Data: providers.ToolStartEvent{ID: "call_1", Name: "pulse_query"},
})
callback(providers.StreamEvent{
Type: "done",
Data: providers.DoneEvent{
ToolCalls: []providers.ToolCall{{
ID: "call_1",
Name: "pulse_query",
Input: map[string]interface{}{
"action": "get",
"resource_type": "vm",
"resource_id": "100",
},
}},
},
})
}).Once()
var secondSystem string
mockProvider.On("ChatStream", mock.Anything, mock.MatchedBy(func(req providers.ChatRequest) bool {
return len(req.Messages) == 3
}), mock.Anything).Return(nil).Run(func(args mock.Arguments) {
req := args.Get(1).(providers.ChatRequest)
secondSystem = req.System
callback := args.Get(2).(providers.StreamCallback)
callback(providers.StreamEvent{
Type: "content",
Data: providers.ContentEvent{Text: "Done."},
})
callback(providers.StreamEvent{
Type: "done",
Data: providers.DoneEvent{},
})
}).Once()
_, err := loop.Execute(ctx, sessionID, messages, func(event StreamEvent) {})
require.NoError(t, err)
require.NotEmpty(t, secondSystem)
assert.Contains(t, secondSystem, "## Known Facts")
assert.Contains(t, secondSystem, "vm-one")
})
}
func TestAgenticLoop_UpdateTools(t *testing.T) {
+216
View File
@@ -0,0 +1,216 @@
package chat
import (
"fmt"
"strings"
"sync"
"time"
)
// FactCategory classifies a knowledge fact for grouping in the rendered output.
type FactCategory string
const (
FactCategoryResource FactCategory = "resource"
FactCategoryStorage FactCategory = "storage"
FactCategoryDiscovery FactCategory = "discovery"
FactCategoryExec FactCategory = "exec"
FactCategoryMetrics FactCategory = "metrics"
FactCategoryFinding FactCategory = "finding"
)
// Fact is a single extracted knowledge entry.
type Fact struct {
Category FactCategory
Key string // Dedup key, e.g. "lxc:delly:106:status"
Value string // Compact value, e.g. "running, Postfix, hostname=patrol-signal-test"
ObservedAt time.Time
Turn int
}
const (
defaultMaxEntries = 60
defaultMaxChars = 2000
maxValueLen = 200
)
// KnowledgeAccumulator stores extracted facts from tool results.
// Per-session, in-memory, bounded. Facts are keyed for upsert semantics.
// Thread-safe: all methods are protected by a mutex.
type KnowledgeAccumulator struct {
mu sync.Mutex
facts map[string]*Fact // key -> fact (upsert: same key updates value)
order []string // insertion order for LRU eviction
totalChars int
maxEntries int
maxChars int
currentTurn int
}
// NewKnowledgeAccumulator creates a new bounded accumulator.
func NewKnowledgeAccumulator() *KnowledgeAccumulator {
return &KnowledgeAccumulator{
facts: make(map[string]*Fact),
maxEntries: defaultMaxEntries,
maxChars: defaultMaxChars,
}
}
// SetTurn updates the current turn number for new facts.
func (ka *KnowledgeAccumulator) SetTurn(turn int) {
ka.mu.Lock()
ka.currentTurn = turn
ka.mu.Unlock()
}
// AddFact upserts a fact. If the key already exists, the value is updated.
// If over budget after insertion, the oldest non-pinned fact is evicted.
func (ka *KnowledgeAccumulator) AddFact(category FactCategory, key, value string) {
if key == "" || value == "" {
return
}
ka.mu.Lock()
defer ka.mu.Unlock()
// Truncate value
if len(value) > maxValueLen {
value = value[:maxValueLen]
}
now := time.Now()
if existing, ok := ka.facts[key]; ok {
// Upsert: update existing fact
ka.totalChars -= len(existing.Value)
existing.Value = value
existing.Category = category
existing.ObservedAt = now
existing.Turn = ka.currentTurn
ka.totalChars += len(value)
} else {
// New fact
fact := &Fact{
Category: category,
Key: key,
Value: value,
ObservedAt: now,
Turn: ka.currentTurn,
}
ka.facts[key] = fact
ka.order = append(ka.order, key)
ka.totalChars += len(value)
}
// Evict until within budget
ka.evict()
}
// evict removes the oldest non-pinned facts until within budget.
// Facts from the current or previous turn are soft-pinned (not evicted).
// Caller must hold ka.mu.
func (ka *KnowledgeAccumulator) evict() {
for (len(ka.facts) > ka.maxEntries || ka.totalChars > ka.maxChars) && len(ka.order) > 0 {
evicted := false
for i, key := range ka.order {
fact, ok := ka.facts[key]
if !ok {
// Stale key in order slice — remove it
ka.order = append(ka.order[:i], ka.order[i+1:]...)
evicted = true
break
}
// Soft-pin: don't evict facts from current or previous turn
if fact.Turn >= ka.currentTurn-1 {
continue
}
// Evict this fact
ka.totalChars -= len(fact.Value)
delete(ka.facts, key)
ka.order = append(ka.order[:i], ka.order[i+1:]...)
evicted = true
break
}
if !evicted {
// All remaining facts are pinned — can't evict more
break
}
}
}
// Len returns the number of facts stored.
func (ka *KnowledgeAccumulator) Len() int {
ka.mu.Lock()
defer ka.mu.Unlock()
return len(ka.facts)
}
// TotalChars returns the total character count of all fact values.
func (ka *KnowledgeAccumulator) TotalChars() int {
ka.mu.Lock()
defer ka.mu.Unlock()
return ka.totalChars
}
// Lookup returns the fact value for a given key, or empty string if not found.
func (ka *KnowledgeAccumulator) Lookup(key string) (string, bool) {
ka.mu.Lock()
defer ka.mu.Unlock()
if fact, ok := ka.facts[key]; ok {
return fact.Value, true
}
return "", false
}
// Render produces the system prompt section with grouped facts.
func (ka *KnowledgeAccumulator) Render() string {
ka.mu.Lock()
defer ka.mu.Unlock()
if len(ka.facts) == 0 {
return ""
}
// Group facts by category
categoryOrder := []FactCategory{
FactCategoryResource,
FactCategoryStorage,
FactCategoryDiscovery,
FactCategoryExec,
FactCategoryMetrics,
FactCategoryFinding,
}
categoryLabels := map[FactCategory]string{
FactCategoryResource: "Resources",
FactCategoryStorage: "Storage",
FactCategoryDiscovery: "Discovery",
FactCategoryExec: "Exec",
FactCategoryMetrics: "Metrics",
FactCategoryFinding: "Findings",
}
grouped := make(map[FactCategory][]*Fact)
for _, key := range ka.order {
if fact, ok := ka.facts[key]; ok {
grouped[fact.Category] = append(grouped[fact.Category], fact)
}
}
var sb strings.Builder
sb.WriteString("## Known Facts (auto-collected — do NOT re-query unless user asks for fresh data)")
for _, cat := range categoryOrder {
facts, ok := grouped[cat]
if !ok || len(facts) == 0 {
continue
}
label := categoryLabels[cat]
sb.WriteString(fmt.Sprintf("\n%s:", label))
for _, fact := range facts {
sb.WriteString(fmt.Sprintf("\n- %s", fact.Value))
}
}
return sb.String()
}
@@ -0,0 +1,183 @@
package chat
import (
"strings"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestKnowledgeAccumulator_Basic(t *testing.T) {
ka := NewKnowledgeAccumulator()
assert.Equal(t, 0, ka.Len())
assert.Equal(t, "", ka.Render())
ka.AddFact(FactCategoryResource, "lxc:delly:106:status", "running, Postfix")
assert.Equal(t, 1, ka.Len())
rendered := ka.Render()
assert.Contains(t, rendered, "Known Facts")
assert.Contains(t, rendered, "Resources:")
assert.Contains(t, rendered, "running, Postfix")
}
func TestKnowledgeAccumulator_Upsert(t *testing.T) {
ka := NewKnowledgeAccumulator()
ka.AddFact(FactCategoryResource, "lxc:delly:106:status", "running, CPU=5%")
assert.Equal(t, 1, ka.Len())
// Upsert with same key should update value, not add a new entry
ka.AddFact(FactCategoryResource, "lxc:delly:106:status", "stopped")
assert.Equal(t, 1, ka.Len())
rendered := ka.Render()
assert.Contains(t, rendered, "stopped")
assert.NotContains(t, rendered, "CPU=5%")
}
func TestKnowledgeAccumulator_EmptyKeyOrValue(t *testing.T) {
ka := NewKnowledgeAccumulator()
ka.AddFact(FactCategoryResource, "", "some value")
assert.Equal(t, 0, ka.Len())
ka.AddFact(FactCategoryResource, "some-key", "")
assert.Equal(t, 0, ka.Len())
}
func TestKnowledgeAccumulator_ValueTruncation(t *testing.T) {
ka := NewKnowledgeAccumulator()
longValue := strings.Repeat("x", 300)
ka.AddFact(FactCategoryExec, "exec:host:cmd", longValue)
fact := ka.facts["exec:host:cmd"]
require.NotNil(t, fact)
assert.Equal(t, maxValueLen, len(fact.Value))
}
func TestKnowledgeAccumulator_MaxEntries(t *testing.T) {
ka := NewKnowledgeAccumulator()
ka.maxEntries = 5
ka.maxChars = 100000 // High limit so entries are the constraint
// Insert 5 facts at turn 0
ka.SetTurn(0)
for i := 0; i < 5; i++ {
ka.AddFact(FactCategoryResource, keyForIndex(i), valForIndex(i))
}
assert.Equal(t, 5, ka.Len())
// Move to turn 2 so turn 0 facts are evictable
ka.SetTurn(2)
ka.AddFact(FactCategoryResource, "new-key", "new-value")
// Should have evicted one old fact to make room
assert.Equal(t, 5, ka.Len())
_, hasNew := ka.facts["new-key"]
assert.True(t, hasNew, "new fact should be present")
}
func TestKnowledgeAccumulator_MaxChars(t *testing.T) {
ka := NewKnowledgeAccumulator()
ka.maxEntries = 1000
ka.maxChars = 50 // Very low char budget
ka.SetTurn(0)
ka.AddFact(FactCategoryResource, "k1", strings.Repeat("a", 20))
ka.AddFact(FactCategoryResource, "k2", strings.Repeat("b", 20))
// Move forward so old facts can be evicted
ka.SetTurn(2)
ka.AddFact(FactCategoryResource, "k3", strings.Repeat("c", 20))
// Should have evicted oldest to stay within budget
assert.LessOrEqual(t, ka.TotalChars(), 50+maxValueLen, "should stay near char budget")
}
func TestKnowledgeAccumulator_SoftPinCurrentTurn(t *testing.T) {
ka := NewKnowledgeAccumulator()
ka.maxEntries = 2
ka.maxChars = 100000
// All facts in current turn should be soft-pinned
ka.SetTurn(5)
ka.AddFact(FactCategoryResource, "k1", "val1")
ka.AddFact(FactCategoryResource, "k2", "val2")
ka.AddFact(FactCategoryResource, "k3", "val3")
// All 3 should survive even though maxEntries=2 because they're all from current turn
assert.Equal(t, 3, ka.Len())
}
func TestKnowledgeAccumulator_CategoryGrouping(t *testing.T) {
ka := NewKnowledgeAccumulator()
ka.AddFact(FactCategoryResource, "res1", "LXC 106 running")
ka.AddFact(FactCategoryStorage, "stor1", "PBS available, 42% used")
ka.AddFact(FactCategoryExec, "exec1", "exit=0, service active")
ka.AddFact(FactCategoryFinding, "find1", "warning: high CPU on vm101")
rendered := ka.Render()
// Verify category headers appear
assert.Contains(t, rendered, "Resources:")
assert.Contains(t, rendered, "Storage:")
assert.Contains(t, rendered, "Exec:")
assert.Contains(t, rendered, "Findings:")
// Verify values appear
assert.Contains(t, rendered, "LXC 106 running")
assert.Contains(t, rendered, "PBS available, 42% used")
assert.Contains(t, rendered, "exit=0, service active")
assert.Contains(t, rendered, "warning: high CPU on vm101")
}
func TestKnowledgeAccumulator_RenderOrder(t *testing.T) {
ka := NewKnowledgeAccumulator()
// Add in non-category order
ka.AddFact(FactCategoryFinding, "f1", "finding")
ka.AddFact(FactCategoryResource, "r1", "resource")
rendered := ka.Render()
// Resources should appear before Findings in the rendered output
resIdx := strings.Index(rendered, "Resources:")
findIdx := strings.Index(rendered, "Findings:")
assert.Greater(t, findIdx, resIdx, "Resources should come before Findings")
}
func TestKnowledgeAccumulator_Lookup(t *testing.T) {
ka := NewKnowledgeAccumulator()
ka.AddFact(FactCategoryDiscovery, "discovery:delly:106", "service=Postfix, hostname=patrol-signal-test")
val, found := ka.Lookup("discovery:delly:106")
assert.True(t, found)
assert.Equal(t, "service=Postfix, hostname=patrol-signal-test", val)
val, found = ka.Lookup("discovery:minipc:200")
assert.False(t, found)
assert.Equal(t, "", val)
}
func TestKnowledgeAccumulator_SetTurn(t *testing.T) {
ka := NewKnowledgeAccumulator()
ka.SetTurn(3)
ka.AddFact(FactCategoryResource, "k1", "val1")
fact := ka.facts["k1"]
require.NotNil(t, fact)
assert.Equal(t, 3, fact.Turn)
}
func keyForIndex(i int) string {
return strings.Repeat("k", i+1)
}
func valForIndex(i int) string {
return strings.Repeat("v", i+1)
}
+582
View File
@@ -0,0 +1,582 @@
package chat
import (
"encoding/json"
"fmt"
"strings"
)
// FactEntry is the output of ExtractFacts — ready to feed into KnowledgeAccumulator.AddFact.
type FactEntry struct {
Category FactCategory
Key string
Value string
}
// ExtractFacts deterministically extracts knowledge facts from a tool result.
// No LLM calls. Parses the JSON text from FormatToolResult() output.
// Returns empty slice on parse errors or unrecognized tools — never panics.
//
// Tool results use NewJSONResult (direct struct marshaling), NOT ToolResponse wrapper.
// So the resultText is the JSON of the response struct directly (e.g. ResourceResponse).
func ExtractFacts(toolName string, toolInput map[string]interface{}, resultText string) []FactEntry {
switch toolName {
case "pulse_query":
return extractQueryFacts(toolInput, resultText)
case "pulse_storage":
return extractStorageFacts(toolInput, resultText)
case "pulse_discovery":
return extractDiscoveryFacts(toolInput, resultText)
case "pulse_read", "pulse_run_command":
return extractExecFacts(toolInput, resultText)
case "pulse_metrics":
return extractMetricsFacts(toolInput, resultText)
case "patrol_report_finding":
return extractFindingFacts(toolInput, resultText)
default:
return nil
}
}
// --- pulse_query ---
func extractQueryFacts(input map[string]interface{}, resultText string) []FactEntry {
action := strFromMap(input, "action")
if action == "" {
action = strFromMap(input, "type")
}
switch action {
case "get":
return extractQueryGetFacts(input, resultText)
case "search":
return extractQuerySearchFacts(input, resultText)
default:
return nil
}
}
func extractQueryGetFacts(input map[string]interface{}, resultText string) []FactEntry {
// Tool results are direct JSON (NewJSONResult), no ToolResponse wrapper.
// ResourceResponse has nested CPU/Memory structs.
var resource struct {
Type string `json:"type"`
Name string `json:"name"`
Status string `json:"status"`
Node string `json:"node"`
ID string `json:"id"`
VMID int `json:"vmid"`
Host string `json:"host"`
CPU struct {
Percent float64 `json:"percent"`
} `json:"cpu"`
Memory struct {
Percent float64 `json:"percent"`
} `json:"memory"`
// Error field for not-found responses
Error string `json:"error"`
}
if err := json.Unmarshal([]byte(resultText), &resource); err != nil {
return nil
}
// Skip error/not-found responses
if resource.Error != "" {
return nil
}
if resource.Name == "" && resource.ID == "" {
return nil
}
resType := resource.Type
if resType == "" {
resType = strFromMap(input, "resource_type")
}
node := resource.Node
if node == "" {
node = resource.Host
}
id := resource.ID
if id == "" && resource.VMID > 0 {
id = fmt.Sprintf("%d", resource.VMID)
}
if id == "" {
id = resource.Name
}
key := fmt.Sprintf("%s:%s:%s:status", resType, node, id)
var parts []string
if resource.Status != "" {
parts = append(parts, resource.Status)
}
if resource.Name != "" {
parts = append(parts, resource.Name)
}
if resource.CPU.Percent > 0 {
parts = append(parts, fmt.Sprintf("CPU=%.1f%%", resource.CPU.Percent))
}
if resource.Memory.Percent > 0 {
parts = append(parts, fmt.Sprintf("Mem=%.1f%%", resource.Memory.Percent))
}
if len(parts) == 0 {
return nil
}
return []FactEntry{{
Category: FactCategoryResource,
Key: key,
Value: strings.Join(parts, ", "),
}}
}
func extractQuerySearchFacts(input map[string]interface{}, resultText string) []FactEntry {
query := strFromMap(input, "query")
if query == "" {
query = strFromMap(input, "search")
}
// ResourceSearchResponse — direct JSON, no wrapper
var resp struct {
Matches []struct {
Name string `json:"name"`
Status string `json:"status"`
Type string `json:"type"`
} `json:"matches"`
Total int `json:"total"`
}
if err := json.Unmarshal([]byte(resultText), &resp); err != nil {
return nil
}
total := resp.Total
if total == 0 && len(resp.Matches) == 0 {
return nil
}
// Summarize first 5 matches
var summaryParts []string
limit := 5
if len(resp.Matches) < limit {
limit = len(resp.Matches)
}
for _, m := range resp.Matches[:limit] {
entry := m.Name
if m.Status != "" {
entry += " (" + m.Status + ")"
}
summaryParts = append(summaryParts, entry)
}
value := fmt.Sprintf("%d results: %s", total, strings.Join(summaryParts, ", "))
return []FactEntry{{
Category: FactCategoryResource,
Key: fmt.Sprintf("search:%s:summary", query),
Value: truncateValue(value),
}}
}
// --- pulse_storage ---
func extractStorageFacts(input map[string]interface{}, resultText string) []FactEntry {
action := strFromMap(input, "action")
if action == "" {
action = strFromMap(input, "type")
}
switch action {
case "pools":
return extractStoragePoolFacts(resultText)
case "backup_tasks":
return extractBackupTaskFacts(resultText)
default:
return nil
}
}
func extractStoragePoolFacts(resultText string) []FactEntry {
// StorageResponse — direct JSON, no wrapper
var resp struct {
Pools []struct {
Name string `json:"name"`
Node string `json:"node"`
Nodes []string `json:"nodes"`
Type string `json:"type"`
Status string `json:"status"`
Active bool `json:"active"`
UsagePercent float64 `json:"usage_percent"`
TotalGB float64 `json:"total_gb"`
UsedGB float64 `json:"used_gb"`
} `json:"pools"`
}
if err := json.Unmarshal([]byte(resultText), &resp); err != nil {
return nil
}
var facts []FactEntry
for _, pool := range resp.Pools {
node := pool.Node
if node == "" && len(pool.Nodes) > 0 {
node = strings.Join(pool.Nodes, "+")
}
freeGB := pool.TotalGB - pool.UsedGB
var parts []string
if pool.Type != "" {
parts = append(parts, pool.Type)
}
if pool.Status != "" {
parts = append(parts, pool.Status)
}
if pool.Active {
parts = append(parts, fmt.Sprintf("active on %s", node))
}
parts = append(parts, fmt.Sprintf("%.1f%% used", pool.UsagePercent))
if freeGB > 0 {
parts = append(parts, fmt.Sprintf("%.0fGB free", freeGB))
}
facts = append(facts, FactEntry{
Category: FactCategoryStorage,
Key: fmt.Sprintf("storage:%s:%s", node, pool.Name),
Value: truncateValue(strings.Join(parts, ", ")),
})
}
return facts
}
func extractBackupTaskFacts(resultText string) []FactEntry {
// BackupTasksListResponse — direct JSON, no wrapper
var resp struct {
Tasks []struct {
VMID string `json:"vmid"`
Node string `json:"node"`
Status string `json:"status"`
StartTime string `json:"start_time"`
Error string `json:"error"`
} `json:"tasks"`
}
if err := json.Unmarshal([]byte(resultText), &resp); err != nil {
return nil
}
var facts []FactEntry
for _, task := range resp.Tasks {
// Only record failures
if task.Status == "ok" || task.Status == "success" || task.Status == "" {
continue
}
var parts []string
parts = append(parts, task.Status)
if task.StartTime != "" {
parts = append(parts, "at "+task.StartTime)
}
if task.Error != "" {
parts = append(parts, "error="+task.Error)
}
facts = append(facts, FactEntry{
Category: FactCategoryStorage,
Key: fmt.Sprintf("backup:%s:%s", task.VMID, task.Node),
Value: truncateValue(strings.Join(parts, ", ")),
})
}
return facts
}
// --- pulse_discovery ---
func extractDiscoveryFacts(input map[string]interface{}, resultText string) []FactEntry {
// ResourceDiscoveryInfo — direct JSON, no wrapper
var disc struct {
ServiceType string `json:"service_type"`
Hostname string `json:"hostname"`
HostID string `json:"host_id"`
ResourceID string `json:"resource_id"`
Ports []struct {
Port int `json:"port"`
} `json:"ports"`
}
if err := json.Unmarshal([]byte(resultText), &disc); err != nil {
return nil
}
host := disc.HostID
if host == "" {
host = strFromMap(input, "host")
}
id := disc.ResourceID
if id == "" {
id = strFromMap(input, "resource_id")
}
var parts []string
if disc.ServiceType != "" {
parts = append(parts, "service="+disc.ServiceType)
}
if disc.Hostname != "" {
parts = append(parts, "hostname="+disc.Hostname)
}
if len(disc.Ports) > 0 {
var portStrs []string
for _, p := range disc.Ports {
portStrs = append(portStrs, fmt.Sprintf("%d", p.Port))
}
parts = append(parts, "ports=["+strings.Join(portStrs, ",")+"]")
}
if len(parts) == 0 {
return nil
}
return []FactEntry{{
Category: FactCategoryDiscovery,
Key: fmt.Sprintf("discovery:%s:%s", host, id),
Value: truncateValue(strings.Join(parts, ", ")),
}}
}
// --- pulse_read / pulse_run_command ---
func extractExecFacts(input map[string]interface{}, resultText string) []FactEntry {
host := strFromMap(input, "target_host")
if host == "" {
host = strFromMap(input, "host")
}
cmd := strFromMap(input, "command")
if cmd == "" {
// For pulse_read file/tail/find actions, use action+path to distinguish
// different file reads on the same host.
action := strFromMap(input, "action")
path := strFromMap(input, "path")
if action != "" && path != "" {
cmd = action + ":" + path
} else if action != "" {
cmd = action
}
}
if cmd == "" {
return nil
}
// Use first 60 chars of command as key prefix (longer to accommodate path)
cmdPrefix := cmd
if len(cmdPrefix) > 60 {
cmdPrefix = cmdPrefix[:60]
}
// Try to parse as CommandResponse (direct JSON, no wrapper)
var cmdResp struct {
Success bool `json:"success"`
ExitCode int `json:"exit_code"`
Output string `json:"output"`
Stdout string `json:"stdout"`
Error string `json:"error"`
}
var value string
if err := json.Unmarshal([]byte(resultText), &cmdResp); err == nil && (cmdResp.Output != "" || cmdResp.Stdout != "" || cmdResp.Error != "") {
output := cmdResp.Output
if output == "" {
output = cmdResp.Stdout
}
if output == "" {
output = cmdResp.Error
}
// Take first 2 lines
lines := strings.SplitN(output, "\n", 3)
summary := strings.Join(lines[:min(2, len(lines))], "; ")
value = fmt.Sprintf("exit=%d, %s", cmdResp.ExitCode, summary)
} else {
// Fallback: use first 2 lines of raw result text
lines := strings.SplitN(resultText, "\n", 3)
summary := strings.Join(lines[:min(2, len(lines))], "; ")
value = summary
}
return []FactEntry{{
Category: FactCategoryExec,
Key: fmt.Sprintf("exec:%s:%s", host, cmdPrefix),
Value: truncateValue(value),
}}
}
// --- pulse_metrics ---
func extractMetricsFacts(input map[string]interface{}, resultText string) []FactEntry {
action := strFromMap(input, "action")
if action == "" {
action = strFromMap(input, "type")
}
if action != "performance" {
return nil
}
resourceID := strFromMap(input, "resource_id")
if resourceID == "" {
return nil
}
// MetricsResponse — direct JSON, no wrapper.
// Summary is map[string]ResourceMetricsSummary keyed by resource ID.
var resp struct {
Summary map[string]struct {
AvgCPU float64 `json:"avg_cpu"`
MaxCPU float64 `json:"max_cpu"`
AvgMemory float64 `json:"avg_memory"`
MaxMemory float64 `json:"max_memory"`
Trend string `json:"trend"`
} `json:"summary"`
}
if err := json.Unmarshal([]byte(resultText), &resp); err != nil {
return nil
}
// Look up the summary for this resource (or take the first entry)
var avgCPU, maxCPU float64
var trend string
if s, ok := resp.Summary[resourceID]; ok {
avgCPU = s.AvgCPU
maxCPU = s.MaxCPU
trend = s.Trend
} else {
// Take first entry if resource ID doesn't match exactly
for _, s := range resp.Summary {
avgCPU = s.AvgCPU
maxCPU = s.MaxCPU
trend = s.Trend
break
}
}
var parts []string
if avgCPU > 0 {
parts = append(parts, fmt.Sprintf("avg_cpu=%.1f%%", avgCPU))
}
if maxCPU > 0 {
parts = append(parts, fmt.Sprintf("max=%.1f%%", maxCPU))
}
if trend != "" {
parts = append(parts, "trend="+trend)
}
if len(parts) == 0 {
return nil
}
return []FactEntry{{
Category: FactCategoryMetrics,
Key: fmt.Sprintf("metrics:%s", resourceID),
Value: truncateValue(strings.Join(parts, ", ")),
}}
}
// --- patrol_report_finding ---
func extractFindingFacts(input map[string]interface{}, resultText string) []FactEntry {
key := strFromMap(input, "key")
if key == "" {
key = strFromMap(input, "finding_key")
}
severity := strFromMap(input, "severity")
title := strFromMap(input, "title")
resourceID := strFromMap(input, "resource_id")
if key == "" || title == "" {
return nil
}
var parts []string
if severity != "" {
parts = append(parts, severity)
}
parts = append(parts, title)
if resourceID != "" {
parts = append(parts, "on "+resourceID)
}
return []FactEntry{{
Category: FactCategoryFinding,
Key: fmt.Sprintf("finding:%s", key),
Value: truncateValue(strings.Join(parts, ": ")),
}}
}
// PredictFactKeys returns the KA fact keys that this tool call would produce,
// based solely on the tool input (without needing the result).
// Used by the gate to check if we already have facts for this call.
// Returns nil if the key can't be predicted from input alone.
func PredictFactKeys(toolName string, toolInput map[string]interface{}) []string {
switch toolName {
case "pulse_discovery":
host := strFromMap(toolInput, "host_id")
if host == "" {
host = strFromMap(toolInput, "host")
}
id := strFromMap(toolInput, "resource_id")
if host != "" && id != "" {
return []string{fmt.Sprintf("discovery:%s:%s", host, id)}
}
case "pulse_read", "pulse_run_command":
host := strFromMap(toolInput, "target_host")
if host == "" {
host = strFromMap(toolInput, "host")
}
cmd := strFromMap(toolInput, "command")
if cmd == "" {
// For file/tail/find actions, include path to distinguish different file reads
action := strFromMap(toolInput, "action")
path := strFromMap(toolInput, "path")
if action != "" && path != "" {
cmd = action + ":" + path
} else if action != "" {
cmd = action
}
}
if host != "" && cmd != "" {
cmdPrefix := cmd
if len(cmdPrefix) > 60 {
cmdPrefix = cmdPrefix[:60]
}
return []string{fmt.Sprintf("exec:%s:%s", host, cmdPrefix)}
}
case "pulse_metrics":
action := strFromMap(toolInput, "action")
if action == "" {
action = strFromMap(toolInput, "type")
}
resourceID := strFromMap(toolInput, "resource_id")
if action == "performance" && resourceID != "" {
return []string{fmt.Sprintf("metrics:%s", resourceID)}
}
}
// pulse_query get/search and pulse_storage: keys depend on result data,
// can't predict from input alone. Return nil — these calls won't be gated.
return nil
}
// --- helpers ---
func strFromMap(m map[string]interface{}, key string) string {
if m == nil {
return ""
}
v, ok := m[key]
if !ok {
return ""
}
s, ok := v.(string)
if !ok {
return ""
}
return s
}
func truncateValue(s string) string {
if len(s) > maxValueLen {
return s[:maxValueLen]
}
return s
}
@@ -0,0 +1,279 @@
package chat
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestExtractFacts_UnknownTool(t *testing.T) {
facts := ExtractFacts("unknown_tool", nil, `{}`)
assert.Empty(t, facts)
}
func TestExtractFacts_InvalidJSON(t *testing.T) {
facts := ExtractFacts("pulse_query", map[string]interface{}{"action": "get"}, "not json")
assert.Empty(t, facts)
}
func TestExtractFacts_QueryGet(t *testing.T) {
input := map[string]interface{}{"action": "get", "resource_type": "lxc"}
// Actual format from NewJSONResult(ResourceResponse): direct JSON, no wrapper.
// CPU/Memory are nested structs.
result := `{"type":"lxc","name":"postfix-server","status":"running","node":"delly","id":"lxc/106","vmid":106,"cpu":{"percent":2.5,"cores":4},"memory":{"percent":45.0,"used_gb":1.2,"total_gb":4.0}}`
facts := ExtractFacts("pulse_query", input, result)
require.Len(t, facts, 1)
f := facts[0]
assert.Equal(t, FactCategoryResource, f.Category)
assert.Equal(t, "lxc:delly:lxc/106:status", f.Key)
assert.Contains(t, f.Value, "running")
assert.Contains(t, f.Value, "postfix-server")
assert.Contains(t, f.Value, "CPU=2.5%")
assert.Contains(t, f.Value, "Mem=45.0%")
}
func TestExtractFacts_QueryGet_NotFound(t *testing.T) {
input := map[string]interface{}{"action": "get", "resource_type": "vm"}
result := `{"error":"not_found","resource_id":"999","type":"vm"}`
facts := ExtractFacts("pulse_query", input, result)
assert.Empty(t, facts)
}
func TestExtractFacts_QueryGet_NoCPU(t *testing.T) {
// Some resources may not have CPU data populated
input := map[string]interface{}{"action": "get", "resource_type": "container"}
result := `{"type":"container","name":"test-ct","status":"stopped","node":"minipc","id":"lxc/200","cpu":{"percent":0},"memory":{"percent":0}}`
facts := ExtractFacts("pulse_query", input, result)
require.Len(t, facts, 1)
assert.Contains(t, facts[0].Value, "stopped")
assert.Contains(t, facts[0].Value, "test-ct")
assert.NotContains(t, facts[0].Value, "CPU=")
}
func TestExtractFacts_QuerySearch(t *testing.T) {
input := map[string]interface{}{"action": "search", "query": "postfix"}
result := `{"query":"postfix","matches":[{"name":"postfix-lxc","status":"running","type":"lxc"},{"name":"mail-server","status":"stopped","type":"vm"}],"total":2}`
facts := ExtractFacts("pulse_query", input, result)
require.Len(t, facts, 1)
f := facts[0]
assert.Equal(t, FactCategoryResource, f.Category)
assert.Equal(t, "search:postfix:summary", f.Key)
assert.Contains(t, f.Value, "2 results")
assert.Contains(t, f.Value, "postfix-lxc (running)")
assert.Contains(t, f.Value, "mail-server (stopped)")
}
func TestExtractFacts_QuerySearch_Empty(t *testing.T) {
input := map[string]interface{}{"action": "search", "query": "nonexistent"}
result := `{"query":"nonexistent","matches":[],"total":0}`
facts := ExtractFacts("pulse_query", input, result)
assert.Empty(t, facts)
}
func TestExtractFacts_StoragePools(t *testing.T) {
input := map[string]interface{}{"action": "pools"}
result := `{"pools":[{"name":"pbs-minipc","node":"","nodes":["delly","minipc"],"type":"PBS","status":"available","active":true,"usage_percent":42.7,"total_gb":1000,"used_gb":427}]}`
facts := ExtractFacts("pulse_storage", input, result)
require.Len(t, facts, 1)
f := facts[0]
assert.Equal(t, FactCategoryStorage, f.Category)
assert.Contains(t, f.Key, "storage:")
assert.Contains(t, f.Key, "pbs-minipc")
assert.Contains(t, f.Value, "PBS")
assert.Contains(t, f.Value, "42.7% used")
assert.Contains(t, f.Value, "573GB free")
}
func TestExtractFacts_BackupTasks_OnlyFailures(t *testing.T) {
input := map[string]interface{}{"action": "backup_tasks"}
result := `{"tasks":[{"vmid":"106","node":"delly","status":"ok"},{"vmid":"200","node":"minipc","status":"failed","start_time":"2024-01-15T03:00","error":"snapshot failed"}]}`
facts := ExtractFacts("pulse_storage", input, result)
require.Len(t, facts, 1)
f := facts[0]
assert.Equal(t, FactCategoryStorage, f.Category)
assert.Equal(t, "backup:200:minipc", f.Key)
assert.Contains(t, f.Value, "failed")
assert.Contains(t, f.Value, "snapshot failed")
}
func TestExtractFacts_Discovery(t *testing.T) {
input := map[string]interface{}{"host": "delly", "resource_id": "106"}
result := `{"service_type":"Postfix","hostname":"patrol-signal-test","host_id":"delly","resource_id":"106","ports":[{"port":25},{"port":22}]}`
facts := ExtractFacts("pulse_discovery", input, result)
require.Len(t, facts, 1)
f := facts[0]
assert.Equal(t, FactCategoryDiscovery, f.Category)
assert.Equal(t, "discovery:delly:106", f.Key)
assert.Contains(t, f.Value, "service=Postfix")
assert.Contains(t, f.Value, "hostname=patrol-signal-test")
assert.Contains(t, f.Value, "ports=[25,22]")
}
func TestExtractFacts_Exec_JSON(t *testing.T) {
input := map[string]interface{}{"command": "pvesm status | grep pbs-minipc", "target_host": "delly"}
result := `{"success":true,"exit_code":0,"output":"pbs-minipc active 42.68%"}`
facts := ExtractFacts("pulse_read", input, result)
require.Len(t, facts, 1)
f := facts[0]
assert.Equal(t, FactCategoryExec, f.Category)
assert.Contains(t, f.Key, "exec:delly:")
assert.Contains(t, f.Value, "exit=0")
assert.Contains(t, f.Value, "pbs-minipc")
}
func TestExtractFacts_Exec_FallbackRaw(t *testing.T) {
input := map[string]interface{}{"command": "some-cmd", "target_host": "host1"}
result := `not json at all`
facts := ExtractFacts("pulse_read", input, result)
require.Len(t, facts, 1)
assert.Contains(t, facts[0].Value, "not json at all")
}
func TestExtractFacts_Metrics(t *testing.T) {
input := map[string]interface{}{"action": "performance", "resource_id": "vm101"}
// Actual format: summary is map[string]ResourceMetricsSummary keyed by resource ID
result := `{"resource_id":"vm101","period":"7d","summary":{"vm101":{"resource_id":"vm101","avg_cpu":12.3,"max_cpu":78.5,"avg_memory":65.0,"max_memory":89.0,"trend":"growing"}}}`
facts := ExtractFacts("pulse_metrics", input, result)
require.Len(t, facts, 1)
f := facts[0]
assert.Equal(t, FactCategoryMetrics, f.Category)
assert.Equal(t, "metrics:vm101", f.Key)
assert.Contains(t, f.Value, "avg_cpu=12.3%")
assert.Contains(t, f.Value, "max=78.5%")
assert.Contains(t, f.Value, "trend=growing")
}
func TestExtractFacts_Metrics_EmptySummary(t *testing.T) {
input := map[string]interface{}{"action": "performance", "resource_id": "vm101"}
result := `{"resource_id":"vm101","period":"7d","summary":{}}`
facts := ExtractFacts("pulse_metrics", input, result)
assert.Empty(t, facts)
}
func TestExtractFacts_Finding(t *testing.T) {
input := map[string]interface{}{
"key": "high-cpu-vm101",
"severity": "warning",
"title": "High CPU usage on vm101",
"resource_id": "vm101",
}
result := `{"id":"abc123","status":"created"}`
facts := ExtractFacts("patrol_report_finding", input, result)
require.Len(t, facts, 1)
f := facts[0]
assert.Equal(t, FactCategoryFinding, f.Category)
assert.Equal(t, "finding:high-cpu-vm101", f.Key)
assert.Contains(t, f.Value, "warning")
assert.Contains(t, f.Value, "High CPU usage on vm101")
assert.Contains(t, f.Value, "on vm101")
}
func TestExtractFacts_Finding_MissingFields(t *testing.T) {
// Missing key and title should return empty
input := map[string]interface{}{"severity": "warning"}
result := `{"id":"abc123"}`
facts := ExtractFacts("patrol_report_finding", input, result)
assert.Empty(t, facts)
}
func TestPredictFactKeys_Discovery(t *testing.T) {
keys := PredictFactKeys("pulse_discovery", map[string]interface{}{
"host_id": "delly",
"resource_id": "106",
})
require.Len(t, keys, 1)
assert.Equal(t, "discovery:delly:106", keys[0])
}
func TestPredictFactKeys_DiscoveryAltFields(t *testing.T) {
keys := PredictFactKeys("pulse_discovery", map[string]interface{}{
"host": "minipc",
"resource_id": "pbs-minipc",
})
require.Len(t, keys, 1)
assert.Equal(t, "discovery:minipc:pbs-minipc", keys[0])
}
func TestPredictFactKeys_Exec(t *testing.T) {
keys := PredictFactKeys("pulse_read", map[string]interface{}{
"target_host": "delly",
"command": "pvesm status",
})
require.Len(t, keys, 1)
assert.Equal(t, "exec:delly:pvesm status", keys[0])
}
func TestPredictFactKeys_Metrics(t *testing.T) {
keys := PredictFactKeys("pulse_metrics", map[string]interface{}{
"action": "performance",
"resource_id": "vm101",
})
require.Len(t, keys, 1)
assert.Equal(t, "metrics:vm101", keys[0])
}
func TestPredictFactKeys_UnpredictableTools(t *testing.T) {
// pulse_query and pulse_storage keys depend on result data
assert.Nil(t, PredictFactKeys("pulse_query", map[string]interface{}{"action": "get"}))
assert.Nil(t, PredictFactKeys("pulse_storage", map[string]interface{}{"action": "pools"}))
assert.Nil(t, PredictFactKeys("unknown_tool", nil))
}
func TestPredictFactKeys_MissingFields(t *testing.T) {
// Discovery without host_id should return nil
assert.Nil(t, PredictFactKeys("pulse_discovery", map[string]interface{}{"resource_id": "106"}))
// Exec without command should return nil
assert.Nil(t, PredictFactKeys("pulse_read", map[string]interface{}{"target_host": "delly"}))
}
func TestPredictFactKeys_FileReadDistinctPaths(t *testing.T) {
// Different file paths on the same host should produce different keys
keys1 := PredictFactKeys("pulse_read", map[string]interface{}{
"target_host": "delly",
"action": "file",
"path": "/etc/pve/storage.cfg",
})
keys2 := PredictFactKeys("pulse_read", map[string]interface{}{
"target_host": "delly",
"action": "file",
"path": "/var/log/pve/tasks/some-task-log",
})
require.Len(t, keys1, 1)
require.Len(t, keys2, 1)
assert.NotEqual(t, keys1[0], keys2[0], "different file paths must produce different keys")
assert.Contains(t, keys1[0], "storage.cfg")
assert.Contains(t, keys2[0], "tasks/some-task-log")
}
func TestExtractFacts_ValueTruncation(t *testing.T) {
input := map[string]interface{}{"command": "long-output-cmd", "target_host": "host1"}
// Create a result with very long output
longOutput := `{"success":true,"exit_code":0,"output":"` + bigContent(500) + `"}`
facts := ExtractFacts("pulse_read", input, longOutput)
require.Len(t, facts, 1)
assert.LessOrEqual(t, len(facts[0].Value), maxValueLen)
}
+19
View File
@@ -440,6 +440,11 @@ func (s *Service) ExecuteStream(ctx context.Context, req ExecuteRequest, callbac
sessionFSM := sessions.GetSessionFSM(session.ID)
loop.SetSessionFSM(sessionFSM)
// Set session-scoped knowledge accumulator for fact extraction across turns.
// For user chat, this persists across messages so facts accumulate during a conversation.
ka := sessions.GetKnowledgeAccumulator(session.ID)
loop.SetKnowledgeAccumulator(ka)
// If the prefetcher resolved mentions, advance FSM past RESOLVING.
// The prefetched context already contains the resource details (type, VMID, node, host)
// so forcing the AI to redundantly call a read tool would be wasteful.
@@ -469,6 +474,13 @@ func (s *Service) ExecuteStream(ctx context.Context, req ExecuteRequest, callbac
}
}
// Override max turns if the caller specified one (e.g. investigations use 15).
// Reset after the call to avoid affecting concurrent sessions on the shared loop.
if req.MaxTurns > 0 {
loop.SetMaxTurns(req.MaxTurns)
defer loop.SetMaxTurns(MaxAgenticTurns)
}
resultMessages, err := loop.ExecuteWithTools(ctx, session.ID, messages, filteredTools, callback)
log.Debug().
@@ -599,6 +611,13 @@ func (s *Service) ExecutePatrolStream(ctx context.Context, req PatrolRequest, ca
sessionFSM := sessions.GetSessionFSM(session.ID)
tempLoop.SetSessionFSM(sessionFSM)
// Create a fresh knowledge accumulator for this patrol run.
// Unlike user chat (which reuses session-scoped KA across messages),
// patrol runs need a clean slate to avoid stale facts from prior runs
// (patrol-main reuses the same session ID across scheduled runs).
ka := sessions.NewKnowledgeAccumulatorForRun(session.ID)
tempLoop.SetKnowledgeAccumulator(ka)
// Set mid-run budget checker if configured
if s.budgetChecker != nil {
tempLoop.SetBudgetChecker(s.budgetChecker)
+19 -4
View File
@@ -400,9 +400,8 @@ func TestService_SettersAndUpdateControlSettings(t *testing.T) {
assert.True(t, hasTool(service.executor.ListTools(), "pulse_control"))
}
func TestService_FilterToolsForPrompt_AllToolsReturned(t *testing.T) {
// All tools should always be returned regardless of prompt content.
// The approval mechanism handles control, not tool filtering.
func TestService_FilterToolsForPrompt_ReadOnlyFiltersWriteTools(t *testing.T) {
// Read-only prompts should not include write/control tools.
service := NewService(Config{
AIConfig: &config.AIConfig{ControlLevel: config.ControlLevelControlled},
StateProvider: &mockStateProvider{},
@@ -414,8 +413,24 @@ func TestService_FilterToolsForPrompt_AllToolsReturned(t *testing.T) {
require.True(t, hasTool(service.executor.ListTools(), "pulse_docker"))
require.True(t, hasTool(service.executor.ListTools(), "pulse_query"))
// All tools should be returned for any prompt
// Read-only prompts should exclude write tools
filtered := service.filterToolsForPrompt(context.Background(), "run uptime")
assert.False(t, hasProviderTool(filtered, "pulse_control"))
assert.False(t, hasProviderTool(filtered, "pulse_docker"))
assert.True(t, hasProviderTool(filtered, "pulse_query"))
}
func TestService_FilterToolsForPrompt_WriteIntentIncludesWriteTools(t *testing.T) {
service := NewService(Config{
AIConfig: &config.AIConfig{ControlLevel: config.ControlLevelControlled},
StateProvider: &mockStateProvider{},
AgentServer: &mockAgentServer{},
})
require.True(t, hasTool(service.executor.ListTools(), "pulse_control"))
require.True(t, hasTool(service.executor.ListTools(), "pulse_docker"))
require.True(t, hasTool(service.executor.ListTools(), "pulse_query"))
filtered := service.filterToolsForPrompt(context.Background(), "restart vm 101")
assert.True(t, hasProviderTool(filtered, "pulse_control"))
assert.True(t, hasProviderTool(filtered, "pulse_docker"))
assert.True(t, hasProviderTool(filtered, "pulse_query"))
+39 -5
View File
@@ -32,6 +32,11 @@ type SessionStore struct {
// sessionToolSets holds per-session tool allowlists (in-memory only).
// These keep tool availability stable across turns while allowing additive expansion.
sessionToolSets map[string]map[string]bool
// knowledgeAccumulators holds per-session knowledge accumulators (in-memory only).
// These extract and preserve key facts from tool results to prevent amnesia
// when old tool results are compacted from the conversation context.
knowledgeAccumulators map[string]*KnowledgeAccumulator
}
// sessionData is the on-disk format for a session
@@ -51,10 +56,11 @@ func NewSessionStore(dataDir string) (*SessionStore, error) {
}
return &SessionStore{
dataDir: sessionsDir,
resolvedContexts: make(map[string]*ResolvedContext),
sessionFSMs: make(map[string]*SessionFSM),
sessionToolSets: make(map[string]map[string]bool),
dataDir: sessionsDir,
resolvedContexts: make(map[string]*ResolvedContext),
sessionFSMs: make(map[string]*SessionFSM),
sessionToolSets: make(map[string]map[string]bool),
knowledgeAccumulators: make(map[string]*KnowledgeAccumulator),
}, nil
}
@@ -160,10 +166,11 @@ func (s *SessionStore) Delete(id string) error {
return fmt.Errorf("failed to delete session: %w", err)
}
// Also clean up resolved context and FSM
// Also clean up resolved context, FSM, and knowledge accumulator
delete(s.resolvedContexts, id)
delete(s.sessionFSMs, id)
delete(s.sessionToolSets, id)
delete(s.knowledgeAccumulators, id)
return nil
}
@@ -355,6 +362,32 @@ func (s *SessionStore) GetSessionFSM(sessionID string) *SessionFSM {
return fsm
}
// GetKnowledgeAccumulator returns the knowledge accumulator for a session, creating one if needed.
// For user chat sessions, this persists across messages (facts accumulate during a conversation).
func (s *SessionStore) GetKnowledgeAccumulator(sessionID string) *KnowledgeAccumulator {
s.mu.Lock()
defer s.mu.Unlock()
ka, ok := s.knowledgeAccumulators[sessionID]
if !ok {
ka = NewKnowledgeAccumulator()
s.knowledgeAccumulators[sessionID] = ka
}
return ka
}
// NewKnowledgeAccumulatorForRun creates a fresh KA for a patrol run.
// Unlike GetKnowledgeAccumulator (which reuses a session-scoped KA),
// this always returns a new instance to avoid stale facts from prior runs.
func (s *SessionStore) NewKnowledgeAccumulatorForRun(sessionID string) *KnowledgeAccumulator {
s.mu.Lock()
defer s.mu.Unlock()
ka := NewKnowledgeAccumulator()
s.knowledgeAccumulators[sessionID] = ka
return ka
}
// ResetSessionFSM resets the FSM for a session (e.g., after context clear)
func (s *SessionStore) ResetSessionFSM(sessionID string, keepProgress bool) {
s.mu.Lock()
@@ -434,6 +467,7 @@ func (s *SessionStore) ClearSessionState(sessionID string, keepPinned bool) {
if !keepPinned {
delete(s.sessionToolSets, sessionID)
delete(s.knowledgeAccumulators, sessionID)
}
// Reset FSM coherently with context state
+11 -8
View File
@@ -70,6 +70,7 @@ type ExecuteRequest struct {
SessionID string `json:"session_id,omitempty"`
Model string `json:"model,omitempty"`
Mentions []StructuredMention `json:"mentions,omitempty"`
MaxTurns int `json:"max_turns,omitempty"` // Override max agentic turns (0 = use default)
}
// QuestionAnswer represents a user's answer to a question
@@ -90,18 +91,20 @@ type ThinkingData struct {
// ToolStartData is the data for "tool_start" events
type ToolStartData struct {
ID string `json:"id"`
Name string `json:"name"`
Input string `json:"input"` // JSON string of input parameters
ID string `json:"id"`
Name string `json:"name"`
Input string `json:"input"` // JSON string of input parameters
RawInput string `json:"raw_input,omitempty"` // Unmodified JSON input
}
// ToolEndData is the data for "tool_end" events
type ToolEndData struct {
ID string `json:"id"`
Name string `json:"name"`
Input string `json:"input,omitempty"`
Output string `json:"output,omitempty"`
Success bool `json:"success"`
ID string `json:"id"`
Name string `json:"name"`
Input string `json:"input,omitempty"`
RawInput string `json:"raw_input,omitempty"`
Output string `json:"output,omitempty"`
Success bool `json:"success"`
}
// ApprovalNeededData is the data for "approval_needed" events