fix: Fix tasks with the base Context class

crates/fluxqueue-worker/src/task.rs

@@ -381,19 +381,22 @@ fn get_registry(module_path: &str, queue_name: &str) -> Result<TasksAndContexts>
             )
             .collect::<Result<HashMap<_, _>, _>>()?;
 
-        let contexts: HashMap<Arc<String>, Arc<Py<PyAny>>> = registry
+        let mut contexts: HashMap<Arc<String>, Arc<Py<PyAny>>> = registry
             .get_item("contexts")?
             .expect("contexts missing")
             .cast::<PyDict>()
             .map_err(|e| anyhow!("contexts is not a dict: {}", e))?
             .iter()
             .filter_map(|(key, value): (Bound<PyAny>, Bound<PyAny>)| {
                 let name: String = key.extract().ok()?;
-                let func: Py<PyAny> = value.unbind();
-                Some((Arc::new(name), Arc::new(func)))
+                let class: Py<PyAny> = value.unbind();
+                Some((Arc::new(name), Arc::new(class)))
             })
             .collect();
 
+        let (base_name, base_class) = get_base_context_class(py)?;
+        contexts.insert(base_name, base_class);
+
         Ok((tasks, contexts))
     })?;
 
@@ -416,6 +419,15 @@ fn get_task_metadata(py: Python<'_>, task: Arc<Task>) -> Result<Py<PyAny>> {
     Ok(task_metadata)
 }
 
+fn get_base_context_class(py: Python<'_>) -> Result<(Arc<String>, Arc<Py<PyAny>>)> {
+    let module = py.import("fluxqueue.context")?;
+    let context_class = module.getattr("Context")?;
+    Ok((
+        Arc::new("_Context".to_string()),
+        Arc::new(context_class.unbind()),
+    ))
+}
+
 fn is_coroutine(py: Python<'_>, func: Arc<Py<PyAny>>) -> Result<bool> {
     let inspect = py.import("inspect")?;
     let is_coro: bool = inspect

python/fluxqueue/context.py

@@ -29,7 +29,7 @@ class Context:
     to provide domain-specific resources.
     """
 
-    __fluxqueue_context__: str | None = None
+    __fluxqueue_context__: str | None = "_Context"
 
     def __init__(self) -> None:
         self._thread_local = threading.local()
@@ -73,7 +73,10 @@ async def _run_async_task(
             self._metadata_var.reset(token)
 
     def __init_subclass__(cls) -> None:
-        if not cls.__fluxqueue_context__:
+        if cls.__name__ == "_Context":
+            raise ValueError("Context subclass cannot be named '_Context'")
+
+        if not cls.__fluxqueue_context__ or cls.__fluxqueue_context__ == "_Context":
             cls.__fluxqueue_context__ = cls.__name__