From 72f0b3908fb041dd59c605103322d9d0cb1750de Mon Sep 17 00:00:00 2001
From: Newman Gao <newmangao@outlook.com>
Date: Fri, 16 Oct 2026 02:10:14 +0800
Subject: [PATCH 1/2] Fix dynamic hook module registration

---
 agentlightning/hooks.py | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/agentlightning/hooks.py b/agentlightning/hooks.py
index fad1f57e9..35a5730cf 100644
--- a/agentlightning/hooks.py
+++ b/agentlightning/hooks.py
@@ -4,6 +4,8 @@
 
 from __future__ import annotations
 
+import sys
+from hashlib import sha256
 from typing import TYPE_CHECKING, Any, Protocol
 
 from agentlightning.schemas import RolloutCreate
@@ -43,10 +45,20 @@ def load_hooks(path: str) -> RolloutHooks:
     if not module_path.exists():
         raise FileNotFoundError(f"Hooks module not found: {module_path}")
 
-    spec = importlib.util.spec_from_file_location("_agl_hooks", str(module_path))
+    module_name = f"_agl_hooks_{sha256(str(module_path).encode()).hexdigest()[:16]}"
+    spec = importlib.util.spec_from_file_location(module_name, str(module_path))
     assert spec is not None and spec.loader is not None
     module = importlib.util.module_from_spec(spec)
-    spec.loader.exec_module(module)
+    previous_module = sys.modules.get(module_name)
+    sys.modules[module_name] = module
+    try:
+        spec.loader.exec_module(module)
+    except BaseException:
+        if previous_module is None:
+            sys.modules.pop(module_name, None)
+        else:
+            sys.modules[module_name] = previous_module
+        raise
 
     hook_classes = [
         obj

From ff513126cfb5447cfd44a4f57652f63ffcc2623d Mon Sep 17 00:00:00 2001
From: Newman Gao <newmangao@outlook.com>
Date: Fri, 16 Oct 2026 02:10:30 +0800
Subject: [PATCH 2/2] Test dynamically loaded hook module registration

---
 tests/test_hooks.py | 43 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 43 insertions(+)
 create mode 100644 tests/test_hooks.py

diff --git a/tests/test_hooks.py b/tests/test_hooks.py
new file mode 100644
index 000000000..70b4e532b
--- /dev/null
+++ b/tests/test_hooks.py
@@ -0,0 +1,43 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Regression coverage for dynamically loaded rollout hooks."""
+
+import pickle
+import sys
+
+from agentlightning.hooks import load_hooks
+from agentlightning.schemas import RolloutCreate
+
+
+def test_load_hooks_registers_dataclass_hook_module(tmp_path) -> None:
+    hooks_path = tmp_path / "custom_hooks.py"
+    hooks_path.write_text(
+        """
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from agentlightning.hooks import RolloutHooks
+from agentlightning.schemas import RolloutCreate
+
+
+@dataclass
+class HookConfig:
+    prefix: str = "checked:"
+
+
+class ConfiguredHooks(RolloutHooks):
+    def __init__(self) -> None:
+        self.config = HookConfig()
+
+    def on_enqueue(self, request: RolloutCreate) -> RolloutCreate:
+        return request.model_copy(update={"input": f"{self.config.prefix}{request.input}"})
+""",
+        encoding="utf-8",
+    )
+
+    hooks = load_hooks(str(hooks_path))
+
+    assert hooks.on_enqueue(RolloutCreate(input="task")).input == "checked:task"
+    assert hooks.__class__.__module__ in sys.modules
+    assert pickle.loads(pickle.dumps(hooks)).config.prefix == "checked:"