Coverage for mcpgateway / plugins / framework / hooks / agents.py: 100%

1# -*- coding: utf-8 -*- 

2"""Location: ./mcpgateway/plugins/models/agents.py 

3Copyright 2025 

4SPDX-License-Identifier: Apache-2.0 

5Authors: Teryl Taylor, Fred Araujo 

6 

7Pydantic models for agent plugins. 

8This module implements the pydantic models associated with 

9the base plugin layer including configurations, and contexts. 

10""" 

11 

12# Standard 

13from enum import Enum 

14from typing import Any, Dict, List, Optional 

15 

16# Third-Party 

17from pydantic import Field, field_validator 

18 

19# First-Party 

20from mcpgateway.plugins.framework.hooks.http import HttpHeaderPayload 

21from mcpgateway.plugins.framework.models import PluginPayload, PluginResult 

22from mcpgateway.plugins.framework.protocols import MessageLike # noqa: F401 # pylint: disable=unused-import 

23from mcpgateway.plugins.framework.utils import coerce_messages 

24 

25 

26class AgentHookType(str, Enum): 

27 """Agent hook points. 

28 

29 Attributes: 

30 AGENT_PRE_INVOKE: Before agent invocation. 

31 AGENT_POST_INVOKE: After agent responds. 

32 

33 Examples: 

34 >>> AgentHookType.AGENT_PRE_INVOKE 

35 <AgentHookType.AGENT_PRE_INVOKE: 'agent_pre_invoke'> 

36 >>> AgentHookType.AGENT_PRE_INVOKE.value 

37 'agent_pre_invoke' 

38 >>> AgentHookType('agent_post_invoke') 

39 <AgentHookType.AGENT_POST_INVOKE: 'agent_post_invoke'> 

40 >>> list(AgentHookType) 

41 [<AgentHookType.AGENT_PRE_INVOKE: 'agent_pre_invoke'>, <AgentHookType.AGENT_POST_INVOKE: 'agent_post_invoke'>] 

42 """ 

43 

44 AGENT_PRE_INVOKE = "agent_pre_invoke" 

45 AGENT_POST_INVOKE = "agent_post_invoke" 

46 

47 

48class AgentPreInvokePayload(PluginPayload): 

49 """Agent payload for pre-invoke hook. 

50 

51 Attributes: 

52 agent_id: The agent identifier (can be modified for routing). 

53 messages: Conversation messages (accepts any MessageLike-satisfying objects). 

54 tools: Optional list of tools available to agent. 

55 headers: Optional HTTP headers. 

56 model: Optional model override. 

57 system_prompt: Optional system instructions. 

58 parameters: Optional LLM parameters (temperature, max_tokens, etc.). 

59 

60 Examples: 

61 >>> payload = AgentPreInvokePayload(agent_id="agent-123", messages=[]) 

62 >>> payload.agent_id 

63 'agent-123' 

64 >>> payload.messages 

65 [] 

66 >>> payload.tools is None 

67 True 

68 """ 

69 

70 agent_id: str 

71 messages: List[Any] # Elements satisfy MessageLike protocol (role, content attributes) 

72 tools: Optional[List[str]] = None 

73 headers: Optional[HttpHeaderPayload] = None 

74 model: Optional[str] = None 

75 system_prompt: Optional[str] = None 

76 parameters: Optional[Dict[str, Any]] = Field(default_factory=dict) 

77 

78 @field_validator("messages", mode="before") 

79 @classmethod 

80 def _coerce_messages(cls, v: Any) -> Any: 

81 """Convert nested dicts in messages list to objects with attribute access. 

82 

83 Args: 

84 v: The raw messages value to coerce. 

85 

86 Returns: 

87 The coerced messages list. 

88 """ 

89 return coerce_messages(v) 

90 

91 

92class AgentPostInvokePayload(PluginPayload): 

93 """Agent payload for post-invoke hook. 

94 

95 Attributes: 

96 agent_id: The agent identifier. 

97 messages: Response messages from agent (accepts any MessageLike-satisfying objects). 

98 tool_calls: Optional tool invocations made by agent. 

99 

100 Examples: 

101 >>> payload = AgentPostInvokePayload(agent_id="agent-123", messages=[]) 

102 >>> payload.agent_id 

103 'agent-123' 

104 >>> payload.messages 

105 [] 

106 >>> payload.tool_calls is None 

107 True 

108 """ 

109 

110 agent_id: str 

111 messages: List[Any] # Elements satisfy MessageLike protocol (role, content attributes) 

112 tool_calls: Optional[List[Dict[str, Any]]] = None 

113 

114 @field_validator("messages", mode="before") 

115 @classmethod 

116 def _coerce_messages(cls, v: Any) -> Any: 

117 """Convert nested dicts in messages list to objects with attribute access. 

118 

119 Args: 

120 v: The raw messages value to coerce. 

121 

122 Returns: 

123 The coerced messages list. 

124 """ 

125 return coerce_messages(v) 

126 

127 

128AgentPreInvokeResult = PluginResult[AgentPreInvokePayload] 

129AgentPostInvokeResult = PluginResult[AgentPostInvokePayload] 

130 

131 

132def _register_agent_hooks() -> None: 

133 """Register agent hooks in the global registry. 

134 

135 This is called lazily to avoid circular import issues. 

136 """ 

137 # Import here to avoid circular dependency at module load time 

138 # First-Party 

139 from mcpgateway.plugins.framework.hooks.registry import get_hook_registry # pylint: disable=import-outside-toplevel 

140 

141 registry = get_hook_registry() 

142 

143 # Only register if not already registered (idempotent) 

144 if not registry.is_registered(AgentHookType.AGENT_PRE_INVOKE): 

145 registry.register_hook(AgentHookType.AGENT_PRE_INVOKE, AgentPreInvokePayload, AgentPreInvokeResult) 

146 registry.register_hook(AgentHookType.AGENT_POST_INVOKE, AgentPostInvokePayload, AgentPostInvokeResult) 

147 

148 

149_register_agent_hooks()