Skip to content

Commit 47db0a1

Browse files
jssmithahmedmustahidclaude
committed
Add streamable HTTP transport support
Add support for the streamable-http transport option alongside the existing stdio and sse transports. This addresses the deprecation of SSE transport in favor of streamable HTTP. Changes: - Add 'streamable-http' as a transport choice - Add --streamable-http-host and --streamable-http-port arguments - Call mcp.run_streamable_http_async() for streamable-http transport - Add comprehensive tests for all transport options Based on work by @ahmedmustahid in #78. Closes #102 Co-Authored-By: Ahmed Mustahid <ahmedmustahid@users.noreply.github.com> Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
1 parent d03cf9f commit 47db0a1

2 files changed

Lines changed: 148 additions & 4 deletions

File tree

src/postgres_mcp/server.py

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -568,9 +568,9 @@ async def main():
568568
parser.add_argument(
569569
"--transport",
570570
type=str,
571-
choices=["stdio", "sse"],
571+
choices=["stdio", "sse", "streamable-http"],
572572
default="stdio",
573-
help="Select MCP transport: stdio (default) or sse",
573+
help="Select MCP transport: stdio (default), sse, or streamable-http",
574574
)
575575
parser.add_argument(
576576
"--sse-host",
@@ -584,6 +584,18 @@ async def main():
584584
default=8000,
585585
help="Port for SSE server (default: 8000)",
586586
)
587+
parser.add_argument(
588+
"--streamable-http-host",
589+
type=str,
590+
default="localhost",
591+
help="Host to bind streamable HTTP server to (default: localhost)",
592+
)
593+
parser.add_argument(
594+
"--streamable-http-port",
595+
type=int,
596+
default=8000,
597+
help="Port for streamable HTTP server (default: 8000)",
598+
)
587599

588600
args = parser.parse_args()
589601

@@ -647,11 +659,14 @@ async def main():
647659
# Run the server with the selected transport (always async)
648660
if args.transport == "stdio":
649661
await mcp.run_stdio_async()
650-
else:
651-
# Update FastMCP settings based on command line arguments
662+
elif args.transport == "sse":
652663
mcp.settings.host = args.sse_host
653664
mcp.settings.port = args.sse_port
654665
await mcp.run_sse_async()
666+
elif args.transport == "streamable-http":
667+
mcp.settings.host = args.streamable_http_host
668+
mcp.settings.port = args.streamable_http_port
669+
await mcp.run_streamable_http_async()
655670

656671

657672
async def shutdown(sig=None):

tests/unit/test_transport.py

Lines changed: 129 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,129 @@
1+
import sys
2+
from unittest.mock import AsyncMock
3+
from unittest.mock import patch
4+
5+
import pytest
6+
7+
8+
@pytest.mark.asyncio
9+
@pytest.mark.parametrize("transport", ["stdio", "sse", "streamable-http"])
10+
async def test_transport_argument_parsing(transport):
11+
"""Test that all transport options are parsed correctly."""
12+
from postgres_mcp.server import main
13+
14+
original_argv = sys.argv
15+
try:
16+
sys.argv = [
17+
"postgres_mcp",
18+
"postgresql://user:password@localhost/db",
19+
f"--transport={transport}",
20+
]
21+
22+
with (
23+
patch("postgres_mcp.server.db_connection.pool_connect", AsyncMock()),
24+
patch("postgres_mcp.server.mcp.run_stdio_async", AsyncMock()) as mock_stdio,
25+
patch("postgres_mcp.server.mcp.run_sse_async", AsyncMock()) as mock_sse,
26+
patch("postgres_mcp.server.mcp.run_streamable_http_async", AsyncMock()) as mock_http,
27+
):
28+
await main()
29+
30+
# Verify the correct transport method was called
31+
if transport == "stdio":
32+
mock_stdio.assert_called_once()
33+
mock_sse.assert_not_called()
34+
mock_http.assert_not_called()
35+
elif transport == "sse":
36+
mock_stdio.assert_not_called()
37+
mock_sse.assert_called_once()
38+
mock_http.assert_not_called()
39+
elif transport == "streamable-http":
40+
mock_stdio.assert_not_called()
41+
mock_sse.assert_not_called()
42+
mock_http.assert_called_once()
43+
finally:
44+
sys.argv = original_argv
45+
46+
47+
@pytest.mark.asyncio
48+
async def test_streamable_http_host_port_arguments():
49+
"""Test that streamable-http host and port arguments are applied correctly."""
50+
from postgres_mcp.server import main
51+
from postgres_mcp.server import mcp
52+
53+
original_argv = sys.argv
54+
try:
55+
sys.argv = [
56+
"postgres_mcp",
57+
"postgresql://user:password@localhost/db",
58+
"--transport=streamable-http",
59+
"--streamable-http-host=0.0.0.0",
60+
"--streamable-http-port=9000",
61+
]
62+
63+
with (
64+
patch("postgres_mcp.server.db_connection.pool_connect", AsyncMock()),
65+
patch("postgres_mcp.server.mcp.run_streamable_http_async", AsyncMock()),
66+
):
67+
await main()
68+
69+
# Verify the host and port were set correctly
70+
assert mcp.settings.host == "0.0.0.0"
71+
assert mcp.settings.port == 9000
72+
finally:
73+
sys.argv = original_argv
74+
75+
76+
@pytest.mark.asyncio
77+
async def test_sse_host_port_arguments():
78+
"""Test that SSE host and port arguments are applied correctly."""
79+
from postgres_mcp.server import main
80+
from postgres_mcp.server import mcp
81+
82+
original_argv = sys.argv
83+
try:
84+
sys.argv = [
85+
"postgres_mcp",
86+
"postgresql://user:password@localhost/db",
87+
"--transport=sse",
88+
"--sse-host=0.0.0.0",
89+
"--sse-port=8080",
90+
]
91+
92+
with (
93+
patch("postgres_mcp.server.db_connection.pool_connect", AsyncMock()),
94+
patch("postgres_mcp.server.mcp.run_sse_async", AsyncMock()),
95+
):
96+
await main()
97+
98+
# Verify the host and port were set correctly
99+
assert mcp.settings.host == "0.0.0.0"
100+
assert mcp.settings.port == 8080
101+
finally:
102+
sys.argv = original_argv
103+
104+
105+
@pytest.mark.asyncio
106+
async def test_default_transport_is_stdio():
107+
"""Test that the default transport is stdio when not specified."""
108+
from postgres_mcp.server import main
109+
110+
original_argv = sys.argv
111+
try:
112+
sys.argv = [
113+
"postgres_mcp",
114+
"postgresql://user:password@localhost/db",
115+
]
116+
117+
with (
118+
patch("postgres_mcp.server.db_connection.pool_connect", AsyncMock()),
119+
patch("postgres_mcp.server.mcp.run_stdio_async", AsyncMock()) as mock_stdio,
120+
patch("postgres_mcp.server.mcp.run_sse_async", AsyncMock()) as mock_sse,
121+
patch("postgres_mcp.server.mcp.run_streamable_http_async", AsyncMock()) as mock_http,
122+
):
123+
await main()
124+
125+
mock_stdio.assert_called_once()
126+
mock_sse.assert_not_called()
127+
mock_http.assert_not_called()
128+
finally:
129+
sys.argv = original_argv

0 commit comments

Comments
 (0)