Skip to content

fix(gateway): fail fast + cache tool lookup on WS upgrade (+ upstream-failure logging) #1255

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mikhailm-coder wants to merge 5 commits into
base: main
Choose a base branch
from
Open

fix(gateway): fail fast + cache tool lookup on WS upgrade (+ upstream-failure logging) #1255

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
eb8c45f
fix(gateway): fail fast and cache tool lookup on WS upgrade
mikhailm-coder Jun 24, 2026
1dfd722
log: surface upstream WS connect/relay failures in the proxy client
mikhailm-coder Jun 24, 2026
4a293bc
log: sanitize upstream URL (host+path) to avoid logging query secrets
mikhailm-coder Jun 24, 2026
0b45e5d
Merge branch 'main' into hotfix/ws-tool-lookup-resilience
mikhailm-coder Jun 24, 2026
8d30f54
Merge branch 'main' into hotfix/ws-tool-lookup-resilience
mikhailm-coder Jun 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 12 additions & 2 deletions ...ore/src/main/java/com/openframe/gateway/config/ws/ProxySessionCleanupWebSocketClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,12 +25,22 @@ public class ProxySessionCleanupWebSocketClient implements WebSocketClient {

@Override
public Mono<Void> execute(URI url, WebSocketHandler handler) {
return delegate.execute(url, wrapHandler(url, handler));
return instrumentUpstream(url, delegate.execute(url, wrapHandler(url, handler)));
}

@Override
public Mono<Void> execute(URI url, HttpHeaders headers, WebSocketHandler handler) {
return delegate.execute(url, headers, wrapHandler(url, handler));
return instrumentUpstream(url, delegate.execute(url, headers, wrapHandler(url, handler)));
}

// Log the upstream proxy outcome: a connect failure/timeout to a dead upstream (e.g. meshcentral or NATS) is otherwise invisible here — the inbound upgrade just never completes and the agent sees "No HTTP response".
private Mono<Void> instrumentUpstream(URI url, Mono<Void> proxy) {
boolean debugPath = loggingProperties.isDebugPath(url.getPath());
// Host+path only: the upstream URI can carry forwarded query params (e.g. an agent key), which must not be logged.
String safeUrl = url.getScheme() + "://" + url.getHost() + (url.getPort() < 0 ? "" : ":" + url.getPort()) + url.getPath();
return proxy
.doOnSubscribe(s -> { if (debugPath) { log.debug("ws proxy: connecting to upstream {}", safeUrl); } })
.doOnError(ex -> log.warn("ws proxy: upstream connect/relay failed for {}: {}", safeUrl, ex.toString()));
}

private WebSocketHandler wrapHandler(URI targetUrl, WebSocketHandler handler) {
Expand Down
39 changes: 37 additions & 2 deletions ...rvice-core/src/main/java/com/openframe/gateway/config/ws/ToolWebSocketProxyUrlFilter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
package com.openframe.gateway.config.ws;

import com.github.benmanes.caffeine.cache.Cache;
import com.github.benmanes.caffeine.cache.Caffeine;
import com.openframe.data.document.tool.IntegratedTool;
import com.openframe.data.reactive.repository.tool.ReactiveIntegratedToolRepository;
import com.openframe.data.service.TenantIdProvider;
Expand All @@ -15,22 +17,34 @@
import org.springframework.core.Ordered;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpStatus;
import org.springframework.http.HttpStatusCode;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.web.server.ResponseStatusException;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.net.URI;
import java.time.Duration;

@RequiredArgsConstructor
@Slf4j
public abstract class ToolWebSocketProxyUrlFilter implements GatewayFilter, Ordered {

public static final String ORIGINAL_AUTHORIZATION_ATTR = "originalAuthorization";

// Bounds the per-upgrade tool lookup so a stalled Mongo query fails fast instead of letting the WS upgrade hang to the gateway response-timeout.
private static final Duration LOOKUP_TIMEOUT = Duration.ofSeconds(5);

private final ReactiveIntegratedToolRepository toolRepository;
private final ToolUpstreamResolverRegistry upstreamRegistry;

// Cache resolved tools briefly so an agent reconnect storm doesn't put a Mongo lookup on every WS upgrade (configs change rarely).
private final Cache<String, IntegratedTool> toolCache = Caffeine.newBuilder()
.maximumSize(10_000)
.expireAfterWrite(Duration.ofSeconds(60))
.build();
Comment on lines +43 to +46

@coderabbitai coderabbitai Bot Jun 24, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔒 Security & Privacy | 🟠 Major | 🏗️ Heavy lift

Avoid caching revocation state without invalidation.

Once an enabled tool is cached, later upgrades bypass the repository and the isEnabled() check for up to 60 seconds. If disabling/deleting a tool is an access-revocation action, add cache invalidation on tool updates or reduce/remove caching for that gate.

Also applies to: 112-132

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In
`@openframe-gateway-service-core/src/main/java/com/openframe/gateway/config/ws/ToolWebSocketProxyUrlFilter.java`
around lines 43 - 46, The ToolWebSocketProxyUrlFilter cache currently keeps
IntegratedTool entries for 60 seconds without any invalidation, so a tool that
has been disabled or deleted can still pass through via the cached value. Update
the caching behavior around the tool lookup and isEnabled() gate in
ToolWebSocketProxyUrlFilter so revocation changes are reflected immediately,
either by invalidating toolCache on tool update/delete events or by
removing/reducing caching for this access check.


/**
* Shared multi-tenant routing mode. When true, a tool lookup with no resolved tenant must NOT
* fall back to an unscoped query. Defaults false so single-tenant / OSS pods keep prior behavior.
Expand Down Expand Up @@ -68,7 +82,17 @@ public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {

ServerWebExchange mutatedExchange = mutateExchange(exchange, tool);
return chain.filter(mutatedExchange);
});
})
// Without this, a lookup error/timeout left the WS upgrade hanging with no response written, so the agent saw "No HTTP response" and retried forever; fail fast with a real status instead.
.onErrorResume(ex -> abortUpgrade(exchange, toolId, ex));
}

private Mono<Void> abortUpgrade(ServerWebExchange exchange, String toolId, Throwable ex) {
HttpStatusCode status = (ex instanceof ResponseStatusException rse) ? rse.getStatusCode() : HttpStatus.SERVICE_UNAVAILABLE;
log.warn("Tool websocket upgrade aborted (tool={}, status={}): {}", toolId, status, ex.toString());
ServerHttpResponse response = exchange.getResponse();
response.setStatusCode(status);
return response.setComplete();
Comment on lines +85 to +95

@coderabbitai coderabbitai Bot Jun 24, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟠 Major

🧩 Analysis chain

🌐 Web query:

Spring WebFlux ServerHttpResponse setStatusCode return value when response is already committed

💡 Result:

When you call setStatusCode(HttpStatusCode) on a Spring WebFlux ServerHttpResponse, the method returns a boolean value indicating whether the status code was successfully applied [1]. The return value behavior is as follows: - Returns true if the status code was successfully set [1]. - Returns false if the status code change could not be processed because the HTTP response is already committed [1]. A response is considered "committed" once the headers and the initial part of the response body have been sent to the client [2][3]. Once the response is in this state, it is no longer possible to modify the status code or headers, as they have already been written to the underlying network stream [1][4]. In practice, developers often use the boolean return value to conditionally handle errors or log status changes, ensuring that they do not attempt to set a status code when it is too late to do so [5][6]. Because attempts to change the status after commitment are ignored (resulting in a false return value), application logic should be designed to handle potential failures or status changes before any part of the response body is emitted [3]. If you need to perform actions or set statuses precisely before the response is committed, you can use the beforeCommit(Supplier) hook provided by the ServerHttpResponse [1][4].

Citations:

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Locate the target file and inspect the surrounding implementation.
git ls-files 'openframe-gateway-service-core/src/main/java/com/openframe/gateway/config/ws/ToolWebSocketProxyUrlFilter.java'
wc -l openframe-gateway-service-core/src/main/java/com/openframe/gateway/config/ws/ToolWebSocketProxyUrlFilter.java
sed -n '1,220p' openframe-gateway-service-core/src/main/java/com/openframe/gateway/config/ws/ToolWebSocketProxyUrlFilter.java

Repository: flamingo-stack/openframe-oss-lib

Length of output: 7260

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Inspect the filter implementation and related WebSocket upgrade flow.
sed -n '1,220p' openframe-gateway-service-core/src/main/java/com/openframe/gateway/config/ws/ToolWebSocketProxyUrlFilter.java

printf '\n--- SEARCH ---\n'
rg -n "chain\.filter|onErrorResume|setStatusCode|isCommitted|upgrade|WebSocket" openframe-gateway-service-core/src/main/java/com/openframe/gateway/config/ws -g '*.java'

Repository: flamingo-stack/openframe-oss-lib

Length of output: 19048

Skip abortUpgrade(...) once the response is committed. setStatusCode(...) is a no-op after commit, so this onErrorResume can still complete successfully while the client gets no usable HTTP status. Guard with response.isCommitted() or the setStatusCode(...) return value and let the error propagate when the response can no longer be modified.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In
`@openframe-gateway-service-core/src/main/java/com/openframe/gateway/config/ws/ToolWebSocketProxyUrlFilter.java`
around lines 85 - 95, The abort path in
ToolWebSocketProxyUrlFilter.abortUpgrade() should skip writing a status once the
ServerHttpResponse is already committed, since setStatusCode() cannot take
effect then. Update the onErrorResume flow to check response.isCommitted() (or
the boolean result of setStatusCode()) before calling setComplete(), and if the
response can no longer be modified, let the error propagate instead of returning
a successful Mono<Void>. Keep the existing abortUpgrade(ServerWebExchange,
String, Throwable) and abortUpgrade(...) logging/status handling as the main
place to make this change.

}

protected ServerWebExchange mutateExchange(ServerWebExchange exchange, IntegratedTool tool) {
Expand All @@ -82,6 +106,14 @@ protected ServerWebExchange mutateExchange(ServerWebExchange exchange, Integrate
* headers are never read; the unscoped {@code findByKey} is correct (one tenant only).
*/
private Mono<IntegratedTool> getTool(String toolId, ServerHttpRequest request) {
String cacheKey = tenantRoutingEnabled
? (TenantRoutingHeaders.tenantId(request) + "::" + toolId)
: toolId;
IntegratedTool cached = toolCache.getIfPresent(cacheKey);
if (cached != null) {
return Mono.just(cached);
}

Mono<IntegratedTool> lookup = tenantRoutingEnabled
? toolRepository.findByTenantIdAndKey(TenantRoutingHeaders.tenantId(request), toolId)
: toolRepository.findByKey(toolId);
Expand All @@ -94,7 +126,10 @@ private Mono<IntegratedTool> getTool(String toolId, ServerHttpRequest request) {
return Mono.error(new IllegalArgumentException("Tool " + tool.getName() + " is not enabled"));
}
return Mono.just(tool);
});
})
.timeout(LOOKUP_TIMEOUT)
// Only successful (present + enabled) tools are cached; not-found / not-enabled stay uncached so they recover immediately once fixed.
.doOnNext(tool -> toolCache.put(cacheKey, tool));
}

protected abstract String getRequestToolId(String path);
Expand Down
Loading