Echo-AI-official's picture
Upload 280 files
0e759d2 verified
import { Logger } from "winston";
import { z, ZodError } from "zod";
import * as Sentry from "@sentry/node";
import { MockState, saveMock } from "./mock";
import { TimeoutSignal } from "../../../controllers/v1/types";
import { fireEngineURL } from "../engines/fire-engine/scrape";
import { fetch, RequestInit, Response, FormData, Agent } from "undici";
export type RobustFetchParams<Schema extends z.Schema<any>> = {
url: string;
logger: Logger;
method: "GET" | "POST" | "DELETE" | "PUT";
body?: any;
headers?: Record<string, string>;
schema?: Schema;
dontParseResponse?: boolean;
ignoreResponse?: boolean;
ignoreFailure?: boolean;
requestId?: string;
tryCount?: number;
tryCooldown?: number;
mock: MockState | null;
abort?: AbortSignal;
};
export async function robustFetch<
Schema extends z.Schema<any>,
Output = z.infer<Schema>,
>({
url,
logger,
method = "GET",
body,
headers,
schema,
ignoreResponse = false,
ignoreFailure = false,
requestId = crypto.randomUUID(),
tryCount = 1,
tryCooldown,
mock,
abort,
}: RobustFetchParams<Schema>): Promise<Output> {
abort?.throwIfAborted();
const params = {
url,
logger,
method,
body,
headers,
schema,
ignoreResponse,
ignoreFailure,
tryCount,
tryCooldown,
abort,
};
let response: {
status: number;
headers: Headers;
body: string;
};
if (mock === null) {
let request: Response;
try {
request = await fetch(url, {
method,
headers: {
...(body instanceof FormData
? {}
: body !== undefined
? {
"Content-Type": "application/json",
}
: {}),
...(headers !== undefined ? headers : {}),
},
signal: abort,
dispatcher: new Agent({
headersTimeout: 0,
bodyTimeout: 0,
}),
...(body instanceof FormData
? {
body,
}
: body !== undefined
? {
body: JSON.stringify(body),
}
: {}),
});
} catch (error) {
if (error instanceof TimeoutSignal) {
throw error;
} else if (!ignoreFailure) {
Sentry.captureException(error);
if (tryCount > 1) {
logger.debug(
"Request failed, trying " + (tryCount - 1) + " more times",
{ params, error, requestId },
);
return await robustFetch({
...params,
requestId,
tryCount: tryCount - 1,
mock,
});
} else {
logger.debug("Request failed", { params, error, requestId });
throw new Error("Request failed", {
cause: {
params,
requestId,
error,
},
});
}
} else {
return null as Output;
}
}
if (ignoreResponse === true) {
return null as Output;
}
const resp = await request.text();
response = {
status: request.status,
headers: request.headers,
body: resp, // NOTE: can this throw an exception?
};
} else {
if (ignoreResponse === true) {
return null as Output;
}
const makeRequestTypeId = (
request: (typeof mock)["requests"][number]["options"],
) => {
let trueUrl = request.url.startsWith(fireEngineURL)
? request.url.replace(fireEngineURL, "<fire-engine>")
: request.url;
let out = trueUrl + ";" + request.method;
if (
trueUrl.startsWith("<fire-engine>") &&
request.method === "POST"
) {
out += "f-e;" + request.body?.engine + ";" + request.body?.url;
}
return out;
};
const thisId = makeRequestTypeId(params);
const matchingMocks = mock.requests
.filter((x) => makeRequestTypeId(x.options) === thisId)
.sort((a, b) => a.time - b.time);
const nextI = mock.tracker[thisId] ?? 0;
mock.tracker[thisId] = nextI + 1;
if (!matchingMocks[nextI]) {
throw new Error("Failed to mock request -- no mock targets found.");
}
response = {
...matchingMocks[nextI].result,
headers: new Headers(matchingMocks[nextI].result.headers),
};
}
if (response.status >= 300) {
if (tryCount > 1) {
logger.debug(
"Request sent failure status, trying " + (tryCount - 1) + " more times",
{ params: { ...params, logger: undefined }, response: { status: response.status, body: response.body }, requestId },
);
if (tryCooldown !== undefined) {
await new Promise((resolve) =>
setTimeout(() => resolve(null), tryCooldown),
);
}
return await robustFetch({
...params,
requestId,
tryCount: tryCount - 1,
mock,
});
} else {
logger.debug("Request sent failure status", {
params: { ...params, logger: undefined },
response: { status: response.status, body: response.body },
requestId,
});
throw new Error("Request sent failure status", {
cause: {
params: { ...params, logger: undefined },
response: { status: response.status, body: response.body },
requestId,
},
});
}
}
if (mock === null) {
await saveMock(
{
...params,
logger: undefined,
schema: undefined,
headers: undefined,
},
response,
);
}
let data: Output;
try {
data = JSON.parse(response.body);
} catch (error) {
logger.debug("Request sent malformed JSON", {
params: { ...params, logger: undefined },
response: { status: response.status, body: response.body },
requestId,
});
throw new Error("Request sent malformed JSON", {
cause: {
params: { ...params, logger: undefined },
response,
requestId,
},
});
}
if (schema) {
try {
data = schema.parse(data);
} catch (error) {
if (error instanceof ZodError) {
logger.debug("Response does not match provided schema", {
params: { ...params, logger: undefined },
response: { status: response.status, body: response.body },
requestId,
error,
schema,
});
throw new Error("Response does not match provided schema", {
cause: {
params: { ...params, logger: undefined },
response,
requestId,
error,
schema,
},
});
} else {
logger.debug("Parsing response with provided schema failed", {
params: { ...params, logger: undefined },
response: { status: response.status, body: response.body },
requestId,
error,
schema,
});
throw new Error("Parsing response with provided schema failed", {
cause: {
params: { ...params, logger: undefined },
response,
requestId,
error,
schema,
},
});
}
}
}
return data;
}