Skip to content

Commit f58145d

Browse files
feat(auth): add create bank validation hook
Add a precise operation-validator hook for bank creation, with a no-op default so deployments without custom validators keep existing behavior. Route lazy bank creation through the hook from retain, imports, MCP create_bank, and the default get_bank_profile auto-create path. This keeps create-bank authorization separate from bank-scoped write validation, which often assumes the target bank already exists. Add regression coverage for rejected creation, existing-bank skips, HTTP create/import paths, async retain, profile auto-create, and MCP create_bank.
1 parent 00968a1 commit f58145d

7 files changed

Lines changed: 223 additions & 17 deletions

File tree

hindsight-api-slim/hindsight_api/api/http.py

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5642,8 +5642,12 @@ async def api_create_or_update_bank(
56425642
):
56435643
"""Create or update an agent with disposition and mission."""
56445644
try:
5645-
# Ensure bank exists by getting profile (auto-creates with defaults)
5646-
await app.state.memory.get_bank_profile(bank_id, request_context=request_context)
5645+
# Ensure bank exists, validating create_bank only when this call
5646+
# actually creates a missing bank.
5647+
await app.state.memory._ensure_bank_exists(
5648+
bank_id,
5649+
request_context,
5650+
)
56475651

56485652
# Update name if provided (stored in DB for display only, deprecated)
56495653
if request.name is not None:
@@ -5835,8 +5839,12 @@ async def api_import_bank_template(
58355839
dry_run=True,
58365840
)
58375841

5838-
# Ensure bank exists (auto-creates with defaults if needed)
5839-
await app.state.memory.get_bank_profile(bank_id, request_context=request_context)
5842+
# Ensure bank exists, validating create_bank only when this import
5843+
# actually creates a missing target bank.
5844+
await app.state.memory._ensure_bank_exists(
5845+
bank_id,
5846+
request_context,
5847+
)
58405848

58415849
return await apply_bank_template_manifest(
58425850
memory=app.state.memory,

hindsight-api-slim/hindsight_api/engine/memory_engine.py

Lines changed: 44 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -3271,6 +3271,8 @@ async def retain_batch_async(
32713271
if result and result.contents is not None:
32723272
contents = cast(list[RetainContentDict], result.contents)
32733273

3274+
await self._ensure_bank_exists(bank_id, request_context)
3275+
32743276
# Engine-owned copy: the orchestrator clears per-item "content" strings
32753277
# after building the document's combined text (memory pressure
32763278
# optimization, see retain/orchestrator.py). Without an internal copy
@@ -3757,6 +3759,14 @@ async def import_bank_async(
37573759
# target bank's config before the restore.
37583760
parsed = parse_bank_archive(archive_bytes)
37593761
bank_id = target_bank_id or parsed.manifest.source_bank_id
3762+
if self._operation_validator and await bank_utils.get_bank_profile_if_exists(backend, bank_id) is None:
3763+
from hindsight_api.extensions import CreateBankContext
3764+
3765+
ctx = CreateBankContext(
3766+
bank_id=bank_id,
3767+
request_context=request_context,
3768+
)
3769+
await self._validate_operation(self._operation_validator.validate_create_bank(ctx))
37603770
resolved_config = await self._config_resolver.resolve_full_config(bank_id, request_context)
37613771
return await import_bank(
37623772
backend=backend,
@@ -8537,16 +8547,12 @@ async def get_bank_profile(
85378547
existing = await bank_utils.get_bank_profile_if_exists(backend, bank_id)
85388548
if existing is None:
85398549
return None
8540-
profile, created = existing, False
8550+
profile = existing
85418551
else:
8542-
result = await bank_utils.get_or_create_bank_profile(backend, bank_id)
8543-
profile, created = result.profile, result.created
8544-
8545-
# Apply HINDSIGHT_API_DEFAULT_BANK_TEMPLATE to freshly-created banks. Done
8546-
# before reading the resolved config below so the template's overrides
8547-
# (e.g. reflect_mission, dispositions) are visible on this very call.
8548-
if created:
8549-
await self._apply_default_bank_template(bank_id, request_context)
8552+
await self._ensure_bank_exists(bank_id, request_context)
8553+
profile = await bank_utils.get_bank_profile_if_exists(backend, bank_id)
8554+
if profile is None:
8555+
raise RuntimeError(f"Bank '{bank_id}' was not found after ensuring it exists")
85508556

85518557
# reflect_mission and disposition in config take precedence over the legacy DB columns
85528558
config_dict = await self._config_resolver.get_bank_config(bank_id, request_context)
@@ -8602,6 +8608,20 @@ async def _ensure_bank_exists(
86028608
True if the bank was freshly created on this call.
86038609
"""
86048610
backend = await self._get_backend()
8611+
if self._operation_validator:
8612+
if conn is not None:
8613+
exists = await conn.fetchval(f"SELECT 1 FROM {fq_table('banks')} WHERE bank_id = $1", bank_id)
8614+
else:
8615+
exists = await bank_utils.get_bank_profile_if_exists(backend, bank_id)
8616+
if not exists:
8617+
from hindsight_api.extensions import CreateBankContext
8618+
8619+
ctx = CreateBankContext(
8620+
bank_id=bank_id,
8621+
request_context=request_context,
8622+
)
8623+
await self._validate_operation(self._operation_validator.validate_create_bank(ctx))
8624+
86058625
if conn is not None:
86068626
result = await bank_utils.get_or_create_bank_profile_on_conn(conn, bank_id, ops=backend.ops)
86078627
return result.created
@@ -10343,7 +10363,11 @@ async def create_mental_model(
1034310363
# outlives a mental-model insert that ultimately fails.
1034410364
async with acquire_with_retry(backend) as conn:
1034510365
async with conn.transaction():
10346-
created = await self._ensure_bank_exists(bank_id, request_context, conn=conn)
10366+
created = await self._ensure_bank_exists(
10367+
bank_id,
10368+
request_context,
10369+
conn=conn,
10370+
)
1034710371
if mental_model_id:
1034810372
row = await conn.fetchrow(
1034910373
f"""
@@ -11964,7 +11988,11 @@ async def create_webhook(
1196411988
# commit (or roll back) atomically.
1196511989
async with acquire_with_retry(backend) as conn:
1196611990
async with conn.transaction():
11967-
created = await self._ensure_bank_exists(bank_id, request_context, conn=conn)
11991+
created = await self._ensure_bank_exists(
11992+
bank_id,
11993+
request_context,
11994+
conn=conn,
11995+
)
1196811996
row = await backend.ops.create_webhook(
1196911997
conn,
1197011998
fq_table("webhooks"),
@@ -12413,7 +12441,11 @@ async def submit_async_retain(
1241312441
# async_operations.bank_id has a FK to banks. Create the bank
1241412442
# lazily inside this same transaction so it is atomic with the
1241512443
# parent + child operation rows.
12416-
created = await self._ensure_bank_exists(bank_id, request_context, conn=conn)
12444+
created = await self._ensure_bank_exists(
12445+
bank_id,
12446+
request_context,
12447+
conn=conn,
12448+
)
1241712449
await conn.execute(
1241812450
f"""
1241912451
INSERT INTO {fq_table("async_operations")} (operation_id, bank_id, operation_type, result_metadata, status)

hindsight-api-slim/hindsight_api/extensions/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
# Consolidation operation
4646
ConsolidateContext,
4747
ConsolidateResult,
48+
CreateBankContext,
4849
# File Conversion
4950
FileConvertResult,
5051
# Mental Model operations
@@ -105,6 +106,7 @@
105106
"BankReadOperation",
106107
"BankWriteContext",
107108
"BankWriteOperation",
109+
"CreateBankContext",
108110
# Operation Validator - Consolidation
109111
"ConsolidateContext",
110112
"ConsolidateResult",

hindsight-api-slim/hindsight_api/extensions/operation_validator.py

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -397,6 +397,14 @@ class BankWriteContext:
397397
request_context: "RequestContext"
398398

399399

400+
@dataclass
401+
class CreateBankContext:
402+
"""Context for validating creation of a new bank."""
403+
404+
bank_id: str
405+
request_context: "RequestContext"
406+
407+
400408
@dataclass
401409
class BankListContext:
402410
"""Context for filtering the bank list (post-query)."""
@@ -881,6 +889,23 @@ async def validate_bank_write(self, ctx: BankWriteContext) -> ValidationResult:
881889
"""
882890
return ValidationResult.accept()
883891

892+
async def validate_create_bank(self, ctx: CreateBankContext) -> ValidationResult:
893+
"""
894+
Validate creation of a new bank before the bank row is inserted.
895+
896+
Override to implement custom validation logic for operations that
897+
explicitly or implicitly create a bank.
898+
899+
Args:
900+
ctx: Context containing:
901+
- bank_id: Bank identifier
902+
- request_context: Request context with auth info
903+
904+
Returns:
905+
ValidationResult indicating whether the bank may be created.
906+
"""
907+
return ValidationResult.accept()
908+
884909
async def filter_bank_list(self, ctx: BankListContext) -> BankListResult:
885910
"""
886911
Filter the bank list after querying.

hindsight-api-slim/hindsight_api/mcp_tools.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1225,7 +1225,9 @@ async def create_bank(bank_id: str, name: str | None = None, mission: str | None
12251225
"""
12261226
try:
12271227
request_context = _get_request_context(config)
1228-
# get_bank_profile auto-creates bank if it doesn't exist
1228+
# create_bank may auto-create the bank; validate that explicit
1229+
# creation permission before reading the resulting profile.
1230+
await memory._ensure_bank_exists(bank_id, request_context)
12291231
profile = await memory.get_bank_profile(bank_id, request_context=request_context)
12301232

12311233
# Update name/mission if provided

hindsight-api-slim/tests/test_extensions.py

Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
# Consolidation operation
1717
ConsolidateContext,
1818
ConsolidateResult,
19+
CreateBankContext,
1920
Extension,
2021
HttpExtension,
2122
OperationValidationError,
@@ -202,6 +203,30 @@ async def on_consolidate_complete(self, result: ConsolidateResult) -> None:
202203
self.post_consolidate_calls.append(result)
203204

204205

206+
class CreateBankRejectingValidator(OperationValidatorExtension):
207+
"""Validator that records create-bank checks and can reject them."""
208+
209+
def __init__(self, *, reject: bool = True):
210+
super().__init__({})
211+
self.reject = reject
212+
self.create_bank_calls: list[CreateBankContext] = []
213+
214+
async def validate_retain(self, ctx: RetainContext) -> ValidationResult:
215+
return ValidationResult.accept()
216+
217+
async def validate_recall(self, ctx: RecallContext) -> ValidationResult:
218+
return ValidationResult.accept()
219+
220+
async def validate_reflect(self, ctx: ReflectContext) -> ValidationResult:
221+
return ValidationResult.accept()
222+
223+
async def validate_create_bank(self, ctx: CreateBankContext) -> ValidationResult:
224+
self.create_bank_calls.append(ctx)
225+
if self.reject:
226+
return ValidationResult.reject("bank creation not allowed", status_code=402)
227+
return ValidationResult.accept()
228+
229+
205230
class TestMemoryEngineValidation:
206231
"""Tests for validation integration with MemoryEngine.
207232
@@ -297,6 +322,109 @@ async def test_reflect_validation(self, memory_with_validator):
297322

298323
assert "limit exceeded" in str(exc_info.value).lower()
299324

325+
@pytest.mark.asyncio
326+
async def test_retain_validates_create_bank_for_missing_bank(self, memory):
327+
"""Retain may lazily create a bank, so missing-bank retain validates create_bank."""
328+
bank_id = "test-retain-create-bank-validation"
329+
ctx = RequestContext()
330+
validator = CreateBankRejectingValidator()
331+
memory._operation_validator = validator
332+
333+
with pytest.raises(OperationValidationError) as exc_info:
334+
await memory.retain_batch_async(
335+
bank_id=bank_id,
336+
contents=[{"content": "Should not create a bank"}],
337+
request_context=ctx,
338+
)
339+
340+
assert "bank creation not allowed" in str(exc_info.value)
341+
assert len(validator.create_bank_calls) == 1
342+
assert validator.create_bank_calls[0].bank_id == bank_id
343+
assert await memory.get_bank_profile(bank_id, request_context=ctx, create_if_missing=False) is None
344+
345+
@pytest.mark.asyncio
346+
async def test_async_retain_validates_create_bank_for_missing_bank(self, memory):
347+
"""Async retain validates create_bank before creating its parent operation."""
348+
bank_id = "test-async-retain-create-bank-validation"
349+
ctx = RequestContext()
350+
validator = CreateBankRejectingValidator()
351+
memory._operation_validator = validator
352+
353+
with pytest.raises(OperationValidationError) as exc_info:
354+
await memory.submit_async_retain(
355+
bank_id=bank_id,
356+
contents=[{"content": "Should not create a bank"}],
357+
request_context=ctx,
358+
)
359+
360+
assert "bank creation not allowed" in str(exc_info.value)
361+
assert len(validator.create_bank_calls) == 1
362+
assert validator.create_bank_calls[0].bank_id == bank_id
363+
assert await memory.get_bank_profile(bank_id, request_context=ctx, create_if_missing=False) is None
364+
365+
@pytest.mark.asyncio
366+
async def test_create_bank_validation_skips_existing_bank(self, memory):
367+
"""Existing banks do not require create_bank validation."""
368+
bank_id = "test-create-bank-existing"
369+
ctx = RequestContext()
370+
await memory.get_bank_profile(bank_id, request_context=ctx)
371+
372+
validator = CreateBankRejectingValidator()
373+
memory._operation_validator = validator
374+
375+
created = await memory._ensure_bank_exists(bank_id, ctx)
376+
377+
assert created is False
378+
assert validator.create_bank_calls == []
379+
380+
@pytest.mark.asyncio
381+
async def test_get_bank_profile_validates_create_bank_for_missing_bank(self, memory):
382+
"""The default auto-create profile path validates create_bank."""
383+
bank_id = "test-profile-create-bank-validation"
384+
ctx = RequestContext()
385+
validator = CreateBankRejectingValidator()
386+
memory._operation_validator = validator
387+
388+
with pytest.raises(OperationValidationError) as exc_info:
389+
await memory.get_bank_profile(bank_id, request_context=ctx)
390+
391+
assert "bank creation not allowed" in str(exc_info.value)
392+
assert len(validator.create_bank_calls) == 1
393+
assert validator.create_bank_calls[0].bank_id == bank_id
394+
assert await memory.get_bank_profile(bank_id, request_context=ctx, create_if_missing=False) is None
395+
396+
@pytest.mark.asyncio
397+
async def test_http_create_or_update_validates_create_bank(self, api_client, memory):
398+
bank_id = "test-http-put-create-bank-validation"
399+
validator = CreateBankRejectingValidator()
400+
memory._operation_validator = validator
401+
402+
resp = await api_client.put(
403+
f"/v1/default/banks/{bank_id}",
404+
json={"name": "Blocked Bank"},
405+
)
406+
407+
assert resp.status_code == 402
408+
assert resp.json()["detail"] == "bank creation not allowed"
409+
assert len(validator.create_bank_calls) == 1
410+
assert validator.create_bank_calls[0].bank_id == bank_id
411+
412+
@pytest.mark.asyncio
413+
async def test_http_template_import_validates_create_bank(self, api_client, memory):
414+
bank_id = "test-http-import-create-bank-validation"
415+
validator = CreateBankRejectingValidator()
416+
memory._operation_validator = validator
417+
418+
resp = await api_client.post(
419+
f"/v1/default/banks/{bank_id}/import",
420+
json={"version": "1"},
421+
)
422+
423+
assert resp.status_code == 402
424+
assert resp.json()["detail"] == "bank creation not allowed"
425+
assert len(validator.create_bank_calls) == 1
426+
assert validator.create_bank_calls[0].bank_id == bank_id
427+
300428

301429
@pytest.fixture
302430
def memory_with_validator(memory):

hindsight-api-slim/tests/test_mcp_tools.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -188,6 +188,7 @@ async def _get_mental_model(**kwargs):
188188

189189
# Tags & bank methods
190190
memory.list_tags = AsyncMock(return_value={"items": ["tag1", "tag2"], "total": 2})
191+
memory._ensure_bank_exists = AsyncMock(return_value=True)
191192
memory.get_bank_profile = AsyncMock(return_value={"id": "test-bank", "name": "Test Bank", "mission": "Testing"})
192193
memory.get_bank_stats = AsyncMock(return_value={"nodes": 100, "links": 50})
193194
memory.update_bank = AsyncMock(return_value={"id": "test-bank", "name": "Updated"})
@@ -1512,6 +1513,14 @@ async def test_get_bank(self, mock_memory):
15121513
result = await _tools(mcp)["get_bank"].fn()
15131514
assert '"test-bank"' in result or "test-bank" in result
15141515

1516+
async def test_create_bank_validates_creation(self, mock_memory):
1517+
mcp = _make_mcp_server(mock_memory, {"create_bank"}, include_bank_id=True)
1518+
result = await _tools(mcp)["create_bank"].fn(bank_id="new-bank")
1519+
assert '"test-bank"' in result or "test-bank" in result
1520+
call = mock_memory._ensure_bank_exists.call_args
1521+
assert call.args[0] == "new-bank"
1522+
assert "operation" not in call.kwargs
1523+
15151524
async def test_get_bank_stats(self, mock_memory):
15161525
mcp = _make_mcp_server(mock_memory, {"get_bank_stats"}, include_bank_id=True)
15171526
result = await _tools(mcp)["get_bank_stats"].fn()

0 commit comments

Comments
 (0)