From e999982848c16b3c2d5458c389aa2132deed1183 Mon Sep 17 00:00:00 2001
From: Taylor Wilsdon <tbarrettwilsdon@gmail.com>
Date: Sun, 1 Mar 2026 18:21:23 -0500
Subject: [PATCH] cache test

---
 ...est_well_known_cache_control_middleware.py | 88 +++++++++++++++++++
 1 file changed, 88 insertions(+)
 create mode 100644 tests/core/test_well_known_cache_control_middleware.py

diff --git a/tests/core/test_well_known_cache_control_middleware.py b/tests/core/test_well_known_cache_control_middleware.py
new file mode 100644
index 0000000..66e9b98
--- /dev/null
+++ b/tests/core/test_well_known_cache_control_middleware.py
@@ -0,0 +1,88 @@
+import importlib
+
+from starlette.applications import Starlette
+from starlette.middleware import Middleware
+from starlette.responses import Response
+from starlette.routing import Route
+from starlette.testclient import TestClient
+
+
+def test_well_known_cache_control_middleware_rewrites_headers():
+    from core.server import WellKnownCacheControlMiddleware, _compute_scope_fingerprint
+
+    async def well_known_endpoint(request):
+        response = Response("ok")
+        response.headers["Cache-Control"] = "public, max-age=3600"
+        response.set_cookie("a", "1")
+        response.set_cookie("b", "2")
+        return response
+
+    async def regular_endpoint(request):
+        response = Response("ok")
+        response.headers["Cache-Control"] = "public, max-age=3600"
+        return response
+
+    app = Starlette(
+        routes=[
+            Route("/.well-known/oauth-authorization-server", well_known_endpoint),
+            Route("/.well-known/oauth-authorization-server-extra", regular_endpoint),
+            Route("/health", regular_endpoint),
+        ],
+        middleware=[Middleware(WellKnownCacheControlMiddleware)],
+    )
+    client = TestClient(app)
+
+    well_known = client.get("/.well-known/oauth-authorization-server")
+    assert well_known.status_code == 200
+    assert well_known.headers["cache-control"] == "no-store, must-revalidate"
+    assert well_known.headers["etag"] == f'"{_compute_scope_fingerprint()}"'
+    assert sorted(well_known.headers.get_list("set-cookie")) == sorted(
+        ["a=1; Path=/; SameSite=lax", "b=2; Path=/; SameSite=lax"]
+    )
+
+    regular = client.get("/health")
+    assert regular.status_code == 200
+    assert regular.headers["cache-control"] == "public, max-age=3600"
+    assert "etag" not in regular.headers
+
+    extra = client.get("/.well-known/oauth-authorization-server-extra")
+    assert extra.status_code == 200
+    assert extra.headers["cache-control"] == "public, max-age=3600"
+    assert "etag" not in extra.headers
+
+
+def test_configured_server_applies_no_cache_to_served_oauth_discovery_routes(monkeypatch):
+    monkeypatch.setenv("MCP_ENABLE_OAUTH21", "true")
+    monkeypatch.setenv("GOOGLE_OAUTH_CLIENT_ID", "dummy-client")
+    monkeypatch.setenv("GOOGLE_OAUTH_CLIENT_SECRET", "dummy-secret")
+    monkeypatch.setenv("WORKSPACE_MCP_BASE_URI", "http://localhost")
+    monkeypatch.setenv("WORKSPACE_MCP_PORT", "8000")
+    monkeypatch.delenv("WORKSPACE_EXTERNAL_URL", raising=False)
+    monkeypatch.setenv("EXTERNAL_OAUTH21_PROVIDER", "false")
+
+    import core.server as core_server
+    from auth.oauth_config import reload_oauth_config
+
+    reload_oauth_config()
+    core_server = importlib.reload(core_server)
+    core_server.set_transport_mode("streamable-http")
+    core_server.configure_server_for_http()
+
+    app = core_server.server.http_app(transport="streamable-http", path="/mcp")
+    client = TestClient(app)
+
+    authorization_server = client.get("/.well-known/oauth-authorization-server")
+    assert authorization_server.status_code == 200
+    assert authorization_server.headers["cache-control"] == "no-store, must-revalidate"
+    assert authorization_server.headers["etag"].startswith('"')
+    assert authorization_server.headers["etag"].endswith('"')
+
+    protected_resource = client.get("/.well-known/oauth-protected-resource/mcp")
+    assert protected_resource.status_code == 200
+    assert protected_resource.headers["cache-control"] == "no-store, must-revalidate"
+    assert protected_resource.headers["etag"].startswith('"')
+    assert protected_resource.headers["etag"].endswith('"')
+
+    # Ensure we did not create a shadow route at the wrong path.
+    wrong_path = client.get("/.well-known/oauth-protected-resource")
+    assert wrong_path.status_code == 404