From 9e8e5b7297756a9ffd384d206e39e8bde3d980bf Mon Sep 17 00:00:00 2001
From: viettranx <viettranx@gmail.com>
Date: Tue, 10 Mar 2026 20:58:24 +0700
Subject: [PATCH] feat(kg): session-based scope picker, dark mode support,
 inline agent select

Replace manual user_id text input with select dropdown populated from
sessions (DM + group chats). Add colorMode prop to ReactFlow for dark
mode Controls/MiniMap. Add agent select in empty state.
---
 .../knowledge-graph/knowledge-graph-page.tsx  | 54 +++++++++++++++++--
 ui/web/src/pages/memory/kg-graph-view.tsx     |  5 ++
 2 files changed, 54 insertions(+), 5 deletions(-)

diff --git a/ui/web/src/pages/knowledge-graph/knowledge-graph-page.tsx b/ui/web/src/pages/knowledge-graph/knowledge-graph-page.tsx
index 4a71f8c4..95dfb514 100644
--- a/ui/web/src/pages/knowledge-graph/knowledge-graph-page.tsx
+++ b/ui/web/src/pages/knowledge-graph/knowledge-graph-page.tsx
@@ -1,8 +1,10 @@
-import { useState } from "react";
+import { useState, useMemo } from "react";
 import { useTranslation } from "react-i18next";
 import { Network } from "lucide-react";
 import { EmptyState } from "@/components/shared/empty-state";
 import { useAgents } from "@/pages/agents/hooks/use-agents";
+import { useSessions } from "@/pages/sessions/hooks/use-sessions";
+import { parseSessionKey } from "@/lib/session-key";
 import { KGEntitiesTab } from "@/pages/memory/kg-entities-tab";
 
 export function KnowledgeGraphPage() {
@@ -11,6 +13,30 @@ export function KnowledgeGraphPage() {
   const [agentId, setAgentId] = useState("");
   const [userIdFilter, setUserIdFilter] = useState("");
 
+  // Resolve agent UUID → agent_key (sessions filter uses agent_key in session key pattern)
+  const selectedAgent = agents.find((a) => a.id === agentId);
+  const agentKey = selectedAgent?.agent_key ?? "";
+
+  // Fetch sessions for selected agent to build scope picker (DM + group chats)
+  const { sessions } = useSessions({ agentFilter: agentKey || undefined, limit: 200 });
+
+  // Dedupe sessions by userID → scope options showing chat title / display name
+  const scopeOptions = useMemo(() => {
+    const seen = new Map<string, string>();
+    for (const s of sessions) {
+      const uid = s.userID || parseSessionKey(s.key).scope;
+      if (!uid || seen.has(uid)) continue;
+      const meta = s.metadata;
+      const label = meta?.chat_title || meta?.display_name
+        || (meta?.username ? `@${meta.username}` : null)
+        || uid;
+      seen.set(uid, label);
+    }
+    return Array.from(seen.entries())
+      .map(([value, label]) => ({ value, label }))
+      .sort((a, b) => a.label.localeCompare(b.label));
+  }, [sessions]);
+
   return (
     <div className="flex h-full flex-col p-4 sm:p-6">
       {/* Header + filters in one row */}
@@ -33,13 +59,17 @@ export function KnowledgeGraphPage() {
           ))}
         </select>
         {agentId && (
-          <input
+          <select
             id="kg-scope"
             value={userIdFilter}
             onChange={(e) => setUserIdFilter(e.target.value)}
-            placeholder={t("filters.allScope")}
-            className="h-8 rounded-md border bg-background px-2 text-sm w-[160px]"
-          />
+            className="h-8 rounded-md border bg-background px-2 text-sm max-w-[240px]"
+          >
+            <option value="">{t("filters.allScope")}</option>
+            {scopeOptions.map((o) => (
+              <option key={o.value} value={o.value}>{o.label}</option>
+            ))}
+          </select>
         )}
       </div>
 
@@ -50,6 +80,20 @@ export function KnowledgeGraphPage() {
             icon={Network}
             title={t("kg.selectAgentTitle")}
             description={t("kg.selectAgentDescription")}
+            action={
+              <select
+                value={agentId}
+                onChange={(e) => { setAgentId(e.target.value); setUserIdFilter(""); }}
+                className="mt-2 h-9 rounded-md border bg-background px-3 text-sm"
+              >
+                <option value="">{t("filters.selectAgent")}</option>
+                {agents.map((a) => (
+                  <option key={a.id} value={a.id}>
+                    {a.display_name || a.agent_key}
+                  </option>
+                ))}
+              </select>
+            }
           />
         ) : (
           <KGEntitiesTab agentId={agentId} userId={userIdFilter || undefined} />
diff --git a/ui/web/src/pages/memory/kg-graph-view.tsx b/ui/web/src/pages/memory/kg-graph-view.tsx
index b6799cdc..3dcb40dc 100644
--- a/ui/web/src/pages/memory/kg-graph-view.tsx
+++ b/ui/web/src/pages/memory/kg-graph-view.tsx
@@ -10,12 +10,14 @@ import {
   MiniMap,
   type Node,
   type Edge,
+  type ColorMode,
   Handle,
   Position,
 } from "@xyflow/react";
 import { forceSimulation, forceLink, forceManyBody, forceCenter, forceCollide, forceX, forceY, type SimulationNodeDatum } from "d3-force";
 import "@xyflow/react/dist/style.css";
 import { useTranslation } from "react-i18next";
+import { useUiStore } from "@/stores/use-ui-store";
 import type { KGEntity, KGRelation } from "@/types/knowledge-graph";
 
 // Color mapping for entity types
@@ -141,6 +143,8 @@ export function KGGraphView(props: KGGraphViewProps) {
 function KGGraphViewInner({ entities, relations, onEntityClick }: KGGraphViewProps) {
   const { t } = useTranslation("memory");
   const { fitView } = useReactFlow();
+  const theme = useUiStore((s) => s.theme);
+  const colorMode: ColorMode = theme === "system" ? "system" : theme;
   const { nodes: initialNodes, edges: initialEdges } = useMemo(
     () => buildGraph(entities, relations),
     [entities, relations],
@@ -190,6 +194,7 @@ function KGGraphViewInner({ entities, relations, onEntityClick }: KGGraphViewPro
           onEdgesChange={onEdgesChange}
           onNodeClick={handleNodeClick}
           nodeTypes={nodeTypes}
+          colorMode={colorMode}
           fitView
           minZoom={0.1}
           maxZoom={3}