@@ -18,6 +18,7 @@ import (
1818 extProcV3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
1919 "github.com/fsnotify/fsnotify"
2020 "github.com/kagenti/mcp-gateway/internal/broker"
21+ "github.com/kagenti/mcp-gateway/internal/clients"
2122 config "github.com/kagenti/mcp-gateway/internal/config"
2223 mcpRouter "github.com/kagenti/mcp-gateway/internal/mcp-router"
2324 "github.com/kagenti/mcp-gateway/internal/session"
@@ -49,20 +50,25 @@ func init() {
4950 _ = gatewayv1 .Install (scheme )
5051}
5152
53+ var (
54+ mcpRouterAddrFlag string
55+ mcpBrokerAddrFlag string
56+ mcpConfigAddrFlag string
57+ mcpRoutePublicHost string
58+ mcpRoutePrivateHost string
59+ mcpRouterKey string
60+ cacheConnectionStringFlag string
61+ mcpConfigFile string
62+ jwtSigningKeyFlag string
63+ sessionDurationInMins int64
64+ loglevel int
65+ logFormat string
66+ controllerMode bool
67+ enforceToolFilteringFlag bool
68+ )
69+
5270func main () {
53- var (
54- mcpRouterAddrFlag string
55- mcpBrokerAddrFlag string
56- mcpConfigAddrFlag string
57- mcpRoutePublicHost string
58- mcpConfigFile string
59- jwtSigningKeyFlag string
60- sessionDurationInHours int64
61- loglevel int
62- logFormat string
63- controllerMode bool
64- enforceToolFilteringFlag bool
65- )
71+
6672 flag .StringVar (
6773 & mcpRouterAddrFlag ,
6874 "mcp-router-address" ,
@@ -81,6 +87,20 @@ func main() {
8187 "" ,
8288 "The public host the MCP Gateway is exposing MCP servers on. The gateway router will always set the :authority header to this value to ensure the broker component cannot be bypassed." ,
8389 )
90+ flag .StringVar (
91+ & mcpRoutePrivateHost ,
92+ "mcp-gateway-private-host" ,
93+ "mcp-gateway-istio.gateway-system.svc.cluster.local:8080" ,
94+ "The private host the MCP Gateway. The gateway router will use this to hairpin request to initialize MCP servers etc." ,
95+ )
96+
97+ // TODO ick not sure how to describe this
98+ flag .StringVar (
99+ & mcpRouterKey ,
100+ "mcp-router-key" ,
101+ goenv .GetDefault ("MCP_ROUTER_API_KEY" , "secret-api-key" ),
102+ "this key is used to allow the router to send request through the gateway and be trusted by the router" ,
103+ )
84104 flag .StringVar (
85105 & mcpConfigAddrFlag ,
86106 "mcp-broker-config-address" ,
@@ -99,9 +119,20 @@ func main() {
99119 int (slog .LevelInfo ),
100120 "set the log level 0=info, 4=warn , 8=error and -4=debug" ,
101121 )
122+ flag .StringVar (& jwtSigningKeyFlag ,
123+ "session-signing-key" ,
124+ goenv .GetDefault ("JWT_SESSION_SIGNING_KEY" , defaultJWTSigningKey ),
125+ "JWT signing key for session tokens (env: JWT_SESSION_SIGNING_KEY)" ,
126+ )
127+ //"redis://redis.mcp-system.svc.cluster.local:6379
128+ flag .StringVar (& cacheConnectionStringFlag ,
129+ "cache-connection-string" ,
130+ goenv .GetDefault ("CACHE_CONNECTION_STRING" , "" ),
131+ "redis based cache connection string redis://<user>:<pass>@localhost:6379/<db> (env: CACHE_CONNECTION_STRING). If not set defaults to in memory storage" ,
132+ )
102133 flag .StringVar (& logFormat , "log-format" , "txt" , "switch to json logs with --log-format=json" )
103- flag . StringVar ( & jwtSigningKeyFlag , "session-signing-key" , goenv . GetDefault ( "JWT_SESSION_SIGNING_KEY" , defaultJWTSigningKey ), "JWT signing key for session tokens (env: JWT_SESSION_SIGNING_KEY)" )
104- flag .Int64Var (& sessionDurationInHours , "session-length" , 24 , "default session length with the gateway in hours " )
134+
135+ flag .Int64Var (& sessionDurationInMins , "session-length" , 60 * 24 , "default session length with the gateway in minutes. Default 24h " )
105136 flag .BoolVar (& controllerMode , "controller" , false , "Run in controller mode" )
106137 flag .BoolVar (& enforceToolFilteringFlag , "enforce-tool-filtering" , false , "when enabled an x-authorized-tools header will be needed to return any tools" )
107138 flag .Parse ()
@@ -140,6 +171,19 @@ func main() {
140171 }
141172
142173 ctx := context .Background ()
174+
175+ sessionCache , err := session .NewCache ()
176+ if err != nil {
177+ panic ("failed to setup session cache" + err .Error ())
178+ }
179+ if cacheConnectionStringFlag != "" {
180+ logger .Info ("session cache using external store" )
181+ sessionCache , err = session .NewCache (session .WithConnectionString (cacheConnectionStringFlag ))
182+ if err != nil {
183+ panic ("failed to setup session cache" + err .Error ())
184+ }
185+ }
186+
143187 var jwtSessionMgr * session.JWTManager
144188 if jwtSigningKeyFlag == "" {
145189 panic ("jwt session signing key is empty. Cannot proceed" )
@@ -148,22 +192,24 @@ func main() {
148192 logger .Warn ("jwt session signing key is set to the default value. This is not recommended for production" )
149193 }
150194
151- jwtmgr , err := session .NewJWTManager (jwtSigningKeyFlag , sessionDurationInHours , logger )
195+ jwtmgr , err := session .NewJWTManager (jwtSigningKeyFlag , sessionDurationInMins , logger , sessionCache )
152196 if err != nil {
153197 panic ("failed to setup jwt manager " + err .Error ())
154198 }
155199 jwtSessionMgr = jwtmgr
156200
157201 configServer := setUpConfigServer (mcpConfigAddrFlag )
158202 brokerServer , mcpBroker , mcpServer := setUpBroker (mcpBrokerAddrFlag , enforceToolFilteringFlag , jwtSessionMgr )
159- routerGRPCServer , router := setUpRouter (mcpBroker , logger , jwtSessionMgr )
203+ routerGRPCServer , router := setUpRouter (mcpBroker , logger , jwtSessionMgr , sessionCache )
160204 mcpConfig .RegisterObserver (router )
161205 mcpConfig .RegisterObserver (mcpBroker )
162206 if mcpRoutePublicHost == "" {
163207 panic ("--mcp-gateway-public-host cannot be empty. The mcp gateway needs to be informed of what public host to expect requests from so it can ensure routing and session mgmt happens. Set --mcp-gateway-public-host" )
164208 }
165209
166- mcpConfig .MCPGatewayHostname = mcpRoutePublicHost
210+ mcpConfig .MCPGatewayExternalHostname = mcpRoutePublicHost
211+ mcpConfig .MCPGatewayInternalHostname = mcpRoutePrivateHost
212+ mcpConfig .RouterAPIKey = mcpRouterKey
167213
168214 // Only load config and run broker/router in standalone mode
169215 LoadConfig (mcpConfigFile )
@@ -296,20 +342,19 @@ func setUpConfigServer(address string) *http.Server {
296342 }
297343}
298344
299- func setUpRouter (broker broker.MCPBroker , logger * slog.Logger , jwtManager * session.JWTManager ) (* grpc.Server , * mcpRouter.ExtProcServer ) {
300- grpcSrv := grpc .NewServer ()
345+ func setUpRouter (broker broker.MCPBroker , logger * slog.Logger , jwtManager * session.JWTManager , sessionCache * session.Cache ) (* grpc.Server , * mcpRouter.ExtProcServer ) {
301346
347+ grpcSrv := grpc .NewServer ()
302348 // Create the ExtProcServer instance
303349 server := & mcpRouter.ExtProcServer {
304350 RoutingConfig : mcpConfig ,
305- // TODO this seems wrong. Why does the router need to be passed an instance of the broker?
306- Broker : broker ,
307- Logger : logger ,
308- JWTManager : jwtManager ,
309- }
351+ Logger : logger ,
352+ JWTManager : jwtManager ,
353+ InitForClient : clients . Initialize ,
354+ SessionCache : sessionCache ,
355+ Broker : broker , // TODO we shouldn't need a handle to broker in the router
310356
311- // Setup the session cache with proper initialization
312- server .SetupSessionCache ()
357+ }
313358
314359 extProcV3 .RegisterExternalProcessorServer (grpcSrv , server )
315360 return grpcSrv , server
0 commit comments