massgen 0.1.0a3__py3-none-any.whl → 0.1.2__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of massgen might be problematic. Click here for more details.

Files changed (120) hide show
  1. massgen/__init__.py +1 -1
  2. massgen/agent_config.py +17 -0
  3. massgen/api_params_handler/_api_params_handler_base.py +1 -0
  4. massgen/api_params_handler/_chat_completions_api_params_handler.py +15 -2
  5. massgen/api_params_handler/_claude_api_params_handler.py +8 -1
  6. massgen/api_params_handler/_gemini_api_params_handler.py +73 -0
  7. massgen/api_params_handler/_response_api_params_handler.py +8 -1
  8. massgen/backend/base.py +83 -0
  9. massgen/backend/{base_with_mcp.py → base_with_custom_tool_and_mcp.py} +286 -15
  10. massgen/backend/capabilities.py +6 -6
  11. massgen/backend/chat_completions.py +200 -103
  12. massgen/backend/claude.py +115 -18
  13. massgen/backend/claude_code.py +378 -14
  14. massgen/backend/docs/CLAUDE_API_RESEARCH.md +3 -3
  15. massgen/backend/gemini.py +1333 -1629
  16. massgen/backend/gemini_mcp_manager.py +545 -0
  17. massgen/backend/gemini_trackers.py +344 -0
  18. massgen/backend/gemini_utils.py +43 -0
  19. massgen/backend/grok.py +39 -6
  20. massgen/backend/response.py +147 -81
  21. massgen/cli.py +605 -110
  22. massgen/config_builder.py +376 -27
  23. massgen/configs/README.md +123 -80
  24. massgen/configs/basic/multi/three_agents_default.yaml +3 -3
  25. massgen/configs/basic/single/single_agent.yaml +1 -1
  26. massgen/configs/providers/openai/gpt5_nano.yaml +3 -3
  27. massgen/configs/tools/custom_tools/claude_code_custom_tool_example.yaml +32 -0
  28. massgen/configs/tools/custom_tools/claude_code_custom_tool_example_no_path.yaml +28 -0
  29. massgen/configs/tools/custom_tools/claude_code_custom_tool_with_mcp_example.yaml +40 -0
  30. massgen/configs/tools/custom_tools/claude_code_custom_tool_with_wrong_mcp_example.yaml +38 -0
  31. massgen/configs/tools/custom_tools/claude_code_wrong_custom_tool_with_mcp_example.yaml +38 -0
  32. massgen/configs/tools/custom_tools/claude_custom_tool_example.yaml +24 -0
  33. massgen/configs/tools/custom_tools/claude_custom_tool_example_no_path.yaml +22 -0
  34. massgen/configs/tools/custom_tools/claude_custom_tool_with_mcp_example.yaml +35 -0
  35. massgen/configs/tools/custom_tools/claude_custom_tool_with_wrong_mcp_example.yaml +33 -0
  36. massgen/configs/tools/custom_tools/claude_wrong_custom_tool_with_mcp_example.yaml +33 -0
  37. massgen/configs/tools/custom_tools/gemini_custom_tool_example.yaml +24 -0
  38. massgen/configs/tools/custom_tools/gemini_custom_tool_example_no_path.yaml +22 -0
  39. massgen/configs/tools/custom_tools/gemini_custom_tool_with_mcp_example.yaml +35 -0
  40. massgen/configs/tools/custom_tools/gemini_custom_tool_with_wrong_mcp_example.yaml +33 -0
  41. massgen/configs/tools/custom_tools/gemini_wrong_custom_tool_with_mcp_example.yaml +33 -0
  42. massgen/configs/tools/custom_tools/github_issue_market_analysis.yaml +94 -0
  43. massgen/configs/tools/custom_tools/gpt5_nano_custom_tool_example.yaml +24 -0
  44. massgen/configs/tools/custom_tools/gpt5_nano_custom_tool_example_no_path.yaml +22 -0
  45. massgen/configs/tools/custom_tools/gpt5_nano_custom_tool_with_mcp_example.yaml +35 -0
  46. massgen/configs/tools/custom_tools/gpt5_nano_custom_tool_with_wrong_mcp_example.yaml +33 -0
  47. massgen/configs/tools/custom_tools/gpt5_nano_wrong_custom_tool_with_mcp_example.yaml +33 -0
  48. massgen/configs/tools/custom_tools/gpt_oss_custom_tool_example.yaml +25 -0
  49. massgen/configs/tools/custom_tools/gpt_oss_custom_tool_example_no_path.yaml +23 -0
  50. massgen/configs/tools/custom_tools/gpt_oss_custom_tool_with_mcp_example.yaml +34 -0
  51. massgen/configs/tools/custom_tools/gpt_oss_custom_tool_with_wrong_mcp_example.yaml +34 -0
  52. massgen/configs/tools/custom_tools/gpt_oss_wrong_custom_tool_with_mcp_example.yaml +34 -0
  53. massgen/configs/tools/custom_tools/grok3_mini_custom_tool_example.yaml +24 -0
  54. massgen/configs/tools/custom_tools/grok3_mini_custom_tool_example_no_path.yaml +22 -0
  55. massgen/configs/tools/custom_tools/grok3_mini_custom_tool_with_mcp_example.yaml +35 -0
  56. massgen/configs/tools/custom_tools/grok3_mini_custom_tool_with_wrong_mcp_example.yaml +33 -0
  57. massgen/configs/tools/custom_tools/grok3_mini_wrong_custom_tool_with_mcp_example.yaml +33 -0
  58. massgen/configs/tools/custom_tools/qwen_api_custom_tool_example.yaml +25 -0
  59. massgen/configs/tools/custom_tools/qwen_api_custom_tool_example_no_path.yaml +23 -0
  60. massgen/configs/tools/custom_tools/qwen_api_custom_tool_with_mcp_example.yaml +36 -0
  61. massgen/configs/tools/custom_tools/qwen_api_custom_tool_with_wrong_mcp_example.yaml +34 -0
  62. massgen/configs/tools/custom_tools/qwen_api_wrong_custom_tool_with_mcp_example.yaml +34 -0
  63. massgen/configs/tools/custom_tools/qwen_local_custom_tool_example.yaml +24 -0
  64. massgen/configs/tools/custom_tools/qwen_local_custom_tool_example_no_path.yaml +22 -0
  65. massgen/configs/tools/custom_tools/qwen_local_custom_tool_with_mcp_example.yaml +35 -0
  66. massgen/configs/tools/custom_tools/qwen_local_custom_tool_with_wrong_mcp_example.yaml +33 -0
  67. massgen/configs/tools/custom_tools/qwen_local_wrong_custom_tool_with_mcp_example.yaml +33 -0
  68. massgen/configs/tools/filesystem/claude_code_context_sharing.yaml +1 -1
  69. massgen/configs/tools/planning/five_agents_discord_mcp_planning_mode.yaml +7 -29
  70. massgen/configs/tools/planning/five_agents_filesystem_mcp_planning_mode.yaml +5 -6
  71. massgen/configs/tools/planning/five_agents_notion_mcp_planning_mode.yaml +4 -4
  72. massgen/configs/tools/planning/five_agents_twitter_mcp_planning_mode.yaml +4 -4
  73. massgen/configs/tools/planning/gpt5_mini_case_study_mcp_planning_mode.yaml +2 -2
  74. massgen/configs/voting/gemini_gpt_voting_sensitivity.yaml +67 -0
  75. massgen/formatter/_chat_completions_formatter.py +104 -0
  76. massgen/formatter/_claude_formatter.py +120 -0
  77. massgen/formatter/_gemini_formatter.py +448 -0
  78. massgen/formatter/_response_formatter.py +88 -0
  79. massgen/frontend/coordination_ui.py +4 -2
  80. massgen/logger_config.py +35 -3
  81. massgen/message_templates.py +56 -6
  82. massgen/orchestrator.py +512 -16
  83. massgen/stream_chunk/base.py +3 -0
  84. massgen/tests/custom_tools_example.py +392 -0
  85. massgen/tests/mcp_test_server.py +17 -7
  86. massgen/tests/test_config_builder.py +423 -0
  87. massgen/tests/test_custom_tools.py +401 -0
  88. massgen/tests/test_intelligent_planning_mode.py +643 -0
  89. massgen/tests/test_tools.py +127 -0
  90. massgen/token_manager/token_manager.py +13 -4
  91. massgen/tool/README.md +935 -0
  92. massgen/tool/__init__.py +39 -0
  93. massgen/tool/_async_helpers.py +70 -0
  94. massgen/tool/_basic/__init__.py +8 -0
  95. massgen/tool/_basic/_two_num_tool.py +24 -0
  96. massgen/tool/_code_executors/__init__.py +10 -0
  97. massgen/tool/_code_executors/_python_executor.py +74 -0
  98. massgen/tool/_code_executors/_shell_executor.py +61 -0
  99. massgen/tool/_exceptions.py +39 -0
  100. massgen/tool/_file_handlers/__init__.py +10 -0
  101. massgen/tool/_file_handlers/_file_operations.py +218 -0
  102. massgen/tool/_manager.py +634 -0
  103. massgen/tool/_registered_tool.py +88 -0
  104. massgen/tool/_result.py +66 -0
  105. massgen/tool/_self_evolution/_github_issue_analyzer.py +369 -0
  106. massgen/tool/docs/builtin_tools.md +681 -0
  107. massgen/tool/docs/exceptions.md +794 -0
  108. massgen/tool/docs/execution_results.md +691 -0
  109. massgen/tool/docs/manager.md +887 -0
  110. massgen/tool/docs/workflow_toolkits.md +529 -0
  111. massgen/tool/workflow_toolkits/__init__.py +57 -0
  112. massgen/tool/workflow_toolkits/base.py +55 -0
  113. massgen/tool/workflow_toolkits/new_answer.py +126 -0
  114. massgen/tool/workflow_toolkits/vote.py +167 -0
  115. {massgen-0.1.0a3.dist-info → massgen-0.1.2.dist-info}/METADATA +87 -129
  116. {massgen-0.1.0a3.dist-info → massgen-0.1.2.dist-info}/RECORD +120 -44
  117. {massgen-0.1.0a3.dist-info → massgen-0.1.2.dist-info}/WHEEL +0 -0
  118. {massgen-0.1.0a3.dist-info → massgen-0.1.2.dist-info}/entry_points.txt +0 -0
  119. {massgen-0.1.0a3.dist-info → massgen-0.1.2.dist-info}/licenses/LICENSE +0 -0
  120. {massgen-0.1.0a3.dist-info → massgen-0.1.2.dist-info}/top_level.txt +0 -0
massgen/backend/gemini.py CHANGED
@@ -19,15 +19,13 @@ TECHNICAL SOLUTION:
19
19
  - Maintains compatibility with existing MassGen workflow
20
20
  """
21
21
 
22
- import asyncio
23
- import enum
24
- import hashlib
25
22
  import json
26
23
  import os
27
- import re
28
24
  import time
29
- from typing import Any, AsyncGenerator, Awaitable, Callable, Dict, List, Optional
25
+ from typing import Any, AsyncGenerator, Dict, List, Optional
30
26
 
27
+ from ..api_params_handler._gemini_api_params_handler import GeminiAPIParamsHandler
28
+ from ..formatter._gemini_formatter import GeminiFormatter
31
29
  from ..logger_config import (
32
30
  log_backend_activity,
33
31
  log_backend_agent_message,
@@ -35,7 +33,11 @@ from ..logger_config import (
35
33
  log_tool_call,
36
34
  logger,
37
35
  )
38
- from .base import FilesystemSupport, LLMBackend, StreamChunk
36
+ from .base import FilesystemSupport, StreamChunk
37
+ from .base_with_custom_tool_and_mcp import CustomToolAndMCPBackend
38
+ from .gemini_mcp_manager import GeminiMCPManager
39
+ from .gemini_trackers import MCPCallTracker, MCPResponseExtractor, MCPResponseTracker
40
+ from .gemini_utils import CoordinationResponse
39
41
 
40
42
  try:
41
43
  from pydantic import BaseModel, Field
@@ -45,10 +47,12 @@ except ImportError:
45
47
 
46
48
  # MCP integration imports
47
49
  try:
48
- from ..mcp_tools import MCPClient, MCPConnectionError, MCPError
49
- from ..mcp_tools.config_validator import MCPConfigValidator
50
- from ..mcp_tools.exceptions import (
50
+ from ..mcp_tools import (
51
+ MCPClient,
51
52
  MCPConfigurationError,
53
+ MCPConfigValidator,
54
+ MCPConnectionError,
55
+ MCPError,
52
56
  MCPServerError,
53
57
  MCPTimeoutError,
54
58
  MCPValidationError,
@@ -82,1106 +86,83 @@ except ImportError:
82
86
  MCPConfigHelper = None # type: ignore[assignment]
83
87
 
84
88
 
85
- class ActionType(enum.Enum):
86
- """Action types for structured output."""
87
-
88
- VOTE = "vote"
89
- NEW_ANSWER = "new_answer"
90
-
91
-
92
- class VoteAction(BaseModel):
93
- """Structured output for voting action."""
94
-
95
- action: ActionType = Field(default=ActionType.VOTE, description="Action type")
96
- agent_id: str = Field(description="Anonymous agent ID to vote for (e.g., 'agent1', 'agent2')")
97
- reason: str = Field(description="Brief reason why this agent has the best answer")
98
-
99
-
100
- class NewAnswerAction(BaseModel):
101
- """Structured output for new answer action."""
102
-
103
- action: ActionType = Field(default=ActionType.NEW_ANSWER, description="Action type")
104
- content: str = Field(description="Your improved answer. If any builtin tools like search or code execution were used, include how they are used here.")
105
-
106
-
107
- class CoordinationResponse(BaseModel):
108
- """Structured response for coordination actions."""
109
-
110
- action_type: ActionType = Field(description="Type of action to take")
111
- vote_data: Optional[VoteAction] = Field(default=None, description="Vote data if action is vote")
112
- answer_data: Optional[NewAnswerAction] = Field(default=None, description="Answer data if action is new_answer")
113
-
114
-
115
- class MCPResponseTracker:
89
+ def format_tool_response_as_json(response_text: str) -> str:
116
90
  """
117
- Tracks MCP tool responses across streaming chunks to handle deduplication.
118
-
119
- Similar to MCPCallTracker but for tracking tool responses to avoid duplicate output.
120
- """
121
-
122
- def __init__(self):
123
- """Initialize the tracker with empty storage."""
124
- self.processed_responses = set() # Store hashes of processed responses
125
- self.response_history = [] # Store all unique responses with timestamps
126
-
127
- def get_response_hash(self, tool_name: str, tool_response: Any) -> str:
128
- """
129
- Generate a unique hash for a tool response based on name and response content.
130
-
131
- Args:
132
- tool_name: Name of the tool that responded
133
- tool_response: Response from the tool
134
-
135
- Returns:
136
- MD5 hash string identifying this specific response
137
- """
138
- # Create a deterministic string representation
139
- content = f"{tool_name}:{str(tool_response)}"
140
- return hashlib.md5(content.encode()).hexdigest()
141
-
142
- def is_new_response(self, tool_name: str, tool_response: Any) -> bool:
143
- """
144
- Check if this is a new tool response we haven't seen before.
145
-
146
- Args:
147
- tool_name: Name of the tool that responded
148
- tool_response: Response from the tool
91
+ Format tool response text as pretty-printed JSON if possible.
149
92
 
150
- Returns:
151
- True if this is a new response, False if already processed
152
- """
153
- response_hash = self.get_response_hash(tool_name, tool_response)
154
- return response_hash not in self.processed_responses
155
-
156
- def add_response(self, tool_name: str, tool_response: Any) -> Dict[str, Any]:
157
- """
158
- Add a new response to the tracker.
159
-
160
- Args:
161
- tool_name: Name of the tool that responded
162
- tool_response: Response from the tool
163
-
164
- Returns:
165
- Dictionary containing response details and timestamp
166
- """
167
- response_hash = self.get_response_hash(tool_name, tool_response)
168
- self.processed_responses.add(response_hash)
169
-
170
- record = {
171
- "tool_name": tool_name,
172
- "response": tool_response,
173
- "hash": response_hash,
174
- "timestamp": time.time(),
175
- }
176
- self.response_history.append(record)
177
- return record
178
-
179
-
180
- class MCPCallTracker:
181
- """
182
- Tracks MCP tool calls across streaming chunks to handle deduplication.
93
+ Args:
94
+ response_text: The raw response text from a tool
183
95
 
184
- Uses hashing to identify unique tool calls and timestamps to track when they occurred.
185
- This ensures we don't double-count the same tool call appearing in multiple chunks.
96
+ Returns:
97
+ Pretty-printed JSON string if response is valid JSON, otherwise original text
186
98
  """
99
+ try:
100
+ # Try to parse as JSON
101
+ parsed = json.loads(response_text)
102
+ # Return pretty-printed JSON with 2-space indentation
103
+ return json.dumps(parsed, indent=2, ensure_ascii=False)
104
+ except (json.JSONDecodeError, TypeError):
105
+ # If not valid JSON, return original text
106
+ return response_text
187
107
 
188
- def __init__(self):
189
- """Initialize the tracker with empty storage."""
190
- self.processed_calls = set() # Store hashes of processed calls
191
- self.call_history = [] # Store all unique calls with timestamps
192
- self.last_chunk_calls = [] # Track calls from the last chunk for deduplication
193
- self.dedup_window = 0.5 # Time window in seconds for deduplication
194
108
 
195
- def get_call_hash(self, tool_name: str, tool_args: Dict[str, Any]) -> str:
196
- """
197
- Generate a unique hash for a tool call based on name and arguments.
198
-
199
- Args:
200
- tool_name: Name of the tool being called
201
- tool_args: Arguments passed to the tool
202
-
203
- Returns:
204
- MD5 hash string identifying this specific call
205
- """
206
- # Create a deterministic string representation
207
- content = f"{tool_name}:{json.dumps(tool_args, sort_keys=True)}"
208
- return hashlib.md5(content.encode()).hexdigest()
209
-
210
- def is_new_call(self, tool_name: str, tool_args: Dict[str, Any]) -> bool:
211
- """
212
- Check if this is a new tool call we haven't seen before.
213
-
214
- Uses a time-window based approach: identical calls within the dedup_window
215
- are considered duplicates (likely from streaming chunks), while those outside
216
- the window are considered new calls (likely intentional repeated calls).
217
-
218
- Args:
219
- tool_name: Name of the tool being called
220
- tool_args: Arguments passed to the tool
221
-
222
- Returns:
223
- True if this is a new call, False if we've seen it before
224
- """
225
- call_hash = self.get_call_hash(tool_name, tool_args)
226
- current_time = time.time()
227
-
228
- # Check if this call exists in recent history within the dedup window
229
- for call in self.call_history[-10:]: # Check last 10 calls for efficiency
230
- if call.get("hash") == call_hash:
231
- time_diff = current_time - call.get("timestamp", 0)
232
- if time_diff < self.dedup_window:
233
- # This is likely a duplicate from streaming chunks
234
- return False
235
- # If outside the window, treat as a new intentional call
236
-
237
- # Mark as processed
238
- self.processed_calls.add(call_hash)
239
- return True
240
-
241
- def add_call(self, tool_name: str, tool_args: Dict[str, Any]) -> Dict[str, Any]:
242
- """
243
- Add a new tool call to the history.
244
-
245
- Args:
246
- tool_name: Name of the tool being called
247
- tool_args: Arguments passed to the tool
248
-
249
- Returns:
250
- Dictionary containing the call details with timestamp and hash
251
- """
252
- call_record = {
253
- "name": tool_name,
254
- "arguments": tool_args,
255
- "timestamp": time.time(),
256
- "hash": self.get_call_hash(tool_name, tool_args),
257
- "sequence": len(self.call_history), # Add sequence number for ordering
258
- }
259
- self.call_history.append(call_record)
260
-
261
- # Clean up old history to prevent memory growth
262
- if len(self.call_history) > 100:
263
- self.call_history = self.call_history[-50:]
264
-
265
- return call_record
266
-
267
- def get_summary(self) -> str:
268
- """
269
- Get a summary of all tracked tool calls.
270
-
271
- Returns:
272
- Human-readable summary of tool usage
273
- """
274
- if not self.call_history:
275
- return "No MCP tools called"
276
-
277
- tool_names = [call["name"] for call in self.call_history]
278
- unique_tools = list(dict.fromkeys(tool_names)) # Preserve order
279
- return f"Used {len(self.call_history)} MCP tool calls: {', '.join(unique_tools)}"
280
-
281
-
282
- class MCPResponseExtractor:
283
- """
284
- Extracts MCP tool calls and responses from Gemini SDK stream chunks.
285
-
286
- This class parses the internal SDK chunks to capture:
287
- - function_call parts (tool invocations)
288
- - function_response parts (tool results)
289
- - Paired call-response data for tracking complete tool executions
290
- """
291
-
292
- def __init__(self):
293
- """Initialize the extractor with empty storage."""
294
- self.mcp_calls = [] # All tool calls
295
- self.mcp_responses = [] # All tool responses
296
- self.call_response_pairs = [] # Matched call-response pairs
297
- self._pending_call = None # Track current call awaiting response
298
-
299
- def extract_function_call(self, function_call) -> Optional[Dict[str, Any]]:
300
- """
301
- Extract tool call information from SDK function_call object.
302
-
303
- Tries multiple methods to extract data from different SDK versions:
304
- 1. Direct attributes (name, args)
305
- 2. Dictionary-like interface (get method)
306
- 3. __dict__ attributes
307
- 4. Protobuf _pb attributes
308
- """
309
- tool_name = None
310
- tool_args = None
311
-
312
- # Method 1: Direct attributes
313
- tool_name = getattr(function_call, "name", None)
314
- tool_args = getattr(function_call, "args", None)
315
-
316
- # Method 2: Dictionary-like object
317
- if tool_name is None:
318
- try:
319
- if hasattr(function_call, "get"):
320
- tool_name = function_call.get("name", None)
321
- tool_args = function_call.get("args", None)
322
- except Exception:
323
- pass
324
-
325
- # Method 3: __dict__ inspection
326
- if tool_name is None:
327
- try:
328
- if hasattr(function_call, "__dict__"):
329
- fc_dict = function_call.__dict__
330
- tool_name = fc_dict.get("name", None)
331
- tool_args = fc_dict.get("args", None)
332
- except Exception:
333
- pass
334
-
335
- # Method 4: Protobuf _pb attribute
336
- if tool_name is None:
337
- try:
338
- if hasattr(function_call, "_pb"):
339
- pb = function_call._pb
340
- if hasattr(pb, "name"):
341
- tool_name = pb.name
342
- if hasattr(pb, "args"):
343
- tool_args = pb.args
344
- except Exception:
345
- pass
346
-
347
- if tool_name:
348
- call_data = {
349
- "name": tool_name,
350
- "arguments": tool_args or {},
351
- "timestamp": time.time(),
352
- "raw": str(function_call)[:200], # Truncate for logging
353
- }
354
- self.mcp_calls.append(call_data)
355
- self._pending_call = call_data
356
- return call_data
357
-
358
- return None
359
-
360
- def extract_function_response(self, function_response) -> Optional[Dict[str, Any]]:
361
- """
362
- Extract tool response information from SDK function_response object.
363
-
364
- Uses same extraction methods as function_call for consistency.
365
- """
366
- tool_name = None
367
- tool_response = None
368
-
369
- # Method 1: Direct attributes
370
- tool_name = getattr(function_response, "name", None)
371
- tool_response = getattr(function_response, "response", None)
372
-
373
- # Method 2: Dictionary-like object
374
- if tool_name is None:
375
- try:
376
- if hasattr(function_response, "get"):
377
- tool_name = function_response.get("name", None)
378
- tool_response = function_response.get("response", None)
379
- except Exception:
380
- pass
381
-
382
- # Method 3: __dict__ inspection
383
- if tool_name is None:
384
- try:
385
- if hasattr(function_response, "__dict__"):
386
- fr_dict = function_response.__dict__
387
- tool_name = fr_dict.get("name", None)
388
- tool_response = fr_dict.get("response", None)
389
- except Exception:
390
- pass
391
-
392
- # Method 4: Protobuf _pb attribute
393
- if tool_name is None:
394
- try:
395
- if hasattr(function_response, "_pb"):
396
- pb = function_response._pb
397
- if hasattr(pb, "name"):
398
- tool_name = pb.name
399
- if hasattr(pb, "response"):
400
- tool_response = pb.response
401
- except Exception:
402
- pass
403
-
404
- if tool_name:
405
- response_data = {
406
- "name": tool_name,
407
- "response": tool_response or {},
408
- "timestamp": time.time(),
409
- "raw": str(function_response)[:500], # Truncate for logging
410
- }
411
- self.mcp_responses.append(response_data)
412
-
413
- # Pair with pending call if names match
414
- if self._pending_call and self._pending_call["name"] == tool_name:
415
- self.call_response_pairs.append(
416
- {
417
- "call": self._pending_call,
418
- "response": response_data,
419
- "duration": response_data["timestamp"] - self._pending_call["timestamp"],
420
- "paired_at": time.time(),
421
- },
422
- )
423
- self._pending_call = None
424
-
425
- return response_data
426
-
427
- return None
428
-
429
- def get_summary(self) -> Dict[str, Any]:
430
- """
431
- Get a summary of all extracted MCP tool interactions.
432
- """
433
- return {
434
- "total_calls": len(self.mcp_calls),
435
- "total_responses": len(self.mcp_responses),
436
- "paired_interactions": len(self.call_response_pairs),
437
- "pending_call": self._pending_call is not None,
438
- "tool_names": list(set(call["name"] for call in self.mcp_calls)),
439
- "average_duration": (sum(pair["duration"] for pair in self.call_response_pairs) / len(self.call_response_pairs) if self.call_response_pairs else 0),
440
- }
109
+ class GeminiBackend(CustomToolAndMCPBackend):
110
+ """Google Gemini backend using structured output for coordination and MCP tool integration."""
441
111
 
442
- def clear(self):
443
- """Clear all stored data."""
444
- self.mcp_calls.clear()
445
- self.mcp_responses.clear()
446
- self.call_response_pairs.clear()
447
- self._pending_call = None
112
+ def __init__(self, api_key: Optional[str] = None, **kwargs):
113
+ # Store Gemini-specific API key before calling parent init
114
+ gemini_api_key = api_key or os.getenv("GOOGLE_API_KEY") or os.getenv("GEMINI_API_KEY")
448
115
 
116
+ # Call parent class __init__ - this initializes custom_tool_manager and MCP-related attributes
117
+ super().__init__(gemini_api_key, **kwargs)
449
118
 
450
- class GeminiBackend(LLMBackend):
451
- """Google Gemini backend using structured output for coordination and MCP tool integration."""
119
+ # Override API key with Gemini-specific value
120
+ self.api_key = gemini_api_key
452
121
 
453
- def __init__(self, api_key: Optional[str] = None, **kwargs):
454
- super().__init__(api_key, **kwargs)
455
- self.api_key = api_key or os.getenv("GOOGLE_API_KEY") or os.getenv("GEMINI_API_KEY")
122
+ # Gemini-specific counters for builtin tools
456
123
  self.search_count = 0
457
124
  self.code_execution_count = 0
458
125
 
459
- # MCP integration (filesystem MCP server may have been injected by base class)
460
- self.mcp_servers = self.config.get("mcp_servers", [])
461
- self.allowed_tools = kwargs.pop("allowed_tools", None)
462
- self.exclude_tools = kwargs.pop("exclude_tools", None)
463
- self._mcp_client: Optional[MCPClient] = None
464
- self._mcp_initialized = False
126
+ # New components for separation of concerns
127
+ self.formatter = GeminiFormatter()
128
+ self.api_params_handler = GeminiAPIParamsHandler(self)
465
129
 
466
- # MCP tool execution monitoring
467
- self._mcp_tool_calls_count = 0
468
- self._mcp_tool_failures = 0
130
+ # Gemini-specific MCP monitoring (additional to parent class)
469
131
  self._mcp_tool_successes = 0
470
-
471
- # MCP Response Extractor for capturing tool interactions
472
- self.mcp_extractor = MCPResponseExtractor()
473
-
474
- # Limit for message history growth within MCP execution loop
475
- self._max_mcp_message_history = kwargs.pop("max_mcp_message_history", 200)
476
132
  self._mcp_connection_retries = 0
477
133
 
478
- # Circuit breaker configuration
479
- self._circuit_breakers_enabled = kwargs.pop("circuit_breaker_enabled", True)
480
- self._mcp_tools_circuit_breaker = None
481
-
482
- # Initialize agent_id for use throughout the class
483
- self.agent_id = kwargs.get("agent_id", None)
484
-
485
- # Initialize circuit breaker if enabled
486
- if self._circuit_breakers_enabled:
487
- # Fail fast if required utilities are missing
488
- if MCPCircuitBreakerManager is None:
489
- raise RuntimeError("Circuit breakers enabled but MCPCircuitBreakerManager is not available")
490
-
491
- try:
492
- from ..mcp_tools.circuit_breaker import MCPCircuitBreaker
134
+ # MCP Response Extractor for capturing tool interactions (Gemini-specific)
135
+ self.mcp_extractor = MCPResponseExtractor()
493
136
 
494
- # Use shared utility to build circuit breaker configuration
495
- if MCPConfigHelper is not None:
496
- mcp_tools_config = MCPConfigHelper.build_circuit_breaker_config("mcp_tools", backend_name="gemini")
497
- else:
498
- mcp_tools_config = None
499
- if mcp_tools_config:
500
- self._mcp_tools_circuit_breaker = MCPCircuitBreaker(mcp_tools_config, backend_name="gemini", agent_id=self.agent_id)
501
- log_backend_activity(
502
- "gemini",
503
- "Circuit breaker initialized for MCP tools",
504
- {"enabled": True},
505
- agent_id=self.agent_id,
506
- )
507
- else:
508
- log_backend_activity(
509
- "gemini",
510
- "Circuit breaker config unavailable",
511
- {"fallback": "disabled"},
512
- agent_id=self.agent_id,
513
- )
514
- self._circuit_breakers_enabled = False
515
- except ImportError:
516
- log_backend_activity(
517
- "gemini",
518
- "Circuit breaker import failed",
519
- {"fallback": "disabled"},
520
- agent_id=self.agent_id,
521
- )
522
- self._circuit_breakers_enabled = False
137
+ # Initialize Gemini MCP manager after all attributes are ready
138
+ self.mcp_manager = GeminiMCPManager(self)
523
139
 
524
140
  def _setup_permission_hooks(self):
525
141
  """Override base class - Gemini uses session-based permissions, not function hooks."""
526
142
  logger.debug("[Gemini] Using session-based permissions, skipping function hook setup")
527
143
 
528
- async def _setup_mcp_with_status_stream(self, agent_id: Optional[str] = None) -> AsyncGenerator[StreamChunk, None]:
529
- """Initialize MCP client with status streaming."""
530
- status_queue: asyncio.Queue[StreamChunk] = asyncio.Queue()
531
-
532
- async def status_callback(status: str, details: Dict[str, Any]) -> None:
533
- """Callback to queue status updates as StreamChunks."""
534
- chunk = StreamChunk(
535
- type="mcp_status",
536
- status=status,
537
- content=details.get("message", ""),
538
- source="mcp_tools",
539
- )
540
- await status_queue.put(chunk)
541
-
542
- # Start the actual setup in background
543
- setup_task = asyncio.create_task(self._setup_mcp_tools_internal(agent_id, status_callback))
544
-
545
- # Yield status updates while setup is running
546
- while not setup_task.done():
547
- try:
548
- chunk = await asyncio.wait_for(status_queue.get(), timeout=0.1)
549
- yield chunk
550
- except asyncio.TimeoutError:
551
- continue
552
-
553
- # Wait for setup to complete and handle any final errors
554
- try:
555
- await setup_task
556
- except Exception as e:
557
- yield StreamChunk(
558
- type="mcp_status",
559
- status="error",
560
- content=f"MCP setup failed: {e}",
561
- source="mcp_tools",
562
- )
563
-
564
- async def _setup_mcp_tools(self, agent_id: Optional[str] = None) -> None:
565
- """Initialize MCP client (sessions only) - backward compatibility."""
566
- if not self.mcp_servers or self._mcp_initialized:
567
- return
568
- # Consume status updates without yielding them
569
- async for _ in self._setup_mcp_with_status_stream(agent_id):
570
- pass
571
-
572
- async def _setup_mcp_tools_internal(
573
- self,
574
- agent_id: Optional[str] = None,
575
- status_callback: Optional[Callable[[str, Dict[str, Any]], Awaitable[None]]] = None,
576
- ) -> None:
577
- """Internal MCP setup logic."""
578
- if not self.mcp_servers or self._mcp_initialized:
579
- return
580
-
581
- if MCPClient is None:
582
- reason = "MCP import failed - MCPClient not available"
583
- log_backend_activity(
584
- "gemini",
585
- "MCP import failed",
586
- {"reason": reason, "fallback": "workflow_tools"},
587
- agent_id=agent_id,
588
- )
589
- if status_callback:
590
- await status_callback(
591
- "error",
592
- {"message": "MCP import failed - falling back to workflow tools"},
593
- )
594
- # Clear MCP servers to prevent further attempts
595
- self.mcp_servers = []
596
- return
597
-
598
- try:
599
- # Validate MCP configuration before initialization
600
- validated_config = {
601
- "mcp_servers": self.mcp_servers,
602
- "allowed_tools": self.allowed_tools,
603
- "exclude_tools": self.exclude_tools,
604
- }
605
-
606
- if MCPConfigValidator is not None:
607
- try:
608
- backend_config = {
609
- "mcp_servers": self.mcp_servers,
610
- "allowed_tools": self.allowed_tools,
611
- "exclude_tools": self.exclude_tools,
612
- }
613
- # Use the comprehensive validator class for enhanced validation
614
- validator = MCPConfigValidator()
615
- validated_config = validator.validate_backend_mcp_config(backend_config)
616
- self.mcp_servers = validated_config.get("mcp_servers", self.mcp_servers)
617
- log_backend_activity(
618
- "gemini",
619
- "MCP configuration validated",
620
- {"server_count": len(self.mcp_servers)},
621
- agent_id=agent_id,
622
- )
623
- if status_callback:
624
- await status_callback(
625
- "info",
626
- {"message": f"MCP configuration validated: {len(self.mcp_servers)} servers"},
627
- )
628
-
629
- # Log validated server names for debugging
630
- if True:
631
- server_names = [server.get("name", "unnamed") for server in self.mcp_servers]
632
- log_backend_activity(
633
- "gemini",
634
- "MCP servers validated",
635
- {"servers": server_names},
636
- agent_id=agent_id,
637
- )
638
- except MCPConfigurationError as e:
639
- log_backend_activity(
640
- "gemini",
641
- "MCP configuration validation failed",
642
- {"error": e.original_message},
643
- agent_id=agent_id,
644
- )
645
- if status_callback:
646
- await status_callback(
647
- "error",
648
- {"message": f"Invalid MCP configuration: {e.original_message}"},
649
- )
650
- self._mcp_client = None # Clear client state for consistency
651
- raise RuntimeError(f"Invalid MCP configuration: {e.original_message}") from e
652
- except MCPValidationError as e:
653
- log_backend_activity(
654
- "gemini",
655
- "MCP validation failed",
656
- {"error": e.original_message},
657
- agent_id=agent_id,
658
- )
659
- if status_callback:
660
- await status_callback(
661
- "error",
662
- {"message": f"MCP validation error: {e.original_message}"},
663
- )
664
- self._mcp_client = None # Clear client state for consistency
665
- raise RuntimeError(f"MCP validation error: {e.original_message}") from e
666
- except Exception as e:
667
- if isinstance(e, (ImportError, AttributeError)):
668
- log_backend_activity(
669
- "gemini",
670
- "MCP validation unavailable",
671
- {"reason": str(e)},
672
- agent_id=agent_id,
673
- )
674
- # Don't clear client for import errors - validation just unavailable
675
- else:
676
- log_backend_activity(
677
- "gemini",
678
- "MCP validation error",
679
- {"error": str(e)},
680
- agent_id=agent_id,
681
- )
682
- self._mcp_client = None # Clear client state for consistency
683
- raise RuntimeError(f"MCP configuration validation failed: {e}") from e
684
- else:
685
- log_backend_activity(
686
- "gemini",
687
- "MCP validation skipped",
688
- {"reason": "validator_unavailable"},
689
- agent_id=agent_id,
690
- )
691
-
692
- # Instead of the current fallback logic
693
- normalized_servers = MCPSetupManager.normalize_mcp_servers(self.mcp_servers)
694
- log_backend_activity(
695
- "gemini",
696
- "Setting up MCP sessions",
697
- {"server_count": len(normalized_servers)},
698
- agent_id=agent_id,
699
- )
700
- if status_callback:
701
- await status_callback(
702
- "info",
703
- {"message": f"Setting up MCP sessions for {len(normalized_servers)} servers"},
704
- )
705
-
706
- # Apply circuit breaker filtering before connection attempts
707
- if self._circuit_breakers_enabled and self._mcp_tools_circuit_breaker:
708
- filtered_servers = MCPCircuitBreakerManager.apply_circuit_breaker_filtering(
709
- normalized_servers,
710
- self._mcp_tools_circuit_breaker,
711
- backend_name="gemini",
712
- agent_id=agent_id,
713
- )
714
- else:
715
- filtered_servers = normalized_servers
716
- if not filtered_servers:
717
- log_backend_activity(
718
- "gemini",
719
- "All MCP servers blocked by circuit breaker",
720
- {},
721
- agent_id=agent_id,
722
- )
723
- if status_callback:
724
- await status_callback(
725
- "warning",
726
- {"message": "All MCP servers blocked by circuit breaker"},
727
- )
728
- return
729
-
730
- if len(filtered_servers) < len(normalized_servers):
731
- log_backend_activity(
732
- "gemini",
733
- "Circuit breaker filtered servers",
734
- {"filtered_count": len(normalized_servers) - len(filtered_servers)},
735
- agent_id=agent_id,
736
- )
737
- if status_callback:
738
- await status_callback(
739
- "warning",
740
- {"message": f"Circuit breaker filtered {len(normalized_servers) - len(filtered_servers)} servers"},
741
- )
742
-
743
- # Extract tool filtering parameters from validated config
744
- allowed_tools = validated_config.get("allowed_tools")
745
- exclude_tools = validated_config.get("exclude_tools")
746
-
747
- # Log tool filtering if configured
748
- if allowed_tools:
749
- log_backend_activity(
750
- "gemini",
751
- "MCP tool filtering configured",
752
- {"allowed_tools": allowed_tools},
753
- agent_id=agent_id,
754
- )
755
- if exclude_tools:
756
- log_backend_activity(
757
- "gemini",
758
- "MCP tool filtering configured",
759
- {"exclude_tools": exclude_tools},
760
- agent_id=agent_id,
761
- )
762
-
763
- # Create client with status callback and hooks
764
- self._mcp_client = MCPClient(
765
- filtered_servers,
766
- timeout_seconds=30,
767
- allowed_tools=allowed_tools,
768
- exclude_tools=exclude_tools,
769
- status_callback=status_callback,
770
- hooks=self.filesystem_manager.get_pre_tool_hooks() if self.filesystem_manager else {},
771
- )
772
-
773
- # Connect the client
774
- await self._mcp_client.connect()
775
-
776
- # Determine which servers actually connected
777
- try:
778
- connected_server_names = self._mcp_client.get_server_names()
779
- except Exception:
780
- connected_server_names = []
781
-
782
- if not connected_server_names:
783
- # Treat as connection failure: no active servers
784
- if self._circuit_breakers_enabled and self._mcp_tools_circuit_breaker:
785
- await MCPCircuitBreakerManager.record_event(
786
- filtered_servers,
787
- self._mcp_tools_circuit_breaker,
788
- "failure",
789
- error_message="No servers connected",
790
- backend_name="gemini",
791
- agent_id=agent_id,
792
- )
793
-
794
- log_backend_activity(
795
- "gemini",
796
- "MCP connection failed: no servers connected",
797
- {},
798
- agent_id=agent_id,
799
- )
800
- if status_callback:
801
- await status_callback(
802
- "error",
803
- {"message": "MCP connection failed: no servers connected"},
804
- )
805
- self._mcp_client = None
806
- return
807
-
808
- # Record success ONLY for servers that actually connected
809
- connected_server_configs = [server for server in filtered_servers if server.get("name") in connected_server_names]
810
- if connected_server_configs:
811
- if self._circuit_breakers_enabled and self._mcp_tools_circuit_breaker:
812
- await MCPCircuitBreakerManager.record_event(
813
- connected_server_configs,
814
- self._mcp_tools_circuit_breaker,
815
- "success",
816
- backend_name="gemini",
817
- agent_id=agent_id,
818
- )
819
-
820
- self._mcp_initialized = True
821
- log_backend_activity("gemini", "MCP sessions initialized successfully", {}, agent_id=agent_id)
822
- if status_callback:
823
- await status_callback(
824
- "success",
825
- {"message": f"MCP sessions initialized successfully with {len(connected_server_names)} servers"},
826
- )
827
-
828
- except Exception as e:
829
- # Record failure for circuit breaker
830
- if self._circuit_breakers_enabled and self._mcp_tools_circuit_breaker:
831
- servers = MCPSetupManager.normalize_mcp_servers(self.mcp_servers)
832
- await MCPCircuitBreakerManager.record_event(
833
- servers,
834
- self._mcp_tools_circuit_breaker,
835
- "failure",
836
- error_message=str(e),
837
- backend_name="gemini",
838
- agent_id=agent_id,
839
- )
840
-
841
- # Enhanced error handling for different MCP error types
842
- if isinstance(e, RuntimeError) and "MCP configuration" in str(e):
843
- raise
844
- elif isinstance(e, MCPConnectionError):
845
- log_backend_activity(
846
- "gemini",
847
- "MCP connection failed during setup",
848
- {"error": str(e)},
849
- agent_id=agent_id,
850
- )
851
- if status_callback:
852
- await status_callback(
853
- "error",
854
- {"message": f"Failed to establish MCP connections: {e}"},
855
- )
856
- self._mcp_client = None
857
- raise RuntimeError(f"Failed to establish MCP connections: {e}") from e
858
- elif isinstance(e, MCPTimeoutError):
859
- log_backend_activity(
860
- "gemini",
861
- "MCP connection timeout during setup",
862
- {"error": str(e)},
863
- agent_id=agent_id,
864
- )
865
- if status_callback:
866
- await status_callback("error", {"message": f"MCP connection timeout: {e}"})
867
- self._mcp_client = None
868
- raise RuntimeError(f"MCP connection timeout: {e}") from e
869
- elif isinstance(e, MCPServerError):
870
- log_backend_activity(
871
- "gemini",
872
- "MCP server error during setup",
873
- {"error": str(e)},
874
- agent_id=agent_id,
875
- )
876
- if status_callback:
877
- await status_callback("error", {"message": f"MCP server error: {e}"})
878
- self._mcp_client = None
879
- raise RuntimeError(f"MCP server error: {e}") from e
880
- elif isinstance(e, MCPError):
881
- log_backend_activity(
882
- "gemini",
883
- "MCP error during setup",
884
- {"error": str(e)},
885
- agent_id=agent_id,
886
- )
887
- if status_callback:
888
- await status_callback("error", {"message": f"MCP error during setup: {e}"})
889
- self._mcp_client = None
890
- return
891
-
892
- else:
893
- log_backend_activity(
894
- "gemini",
895
- "MCP session setup failed",
896
- {"error": str(e)},
897
- agent_id=agent_id,
898
- )
899
- if status_callback:
900
- await status_callback("error", {"message": f"MCP session setup failed: {e}"})
901
- self._mcp_client = None
902
-
903
- def detect_coordination_tools(self, tools: List[Dict[str, Any]]) -> bool:
904
- """Detect if tools contain vote/new_answer coordination tools."""
905
- if not tools:
906
- return False
907
-
908
- tool_names = set()
909
- for tool in tools:
910
- if tool.get("type") == "function":
911
- if "function" in tool:
912
- tool_names.add(tool["function"].get("name", ""))
913
- elif "name" in tool:
914
- tool_names.add(tool.get("name", ""))
915
-
916
- return "vote" in tool_names and "new_answer" in tool_names
917
-
918
- def build_structured_output_prompt(self, base_content: str, valid_agent_ids: Optional[List[str]] = None) -> str:
919
- """Build prompt that encourages structured output for coordination."""
920
- agent_list = ""
921
- if valid_agent_ids:
922
- agent_list = f"Valid agents: {', '.join(valid_agent_ids)}"
923
-
924
- return f"""{base_content}
925
-
926
- IMPORTANT: You must respond with a structured JSON decision at the end of your response.
927
-
928
- If you want to VOTE for an existing agent's answer:
929
- {{
930
- "action_type": "vote",
931
- "vote_data": {{
932
- "action": "vote",
933
- "agent_id": "agent1", // Choose from: {agent_list or "agent1, agent2, agent3, etc."}
934
- "reason": "Brief reason for your vote"
935
- }}
936
- }}
937
-
938
- If you want to provide a NEW ANSWER:
939
- {{
940
- "action_type": "new_answer",
941
- "answer_data": {{
942
- "action": "new_answer",
943
- "content": "Your complete improved answer here"
944
- }}
945
- }}
946
-
947
- Make your decision and include the JSON at the very end of your response."""
948
-
949
- def extract_structured_response(self, response_text: str) -> Optional[Dict[str, Any]]:
950
- """Extract structured JSON response from model output."""
951
- try:
952
- # Strategy 0: Look for JSON inside markdown code blocks first
953
- markdown_json_pattern = r"```json\s*(\{.*?\})\s*```"
954
- markdown_matches = re.findall(markdown_json_pattern, response_text, re.DOTALL)
955
-
956
- for match in reversed(markdown_matches):
957
- try:
958
- parsed = json.loads(match.strip())
959
- if isinstance(parsed, dict) and "action_type" in parsed:
960
- return parsed
961
- except json.JSONDecodeError:
962
- continue
963
-
964
- # Strategy 1: Look for complete JSON blocks with proper braces
965
- json_pattern = r"\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}"
966
- json_matches = re.findall(json_pattern, response_text, re.DOTALL)
967
-
968
- # Try parsing each match (in reverse order - last one first)
969
- for match in reversed(json_matches):
970
- try:
971
- cleaned_match = match.strip()
972
- parsed = json.loads(cleaned_match)
973
- if isinstance(parsed, dict) and "action_type" in parsed:
974
- return parsed
975
- except json.JSONDecodeError:
976
- continue
977
-
978
- # Strategy 2: Look for JSON blocks with nested braces (more complex)
979
- brace_count = 0
980
- json_start = -1
981
-
982
- for i, char in enumerate(response_text):
983
- if char == "{":
984
- if brace_count == 0:
985
- json_start = i
986
- brace_count += 1
987
- elif char == "}":
988
- brace_count -= 1
989
- if brace_count == 0 and json_start >= 0:
990
- # Found a complete JSON block
991
- json_block = response_text[json_start : i + 1]
992
- try:
993
- parsed = json.loads(json_block)
994
- if isinstance(parsed, dict) and "action_type" in parsed:
995
- return parsed
996
- except json.JSONDecodeError:
997
- pass
998
- json_start = -1
999
-
1000
- # Strategy 3: Line-by-line approach (fallback)
1001
- lines = response_text.strip().split("\n")
1002
- json_candidates = []
1003
-
1004
- for i, line in enumerate(lines):
1005
- stripped = line.strip()
1006
- if stripped.startswith("{") and stripped.endswith("}"):
1007
- json_candidates.append(stripped)
1008
- elif stripped.startswith("{"):
1009
- # Multi-line JSON - collect until closing brace
1010
- json_text = stripped
1011
- for j in range(i + 1, len(lines)):
1012
- json_text += "\n" + lines[j].strip()
1013
- if lines[j].strip().endswith("}"):
1014
- json_candidates.append(json_text)
1015
- break
1016
-
1017
- # Try to parse each candidate
1018
- for candidate in reversed(json_candidates):
1019
- try:
1020
- parsed = json.loads(candidate)
1021
- if isinstance(parsed, dict) and "action_type" in parsed:
1022
- return parsed
1023
- except json.JSONDecodeError:
1024
- continue
1025
-
1026
- return None
1027
-
1028
- except Exception:
1029
- return None
1030
-
1031
- def convert_structured_to_tool_calls(self, structured_response: Dict[str, Any]) -> List[Dict[str, Any]]:
1032
- """Convert structured response to tool call format."""
1033
- action_type = structured_response.get("action_type")
1034
-
1035
- if action_type == "vote":
1036
- vote_data = structured_response.get("vote_data", {})
1037
- return [
1038
- {
1039
- "id": f"vote_{abs(hash(str(vote_data))) % 10000 + 1}",
1040
- "type": "function",
1041
- "function": {
1042
- "name": "vote",
1043
- "arguments": {
1044
- "agent_id": vote_data.get("agent_id", ""),
1045
- "reason": vote_data.get("reason", ""),
1046
- },
1047
- },
1048
- },
1049
- ]
1050
-
1051
- elif action_type == "new_answer":
1052
- answer_data = structured_response.get("answer_data", {})
1053
- return [
1054
- {
1055
- "id": f"new_answer_{abs(hash(str(answer_data))) % 10000 + 1}",
1056
- "type": "function",
1057
- "function": {
1058
- "name": "new_answer",
1059
- "arguments": {"content": answer_data.get("content", "")},
1060
- },
1061
- },
1062
- ]
1063
-
1064
- return []
1065
-
1066
- async def _handle_mcp_retry_error(self, error: Exception, retry_count: int, max_retries: int) -> tuple[bool, AsyncGenerator[StreamChunk, None]]:
1067
- """Handle MCP retry errors with specific messaging and fallback logic.
1068
-
1069
- Returns:
1070
- tuple: (should_continue_retrying, error_chunks_generator)
144
+ async def _process_stream(self, stream, all_params, agent_id: Optional[str] = None) -> AsyncGenerator[StreamChunk, None]:
1071
145
  """
1072
- log_type, user_message, _ = MCPErrorHandler.get_error_details(error, None, log=False)
1073
-
1074
- # Log the retry attempt
1075
- log_backend_activity(
1076
- "gemini",
1077
- f"MCP {log_type} on retry",
1078
- {"attempt": retry_count, "error": str(error)},
1079
- agent_id=self.agent_id,
1080
- )
1081
-
1082
- # Check if we've exhausted retries
1083
- if retry_count >= max_retries:
1084
-
1085
- async def error_chunks():
1086
- yield StreamChunk(
1087
- type="content",
1088
- content=f"\n⚠️ {user_message} after {max_retries} attempts; falling back to workflow tools\n",
1089
- )
1090
-
1091
- return False, error_chunks()
1092
-
1093
- # Continue retrying
1094
- async def empty_chunks():
1095
- # Empty generator - just return without yielding anything
1096
- if False: # Make this a generator without actually yielding
1097
- yield
1098
-
1099
- return True, empty_chunks()
1100
-
1101
- async def _handle_mcp_error_and_fallback(
1102
- self,
1103
- error: Exception,
1104
- ) -> AsyncGenerator[StreamChunk, None]:
1105
- """Handle MCP errors with specific messaging"""
1106
- self._mcp_tool_failures += 1
1107
-
1108
- log_type, user_message, _ = MCPErrorHandler.get_error_details(error, None, log=False)
1109
-
1110
- # Log with specific error type
1111
- log_backend_activity(
1112
- "gemini",
1113
- "MCP tool call failed",
1114
- {
1115
- "call_number": self._mcp_tool_calls_count,
1116
- "error_type": log_type,
1117
- "error": str(error),
1118
- },
1119
- agent_id=self.agent_id,
1120
- )
1121
-
1122
- # Yield user-friendly error message
1123
- yield StreamChunk(
1124
- type="content",
1125
- content=f"\n⚠️ {user_message} ({error}); continuing without MCP tools\n",
1126
- )
146
+ Required by CustomToolAndMCPBackend abstract method.
147
+ Not used by Gemini - Gemini SDK handles streaming directly in stream_with_tools().
148
+ """
149
+ if False:
150
+ yield # Make this an async generator
151
+ raise NotImplementedError("Gemini uses custom streaming logic in stream_with_tools()")
1127
152
 
1128
- async def _execute_mcp_function_with_retry(self, function_name: str, args: Dict[str, Any], agent_id: Optional[str] = None) -> Any:
1129
- """Execute MCP function with exponential backoff retry logic."""
1130
- if MCPExecutionManager is None:
1131
- raise RuntimeError("MCPExecutionManager is not available - MCP backend utilities are missing")
1132
-
1133
- # Stats callback for tracking
1134
- async def stats_callback(action: str) -> int:
1135
- if action == "increment_calls":
1136
- self._mcp_tool_calls_count += 1
1137
- return self._mcp_tool_calls_count
1138
- elif action == "increment_failures":
1139
- self._mcp_tool_failures += 1
1140
- return self._mcp_tool_failures
1141
- return 0
1142
-
1143
- # Circuit breaker callback
1144
- async def circuit_breaker_callback(event: str, error_msg: str) -> None:
1145
- if event == "failure":
1146
- if self._circuit_breakers_enabled and self._mcp_tools_circuit_breaker:
1147
- servers = MCPSetupManager.normalize_mcp_servers(self.mcp_servers)
1148
- await MCPCircuitBreakerManager.record_event(
1149
- servers,
1150
- self._mcp_tools_circuit_breaker,
1151
- "failure",
1152
- error_message=error_msg,
1153
- backend_name="gemini",
1154
- agent_id=agent_id,
1155
- )
1156
- else:
1157
- # Record success only for currently connected servers
1158
- connected_names: List[str] = []
1159
- try:
1160
- if self._mcp_client:
1161
- connected_names = self._mcp_client.get_server_names()
1162
- except Exception:
1163
- connected_names = []
1164
-
1165
- if connected_names:
1166
- servers_to_record = [{"name": name} for name in connected_names]
1167
- if self._circuit_breakers_enabled and self._mcp_tools_circuit_breaker:
1168
- await MCPCircuitBreakerManager.record_event(
1169
- servers_to_record,
1170
- self._mcp_tools_circuit_breaker,
1171
- "success",
1172
- backend_name="gemini",
1173
- agent_id=agent_id,
1174
- )
153
+ async def _setup_mcp_tools(self) -> None:
154
+ """
155
+ Override parent class - Gemini uses GeminiMCPManager for MCP setup.
156
+ This method is called by the parent class's __aenter__() context manager.
157
+ """
158
+ await self.mcp_manager.setup_mcp_tools(agent_id=self.agent_id)
1175
159
 
1176
- return await MCPExecutionManager.execute_function_with_retry(
1177
- function_name=function_name,
1178
- args=args,
1179
- functions=self.functions,
1180
- max_retries=3,
1181
- stats_callback=stats_callback,
1182
- circuit_breaker_callback=circuit_breaker_callback,
1183
- logger_instance=logger,
1184
- )
160
+ def supports_upload_files(self) -> bool:
161
+ """
162
+ Override parent class - Gemini does not support upload_files preprocessing.
163
+ Returns False to skip upload_files processing in parent class methods.
164
+ """
165
+ return False
1185
166
 
1186
167
  async def stream_with_tools(self, messages: List[Dict[str, Any]], tools: List[Dict[str, Any]], **kwargs) -> AsyncGenerator[StreamChunk, None]:
1187
168
  """Stream response using Gemini API with structured output for coordination and MCP tool support."""
@@ -1216,262 +197,73 @@ Make your decision and include the JSON at the very end of your response."""
1216
197
  try:
1217
198
  from google import genai
1218
199
 
1219
- # Setup MCP with status streaming if not already initialized
200
+ # Setup MCP with status streaming via manager if not already initialized
1220
201
  if not self._mcp_initialized and self.mcp_servers:
1221
- async for chunk in self._setup_mcp_with_status_stream(agent_id):
202
+ async for chunk in self.mcp_manager.setup_mcp_with_status_stream(agent_id):
1222
203
  yield chunk
1223
204
  elif not self._mcp_initialized:
1224
205
  # Setup MCP without streaming for backward compatibility
1225
- await self._setup_mcp_tools(agent_id)
206
+ await self.mcp_manager.setup_mcp_tools(agent_id)
1226
207
 
1227
208
  # Merge constructor config with stream kwargs (stream kwargs take priority)
1228
209
  all_params = {**self.config, **kwargs}
1229
210
 
1230
211
  # Extract framework-specific parameters
1231
- enable_web_search = all_params.get("enable_web_search", False)
212
+ all_params.get("enable_web_search", False)
1232
213
  enable_code_execution = all_params.get("enable_code_execution", False)
1233
214
 
1234
215
  # Always use SDK MCP sessions when mcp_servers are configured
1235
216
  using_sdk_mcp = bool(self.mcp_servers)
1236
217
 
1237
- # Analyze tool types
1238
- is_coordination = self.detect_coordination_tools(tools)
1239
-
1240
- valid_agent_ids = None
1241
-
1242
- if is_coordination:
1243
- # Extract valid agent IDs from vote tool enum if available
1244
- for tool in tools:
1245
- if tool.get("type") == "function":
1246
- func_def = tool.get("function", {})
1247
- if func_def.get("name") == "vote":
1248
- agent_id_param = func_def.get("parameters", {}).get("properties", {}).get("agent_id", {})
1249
- if "enum" in agent_id_param:
1250
- valid_agent_ids = agent_id_param["enum"]
1251
- break
1252
-
1253
- # Build content string from messages (include tool results for multi-turn tool calling)
1254
- conversation_content = ""
1255
- system_message = ""
1256
-
1257
- for msg in messages:
1258
- role = msg.get("role")
1259
- if role == "system":
1260
- system_message = msg.get("content", "")
1261
- elif role == "user":
1262
- conversation_content += f"User: {msg.get('content', '')}\n"
1263
- elif role == "assistant":
1264
- conversation_content += f"Assistant: {msg.get('content', '')}\n"
1265
- elif role == "tool":
1266
- # Ensure tool outputs are visible to the model on the next turn
1267
- tool_output = msg.get("content", "")
1268
- conversation_content += f"Tool Result: {tool_output}\n"
1269
-
1270
- # For coordination requests, modify the prompt to use structured output
1271
- if is_coordination:
1272
- conversation_content = self.build_structured_output_prompt(conversation_content, valid_agent_ids)
1273
-
1274
- # Combine system message and conversation
1275
- full_content = ""
1276
- if system_message:
1277
- full_content += f"{system_message}\n\n"
1278
- full_content += conversation_content
1279
-
1280
- # Use google-genai package
1281
- client = genai.Client(api_key=self.api_key)
1282
-
1283
- # Setup builtin tools (only when not using SDK MCP sessions)
1284
- builtin_tools = []
1285
- if enable_web_search:
1286
- try:
1287
- from google.genai import types
1288
-
1289
- grounding_tool = types.Tool(google_search=types.GoogleSearch())
1290
- builtin_tools.append(grounding_tool)
1291
- except ImportError:
1292
- yield StreamChunk(
1293
- type="content",
1294
- content="\n⚠️ Web search requires google.genai.types\n",
1295
- )
1296
-
1297
- if enable_code_execution:
1298
- try:
1299
- from google.genai import types
1300
-
1301
- code_tool = types.Tool(code_execution=types.ToolCodeExecution())
1302
- builtin_tools.append(code_tool)
1303
- except ImportError:
1304
- yield StreamChunk(
1305
- type="content",
1306
- content="\n⚠️ Code execution requires google.genai.types\n",
1307
- )
1308
-
1309
- # Build config with direct parameter passthrough
1310
- config = {}
1311
-
1312
- # Direct passthrough of all parameters except those handled separately
1313
- excluded_params = self.get_base_excluded_config_params() | {
1314
- # Gemini specific exclusions
1315
- "enable_web_search",
1316
- "enable_code_execution",
1317
- "use_multi_mcp",
1318
- "mcp_sdk_auto",
1319
- "allowed_tools",
1320
- "exclude_tools",
1321
- }
1322
- for key, value in all_params.items():
1323
- if key not in excluded_params and value is not None:
1324
- # Handle Gemini-specific parameter mappings
1325
- if key == "max_tokens":
1326
- config["max_output_tokens"] = value
1327
- elif key == "model":
1328
- model_name = value
1329
- else:
1330
- config[key] = value
1331
-
1332
- # Setup tools configuration (builtins only when not using sessions)
1333
- all_tools = []
1334
-
1335
- # Branch 1: SDK auto-calling via MCP sessions (reuse existing MCPClient sessions)
1336
- if using_sdk_mcp and self.mcp_servers:
1337
- if not self._mcp_client or not getattr(self._mcp_client, "is_connected", lambda: False)():
1338
- # Retry MCP connection up to 5 times before falling back
1339
- max_mcp_retries = 5
1340
- mcp_connected = False
1341
-
1342
- for retry_count in range(1, max_mcp_retries + 1):
1343
- try:
1344
- # Track retry attempts
1345
- self._mcp_connection_retries = retry_count
218
+ # Custom tool handling - add custom tools if any
219
+ using_custom_tools = bool(self.custom_tool_manager and len(self._custom_tool_names) > 0)
1346
220
 
1347
- if retry_count > 1:
1348
- log_backend_activity(
1349
- "gemini",
1350
- "MCP connection retry",
1351
- {
1352
- "attempt": retry_count,
1353
- "max_retries": max_mcp_retries,
1354
- },
1355
- agent_id=agent_id,
1356
- )
1357
- # Yield retry status
1358
- yield StreamChunk(
1359
- type="mcp_status",
1360
- status="mcp_retry",
1361
- content=f"Retrying MCP connection (attempt {retry_count}/{max_mcp_retries})",
1362
- source="mcp_tools",
1363
- )
1364
- # Brief delay between retries
1365
- await asyncio.sleep(0.5 * retry_count) # Progressive backoff
1366
-
1367
- # Apply circuit breaker filtering before retry attempts
1368
- if self._circuit_breakers_enabled and self._mcp_tools_circuit_breaker:
1369
- filtered_retry_servers = MCPCircuitBreakerManager.apply_circuit_breaker_filtering(
1370
- self.mcp_servers,
1371
- self._mcp_tools_circuit_breaker,
1372
- backend_name="gemini",
1373
- agent_id=agent_id,
1374
- )
1375
- else:
1376
- filtered_retry_servers = self.mcp_servers
1377
- if not filtered_retry_servers:
1378
- log_backend_activity(
1379
- "gemini",
1380
- "All MCP servers blocked during retry",
1381
- {},
1382
- agent_id=agent_id,
1383
- )
1384
- # Yield blocked status
1385
- yield StreamChunk(
1386
- type="mcp_status",
1387
- status="mcp_blocked",
1388
- content="All MCP servers blocked by circuit breaker",
1389
- source="mcp_tools",
1390
- )
1391
- using_sdk_mcp = False
1392
- break
1393
-
1394
- # Get validated config for tool filtering parameters
1395
- backend_config = {"mcp_servers": self.mcp_servers}
1396
- if MCPConfigValidator is not None:
1397
- try:
1398
- validator = MCPConfigValidator()
1399
- validated_config_retry = validator.validate_backend_mcp_config(backend_config)
1400
- allowed_tools_retry = validated_config_retry.get("allowed_tools")
1401
- exclude_tools_retry = validated_config_retry.get("exclude_tools")
1402
- except Exception:
1403
- allowed_tools_retry = None
1404
- exclude_tools_retry = None
1405
- else:
1406
- allowed_tools_retry = None
1407
- exclude_tools_retry = None
1408
-
1409
- self._mcp_client = await MCPClient.create_and_connect(
1410
- filtered_retry_servers,
1411
- timeout_seconds=30,
1412
- allowed_tools=allowed_tools_retry,
1413
- exclude_tools=exclude_tools_retry,
1414
- )
1415
-
1416
- # Record success for circuit breaker
1417
- if self._circuit_breakers_enabled and self._mcp_tools_circuit_breaker:
1418
- await MCPCircuitBreakerManager.record_event(
1419
- filtered_retry_servers,
1420
- self._mcp_tools_circuit_breaker,
1421
- "success",
1422
- backend_name="gemini",
1423
- agent_id=agent_id,
1424
- )
1425
- mcp_connected = True
1426
- log_backend_activity(
1427
- "gemini",
1428
- "MCP connection successful on retry",
1429
- {"attempt": retry_count},
1430
- agent_id=agent_id,
1431
- )
1432
- # Yield success status
1433
- yield StreamChunk(
1434
- type="mcp_status",
1435
- status="mcp_connected",
1436
- content=f"MCP connection successful on attempt {retry_count}",
1437
- source="mcp_tools",
1438
- )
221
+ # Analyze tool types
222
+ is_coordination = self.formatter.has_coordination_tools(tools)
223
+
224
+ valid_agent_ids = None
225
+
226
+ if is_coordination:
227
+ # Extract valid agent IDs from vote tool enum if available
228
+ for tool in tools:
229
+ if tool.get("type") == "function":
230
+ func_def = tool.get("function", {})
231
+ if func_def.get("name") == "vote":
232
+ agent_id_param = func_def.get("parameters", {}).get("properties", {}).get("agent_id", {})
233
+ if "enum" in agent_id_param:
234
+ valid_agent_ids = agent_id_param["enum"]
1439
235
  break
1440
236
 
1441
- except (
1442
- MCPConnectionError,
1443
- MCPTimeoutError,
1444
- MCPServerError,
1445
- MCPError,
1446
- Exception,
1447
- ) as e:
1448
- # Record failure for circuit breaker
1449
- if self._circuit_breakers_enabled and self._mcp_tools_circuit_breaker:
1450
- servers = MCPSetupManager.normalize_mcp_servers(self.mcp_servers)
1451
- await MCPCircuitBreakerManager.record_event(
1452
- servers,
1453
- self._mcp_tools_circuit_breaker,
1454
- "failure",
1455
- error_message=str(e),
1456
- backend_name="gemini",
1457
- agent_id=agent_id,
1458
- )
237
+ # Build content string from messages using formatter
238
+ full_content = self.formatter.format_messages(messages)
239
+ # For coordination requests, modify the prompt to use structured output
240
+ if is_coordination:
241
+ full_content = self.formatter.build_structured_output_prompt(full_content, valid_agent_ids)
242
+
243
+ # Use google-genai package
244
+ client = genai.Client(api_key=self.api_key)
1459
245
 
1460
- (
1461
- should_continue,
1462
- error_chunks,
1463
- ) = await self._handle_mcp_retry_error(e, retry_count, max_mcp_retries)
1464
- if not should_continue:
1465
- async for chunk in error_chunks:
1466
- yield chunk
1467
- using_sdk_mcp = False
246
+ # Setup builtin tools via API params handler (SDK Tool objects)
247
+ builtin_tools = self.api_params_handler.get_provider_tools(all_params)
248
+ # Build config via API params handler (maps params, excludes backend-managed ones)
249
+ config = await self.api_params_handler.build_api_params(messages, tools, all_params)
250
+ # Extract model name (not included in config)
251
+ model_name = all_params.get("model")
1468
252
 
1469
- # If all retries failed, ensure we fall back gracefully
253
+ # Setup tools configuration (builtins only when not using sessions)
254
+ all_tools = []
255
+
256
+ # Branch 1: SDK auto-calling via MCP sessions (reuse existing MCPClient sessions)
257
+ if using_sdk_mcp and self.mcp_servers:
258
+ if not self._mcp_client or not getattr(self._mcp_client, "is_connected", lambda: False)():
259
+ mcp_connected, status_chunks = await self.mcp_manager.setup_mcp_sessions_with_retry(agent_id, max_retries=5)
260
+ async for chunk in status_chunks:
261
+ yield chunk
1470
262
  if not mcp_connected:
1471
263
  using_sdk_mcp = False
1472
264
  self._mcp_client = None
1473
265
 
1474
- if not using_sdk_mcp:
266
+ if not using_sdk_mcp and not using_custom_tools:
1475
267
  all_tools.extend(builtin_tools)
1476
268
  if all_tools:
1477
269
  config["tools"] = all_tools
@@ -1479,7 +271,7 @@ Make your decision and include the JSON at the very end of your response."""
1479
271
  # For coordination requests, use JSON response format (may conflict with tools/sessions)
1480
272
  if is_coordination:
1481
273
  # Only request JSON schema when no tools are present
1482
- if (not using_sdk_mcp) and (not all_tools):
274
+ if (not using_sdk_mcp) and (not using_custom_tools) and (not all_tools):
1483
275
  config["response_mime_type"] = "application/json"
1484
276
  config["response_schema"] = CoordinationResponse.model_json_schema()
1485
277
  else:
@@ -1499,107 +291,363 @@ Make your decision and include the JSON at the very end of your response."""
1499
291
  # Use streaming for real-time response
1500
292
  full_content_text = ""
1501
293
  final_response = None
1502
- if using_sdk_mcp and self.mcp_servers:
1503
- # Reuse active sessions from MCPClient
294
+ # Buffer the last response chunk that contains candidate metadata so we can
295
+ # inspect builtin tool usage (grounding/code execution) after streaming
296
+ last_response_with_candidates = None
297
+ if (using_sdk_mcp and self.mcp_servers) or using_custom_tools:
298
+ # Process MCP and/or custom tools
1504
299
  try:
1505
- if not self._mcp_client:
1506
- raise RuntimeError("MCP client not initialized")
1507
- mcp_sessions = self._mcp_client.get_active_sessions()
1508
- if not mcp_sessions:
1509
- raise RuntimeError("No active MCP sessions available")
1510
-
1511
- # Convert sessions to permission sessions if filesystem manager is available
1512
- if self.filesystem_manager:
1513
- logger.info(f"[Gemini] Converting {len(mcp_sessions)} MCP sessions to permission sessions")
300
+ # ====================================================================
301
+ # Preparation phase: Initialize MCP and custom tools
302
+ # ====================================================================
303
+ mcp_sessions = []
304
+ mcp_error = None
305
+ custom_tools_functions = []
306
+ custom_tools_error = None
307
+
308
+ # Try to initialize MCP sessions
309
+ if using_sdk_mcp and self.mcp_servers:
1514
310
  try:
1515
- from ..mcp_tools.hooks import (
1516
- convert_sessions_to_permission_sessions,
311
+ if not self._mcp_client:
312
+ raise RuntimeError("MCP client not initialized")
313
+ mcp_sessions = self.mcp_manager.get_active_mcp_sessions(
314
+ convert_to_permission_sessions=bool(self.filesystem_manager),
1517
315
  )
316
+ if not mcp_sessions:
317
+ # If no MCP sessions, record error but don't interrupt (may still have custom tools)
318
+ mcp_error = RuntimeError("No active MCP sessions available")
319
+ logger.warning(f"[Gemini] MCP sessions unavailable: {mcp_error}")
320
+ except Exception as e:
321
+ mcp_error = e
322
+ logger.warning(f"[Gemini] Failed to initialize MCP sessions: {e}")
323
+
324
+ # Try to initialize custom tools
325
+ if using_custom_tools:
326
+ try:
327
+ # Get custom tools schemas (in OpenAI format)
328
+ custom_tools_schemas = self._get_custom_tools_schemas()
329
+ if custom_tools_schemas:
330
+ # Convert to Gemini SDK format using formatter
331
+ # formatter handles: OpenAI format -> Gemini dict -> FunctionDeclaration objects
332
+ custom_tools_functions = self.formatter.format_custom_tools(
333
+ custom_tools_schemas,
334
+ return_sdk_objects=True,
335
+ )
1518
336
 
1519
- mcp_sessions = convert_sessions_to_permission_sessions(mcp_sessions, self.filesystem_manager.path_permission_manager)
337
+ if custom_tools_functions:
338
+ logger.debug(
339
+ f"[Gemini] Loaded {len(custom_tools_functions)} custom tools " f"as FunctionDeclarations",
340
+ )
341
+ else:
342
+ custom_tools_error = RuntimeError("Custom tools conversion failed")
343
+ logger.warning(f"[Gemini] Custom tools unavailable: {custom_tools_error}")
344
+ else:
345
+ custom_tools_error = RuntimeError("No custom tools available")
346
+ logger.warning(f"[Gemini] Custom tools unavailable: {custom_tools_error}")
1520
347
  except Exception as e:
1521
- logger.error(f"[Gemini] Failed to convert sessions to permission sessions: {e}")
1522
- # Continue with regular sessions on error
1523
- else:
1524
- logger.debug("[Gemini] No filesystem manager found, using standard sessions")
348
+ custom_tools_error = e
349
+ logger.warning(f"[Gemini] Failed to initialize custom tools: {e}")
350
+
351
+ # Check if at least one tool system is available
352
+ has_mcp = bool(mcp_sessions and not mcp_error)
353
+ has_custom_tools = bool(custom_tools_functions and not custom_tools_error)
354
+
355
+ if not has_mcp and not has_custom_tools:
356
+ # Both failed, raise error to enter fallback
357
+ raise RuntimeError(
358
+ f"Both MCP and custom tools unavailable. " f"MCP error: {mcp_error}. Custom tools error: {custom_tools_error}",
359
+ )
1525
360
 
1526
- # Apply sessions as tools, do not mix with builtin or function_declarations
361
+ # ====================================================================
362
+ # Configuration phase: Build session_config
363
+ # ====================================================================
1527
364
  session_config = dict(config)
1528
365
 
1529
- # Get available tools from MCP client for logging
1530
- available_tools = []
1531
- if self._mcp_client:
1532
- available_tools = list(self._mcp_client.tools.keys())
366
+ # Collect all available tool information
367
+ available_mcp_tools = []
368
+ if has_mcp and self._mcp_client:
369
+ available_mcp_tools = list(self._mcp_client.tools.keys())
370
+
371
+ available_custom_tool_names = list(self._custom_tool_names) if has_custom_tools else []
372
+
373
+ # Apply tools to config
374
+ tools_to_apply = []
375
+ sessions_applied = False
376
+ custom_tools_applied = False
377
+
378
+ # Add MCP sessions (if available and not blocked by planning mode)
379
+ if has_mcp:
380
+ if not self.mcp_manager.should_block_mcp_tools_in_planning_mode(
381
+ self.is_planning_mode_enabled(),
382
+ available_mcp_tools,
383
+ ):
384
+ logger.debug(
385
+ f"[Gemini] Passing {len(mcp_sessions)} MCP sessions to SDK: " f"{[type(s).__name__ for s in mcp_sessions]}",
386
+ )
387
+ tools_to_apply.extend(mcp_sessions)
388
+ sessions_applied = True
389
+
390
+ if self.is_planning_mode_enabled():
391
+ blocked_tools = self.get_planning_mode_blocked_tools()
392
+
393
+ if not blocked_tools:
394
+ # Empty set means block ALL MCP tools (backward compatible)
395
+ logger.info("[Gemini] Planning mode enabled - blocking ALL MCP tools during coordination")
396
+ # Don't set tools at all - this prevents any MCP tool execution
397
+ log_backend_activity(
398
+ "gemini",
399
+ "All MCP tools blocked in planning mode",
400
+ {
401
+ "blocked_tools": len(available_mcp_tools),
402
+ "session_count": len(mcp_sessions),
403
+ },
404
+ agent_id=agent_id,
405
+ )
406
+ else:
407
+ # Selective blocking - allow non-blocked tools to be called
408
+ # The execution layer (_execute_mcp_function_with_retry) will enforce blocking
409
+ # but we still register all tools so non-blocked ones can be used
410
+ logger.info(f"[Gemini] Planning mode enabled - allowing non-blocked MCP tools, blocking {len(blocked_tools)} specific tools")
411
+
412
+ # Pass all sessions - the backend's is_mcp_tool_blocked() will handle selective blocking
413
+ session_config["tools"] = mcp_sessions
414
+
415
+ log_backend_activity(
416
+ "gemini",
417
+ "Selective MCP tools blocked in planning mode",
418
+ {
419
+ "total_tools": len(available_mcp_tools),
420
+ "blocked_tools": len(blocked_tools),
421
+ "allowed_tools": len(available_mcp_tools) - len(blocked_tools),
422
+ },
423
+ agent_id=agent_id,
424
+ )
425
+
426
+ # Add custom tools (if available)
427
+ if has_custom_tools:
428
+ # Wrap FunctionDeclarations in a Tool object for Gemini SDK
429
+ try:
430
+ from google.genai import types
431
+
432
+ # Create a Tool object containing all custom function declarations
433
+ custom_tool = types.Tool(function_declarations=custom_tools_functions)
434
+
435
+ logger.debug(
436
+ f"[Gemini] Wrapped {len(custom_tools_functions)} custom tools " f"in Tool object for SDK",
437
+ )
438
+ tools_to_apply.append(custom_tool)
439
+ custom_tools_applied = True
440
+ except Exception as e:
441
+ logger.error(f"[Gemini] Failed to wrap custom tools in Tool object: {e}")
442
+ custom_tools_error = e
443
+
444
+ # Apply tool configuration
445
+ if tools_to_apply:
446
+ session_config["tools"] = tools_to_apply
447
+
448
+ # Disable automatic function calling for custom tools
449
+ # MassGen uses declarative mode: SDK should return function call requests
450
+ # instead of automatically executing them
451
+ if has_custom_tools:
452
+ from google.genai import types
453
+
454
+ session_config["automatic_function_calling"] = types.AutomaticFunctionCallingConfig(
455
+ disable=True,
456
+ )
457
+ logger.debug("[Gemini] Disabled automatic function calling for custom tools")
458
+
459
+ # ====================================================================
460
+ # Logging and status output
461
+ # ====================================================================
462
+ if sessions_applied:
463
+ # Track MCP tool usage attempt
464
+ self._mcp_tool_calls_count += 1
1533
465
 
1534
- # Check planning mode - block MCP tools during coordination phase
1535
- if self.is_planning_mode_enabled():
1536
- logger.info("[Gemini] Planning mode enabled - blocking MCP tools during coordination")
1537
- # Don't set tools, which prevents automatic function calling
1538
466
  log_backend_activity(
1539
467
  "gemini",
1540
- "MCP tools blocked in planning mode",
468
+ "MCP tool call initiated",
1541
469
  {
1542
- "blocked_tools": len(available_tools),
470
+ "call_number": self._mcp_tool_calls_count,
1543
471
  "session_count": len(mcp_sessions),
472
+ "available_tools": available_mcp_tools[:],
473
+ "total_tools": len(available_mcp_tools),
1544
474
  },
1545
475
  agent_id=agent_id,
1546
476
  )
1547
- else:
1548
- # Log session types for debugging if needed
1549
- logger.debug(f"[Gemini] Passing {len(mcp_sessions)} sessions to SDK: {[type(s).__name__ for s in mcp_sessions]}")
1550
-
1551
- session_config["tools"] = mcp_sessions
1552
-
1553
- # Track MCP tool usage attempt
1554
- self._mcp_tool_calls_count += 1
1555
-
1556
- log_backend_activity(
1557
- "gemini",
1558
- "MCP tool call initiated",
1559
- {
1560
- "call_number": self._mcp_tool_calls_count,
1561
- "session_count": len(mcp_sessions),
1562
- "available_tools": available_tools[:], # Log first 10 tools for brevity
1563
- "total_tools": len(available_tools),
1564
- },
1565
- agent_id=agent_id,
1566
- )
1567
477
 
1568
- # Log MCP tool usage (SDK handles actual tool calling automatically)
1569
- log_tool_call(
1570
- agent_id,
1571
- "mcp_session_tools",
1572
- {
1573
- "session_count": len(mcp_sessions),
1574
- "call_number": self._mcp_tool_calls_count,
1575
- "available_tools": available_tools,
1576
- },
1577
- backend_name="gemini",
1578
- )
478
+ log_tool_call(
479
+ agent_id,
480
+ "mcp_session_tools",
481
+ {
482
+ "session_count": len(mcp_sessions),
483
+ "call_number": self._mcp_tool_calls_count,
484
+ "available_tools": available_mcp_tools,
485
+ },
486
+ backend_name="gemini",
487
+ )
1579
488
 
1580
- # Yield detailed MCP status as StreamChunk
1581
- tools_info = f" ({len(available_tools)} tools available)" if available_tools else ""
1582
- yield StreamChunk(
1583
- type="mcp_status",
1584
- status="mcp_tools_initiated",
1585
- content=f"MCP tool call initiated (call #{self._mcp_tool_calls_count}){tools_info}: {', '.join(available_tools[:5])}{'...' if len(available_tools) > 5 else ''}",
1586
- source="mcp_tools",
1587
- )
489
+ tools_info = f" ({len(available_mcp_tools)} tools available)" if available_mcp_tools else ""
490
+ yield StreamChunk(
491
+ type="mcp_status",
492
+ status="mcp_tools_initiated",
493
+ content=f"MCP tool call initiated (call #{self._mcp_tool_calls_count}){tools_info}: {', '.join(available_mcp_tools[:5])}{'...' if len(available_mcp_tools) > 5 else ''}",
494
+ source="mcp_tools",
495
+ )
496
+
497
+ if custom_tools_applied:
498
+ # Track custom tool usage attempt
499
+ log_backend_activity(
500
+ "gemini",
501
+ "Custom tools initiated",
502
+ {
503
+ "tool_count": len(custom_tools_functions),
504
+ "available_tools": available_custom_tool_names,
505
+ },
506
+ agent_id=agent_id,
507
+ )
1588
508
 
1589
- # Use async streaming call with sessions (SDK supports auto-calling MCP here)
1590
- # The SDK's session feature will still handle tool calling automatically
1591
- stream = await client.aio.models.generate_content_stream(model=model_name, contents=full_content, config=session_config)
509
+ tools_preview = ", ".join(available_custom_tool_names[:5])
510
+ tools_suffix = "..." if len(available_custom_tool_names) > 5 else ""
511
+ yield StreamChunk(
512
+ type="custom_tool_status",
513
+ status="custom_tools_initiated",
514
+ content=f"Custom tools initiated ({len(custom_tools_functions)} tools available): {tools_preview}{tools_suffix}",
515
+ source="custom_tools",
516
+ )
517
+
518
+ # ====================================================================
519
+ # Streaming phase
520
+ # ====================================================================
521
+ # Use async streaming call with sessions/tools
522
+ stream = await client.aio.models.generate_content_stream(
523
+ model=model_name,
524
+ contents=full_content,
525
+ config=session_config,
526
+ )
1592
527
 
1593
- # Initialize MCPCallTracker and MCPResponseTracker for deduplication across chunks
528
+ # Initialize trackers for both MCP and custom tools
1594
529
  mcp_tracker = MCPCallTracker()
1595
530
  mcp_response_tracker = MCPResponseTracker()
531
+ custom_tracker = MCPCallTracker() # Reuse MCPCallTracker for custom tools
532
+ custom_response_tracker = MCPResponseTracker() # Reuse for custom tools
533
+
1596
534
  mcp_tools_used = [] # Keep for backward compatibility
535
+ custom_tools_used = [] # Track custom tool usage
1597
536
 
1598
537
  # Iterate over the asynchronous stream to get chunks as they arrive
1599
538
  async for chunk in stream:
1600
539
  # ============================================
1601
- # 1. Process MCP function calls/responses
540
+ # 1. Process function calls/responses
1602
541
  # ============================================
542
+
543
+ # First check for function calls in the current chunk's candidates
544
+ # (this is where custom tool calls appear, not in automatic_function_calling_history)
545
+ if hasattr(chunk, "candidates") and chunk.candidates:
546
+ for candidate in chunk.candidates:
547
+ if hasattr(candidate, "content") and candidate.content:
548
+ if hasattr(candidate.content, "parts") and candidate.content.parts:
549
+ for part in candidate.content.parts:
550
+ # Check for function_call part
551
+ if hasattr(part, "function_call") and part.function_call:
552
+ # Extract call data
553
+ call_data = self.mcp_extractor.extract_function_call(part.function_call)
554
+
555
+ if call_data:
556
+ tool_name = call_data["name"]
557
+ tool_args = call_data["arguments"]
558
+
559
+ # DEBUG: Log tool matching
560
+ logger.info(f"🔍 [DEBUG] Function call detected: tool_name='{tool_name}'")
561
+ logger.info(f"🔍 [DEBUG] Available MCP tools: {available_mcp_tools}")
562
+ logger.info(f"🔍 [DEBUG] Available custom tools: {list(self._custom_tool_names) if has_custom_tools else []}")
563
+
564
+ # Determine if it's MCP tool or custom tool
565
+ # MCP tools may come from SDK without prefix, so we need to check both:
566
+ # 1. Direct match (tool_name in list)
567
+ # 2. Prefixed match (mcp__server__tool_name in list)
568
+ is_mcp_tool = False
569
+ if has_mcp:
570
+ # Direct match
571
+ if tool_name in available_mcp_tools:
572
+ is_mcp_tool = True
573
+ else:
574
+ # Try matching with MCP prefix format: mcp__<server>__<tool>
575
+ # Check if any available MCP tool ends with the current tool_name
576
+ for mcp_tool in available_mcp_tools:
577
+ # Format: mcp__server__toolname
578
+ if mcp_tool.startswith("mcp__") and mcp_tool.endswith(f"__{tool_name}"):
579
+ is_mcp_tool = True
580
+ logger.info(f"🔍 [DEBUG] Matched MCP tool: {tool_name} -> {mcp_tool}")
581
+ break
582
+
583
+ is_custom_tool = has_custom_tools and tool_name in self._custom_tool_names
584
+
585
+ logger.info(f"🔍 [DEBUG] Tool matching result: is_mcp_tool={is_mcp_tool}, is_custom_tool={is_custom_tool}")
586
+
587
+ if is_custom_tool:
588
+ # Process custom tool call
589
+ if custom_tracker.is_new_call(tool_name, tool_args):
590
+ call_record = custom_tracker.add_call(tool_name, tool_args)
591
+
592
+ custom_tools_used.append(
593
+ {
594
+ "name": tool_name,
595
+ "arguments": tool_args,
596
+ "timestamp": call_record["timestamp"],
597
+ },
598
+ )
599
+
600
+ timestamp_str = time.strftime(
601
+ "%H:%M:%S",
602
+ time.localtime(call_record["timestamp"]),
603
+ )
604
+
605
+ yield StreamChunk(
606
+ type="custom_tool_status",
607
+ status="custom_tool_called",
608
+ content=f"🔧 Custom Tool Called: {tool_name} at {timestamp_str} with args: {json.dumps(tool_args, indent=2)}",
609
+ source="custom_tools",
610
+ )
611
+
612
+ log_tool_call(
613
+ agent_id,
614
+ tool_name,
615
+ tool_args,
616
+ backend_name="gemini",
617
+ )
618
+ elif is_mcp_tool:
619
+ # Process MCP tool call
620
+ if mcp_tracker.is_new_call(tool_name, tool_args):
621
+ call_record = mcp_tracker.add_call(tool_name, tool_args)
622
+
623
+ mcp_tools_used.append(
624
+ {
625
+ "name": tool_name,
626
+ "arguments": tool_args,
627
+ "timestamp": call_record["timestamp"],
628
+ },
629
+ )
630
+
631
+ timestamp_str = time.strftime(
632
+ "%H:%M:%S",
633
+ time.localtime(call_record["timestamp"]),
634
+ )
635
+
636
+ yield StreamChunk(
637
+ type="mcp_status",
638
+ status="mcp_tool_called",
639
+ content=f"🔧 MCP Tool Called: {tool_name} at {timestamp_str} with args: {json.dumps(tool_args, indent=2)}",
640
+ source="mcp_tools",
641
+ )
642
+
643
+ log_tool_call(
644
+ agent_id,
645
+ tool_name,
646
+ tool_args,
647
+ backend_name="gemini",
648
+ )
649
+
650
+ # Then check automatic_function_calling_history (for MCP tools that were auto-executed)
1603
651
  if hasattr(chunk, "automatic_function_calling_history") and chunk.automatic_function_calling_history:
1604
652
  for history_item in chunk.automatic_function_calling_history:
1605
653
  if hasattr(history_item, "parts") and history_item.parts is not None:
@@ -1613,127 +661,190 @@ Make your decision and include the JSON at the very end of your response."""
1613
661
  tool_name = call_data["name"]
1614
662
  tool_args = call_data["arguments"]
1615
663
 
1616
- # Check if this is a new call using the tracker
1617
- if mcp_tracker.is_new_call(tool_name, tool_args):
1618
- # Add to tracker history
1619
- call_record = mcp_tracker.add_call(tool_name, tool_args)
1620
-
1621
- # Add to legacy list for compatibility
1622
- mcp_tools_used.append(
1623
- {
1624
- "name": tool_name,
1625
- "arguments": tool_args,
1626
- "timestamp": call_record["timestamp"],
1627
- },
1628
- )
1629
-
1630
- # Format timestamp for display
1631
- timestamp_str = time.strftime(
1632
- "%H:%M:%S",
1633
- time.localtime(call_record["timestamp"]),
1634
- )
1635
-
1636
- # Yield detailed MCP tool call information
1637
- yield StreamChunk(
1638
- type="mcp_status",
1639
- status="mcp_tool_called",
1640
- content=f"🔧 MCP Tool Called: {tool_name} at {timestamp_str} with args: {json.dumps(tool_args, indent=2)}",
1641
- source="mcp_tools",
1642
- )
1643
-
1644
- # Log the specific tool call
1645
- log_tool_call(
1646
- agent_id,
1647
- tool_name,
1648
- tool_args,
1649
- backend_name="gemini",
1650
- )
664
+ # DEBUG: Log tool matching (from automatic_function_calling_history)
665
+ logger.info(f"🔍 [DEBUG-AUTO] Function call in history: tool_name='{tool_name}'")
666
+ logger.info(f"🔍 [DEBUG-AUTO] Available MCP tools: {available_mcp_tools}")
667
+ logger.info(f"🔍 [DEBUG-AUTO] Available custom tools: {list(self._custom_tool_names) if has_custom_tools else []}")
668
+
669
+ # Determine if it's MCP tool or custom tool
670
+ # MCP tools may come from SDK without prefix, so we need to check both
671
+ is_mcp_tool = False
672
+ if has_mcp:
673
+ if tool_name in available_mcp_tools:
674
+ is_mcp_tool = True
675
+ else:
676
+ # Try matching with MCP prefix format
677
+ for mcp_tool in available_mcp_tools:
678
+ if mcp_tool.startswith("mcp__") and mcp_tool.endswith(f"__{tool_name}"):
679
+ is_mcp_tool = True
680
+ logger.info(f"🔍 [DEBUG-AUTO] Matched MCP tool: {tool_name} -> {mcp_tool}")
681
+ break
682
+
683
+ is_custom_tool = has_custom_tools and tool_name in self._custom_tool_names
684
+
685
+ logger.info(f"🔍 [DEBUG-AUTO] Tool matching result: is_mcp_tool={is_mcp_tool}, is_custom_tool={is_custom_tool}")
686
+
687
+ if is_mcp_tool:
688
+ # Process MCP tool call
689
+ if mcp_tracker.is_new_call(tool_name, tool_args):
690
+ call_record = mcp_tracker.add_call(tool_name, tool_args)
691
+
692
+ mcp_tools_used.append(
693
+ {
694
+ "name": tool_name,
695
+ "arguments": tool_args,
696
+ "timestamp": call_record["timestamp"],
697
+ },
698
+ )
699
+
700
+ timestamp_str = time.strftime(
701
+ "%H:%M:%S",
702
+ time.localtime(call_record["timestamp"]),
703
+ )
704
+
705
+ yield StreamChunk(
706
+ type="mcp_status",
707
+ status="mcp_tool_called",
708
+ content=f"🔧 MCP Tool Called: {tool_name} at {timestamp_str} with args: {json.dumps(tool_args, indent=2)}",
709
+ source="mcp_tools",
710
+ )
711
+
712
+ log_tool_call(
713
+ agent_id,
714
+ tool_name,
715
+ tool_args,
716
+ backend_name="gemini",
717
+ )
718
+
719
+ elif is_custom_tool:
720
+ # Process custom tool call
721
+ if custom_tracker.is_new_call(tool_name, tool_args):
722
+ call_record = custom_tracker.add_call(tool_name, tool_args)
723
+
724
+ custom_tools_used.append(
725
+ {
726
+ "name": tool_name,
727
+ "arguments": tool_args,
728
+ "timestamp": call_record["timestamp"],
729
+ },
730
+ )
731
+
732
+ timestamp_str = time.strftime(
733
+ "%H:%M:%S",
734
+ time.localtime(call_record["timestamp"]),
735
+ )
736
+
737
+ yield StreamChunk(
738
+ type="custom_tool_status",
739
+ status="custom_tool_called",
740
+ content=f"🔧 Custom Tool Called: {tool_name} at {timestamp_str} with args: {json.dumps(tool_args, indent=2)}",
741
+ source="custom_tools",
742
+ )
743
+
744
+ log_tool_call(
745
+ agent_id,
746
+ tool_name,
747
+ tool_args,
748
+ backend_name="gemini",
749
+ )
1651
750
 
1652
751
  # Check for function_response part
1653
752
  elif hasattr(part, "function_response") and part.function_response:
1654
- # Use MCPResponseExtractor to extract response data
1655
753
  response_data = self.mcp_extractor.extract_function_response(part.function_response)
1656
754
 
1657
755
  if response_data:
1658
756
  tool_name = response_data["name"]
1659
757
  tool_response = response_data["response"]
1660
758
 
1661
- # Check if this is a new response using the tracker
1662
- if mcp_response_tracker.is_new_response(tool_name, tool_response):
1663
- # Add to tracker history
1664
- response_record = mcp_response_tracker.add_response(tool_name, tool_response)
1665
-
1666
- # Extract text content from CallToolResult
1667
- response_text = None
1668
- if isinstance(tool_response, dict) and "result" in tool_response:
1669
- result = tool_response["result"]
1670
- # Check if result has content attribute (CallToolResult object)
1671
- if hasattr(result, "content") and result.content:
1672
- # Get the first content item (TextContent object)
1673
- first_content = result.content[0]
1674
- # Extract the text attribute
1675
- if hasattr(first_content, "text"):
1676
- response_text = first_content.text
1677
-
1678
- # Use extracted text or fallback to string representation
1679
- if response_text is None:
759
+ # Determine if it's MCP tool or custom tool
760
+ # MCP tools may come from SDK without prefix
761
+ is_mcp_tool = False
762
+ if has_mcp:
763
+ if tool_name in available_mcp_tools:
764
+ is_mcp_tool = True
765
+ else:
766
+ # Try matching with MCP prefix format
767
+ for mcp_tool in available_mcp_tools:
768
+ if mcp_tool.startswith("mcp__") and mcp_tool.endswith(f"__{tool_name}"):
769
+ is_mcp_tool = True
770
+ break
771
+
772
+ is_custom_tool = has_custom_tools and tool_name in self._custom_tool_names
773
+
774
+ if is_mcp_tool:
775
+ # Process MCP tool response
776
+ if mcp_response_tracker.is_new_response(tool_name, tool_response):
777
+ response_record = mcp_response_tracker.add_response(tool_name, tool_response)
778
+
779
+ # Extract text content from CallToolResult
780
+ response_text = None
781
+ if isinstance(tool_response, dict) and "result" in tool_response:
782
+ result = tool_response["result"]
783
+ if hasattr(result, "content") and result.content:
784
+ first_content = result.content[0]
785
+ if hasattr(first_content, "text"):
786
+ response_text = first_content.text
787
+
788
+ if response_text is None:
789
+ response_text = str(tool_response)
790
+
791
+ timestamp_str = time.strftime(
792
+ "%H:%M:%S",
793
+ time.localtime(response_record["timestamp"]),
794
+ )
795
+
796
+ # Format response as JSON if possible
797
+ formatted_response = format_tool_response_as_json(response_text)
798
+
799
+ yield StreamChunk(
800
+ type="mcp_status",
801
+ status="mcp_tool_response",
802
+ content=f"✅ MCP Tool Response from {tool_name} at {timestamp_str}: {formatted_response}",
803
+ source="mcp_tools",
804
+ )
805
+
806
+ log_backend_activity(
807
+ "gemini",
808
+ "MCP tool response received",
809
+ {
810
+ "tool_name": tool_name,
811
+ "response_preview": str(tool_response)[:],
812
+ },
813
+ agent_id=agent_id,
814
+ )
815
+
816
+ elif is_custom_tool:
817
+ # Process custom tool response
818
+ if custom_response_tracker.is_new_response(tool_name, tool_response):
819
+ response_record = custom_response_tracker.add_response(tool_name, tool_response)
820
+
821
+ # Extract text from response
1680
822
  response_text = str(tool_response)
1681
823
 
1682
- # Format timestamp for display
1683
- timestamp_str = time.strftime(
1684
- "%H:%M:%S",
1685
- time.localtime(response_record["timestamp"]),
1686
- )
1687
-
1688
- # Yield MCP tool response information
1689
- yield StreamChunk(
1690
- type="mcp_status",
1691
- status="mcp_tool_response",
1692
- content=f"✅ MCP Tool Response from {tool_name} at {timestamp_str}: {response_text}",
1693
- source="mcp_tools",
1694
- )
1695
-
1696
- # Log the tool response
1697
- log_backend_activity(
1698
- "gemini",
1699
- "MCP tool response received",
1700
- {
1701
- "tool_name": tool_name,
1702
- "response_preview": str(tool_response)[:],
1703
- },
1704
- agent_id=agent_id,
1705
- )
1706
-
1707
- # Track successful MCP tool execution (only on first chunk with MCP history)
1708
- if not hasattr(self, "_mcp_stream_started"):
1709
- self._mcp_tool_successes += 1
1710
- self._mcp_stream_started = True
1711
- log_backend_activity(
1712
- "gemini",
1713
- "MCP tool call succeeded",
1714
- {"call_number": self._mcp_tool_calls_count},
1715
- agent_id=agent_id,
1716
- )
1717
-
1718
- # Log MCP tool success as a tool call event
1719
- log_tool_call(
1720
- agent_id,
1721
- "mcp_session_tools",
1722
- {
1723
- "session_count": len(mcp_sessions),
1724
- "call_number": self._mcp_tool_calls_count,
1725
- },
1726
- result="success",
1727
- backend_name="gemini",
1728
- )
1729
-
1730
- # Yield MCP success status as StreamChunk
1731
- yield StreamChunk(
1732
- type="mcp_status",
1733
- status="mcp_tools_success",
1734
- content=f"MCP tool call succeeded (call #{self._mcp_tool_calls_count})",
1735
- source="mcp_tools",
1736
- )
824
+ timestamp_str = time.strftime(
825
+ "%H:%M:%S",
826
+ time.localtime(response_record["timestamp"]),
827
+ )
828
+
829
+ # Format response as JSON if possible
830
+ formatted_response = format_tool_response_as_json(response_text)
831
+
832
+ yield StreamChunk(
833
+ type="custom_tool_status",
834
+ status="custom_tool_response",
835
+ content=f"✅ Custom Tool Response from {tool_name} at {timestamp_str}: {formatted_response}",
836
+ source="custom_tools",
837
+ )
838
+
839
+ log_backend_activity(
840
+ "gemini",
841
+ "Custom tool response received",
842
+ {
843
+ "tool_name": tool_name,
844
+ "response_preview": str(tool_response),
845
+ },
846
+ agent_id=agent_id,
847
+ )
1737
848
 
1738
849
  # ============================================
1739
850
  # 2. Process text content
@@ -1750,24 +861,595 @@ Make your decision and include the JSON at the very end of your response."""
1750
861
  log_stream_chunk("backend.gemini", "content", chunk_text, agent_id)
1751
862
  yield StreamChunk(type="content", content=chunk_text)
1752
863
 
864
+ # ============================================
865
+ # 3. Buffer last chunk with candidates
866
+ # ============================================
867
+ if hasattr(chunk, "candidates") and chunk.candidates:
868
+ last_response_with_candidates = chunk
869
+
1753
870
  # Reset stream tracking
1754
871
  if hasattr(self, "_mcp_stream_started"):
1755
872
  delattr(self, "_mcp_stream_started")
1756
873
 
1757
- # Add MCP usage indicator with detailed summary using tracker
1758
- tools_summary = mcp_tracker.get_summary()
1759
- if not tools_summary or tools_summary == "No MCP tools called":
1760
- tools_summary = "MCP session completed (no tools explicitly called)"
1761
- else:
1762
- tools_summary = f"MCP session complete - {tools_summary}"
1763
-
1764
- log_stream_chunk("backend.gemini", "mcp_indicator", tools_summary, agent_id)
1765
- yield StreamChunk(
1766
- type="mcp_status",
1767
- status="mcp_session_complete",
1768
- content=f"MCP session complete - {tools_summary}",
1769
- source="mcp_tools",
1770
- )
874
+ # ====================================================================
875
+ # Tool execution loop: Execute tools until model stops calling them
876
+ # ====================================================================
877
+ # Note: When automatic_function_calling is disabled, BOTH custom and MCP tools
878
+ # need to be manually executed. The model may make multiple rounds of tool calls
879
+ # (e.g., call custom tool first, then MCP tool after seeing the result).
880
+
881
+ executed_tool_calls = set() # Track which tools we've already executed
882
+ max_tool_rounds = 10 # Prevent infinite loops
883
+ tool_round = 0
884
+
885
+ while tool_round < max_tool_rounds:
886
+ # Find new tool calls that haven't been executed yet
887
+ new_custom_tools = []
888
+ new_mcp_tools = []
889
+
890
+ for tool_call in custom_tools_used:
891
+ call_signature = f"custom_{tool_call['name']}_{json.dumps(tool_call['arguments'], sort_keys=True)}"
892
+ if call_signature not in executed_tool_calls:
893
+ new_custom_tools.append(tool_call)
894
+ executed_tool_calls.add(call_signature)
895
+
896
+ for tool_call in mcp_tools_used:
897
+ call_signature = f"mcp_{tool_call['name']}_{json.dumps(tool_call['arguments'], sort_keys=True)}"
898
+ if call_signature not in executed_tool_calls:
899
+ new_mcp_tools.append(tool_call)
900
+ executed_tool_calls.add(call_signature)
901
+
902
+ # If no new tools to execute, break the loop
903
+ if not new_custom_tools and not new_mcp_tools:
904
+ break
905
+
906
+ tool_round += 1
907
+ logger.debug(f"[Gemini] Tool execution round {tool_round}: {len(new_custom_tools)} custom, {len(new_mcp_tools)} MCP")
908
+
909
+ # Execute tools and collect results for this round
910
+ tool_responses = []
911
+
912
+ # Execute custom tools
913
+ for tool_call in new_custom_tools:
914
+ tool_name = tool_call["name"]
915
+ tool_args = tool_call["arguments"]
916
+
917
+ try:
918
+ # Execute the custom tool
919
+ result_str = await self._execute_custom_tool(
920
+ {
921
+ "name": tool_name,
922
+ "arguments": json.dumps(tool_args) if isinstance(tool_args, dict) else tool_args,
923
+ },
924
+ )
925
+
926
+ # Format result as JSON if possible
927
+ formatted_result = format_tool_response_as_json(result_str)
928
+
929
+ # Yield execution status
930
+ yield StreamChunk(
931
+ type="custom_tool_status",
932
+ status="custom_tool_executed",
933
+ content=f"✅ Custom Tool Executed: {tool_name} -> {formatted_result}",
934
+ source="custom_tools",
935
+ )
936
+
937
+ # Build function response in Gemini format
938
+ tool_responses.append(
939
+ {
940
+ "name": tool_name,
941
+ "response": {"result": result_str},
942
+ },
943
+ )
944
+
945
+ except Exception as e:
946
+ error_msg = f"Error executing custom tool {tool_name}: {str(e)}"
947
+ logger.error(error_msg)
948
+ yield StreamChunk(
949
+ type="custom_tool_status",
950
+ status="custom_tool_error",
951
+ content=f"❌ {error_msg}",
952
+ source="custom_tools",
953
+ )
954
+ # Add error response
955
+ tool_responses.append(
956
+ {
957
+ "name": tool_name,
958
+ "response": {"error": str(e)},
959
+ },
960
+ )
961
+
962
+ # Execute MCP tools manually (since automatic_function_calling is disabled)
963
+ for tool_call in new_mcp_tools:
964
+ tool_name = tool_call["name"]
965
+ tool_args = tool_call["arguments"]
966
+
967
+ try:
968
+ # Execute the MCP tool via MCP client
969
+ if not self._mcp_client:
970
+ raise RuntimeError("MCP client not initialized")
971
+
972
+ # Convert tool name to prefixed format if needed
973
+ # MCP client expects: mcp__server__toolname
974
+ # Gemini SDK returns: toolname (without prefix)
975
+ prefixed_tool_name = tool_name
976
+ if not tool_name.startswith("mcp__"):
977
+ # Find the matching prefixed tool name
978
+ for mcp_tool in available_mcp_tools:
979
+ if mcp_tool.startswith("mcp__") and mcp_tool.endswith(f"__{tool_name}"):
980
+ prefixed_tool_name = mcp_tool
981
+ logger.info(f"🔧 [DEBUG] Converting tool name for execution: {tool_name} -> {prefixed_tool_name}")
982
+ break
983
+
984
+ mcp_result = await self._mcp_client.call_tool(prefixed_tool_name, tool_args)
985
+
986
+ # Extract text from CallToolResult object
987
+ result_str = None
988
+ if mcp_result:
989
+ if hasattr(mcp_result, "content") and mcp_result.content:
990
+ first_content = mcp_result.content[0]
991
+ if hasattr(first_content, "text"):
992
+ result_str = first_content.text
993
+
994
+ if result_str is None:
995
+ result_str = str(mcp_result) if mcp_result else "None"
996
+
997
+ # Format result as JSON if possible
998
+ formatted_result = format_tool_response_as_json(result_str)
999
+ result_preview = formatted_result
1000
+
1001
+ # Yield execution status
1002
+ yield StreamChunk(
1003
+ type="mcp_status",
1004
+ status="mcp_tool_executed",
1005
+ content=f"✅ MCP Tool Executed: {tool_name} -> {result_preview}{'...' if len(formatted_result) > 200 else ''}",
1006
+ source="mcp_tools",
1007
+ )
1008
+
1009
+ # Build function response in Gemini format
1010
+ tool_responses.append(
1011
+ {
1012
+ "name": tool_name,
1013
+ "response": {"result": mcp_result},
1014
+ },
1015
+ )
1016
+
1017
+ except Exception as e:
1018
+ error_msg = f"Error executing MCP tool {tool_name}: {str(e)}"
1019
+ logger.error(error_msg)
1020
+ yield StreamChunk(
1021
+ type="mcp_status",
1022
+ status="mcp_tool_error",
1023
+ content=f"❌ {error_msg}",
1024
+ source="mcp_tools",
1025
+ )
1026
+ # Add error response
1027
+ tool_responses.append(
1028
+ {
1029
+ "name": tool_name,
1030
+ "response": {"error": str(e)},
1031
+ },
1032
+ )
1033
+
1034
+ # Make continuation call with tool results from this round
1035
+ if tool_responses:
1036
+ try:
1037
+ from google.genai import types
1038
+
1039
+ # Build conversation history for continuation
1040
+ # Track all function calls from this round
1041
+ round_function_calls = new_custom_tools + new_mcp_tools
1042
+
1043
+ # Build conversation history
1044
+ conversation_history = []
1045
+
1046
+ # Add original user content
1047
+ conversation_history.append(
1048
+ types.Content(
1049
+ parts=[types.Part(text=full_content)],
1050
+ role="user",
1051
+ ),
1052
+ )
1053
+
1054
+ # Add model's function call response (tools from THIS round)
1055
+ model_parts = []
1056
+ for tool_call in round_function_calls:
1057
+ model_parts.append(
1058
+ types.Part.from_function_call(
1059
+ name=tool_call["name"],
1060
+ args=tool_call["arguments"],
1061
+ ),
1062
+ )
1063
+
1064
+ conversation_history.append(
1065
+ types.Content(
1066
+ parts=model_parts,
1067
+ role="model",
1068
+ ),
1069
+ )
1070
+
1071
+ # Add function response (as user message with function_response parts)
1072
+ response_parts = []
1073
+ for resp in tool_responses:
1074
+ response_parts.append(
1075
+ types.Part.from_function_response(
1076
+ name=resp["name"],
1077
+ response=resp["response"],
1078
+ ),
1079
+ )
1080
+
1081
+ conversation_history.append(
1082
+ types.Content(
1083
+ parts=response_parts,
1084
+ role="user",
1085
+ ),
1086
+ )
1087
+
1088
+ # Make continuation call
1089
+ yield StreamChunk(
1090
+ type="custom_tool_status",
1091
+ status="continuation_call",
1092
+ content=f"🔄 Making continuation call with {len(tool_responses)} tool results...",
1093
+ source="custom_tools",
1094
+ )
1095
+
1096
+ # Use same session_config as before
1097
+ continuation_stream = await client.aio.models.generate_content_stream(
1098
+ model=model_name,
1099
+ contents=conversation_history,
1100
+ config=session_config,
1101
+ )
1102
+
1103
+ # Process continuation stream (same processing as main stream)
1104
+ async for chunk in continuation_stream:
1105
+ # ============================================
1106
+ # Process function calls/responses in continuation
1107
+ # ============================================
1108
+ # Check for function calls in current chunk's candidates
1109
+ if hasattr(chunk, "candidates") and chunk.candidates:
1110
+ for candidate in chunk.candidates:
1111
+ if hasattr(candidate, "content") and candidate.content:
1112
+ if hasattr(candidate.content, "parts") and candidate.content.parts:
1113
+ for part in candidate.content.parts:
1114
+ # Check for function_call part
1115
+ if hasattr(part, "function_call") and part.function_call:
1116
+ call_data = self.mcp_extractor.extract_function_call(part.function_call)
1117
+
1118
+ if call_data:
1119
+ tool_name = call_data["name"]
1120
+ tool_args = call_data["arguments"]
1121
+
1122
+ # Determine if it's MCP tool or custom tool
1123
+ # MCP tools may come from SDK without prefix
1124
+ is_mcp_tool = False
1125
+ if has_mcp:
1126
+ if tool_name in available_mcp_tools:
1127
+ is_mcp_tool = True
1128
+ else:
1129
+ # Try matching with MCP prefix format
1130
+ for mcp_tool in available_mcp_tools:
1131
+ if mcp_tool.startswith("mcp__") and mcp_tool.endswith(f"__{tool_name}"):
1132
+ is_mcp_tool = True
1133
+ break
1134
+
1135
+ is_custom_tool = has_custom_tools and tool_name in self._custom_tool_names
1136
+
1137
+ if is_custom_tool:
1138
+ # Process custom tool call
1139
+ if custom_tracker.is_new_call(tool_name, tool_args):
1140
+ call_record = custom_tracker.add_call(tool_name, tool_args)
1141
+
1142
+ custom_tools_used.append(
1143
+ {
1144
+ "name": tool_name,
1145
+ "arguments": tool_args,
1146
+ "timestamp": call_record["timestamp"],
1147
+ },
1148
+ )
1149
+
1150
+ timestamp_str = time.strftime(
1151
+ "%H:%M:%S",
1152
+ time.localtime(call_record["timestamp"]),
1153
+ )
1154
+
1155
+ yield StreamChunk(
1156
+ type="custom_tool_status",
1157
+ status="custom_tool_called",
1158
+ content=f"🔧 Custom Tool Called: {tool_name} at {timestamp_str} with args: {json.dumps(tool_args, indent=2)}",
1159
+ source="custom_tools",
1160
+ )
1161
+
1162
+ log_tool_call(
1163
+ agent_id,
1164
+ tool_name,
1165
+ tool_args,
1166
+ backend_name="gemini",
1167
+ )
1168
+ elif is_mcp_tool:
1169
+ # Process MCP tool call
1170
+ if mcp_tracker.is_new_call(tool_name, tool_args):
1171
+ call_record = mcp_tracker.add_call(tool_name, tool_args)
1172
+
1173
+ mcp_tools_used.append(
1174
+ {
1175
+ "name": tool_name,
1176
+ "arguments": tool_args,
1177
+ "timestamp": call_record["timestamp"],
1178
+ },
1179
+ )
1180
+
1181
+ timestamp_str = time.strftime(
1182
+ "%H:%M:%S",
1183
+ time.localtime(call_record["timestamp"]),
1184
+ )
1185
+
1186
+ yield StreamChunk(
1187
+ type="mcp_status",
1188
+ status="mcp_tool_called",
1189
+ content=f"🔧 MCP Tool Called: {tool_name} at {timestamp_str} with args: {json.dumps(tool_args, indent=2)}",
1190
+ source="mcp_tools",
1191
+ )
1192
+
1193
+ log_tool_call(
1194
+ agent_id,
1195
+ tool_name,
1196
+ tool_args,
1197
+ backend_name="gemini",
1198
+ )
1199
+
1200
+ # Check automatic_function_calling_history (for auto-executed MCP tools)
1201
+ if hasattr(chunk, "automatic_function_calling_history") and chunk.automatic_function_calling_history:
1202
+ for history_item in chunk.automatic_function_calling_history:
1203
+ if hasattr(history_item, "parts") and history_item.parts is not None:
1204
+ for part in history_item.parts:
1205
+ # Check for function_call part
1206
+ if hasattr(part, "function_call") and part.function_call:
1207
+ call_data = self.mcp_extractor.extract_function_call(part.function_call)
1208
+
1209
+ if call_data:
1210
+ tool_name = call_data["name"]
1211
+ tool_args = call_data["arguments"]
1212
+
1213
+ # Determine if it's MCP tool or custom tool
1214
+ # MCP tools may come from SDK without prefix
1215
+ is_mcp_tool = False
1216
+ if has_mcp:
1217
+ if tool_name in available_mcp_tools:
1218
+ is_mcp_tool = True
1219
+ else:
1220
+ # Try matching with MCP prefix format
1221
+ for mcp_tool in available_mcp_tools:
1222
+ if mcp_tool.startswith("mcp__") and mcp_tool.endswith(f"__{tool_name}"):
1223
+ is_mcp_tool = True
1224
+ break
1225
+
1226
+ is_custom_tool = has_custom_tools and tool_name in self._custom_tool_names
1227
+
1228
+ if is_mcp_tool:
1229
+ # Process MCP tool call
1230
+ if mcp_tracker.is_new_call(tool_name, tool_args):
1231
+ call_record = mcp_tracker.add_call(tool_name, tool_args)
1232
+
1233
+ mcp_tools_used.append(
1234
+ {
1235
+ "name": tool_name,
1236
+ "arguments": tool_args,
1237
+ "timestamp": call_record["timestamp"],
1238
+ },
1239
+ )
1240
+
1241
+ timestamp_str = time.strftime(
1242
+ "%H:%M:%S",
1243
+ time.localtime(call_record["timestamp"]),
1244
+ )
1245
+
1246
+ yield StreamChunk(
1247
+ type="mcp_status",
1248
+ status="mcp_tool_called",
1249
+ content=f"🔧 MCP Tool Called: {tool_name} at {timestamp_str} with args: {json.dumps(tool_args, indent=2)}",
1250
+ source="mcp_tools",
1251
+ )
1252
+
1253
+ log_tool_call(
1254
+ agent_id,
1255
+ tool_name,
1256
+ tool_args,
1257
+ backend_name="gemini",
1258
+ )
1259
+
1260
+ elif is_custom_tool:
1261
+ # Process custom tool call
1262
+ if custom_tracker.is_new_call(tool_name, tool_args):
1263
+ call_record = custom_tracker.add_call(tool_name, tool_args)
1264
+
1265
+ custom_tools_used.append(
1266
+ {
1267
+ "name": tool_name,
1268
+ "arguments": tool_args,
1269
+ "timestamp": call_record["timestamp"],
1270
+ },
1271
+ )
1272
+
1273
+ timestamp_str = time.strftime(
1274
+ "%H:%M:%S",
1275
+ time.localtime(call_record["timestamp"]),
1276
+ )
1277
+
1278
+ yield StreamChunk(
1279
+ type="custom_tool_status",
1280
+ status="custom_tool_called",
1281
+ content=f"🔧 Custom Tool Called: {tool_name} at {timestamp_str} with args: {json.dumps(tool_args, indent=2)}",
1282
+ source="custom_tools",
1283
+ )
1284
+
1285
+ log_tool_call(
1286
+ agent_id,
1287
+ tool_name,
1288
+ tool_args,
1289
+ backend_name="gemini",
1290
+ )
1291
+
1292
+ # Check for function_response part
1293
+ elif hasattr(part, "function_response") and part.function_response:
1294
+ response_data = self.mcp_extractor.extract_function_response(part.function_response)
1295
+
1296
+ if response_data:
1297
+ tool_name = response_data["name"]
1298
+ tool_response = response_data["response"]
1299
+
1300
+ # Determine if it's MCP tool or custom tool
1301
+ # MCP tools may come from SDK without prefix
1302
+ is_mcp_tool = False
1303
+ if has_mcp:
1304
+ if tool_name in available_mcp_tools:
1305
+ is_mcp_tool = True
1306
+ else:
1307
+ # Try matching with MCP prefix format
1308
+ for mcp_tool in available_mcp_tools:
1309
+ if mcp_tool.startswith("mcp__") and mcp_tool.endswith(f"__{tool_name}"):
1310
+ is_mcp_tool = True
1311
+ break
1312
+
1313
+ is_custom_tool = has_custom_tools and tool_name in self._custom_tool_names
1314
+
1315
+ if is_mcp_tool:
1316
+ # Process MCP tool response
1317
+ if mcp_response_tracker.is_new_response(tool_name, tool_response):
1318
+ response_record = mcp_response_tracker.add_response(tool_name, tool_response)
1319
+
1320
+ # Extract text content from CallToolResult
1321
+ response_text = None
1322
+ if isinstance(tool_response, dict) and "result" in tool_response:
1323
+ result = tool_response["result"]
1324
+ if hasattr(result, "content") and result.content:
1325
+ first_content = result.content[0]
1326
+ if hasattr(first_content, "text"):
1327
+ response_text = first_content.text
1328
+
1329
+ if response_text is None:
1330
+ response_text = str(tool_response)
1331
+
1332
+ timestamp_str = time.strftime(
1333
+ "%H:%M:%S",
1334
+ time.localtime(response_record["timestamp"]),
1335
+ )
1336
+
1337
+ # Format response as JSON if possible
1338
+ formatted_response = format_tool_response_as_json(response_text)
1339
+
1340
+ yield StreamChunk(
1341
+ type="mcp_status",
1342
+ status="mcp_tool_response",
1343
+ content=f"✅ MCP Tool Response from {tool_name} at {timestamp_str}: {formatted_response}",
1344
+ source="mcp_tools",
1345
+ )
1346
+
1347
+ log_backend_activity(
1348
+ "gemini",
1349
+ "MCP tool response received",
1350
+ {
1351
+ "tool_name": tool_name,
1352
+ "response_preview": str(tool_response)[:],
1353
+ },
1354
+ agent_id=agent_id,
1355
+ )
1356
+
1357
+ elif is_custom_tool:
1358
+ # Process custom tool response
1359
+ if custom_response_tracker.is_new_response(tool_name, tool_response):
1360
+ response_record = custom_response_tracker.add_response(tool_name, tool_response)
1361
+
1362
+ # Extract text from response
1363
+ response_text = str(tool_response)
1364
+
1365
+ timestamp_str = time.strftime(
1366
+ "%H:%M:%S",
1367
+ time.localtime(response_record["timestamp"]),
1368
+ )
1369
+
1370
+ # Format response as JSON if possible
1371
+ formatted_response = format_tool_response_as_json(response_text)
1372
+
1373
+ yield StreamChunk(
1374
+ type="custom_tool_status",
1375
+ status="custom_tool_response",
1376
+ content=f"✅ Custom Tool Response from {tool_name} at {timestamp_str}: {formatted_response}",
1377
+ source="custom_tools",
1378
+ )
1379
+
1380
+ log_backend_activity(
1381
+ "gemini",
1382
+ "Custom tool response received",
1383
+ {
1384
+ "tool_name": tool_name,
1385
+ "response_preview": str(tool_response),
1386
+ },
1387
+ agent_id=agent_id,
1388
+ )
1389
+
1390
+ # ============================================
1391
+ # Process text content
1392
+ # ============================================
1393
+ if hasattr(chunk, "text") and chunk.text:
1394
+ chunk_text = chunk.text
1395
+ full_content_text += chunk_text
1396
+ log_stream_chunk("backend.gemini", "continuation_content", chunk_text, agent_id)
1397
+ yield StreamChunk(type="content", content=chunk_text)
1398
+
1399
+ # ============================================
1400
+ # Buffer last chunk
1401
+ # ============================================
1402
+ if hasattr(chunk, "candidates") and chunk.candidates:
1403
+ last_response_with_candidates = chunk
1404
+
1405
+ except Exception as e:
1406
+ error_msg = f"Error in continuation call: {str(e)}"
1407
+ logger.error(error_msg)
1408
+ yield StreamChunk(
1409
+ type="custom_tool_status",
1410
+ status="continuation_error",
1411
+ content=f"❌ {error_msg}",
1412
+ source="custom_tools",
1413
+ )
1414
+
1415
+ # ====================================================================
1416
+ # Completion phase: Output summary
1417
+ # ====================================================================
1418
+
1419
+ # Add MCP usage indicator with detailed summary
1420
+ if has_mcp:
1421
+ mcp_summary = mcp_tracker.get_summary()
1422
+ if not mcp_summary or mcp_summary == "No MCP tools called":
1423
+ mcp_summary = "MCP session completed (no tools explicitly called)"
1424
+ else:
1425
+ mcp_summary = f"MCP session complete - {mcp_summary}"
1426
+
1427
+ log_stream_chunk("backend.gemini", "mcp_indicator", mcp_summary, agent_id)
1428
+ yield StreamChunk(
1429
+ type="mcp_status",
1430
+ status="mcp_session_complete",
1431
+ content=mcp_summary,
1432
+ source="mcp_tools",
1433
+ )
1434
+
1435
+ # Add custom tool usage indicator with detailed summary
1436
+ if has_custom_tools:
1437
+ custom_summary = custom_tracker.get_summary()
1438
+ if not custom_summary or custom_summary == "No MCP tools called":
1439
+ custom_summary = "Custom tools session completed (no tools explicitly called)"
1440
+ else:
1441
+ # Replace "MCP tool" with "Custom tool"
1442
+ custom_summary = custom_summary.replace("MCP tool", "Custom tool")
1443
+ custom_summary = f"Custom tools session complete - {custom_summary}"
1444
+
1445
+ log_stream_chunk("backend.gemini", "custom_tools_indicator", custom_summary, agent_id)
1446
+ yield StreamChunk(
1447
+ type="custom_tool_status",
1448
+ status="custom_tools_session_complete",
1449
+ content=custom_summary,
1450
+ source="custom_tools",
1451
+ )
1452
+
1771
1453
  except (
1772
1454
  MCPConnectionError,
1773
1455
  MCPTimeoutError,
@@ -1775,26 +1457,102 @@ Make your decision and include the JSON at the very end of your response."""
1775
1457
  MCPError,
1776
1458
  Exception,
1777
1459
  ) as e:
1778
- log_stream_chunk("backend.gemini", "mcp_error", str(e), agent_id)
1460
+ log_stream_chunk("backend.gemini", "tools_error", str(e), agent_id)
1461
+
1462
+ # ====================================================================
1463
+ # Error handling: Distinguish MCP and custom tools errors
1464
+ # ====================================================================
1465
+
1466
+ # Determine error type
1467
+ is_mcp_error = isinstance(e, (MCPConnectionError, MCPTimeoutError, MCPServerError, MCPError))
1468
+ is_custom_tool_error = not is_mcp_error and using_custom_tools
1779
1469
 
1780
1470
  # Emit user-friendly error message
1781
- async for chunk in self._handle_mcp_error_and_fallback(e):
1782
- yield chunk
1471
+ if is_mcp_error:
1472
+ async for chunk in self.mcp_manager.handle_mcp_error_and_fallback(e):
1473
+ yield chunk
1474
+ elif is_custom_tool_error:
1475
+ yield StreamChunk(
1476
+ type="custom_tool_status",
1477
+ status="custom_tools_error",
1478
+ content=f"⚠️ [Custom Tools] Error: {str(e)}; falling back to non-custom-tool mode",
1479
+ source="custom_tools",
1480
+ )
1481
+ else:
1482
+ yield StreamChunk(
1483
+ type="mcp_status",
1484
+ status="tools_error",
1485
+ content=f"⚠️ [Tools] Error: {str(e)}; falling back",
1486
+ source="tools",
1487
+ )
1783
1488
 
1784
- # Fallback to non-MCP streaming with manual configuration
1489
+ # Fallback configuration
1785
1490
  manual_config = dict(config)
1786
- if all_tools:
1787
- manual_config["tools"] = all_tools
1788
1491
 
1789
- # Need to create a new stream for fallback since stream is None
1790
- stream = await client.aio.models.generate_content_stream(model=model_name, contents=full_content, config=manual_config)
1492
+ # Decide fallback configuration based on error type
1493
+ if is_mcp_error and using_custom_tools:
1494
+ # MCP error but custom tools available: exclude MCP, keep custom tools
1495
+ try:
1496
+ custom_tools_schemas = self._get_custom_tools_schemas()
1497
+ if custom_tools_schemas:
1498
+ # Convert to Gemini format using formatter
1499
+ custom_tools_functions = self.formatter.format_custom_tools(
1500
+ custom_tools_schemas,
1501
+ return_sdk_objects=True,
1502
+ )
1503
+ # Wrap FunctionDeclarations in a Tool object for Gemini SDK
1504
+ from google.genai import types
1505
+
1506
+ custom_tool = types.Tool(function_declarations=custom_tools_functions)
1507
+ manual_config["tools"] = [custom_tool]
1508
+ logger.info("[Gemini] Fallback: using custom tools only (MCP failed)")
1509
+ else:
1510
+ # Custom tools also unavailable, use builtin tools
1511
+ if all_tools:
1512
+ manual_config["tools"] = all_tools
1513
+ logger.info("[Gemini] Fallback: using builtin tools only (both MCP and custom tools failed)")
1514
+ except Exception:
1515
+ if all_tools:
1516
+ manual_config["tools"] = all_tools
1517
+ logger.info("[Gemini] Fallback: using builtin tools only (custom tools also failed)")
1518
+
1519
+ elif is_custom_tool_error and using_sdk_mcp:
1520
+ # Custom tools error but MCP available: exclude custom tools, keep MCP
1521
+ try:
1522
+ if self._mcp_client:
1523
+ mcp_sessions = self.mcp_manager.get_active_mcp_sessions(
1524
+ convert_to_permission_sessions=bool(self.filesystem_manager),
1525
+ )
1526
+ if mcp_sessions:
1527
+ manual_config["tools"] = mcp_sessions
1528
+ logger.info("[Gemini] Fallback: using MCP only (custom tools failed)")
1529
+ else:
1530
+ if all_tools:
1531
+ manual_config["tools"] = all_tools
1532
+ logger.info("[Gemini] Fallback: using builtin tools only (both custom tools and MCP failed)")
1533
+ except Exception:
1534
+ if all_tools:
1535
+ manual_config["tools"] = all_tools
1536
+ logger.info("[Gemini] Fallback: using builtin tools only (MCP also failed)")
1537
+
1538
+ else:
1539
+ # Both failed or cannot determine: use builtin tools
1540
+ if all_tools:
1541
+ manual_config["tools"] = all_tools
1542
+ logger.info("[Gemini] Fallback: using builtin tools only (all advanced tools failed)")
1543
+
1544
+ # Create new stream for fallback
1545
+ stream = await client.aio.models.generate_content_stream(
1546
+ model=model_name,
1547
+ contents=full_content,
1548
+ config=manual_config,
1549
+ )
1791
1550
 
1792
1551
  async for chunk in stream:
1793
1552
  # Process text content
1794
1553
  if hasattr(chunk, "text") and chunk.text:
1795
1554
  chunk_text = chunk.text
1796
1555
  full_content_text += chunk_text
1797
- # Log fallback content chunks
1798
1556
  log_stream_chunk(
1799
1557
  "backend.gemini",
1800
1558
  "fallback_content",
@@ -1802,30 +1560,46 @@ Make your decision and include the JSON at the very end of your response."""
1802
1560
  agent_id,
1803
1561
  )
1804
1562
  yield StreamChunk(type="content", content=chunk_text)
1805
-
1563
+ # Buffer last chunk with candidates for fallback path
1564
+ if hasattr(chunk, "candidates") and chunk.candidates:
1565
+ last_response_with_candidates = chunk
1806
1566
  else:
1807
- # Non-MCP path (existing behavior)
1808
- # Create stream for non-MCP path
1809
- stream = await client.aio.models.generate_content_stream(model=model_name, contents=full_content, config=config)
1810
-
1811
- async for chunk in stream:
1812
- # ============================================
1813
- # 1. Process text content
1814
- # ============================================
1815
- if hasattr(chunk, "text") and chunk.text:
1816
- chunk_text = chunk.text
1817
- full_content_text += chunk_text
1818
-
1819
- # Enhanced logging for non-MCP streaming chunks
1820
- log_stream_chunk("backend.gemini", "content", chunk_text, agent_id)
1821
- log_backend_agent_message(
1822
- agent_id,
1823
- "RECV",
1824
- {"content": chunk_text},
1825
- backend_name="gemini",
1826
- )
1567
+ # Non-MCP streaming path: execute when MCP is disabled
1568
+ try:
1569
+ # Use the standard config (with builtin tools if configured)
1570
+ stream = await client.aio.models.generate_content_stream(
1571
+ model=model_name,
1572
+ contents=full_content,
1573
+ config=config,
1574
+ )
1575
+
1576
+ # Process streaming chunks
1577
+ async for chunk in stream:
1578
+ # Process text content
1579
+ if hasattr(chunk, "text") and chunk.text:
1580
+ chunk_text = chunk.text
1581
+ full_content_text += chunk_text
1582
+ log_backend_agent_message(
1583
+ agent_id,
1584
+ "RECV",
1585
+ {"content": chunk_text},
1586
+ backend_name="gemini",
1587
+ )
1588
+ log_stream_chunk("backend.gemini", "content", chunk_text, agent_id)
1589
+ yield StreamChunk(type="content", content=chunk_text)
1590
+ # Buffer last chunk with candidates for non-MCP path
1591
+ if hasattr(chunk, "candidates") and chunk.candidates:
1592
+ last_response_with_candidates = chunk
1827
1593
 
1828
- yield StreamChunk(type="content", content=chunk_text)
1594
+ except Exception as e:
1595
+ error_msg = f"Non-MCP streaming error: {e}"
1596
+ log_stream_chunk(
1597
+ "backend.gemini",
1598
+ "non_mcp_stream_error",
1599
+ {"error_type": type(e).__name__, "error_message": str(e)},
1600
+ agent_id,
1601
+ )
1602
+ yield StreamChunk(type="error", error=error_msg)
1829
1603
 
1830
1604
  content = full_content_text
1831
1605
 
@@ -1842,11 +1616,11 @@ Make your decision and include the JSON at the very end of your response."""
1842
1616
  structured_response = json.loads(content.strip())
1843
1617
  except json.JSONDecodeError:
1844
1618
  # Strategy 2: Extract JSON from mixed text content (handles markdown-wrapped JSON)
1845
- structured_response = self.extract_structured_response(content)
1619
+ structured_response = self.formatter.extract_structured_response(content)
1846
1620
 
1847
1621
  if structured_response and isinstance(structured_response, dict) and "action_type" in structured_response:
1848
1622
  # Convert to tool calls
1849
- tool_calls = self.convert_structured_to_tool_calls(structured_response)
1623
+ tool_calls = self.formatter.convert_structured_to_tool_calls(structured_response)
1850
1624
  if tool_calls:
1851
1625
  tool_calls_detected = tool_calls
1852
1626
  # Log conversion to tool calls (summary)
@@ -1866,6 +1640,10 @@ Make your decision and include the JSON at the very end of your response."""
1866
1640
  # Ensure logging does not interrupt flow
1867
1641
  pass
1868
1642
 
1643
+ # Assign buffered final response (if available) so builtin tool indicators can be emitted
1644
+ if last_response_with_candidates is not None:
1645
+ final_response = last_response_with_candidates
1646
+
1869
1647
  # Process builtin tool results if any tools were used
1870
1648
  if builtin_tools and final_response and hasattr(final_response, "candidates") and final_response.candidates:
1871
1649
  # Check for grounding or code execution results
@@ -2054,7 +1832,7 @@ Make your decision and include the JSON at the very end of your response."""
2054
1832
  yield StreamChunk(type="error", error=error_msg)
2055
1833
  finally:
2056
1834
  # Cleanup resources
2057
- await self._cleanup_resources(stream, client)
1835
+ await self.mcp_manager.cleanup_genai_resources(stream, client)
2058
1836
  # Ensure context manager exit for MCP cleanup
2059
1837
  try:
2060
1838
  await self.__aexit__(None, None, None)
@@ -2132,7 +1910,7 @@ Make your decision and include the JSON at the very end of your response."""
2132
1910
  super().reset_token_usage()
2133
1911
 
2134
1912
  async def cleanup_mcp(self):
2135
- """Cleanup MCP connections."""
1913
+ """Cleanup MCP connections - override parent class to use Gemini-specific cleanup."""
2136
1914
  if self._mcp_client:
2137
1915
  try:
2138
1916
  await self._mcp_client.disconnect()
@@ -2148,90 +1926,16 @@ Make your decision and include the JSON at the very end of your response."""
2148
1926
  finally:
2149
1927
  self._mcp_client = None
2150
1928
  self._mcp_initialized = False
2151
-
2152
- async def _cleanup_resources(self, stream, client):
2153
- """Cleanup google-genai resources to avoid unclosed aiohttp sessions."""
2154
- # Close stream
2155
- try:
2156
- if stream is not None:
2157
- close_fn = getattr(stream, "aclose", None) or getattr(stream, "close", None)
2158
- if close_fn is not None:
2159
- maybe = close_fn()
2160
- if hasattr(maybe, "__await__"):
2161
- await maybe
2162
- except Exception as e:
2163
- log_backend_activity(
2164
- "gemini",
2165
- "Stream cleanup failed",
2166
- {"error": str(e)},
2167
- agent_id=self.agent_id,
2168
- )
2169
- # Close internal aiohttp session held by google-genai BaseApiClient
2170
- try:
2171
- if client is not None:
2172
- base_client = getattr(client, "_api_client", None)
2173
- if base_client is not None:
2174
- session = getattr(base_client, "_aiohttp_session", None)
2175
- if session is not None and hasattr(session, "close"):
2176
- if not session.closed:
2177
- await session.close()
2178
- log_backend_activity(
2179
- "gemini",
2180
- "Closed google-genai aiohttp session",
2181
- {},
2182
- agent_id=self.agent_id,
2183
- )
2184
- base_client._aiohttp_session = None
2185
- # Yield control to allow connector cleanup
2186
- await asyncio.sleep(0)
2187
- except Exception as e:
2188
- log_backend_activity(
2189
- "gemini",
2190
- "Failed to close google-genai aiohttp session",
2191
- {"error": str(e)},
2192
- agent_id=self.agent_id,
2193
- )
2194
- # Close internal async transport if exposed
2195
- try:
2196
- if client is not None and hasattr(client, "aio") and client.aio is not None:
2197
- aio_obj = client.aio
2198
- for method_name in ("close", "stop"):
2199
- method = getattr(aio_obj, method_name, None)
2200
- if method:
2201
- maybe = method()
2202
- if hasattr(maybe, "__await__"):
2203
- await maybe
2204
- break
2205
- except Exception as e:
2206
- log_backend_activity(
2207
- "gemini",
2208
- "Client AIO cleanup failed",
2209
- {"error": str(e)},
2210
- agent_id=self.agent_id,
2211
- )
2212
-
2213
- # Close client
2214
- try:
2215
- if client is not None:
2216
- for method_name in ("aclose", "close"):
2217
- method = getattr(client, method_name, None)
2218
- if method:
2219
- maybe = method()
2220
- if hasattr(maybe, "__await__"):
2221
- await maybe
2222
- break
2223
- except Exception as e:
2224
- log_backend_activity(
2225
- "gemini",
2226
- "Client cleanup failed",
2227
- {"error": str(e)},
2228
- agent_id=self.agent_id,
2229
- )
1929
+ # Also clear parent class attributes if they exist (for compatibility)
1930
+ if hasattr(self, "_mcp_functions"):
1931
+ self._mcp_functions.clear()
1932
+ if hasattr(self, "_mcp_function_names"):
1933
+ self._mcp_function_names.clear()
2230
1934
 
2231
1935
  async def __aenter__(self) -> "GeminiBackend":
2232
1936
  """Async context manager entry."""
2233
1937
  try:
2234
- await self._setup_mcp_tools(agent_id=self.agent_id)
1938
+ await self.mcp_manager.setup_mcp_tools(agent_id=self.agent_id)
2235
1939
  except Exception as e:
2236
1940
  log_backend_activity(
2237
1941
  "gemini",