Skip to content

Commit 3de6bc8

Browse files
teryltTeryl Taylormonshricrivetimihai
authored andcommitted
feat: add tool metadata and http headers to plugin tool hooks (IBM#854)
* rebase: rebased with main, fixing merge conflicts Signed-off-by: Teryl Taylor <[email protected]> * fix: plugin cleanup to support multiple external plugins. Signed-off-by: Teryl Taylor <[email protected]> * fix(lint): fixed linting issues Signed-off-by: Teryl Taylor <[email protected]> * feat(error): update error handling with enforce_ignore_error Signed-off-by: Teryl Taylor <[email protected]> * fix(plugins): updated documentation and addressed PR comments. Signed-off-by: Teryl Taylor <[email protected]> * fix(lint): fixed linting issue Signed-off-by: Teryl Taylor <[email protected]> * feat(plugins): added initial http header hooks. Signed-off-by: Teryl Taylor <[email protected]> * fix(comments): update docstrings to fix linting. Signed-off-by: Teryl Taylor <[email protected]> * fix: linting issue. Signed-off-by: Teryl Taylor <[email protected]> * feat: added hooks to the plugin manager for http pre/post header requests. Signed-off-by: Teryl Taylor <[email protected]> * feat: added tool metadata and headers to tool payloads. Signed-off-by: Teryl Taylor <[email protected]> * fix: fixed model to support passing tool metadata. Signed-off-by: Teryl Taylor <[email protected]> * feat: added example header plugin for tools. Signed-off-by: Teryl Taylor <[email protected]> * fix: refactored ToolMetaData, GatewayMetadata, removed http hooks, fixed test cases Signed-off-by: Teryl Taylor <[email protected]> * adding handlers for pluginerror and pluginviolationerror Signed-off-by: Shriti Priya <[email protected]> * fix for headers pydantic error in tool, plugin violation error handler Signed-off-by: Shriti Priya <[email protected]> * Error handling changes with test cases modification Signed-off-by: Shriti Priya <[email protected]> * fixing flake8 issues Signed-off-by: Shriti Priya <[email protected]> * refactored error handling in prompt and resource services, added unit tests for meta data, fixed existing tests. Signed-off-by: Teryl Taylor <[email protected]> * fix: made original_name optional Signed-off-by: Teryl Taylor <[email protected]> * tests(tools): added test to check both gateway and tool metadata Signed-off-by: Teryl Taylor <[email protected]> * tests(headers): added tool header tests Signed-off-by: Teryl Taylor <[email protected]> * tests(tool_post_invoke): tests cases for tool post invoke metadata. Signed-off-by: Teryl Taylor <[email protected]> * fix(tool): check whether tools payload headers are None Signed-off-by: Teryl Taylor <[email protected]> * docs(plugins): added some documentation on the headers and meta data. Signed-off-by: Teryl Taylor <[email protected]> * fix: updated error response values Signed-off-by: Teryl Taylor <[email protected]> * Rebase Signed-off-by: Mihai Criveti <[email protected]> --------- Signed-off-by: Teryl Taylor <[email protected]> Signed-off-by: Shriti Priya <[email protected]> Signed-off-by: Mihai Criveti <[email protected]> Co-authored-by: Teryl Taylor <[email protected]> Co-authored-by: Shriti Priya <[email protected]> Co-authored-by: Mihai Criveti <[email protected]>
1 parent 93d7033 commit 3de6bc8

File tree

20 files changed

+956
-291
lines changed

20 files changed

+956
-291
lines changed

docs/docs/using/plugins/index.md

Lines changed: 119 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -243,7 +243,7 @@ Each plugin can operate in one of four modes:
243243
| Mode | Description | Use Case |
244244
|------|-------------|----------|
245245
| **enforce** | Blocks requests on policy violations and plugin errors | Production guardrails |
246-
| **enforce_ignore_errors** | Blocks requests on policy violations; logs errors and continues | Guardrails with fault tolerance |
246+
| **enforce_ignore_errors** | Blocks requests on policy violations but only logs errors | Production guardrails |
247247
| **permissive** | Logs violations but allows requests | Testing and monitoring |
248248
| **disabled** | Plugin loaded but not executed | Temporary deactivation |
249249

@@ -308,6 +308,44 @@ The plugin framework provides comprehensive hook coverage across the entire MCP
308308
| `federation_pre_sync` | Gateway federation validation and filtering | v0.8.0 |
309309
| `federation_post_sync` | Post-federation data processing and reconciliation | v0.8.0 |
310310

311+
### Prompt Hooks Details
312+
313+
The prompt hooks allow plugins to intercept and modify prompt retrieval and rendering:
314+
315+
- **`prompt_pre_fetch`**: Receives the prompt name and arguments before prompt template retrieval. Can modify the arguments.
316+
- **`prompt_post_fetch`**: Receives the completed prompt after rendering. Can modify the prompt text or block it from being returned.
317+
318+
Example Use Cases:
319+
- Detect prompt injection attacks
320+
- Sanitize or anonymize prompts
321+
- Search and replace
322+
323+
#### Prompt Hook Payloads
324+
325+
**PromptPrehookPayload**: Payload for prompt pre-fetch hooks.
326+
327+
```python
328+
class PromptPrehookPayload(BaseModel):
329+
name: str # Prompt template name
330+
args: Optional[dict[str, str]] = Field(default_factory=dict) # Template arguments
331+
```
332+
333+
**Example**:
334+
```python
335+
payload = PromptPrehookPayload(
336+
name="user_greeting",
337+
args={"user_name": "Alice", "time_of_day": "morning"}
338+
)
339+
```
340+
341+
**PromptPosthookPayload**: Payload for prompt post-fetch hooks.
342+
343+
```python
344+
class PromptPosthookPayload(BaseModel):
345+
name: str # Prompt name
346+
result: PromptResult # Rendered prompt result
347+
```
348+
311349
### Tool Hooks Details
312350

313351
The tool hooks enable plugins to intercept and modify tool invocations:
@@ -322,6 +360,42 @@ Example use cases:
322360
- Input validation and sanitization
323361
- Output filtering and transformation
324362

363+
#### Tool Hook Payloads
364+
365+
**ToolPreInvokePayload**: Payload for tool pre-invoke hooks.
366+
367+
```python
368+
class ToolPreInvokePayload(BaseModel):
369+
name: str # Tool name
370+
args: Optional[dict[str, Any]] = Field(default_factory=dict) # Tool arguments
371+
headers: Optional[HttpHeaderPayload] = None # HTTP pass-through headers
372+
```
373+
374+
**ToolPostInvokePayload**: Payload for tool post-invoke hooks.
375+
376+
```python
377+
class ToolPostInvokePayload(BaseModel):
378+
name: str # Tool name
379+
result: Any # Tool execution result
380+
```
381+
382+
The associated `HttpHeaderPayload` object for the `ToolPreInvokePayload` is as follows:
383+
384+
Special payload for HTTP header manipulation.
385+
386+
```python
387+
class HttpHeaderPayload(RootModel[dict[str, str]]):
388+
# Provides dictionary-like access to HTTP headers
389+
# Supports: __iter__, __getitem__, __setitem__, __len__
390+
```
391+
392+
**Usage**:
393+
```python
394+
headers = HttpHeaderPayload({"Authorization": "Bearer token", "Content-Type": "application/json"})
395+
headers["X-Custom-Header"] = "custom_value"
396+
auth_header = headers["Authorization"]
397+
```
398+
325399
### Resource Hooks Details
326400

327401
The resource hooks enable plugins to intercept and modify resource fetching:
@@ -337,6 +411,24 @@ Example use cases:
337411
- Content transformation and filtering
338412
- Resource caching metadata
339413

414+
#### Resource Hook Payloads
415+
416+
**ResourcePreFetchPayload**: Payload for resource pre-fetch hooks.
417+
418+
```python
419+
class ResourcePreFetchPayload(BaseModel):
420+
uri: str # Resource URI
421+
metadata: Optional[dict[str, Any]] = Field(default_factory=dict) # Request metadata
422+
```
423+
424+
**ResourcePostFetchPayload**: Payload for resource post-fetch hooks.
425+
426+
```python
427+
class ResourcePostFetchPayload(BaseModel):
428+
uri: str # Resource URI
429+
content: Any # Fetched resource content
430+
```
431+
340432
Planned hooks (not yet implemented):
341433

342434
- `server_pre_register` / `server_post_register` - Server validation
@@ -611,6 +703,32 @@ async def prompt_post_fetch(self, payload, context):
611703
return PromptPosthookResult()
612704
```
613705

706+
#### Tool and Gateway Metadata
707+
708+
Currently, the tool pre/post hooks have access to tool and gateway metadata through the global context metadata dictionary. They are accessible as follows:
709+
710+
It can be accessed inside of the tool hooks through:
711+
712+
```python
713+
from mcpgateway.plugins.framework.constants import GATEWAY_METADATA, TOOL_METADATA
714+
715+
tool_meta = context.global_context.metadata[TOOL_METADATA]
716+
assert tool_meta.original_name == "test_tool"
717+
assert tool_meta.url.host == "example.com"
718+
assert tool_meta.integration_type == "REST" or tool_meta.integration_type == "MCP"
719+
```
720+
721+
Note, if the integration type is `MCP` the gateway information may also be available as follows.
722+
723+
```python
724+
gateway_meta = context.global_context.metadata[GATEWAY_METADATA]
725+
assert gateway_meta.name == "test_gateway"
726+
assert gateway_meta.transport == "sse"
727+
assert gateway_meta.url.host == "example.com"
728+
```
729+
730+
Metadata for other entities such as prompts and resources will be added in future versions of the gateway.
731+
614732
### External Service Plugin Example
615733

616734
```python

mcpgateway/main.py

Lines changed: 85 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,7 @@
6969
from mcpgateway.middleware.token_scoping import token_scoping_middleware
7070
from mcpgateway.models import InitializeResult, ListResourceTemplatesResult, LogLevel, Root
7171
from mcpgateway.observability import init_telemetry
72-
from mcpgateway.plugins.framework import PluginManager, PluginViolationError
72+
from mcpgateway.plugins.framework import PluginError, PluginManager, PluginViolationError
7373
from mcpgateway.routers.well_known import router as well_known_router
7474
from mcpgateway.schemas import (
7575
A2AAgentCreate,
@@ -494,6 +494,81 @@ async def database_exception_handler(_request: Request, exc: IntegrityError):
494494
return JSONResponse(status_code=409, content=ErrorFormatter.format_database_error(exc))
495495

496496

497+
@app.exception_handler(PluginViolationError)
498+
async def plugin_violation_exception_handler(_request: Request, exc: PluginViolationError):
499+
"""Handle plugins violations globally.
500+
501+
Intercepts PluginViolationError exceptions (e.g., OPA policy violation) and returns a properly formatted JSON error response.
502+
This provides consistent error handling for plugin violation across the entire application.
503+
504+
Args:
505+
_request: The FastAPI request object that triggered the database error.
506+
(Unused but required by FastAPI's exception handler interface)
507+
exc: The PluginViolationError exception containing constraint
508+
violation details.
509+
510+
Returns:
511+
JSONResponse: A 403 response with access forbidden.
512+
513+
Examples:
514+
>>> from mcpgateway.plugins.framework import PluginViolationError
515+
>>> from mcpgateway.plugins.framework.models import PluginViolation
516+
>>> from fastapi import Request
517+
>>> import asyncio
518+
>>>
519+
>>> # Create a mock integrity error
520+
>>> mock_error = PluginViolationError(message="plugin violation",violation = PluginViolation(
521+
... reason="Invalid input",
522+
... description="The input contains prohibited content",
523+
... code="PROHIBITED_CONTENT",
524+
... details={"field": "message", "value": "test"}
525+
... ))
526+
>>> result = asyncio.run(plugin_violation_exception_handler(None, mock_error))
527+
>>> result.status_code
528+
403
529+
"""
530+
policy_violation = exc.violation.model_dump() if exc.violation else {}
531+
policy_violation["message"] = exc.message
532+
return JSONResponse(status_code=403, content=policy_violation)
533+
534+
535+
@app.exception_handler(PluginError)
536+
async def plugin_exception_handler(_request: Request, exc: PluginError):
537+
"""Handle plugins errors globally.
538+
539+
Intercepts PluginError exceptions and returns a properly formatted JSON error response.
540+
This provides consistent error handling for plugin error across the entire application.
541+
542+
Args:
543+
_request: The FastAPI request object that triggered the database error.
544+
(Unused but required by FastAPI's exception handler interface)
545+
exc: The PluginError exception containing constraint
546+
violation details.
547+
548+
Returns:
549+
JSONResponse: A 500 response with internal server error.
550+
551+
Examples:
552+
>>> from mcpgateway.plugins.framework import PluginViolationError
553+
>>> from mcpgateway.plugins.framework.models import PluginErrorModel
554+
>>> from fastapi import Request
555+
>>> import asyncio
556+
>>>
557+
>>> # Create a mock integrity error
558+
>>> mock_error = PluginError(error = PluginErrorModel(
559+
... message="plugin error",
560+
... code="timeout",
561+
... plugin_name="abc",
562+
... details={"field": "message", "value": "test"}
563+
... ))
564+
>>> result = asyncio.run(plugin_exception_handler(None, mock_error))
565+
>>> result.status_code
566+
500
567+
"""
568+
error_obj = exc.error.model_dump() if exc.error else {}
569+
return JSONResponse(status_code=500, content=error_obj)
570+
571+
497572
class DocsAuthMiddleware(BaseHTTPMiddleware):
498573
"""
499574
Middleware to protect FastAPI's auto-generated documentation routes
@@ -3214,6 +3289,10 @@ async def handle_rpc(request: Request, db: Session = Depends(get_db), user=Depen
32143289
32153290
Returns:
32163291
Response with the RPC result or error.
3292+
3293+
Raises:
3294+
PluginError: If encounters issue with plugin
3295+
PluginViolationError: If plugin violated the request. Example - In case of OPA plugin, if the request is denied by policy.
32173296
"""
32183297
try:
32193298
# Extract user identifier from either RBAC user object or JWT payload
@@ -3324,13 +3403,14 @@ async def handle_rpc(request: Request, db: Session = Depends(get_db), user=Depen
33243403
else:
33253404
# Backward compatibility: Try to invoke as a tool directly
33263405
# This allows both old format (method=tool_name) and new format (method=tools/call)
3406+
# Standard
33273407
headers = {k.lower(): v for k, v in request.headers.items()}
33283408
try:
33293409
result = await tool_service.invoke_tool(db=db, name=method, arguments=params, request_headers=headers)
33303410
if hasattr(result, "model_dump"):
33313411
result = result.model_dump(by_alias=True, exclude_none=True)
3332-
except PluginViolationError:
3333-
return JSONResponse(status_code=403, content={"detail": "policy_deny"})
3412+
except (PluginError, PluginViolationError):
3413+
raise
33343414
except (ValueError, Exception):
33353415
# If not a tool, try forwarding to gateway
33363416
try:
@@ -3343,6 +3423,8 @@ async def handle_rpc(request: Request, db: Session = Depends(get_db), user=Depen
33433423

33443424
return {"jsonrpc": "2.0", "result": result, "id": req_id}
33453425

3426+
except (PluginError, PluginViolationError):
3427+
raise
33463428
except JSONRPCError as e:
33473429
error = e.to_dict()
33483430
return {"jsonrpc": "2.0", "error": error["error"], "id": req_id}

0 commit comments

Comments
 (0)