Skip to content

Commit e887261

Browse files
itsskofficialitsskofficial
committed
feat (multi-tenancy): added dynamic polling for gmail
1 parent dd321cc commit e887261

File tree

8 files changed

+894
-344
lines changed

8 files changed

+894
-344
lines changed

‎src/client/app/layout.js‎

Lines changed: 39 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -24,16 +24,42 @@ export const metadata = {
2424
* @returns {React.ReactNode} - The RootLayout component UI.
2525
*/
2626
export default async function RootLayout({ children }) {
27-
return (
28-
<html lang="en" suppressHydrationWarning>
29-
{/* Root html element with language set to English and hydration warning suppressed */}
30-
<body className="bg-black">
31-
{/* Body element with a black background using global styles */}
32-
<Toaster position="bottom-right" />
33-
{/* Toaster component for displaying notifications, positioned at the bottom-right */}
34-
{children}
35-
{/* Render the child components, which is the main content of the application */}
36-
</body>
37-
</html>
38-
)
39-
}
27+
useEffect(() => {
28+
let intervalId;
29+
const sendHeartbeat = () => {
30+
if (document.hasFocus() && window.electron && typeof window.electron.sendUserActivityHeartbeat === 'function') {
31+
console.log("Client: Sending activity heartbeat...");
32+
window.electron.sendUserActivityHeartbeat().catch(err => console.error("Heartbeat IPC error:", err));
33+
}
34+
};
35+
36+
// Send heartbeat immediately on mount (if focused) and then every 5 minutes
37+
sendHeartbeat();
38+
intervalId = setInterval(sendHeartbeat, 5 * 60 * 1000); // 5 minutes
39+
40+
// Also send on window focus
41+
const handleFocus = () => {
42+
console.log("Client: Window focused, sending heartbeat.");
43+
sendHeartbeat();
44+
};
45+
window.addEventListener('focus', handleFocus);
46+
47+
return () => {
48+
clearInterval(intervalId);
49+
window.removeEventListener('focus', handleFocus);
50+
};
51+
}, []); // Empty dependency array means this runs once on mount and cleans up on unmount
52+
53+
return (
54+
<html lang="en" suppressHydrationWarning>
55+
{/* Root html element with language set to English and hydration warning suppressed */}
56+
<body className="bg-black">
57+
{/* Body element with a black background using global styles */}
58+
<Toaster position="bottom-right" />
59+
{/* Toaster component for displaying notifications, positioned at the bottom-right */}
60+
{children}
61+
{/* Render the child components, which is the main content of the application */}
62+
</body>
63+
</html>
64+
)
65+
}

‎src/client/main/index.js‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2920,5 +2920,81 @@ ipcMain.handle("fetch-memory-categories", async () => {
29202920
}
29212921
})
29222922

2923+
ipcMain.handle("user-activity-heartbeat", async () => {
2924+
const token = getAccessToken() // Ensure you have access to this function
2925+
if (!token) {
2926+
console.error("[IPC HEARTBEAT] No access token available.")
2927+
return { success: false, error: "Not authenticated" }
2928+
}
2929+
try {
2930+
const backendUrl = process.env.APP_SERVER_URL || "http://localhost:5000"
2931+
const response = await fetch(`${backendUrl}/users/activity/heartbeat`, {
2932+
method: "POST",
2933+
headers: {
2934+
"Content-Type": "application/json",
2935+
Authorization: `Bearer ${token}`
2936+
}
2937+
})
2938+
if (!response.ok) {
2939+
const errorText = await response.text()
2940+
console.error(
2941+
`[IPC HEARTBEAT] Backend error: ${response.status} - ${errorText}`
2942+
)
2943+
return {
2944+
success: false,
2945+
error: `Backend error: ${response.status}`
2946+
}
2947+
}
2948+
// console.log('[IPC HEARTBEAT] Heartbeat sent successfully.');
2949+
return { success: true }
2950+
} catch (error) {
2951+
console.error("[IPC HEARTBEAT] Error sending heartbeat:", error)
2952+
return { success: false, error: error.message }
2953+
}
2954+
})
2955+
2956+
ipcMain.handle("force-sync-service", async (event, serviceName) => {
2957+
const token = getAccessToken()
2958+
if (!token) {
2959+
console.error("[IPC FORCE_SYNC] No access token available.")
2960+
return { success: false, error: "Not authenticated" }
2961+
}
2962+
if (!serviceName) {
2963+
console.error("[IPC FORCE_SYNC] Service name not provided.")
2964+
return { success: false, error: "Service name required" }
2965+
}
2966+
try {
2967+
const backendUrl = process.env.APP_SERVER_URL || "http://localhost:5000"
2968+
const response = await fetch(
2969+
`${backendUrl}/users/force-sync/${serviceName}`,
2970+
{
2971+
method: "POST",
2972+
headers: {
2973+
"Content-Type": "application/json",
2974+
Authorization: `Bearer ${token}`
2975+
}
2976+
}
2977+
)
2978+
if (!response.ok) {
2979+
const errorText = await response.text()
2980+
console.error(
2981+
`[IPC FORCE_SYNC] Backend error for ${serviceName}: ${response.status} - ${errorText}`
2982+
)
2983+
return {
2984+
success: false,
2985+
error: `Backend error: ${response.status}`
2986+
}
2987+
}
2988+
console.log(`[IPC FORCE_SYNC] Force sync requested for ${serviceName}.`)
2989+
return { success: true, message: `Sync requested for ${serviceName}.` }
2990+
} catch (error) {
2991+
console.error(
2992+
`[IPC FORCE_SYNC] Error requesting force sync for ${serviceName}:`,
2993+
error
2994+
)
2995+
return { success: false, error: error.message }
2996+
}
2997+
})
2998+
29232999
// --- End of File ---
29243000
console.log("Electron main process script execution completed initial setup.")

‎src/client/main/preload.js‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -175,5 +175,9 @@ contextBridge.exposeInMainWorld("electron", {
175175
ipcRenderer.removeListener("update-progress", listener)
176176
console.log("Removed update-progress listener") // For debugging
177177
}
178-
}
178+
},
179+
sendUserActivityHeartbeat: () =>
180+
ipcRenderer.invoke("user-activity-heartbeat"),
181+
forceSyncService: (serviceName) =>
182+
ipcRenderer.invoke("force-sync-service", serviceName)
179183
})

‎src/server/app/app.py‎

Lines changed: 134 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,11 @@
1212
import multiprocessing
1313
import traceback # For detailed error printing
1414
import json # Keep json for websocket messages
15+
from server.context.base import BaseContextEngine, POLLING_INTERVALS # Import POLLING_INTERVALS
16+
from server.context.gmail import GmailContextEngine
17+
from server.context.gcalendar import GCalendarContextEngine
18+
from server.context.internet import InternetSearchContextEngine
19+
from datetime import timezone, timedelta # For scheduling
1520

1621
# --- MongoDB Manager Import ---
1722
from server.db import MongoManager
@@ -382,6 +387,13 @@ async def broadcast_json(self, data: dict):
382387
mongo_manager = MongoManager()
383388
print(f"[INIT] {datetime.datetime.now()}: MongoManager initialized.")
384389

390+
# Dictionary to hold active context engine instances per user
391+
# Structure: { "user_id1": {"gmail": GmailEngineInstance, "gcalendar": GCalendarEngineInstance}, ... }
392+
active_context_engines: Dict[str, Dict[str, BaseContextEngine]] = {}
393+
polling_scheduler_task_handle: Optional[asyncio.Task] = None
394+
395+
POLLING_SCHEDULER_INTERVAL_SECONDS = 30
396+
385397
# --- Helper Functions (User-Specific DB Operations) ---
386398
async def load_user_profile(user_id: str) -> Dict[str, Any]:
387399
profile = await mongo_manager.get_user_profile(user_id)
@@ -638,10 +650,110 @@ async def gdrive_tool(tool_call_input: dict) -> Dict[str, Any]:
638650

639651
agent_task_processor_instance: Optional[AgentTaskProcessor] = None
640652

653+
async def polling_scheduler_loop():
654+
print(f"[POLLING_SCHEDULER] Starting polling scheduler loop (interval: {POLLING_SCHEDULER_INTERVAL_SECONDS}s)")
655+
while True:
656+
try:
657+
print(f"[POLLING_SCHEDULER] Checking for due polling tasks at {datetime.now(timezone.utc).isoformat()}")
658+
due_tasks_states = await mongo_manager.get_due_polling_tasks()
659+
660+
if not due_tasks_states:
661+
# print(f"[POLLING_SCHEDULER] No tasks due at this time.")
662+
pass
663+
else:
664+
print(f"[POLLING_SCHEDULER] Found {len(due_tasks_states)} due tasks.")
665+
666+
for task_state in due_tasks_states:
667+
user_id = task_state["user_id"]
668+
engine_category = task_state["engine_category"]
669+
670+
print(f"[POLLING_SCHEDULER] Attempting to process task for {user_id}/{engine_category}")
671+
672+
# Try to acquire lock and get the task. If successful, proceed.
673+
locked_task_state = await mongo_manager.set_polling_status_and_get(user_id, engine_category)
674+
675+
if locked_task_state:
676+
print(f"[POLLING_SCHEDULER] Acquired lock for {user_id}/{engine_category}. Triggering poll.")
677+
engine_instance = active_context_engines.get(user_id, {}).get(engine_category)
678+
679+
if not engine_instance:
680+
engine_class = DATA_SOURCES_CONFIG.get(engine_category, {}).get("engine_class")
681+
if engine_class:
682+
print(f"[POLLING_SCHEDULER] Creating new instance for {user_id}/{engine_category}")
683+
engine_instance = engine_class(
684+
user_id=user_id,
685+
task_queue=task_queue, # global
686+
memory_backend=memory_backend, # global
687+
websocket_manager=manager, # global websocket_manager
688+
mongo_manager_instance=mongo_manager # global
689+
)
690+
if user_id not in active_context_engines:
691+
active_context_engines[user_id] = {}
692+
active_context_engines[user_id][engine_category] = engine_instance
693+
# No need to call engine_instance.initialize_polling_state() here,
694+
# run_poll_cycle will handle it if state is missing or needs reset.
695+
else:
696+
print(f"[POLLING_SCHEDULER_ERROR] No engine class configured for category: {engine_category}")
697+
# Release lock if instance can't be created
698+
await mongo_manager.update_polling_state(user_id, engine_category, {"is_currently_polling": False})
699+
continue
700+
701+
# Run the poll cycle in a new task so the scheduler doesn't block
702+
asyncio.create_task(engine_instance.run_poll_cycle())
703+
else:
704+
print(f"[POLLING_SCHEDULER] Could not acquire lock for {user_id}/{engine_category} (already processing or not due).")
705+
706+
except Exception as e:
707+
print(f"[POLLING_SCHEDULER_ERROR] Error in scheduler loop: {e}")
708+
traceback.print_exc()
709+
710+
await asyncio.sleep(POLLING_SCHEDULER_INTERVAL_SECONDS)
711+
712+
713+
async def start_user_context_engines(user_id: str):
714+
"""Starts context engines for a given user if not already running and initializes polling state."""
715+
if user_id not in active_context_engines:
716+
active_context_engines[user_id] = {}
717+
718+
user_profile = await load_user_profile(user_id) # load_user_profile is your existing helper
719+
user_settings = user_profile.get("userData", {})
720+
721+
for source_name, config in DATA_SOURCES_CONFIG.items():
722+
is_enabled = user_settings.get(f"{source_name}Enabled", config["enabled_by_default"])
723+
724+
if is_enabled:
725+
if source_name not in active_context_engines[user_id]:
726+
print(f"[CONTEXT_ENGINE_MGR] Starting {source_name} engine for user {user_id}...")
727+
engine_class = config["engine_class"]
728+
engine_instance = engine_class(
729+
user_id=user_id,
730+
task_queue=task_queue,
731+
memory_backend=memory_backend,
732+
websocket_manager=manager,
733+
mongo_manager_instance=mongo_manager
734+
)
735+
active_context_engines[user_id][source_name] = engine_instance
736+
# Initialize polling state (will set next_poll_at to now if new)
737+
await engine_instance.initialize_polling_state()
738+
print(f"[CONTEXT_ENGINE_MGR] {source_name} engine started and polling state initialized for user {user_id}.")
739+
else:
740+
# Engine already active, ensure its polling state is initialized (e.g. if server restarted)
741+
await active_context_engines[user_id][source_name].initialize_polling_state()
742+
print(f"[CONTEXT_ENGINE_MGR] {source_name} engine for user {user_id} already active. Ensured polling state.")
743+
744+
elif not is_enabled and source_name in active_context_engines[user_id]:
745+
print(f"[CONTEXT_ENGINE_MGR] {source_name} engine for user {user_id} is disabled. Stopping (if implemented) and removing.")
746+
# Add engine stop logic if available:
747+
# if hasattr(active_context_engines[user_id][source_name], 'stop'):
748+
# await active_context_engines[user_id][source_name].stop()
749+
del active_context_engines[user_id][source_name]
750+
751+
641752
@app.on_event("startup")
642753
async def startup_event():
754+
# ... (your existing startup code: STT, TTS, DB init, agent_task_processor, memory_backend) ...
643755
print(f"[FASTAPI_LIFECYCLE] App startup.")
644-
global stt_model, tts_model, agent_task_processor_instance
756+
global stt_model, tts_model, agent_task_processor_instance, polling_scheduler_task_handle
645757

646758
agent_task_processor_instance = agent_task_processor
647759

@@ -665,15 +777,15 @@ async def startup_event():
665777
elif TTS_PROVIDER == "ELEVENLABS":
666778
tts_model = ElevenLabsTTS()
667779
print("[LIFECYCLE] ElevenLabs TTS loaded (Production Mode).")
668-
else: # Default to ORPHEUS
780+
else:
669781
tts_model = OrpheusTTS(verbose=False, default_voice_id=SELECTED_TTS_VOICE)
670782
print(f"[WARN] Unknown TTS_PROVIDER '{TTS_PROVIDER}'. Defaulting to Orpheus TTS.")
671783
except Exception as e:
672784
print(f"[ERROR] TTS model failed to load. Voice features will be unavailable. Details: {e}")
673785
tts_model = None
674786

675787
await task_queue.initialize_db()
676-
await mongo_manager.initialize_db()
788+
await mongo_manager.initialize_db() # This now includes polling state indexes
677789
await memory_backend.memory_queue.initialize_db()
678790

679791
if agent_task_processor_instance:
@@ -684,18 +796,35 @@ async def startup_event():
684796

685797
asyncio.create_task(memory_backend.process_memory_operations())
686798

799+
# Start the central polling scheduler
800+
polling_scheduler_task_handle = asyncio.create_task(polling_scheduler_loop())
801+
print(f"[FASTAPI_LIFECYCLE] Central polling scheduler started.")
802+
687803
print(f"[FASTAPI_LIFECYCLE] App startup complete.")
688804

805+
689806
@app.on_event("shutdown")
690807
async def shutdown_event():
808+
# ... (your existing shutdown code) ...
809+
global polling_scheduler_task_handle
810+
if polling_scheduler_task_handle and not polling_scheduler_task_handle.done():
811+
print("[FASTAPI_LIFECYCLE] Cancelling polling scheduler task...")
812+
polling_scheduler_task_handle.cancel()
813+
try:
814+
await polling_scheduler_task_handle
815+
except asyncio.CancelledError:
816+
print("[FASTAPI_LIFECYCLE] Polling scheduler task cancelled.")
817+
except Exception as e:
818+
print(f"[FASTAPI_LIFECYCLE_ERROR] Error during polling scheduler shutdown: {e}")
819+
# ... (rest of your shutdown) ...
691820
print(f"[FASTAPI_LIFECYCLE] App shutdown.")
692821
if mongo_manager.client:
693822
mongo_manager.client.close()
694823
print("[FASTAPI_LIFECYCLE] MongoManager client closed.")
695-
if task_queue.client: # TaskQueue now has its own client
824+
if task_queue.client:
696825
task_queue.client.close()
697826
print("[FASTAPI_LIFECYCLE] TaskQueue MongoDB client closed.")
698-
if memory_backend.memory_queue.client: # MemoryQueue now has its own client
827+
if memory_backend.memory_queue.client:
699828
memory_backend.memory_queue.client.close()
700829
print("[FASTAPI_LIFECYCLE] MemoryQueue MongoDB client closed.")
701830
print(f"[FASTAPI_LIFECYCLE] All MongoDB clients known to app.py closed.")

0 commit comments

Comments
 (0)