51 lines
2.0 KiB
Python
51 lines
2.0 KiB
Python
"""
|
|
Request context management for ComfyUI Copilot
|
|
Uses contextvars to provide request-scoped context variables with async safety
|
|
"""
|
|
|
|
import contextvars
|
|
from typing import Optional, Dict, Any
|
|
|
|
# Define context variables
|
|
_session_id: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar('session_id', default=None)
|
|
_workflow_checkpoint_id: contextvars.ContextVar[Optional[int]] = contextvars.ContextVar('workflow_checkpoint_id', default=None)
|
|
_config: contextvars.ContextVar[Optional[Dict[str, Any]]] = contextvars.ContextVar('config', default=None)
|
|
|
|
def set_session_id(session_id: str) -> None:
|
|
"""Set the session ID for the current request context"""
|
|
_session_id.set(session_id)
|
|
|
|
def get_session_id() -> Optional[str]:
|
|
"""Get the session ID from the current request context"""
|
|
return _session_id.get()
|
|
|
|
def set_workflow_checkpoint_id(checkpoint_id: Optional[int]) -> None:
|
|
"""Set the workflow checkpoint ID for the current request context"""
|
|
_workflow_checkpoint_id.set(checkpoint_id)
|
|
|
|
def get_workflow_checkpoint_id() -> Optional[int]:
|
|
"""Get the workflow checkpoint ID from the current request context"""
|
|
return _workflow_checkpoint_id.get()
|
|
|
|
def set_config(config: Dict[str, Any]) -> None:
|
|
"""Set the request config for the current request context"""
|
|
_config.set(config)
|
|
|
|
def get_config() -> Optional[Dict[str, Any]]:
|
|
"""Get the request config from the current request context"""
|
|
return _config.get()
|
|
|
|
def set_request_context(session_id: str, workflow_checkpoint_id: Optional[int] = None, config: Optional[Dict[str, Any]] = None) -> None:
|
|
"""Set all request context variables at once"""
|
|
set_session_id(session_id)
|
|
if workflow_checkpoint_id is not None:
|
|
set_workflow_checkpoint_id(workflow_checkpoint_id)
|
|
if config is not None:
|
|
set_config(config)
|
|
|
|
def clear_request_context() -> None:
|
|
"""Clear all request context variables"""
|
|
_session_id.set(None)
|
|
_workflow_checkpoint_id.set(None)
|
|
_config.set(None)
|