| 1 | """Shared context helper used by both ApiHandler and WsHandler.""" |
| 2 | |
| 3 | import threading |
| 4 | from typing import Union |
| 5 | |
| 6 | ThreadLockType = Union[threading.Lock, threading.RLock] |
| 7 | |
| 8 | |
| 9 | def use_context(lock: ThreadLockType, ctxid: str, create_if_not_exists: bool = True): |
| 10 | from agent import AgentContext |
| 11 | from helpers import projects |
| 12 | from initialize import initialize_agent |
| 13 | |
| 14 | with lock: |
| 15 | if not ctxid: |
| 16 | first = AgentContext.first() |
| 17 | if first: |
| 18 | AgentContext.use(first.id) |
| 19 | return first |
| 20 | context = AgentContext(config=initialize_agent(), set_current=True) |
| 21 | projects.reconcile_agent_profile( |
| 22 | context, projects.get_context_project_name(context) |
| 23 | ) |
| 24 | return context |
| 25 | got = AgentContext.use(ctxid) |
| 26 | if got: |
| 27 | return got |
| 28 | if create_if_not_exists: |
| 29 | context = AgentContext( |
| 30 | config=initialize_agent(), id=ctxid, set_current=True |
| 31 | ) |
| 32 | projects.reconcile_agent_profile( |
| 33 | context, projects.get_context_project_name(context) |
| 34 | ) |
| 35 | return context |
| 36 | raise Exception(f"Context {ctxid} not found") |