Coverage for src / agent_contracts / runtime / hooks.py: 100%

1"""Runtime hooks for customization.""" 

2from __future__ import annotations 

3 

4from typing import Protocol, runtime_checkable 

5 

6from agent_contracts.runtime.context import RequestContext, ExecutionResult 

7 

8 

9@runtime_checkable 

10class RuntimeHooks(Protocol): 

11 """Protocol for runtime customization hooks. 

12  

13 Implement this protocol to customize the execution lifecycle. 

14 Hooks are called at specific points during execution: 

15  

16 1. prepare_state: Called before graph execution to customize initial state 

17 2. after_execution: Called after graph execution for cleanup/persistence 

18  

19 Example: 

20 >>> class MyHooks: 

21 ... async def prepare_state(self, state, request): 

22 ... # Add custom data to state 

23 ... state = Internal.active_mode.set(state, "shopping") 

24 ... return state 

25 ...  

26 ... async def after_execution(self, state, result): 

27 ... # Persist session if needed 

28 ... await self.session_store.save(...) 

29 """ 

30 

31 async def prepare_state( 

32 self, 

33 state: dict, 

34 request: RequestContext, 

35 ) -> dict: 

36 """Prepare state before graph execution. 

37  

38 Called after initial state creation and session restoration. 

39 Use this to add app-specific state modifications. 

40  

41 Args: 

42 state: Initial state (may include restored session data) 

43 request: The execution request context 

44  

45 Returns: 

46 Modified state (should be immutable - return new dict) 

47 """ 

48 ... 

49 

50 async def after_execution( 

51 self, 

52 state: dict, 

53 result: ExecutionResult, 

54 ) -> None: 

55 """Handle post-execution tasks. 

56  

57 Called after graph execution completes. 

58 Use this for session persistence, cleanup, logging, etc. 

59  

60 Args: 

61 state: Final state after graph execution 

62 result: The execution result 

63 """ 

64 ... 

65 

66 

67class DefaultHooks: 

68 """Default implementation of RuntimeHooks (no-op). 

69  

70 Use this when no customization is needed, or as a base class 

71 for partial implementations. 

72 """ 

73 

74 async def prepare_state( 

75 self, 

76 state: dict, 

77 request: RequestContext, 

78 ) -> dict: 

79 """Default: return state unchanged.""" 

80 return state 

81 

82 async def after_execution( 

83 self, 

84 state: dict, 

85 result: ExecutionResult, 

86 ) -> None: 

87 """Default: do nothing.""" 

88 pass