diff --git a/gateway/mcp/mcp.py b/gateway/mcp/mcp.py index 865d11c..b54c6e9 100644 --- a/gateway/mcp/mcp.py +++ b/gateway/mcp/mcp.py @@ -9,7 +9,7 @@ import select import socket import subprocess import sys - +from datetime import datetime from invariant_sdk.async_client import AsyncClient from invariant_sdk.types.append_messages import AppendMessagesRequest from invariant_sdk.types.push_traces import PushTracesRequest @@ -74,6 +74,11 @@ def deduplicate_annotations(ctx: McpContext, new_annotations: list) -> list: return deduped_annotations +def generate_message_timestamp() -> str: + """Generate a timestamp for a message.""" + return datetime.now().isoformat() + + def check_if_new_errors(ctx: McpContext, guardrails_result: dict) -> bool: """Checks if there are new errors in the guardrails result.""" annotations = create_annotations_from_guardrails_errors( @@ -168,6 +173,10 @@ async def get_guardrails_check_result( or (ctx.guardrails.logging_guardrails and action == GuardrailAction.LOG) ): return {} + + # Add timestamp to message if it doesn't already have one + if "timestamp" not in message: + message["timestamp"] = generate_message_timestamp() # Prepare context and select appropriate guardrails context = RequestContext.create( @@ -471,6 +480,7 @@ async def process_line( }, } ], + "timestamp": generate_message_timestamp(), } ) mcp_process.stdin.write(write_as_utf8_bytes(parsed_json))