Skip to content

Commit 6a7f8f4

Browse files
ivnvxdivnvxd
committed
refactor: Improve YOLO mode response structure
1 parent 29e0f1a commit 6a7f8f4

File tree

3 files changed

+108
-77
lines changed

3 files changed

+108
-77
lines changed

‎mcp_server_odoo/tools.py

Lines changed: 36 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -763,7 +763,7 @@ async def _handle_get_record_tool(
763763
sanitized_msg = ErrorSanitizer.sanitize_message(str(e))
764764
raise ToolError(f"Failed to get record: {sanitized_msg}") from e
765765

766-
async def _handle_list_models_tool(self) -> Dict[str, List[Dict[str, Any]]]:
766+
async def _handle_list_models_tool(self) -> Dict[str, Any]:
767767
"""Handle list models tool request with permissions."""
768768
try:
769769
with perf_logger.track_operation("tool_list_models"):
@@ -800,55 +800,66 @@ async def _handle_list_models_tool(self) -> Dict[str, List[Dict[str, Any]]]:
800800
limit=200, # Reasonable limit for practical use
801801
)
802802

803-
# Prepare response with YOLO mode indicator
804-
enriched_models = []
803+
# Prepare response with YOLO mode metadata
805804
mode_desc = (
806805
"READ-ONLY" if self.config.yolo_mode == "read" else "FULL ACCESS"
807806
)
808807

809-
# Add warning as first "model" in the list
810-
warning_model = {
811-
"model": "⚠️ YOLO MODE",
812-
"name": f"🚨 {mode_desc} - All models accessible without MCP security!",
808+
# Create metadata about YOLO mode
809+
yolo_metadata = {
810+
"enabled": True,
811+
"level": self.config.yolo_mode, # "read" or "true"
812+
"description": mode_desc,
813+
"warning": "🚨 All models accessible without MCP security!",
813814
"operations": {
814815
"read": True,
815816
"write": self.config.yolo_mode == "true",
816817
"create": self.config.yolo_mode == "true",
817818
"unlink": self.config.yolo_mode == "true",
818819
},
819820
}
820-
enriched_models.append(warning_model)
821821

822-
# Process actual models (no need to repeat permissions in YOLO mode)
822+
# Process actual models (clean data without permissions)
823+
models_list = []
823824
for record in model_records:
824-
enriched_model = {
825+
model_entry = {
825826
"model": record["model"],
826827
"name": record["name"] or record["model"],
827828
}
828-
enriched_models.append(enriched_model)
829+
models_list.append(model_entry)
829830

830831
logger.info(
831832
f"YOLO mode ({mode_desc}): Listed {len(model_records)} models from database"
832833
)
833834

834-
return {"models": enriched_models}
835+
return {
836+
"yolo_mode": yolo_metadata,
837+
"models": models_list,
838+
"total": len(models_list),
839+
}
835840

836841
except Exception as e:
837842
logger.error(f"Failed to query models in YOLO mode: {e}")
838-
# Fall back to returning a warning with error
843+
# Return error in consistent structure
844+
mode_desc = (
845+
"READ-ONLY" if self.config.yolo_mode == "read" else "FULL ACCESS"
846+
)
839847
return {
840-
"models": [
841-
{
842-
"model": "⚠️ ERROR",
843-
"name": f"Failed to query models: {str(e)}",
844-
"operations": {
845-
"read": False,
846-
"write": False,
847-
"create": False,
848-
"unlink": False,
849-
},
850-
}
851-
]
848+
"yolo_mode": {
849+
"enabled": True,
850+
"level": self.config.yolo_mode,
851+
"description": mode_desc,
852+
"warning": f"⚠️ Error querying models: {str(e)}",
853+
"operations": {
854+
"read": False,
855+
"write": False,
856+
"create": False,
857+
"unlink": False,
858+
},
859+
},
860+
"models": [],
861+
"total": 0,
862+
"error": str(e),
852863
}
853864

854865
# Standard mode: Get models from MCP access controller

‎tests/test_e2e_yolo.py

Lines changed: 22 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -86,12 +86,14 @@ async def test_yolo_complete_workflow_read_only(self, config_read_only):
8686
# 2. List models - should work and show indicator
8787
models_result = await handler._handle_list_models_tool()
8888
assert "models" in models_result
89+
assert "yolo_mode" in models_result
8990
assert len(models_result["models"]) > 0
9091

91-
# Check for YOLO warning
92-
first_model = models_result["models"][0]
93-
assert "YOLO MODE" in first_model["model"]
94-
assert "READ-ONLY" in first_model["name"]
92+
# Check for YOLO metadata
93+
yolo_meta = models_result["yolo_mode"]
94+
assert yolo_meta["enabled"] is True
95+
assert yolo_meta["level"] == "read"
96+
assert "READ-ONLY" in yolo_meta["description"]
9597

9698
# 3. Search records - should work
9799
search_result = await handler._handle_search_tool(
@@ -165,11 +167,13 @@ async def test_yolo_complete_workflow_full_access(self, config_full_access):
165167
# 2. List models - should work and show warning
166168
models_result = await handler._handle_list_models_tool()
167169
assert "models" in models_result
170+
assert "yolo_mode" in models_result
168171

169-
# Check for YOLO warning
170-
first_model = models_result["models"][0]
171-
assert "YOLO MODE" in first_model["model"]
172-
assert "FULL ACCESS" in first_model["name"]
172+
# Check for YOLO metadata
173+
yolo_meta = models_result["yolo_mode"]
174+
assert yolo_meta["enabled"] is True
175+
assert yolo_meta["level"] == "true"
176+
assert "FULL ACCESS" in yolo_meta["description"]
173177

174178
# 3. Create a test record - should work
175179
create_result = await handler._handle_create_record_tool(
@@ -367,10 +371,10 @@ async def test_mode_indicators_in_responses(self, config_read_only, config_full_
367371

368372
# Check list_models indicator
369373
models_result = await handler._handle_list_models_tool()
370-
first_model = models_result["models"][0]
371-
assert "READ-ONLY" in first_model["name"]
372-
assert first_model["operations"]["read"] is True
373-
assert first_model["operations"]["write"] is False
374+
yolo_meta = models_result["yolo_mode"]
375+
assert "READ-ONLY" in yolo_meta["description"]
376+
assert yolo_meta["operations"]["read"] is True
377+
assert yolo_meta["operations"]["write"] is False
374378

375379
connection.disconnect()
376380

@@ -384,14 +388,14 @@ async def test_mode_indicators_in_responses(self, config_read_only, config_full_
384388

385389
# Check list_models indicator
386390
models_result = await handler._handle_list_models_tool()
387-
first_model = models_result["models"][0]
388-
assert "FULL ACCESS" in first_model["name"]
391+
yolo_meta = models_result["yolo_mode"]
392+
assert "FULL ACCESS" in yolo_meta["description"]
389393
assert all(
390394
[
391-
first_model["operations"]["read"],
392-
first_model["operations"]["write"],
393-
first_model["operations"]["create"],
394-
first_model["operations"]["unlink"],
395+
yolo_meta["operations"]["read"],
396+
yolo_meta["operations"]["write"],
397+
yolo_meta["operations"]["create"],
398+
yolo_meta["operations"]["unlink"],
395399
]
396400
)
397401

‎tests/test_tools_yolo.py

Lines changed: 50 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -96,21 +96,25 @@ async def test_list_models_yolo_read_mode(
9696
assert "transient" in str(call_args[0][1]) # Domain includes transient filter
9797

9898
# Check result structure
99+
assert "yolo_mode" in result
99100
assert "models" in result
101+
assert "total" in result
102+
103+
# Check YOLO mode metadata
104+
yolo_meta = result["yolo_mode"]
105+
assert yolo_meta["enabled"] is True
106+
assert yolo_meta["level"] == "read"
107+
assert "READ-ONLY" in yolo_meta["description"]
108+
assert "🚨" in yolo_meta["warning"]
109+
assert yolo_meta["operations"]["read"] is True
110+
assert yolo_meta["operations"]["write"] is False
111+
assert yolo_meta["operations"]["create"] is False
112+
assert yolo_meta["operations"]["unlink"] is False
113+
114+
# Check actual models are clean (no operations field)
100115
models = result["models"]
101-
assert len(models) > 0
102-
103-
# Check for YOLO warning as first model
104-
warning_model = models[0]
105-
assert "YOLO MODE" in warning_model["model"]
106-
assert "READ-ONLY" in warning_model["name"]
107-
assert warning_model["operations"]["read"] is True
108-
assert warning_model["operations"]["write"] is False
109-
assert warning_model["operations"]["create"] is False
110-
assert warning_model["operations"]["unlink"] is False
111-
112-
# Check actual models don't have operations field (only in warning model)
113-
for model in models[1:]:
116+
assert len(models) == 3 # Should match mock data
117+
for model in models:
114118
assert "operations" not in model
115119
assert "model" in model
116120
assert "name" in model
@@ -135,20 +139,25 @@ async def test_list_models_yolo_full_mode(
135139
result = await handler._handle_list_models_tool()
136140

137141
# Check result structure
142+
assert "yolo_mode" in result
138143
assert "models" in result
144+
assert "total" in result
145+
146+
# Check YOLO mode metadata
147+
yolo_meta = result["yolo_mode"]
148+
assert yolo_meta["enabled"] is True
149+
assert yolo_meta["level"] == "true"
150+
assert "FULL ACCESS" in yolo_meta["description"]
151+
assert "🚨" in yolo_meta["warning"]
152+
assert yolo_meta["operations"]["read"] is True
153+
assert yolo_meta["operations"]["write"] is True
154+
assert yolo_meta["operations"]["create"] is True
155+
assert yolo_meta["operations"]["unlink"] is True
156+
157+
# Check actual models are clean (no operations field)
139158
models = result["models"]
140-
141-
# Check for YOLO warning as first model
142-
warning_model = models[0]
143-
assert "YOLO MODE" in warning_model["model"]
144-
assert "FULL ACCESS" in warning_model["name"]
145-
assert warning_model["operations"]["read"] is True
146-
assert warning_model["operations"]["write"] is True
147-
assert warning_model["operations"]["create"] is True
148-
assert warning_model["operations"]["unlink"] is True
149-
150-
# Check actual models don't have operations field (only in warning model)
151-
for model in models[1:]:
159+
assert len(models) == 2 # Should match mock data
160+
for model in models:
152161
assert "operations" not in model
153162
assert "model" in model
154163
assert "name" in model
@@ -214,16 +223,23 @@ async def test_list_models_yolo_error_handling(
214223
# Call the method
215224
result = await handler._handle_list_models_tool()
216225

217-
# Check error response
226+
# Check error response structure
227+
assert "yolo_mode" in result
218228
assert "models" in result
219-
models = result["models"]
220-
assert len(models) == 1
229+
assert "error" in result
230+
231+
# Check YOLO mode metadata in error case
232+
yolo_meta = result["yolo_mode"]
233+
assert yolo_meta["enabled"] is True
234+
assert yolo_meta["level"] == "read"
235+
assert "Error querying models" in yolo_meta["warning"]
236+
assert yolo_meta["operations"]["read"] is False
237+
assert yolo_meta["operations"]["write"] is False
221238

222-
error_model = models[0]
223-
assert "ERROR" in error_model["model"]
224-
assert "Failed to query models" in error_model["name"]
225-
assert error_model["operations"]["read"] is False
226-
assert error_model["operations"]["write"] is False
239+
# Models should be empty on error
240+
assert result["models"] == []
241+
assert result["total"] == 0
242+
assert "Database connection failed" in result["error"]
227243

228244
@pytest.mark.asyncio
229245
async def test_list_models_yolo_domain_construction(
@@ -272,7 +288,7 @@ async def test_list_models_yolo_includes_common_system_models(
272288

273289
# Check that system models are included
274290
models = result["models"]
275-
model_names = [m["model"] for m in models[1:]] # Skip warning model
291+
model_names = [m["model"] for m in models]
276292
assert "ir.attachment" in model_names
277293
assert "ir.model" in model_names
278294

0 commit comments

Comments
 (0)