aeonshift commited on
Commit
45eef8e
·
verified ·
1 Parent(s): bda84dd

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +36 -17
app.py CHANGED
@@ -25,7 +25,7 @@ app.add_middleware(
25
  CORSMiddleware,
26
  allow_origins=["*"], # Adjust for production
27
  allow_credentials=True,
28
- allow_methods=["*"],
29
  allow_headers=["*"],
30
  )
31
 
@@ -111,12 +111,15 @@ mcp_server = Server(name="airtable-mcp")
111
  mcp_server.tool_handlers = tool_handlers # Set as attribute
112
  mcp_server.tools = tools # Set tools as attribute for Deep Agent to discover
113
 
114
- # Initialize SseServerTransport
115
- transport = SseServerTransport("/airtable/mcp")
116
-
117
  # Store write streams for each session ID
118
  write_streams: Dict[str, anyio.streams.memory.MemoryObjectSendStream] = {}
119
 
 
 
 
 
 
 
120
  # SSE endpoint for GET requests
121
  @app.get("/airtable/mcp")
122
  async def handle_sse(request: Request):
@@ -124,36 +127,42 @@ async def handle_sse(request: Request):
124
  async def sse_writer():
125
  logger.debug("Starting SSE writer")
126
  async with sse_stream_writer, write_stream_reader:
127
- # Extract session_id from endpoint_data
128
  session_id = None
129
  endpoint_data = "/airtable/mcp?session_id={session_id}"
130
  await sse_stream_writer.send(
131
  {"event": "endpoint", "data": endpoint_data}
132
  )
133
  logger.debug(f"Sent endpoint event: {endpoint_data}")
134
- # Extract session_id from the sent message (simplified)
135
- # In a real implementation, we'd parse the session_id from transport
136
- # For now, we'll set it after the first message
137
  async for session_message in write_stream_reader:
138
  logger.debug(f"Sending message via SSE: {session_message}")
 
139
  await sse_stream_writer.send(
140
  {
141
  "event": "message",
142
- "data": session_message.message.model_dump_json(
143
- by_alias=True, exclude_none=True
144
- ),
145
  }
146
  )
 
 
 
 
 
 
 
 
 
 
 
 
147
 
148
  sse_stream_writer, sse_stream_reader = anyio.create_memory_object_stream(0)
149
  try:
150
  async with transport.connect_sse(request.scope, request.receive, request._send) as streams:
151
  read_stream, write_stream = streams
152
  write_stream_reader = write_stream # Since streams are MemoryObject streams
153
- # Extract session_id (this is a simplification; ideally, we'd get it from transport)
154
- session_id = request.query_params.get("session_id", "unknown")
155
- write_streams[session_id] = write_stream # Store write_stream for this session
156
- logger.debug(f"Stored write_stream for session_id: {session_id}")
157
  logger.debug("Running MCP server with streams")
158
  await mcp_server.run(read_stream, write_stream, mcp_server.create_initialization_options())
159
  except Exception as e:
@@ -181,7 +190,16 @@ async def handle_post_message(request: Request):
181
  "capabilities": {
182
  "tools": {
183
  "listChanged": True
184
- }
 
 
 
 
 
 
 
 
 
185
  },
186
  "serverInfo": {
187
  "name": "airtable-mcp",
@@ -203,7 +221,8 @@ async def handle_post_message(request: Request):
203
  "jsonrpc": "2.0",
204
  "id": message.get("id"),
205
  "result": {
206
- "tools": [tool.model_dump(by_alias=True, exclude_none=True) for tool in tools]
 
207
  }
208
  }
209
  logger.debug(f"Manual tools/list response: {response}")
 
25
  CORSMiddleware,
26
  allow_origins=["*"], # Adjust for production
27
  allow_credentials=True,
28
+ allow_methods=["GET", "POST", "OPTIONS"],
29
  allow_headers=["*"],
30
  )
31
 
 
111
  mcp_server.tool_handlers = tool_handlers # Set as attribute
112
  mcp_server.tools = tools # Set tools as attribute for Deep Agent to discover
113
 
 
 
 
114
  # Store write streams for each session ID
115
  write_streams: Dict[str, anyio.streams.memory.MemoryObjectSendStream] = {}
116
 
117
+ # Store session IDs for each connection
118
+ session_ids: Dict[str, str] = {}
119
+
120
+ # Initialize SseServerTransport
121
+ transport = SseServerTransport("/airtable/mcp")
122
+
123
  # SSE endpoint for GET requests
124
  @app.get("/airtable/mcp")
125
  async def handle_sse(request: Request):
 
127
  async def sse_writer():
128
  logger.debug("Starting SSE writer")
129
  async with sse_stream_writer, write_stream_reader:
130
+ # Extract session_id dynamically
131
  session_id = None
132
  endpoint_data = "/airtable/mcp?session_id={session_id}"
133
  await sse_stream_writer.send(
134
  {"event": "endpoint", "data": endpoint_data}
135
  )
136
  logger.debug(f"Sent endpoint event: {endpoint_data}")
 
 
 
137
  async for session_message in write_stream_reader:
138
  logger.debug(f"Sending message via SSE: {session_message}")
139
+ message_data = session_message.message.model_dump_json(by_alias=True, exclude_none=True)
140
  await sse_stream_writer.send(
141
  {
142
  "event": "message",
143
+ "data": message_data
 
 
144
  }
145
  )
146
+ # Extract session_id from the endpoint event
147
+ if not session_id and session_message.message.method == "endpoint":
148
+ try:
149
+ message = json.loads(message_data)
150
+ endpoint_url = message.get("data", "")
151
+ if "session_id=" in endpoint_url:
152
+ session_id = endpoint_url.split("session_id=")[1]
153
+ session_ids[id(write_stream)] = session_id
154
+ write_streams[session_id] = write_stream
155
+ logger.debug(f"Extracted session_id: {session_id}")
156
+ except Exception as e:
157
+ logger.error(f"Error extracting session_id: {str(e)}")
158
 
159
  sse_stream_writer, sse_stream_reader = anyio.create_memory_object_stream(0)
160
  try:
161
  async with transport.connect_sse(request.scope, request.receive, request._send) as streams:
162
  read_stream, write_stream = streams
163
  write_stream_reader = write_stream # Since streams are MemoryObject streams
164
+ # Use a unique identifier for the stream until session_id is extracted
165
+ stream_id = id(write_stream)
 
 
166
  logger.debug("Running MCP server with streams")
167
  await mcp_server.run(read_stream, write_stream, mcp_server.create_initialization_options())
168
  except Exception as e:
 
190
  "capabilities": {
191
  "tools": {
192
  "listChanged": True
193
+ },
194
+ "prompts": {
195
+ "listChanged": False
196
+ },
197
+ "resources": {
198
+ "subscribe": False,
199
+ "listChanged": False
200
+ },
201
+ "logging": {},
202
+ "experimental": {}
203
  },
204
  "serverInfo": {
205
  "name": "airtable-mcp",
 
221
  "jsonrpc": "2.0",
222
  "id": message.get("id"),
223
  "result": {
224
+ "tools": [tool.model_dump(by_alias=True, exclude_none=True) for tool in tools],
225
+ "nextCursor": None
226
  }
227
  }
228
  logger.debug(f"Manual tools/list response: {response}")