src/va_agent/agent.py

@@ -10,17 +10,20 @@ from google.cloud.firestore_v1.async_client import AsyncClient
 from va_agent.auth import auth_headers_provider
 from va_agent.config import settings
 from va_agent.session import FirestoreSessionService
+from va_agent.governance import GovernancePlugin
 
 toolset = McpToolset(
     connection_params=StreamableHTTPConnectionParams(url=settings.mcp_remote_url),
     header_provider=auth_headers_provider,
 )
 
+governance = GovernancePlugin()
 agent = Agent(
     model=settings.agent_model,
     name=settings.agent_name,
     instruction=settings.agent_instructions,
     tools=[toolset],
+    after_model_callback=governance.after_model_callback,
 )
 
 session_service = FirestoreSessionService(