From 0d07349de04dcf7aa3362d818d1410403925418f Mon Sep 17 00:00:00 2001
From: nessi <sascha@nesterovic.cc>
Date: Thu, 9 Jul 2026 21:19:26 +0200
Subject: [PATCH] feat: add traffic flow deduplication with 24h retention and
 update-in-place for existing flows

Add traffic_flow_key helper to generate unique flow identifier from node/IPs/protocol/ports/decision, implement 24-hour retention cutoff to delete old flows instead of all flows on heartbeat, build existing_flows lookup map from database with composite key matching, update agent_heartbeat to check for existing flows and update bytes/packets/state/observed_at/raw in-place instead of creating duplicates, extend
---
 backend/app/api/v1/router.py | 69 +++++++++++++++++++++++++++---------
 1 file changed, 53 insertions(+), 16 deletions(-)

diff --git a/backend/app/api/v1/router.py b/backend/app/api/v1/router.py
index b2955df..7ec6046 100644
--- a/backend/app/api/v1/router.py
+++ b/backend/app/api/v1/router.py
@@ -536,6 +536,19 @@ def raw_flow_decision(flow: TrafficFlow) -> str | None:
     return None
 
 
+def traffic_flow_key(node_id: str, raw_flow: dict[str, object]) -> tuple[object, ...]:
+    decision = str(raw_flow.get("decision") or raw_flow.get("state") or "observed").lower()
+    return (
+        node_id,
+        str(raw_flow.get("source_ip") or ""),
+        str(raw_flow.get("destination_ip") or ""),
+        str(raw_flow.get("protocol") or "unknown"),
+        flow_int(raw_flow.get("source_port"), 0) or None,
+        flow_int(raw_flow.get("destination_port"), 0) or None,
+        decision,
+    )
+
+
 def policy_read_payload(policy: Policy, deployment_status: dict[str, object] | None = None) -> dict[str, object]:
     return {
         "id": policy.id,
@@ -1357,29 +1370,53 @@ def agent_heartbeat(payload: AgentHeartbeat, authorization: str | None = Header(
     agent.version = payload.version
     agent.last_seen_at = datetime.utcnow()
     agent.last_payload = payload.model_dump(mode="json")
-    for old_flow in db.scalars(select(TrafficFlow).where(TrafficFlow.node_id == node.id)).all():
+    retention_cutoff = datetime.utcnow() - timedelta(hours=24)
+    for old_flow in db.scalars(select(TrafficFlow).where(TrafficFlow.node_id == node.id, TrafficFlow.updated_at < retention_cutoff)).all():
         db.delete(old_flow)
+    existing_flows = {
+        (
+            flow.node_id,
+            flow.source_ip,
+            flow.destination_ip,
+            flow.protocol,
+            flow.source_port,
+            flow.destination_port,
+            str(flow.state or "observed").lower(),
+        ): flow
+        for flow in db.scalars(select(TrafficFlow).where(TrafficFlow.node_id == node.id)).all()
+    }
     for raw_flow in payload.flows[:1000]:
         source_ip = str(raw_flow.get("source_ip") or "")
         destination_ip = str(raw_flow.get("destination_ip") or "")
         if not source_ip or not destination_ip:
             continue
         observed_at = payload.collected_at or datetime.utcnow()
-        db.add(
-            TrafficFlow(
-                node_id=node.id,
-                source_ip=source_ip,
-                destination_ip=destination_ip,
-                protocol=str(raw_flow.get("protocol") or "unknown"),
-                source_port=flow_int(raw_flow.get("source_port"), 0) or None,
-                destination_port=flow_int(raw_flow.get("destination_port"), 0) or None,
-                bytes=flow_int(raw_flow.get("bytes")),
-                packets=flow_int(raw_flow.get("packets")),
-                state=str(raw_flow.get("decision") or raw_flow.get("state") or "") or None,
-                observed_at=observed_at.replace(tzinfo=None) if observed_at.tzinfo else observed_at,
-                raw=raw_flow,
-            )
+        observed_at = observed_at.replace(tzinfo=None) if observed_at.tzinfo else observed_at
+        key = traffic_flow_key(node.id, raw_flow)
+        existing = existing_flows.get(key)
+        if existing:
+            existing.bytes = flow_int(raw_flow.get("bytes"))
+            existing.packets = flow_int(raw_flow.get("packets"))
+            existing.state = str(raw_flow.get("decision") or raw_flow.get("state") or "") or None
+            existing.observed_at = observed_at
+            existing.raw = raw_flow
+            existing.updated_at = datetime.utcnow()
+            continue
+        flow = TrafficFlow(
+            node_id=node.id,
+            source_ip=source_ip,
+            destination_ip=destination_ip,
+            protocol=str(raw_flow.get("protocol") or "unknown"),
+            source_port=flow_int(raw_flow.get("source_port"), 0) or None,
+            destination_port=flow_int(raw_flow.get("destination_port"), 0) or None,
+            bytes=flow_int(raw_flow.get("bytes")),
+            packets=flow_int(raw_flow.get("packets")),
+            state=str(raw_flow.get("decision") or raw_flow.get("state") or "") or None,
+            observed_at=observed_at,
+            raw=raw_flow,
         )
+        db.add(flow)
+        existing_flows[key] = flow
     commit_or_400(db)
     return {"status": "ok", "node_id": node.id}
 
@@ -1416,7 +1453,7 @@ async def workload_insights(workload_id: str, _: CurrentUser, db: Session = Depe
         flows = db.scalars(
             select(TrafficFlow)
             .where((TrafficFlow.source_ip.in_(workload_ips)) | (TrafficFlow.destination_ip.in_(workload_ips)))
-            .order_by(TrafficFlow.updated_at.desc())
+            .order_by((TrafficFlow.state == "blocked").desc(), TrafficFlow.updated_at.desc())
             .limit(50)
         ).all()
         for flow in flows: