fix(mcp): validate UUID inputs and add sync_retain tool (#906)
* fix(mcp): validate UUID inputs at engine level and add sync_retain tool (#888) - Add UUID validation in memory_engine for get_memory_unit, delete_memory_unit, get_mental_model, delete_mental_model, get_mental_model_history (raises ValueError) - Catch ValueError → 400 in HTTP route handlers - Add sync_retain MCP tool that calls retain_batch_async directly for immediate availability (no polling needed) - Register sync_retain in _ALL_TOOLS, _SINGLE_BANK_TOOLS, UI MCP_TOOL_GROUPS - Add code-review check for MCP tool registration completeness * fix: remove UUID validation for mental model IDs (column is TEXT, not UUID) Mental model IDs are TEXT columns that accept arbitrary string IDs (e.g., 'team-communication-preferences'). UUID validation was incorrectly added to get_mental_model, delete_mental_model, and get_mental_model_history.
This commit is contained in:
Nicolò Boschi
GitHub
parent
7e23f8e149
commit
48185a4bee
7 changed files with 218 additions and 6 deletions
|
|
@ -157,7 +157,16 @@ If any files in `hindsight-integrations/` were added or changed, verify:
|
|||
- **Release process** — check that the integration name is in the `VALID_INTEGRATIONS` array in `scripts/release-integration.sh`. If missing, flag it.
|
||||
- **Code standards** — the integration code must follow all Python style rules (type hints, no raw dicts, no tuple returns, etc.).
|
||||
|
||||
### 10. Review against other coding standards
|
||||
### 10. Check MCP tool registration completeness
|
||||
|
||||
If any new MCP tools were added or existing tools renamed in `hindsight-api-slim/hindsight_api/mcp_tools.py`:
|
||||
- **`_ALL_TOOLS` set** in `mcp_tools.py` — must include the new tool name
|
||||
- **`tools_to_register` default set** in `register_mcp_tools()` in `mcp_tools.py` — must include the new tool name
|
||||
- **`_SINGLE_BANK_TOOLS` set** in `hindsight-api-slim/hindsight_api/api/mcp.py` — must include the new tool if it is bank-scoped (not a bank-management tool like `list_banks`/`create_bank`)
|
||||
- **`MCP_TOOL_GROUPS`** in `hindsight-control-plane/src/components/bank-config-view.tsx` — must include the new tool in the appropriate group for the UI tool selector
|
||||
- **Tool count assertions** in tests (e.g., `test_mcp_tools.py`) — must be updated to reflect the new count
|
||||
|
||||
### 11. Review against other coding standards
|
||||
|
||||
Check the diff for violations of the standards listed above:
|
||||
- Python files at project root (not allowed)
|
||||
|
|
@ -169,7 +178,7 @@ Check the diff for violations of the standards listed above:
|
|||
- Premature abstractions or speculative helpers
|
||||
- Backwards-compatibility hacks (unused vars, re-exports, "removed" comments)
|
||||
|
||||
### 11. Report findings
|
||||
### 12. Report findings
|
||||
|
||||
Present a clear summary organized by severity:
|
||||
|
||||
|
|
|
|||
|
|
@ -2679,6 +2679,8 @@ def _register_routes(app: FastAPI):
|
|||
return data
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except Exception as e:
|
||||
|
|
@ -3289,6 +3291,8 @@ def _register_routes(app: FastAPI):
|
|||
raise
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
|
|
@ -3322,6 +3326,8 @@ def _register_routes(app: FastAPI):
|
|||
raise
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
|
|
@ -3486,6 +3492,8 @@ def _register_routes(app: FastAPI):
|
|||
return {"status": "deleted"}
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -97,6 +97,7 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
|||
_SINGLE_BANK_TOOLS: frozenset[str] = frozenset(
|
||||
{
|
||||
"retain",
|
||||
"sync_retain",
|
||||
"recall",
|
||||
"reflect",
|
||||
"list_mental_models",
|
||||
|
|
|
|||
|
|
@ -3792,7 +3792,14 @@ class MemoryEngine(MemoryEngineInterface):
|
|||
|
||||
Returns:
|
||||
Dictionary with deletion result
|
||||
|
||||
Raises:
|
||||
ValueError: If unit_id is not a valid UUID
|
||||
"""
|
||||
try:
|
||||
unit_uuid = uuid.UUID(unit_id)
|
||||
except ValueError:
|
||||
raise ValueError(f"Invalid unit_id: '{unit_id}' is not a valid UUID")
|
||||
await self._authenticate_tenant(request_context)
|
||||
pool = await self._get_pool()
|
||||
invalidated_obs = 0
|
||||
|
|
@ -3802,7 +3809,7 @@ class MemoryEngine(MemoryEngineInterface):
|
|||
# Get bank_id and fact_type before deletion
|
||||
row = await conn.fetchrow(
|
||||
f"SELECT bank_id, fact_type FROM {fq_table('memory_units')} WHERE id = $1",
|
||||
unit_id,
|
||||
str(unit_uuid),
|
||||
)
|
||||
bank_id = row["bank_id"] if row else None
|
||||
fact_type = row["fact_type"] if row else None
|
||||
|
|
@ -4697,7 +4704,14 @@ class MemoryEngine(MemoryEngineInterface):
|
|||
|
||||
Returns:
|
||||
Dict with memory unit data or None if not found
|
||||
|
||||
Raises:
|
||||
ValueError: If memory_id is not a valid UUID
|
||||
"""
|
||||
try:
|
||||
memory_uuid = uuid.UUID(memory_id)
|
||||
except ValueError:
|
||||
raise ValueError(f"Invalid memory_id: '{memory_id}' is not a valid UUID")
|
||||
await self._authenticate_tenant(request_context)
|
||||
if self._operation_validator:
|
||||
from hindsight_api.extensions import BankReadContext
|
||||
|
|
@ -4715,7 +4729,7 @@ class MemoryEngine(MemoryEngineInterface):
|
|||
FROM {fq_table("memory_units")}
|
||||
WHERE id = $1 AND bank_id = $2
|
||||
""",
|
||||
memory_id,
|
||||
str(memory_uuid),
|
||||
bank_id,
|
||||
)
|
||||
|
||||
|
|
@ -6497,6 +6511,7 @@ class MemoryEngine(MemoryEngineInterface):
|
|||
|
||||
Returns None if the mental model is not found.
|
||||
Returns a list of history entries (most recent first), each with previous_content and changed_at.
|
||||
|
||||
"""
|
||||
await self._authenticate_tenant(request_context)
|
||||
pool = await self._get_pool()
|
||||
|
|
|
|||
|
|
@ -29,6 +29,7 @@ from hindsight_api.models import RequestContext
|
|||
_ALL_TOOLS: frozenset[str] = frozenset(
|
||||
{
|
||||
"retain",
|
||||
"sync_retain",
|
||||
"recall",
|
||||
"reflect",
|
||||
"list_banks",
|
||||
|
|
@ -202,6 +203,7 @@ def register_mcp_tools(
|
|||
"""
|
||||
tools_to_register = config.tools or {
|
||||
"retain",
|
||||
"sync_retain",
|
||||
"recall",
|
||||
"reflect",
|
||||
"list_banks",
|
||||
|
|
@ -235,6 +237,9 @@ def register_mcp_tools(
|
|||
if "retain" in tools_to_register:
|
||||
_register_retain(mcp, memory, config)
|
||||
|
||||
if "sync_retain" in tools_to_register:
|
||||
_register_sync_retain(mcp, memory, config)
|
||||
|
||||
if "recall" in tools_to_register:
|
||||
_register_recall(mcp, memory, config)
|
||||
|
||||
|
|
@ -630,6 +635,124 @@ def _register_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
|||
return {"status": "error", "message": str(e)}
|
||||
|
||||
|
||||
def _register_sync_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
"""Register the sync_retain tool (synchronous retain that waits for completion)."""
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
async def sync_retain(
|
||||
content: str,
|
||||
context: str = "general",
|
||||
timestamp: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
metadata: dict[str, str] | None = None,
|
||||
document_id: str | None = None,
|
||||
bank_id: str | None = None,
|
||||
strategy: str | None = None,
|
||||
) -> dict:
|
||||
"""Store information to long-term memory and wait for completion.
|
||||
|
||||
Unlike retain (which is asynchronous), this tool blocks until the memory
|
||||
is fully stored and immediately available for recall.
|
||||
|
||||
Args:
|
||||
content: The fact/memory to store (be specific and include relevant details)
|
||||
context: Category for the memory (e.g., 'preferences', 'work', 'hobbies', 'family'). Default: 'general'
|
||||
timestamp: When this event/fact occurred (ISO format, e.g., '2024-01-15T10:30:00Z'). Useful for timeline tracking.
|
||||
tags: Optional tags for scoped visibility filtering (e.g., ['project:alpha', 'user:123'])
|
||||
metadata: Optional key-value metadata to attach (e.g., {'source': 'slack', 'channel': 'general'})
|
||||
document_id: Optional document ID to associate this memory with
|
||||
bank_id: Optional bank to store in (defaults to session bank). Use for cross-bank operations.
|
||||
strategy: Optional named retain strategy (e.g., 'exact' for verbatim storage). Strategies are defined in the bank config.
|
||||
"""
|
||||
target_bank = bank_id or config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return {"status": "error", "message": "No bank_id configured"}
|
||||
|
||||
content_dict, error = build_content_dict(content, context, timestamp, tags, metadata, document_id, strategy)
|
||||
if error:
|
||||
return {"status": "error", "message": error}
|
||||
|
||||
request_context = _get_request_context(config)
|
||||
|
||||
try:
|
||||
result = await memory.retain_batch_async(
|
||||
bank_id=target_bank,
|
||||
contents=[content_dict],
|
||||
request_context=request_context,
|
||||
strategy=content_dict.pop("strategy", None),
|
||||
)
|
||||
memory_ids = [uid for batch in result for uid in batch]
|
||||
return {
|
||||
"status": "completed",
|
||||
"message": "Memory stored successfully",
|
||||
"memory_ids": memory_ids,
|
||||
}
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Sync retain rejected: {e}")
|
||||
return {"status": "error", "message": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error in sync retain: {e}", exc_info=True)
|
||||
return {"status": "error", "message": str(e)}
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
async def sync_retain(
|
||||
content: str,
|
||||
context: str = "general",
|
||||
timestamp: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
metadata: dict[str, str] | None = None,
|
||||
document_id: str | None = None,
|
||||
strategy: str | None = None,
|
||||
) -> dict:
|
||||
"""Store information to long-term memory and wait for completion.
|
||||
|
||||
Unlike retain (which is asynchronous), this tool blocks until the memory
|
||||
is fully stored and immediately available for recall.
|
||||
|
||||
Args:
|
||||
content: The fact/memory to store (be specific and include relevant details)
|
||||
context: Category for the memory (e.g., 'preferences', 'work', 'hobbies', 'family'). Default: 'general'
|
||||
timestamp: When this event/fact occurred (ISO format, e.g., '2024-01-15T10:30:00Z'). Useful for timeline tracking.
|
||||
tags: Optional tags for scoped visibility filtering (e.g., ['project:alpha', 'user:123'])
|
||||
metadata: Optional key-value metadata to attach (e.g., {'source': 'slack', 'channel': 'general'})
|
||||
document_id: Optional document ID to associate this memory with
|
||||
strategy: Optional named retain strategy (e.g., 'exact' for verbatim storage). Strategies are defined in the bank config.
|
||||
"""
|
||||
target_bank = config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return {"status": "error", "message": "No bank_id configured"}
|
||||
|
||||
content_dict, error = build_content_dict(content, context, timestamp, tags, metadata, document_id, strategy)
|
||||
if error:
|
||||
return {"status": "error", "message": error}
|
||||
|
||||
request_context = _get_request_context(config)
|
||||
|
||||
try:
|
||||
result = await memory.retain_batch_async(
|
||||
bank_id=target_bank,
|
||||
contents=[content_dict],
|
||||
request_context=request_context,
|
||||
strategy=content_dict.pop("strategy", None),
|
||||
)
|
||||
memory_ids = [uid for batch in result for uid in batch]
|
||||
return {
|
||||
"status": "completed",
|
||||
"message": "Memory stored successfully",
|
||||
"memory_ids": memory_ids,
|
||||
}
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Sync retain rejected: {e}")
|
||||
return {"status": "error", "message": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error in sync retain: {e}", exc_info=True)
|
||||
return {"status": "error", "message": str(e)}
|
||||
|
||||
|
||||
def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
"""Register the recall tool."""
|
||||
description = config.recall_description or DEFAULT_MCP_RECALL_DESCRIPTION
|
||||
|
|
|
|||
|
|
@ -342,7 +342,8 @@ class TestMentalModelToolRegistration:
|
|||
assert "update_bank" in tools
|
||||
assert "delete_bank" in tools
|
||||
assert "clear_memories" in tools
|
||||
assert len(tools) == 29
|
||||
assert "sync_retain" in tools
|
||||
assert len(tools) == 30
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -1107,12 +1108,67 @@ class TestMemoryBrowsingTools:
|
|||
assert '"deleted"' in result
|
||||
assert mock_memory.delete_memory_unit.call_args.kwargs["unit_id"] == "mem-1"
|
||||
|
||||
async def test_get_memory_invalid_uuid(self, mock_memory):
|
||||
mock_memory.get_memory_unit.side_effect = ValueError("Invalid memory_id: 'nonexistent' is not a valid UUID")
|
||||
mcp = _make_mcp_server(mock_memory, {"get_memory"}, include_bank_id=True)
|
||||
result = await _tools(mcp)["get_memory"].fn(memory_id="nonexistent")
|
||||
assert "not a valid UUID" in result
|
||||
|
||||
async def test_get_memory_invalid_uuid_single_bank(self, mock_memory):
|
||||
mock_memory.get_memory_unit.side_effect = ValueError("Invalid memory_id: 'bad' is not a valid UUID")
|
||||
mcp = _make_mcp_server(mock_memory, {"get_memory"}, include_bank_id=False)
|
||||
result = await _tools(mcp)["get_memory"].fn(memory_id="bad")
|
||||
assert "not a valid UUID" in result["error"]
|
||||
|
||||
async def test_delete_memory_invalid_uuid(self, mock_memory):
|
||||
mock_memory.delete_memory_unit.side_effect = ValueError("Invalid unit_id: 'bad' is not a valid UUID")
|
||||
mcp = _make_mcp_server(mock_memory, {"delete_memory"}, include_bank_id=True)
|
||||
result = await _tools(mcp)["delete_memory"].fn(memory_id="bad")
|
||||
assert "not a valid UUID" in result
|
||||
|
||||
async def test_list_memories_single_bank(self, mock_memory):
|
||||
mcp = _make_mcp_server(mock_memory, {"list_memories"}, include_bank_id=False)
|
||||
result = await _tools(mcp)["list_memories"].fn()
|
||||
assert isinstance(result, dict)
|
||||
|
||||
|
||||
# =========================================================================
|
||||
# Sync Retain Tool Tests
|
||||
# =========================================================================
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestSyncRetainTool:
|
||||
async def test_sync_retain_basic(self, mock_memory):
|
||||
mock_memory.retain_batch_async.return_value = [["unit-1", "unit-2"]]
|
||||
mcp = _make_mcp_server(mock_memory, {"sync_retain"}, include_bank_id=True)
|
||||
result = await _tools(mcp)["sync_retain"].fn(content="test memory")
|
||||
assert result["status"] == "completed"
|
||||
assert result["memory_ids"] == ["unit-1", "unit-2"]
|
||||
|
||||
async def test_sync_retain_single_bank(self, mock_memory):
|
||||
mock_memory.retain_batch_async.return_value = [["unit-1"]]
|
||||
mcp = _make_mcp_server(mock_memory, {"sync_retain"}, include_bank_id=False)
|
||||
result = await _tools(mcp)["sync_retain"].fn(content="test memory")
|
||||
assert result["status"] == "completed"
|
||||
assert result["memory_ids"] == ["unit-1"]
|
||||
|
||||
async def test_sync_retain_with_tags(self, mock_memory):
|
||||
mock_memory.retain_batch_async.return_value = [["unit-1"]]
|
||||
mcp = _make_mcp_server(mock_memory, {"sync_retain"}, include_bank_id=True)
|
||||
result = await _tools(mcp)["sync_retain"].fn(content="test", tags=["project:alpha"])
|
||||
assert result["status"] == "completed"
|
||||
call_kwargs = mock_memory.retain_batch_async.call_args.kwargs
|
||||
assert call_kwargs["contents"][0]["tags"] == ["project:alpha"]
|
||||
|
||||
async def test_sync_retain_error(self, mock_memory):
|
||||
mock_memory.retain_batch_async.side_effect = Exception("DB error")
|
||||
mcp = _make_mcp_server(mock_memory, {"sync_retain"}, include_bank_id=True)
|
||||
result = await _tools(mcp)["sync_retain"].fn(content="test")
|
||||
assert result["status"] == "error"
|
||||
assert "DB error" in result["message"]
|
||||
|
||||
|
||||
# =========================================================================
|
||||
# Document Tool Tests
|
||||
# =========================================================================
|
||||
|
|
|
|||
|
|
@ -112,7 +112,7 @@ const DEFAULT_GEMINI_SAFETY_SETTINGS: GeminiSafetySetting[] = GEMINI_HARM_CATEGO
|
|||
// ─── MCP tool catalogue ───────────────────────────────────────────────────────
|
||||
|
||||
const MCP_TOOL_GROUPS: { label: string; tools: string[] }[] = [
|
||||
{ label: "Core", tools: ["retain", "recall", "reflect"] },
|
||||
{ label: "Core", tools: ["retain", "sync_retain", "recall", "reflect"] },
|
||||
{
|
||||
label: "Bank management",
|
||||
tools: [
|
||||
|
|
|
|||
Loading…
Reference in a new issue