diff --git a/frontend-modern/src/components/AI/Chat/hooks/useChat.ts b/frontend-modern/src/components/AI/Chat/hooks/useChat.ts index 58b3a12d4..7e2c52ff5 100644 --- a/frontend-modern/src/components/AI/Chat/hooks/useChat.ts +++ b/frontend-modern/src/components/AI/Chat/hooks/useChat.ts @@ -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; + 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: [] }; } diff --git a/frontend-modern/src/components/AI/Chat/index.tsx b/frontend-modern/src/components/AI/Chat/index.tsx index 0055918ae..64f4deae3 100644 --- a/frontend-modern/src/components/AI/Chat/index.tsx +++ b/frontend-modern/src/components/AI/Chat/index.tsx @@ -104,6 +104,7 @@ export const AIChat: Component = (props) => { // Chat hook const chat = useChat({ model: initialModelSelections[DEFAULT_SESSION_KEY] || '' }); + const defaultModelLabel = createMemo(() => { const fallback = defaultModel().trim(); if (!fallback) return ''; diff --git a/frontend-modern/src/components/AI/Chat/types.ts b/frontend-modern/src/components/AI/Chat/types.ts index 1666df0ef..454b06c95 100644 --- a/frontend-modern/src/components/AI/Chat/types.ts +++ b/frontend-modern/src/components/AI/Chat/types.ts @@ -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; +} diff --git a/internal/ai/chat/agentic.go b/internal/ai/chat/agentic.go index f2d904e86..f3be021c1 100644 --- a/internal/ai/chat/agentic.go +++ b/internal/ai/chat/agentic.go @@ -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 diff --git a/internal/ai/chat/agentic_test.go b/internal/ai/chat/agentic_test.go index dbb6e8ff3..d713d1546 100644 --- a/internal/ai/chat/agentic_test.go +++ b/internal/ai/chat/agentic_test.go @@ -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) { diff --git a/internal/ai/chat/knowledge_accumulator.go b/internal/ai/chat/knowledge_accumulator.go new file mode 100644 index 000000000..94ff2c6cc --- /dev/null +++ b/internal/ai/chat/knowledge_accumulator.go @@ -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() +} diff --git a/internal/ai/chat/knowledge_accumulator_test.go b/internal/ai/chat/knowledge_accumulator_test.go new file mode 100644 index 000000000..a3dc2ab3e --- /dev/null +++ b/internal/ai/chat/knowledge_accumulator_test.go @@ -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) +} diff --git a/internal/ai/chat/knowledge_extractor.go b/internal/ai/chat/knowledge_extractor.go new file mode 100644 index 000000000..e69ea22cd --- /dev/null +++ b/internal/ai/chat/knowledge_extractor.go @@ -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 +} diff --git a/internal/ai/chat/knowledge_extractor_test.go b/internal/ai/chat/knowledge_extractor_test.go new file mode 100644 index 000000000..ff2461970 --- /dev/null +++ b/internal/ai/chat/knowledge_extractor_test.go @@ -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) +} diff --git a/internal/ai/chat/service.go b/internal/ai/chat/service.go index 383eb0995..42ea6dcf5 100644 --- a/internal/ai/chat/service.go +++ b/internal/ai/chat/service.go @@ -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) diff --git a/internal/ai/chat/service_test.go b/internal/ai/chat/service_test.go index a347b5ccc..3c232bec7 100644 --- a/internal/ai/chat/service_test.go +++ b/internal/ai/chat/service_test.go @@ -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")) diff --git a/internal/ai/chat/session.go b/internal/ai/chat/session.go index cdfea963b..985490adf 100644 --- a/internal/ai/chat/session.go +++ b/internal/ai/chat/session.go @@ -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 diff --git a/internal/ai/chat/types.go b/internal/ai/chat/types.go index b9e395491..ded9083dc 100644 --- a/internal/ai/chat/types.go +++ b/internal/ai/chat/types.go @@ -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