mirror of
https://github.com/vercel-labs/fx.git
synced 2026-10-28 10:15:17 +08:00
Reconnect when an MCP server drops an idle connection
- An MCP tool call over HTTP after the server dropped an idle connection now goes out on a new connection, instead of failing with "the server went away" on every call until fx restarts. - Renewing an MCP login after the sign-in server dropped an idle connection now goes through, instead of failing with "the sign-in server didn't answer" until fx restarts. - A connection the server closed cleanly while idle is replaced before the call is sent, so the call no longer fails once with "the tool may have run". - Connections unused for more than a minute are replaced instead of reused, so a call after the network silently dropped one no longer waits out its timeout. - A request whose send fails goes out once more on a new connection. A request that was fully sent is still never sent twice.
This commit is contained in:
4 files changed
+483
-22
No files matched your search
+254
-22
@@ -26,6 +26,11 @@ pub const Options = struct {
|
||||
max_in_flight: usize = 32,
|
||||
/// Rung after each post, for a host that reads many clients.
|
||||
bell: ?*Bell = null,
|
||||
/// Once the client has sent nothing for this long, the connections it
|
||||
/// kept are closed rather than used again: a network can drop one that
|
||||
/// sat idle without either end hearing of it, and a request sent on it
|
||||
/// would wait out its timeout.
|
||||
max_idle_ms: u32 = 60_000,
|
||||
};
|
||||
|
||||
pub const Method = enum { post, get, delete };
|
||||
@@ -137,6 +142,10 @@ pub const Client = struct {
|
||||
/// Bytes of messages posted and not yet taken: at most one largest
|
||||
/// message, so a fast server waits instead of growing fx.
|
||||
budget: Budget,
|
||||
/// When a request last started or ended, in the awake clock's
|
||||
/// milliseconds; under `used_lock`.
|
||||
last_used_ms: i64 = 0,
|
||||
used_lock: Io.Mutex = .init,
|
||||
|
||||
/// `options.url` must outlive the client.
|
||||
pub fn init(c: *Client, io: Io, gpa: std.mem.Allocator, options: Options) Error!void {
|
||||
@@ -408,21 +417,13 @@ pub const Client = struct {
|
||||
fn run(c: *Client, key: u64, method: Method, body: ?[]u8, headers: []const std.http.Header) Io.Cancelable!Result {
|
||||
// One that can't go out as it is never does: nothing takes it.
|
||||
if (!http_headers.sendable(headers)) return .{ .how = .http_failure };
|
||||
var req = c.http.request(switch (method) {
|
||||
.post => .POST,
|
||||
.get => .GET,
|
||||
.delete => .DELETE,
|
||||
}, c.uri, .{
|
||||
defer _ = c.markUsed();
|
||||
var req = (try c.openSent(method, c.uri, .{
|
||||
.headers = .{ .content_type = if (body != null) .{ .override = "application/json" } else .default, .accept_encoding = .{ .override = "identity" } },
|
||||
.extra_headers = headers,
|
||||
.redirect_behavior = .unhandled,
|
||||
}) catch return .{ .how = .broken };
|
||||
}, body)) orelse return .{ .how = .broken };
|
||||
defer req.deinit();
|
||||
req.accept_encoding = identity_only;
|
||||
if (body) |b| {
|
||||
req.transfer_encoding = .{ .content_length = b.len };
|
||||
req.sendBodyComplete(b) catch return .{ .how = .broken };
|
||||
} else req.sendBodiless() catch return .{ .how = .broken };
|
||||
var response = req.receiveHead(&.{}) catch return .{ .how = .broken };
|
||||
const status: u16 = @backingInt(response.head.status);
|
||||
// The head's bytes don't outlive reading the body.
|
||||
@@ -484,6 +485,79 @@ pub const Client = struct {
|
||||
}
|
||||
}
|
||||
|
||||
/// Opens `method` on `uri`, on a connection kept for its server or a new
|
||||
/// one, and sends `body`. A send that fails never brought the server the
|
||||
/// whole request, so it is made once more, on a new connection; and the
|
||||
/// connection it failed on is closed, where `std.http` would keep it for
|
||||
/// the next request to fail on too. Null when it couldn't go out.
|
||||
fn openSent(c: *Client, method: Method, uri: std.Uri, options: std.http.Client.RequestOptions, body: ?[]u8) Io.Cancelable!?std.http.Client.Request {
|
||||
var fresh = false;
|
||||
while (true) {
|
||||
var with = options;
|
||||
with.connection = try c.takeKept(uri, fresh);
|
||||
var req = c.http.request(switch (method) {
|
||||
.post => .POST,
|
||||
.get => .GET,
|
||||
.delete => .DELETE,
|
||||
}, uri, with) catch {
|
||||
if (with.connection) |kept| c.http.connection_pool.release(kept, c.io);
|
||||
return null;
|
||||
};
|
||||
req.accept_encoding = identity_only;
|
||||
const sending = if (body) |b| send: {
|
||||
req.transfer_encoding = .{ .content_length = b.len };
|
||||
break :send req.sendBodyComplete(b);
|
||||
} else req.sendBodiless();
|
||||
sending catch {
|
||||
req.connection.?.closing = true;
|
||||
req.deinit();
|
||||
if (fresh) return null;
|
||||
fresh = true;
|
||||
continue;
|
||||
};
|
||||
return req;
|
||||
}
|
||||
}
|
||||
|
||||
/// A connection kept for `uri`'s server, taken from the pool, or null for
|
||||
/// a new one. Those kept that the server closed or reset while they sat
|
||||
/// idle are closed here, before anything is sent on them, and so is
|
||||
/// every one kept when `fresh` is set or the client sat idle past
|
||||
/// `max_idle_ms`.
|
||||
fn takeKept(c: *Client, uri: std.Uri, fresh: bool) Io.Cancelable!?*std.http.Client.Connection {
|
||||
// Through a proxy, `std.http` finds its own.
|
||||
if (c.http.http_proxy != null or c.http.https_proxy != null) return null;
|
||||
const protocol = std.http.Client.Protocol.fromUri(uri) orelse return null;
|
||||
var host: [Io.net.HostName.max_len]u8 = undefined;
|
||||
const criteria: std.http.Client.ConnectionPool.Criteria = .{
|
||||
.host = Io.net.HostName.fromUri(uri, &host) catch return null,
|
||||
.port = uri.port orelse switch (protocol) {
|
||||
.plain => 80,
|
||||
.tls => 443,
|
||||
},
|
||||
.protocol = protocol,
|
||||
};
|
||||
const stale = fresh or c.markUsed() > c.options.max_idle_ms;
|
||||
const pool = &c.http.connection_pool;
|
||||
while (try pool.findConnection(c.io, criteria)) |kept| {
|
||||
if (!stale and !closedByServer(kept)) return kept;
|
||||
kept.closing = true;
|
||||
pool.release(kept, c.io);
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
/// Notes that a request starts or ends now, and returns how long none
|
||||
/// had before.
|
||||
fn markUsed(c: *Client) i64 {
|
||||
const at = now(c.io);
|
||||
c.used_lock.lockUncancelable(c.io);
|
||||
defer c.used_lock.unlock(c.io);
|
||||
const idle = at - c.last_used_ms;
|
||||
c.last_used_ms = at;
|
||||
return idle;
|
||||
}
|
||||
|
||||
/// Reads a JSON body into `body` a step at a time, counting each step
|
||||
/// before it is read. One past `max_message` is `TooLong`.
|
||||
fn readBody(c: *Client, reader: *Io.Reader, read: *Read, body: *std.ArrayList(u8)) (Io.Cancelable || error{ TooLong, ReadFailed, OutOfMemory })!void {
|
||||
@@ -519,21 +593,13 @@ pub const Client = struct {
|
||||
fn fetchRun(c: *Client, key: u64, method: Method, uri: std.Uri, body: ?[]u8, content_type: []const u8, headers: []const std.http.Header) Io.Cancelable!@FieldType(Posted, "fetched") {
|
||||
const none: @FieldType(Posted, "fetched") = .{ .key = key };
|
||||
if (!http_headers.sendable(headers)) return none;
|
||||
var req = c.http.request(switch (method) {
|
||||
.post => .POST,
|
||||
.get => .GET,
|
||||
.delete => .DELETE,
|
||||
}, uri, .{
|
||||
defer _ = c.markUsed();
|
||||
var req = (try c.openSent(method, uri, .{
|
||||
.headers = .{ .content_type = if (body != null) .{ .override = content_type } else .default, .accept_encoding = .{ .override = "identity" } },
|
||||
.extra_headers = headers,
|
||||
.redirect_behavior = .unhandled,
|
||||
}) catch return none;
|
||||
}, body)) orelse return none;
|
||||
defer req.deinit();
|
||||
req.accept_encoding = identity_only;
|
||||
if (body) |b| {
|
||||
req.transfer_encoding = .{ .content_length = b.len };
|
||||
req.sendBodyComplete(b) catch return none;
|
||||
} else req.sendBodiless() catch return none;
|
||||
var response = req.receiveHead(&.{}) catch return none;
|
||||
var got: @FieldType(Posted, "fetched") = .{ .key = key, .status = @backingInt(response.head.status) };
|
||||
var challenge: std.ArrayList(u8) = .empty;
|
||||
@@ -636,6 +702,15 @@ fn now(io: Io) i64 {
|
||||
return Io.Clock.awake.now(io).toMilliseconds();
|
||||
}
|
||||
|
||||
/// Whether the server closed or reset a kept connection while it sat idle.
|
||||
/// Nothing is due on one between requests, so anything to read, its end
|
||||
/// included, means it can't carry another.
|
||||
fn closedByServer(kept: *std.http.Client.Connection) bool {
|
||||
if (kept.reader().bufferedLen() > 0 or kept.stream_reader.interface.bufferedLen() > 0) return true;
|
||||
var ready = [_]std.posix.pollfd{.{ .fd = kept.stream_reader.stream.socket.handle, .events = std.posix.POLL.IN, .revents = 0 }};
|
||||
return (std.posix.poll(&ready, 0) catch return true) > 0;
|
||||
}
|
||||
|
||||
/// A Content-Type naming JSON, with any parameters such as `charset`.
|
||||
pub fn isJson(media: []const u8) bool {
|
||||
const end = std.mem.findScalar(u8, media, ';') orelse media.len;
|
||||
@@ -962,3 +1037,160 @@ test "a timed wait starts no task, so it works where none can run" {
|
||||
defer c.deinit();
|
||||
try testing.expect((try c.next(20)) == .timeout);
|
||||
}
|
||||
|
||||
/// What a test server does to each connection once it has answered on it.
|
||||
const Drop = enum { keep, reset, close };
|
||||
|
||||
/// Reads one request on `r`, head and body, so none is left unread when the
|
||||
/// connection closes: a close with unread bytes is a reset.
|
||||
fn readRequest(r: *Io.Reader) !void {
|
||||
const name = "content-length:";
|
||||
var length: usize = 0;
|
||||
while (true) {
|
||||
const line = try r.takeDelimiterInclusive('\n');
|
||||
if (line.len <= 2) break;
|
||||
if (std.ascii.startsWithIgnoreCase(line, name)) length = try std.fmt.parseInt(usize, std.mem.trim(u8, line[name.len..], " \r\n"), 10);
|
||||
}
|
||||
try r.discardAll(length);
|
||||
}
|
||||
|
||||
/// Answers each request on `stream` with a result, keeping the connection
|
||||
/// for the next, until `drop` ends it after the first answer. One the client
|
||||
/// ends is held open, unanswered, until the test does.
|
||||
fn answerKept(stream: Io.net.Stream, drop: Drop, requests: *std.atomic.Value(u32)) !void {
|
||||
defer stream.close(testing.io);
|
||||
var in: [4096]u8 = undefined;
|
||||
var r = stream.reader(testing.io, &in);
|
||||
var out: [1024]u8 = undefined;
|
||||
var w = stream.writer(testing.io, &out);
|
||||
const answer = "{\"jsonrpc\":\"2.0\",\"id\":1,\"result\":{}}";
|
||||
while (true) {
|
||||
readRequest(&r.interface) catch return Io.sleep(testing.io, .fromSeconds(30), .awake);
|
||||
_ = requests.fetchAdd(1, .monotonic);
|
||||
try w.interface.print("HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {d}\r\n\r\n", .{answer.len});
|
||||
try w.interface.writeAll(answer);
|
||||
try w.interface.flush();
|
||||
switch (drop) {
|
||||
.keep => {},
|
||||
.close => return,
|
||||
.reset => {
|
||||
// Once the client has read the answer: a reset, not an end.
|
||||
try Io.sleep(testing.io, .fromMilliseconds(50), .awake);
|
||||
try std.posix.setsockopt(stream.socket.handle, std.posix.SOL.SOCKET, std.posix.SO.LINGER, std.mem.asBytes(&std.posix.linger{ .onoff = 1, .linger = 0 }));
|
||||
return;
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn serveKept(server: *Io.net.Server, drop: Drop, connections: *std.atomic.Value(u32), requests: *std.atomic.Value(u32)) !void {
|
||||
var answering: [4]?Io.Future(@typeInfo(@TypeOf(answerKept)).@"fn".return_type.?) = @splat(null);
|
||||
defer for (&answering) |*a| if (a.*) |*f| f.cancel(testing.io) catch {};
|
||||
for (&answering) |*a| {
|
||||
const stream = try server.accept(testing.io);
|
||||
_ = connections.fetchAdd(1, .monotonic);
|
||||
a.* = try testing.io.concurrent(answerKept, .{ stream, drop, requests });
|
||||
}
|
||||
try Io.sleep(testing.io, .fromSeconds(30), .awake);
|
||||
}
|
||||
|
||||
/// A server that keeps connections alive, counting them and the requests on
|
||||
/// them. It must not move once started.
|
||||
const Kept = struct {
|
||||
server: Io.net.Server,
|
||||
connections: std.atomic.Value(u32) = .init(0),
|
||||
requests: std.atomic.Value(u32) = .init(0),
|
||||
serving: ?Io.Future(@typeInfo(@TypeOf(serveKept)).@"fn".return_type.?) = null,
|
||||
url_buffer: [64]u8 = undefined,
|
||||
url: []const u8 = "",
|
||||
|
||||
fn start(k: *Kept, drop: Drop) !void {
|
||||
var address: Io.net.IpAddress = .{ .ip4 = .loopback(0) };
|
||||
k.* = .{ .server = try address.listen(testing.io, .{ .reuse_address = true }) };
|
||||
k.url = try std.mem.print(&k.url_buffer, "http://127.0.0.1:{d}/mcp", .{k.server.socket.address.getPort()});
|
||||
k.serving = try testing.io.concurrent(serveKept, .{ &k.server, drop, &k.connections, &k.requests });
|
||||
}
|
||||
|
||||
fn stop(k: *Kept) void {
|
||||
if (k.serving) |*f| f.cancel(testing.io) catch {};
|
||||
k.server.deinit(testing.io);
|
||||
}
|
||||
};
|
||||
|
||||
/// How the request `c` sent last ended, past its messages.
|
||||
fn endingOf(c: *Client) !Ending {
|
||||
while (true) switch (try c.next(5_000)) {
|
||||
.ended => |e| return e.how,
|
||||
.timeout => return error.TestTimeout,
|
||||
else => {},
|
||||
};
|
||||
}
|
||||
|
||||
fn fetchedStatus(c: *Client) !u16 {
|
||||
while (true) switch (try c.next(5_000)) {
|
||||
.fetched => |f| return f.status,
|
||||
.timeout => return error.TestTimeout,
|
||||
else => {},
|
||||
};
|
||||
}
|
||||
|
||||
test "a kept connection the server reset or closed while idle isn't sent on again: each request goes out once, on a new one" {
|
||||
for ([_]Drop{ .reset, .close }) |drop| {
|
||||
var kept: Kept = undefined;
|
||||
try kept.start(drop);
|
||||
defer kept.stop();
|
||||
var c: Client = undefined;
|
||||
try c.init(testing.io, testing.allocator, .{ .url = kept.url });
|
||||
defer c.deinit();
|
||||
try c.send(1, .post, "{}", &.{});
|
||||
try testing.expectEqual(Ending.complete, try endingOf(&c));
|
||||
// The server drops each connection after one answer, so the next
|
||||
// request and the sign-in request after it find theirs gone.
|
||||
try Io.sleep(testing.io, .fromMilliseconds(300), .awake);
|
||||
try c.send(2, .post, "{}", &.{});
|
||||
try testing.expectEqual(Ending.complete, try endingOf(&c));
|
||||
try Io.sleep(testing.io, .fromMilliseconds(300), .awake);
|
||||
try c.fetch(3, .post, kept.url, "{}", "application/json", &.{});
|
||||
try testing.expectEqual(@as(u16, 200), try fetchedStatus(&c));
|
||||
try testing.expectEqual(@as(u32, 3), kept.connections.load(.monotonic));
|
||||
try testing.expectEqual(@as(u32, 3), kept.requests.load(.monotonic));
|
||||
}
|
||||
}
|
||||
|
||||
test "a request whose send fails goes out once more, on a new connection, and reaches the server once" {
|
||||
var kept: Kept = undefined;
|
||||
try kept.start(.keep);
|
||||
defer kept.stop();
|
||||
var c: Client = undefined;
|
||||
try c.init(testing.io, testing.allocator, .{ .url = kept.url });
|
||||
defer c.deinit();
|
||||
try c.send(1, .post, "{}", &.{});
|
||||
try testing.expectEqual(Ending.complete, try endingOf(&c));
|
||||
// The connection kept for the next request can't send any more, and
|
||||
// nothing says so until a send on it fails.
|
||||
const node = c.http.connection_pool.free.last orelse return error.TestUnexpectedResult;
|
||||
const connection: *std.http.Client.Connection = @alignCast(@fieldParentPtr("pool_node", node));
|
||||
try connection.stream_reader.stream.shutdown(testing.io, .send);
|
||||
try c.send(2, .post, "{}", &.{});
|
||||
try testing.expectEqual(Ending.complete, try endingOf(&c));
|
||||
try testing.expectEqual(@as(u32, 2), kept.connections.load(.monotonic));
|
||||
try testing.expectEqual(@as(u32, 2), kept.requests.load(.monotonic));
|
||||
}
|
||||
|
||||
test "a kept connection is used again, unless the client sat idle past max_idle_ms" {
|
||||
for ([_]u32{ 60_000, 0 }, [_]u32{ 1, 2 }) |max_idle_ms, connections| {
|
||||
var kept: Kept = undefined;
|
||||
try kept.start(.keep);
|
||||
defer kept.stop();
|
||||
var c: Client = undefined;
|
||||
try c.init(testing.io, testing.allocator, .{ .url = kept.url, .max_idle_ms = max_idle_ms });
|
||||
defer c.deinit();
|
||||
for (1..3) |key| {
|
||||
try c.send(key, .post, "{}", &.{});
|
||||
try testing.expectEqual(Ending.complete, try endingOf(&c));
|
||||
try Io.sleep(testing.io, .fromMilliseconds(20), .awake);
|
||||
}
|
||||
try testing.expectEqual(connections, kept.connections.load(.monotonic));
|
||||
try testing.expectEqual(@as(u32, 2), kept.requests.load(.monotonic));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,134 @@
|
||||
import type { Socket } from "bun";
|
||||
|
||||
/**
|
||||
* How the proxy drops a connection that sat idle: `reset` aborts it, so the
|
||||
* client's next send on it fails at once, and `close` closes it cleanly, so
|
||||
* the send goes out and only the answer is missing.
|
||||
*/
|
||||
export type IdleDrop = "reset" | "close";
|
||||
|
||||
type Link = {
|
||||
upstream: Socket<undefined> | null;
|
||||
/** What the client sent before the upstream connection opened. */
|
||||
early: Uint8Array[];
|
||||
/** What each side's socket hasn't taken yet. */
|
||||
toClient: Uint8Array[];
|
||||
toUpstream: Uint8Array[];
|
||||
timer: ReturnType<typeof setTimeout> | null;
|
||||
dropped: boolean;
|
||||
};
|
||||
|
||||
/**
|
||||
* A TCP proxy to `target`'s host and port that drops each client connection
|
||||
* once it has carried nothing either way for `idleMs`, as a load balancer
|
||||
* drops an idle keep-alive connection. `url` is `target` through the proxy.
|
||||
*/
|
||||
export async function startIdleDropProxy(target: string, drop: IdleDrop, idleMs = 150) {
|
||||
const upstreamUrl = new URL(target);
|
||||
let connections = 0;
|
||||
let dropped = 0;
|
||||
const clients = new Set<Socket<Link>>();
|
||||
|
||||
const send = (socket: Socket<any>, queue: Uint8Array[], chunk: Uint8Array) => {
|
||||
if (queue.length > 0) {
|
||||
queue.push(chunk);
|
||||
return;
|
||||
}
|
||||
const written = socket.write(chunk);
|
||||
if (written < chunk.length) queue.push(chunk.subarray(Math.max(written, 0)));
|
||||
};
|
||||
const flush = (socket: Socket<any>, queue: Uint8Array[]) => {
|
||||
while (queue.length > 0) {
|
||||
const chunk = queue[0]!;
|
||||
const written = socket.write(chunk);
|
||||
if (written < chunk.length) {
|
||||
queue[0] = chunk.subarray(Math.max(written, 0));
|
||||
return;
|
||||
}
|
||||
queue.shift();
|
||||
}
|
||||
};
|
||||
const quiet = (client: Socket<Link>) => {
|
||||
const link = client.data;
|
||||
if (link.timer) clearTimeout(link.timer);
|
||||
link.timer = setTimeout(() => {
|
||||
link.dropped = true;
|
||||
dropped += 1;
|
||||
if (drop === "reset") client.terminate();
|
||||
else client.end();
|
||||
link.upstream?.end();
|
||||
}, idleMs);
|
||||
};
|
||||
|
||||
const server = Bun.listen<Link>({
|
||||
hostname: "127.0.0.1",
|
||||
port: 0,
|
||||
socket: {
|
||||
open(client) {
|
||||
connections += 1;
|
||||
clients.add(client);
|
||||
client.data = { upstream: null, early: [], toClient: [], toUpstream: [], timer: null, dropped: false };
|
||||
quiet(client);
|
||||
Bun.connect({
|
||||
hostname: upstreamUrl.hostname,
|
||||
port: Number(upstreamUrl.port),
|
||||
socket: {
|
||||
open(upstream) {
|
||||
client.data.upstream = upstream;
|
||||
for (const chunk of client.data.early) send(upstream, client.data.toUpstream, chunk);
|
||||
client.data.early = [];
|
||||
},
|
||||
data(_upstream, chunk) {
|
||||
if (client.data.dropped) return;
|
||||
send(client, client.data.toClient, chunk);
|
||||
quiet(client);
|
||||
},
|
||||
drain(upstream) {
|
||||
flush(upstream, client.data.toUpstream);
|
||||
},
|
||||
close() {
|
||||
if (!client.data.dropped) client.end();
|
||||
},
|
||||
error() {
|
||||
if (!client.data.dropped) client.end();
|
||||
},
|
||||
},
|
||||
}).catch(() => client.end());
|
||||
},
|
||||
data(client, chunk) {
|
||||
if (client.data.dropped) return;
|
||||
const upstream = client.data.upstream;
|
||||
if (upstream) send(upstream, client.data.toUpstream, chunk);
|
||||
else client.data.early.push(chunk);
|
||||
quiet(client);
|
||||
},
|
||||
drain(client) {
|
||||
flush(client, client.data.toClient);
|
||||
},
|
||||
close(client) {
|
||||
clients.delete(client);
|
||||
if (client.data.timer) clearTimeout(client.data.timer);
|
||||
client.data.upstream?.end();
|
||||
},
|
||||
error() {},
|
||||
},
|
||||
});
|
||||
|
||||
const origin = `http://127.0.0.1:${server.port}`;
|
||||
return {
|
||||
origin,
|
||||
url: `${origin}${upstreamUrl.pathname}${upstreamUrl.search}`,
|
||||
/** Client connections accepted so far. */
|
||||
get connections() {
|
||||
return connections;
|
||||
},
|
||||
/** Client connections dropped for sitting idle. */
|
||||
get dropped() {
|
||||
return dropped;
|
||||
},
|
||||
stop() {
|
||||
for (const client of clients) client.terminate();
|
||||
server.stop(true);
|
||||
},
|
||||
};
|
||||
}
|
||||
@@ -15,6 +15,7 @@ import {
|
||||
import { tmpdir } from "node:os";
|
||||
import { delimiter, join } from "node:path";
|
||||
import { FX_BIN, runFx } from "../evals/eval-helpers";
|
||||
import { startIdleDropProxy } from "./fixtures/idle-drop-proxy";
|
||||
import { startLegacyStreamableHttpFixture } from "./fixtures/mcp-legacy-remote";
|
||||
import { startModernMcpHttpFixture } from "./fixtures/mcp-modern-http";
|
||||
import { stdoutFrames } from "./render-lab/tape";
|
||||
@@ -47,6 +48,7 @@ let legacyStreamable:
|
||||
let auth: AuthFixture | null = null;
|
||||
let gateway: ReturnType<typeof startFakeGateway> | null = null;
|
||||
let tui: TmuxSession | null = null;
|
||||
let tokenProxy: Awaited<ReturnType<typeof startIdleDropProxy>> | null = null;
|
||||
|
||||
beforeAll(() => {
|
||||
process.env.FX_DISABLE_KEYCHAIN = "1";
|
||||
@@ -101,16 +103,19 @@ afterEach(async () => {
|
||||
const activeAuth = auth;
|
||||
const activeUpstream = upstream;
|
||||
const activeLegacyStreamable = legacyStreamable;
|
||||
const activeTokenProxy = tokenProxy;
|
||||
const activeCleanupRoot = cleanupRoot;
|
||||
tui = null;
|
||||
gateway = null;
|
||||
auth = null;
|
||||
upstream = null;
|
||||
legacyStreamable = null;
|
||||
tokenProxy = null;
|
||||
cleanupRoot = null;
|
||||
|
||||
if (activeTui) await activeTui.kill();
|
||||
activeGateway?.stop();
|
||||
activeTokenProxy?.stop();
|
||||
activeAuth?.stop();
|
||||
activeUpstream?.stop();
|
||||
activeLegacyStreamable?.stop();
|
||||
@@ -1948,6 +1953,47 @@ describe("MCP login", () => {
|
||||
expect(auth.replays).toBe(0);
|
||||
}, 90_000);
|
||||
|
||||
// A load balancer resets the token endpoint's keep-alive connection once it
|
||||
// sits idle. fx keeps it for the next renewal, which must find it gone and
|
||||
// go out on a new connection instead of failing every renewal until fx
|
||||
// restarts.
|
||||
test("a renewal after the token endpoint reset its idle connection still goes through", async () => {
|
||||
upstream = startModernMcpHttpFixture("json");
|
||||
auth = startAuthFixture(upstream.url, { singleUseRefresh: { graceMs: 0 } });
|
||||
tokenProxy = await startIdleDropProxy(auth.url, "reset");
|
||||
const root = createRoot(auth);
|
||||
seedV2Credential(root, auth, { token_endpoint: `${tokenProxy.origin}/token` });
|
||||
const activeAuth = auth;
|
||||
// Each call finds its access token withdrawn, so fx renews the login
|
||||
// first. The second comes after the proxy reset the first one's
|
||||
// connection.
|
||||
const renewed = (withdrawn: string, id: string, waitMs: number) => async () => {
|
||||
await Bun.sleep(waitMs);
|
||||
activeAuth.validTokens.delete(withdrawn);
|
||||
return fakeGatewayToolCall(id, TOOL_NAME, { text: id });
|
||||
};
|
||||
gateway = startFakeGateway([
|
||||
fakeGatewayToolCall("select_mcp", "mcp_select_tool", { name: TOOL_NAME }),
|
||||
renewed(ACCESS_INITIAL, "first", 0),
|
||||
renewed(`${ACCESS_REFRESHED}-1`, "second", 600),
|
||||
fakeGatewayFinalText("Renewed twice."),
|
||||
], { models: [{ id: MODEL, type: "language", tags: ["tool-use"] }] });
|
||||
const result = await runFx(["ask", "--json", "--auto", "--no-save", "Call the protected fixture twice."], {
|
||||
cwd: root.workspace,
|
||||
env: { ...v2Env(root), FX_GATEWAY_BASE_URL: gateway.baseUrl, FX_GATEWAY_CHAT_URL: gateway.chatUrl },
|
||||
timeoutMs: 30_000,
|
||||
});
|
||||
expectExit0(result);
|
||||
const last = gateway.requests.at(-1)?.body ?? "";
|
||||
expect(toolResultText(last, "first")).toContain("first");
|
||||
expect(toolResultText(last, "second")).toContain("second");
|
||||
expect(auth.refreshes).toBe(2);
|
||||
expect(auth.replays).toBe(0);
|
||||
// Each renewal had its own connection: the reset one wasn't used again.
|
||||
expect(tokenProxy.dropped).toBeGreaterThanOrEqual(1);
|
||||
expect(tokenProxy.connections).toBe(2);
|
||||
}, 40_000);
|
||||
|
||||
// A login that can't be saved works only where it was made, so fx says so.
|
||||
test("a login that can't be saved says so, and fx mcp login fails", async () => {
|
||||
upstream = startModernMcpHttpFixture("json");
|
||||
|
||||
@@ -10,6 +10,7 @@ import {
|
||||
import { tmpdir } from "node:os";
|
||||
import { join } from "node:path";
|
||||
import { runFx } from "../evals/eval-helpers";
|
||||
import { startIdleDropProxy, type IdleDrop } from "./fixtures/idle-drop-proxy";
|
||||
import {
|
||||
startContentLengthMcpHttpFixture,
|
||||
type ContentLengthResponseType,
|
||||
@@ -39,21 +40,25 @@ let contentLengthFixture: Awaited<
|
||||
> | null = null;
|
||||
let gateway: ReturnType<typeof startFakeGateway> | null = null;
|
||||
let tui: TmuxSession | null = null;
|
||||
let proxy: Awaited<ReturnType<typeof startIdleDropProxy>> | null = null;
|
||||
|
||||
afterEach(async () => {
|
||||
const activeTui = tui;
|
||||
const activeGateway = gateway;
|
||||
const activeFixture = fixture;
|
||||
const activeContentLengthFixture = contentLengthFixture;
|
||||
const activeProxy = proxy;
|
||||
const activeCleanupRoot = cleanupRoot;
|
||||
tui = null;
|
||||
gateway = null;
|
||||
fixture = null;
|
||||
contentLengthFixture = null;
|
||||
proxy = null;
|
||||
cleanupRoot = null;
|
||||
|
||||
if (activeTui) await activeTui.kill();
|
||||
activeGateway?.stop();
|
||||
activeProxy?.stop();
|
||||
activeFixture?.stop();
|
||||
if (activeContentLengthFixture) await activeContentLengthFixture.stop();
|
||||
if (activeCleanupRoot) {
|
||||
@@ -449,6 +454,50 @@ describe("modern MCP Streamable HTTP", () => {
|
||||
}, 30_000);
|
||||
}
|
||||
|
||||
// A load balancer drops a keep-alive connection once it sits idle, by
|
||||
// resetting it or by closing it. fx keeps connections for the next request,
|
||||
// so each later call must find its connection gone and go out on a new one,
|
||||
// once, instead of failing until fx restarts.
|
||||
for (const drop of ["reset", "close"] as const satisfies readonly IdleDrop[]) {
|
||||
test(`tool calls after the server ${drop === "reset" ? "reset" : "closed"} an idle connection go through, each once`, async () => {
|
||||
contentLengthFixture = await startContentLengthMcpHttpFixture("json");
|
||||
proxy = await startIdleDropProxy(contentLengthFixture.url, drop);
|
||||
const root = createRoot(`idle-${drop}`, proxy);
|
||||
// Each call comes after the proxy dropped every connection fx kept.
|
||||
const later = (id: string, text: string) => async () => {
|
||||
await Bun.sleep(600);
|
||||
return fakeGatewayToolCall(id, TOOL_NAME, { text });
|
||||
};
|
||||
gateway = startFakeGateway([
|
||||
fakeGatewayToolCall("select_mcp", "mcp_select_tool", { name: TOOL_NAME }),
|
||||
later("call_one", "one"),
|
||||
later("call_two", "two"),
|
||||
later("call_three", "three"),
|
||||
fakeGatewayFinalText("Idle connections replaced."),
|
||||
], {
|
||||
models: [{ id: MODEL, type: "language", tags: ["tool-use"] }],
|
||||
});
|
||||
|
||||
const result = await runFx(
|
||||
["ask", "--json", "--auto", "--no-save", "Call the echo tool three times."],
|
||||
{ cwd: root.workspace, env: fixtureEnv(root, gateway), timeoutMs: 30_000 },
|
||||
);
|
||||
|
||||
expect(result.code).toBe(0);
|
||||
const last = gateway.requests.at(-1)?.body ?? "";
|
||||
for (const [id, text] of [["call_one", "one"], ["call_two", "two"], ["call_three", "three"]]) {
|
||||
expect(toolResultText(last, id)).toContain(`${MODERN_HTTP_TOOL_RESULT}:${text}`);
|
||||
}
|
||||
// None was sent twice.
|
||||
expect(
|
||||
contentLengthFixture.requests
|
||||
.filter((entry) => entry.message.method === "tools/call")
|
||||
.map((entry) => entry.message.params?.arguments?.text),
|
||||
).toEqual(["one", "two", "three"]);
|
||||
expect(proxy.dropped).toBeGreaterThanOrEqual(3);
|
||||
}, 40_000);
|
||||
}
|
||||
|
||||
test("live Tools subscription coalesces a storm and atomically exposes the changed tool", async () => {
|
||||
fixture = startModernMcpHttpFixture("cache_subscription");
|
||||
const root = createRoot("cache-subscription", fixture);
|
||||
|
||||
Reference in new issue
Block a user