@@ -122,55 +122,53 @@ def run_async(coro):
122122 loop .close ()
123123 asyncio .set_event_loop (None )
124124
125- @celery_app .task (name = "cud_memory_task" )
126- def cud_memory_task (user_id : str , information : str , source : Optional [str ] = None ):
127- """
128- Celery task wrapper for the CUD (Create, Update, Delete) memory operation.
129- This runs the core memory management logic asynchronously.
130- """
131- logger .info (f"Celery worker received cud_memory_task for user_id: { user_id } )
125+ async def async_cud_memory_task (user_id : str , information : str , source : Optional [str ] = None ):
126+ """The async logic for the CUD memory task."""
132127 db_manager = MongoManager ()
128+ username = user_id # Default fallback
133129 try :
134130 # --- Enforce Memory Limit ---
135- user_profile = run_async ( db_manager .get_user_profile (user_id ) )
131+ user_profile = await db_manager .get_user_profile (user_id )
136132 plan = user_profile .get ("userData" , {}).get ("plan" , "free" ) if user_profile else "free"
137133 limit = PLAN_LIMITS [plan ].get ("memories_total" , 0 )
138134
139135 if limit != float ('inf' ):
140136 from mcp_hub .memory import db as memory_db
141- pool = run_async (memory_db .get_db_pool ())
142- async def count_facts ():
143- async with pool .acquire () as conn :
144- return await conn .fetchval ("SELECT COUNT(*) FROM facts WHERE user_id = $1" , user_id )
145-
146- current_count = run_async (count_facts ())
147-
137+ pool = await memory_db .get_db_pool ()
138+ async with pool .acquire () as conn :
139+ current_count = await conn .fetchval ("SELECT COUNT(*) FROM facts WHERE user_id = $1" , user_id )
148140 if current_count >= limit :
149141 logger .warning (f"User { user_id } { plan } { limit } )
150- # Optionally notify the user
151- run_async (notify_user (user_id , f"You've reached your memory limit of { limit } ))
142+ await notify_user (user_id , f"You've reached your memory limit of { limit } )
152143 return
153- except Exception as e :
154- logger .error (f"Error checking memory limit for user { user_id } { e } , exc_info = True )
155- finally :
156- run_async (db_manager .close ())
157144
158- # --- NEW: Fetch user's name before calling cud_memory ---
159- username = user_id # Default fallback
160- try :
161- user_profile = run_async (db_manager .get_user_profile (user_id ))
145+ # --- Fetch user's name before calling cud_memory ---
162146 if user_profile :
147+ # Use the name from personalInfo, which is set during onboarding and can be updated in settings.
163148 username = user_profile .get ("userData" , {}).get ("personalInfo" , {}).get ("name" , user_id )
149+
164150 except Exception as e :
165- logger .error (f"Failed to fetch user profile for { user_id } { e } )
151+ logger .error (f"Error during pre-CUD setup for user { user_id } { e } , exc_info = True )
152+ # We can still proceed with the CUD operation, just using the user_id as the name.
166153 finally :
167- run_async (db_manager .close ())
168- # --- END NEW ---
154+ await db_manager .close ()
169155
156+ # Initialize models required for the CUD operation
170157 initialize_embedding_model ()
171158 initialize_agents ()
172159 # Pass the fetched username to the cud_memory function
173- run_async (cud_memory (user_id , information , source , username ))
160+ await cud_memory (user_id , information , source , username )
161+
162+ @celery_app .task (name = "cud_memory_task" )
163+ def cud_memory_task (user_id : str , information : str , source : Optional [str ] = None ):
164+ """
165+ Celery task wrapper for the CUD (Create, Update, Delete) memory operation.
166+ This runs the core memory management logic asynchronously.
167+ """
168+ logger .info (f"Celery worker received cud_memory_task for user_id: { user_id } )
169+ # This single call to run_async wraps the entire asynchronous logic,
170+ # ensuring the event loop and DB connections are managed correctly for the task's lifecycle.
171+ run_async (async_cud_memory_task (user_id , information , source ))
174172
175173@celery_app .task (name = "orchestrate_swarm_task" )
176174def orchestrate_swarm_task (task_id : str , user_id : str ):
0 commit comments