feat: network diagram improvements — node limit, filter persistence, report accuracy (#195)

* feat: story improvements — pre-gen screen, context-length error UX, prompt caps, timeline z-index, report filters

- Remove auto-generation on Story tab visit; show a pre-generation screen
  with StoryInfoCard so users can configure settings before generating
- Detect LLM context-length exceeded errors (HTTP 422 / errorCode
  CONTEXT_LENGTH_EXCEEDED) and surface the full prompt in an editable
  textarea so the user can trim it and retry with a custom prompt
- Add prompt-cap controls (Max findings / Max risk matrix rows) to
  StoryInfoCard: preset buttons, custom number input, and an "All" button
  that shows the total count once a story has been generated
- Backend: new ContextLengthExceededException, LlmClient parses token
  counts from 400 body; GenerateStoryRequest accepts customPrompt,
  maxFindings, maxRiskMatrix; StoryService caps findings/risk-matrix rows
  and short-circuits to custom-prompt path on retry
- Fix Recharts Tooltip z-index in TrafficTimeline so modal renders on top
- Pass user session filters to captureNetworkDiagrams in report generation

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* feat: add LLM_CONTEXT_LENGTH env var and clarify LLM_MAX_TOKENS

- Add LLM_CONTEXT_LENGTH config property so operators can explicitly set
  the model's context window size instead of relying on /v1/models
  auto-detection
- LlmClient now prefers the configured context length over auto-detect,
  with auto-detect as fallback and a clearer log message for each path
- Clarify the 80% guard comment: it caps effectiveMaxTokens as a safety
  measure, not the primary purpose of LLM_MAX_TOKENS
- LLM_MAX_TOKENS comment in .env rewritten to make clear it controls
  response length only (not context window size)

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* docs: update .env.example with LLM_CONTEXT_LENGTH and corrected LLM_MAX_TOKENS comment

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* feat: wire LLM_CONTEXT_LENGTH into docker-compose and add pre-flight token check

- Add LLM_CONTEXT_LENGTH to docker-compose.yml so operators can set it
  without rebuilding the image
- LlmClient now performs a pre-flight estimate (chars/4 heuristic) before
  calling the LLM server; if estimated prompt + response reserve exceeds
  the known context length, throws ContextLengthExceededException
  immediately — no LLM call is made

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* fix: keep suggested question buttons visible until send is clicked

Previously clicking a suggestion immediately cleared the buttons.
Now they persist until the question is actually submitted (they
still disappear during loading as before).

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* fix: address gemini review — transaction scope, pattern precompile, dead code, timeout check

- Remove @transactional from generateStory to avoid holding a DB
  connection open during multi-minute LLM calls
- Pre-compile regex Patterns as static constants in LlmClient
- Delete unused StoryCountsResponse DTO and STORY_COUNTS endpoint key
- Remove unused loadingLimits state from StoryPage
- Narrow isTimeout check: drop includes('exceeded') to avoid masking
  CONTEXT_LENGTH_EXCEEDED errors as timeout messages

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* feat: network diagram improvements — node limit banner, filter persistence, report accuracy

- Add node-limit banner to compare topology view (mirrors single-file view),
  showing "Top N / All N" controls with presets and custom input; banner stays
  visible even when all nodes are shown
- Per-file buildNetworkGraph in compare mode now passes maxNodes=0 so
  significance filtering is applied post-merge rather than per-file
- Lift network diagram filter state into AnalysisPage so filters survive
  tab navigation (session-level cache)
- Report diagram capture now uses the already-filtered nodes/edges visible
  on screen instead of re-fetching, so the PDF matches exactly what the user sees
- Active filter labels are sent to the backend and rendered as a dedicated
  PDF section before the topology diagrams

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* fix: address gemini review — empty diagrams fallback, null filter guard, token heuristic, context error handling, investigation steps in retry

- AnalysisPage: fall back to a fresh conversation fetch when networkGraphStateRef
  is empty (user never visited Network Diagram tab), so PDF always has diagrams
- ReportService: skip null filter strings in addNetworkDiagramFilters to prevent NPE
- LlmClient: tighten pre-flight token estimate from chars/4 to chars/2 for denser
  technical content; broaden context-length error detection to catch provider variants
  (Ollama, LM Studio); use estimated counts as fallback when regex finds no match
- StoryService: re-run Phase 1 investigation during custom prompt retry so
  investigationSteps are preserved in the response instead of being silently dropped

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

---------

Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: NotYuSheng

backend/src/main/java/com/tracepcap/report/ReportRequest.java

@@ -1,12 +1,15 @@
 package com.tracepcap.report;
 
+import java.util.List;
 import lombok.Data;
 import lombok.NoArgsConstructor;
 
-/** Optional diagram images (base64-encoded PNGs) sent by the frontend. */
+/** Optional diagram images (base64-encoded PNGs) and active filter labels sent by the frontend. */
 @Data
 @NoArgsConstructor
 public class ReportRequest {
   private String forceDirectedImage;
   private String hierarchicalImage;
+  /** Human-readable labels for each active network-diagram filter, e.g. "Protocol: HTTPS". */
+  private List<String> activeFilters;
 }

backend/src/main/java/com/tracepcap/report/ReportService.java

@@ -202,6 +202,10 @@ public void generateReport(UUID fileId, ReportRequest request, OutputStream out)
         addExtractedFiles(document, extractedFiles, sec++);
       }
 
+      if (request.getActiveFilters() != null && !request.getActiveFilters().isEmpty()) {
+        addNetworkDiagramFilters(document, request.getActiveFilters(), sec++);
+      }
+
       addTopologyDiagram(document, request.getForceDirectedImage(), "Force-Directed Layout", sec++);
       addTopologyDiagram(
           document, request.getHierarchicalImage(), "Hierarchical Layout (Top-Down)", sec++);
@@ -775,6 +779,50 @@ private void addExtractedFiles(Document doc, List<ExtractedFileEntity> files, in
     doc.add(table);
   }
 
+  // ══════════════════════════════════════════════════════════════════════════
+  // Section: Network Diagram Filters Applied
+  // ══════════════════════════════════════════════════════════════════════════
+
+  private void addNetworkDiagramFilters(Document doc, List<String> filters, int sec)
+      throws Exception {
+    addSectionHeader(doc, sec + ". Network Diagram — Active Filters");
+
+    Font bodyFont = new Font(Font.HELVETICA, 10, Font.NORMAL, C_TEXT);
+    Paragraph intro =
+        new Paragraph(
+            "The network topology diagrams in this report were generated with the following filters applied:",
+            bodyFont);
+    intro.setSpacingBefore(6);
+    intro.setSpacingAfter(8);
+    doc.add(intro);
+
+    PdfPTable table = new PdfPTable(1);
+    table.setWidthPercentage(100);
+    table.setSpacingAfter(12);
+    Font labelFont = new Font(Font.HELVETICA, 10, Font.BOLD, C_LABEL);
+    Font valueFont = new Font(Font.HELVETICA, 10, Font.NORMAL, C_TEXT);
+    for (int i = 0; i < filters.size(); i++) {
+      String filter = filters.get(i);
+      if (filter == null) continue;
+      int colon = filter.indexOf(':');
+      Color bg = i % 2 == 0 ? Color.WHITE : C_ROW_ALT;
+      PdfPCell cell = new PdfPCell();
+      cell.setBackgroundColor(bg);
+      cell.setPadding(6);
+      cell.setBorder(Rectangle.NO_BORDER);
+      if (colon > 0) {
+        Phrase phrase = new Phrase();
+        phrase.add(new Phrase(filter.substring(0, colon + 1) + " ", labelFont));
+        phrase.add(new Phrase(filter.substring(colon + 1).trim(), valueFont));
+        cell.setPhrase(phrase);
+      } else {
+        cell.setPhrase(new Phrase(filter, valueFont));
+      }
+      table.addCell(cell);
+    }
+    doc.add(table);
+  }
+
   // ══════════════════════════════════════════════════════════════════════════
   // Section: Network Topology Diagrams (frontend-rendered PNG)
   // ══════════════════════════════════════════════════════════════════════════

backend/src/main/java/com/tracepcap/story/service/LlmClient.java

@@ -157,10 +157,12 @@ public Integer getModelContextLength() {
    */
   public String generateCompletion(String systemPrompt, String userPrompt) {
     // Pre-flight context-length check — fail immediately without calling the LLM server.
-    // Estimate token count using the ~4 chars/token heuristic (good enough for a guard).
+    // Use ~2 chars/token (conservative) rather than 4 — technical content such as network
+    // logs, JSON, and hex strings tokenises more densely and can easily fall below 3 chars/token.
+    // A tighter estimate means fewer false passes that still fail at the server side.
     // Only fires when modelContextLength is known (LLM_CONTEXT_LENGTH set or auto-detected).
     if (modelContextLength != null) {
-      int estimatedPromptTokens = (systemPrompt.length() + userPrompt.length()) / 4;
+      int estimatedPromptTokens = (systemPrompt.length() + userPrompt.length()) / 2;
       int responseReserve = getEffectiveMaxTokens();
       if (estimatedPromptTokens + responseReserve > modelContextLength) {
         log.warn(
@@ -214,15 +216,36 @@ public String generateCompletion(String systemPrompt, String userPrompt) {
     } catch (LlmException e) {
       throw e;
     } catch (Exception e) {
-      // Detect context-length exceeded (OpenAI-compatible 400 response)
+      // Detect context-length exceeded errors.
+      // Primary: OpenAI-compatible "maximum context length" message with token counts.
+      // Fallback: any HTTP 400 whose message hints at context/token overflow — different
+      // providers (Ollama, LM Studio, vLLM) use varying formats so regex may not match.
       String msg = e.getMessage() != null ? e.getMessage() : "";
-      if (msg.contains("maximum context length")) {
+      boolean isHttp400 = msg.contains("400");
+      boolean looksLikeContextError = msg.contains("maximum context length")
+          || msg.contains("context_length_exceeded")
+          || msg.contains("context window")
+          || msg.contains("token limit");
+      if (looksLikeContextError) {
         int promptTokens = parseGroup(msg, PATTERN_PROMPT_TOKENS);
         int contextTokens = parseGroup(msg, PATTERN_CONTEXT_LENGTH);
         if (contextTokens == 0) contextTokens = parseGroup(msg, PATTERN_CONTEXT_LENGTH_ALT);
+        // If regex couldn't extract counts, use the pre-flight estimate as a best-effort value
+        // so the UI shows something meaningful rather than "0 tokens".
+        if (promptTokens == 0) {
+          promptTokens = (systemPrompt.length() + userPrompt.length()) / 2;
+        }
+        if (contextTokens == 0 && modelContextLength != null) {
+          contextTokens = modelContextLength;
+        }
         throw new ContextLengthExceededException(promptTokens, contextTokens, userPrompt);
       }
-      log.error("Error calling LLM API", e);
+      // Re-classify ambiguous HTTP 400s that don't match the patterns above
+      if (isHttp400) {
+        log.warn("LLM returned HTTP 400 — likely context length or bad request: {}", msg);
+      } else {
+        log.error("Error calling LLM API", e);
+      }
       throw new LlmException("Failed to reach the LLM service: " + e.getMessage(), e);
     }
   }

backend/src/main/java/com/tracepcap/story/service/StoryService.java

@@ -86,14 +86,40 @@ public StoryResponse generateStory(UUID fileId, String additionalContext, String
       // skip all prompt-building phases and send it directly to the LLM.
       if (customPrompt != null && !customPrompt.isBlank()) {
         log.info("Using user-supplied custom prompt for file: {}", fileId);
+        long totalConvsCustom = conversationRepository.countByFileId(fileId);
         StoryAggregates aggregates = storyAggregatesService.compute(
-            fileId, List.of(), conversationRepository.countByFileId(fileId));
+            fileId, List.of(), totalConvsCustom);
         List<Finding> findings = findingsService.detectAll(
-            fileId, conversationRepository.countByFileId(fileId), analysis.getTotalBytes());
+            fileId, totalConvsCustom, analysis.getTotalBytes());
+
+        // Re-run Phase 1 so investigation steps are preserved in the response.
+        // The custom prompt already encodes the narrative context, but structured
+        // investigation results are attached separately and displayed in the UI.
+        List<InvestigationStep> investigationSteps = List.of();
+        List<TimelineDataDto> timelineBinsCustom = List.of();
+        try {
+          timelineBinsCustom = timelineService.getTimelineData(fileId, 1, 50);
+        } catch (Exception e) {
+          log.warn("Failed to fetch timeline bins for custom prompt retry: {}", e.getMessage());
+        }
+        try {
+          String phase1Json = llmClient.generateCompletion(
+              buildHypothesisSystemPrompt(),
+              buildHypothesisUserPrompt(file, analysis, additionalContext, aggregates, findings,
+                  timelineBinsCustom, maxFindings, maxRiskMatrix));
+          var phase1 = parseHypothesesAndQueries(phase1Json);
+          investigationSteps =
+              investigationService.executeQueries(fileId, phase1.queries(), phase1.hypotheses());
+          log.info("Custom prompt retry investigation complete: {} steps", investigationSteps.size());
+        } catch (Exception e) {
+          log.warn("Investigation phase skipped during custom prompt retry: {}", e.getMessage());
+        }
+
         String storyContent = llmClient.generateCompletion(buildSystemPrompt(), customPrompt);
         StoryResponse storyResponse = parseStoryContent(storyContent, storyId, fileId);
         storyResponse.setAggregates(aggregates);
         storyResponse.setFindings(findings);
+        storyResponse.setInvestigationSteps(investigationSteps.isEmpty() ? null : investigationSteps);
         StoryEntity story = StoryEntity.builder()
             .id(storyId).fileId(fileId).generatedAt(generatedAt)
             .status(StoryEntity.StoryStatus.COMPLETED)

frontend/src/features/network/hooks/useCompareData.ts

@@ -1,6 +1,6 @@
-import { useState, useEffect } from 'react';
+import { useState, useEffect, useMemo } from 'react';
 import { conversationService } from '@/features/conversation/services/conversationService';
-import { networkService } from '../services/networkService';
+import { networkService, selectSignificantNodes } from '../services/networkService';
 import { mergeGraphs } from '../services/mergeGraphs';
 import type { GraphNode, GraphEdge } from '../types';
 
@@ -21,6 +21,8 @@ export interface FileStats {
 export interface UseCompareDataReturn {
   mergedNodes: GraphNode[];
   mergedEdges: GraphEdge[];
+  totalNodes: number;
+  hiddenNodes: number;
   perFileStats: FileStats[];
   labels: string[];
   loading: boolean;
@@ -55,13 +57,18 @@ async function fetchGraphForFile(fileId: string) {
     response.data,
     undefined,
     MAX_CONVERSATIONS,
-    hostClassifications
+    hostClassifications,
+    0 // no per-file limit — significance filtering is applied after merge
   );
 }
 
-export function useCompareData(fileIds: string[], labels: string[]): UseCompareDataReturn {
-  const [mergedNodes, setMergedNodes] = useState<GraphNode[]>([]);
-  const [mergedEdges, setMergedEdges] = useState<GraphEdge[]>([]);
+export function useCompareData(
+  fileIds: string[],
+  labels: string[],
+  nodeLimit: number
+): UseCompareDataReturn {
+  const [allMergedNodes, setAllMergedNodes] = useState<GraphNode[]>([]);
+  const [allMergedEdges, setAllMergedEdges] = useState<GraphEdge[]>([]);
   const [perFileStats, setPerFileStats] = useState<FileStats[]>([]);
   const [loading, setLoading] = useState(true);
   const [error, setError] = useState<string | null>(null);
@@ -87,8 +94,8 @@ export function useCompareData(fileIds: string[], labels: string[]): UseCompareD
 
         const merged = mergeGraphs(graphs, labels);
 
-        setMergedNodes(merged.nodes);
-        setMergedEdges(merged.edges);
+        setAllMergedNodes(merged.nodes);
+        setAllMergedEdges(merged.edges);
         setPerFileStats(
           graphs.map((g, i) => ({
             label: labels[i],
@@ -115,5 +122,27 @@ export function useCompareData(fileIds: string[], labels: string[]): UseCompareD
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [fileIdsKey, labelsKey]);
 
-  return { mergedNodes, mergedEdges, perFileStats, labels, loading, error };
+  const { mergedNodes, mergedEdges, hiddenNodes } = useMemo(() => {
+    const { significantNodes, hiddenCount } = selectSignificantNodes(
+      allMergedNodes,
+      allMergedEdges,
+      nodeLimit
+    );
+    const sigIds = new Set(significantNodes.map(n => n.id));
+    const visibleEdges = allMergedEdges.filter(
+      e => sigIds.has(e.source) && sigIds.has(e.target)
+    );
+    return { mergedNodes: significantNodes, mergedEdges: visibleEdges, hiddenNodes: hiddenCount };
+  }, [allMergedNodes, allMergedEdges, nodeLimit]);
+
+  return {
+    mergedNodes,
+    mergedEdges,
+    totalNodes: allMergedNodes.length,
+    hiddenNodes,
+    perFileStats,
+    labels,
+    loading,
+    error,
+  };
 }

frontend/src/features/network/services/networkService.ts

@@ -243,7 +243,7 @@ function finalizeNodeRole(node: GraphNode, srcPort: number, dstPort: number) {
  *
  * Returns the selected nodes and the count of nodes that were hidden.
  */
-function selectSignificantNodes(
+export function selectSignificantNodes(
   nodes: GraphNode[],
   edges: GraphEdge[],
   limit: number

frontend/src/features/report/captureNetworkDiagrams.ts

@@ -26,17 +26,8 @@
 import { createElement } from 'react';
 import { createRoot } from 'react-dom/client';
 import { toPng } from 'html-to-image';
-import { conversationService } from '@/features/conversation/services/conversationService';
-import { networkService } from '@/features/network/services/networkService';
 import { NetworkGraph } from '@/components/network/NetworkGraph/NetworkGraph';
-import { CONVERSATION_LIMIT_ENABLED } from '@/features/network/hooks/useNetworkData';
 import type { GraphNode, GraphEdge } from '@/features/network/types';
-import type { ConversationFilters } from '@/features/conversation/types';
-import type { AnalysisSummary } from '@/types';
-
-// Match the same conversation cap the Network Diagram page uses so the report
-// shows identical edges.  If the env flag disables the limit, capture all.
-const MAX_CONVERSATIONS = CONVERSATION_LIMIT_ENABLED ? 500 : Infinity;
 const CAPTURE_W = 1400;
 const CAPTURE_H = 860;
 
@@ -161,46 +152,15 @@ export interface DiagramImages {
   hierarchical: string; // base64 PNG
 }
 
+/**
+ * Captures both ELK layouts for the given pre-filtered nodes and edges.
+ * The caller is responsible for passing exactly the nodes/edges currently
+ * visible on screen so the report matches what the user sees.
+ */
 export async function captureNetworkDiagrams(
-  fileId: string,
-  analysisSummary?: AnalysisSummary,
-  sessionFilters?: Partial<ConversationFilters>
+  nodes: GraphNode[],
+  edges: GraphEdge[]
 ): Promise<DiagramImages> {
-  const response = await conversationService.getConversations(fileId, {
-    ip: '',
-    port: '',
-    payloadContains: '',
-    protocols: [],
-    l7Protocols: [],
-    apps: [],
-    categories: [],
-    hasRisks: false,
-    fileTypes: [],
-    riskTypes: [],
-    customSignatures: [],
-    deviceTypes: [],
-    countries: [],
-    ...sessionFilters,
-    sortBy: '',
-    sortDir: 'asc',
-    page: 1,
-    pageSize: 10000,
-  });
-
-  let hostClassifications;
-  try {
-    hostClassifications = await conversationService.getHostClassifications(fileId);
-  } catch {
-    /* optional — best effort */
-  }
-
-  const { nodes, edges } = networkService.buildNetworkGraph(
-    response.data,
-    analysisSummary,
-    MAX_CONVERSATIONS,
-    hostClassifications
-  );
-
   // Sequential — both layouts share the module-level ELK singleton inside
   // NetworkGraph.tsx; running them in parallel risks a layout race condition.
   const forceDirected = await captureLayout(nodes, edges, 'forceDirected2d');