From c0147248962cdd9c661e4735a045c198e35606a7 Mon Sep 17 00:00:00 2001 From: shinchan-zhai Date: Mon, 23 Mar 2026 15:37:36 +0800 Subject: [PATCH] fix: context propagation in LLM calls, SSE error handling, security hardening MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Propagate request context through MCP client (Request.Ctx field) — Client disconnects now cancel in-flight LLM API calls (both streaming and non-streaming) — Prevents wasted LLM tokens when user navigates away mid-response - proxyBinance: use http.NewRequestWithContext for context-aware upstream calls — Client disconnects cancel Binance proxy requests — Distinguishes client cancellation from upstream failures - Fix SSE parse bug in AgentChatPage: catch block was swallowing error events — throw new Error(data) for 'error' events was caught by the same try/catch — Now parse JSON separately, then handle events outside try/catch - Add io.LimitReader to Hyperliquid coins.go JSON decoder (4MB limit) - Use safe.ReadAllLimited in batch ticker handler for consistency - Sanitize URL validation error in handler_ai_model.go (was leaking raw error) - Cache agent tool definitions (built once, reused per message) - Sort trade history by exit time for consistent ordering --- agent/agent.go | 34 ++++++++-- agent/tools.go | 17 ++++- agent/web.go | 109 ++++++++++++++++++++++++-------- api/handler_ai_model.go | 3 +- mcp/client.go | 15 ++++- mcp/request.go | 7 ++ provider/hyperliquid/coins.go | 4 +- web/src/pages/AgentChatPage.tsx | 77 +++++++++++----------- 8 files changed, 189 insertions(+), 77 deletions(-) diff --git a/agent/agent.go b/agent/agent.go index 11d0f349..6a0e20bf 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -11,6 +11,7 @@ import ( "fmt" "log/slog" "net/http" + "sort" "strconv" "strings" "time" @@ -18,6 +19,7 @@ import ( "nofx/manager" "nofx/market" "nofx/mcp" + "nofx/safe" "nofx/store" ) @@ -95,6 +97,16 @@ func (a *Agent) Start() { if a.config.EnableBriefs { a.brain.StartMarketBriefs(a.config.BriefTimes) } a.scheduler = NewScheduler(a, a.logger) a.scheduler.Start(context.Background()) + + // Periodic cleanup of stale chat sessions (older than 4 hours) + safe.GoNamed("chat-history-cleanup", func() { + ticker := time.NewTicker(30 * time.Minute) + defer ticker.Stop() + for range ticker.C { + a.history.CleanOld(4 * time.Hour) + } + }) + a.logger.Info("NOFXi agent is online 🚀") } @@ -234,6 +246,7 @@ func (a *Agent) thinkAndAct(ctx context.Context, userID int64, lang, text string Messages: messages, Tools: tools, ToolChoice: "auto", + Ctx: ctx, } resp, err := a.aiClient.CallWithRequestFull(req) @@ -241,7 +254,7 @@ func (a *Agent) thinkAndAct(ctx context.Context, userID int64, lang, text string a.logger.Error("LLM call failed", "error", err, "round", round) if round == 0 { // First round failed — try without tools as fallback - plainReq := &mcp.Request{Messages: messages} + plainReq := &mcp.Request{Messages: messages, Ctx: ctx} plainResp, plainErr := a.aiClient.CallWithRequest(plainReq) if plainErr != nil { return a.noAIFallback(lang, text) @@ -293,7 +306,7 @@ func (a *Agent) thinkAndAct(ctx context.Context, userID int64, lang, text string } // If we exhausted all rounds, ask LLM for a final text response without tools - finalReq := &mcp.Request{Messages: messages} + finalReq := &mcp.Request{Messages: messages, Ctx: ctx} finalResp, err := a.aiClient.CallWithRequest(finalReq) if err != nil { return a.noAIFallback(lang, text) @@ -343,13 +356,13 @@ func (a *Agent) thinkAndActStream(ctx context.Context, userID int64, lang, text toolsUsed := false for round := 0; round < maxToolRounds; round++ { - req := &mcp.Request{Messages: messages, Tools: tools, ToolChoice: "auto"} + req := &mcp.Request{Messages: messages, Tools: tools, ToolChoice: "auto", Ctx: ctx} resp, err := a.aiClient.CallWithRequestFull(req) if err != nil { a.logger.Error("LLM call failed (stream)", "error", err, "round", round) if round == 0 { // First round failed — try streaming without tools as fallback - streamReq := &mcp.Request{Messages: messages} + streamReq := &mcp.Request{Messages: messages, Ctx: ctx} streamText, streamErr := a.aiClient.CallWithRequestStream(streamReq, func(chunk string) { onEvent(StreamEventDelta, chunk) }) @@ -396,7 +409,7 @@ func (a *Agent) thinkAndActStream(ctx context.Context, userID int64, lang, text // After tool execution, stream the next LLM response for real-time UX. // Omit tools so LLM can't start another tool round — it must produce text. - streamReq := &mcp.Request{Messages: messages} + streamReq := &mcp.Request{Messages: messages, Ctx: ctx} streamText, streamErr := a.aiClient.CallWithRequestStream(streamReq, func(chunk string) { onEvent(StreamEventDelta, chunk) }) @@ -409,7 +422,7 @@ func (a *Agent) thinkAndActStream(ctx context.Context, userID int64, lang, text } // Exhausted all tool rounds — stream the final synthesis response - finalReq := &mcp.Request{Messages: messages} + finalReq := &mcp.Request{Messages: messages, Ctx: ctx} finalText, err := a.aiClient.CallWithRequestStream(finalReq, func(chunk string) { onEvent(StreamEventDelta, chunk) }) @@ -582,9 +595,16 @@ func (a *Agent) gatherContext(text string) string { } } } + // Collect and sort matched symbols for deterministic selection + sortedSymbols := make([]string, 0, len(matched)) + for sym := range matched { + sortedSymbols = append(sortedSymbols, sym) + } + sort.Strings(sortedSymbols) + // Cap at 5 symbols to avoid slow context gathering count := 0 - for sym := range matched { + for _, sym := range sortedSymbols { if count >= 5 { break } md, err := market.Get(sym + "USDT") if err == nil && md.CurrentPrice > 0 { diff --git a/agent/tools.go b/agent/tools.go index d934d7fd..3640b885 100644 --- a/agent/tools.go +++ b/agent/tools.go @@ -4,14 +4,20 @@ import ( "context" "encoding/json" "fmt" + "sort" "strings" "time" "nofx/mcp" ) -// agentTools defines the tools available to the LLM for autonomous action. -func agentTools() []mcp.Tool { +// cachedTools holds the static tool definitions (built once, reused per message). +var cachedTools = buildAgentTools() + +// agentTools returns the tools available to the LLM for autonomous action. +func agentTools() []mcp.Tool { return cachedTools } + +func buildAgentTools() []mcp.Tool { return []mcp.Tool{ { Type: "function", @@ -423,6 +429,13 @@ func (a *Agent) toolGetTradeHistory(argsJSON string) string { return `{"trades": [], "message": "no closed trades found"}` } + // Sort trades by exit time (most recent first) for consistent ordering across traders + sort.Slice(trades, func(i, j int) bool { + ti, _ := trades[i]["exit_time"].(string) + tj, _ := trades[j]["exit_time"].(string) + return ti > tj // reverse chronological + }) + // Only return up to the limit if len(trades) > args.Limit { trades = trades[:args.Limit] diff --git a/agent/web.go b/agent/web.go index b4b42132..926052c9 100644 --- a/agent/web.go +++ b/agent/web.go @@ -7,6 +7,7 @@ import ( "io" "log/slog" "net/http" + "nofx/safe" "regexp" "time" ) @@ -171,7 +172,7 @@ func (w *WebHandler) HandleKlines(rw http.ResponseWriter, r *http.Request) { return } - proxyBinance(rw, fmt.Sprintf("https://fapi.binance.com/fapi/v1/klines?symbol=%s&interval=%s&limit=300", symbol, interval)) + proxyBinance(rw, r.Context(), fmt.Sprintf("https://fapi.binance.com/fapi/v1/klines?symbol=%s&interval=%s&limit=300", symbol, interval)) } // HandleTicker proxies ticker data from Binance. @@ -184,22 +185,22 @@ func (w *WebHandler) HandleTicker(rw http.ResponseWriter, r *http.Request) { return } - proxyBinance(rw, fmt.Sprintf("https://fapi.binance.com/fapi/v1/ticker/24hr?symbol=%s", symbol)) + proxyBinance(rw, r.Context(), fmt.Sprintf("https://fapi.binance.com/fapi/v1/ticker/24hr?symbol=%s", symbol)) } // HandleTickers handles GET /api/agent/tickers?symbols=BTCUSDT,ETHUSDT,SOLUSDT -// Batch endpoint: fetches multiple tickers in one API call to Binance. +// Batch endpoint: fetches multiple tickers concurrently, returns array. func (w *WebHandler) HandleTickers(rw http.ResponseWriter, r *http.Request) { symbolsParam := r.URL.Query().Get("symbols") if symbolsParam == "" { symbolsParam = "BTCUSDT,ETHUSDT,SOLUSDT" } - // Validate and build JSON array of symbols - symbols := []string{} - for _, s := range splitAndTrim(symbolsParam) { + // Validate symbols + var symbols []string + for _, s := range splitComma(symbolsParam) { if validSymbolRe.MatchString(s) { - symbols = append(symbols, `"`+s+`"`) + symbols = append(symbols, s) } } if len(symbols) == 0 { @@ -211,15 +212,70 @@ func (w *WebHandler) HandleTickers(rw http.ResponseWriter, r *http.Request) { return } - // Binance supports batch ticker: /fapi/v1/ticker/24hr with multiple symbols - url := fmt.Sprintf("https://fapi.binance.com/fapi/v1/ticker/24hr?symbols=[%s]", joinStrings(symbols, ",")) - proxyBinance(rw, url) + // Fetch all tickers concurrently with context propagation + ctx, cancel := context.WithTimeout(r.Context(), 10*time.Second) + defer cancel() + + type result struct { + idx int + data json.RawMessage + } + results := make(chan result, len(symbols)) + for i, sym := range symbols { + idx, s := i, sym + safe.GoNamed("ticker-fetch-"+s, func() { + req, err := http.NewRequestWithContext(ctx, "GET", + fmt.Sprintf("https://fapi.binance.com/fapi/v1/ticker/24hr?symbol=%s", s), nil) + if err != nil { + results <- result{idx: idx} + return + } + resp, err := binanceClient.Do(req) + if err != nil { + results <- result{idx: idx} + return + } + defer resp.Body.Close() + if resp.StatusCode != 200 { + results <- result{idx: idx} + return + } + body, err := safe.ReadAllLimited(resp.Body, 16*1024) + if err != nil { + results <- result{idx: idx} + return + } + results <- result{idx: idx, data: body} + }) + } + + // Collect results in order + ordered := make([]json.RawMessage, len(symbols)) + for range symbols { + r := <-results + if r.data != nil { + ordered[r.idx] = r.data + } + } + + // Filter out nil entries and write response + out := make([]json.RawMessage, 0, len(ordered)) + for _, d := range ordered { + if d != nil { + out = append(out, d) + } + } + rw.Header().Set("Content-Type", "application/json") + json.NewEncoder(rw).Encode(out) } -// splitAndTrim splits a comma-separated string and trims whitespace. -func splitAndTrim(s string) []string { - parts := []string{} - for _, p := range regexp.MustCompile(`\s*,\s*`).Split(s, -1) { +// commaRe is pre-compiled for splitComma — avoids recompiling on every call. +var commaRe = regexp.MustCompile(`\s*,\s*`) + +// splitComma splits a comma-separated string, trims whitespace, skips empty. +func splitComma(s string) []string { + var parts []string + for _, p := range commaRe.Split(s, -1) { if p != "" { parts = append(parts, p) } @@ -227,21 +283,18 @@ func splitAndTrim(s string) []string { return parts } -// joinStrings joins strings with a separator. -func joinStrings(ss []string, sep string) string { - result := "" - for i, s := range ss { - if i > 0 { - result += sep - } - result += s - } - return result -} - -func proxyBinance(rw http.ResponseWriter, url string) { - resp, err := binanceClient.Get(url) +func proxyBinance(rw http.ResponseWriter, ctx context.Context, url string) { + req, err := http.NewRequestWithContext(ctx, "GET", url, nil) if err != nil { + writeJSON(rw, 500, map[string]string{"error": "failed to create request"}) + return + } + resp, err := binanceClient.Do(req) + if err != nil { + // Distinguish client cancellation from upstream failures + if ctx.Err() != nil { + return // Client disconnected, no point writing response + } writeJSON(rw, 502, map[string]string{"error": "upstream request failed"}) return } diff --git a/api/handler_ai_model.go b/api/handler_ai_model.go index 48889d6a..0e4be79b 100644 --- a/api/handler_ai_model.go +++ b/api/handler_ai_model.go @@ -155,7 +155,8 @@ func (s *Server) handleUpdateModelConfigs(c *gin.Context) { if modelData.CustomAPIURL != "" { cleanURL := strings.TrimSuffix(modelData.CustomAPIURL, "#") if err := security.ValidateURL(cleanURL); err != nil { - c.JSON(http.StatusBadRequest, gin.H{"error": fmt.Sprintf("Invalid custom_api_url for model %s: %s", modelID, err.Error())}) + logger.Warnf("Invalid custom_api_url for model %s: %v", modelID, err) + c.JSON(http.StatusBadRequest, gin.H{"error": fmt.Sprintf("Invalid custom_api_url for model %s: URL must be a valid HTTPS endpoint", modelID)}) return } } diff --git a/mcp/client.go b/mcp/client.go index d7abff6f..e24c491a 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -501,6 +501,10 @@ func (client *Client) callWithRequestFull(req *Request) (*LLMResponse, error) { if err != nil { return nil, fmt.Errorf("failed to create request: %w", err) } + // Propagate request context so client disconnects cancel in-flight LLM calls + if req.Ctx != nil { + httpReq = httpReq.WithContext(req.Ctx) + } resp, err := client.HTTPClient.Do(httpReq) if err != nil { @@ -539,6 +543,10 @@ func (client *Client) callWithRequest(req *Request) (string, error) { if err != nil { return "", fmt.Errorf("failed to create request: %w", err) } + // Propagate request context so client disconnects cancel in-flight LLM calls + if req.Ctx != nil { + httpReq = httpReq.WithContext(req.Ctx) + } resp, err := client.HTTPClient.Do(httpReq) if err != nil { @@ -684,8 +692,13 @@ func (client *Client) CallWithRequestStream(req *Request, onChunk func(string)) // Idle-timeout watchdog: cancel the request if no SSE line arrives for 60 seconds. // This breaks the scanner out of an indefinitely blocking Read on a hung connection. + // Use request's Ctx as parent so client disconnects propagate cancellation. const idleTimeout = 60 * time.Second - ctx, cancel := context.WithCancel(context.Background()) + parentCtx := req.Ctx + if parentCtx == nil { + parentCtx = context.Background() + } + ctx, cancel := context.WithCancel(parentCtx) defer cancel() resetCh := make(chan struct{}, 1) safe.GoNamed("mcp-stream-idle-watchdog", func() { diff --git a/mcp/request.go b/mcp/request.go index 548ef094..a0211906 100644 --- a/mcp/request.go +++ b/mcp/request.go @@ -1,5 +1,7 @@ package mcp +import "context" + // Message represents a conversation message. // Supports plain messages (Role+Content), assistant tool-call messages (ToolCalls), // and tool result messages (Role="tool", ToolCallID, Content). @@ -62,6 +64,11 @@ type Request struct { // Advanced features Tools []Tool `json:"tools,omitempty"` // Available tools list ToolChoice string `json:"tool_choice,omitempty"` // Tool choice strategy ("auto", "none", {"type": "function", "function": {"name": "xxx"}}) + + // Context for cancellation — not serialized to JSON. + // If set, HTTP requests use this as the parent context so client disconnects + // propagate and cancel in-flight LLM calls. + Ctx context.Context `json:"-"` } // NewMessage creates a message diff --git a/provider/hyperliquid/coins.go b/provider/hyperliquid/coins.go index 7e5b0d45..ddfe5b08 100644 --- a/provider/hyperliquid/coins.go +++ b/provider/hyperliquid/coins.go @@ -5,6 +5,7 @@ import ( "context" "encoding/json" "fmt" + "io" "net/http" "nofx/logger" "sort" @@ -82,8 +83,9 @@ func (p *CoinProvider) fetchCoins(ctx context.Context) error { } // Response is an array: [meta, [assetCtxs...]] + // Limit read to 4MB to prevent memory exhaustion from oversized responses var rawResp []json.RawMessage - if err := json.NewDecoder(resp.Body).Decode(&rawResp); err != nil { + if err := json.NewDecoder(io.LimitReader(resp.Body, 4*1024*1024)).Decode(&rawResp); err != nil { return fmt.Errorf("failed to decode response: %w", err) } diff --git a/web/src/pages/AgentChatPage.tsx b/web/src/pages/AgentChatPage.tsx index 7fa32c5f..cce651bf 100644 --- a/web/src/pages/AgentChatPage.tsx +++ b/web/src/pages/AgentChatPage.tsx @@ -145,45 +145,48 @@ export function AgentChatPage() { eventType = line.slice(7).trim() } else if (line.startsWith('data: ') && eventType) { const rawData = line.slice(6) + let data: string try { - const data = JSON.parse(rawData) - if (eventType === 'delta') { - // data is the accumulated text so far - finalText = data - setMessages((prev) => - prev.map((m) => - m.id === botId - ? { ...m, text: data, time: now() } - : m - ) - ) - } else if (eventType === 'tool') { - // Show tool being called as a status indicator - setMessages((prev) => - prev.map((m) => - m.id === botId - ? { - ...m, - text: m.text || `🔧 _Calling ${data}..._`, - time: now(), - } - : m - ) - ) - } else if (eventType === 'done') { - finalText = data - setMessages((prev) => - prev.map((m) => - m.id === botId - ? { ...m, text: data, time: now(), streaming: false } - : m - ) - ) - } else if (eventType === 'error') { - throw new Error(data) - } - } catch (parseErr) { + data = JSON.parse(rawData) + } catch { // Ignore malformed SSE data lines + eventType = '' + continue + } + if (eventType === 'delta') { + // data is the accumulated text so far + finalText = data + setMessages((prev) => + prev.map((m) => + m.id === botId + ? { ...m, text: data, time: now() } + : m + ) + ) + } else if (eventType === 'tool') { + // Show tool being called as a status indicator + setMessages((prev) => + prev.map((m) => + m.id === botId + ? { + ...m, + text: m.text || `🔧 _Calling ${data}..._`, + time: now(), + } + : m + ) + ) + } else if (eventType === 'done') { + finalText = data + setMessages((prev) => + prev.map((m) => + m.id === botId + ? { ...m, text: data, time: now(), streaming: false } + : m + ) + ) + } else if (eventType === 'error') { + throw new Error(data) } eventType = '' }