diff --git a/backend/modules/socai/handler/chat.go b/backend/modules/socai/handler/chat.go index 59cb00b3a..485dbd8bf 100644 --- a/backend/modules/socai/handler/chat.go +++ b/backend/modules/socai/handler/chat.go @@ -67,10 +67,16 @@ func NewChatHandler(client socAIStreamer) *ChatHandler { return &ChatHandler{client: client} } +type chatTurn struct { + Role string `json:"role"` + Content string `json:"content"` +} + type chatRequest struct { - Task string `json:"task" binding:"required"` - Page string `json:"page"` - Lang string `json:"lang"` + Task string `json:"task" binding:"required"` + Page string `json:"page"` + Lang string `json:"lang"` + History []chatTurn `json:"history,omitempty"` } // Chat godoc @@ -101,7 +107,7 @@ func (h *ChatHandler) Chat(c *gin.Context) { return } - body, err := json.Marshal(map[string]string{"task": req.Task, "page": req.Page, "lang": req.Lang}) + body, err := json.Marshal(req) if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"status": "error", "message": err.Error()}) return diff --git a/frontend/src/features/soc-ai/SocAiProvider.tsx b/frontend/src/features/soc-ai/SocAiProvider.tsx index 5c21f8579..bb7c206c7 100644 --- a/frontend/src/features/soc-ai/SocAiProvider.tsx +++ b/frontend/src/features/soc-ai/SocAiProvider.tsx @@ -1,7 +1,11 @@ import { createContext, useCallback, useContext, useMemo, useRef, useState, type Dispatch, type ReactNode, type SetStateAction } from 'react' import { useLocation } from 'react-router-dom' import { useTranslation } from 'react-i18next' -import { extractNavigation, streamChat, type NavAction } from './lib/chat-stream' +import { extractNavigation, streamChat, type ChatHistoryTurn, type NavAction } from './lib/chat-stream' + +// How many prior text turns to replay to the backend as chat memory. Server-side +// compaction will still trim if this exceeds the model context window. +const HISTORY_LIMIT = 10 export interface ToolStep { tool: string @@ -114,8 +118,13 @@ export function SocAiProvider({ children }: { children: ReactNode }) { const page = pageContext(location.pathname) const lang = (i18n.language || 'en').split('-')[0] + const history: ChatHistoryTurn[] = current + .filter((m) => m.text && !m.error && !m.pending) + .slice(-HISTORY_LIMIT) + .map((m) => ({ role: m.role === 'user' ? 'user' : 'assistant', content: m.text })) + streamChat( - { task: text, page, lang }, + { task: text, page, lang, history }, (ev) => { patchMsg(scope, aiId, (msg) => { switch (ev.kind) { diff --git a/frontend/src/features/soc-ai/lib/chat-stream.ts b/frontend/src/features/soc-ai/lib/chat-stream.ts index 73615c60b..6e3c687bf 100644 --- a/frontend/src/features/soc-ai/lib/chat-stream.ts +++ b/frontend/src/features/soc-ai/lib/chat-stream.ts @@ -32,13 +32,21 @@ export interface NavAction { time?: string } +/** A single prior chat turn replayed to the backend as context. Only user and + * assistant text turns are forwarded — tool_use/tool_result blocks are internal + * to a single Run() on the server and must not be replayed. */ +export interface ChatHistoryTurn { + role: 'user' | 'assistant' + content: string +} + /** * Streams the SOC-AI chat agent over SSE. The backend (/soc-ai/chat) proxies the * plugin's agent and emits tool_call / tool_result / final / error events. Uses * fetch + ReadableStream because the shared axios client can't stream. */ export async function streamChat( - body: { task: string; page?: string; lang?: string }, + body: { task: string; page?: string; lang?: string; history?: ChatHistoryTurn[] }, onEvent: (e: ChatEvent) => void, signal?: AbortSignal, ): Promise { diff --git a/plugins/soc-ai/internal/agent/loop.go b/plugins/soc-ai/internal/agent/loop.go index c135c4257..e3fa652a4 100644 --- a/plugins/soc-ai/internal/agent/loop.go +++ b/plugins/soc-ai/internal/agent/loop.go @@ -14,6 +14,8 @@ const ( defaultMaxIters = 12 compactionThreshold = 0.80 summaryMaxTokens = 400 // ~200 words + slack + keepTailMessages = 4 // messages kept raw when compacting + genericErrorMsg = "An error has occurred while processing your request." ) var modelContextWindow = []struct { @@ -69,6 +71,7 @@ func (s EventSink) emit(e Event) { type RunTask struct { System string // system prompt Input string // the user turn (alert JSON for triage, free task for ops) + History []Message EnabledGroups []string AlwaysAllow []string MaxIters int @@ -103,7 +106,10 @@ func (a *Agent) Broker() *ToolBroker { return a.broker } func (a *Agent) Run(ctx context.Context, task RunTask, sink EventSink) (RunResult, error) { specs, err := a.broker.ListSpecs(ctx) if err != nil { - sink.emit(Event{Kind: EventError, Text: "could not load tools: " + err.Error()}) + _ = catcher.Error("could not load tools", err, map[string]any{ + "process": "plugin_com.utmstack.soc-ai", + }) + sink.emit(Event{Kind: EventError, Text: genericErrorMsg}) return RunResult{}, fmt.Errorf("list tools: %w", err) } allowed := filterTools(specs, task) @@ -117,14 +123,19 @@ func (a *Agent) Run(ctx context.Context, task RunTask, sink EventSink) (RunResul maxIters = defaultMaxIters } - msgs := []Message{{Role: RoleUser, Content: task.Input}} + msgs := append([]Message{}, task.History...) + msgs = append(msgs, Message{Role: RoleUser, Content: task.Input}) result := RunResult{} + // Skipped in-batch dedup; upgrade to singleflight if same-batch duplicates become measurable. + toolCache := map[string]tcOut{} + var cacheMu sync.Mutex + for step := 1; step <= maxIters; step++ { result.Steps = step if a.contextWindow > 0 && len(msgs) > 1 && - estimateTokens(task.System, msgs) >= int(compactionThreshold*float64(a.contextWindow)) { + estimateTokens(task.System, msgs, allowed) >= int(compactionThreshold*float64(a.contextWindow)) { newMsgs, cErr := a.compact(ctx, task.Input, msgs) if cErr != nil { _ = catcher.Error("context compaction failed, continuing with full history", cErr, map[string]any{ @@ -144,7 +155,10 @@ func (a *Agent) Run(ctx context.Context, task RunTask, sink EventSink) (RunResul MaxTokens: a.maxTokens, }) if err != nil { - sink.emit(Event{Kind: EventError, Text: err.Error()}) + _ = catcher.Error("llm completion failed", err, map[string]any{ + "process": "plugin_com.utmstack.soc-ai", + }) + sink.emit(Event{Kind: EventError, Text: genericErrorMsg}) return result, err } @@ -156,33 +170,80 @@ func (a *Agent) Run(ctx context.Context, task RunTask, sink EventSink) (RunResul msgs = append(msgs, Message{Role: RoleAssistant, Content: resp.Content, ToolCalls: resp.ToolCalls}) - for _, tc := range resp.ToolCalls { + outs := make([]tcOut, len(resp.ToolCalls)) + var wg sync.WaitGroup + for i, tc := range resp.ToolCalls { result.ToolCalls++ sink.emit(Event{Kind: EventToolCall, Step: step, Tool: tc.Name, Args: tc.Args}) if !allowedSet[tc.Name] { - const msg = "tool not permitted in this mode" - msgs = append(msgs, Message{Role: RoleTool, ToolResult: &ToolResult{ID: tc.ID, Name: tc.Name, Content: msg, IsError: true}}) - sink.emit(Event{Kind: EventToolResult, Step: step, Tool: tc.Name, Output: msg, IsError: true}) + outs[i] = tcOut{out: "tool not permitted in this mode", isErr: true} continue } - out, isErr, callErr := a.broker.Call(ctx, tc.Name, tc.Args) - if callErr != nil { - out = callErr.Error() - isErr = true + key := tc.Name + "|" + string(tc.Args) + cacheMu.Lock() + cached, ok := toolCache[key] + cacheMu.Unlock() + if ok { + outs[i] = cached + continue } - msgs = append(msgs, Message{Role: RoleTool, ToolResult: &ToolResult{ID: tc.ID, Name: tc.Name, Content: out, IsError: isErr}}) - sink.emit(Event{Kind: EventToolResult, Step: step, Tool: tc.Name, Output: out, IsError: isErr}) + + wg.Add(1) + go func(i int, tc ToolCall, key string) { + defer wg.Done() + out, isErr, callErr := a.broker.Call(ctx, tc.Name, tc.Args) + if callErr != nil { + out = callErr.Error() + isErr = true + } + r := tcOut{out: out, isErr: isErr} + outs[i] = r + cacheMu.Lock() + toolCache[key] = r + cacheMu.Unlock() + }(i, tc, key) + } + wg.Wait() + + for i, tc := range resp.ToolCalls { + r := outs[i] + msgs = append(msgs, Message{Role: RoleTool, ToolResult: &ToolResult{ID: tc.ID, Name: tc.Name, Content: r.out, IsError: r.isErr}}) + sink.emit(Event{Kind: EventToolResult, Step: step, Tool: tc.Name, Output: r.out, IsError: r.isErr}) } } - const exhausted = "Reached the maximum number of tool iterations before finishing." - sink.emit(Event{Kind: EventFinal, Text: exhausted}) - result.Final = exhausted + // Loop exhausted: give the model one last chance to finalize with no tools. + msgs = append(msgs, Message{ + Role: RoleUser, + Content: "You have reached the maximum number of tool iterations. Do not call any more tools. Provide your final assessment now based on what you have gathered so far.", + }) + finalResp, ferr := a.llm.Complete(ctx, CompletionRequest{ + System: task.System, + Messages: msgs, + Model: a.model, + MaxTokens: a.maxTokens, + }) + if ferr != nil { + _ = catcher.Error("max-iters finalization llm call failed", ferr, map[string]any{ + "process": "plugin_com.utmstack.soc-ai", + }) + const msg = "Reached the maximum number of tool iterations and could not finalize." + sink.emit(Event{Kind: EventFinal, Text: msg}) + result.Final = msg + return result, nil + } + sink.emit(Event{Kind: EventFinal, Text: finalResp.Content}) + result.Final = finalResp.Content return result, nil } +type tcOut struct { + out string + isErr bool +} + func filterTools(specs []ToolSpec, task RunTask) []ToolSpec { enabled := make(map[string]bool, len(task.EnabledGroups)) for _, g := range task.EnabledGroups { @@ -205,8 +266,16 @@ func filterTools(specs []ToolSpec, task RunTask) []ToolSpec { return out } -func estimateTokens(system string, msgs []Message) int { +func estimateTokens(system string, msgs []Message, tools []ToolSpec) int { n := len(system) + for _, t := range tools { + n += len(t.Name) + len(t.Description) + if t.InputSchema != nil { + if b, err := json.Marshal(t.InputSchema); err == nil { + n += len(b) + } + } + } for _, m := range msgs { n += len(m.Content) for _, tc := range m.ToolCalls { @@ -220,8 +289,24 @@ func estimateTokens(system string, msgs []Message) int { } func (a *Agent) compact(ctx context.Context, userInput string, msgs []Message) ([]Message, error) { + // Keep the last keepTailMessages raw. Advance the cut point forward past any + // tool messages so the preserved tail never starts with an orphan tool_result + // (which would reference an assistant tool_call left in the summarized head). + cut := len(msgs) - keepTailMessages + if cut < 1 { + cut = 1 + } + for cut < len(msgs) && msgs[cut].Role == RoleTool { + cut++ + } + head := msgs[:cut] + var tail []Message + if cut < len(msgs) { + tail = msgs[cut:] + } + var b strings.Builder - for _, m := range msgs { + for _, m := range head { fmt.Fprintf(&b, "[%s] %s\n", m.Role, m.Content) for _, tc := range m.ToolCalls { fmt.Fprintf(&b, " tool_call %s(%s)\n", tc.Name, string(tc.Args)) @@ -242,10 +327,11 @@ func (a *Agent) compact(ctx context.Context, userInput string, msgs []Message) ( if strings.TrimSpace(resp.Content) == "" { return msgs, fmt.Errorf("empty summary") } - return []Message{{ + summary := Message{ Role: RoleUser, Content: "Original task:\n" + userInput + "\n\nProgress so far (summary of prior context):\n" + resp.Content + "\n\nContinue the task.", - }}, nil + } + return append([]Message{summary}, tail...), nil } type registry struct { diff --git a/plugins/soc-ai/internal/agent/prompt.go b/plugins/soc-ai/internal/agent/prompt.go index b3826b04a..f13b7fc37 100644 --- a/plugins/soc-ai/internal/agent/prompt.go +++ b/plugins/soc-ai/internal/agent/prompt.go @@ -37,7 +37,7 @@ func OpsPrompt(page, lang string, enabledGroups []string) string { } langLine := "Answer in the same language as the user's message." if strings.TrimSpace(lang) != "" { - langLine = `Always write your reply in the user's interface language, identified by the code "` + strings.TrimSpace(lang) + `" (e.g. es=Spanish, pt=Portuguese, en=English), regardless of the language of their message.` + langLine = `Always write your reply in the user's interface language, identified by the ISO code "` + strings.TrimSpace(lang) + `", regardless of the language of their message.` } return `You are the UTMStack operations agent — an autonomous SOC assistant embedded in the UTMStack SIEM. The user chats with you, and you operate the SIEM on their behalf through the available tools (alerts, incidents, log/alert search, SOAR response actions, datasources, compliance, and more). @@ -52,7 +52,7 @@ Use this to choose the most relevant tools and to craft navigation. For example, ` + permissionsBlock(enabledGroups) + ` ## How to work -- Plan briefly, then act. Carry the task end to end. +- Carry the task end to end. - Use tools ONLY when you need data or actions you don't already have. Many messages need few or no tools — do not over-call; prefer the smallest set of tools that answers the question. - Prefer read-only tools to investigate before any mutating or response action. Mutating/response actions (changing status, creating incidents, running SOAR jobs, etc.) take effect immediately — only perform them when the task clearly asks for them. - Never invent data; rely on tool results. If a tool fails, adapt or report it plainly. diff --git a/plugins/soc-ai/internal/api/server.go b/plugins/soc-ai/internal/api/server.go index cf4152f41..bd5b0423a 100644 --- a/plugins/soc-ai/internal/api/server.go +++ b/plugins/soc-ai/internal/api/server.go @@ -22,6 +22,16 @@ type AgentTaskRequest struct { // Lang is the user's interface language code (en/es/pt/…) so the agent // replies in that language regardless of the message language. Lang string `json:"lang"` + // History is prior chat turns from the client (text only). Only user and + // assistant roles are accepted; tool_use/tool_result turns are internal to + // a single Run() and must not be replayed. + History []AgentTurn `json:"history,omitempty"` +} + +// AgentTurn is a single prior chat message forwarded by the client. +type AgentTurn struct { + Role string `json:"role"` // "user" | "assistant" + Content string `json:"content"` } // AnalyzeRequest represents the request body for manual alert analysis @@ -214,11 +224,38 @@ func handleAgentTask(w http.ResponseWriter, r *http.Request) { _, _ = ag.Run(r.Context(), agent.RunTask{ System: agent.OpsPrompt(req.Page, req.Lang, capabilities), Input: req.Task, + History: toHistory(req.History), EnabledGroups: capabilities, MaxIters: maxIters, }, sink) } +// toHistory converts client-supplied turns into agent.Message. Unknown roles +// and empty content are dropped so a malformed client can't inject tool turns +// or blank rows. +func toHistory(turns []AgentTurn) []agent.Message { + if len(turns) == 0 { + return nil + } + out := make([]agent.Message, 0, len(turns)) + for _, t := range turns { + if t.Content == "" { + continue + } + var role agent.Role + switch t.Role { + case "user": + role = agent.RoleUser + case "assistant": + role = agent.RoleAssistant + default: + continue + } + out = append(out, agent.Message{Role: role, Content: t.Content}) + } + return out +} + func writeJSONError(w http.ResponseWriter, status int, msg string) { w.Header().Set("Content-Type", "application/json") w.WriteHeader(status)