chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,128 @@
|
||||
import { getAdminClickhouse } from "~/services/clickhouse/clickhouseFactory.server";
|
||||
import { llmPricingRegistry } from "~/v3/llmPricingRegistry.server";
|
||||
|
||||
export type MissingLlmModel = {
|
||||
model: string;
|
||||
system: string;
|
||||
count: number;
|
||||
};
|
||||
|
||||
export async function getMissingLlmModels(
|
||||
opts: {
|
||||
lookbackHours?: number;
|
||||
} = {}
|
||||
): Promise<MissingLlmModel[]> {
|
||||
const lookbackHours = opts.lookbackHours ?? 24;
|
||||
const since = new Date(Date.now() - lookbackHours * 60 * 60 * 1000);
|
||||
|
||||
const adminClickhouse = getAdminClickhouse();
|
||||
|
||||
// queryBuilderFast returns a factory function — call it to get the builder
|
||||
const createBuilder = adminClickhouse.reader.queryBuilderFast<{
|
||||
model: string;
|
||||
system: string;
|
||||
cnt: string;
|
||||
}>({
|
||||
name: "missingLlmModels",
|
||||
table: "trigger_dev.task_events_v2",
|
||||
columns: [
|
||||
{ name: "model", expression: "attributes.gen_ai.response.model.:String" },
|
||||
{ name: "system", expression: "attributes.gen_ai.system.:String" },
|
||||
{ name: "cnt", expression: "count()" },
|
||||
],
|
||||
});
|
||||
const qb = createBuilder();
|
||||
|
||||
// Partition pruning on inserted_at (partition key is toDate(inserted_at))
|
||||
qb.where("inserted_at >= {since: DateTime64(3)}", {
|
||||
since: formatDateTime(since),
|
||||
});
|
||||
|
||||
// Only spans that have a model set
|
||||
qb.where("attributes.gen_ai.response.model.:String != {empty: String}", { empty: "" });
|
||||
|
||||
// Only spans that were NOT cost-enriched (trigger.llm.total_cost is NULL)
|
||||
qb.where("attributes.trigger.llm.total_cost.:Float64 IS NULL", {});
|
||||
|
||||
// Only completed spans
|
||||
qb.where("kind = {kind: String}", { kind: "SPAN" });
|
||||
qb.where("status = {status: String}", { status: "OK" });
|
||||
|
||||
qb.groupBy("model, system");
|
||||
qb.orderBy("cnt DESC");
|
||||
qb.limit(100);
|
||||
|
||||
const [err, rows] = await qb.execute();
|
||||
|
||||
if (err) {
|
||||
throw err;
|
||||
}
|
||||
|
||||
if (!rows) {
|
||||
return [];
|
||||
}
|
||||
|
||||
const candidates = rows
|
||||
.filter((r) => r.model)
|
||||
.map((r) => ({
|
||||
model: r.model,
|
||||
system: r.system,
|
||||
count: parseInt(r.cnt, 10),
|
||||
}));
|
||||
|
||||
if (candidates.length === 0) return [];
|
||||
|
||||
// Filter out models that now have pricing in the database (added after spans were inserted).
|
||||
// The registry's match() handles prefix stripping for gateway/openrouter models.
|
||||
if (!llmPricingRegistry || !llmPricingRegistry.isLoaded) return candidates;
|
||||
const registry = llmPricingRegistry;
|
||||
return candidates.filter((c) => !registry.match(c.model));
|
||||
}
|
||||
|
||||
export type MissingModelSample = {
|
||||
span_id: string;
|
||||
run_id: string;
|
||||
message: string;
|
||||
attributes_text: string;
|
||||
duration: string;
|
||||
start_time: string;
|
||||
};
|
||||
|
||||
export async function getMissingModelSamples(opts: {
|
||||
model: string;
|
||||
lookbackHours?: number;
|
||||
limit?: number;
|
||||
}): Promise<MissingModelSample[]> {
|
||||
const lookbackHours = opts.lookbackHours ?? 24;
|
||||
const limit = opts.limit ?? 10;
|
||||
const since = new Date(Date.now() - lookbackHours * 60 * 60 * 1000);
|
||||
|
||||
const adminClickhouse = getAdminClickhouse();
|
||||
|
||||
const createBuilder = adminClickhouse.reader.queryBuilderFast<MissingModelSample>({
|
||||
name: "missingModelSamples",
|
||||
table: "trigger_dev.task_events_v2",
|
||||
columns: ["span_id", "run_id", "message", "attributes_text", "duration", "start_time"],
|
||||
});
|
||||
const qb = createBuilder();
|
||||
|
||||
qb.where("inserted_at >= {since: DateTime64(3)}", { since: formatDateTime(since) });
|
||||
qb.where("attributes.gen_ai.response.model.:String = {model: String}", { model: opts.model });
|
||||
qb.where("attributes.trigger.llm.total_cost.:Float64 IS NULL", {});
|
||||
qb.where("kind = {kind: String}", { kind: "SPAN" });
|
||||
qb.where("status = {status: String}", { status: "OK" });
|
||||
qb.orderBy("start_time DESC");
|
||||
qb.limit(limit);
|
||||
|
||||
const [err, rows] = await qb.execute();
|
||||
|
||||
if (err) {
|
||||
throw err;
|
||||
}
|
||||
|
||||
return rows ?? [];
|
||||
}
|
||||
|
||||
function formatDateTime(date: Date): string {
|
||||
return date.toISOString().replace("T", " ").replace("Z", "");
|
||||
}
|
||||
@@ -0,0 +1,794 @@
|
||||
import { json } from "@remix-run/server-runtime";
|
||||
import { SignJWT, errors, jwtVerify } from "jose";
|
||||
import { z } from "zod";
|
||||
|
||||
import { $replica } from "~/db.server";
|
||||
import { env } from "~/env.server";
|
||||
import { findProjectByRef } from "~/models/project.server";
|
||||
import {
|
||||
authIncludeBase,
|
||||
authIncludeWithParent,
|
||||
findEnvironmentByApiKey,
|
||||
findEnvironmentByPublicApiKey,
|
||||
toAuthenticated,
|
||||
} from "~/models/runtimeEnvironment.server";
|
||||
import { type RuntimeEnvironmentForEnvRepo } from "~/v3/environmentVariables/environmentVariablesRepository.server";
|
||||
import { logger } from "./logger.server";
|
||||
import { safeEnvironmentLogFields } from "./safeEnvironmentLog";
|
||||
import { missingJwtLogContext } from "./safeRequestLogContext";
|
||||
import {
|
||||
type PersonalAccessTokenAuthenticationResult,
|
||||
authenticateApiRequestWithPersonalAccessToken,
|
||||
isPersonalAccessToken,
|
||||
} from "./personalAccessToken.server";
|
||||
import {
|
||||
type OrganizationAccessTokenAuthenticationResult,
|
||||
authenticateApiRequestWithOrganizationAccessToken,
|
||||
isOrganizationAccessToken,
|
||||
} from "./organizationAccessToken.server";
|
||||
import { isPublicJWT, validatePublicJwtKey } from "./realtime/jwtAuth.server";
|
||||
import { isDefaultDevBranch, sanitizeBranchName } from "@trigger.dev/core/v3/utils/gitBranch";
|
||||
|
||||
const ClaimsSchema = z.object({
|
||||
scopes: z.array(z.string()).optional(),
|
||||
// One-time use token
|
||||
otu: z.boolean().optional(),
|
||||
realtime: z
|
||||
.object({
|
||||
skipColumns: z.array(z.string()).optional(),
|
||||
})
|
||||
.optional(),
|
||||
});
|
||||
|
||||
// Re-export the slim shape defined in @trigger.dev/core. Single source of
|
||||
// truth across the auth boundary (RBAC plugin contract → webapp handlers).
|
||||
export type { AuthenticatedEnvironment } from "@trigger.dev/core/v3/auth/environment";
|
||||
import type { AuthenticatedEnvironment } from "@trigger.dev/core/v3/auth/environment";
|
||||
|
||||
export type ApiAuthenticationResult =
|
||||
| ApiAuthenticationResultSuccess
|
||||
| ApiAuthenticationResultFailure;
|
||||
|
||||
export type ApiAuthenticationResultSuccess = {
|
||||
ok: true;
|
||||
apiKey: string;
|
||||
type: "PUBLIC" | "PRIVATE" | "PUBLIC_JWT";
|
||||
environment: AuthenticatedEnvironment;
|
||||
oneTimeUse?: boolean;
|
||||
realtime?: {
|
||||
skipColumns?: string[];
|
||||
};
|
||||
// Present when the request used a public JWT minted from a PAT/UAT exchange
|
||||
// that stamped an `act` delegation claim. `actor.sub` is the acting user id,
|
||||
// used for attribution (e.g. who resolved an error). Absent for plain env
|
||||
// API keys (no user) and JWTs minted without delegation.
|
||||
actor?: {
|
||||
sub: string;
|
||||
};
|
||||
};
|
||||
|
||||
export type ApiAuthenticationResultFailure = {
|
||||
ok: false;
|
||||
error: string;
|
||||
};
|
||||
|
||||
/**
|
||||
* @deprecated Use `authenticateApiRequestWithFailure` instead.
|
||||
*/
|
||||
export async function authenticateApiRequest(
|
||||
request: Request,
|
||||
options: { allowPublicKey?: boolean; allowJWT?: boolean } = {}
|
||||
): Promise<ApiAuthenticationResultSuccess | undefined> {
|
||||
const { apiKey, branchName } = getApiKeyFromRequest(request);
|
||||
|
||||
if (!apiKey) {
|
||||
return;
|
||||
}
|
||||
|
||||
const authentication = await authenticateApiKey(apiKey, { ...options, branchName });
|
||||
|
||||
return authentication;
|
||||
}
|
||||
|
||||
/**
|
||||
* This method is the same as `authenticateApiRequest` but it returns a failure result instead of undefined.
|
||||
* It should be used from now on to ensure that the API key is always validated and provide a failure result.
|
||||
*/
|
||||
export async function authenticateApiRequestWithFailure(
|
||||
request: Request,
|
||||
options: { allowPublicKey?: boolean; allowJWT?: boolean } = {}
|
||||
): Promise<ApiAuthenticationResult> {
|
||||
const { apiKey, branchName } = getApiKeyFromRequest(request);
|
||||
|
||||
if (!apiKey) {
|
||||
return {
|
||||
ok: false,
|
||||
error: "Invalid API Key",
|
||||
};
|
||||
}
|
||||
|
||||
const authentication = await authenticateApiKeyWithFailure(apiKey, { ...options, branchName });
|
||||
|
||||
return authentication;
|
||||
}
|
||||
|
||||
/**
|
||||
* @deprecated Use `authenticateApiKeyWithFailure` instead.
|
||||
*/
|
||||
export async function authenticateApiKey(
|
||||
apiKey: string,
|
||||
options: { allowPublicKey?: boolean; allowJWT?: boolean; branchName?: string } = {}
|
||||
): Promise<ApiAuthenticationResultSuccess | undefined> {
|
||||
const result = getApiKeyResult(apiKey);
|
||||
|
||||
if (!result) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (!options.allowPublicKey && result.type === "PUBLIC") {
|
||||
return;
|
||||
}
|
||||
|
||||
if (!options.allowJWT && result.type === "PUBLIC_JWT") {
|
||||
return;
|
||||
}
|
||||
|
||||
switch (result.type) {
|
||||
case "PUBLIC": {
|
||||
const environment = await findEnvironmentByPublicApiKey(result.apiKey, options.branchName);
|
||||
if (!environment) {
|
||||
return;
|
||||
}
|
||||
|
||||
return {
|
||||
ok: true,
|
||||
...result,
|
||||
environment,
|
||||
};
|
||||
}
|
||||
case "PRIVATE": {
|
||||
const environment = await findEnvironmentByApiKey(result.apiKey, options.branchName);
|
||||
if (!environment) {
|
||||
return;
|
||||
}
|
||||
|
||||
return {
|
||||
ok: true,
|
||||
...result,
|
||||
environment,
|
||||
};
|
||||
}
|
||||
case "PUBLIC_JWT": {
|
||||
const validationResults = await validatePublicJwtKey(result.apiKey);
|
||||
|
||||
if (!validationResults.ok) {
|
||||
return;
|
||||
}
|
||||
|
||||
const parsedClaims = ClaimsSchema.safeParse(validationResults.claims);
|
||||
|
||||
return {
|
||||
ok: true,
|
||||
...result,
|
||||
environment: validationResults.environment,
|
||||
oneTimeUse: parsedClaims.success ? parsedClaims.data.otu : false,
|
||||
realtime: parsedClaims.success ? parsedClaims.data.realtime : undefined,
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* This method is the same as `authenticateApiKey` but it returns a failure result instead of undefined.
|
||||
* It should be used from now on to ensure that the API key is always validated and provide a failure result.
|
||||
*/
|
||||
async function authenticateApiKeyWithFailure(
|
||||
apiKey: string,
|
||||
options: { allowPublicKey?: boolean; allowJWT?: boolean; branchName?: string } = {}
|
||||
): Promise<ApiAuthenticationResult> {
|
||||
const result = getApiKeyResult(apiKey);
|
||||
|
||||
if (!result) {
|
||||
return {
|
||||
ok: false,
|
||||
error: "Invalid API Key",
|
||||
};
|
||||
}
|
||||
|
||||
if (!options.allowPublicKey && result.type === "PUBLIC") {
|
||||
return {
|
||||
ok: false,
|
||||
error: "Public API keys are not allowed for this request",
|
||||
};
|
||||
}
|
||||
|
||||
if (!options.allowJWT && result.type === "PUBLIC_JWT") {
|
||||
return {
|
||||
ok: false,
|
||||
error: "Public JWT API keys are not allowed for this request",
|
||||
};
|
||||
}
|
||||
|
||||
switch (result.type) {
|
||||
case "PUBLIC": {
|
||||
const environment = await findEnvironmentByPublicApiKey(result.apiKey, options.branchName);
|
||||
if (!environment) {
|
||||
return {
|
||||
ok: false,
|
||||
error: "Invalid API Key",
|
||||
};
|
||||
}
|
||||
|
||||
return {
|
||||
ok: true,
|
||||
...result,
|
||||
environment,
|
||||
};
|
||||
}
|
||||
case "PRIVATE": {
|
||||
const environment = await findEnvironmentByApiKey(result.apiKey, options.branchName);
|
||||
if (!environment) {
|
||||
return {
|
||||
ok: false,
|
||||
error: "Invalid API Key",
|
||||
};
|
||||
}
|
||||
|
||||
return {
|
||||
ok: true,
|
||||
...result,
|
||||
environment,
|
||||
};
|
||||
}
|
||||
case "PUBLIC_JWT": {
|
||||
const validationResults = await validatePublicJwtKey(result.apiKey);
|
||||
|
||||
if (!validationResults.ok) {
|
||||
return validationResults;
|
||||
}
|
||||
|
||||
const parsedClaims = ClaimsSchema.safeParse(validationResults.claims);
|
||||
|
||||
return {
|
||||
ok: true,
|
||||
...result,
|
||||
environment: validationResults.environment,
|
||||
oneTimeUse: parsedClaims.success ? parsedClaims.data.otu : false,
|
||||
realtime: parsedClaims.success ? parsedClaims.data.realtime : undefined,
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
export async function authenticateAuthorizationHeader(
|
||||
authorization: string,
|
||||
{
|
||||
allowPublicKey = false,
|
||||
allowJWT = false,
|
||||
}: { allowPublicKey?: boolean; allowJWT?: boolean } = {}
|
||||
): Promise<ApiAuthenticationResult | undefined> {
|
||||
const apiKey = getApiKeyFromHeader(authorization);
|
||||
|
||||
if (!apiKey) {
|
||||
return;
|
||||
}
|
||||
|
||||
return authenticateApiKey(apiKey, { allowPublicKey, allowJWT });
|
||||
}
|
||||
|
||||
function isPublicApiKey(key: string) {
|
||||
return key.startsWith("pk_");
|
||||
}
|
||||
|
||||
function isSecretApiKey(key: string) {
|
||||
return key.startsWith("tr_");
|
||||
}
|
||||
|
||||
/**
|
||||
* Reads the branch off the `x-trigger-branch` header and sanitizes it.
|
||||
* Every server-side reader should go through here so sanitization is applied uniformly.
|
||||
* The dev `"default"` sentinel is intentionally NOT resolved here:
|
||||
* that translation is environment type-dependent.
|
||||
*/
|
||||
export function branchNameFromRequest(request: Request): string | undefined {
|
||||
return sanitizeBranchName(request.headers.get("x-trigger-branch")) ?? undefined;
|
||||
}
|
||||
|
||||
function getApiKeyFromRequest(request: Request): {
|
||||
apiKey: string | undefined;
|
||||
branchName: string | undefined;
|
||||
} {
|
||||
const apiKey = getApiKeyFromHeader(request.headers.get("Authorization"));
|
||||
const branchName = branchNameFromRequest(request);
|
||||
|
||||
return { apiKey, branchName };
|
||||
}
|
||||
|
||||
function getApiKeyFromHeader(authorization?: string | null) {
|
||||
if (typeof authorization !== "string" || !authorization) {
|
||||
return;
|
||||
}
|
||||
|
||||
const apiKey = authorization.replace(/^Bearer /, "");
|
||||
return apiKey;
|
||||
}
|
||||
|
||||
function getApiKeyResult(apiKey: string): {
|
||||
apiKey: string;
|
||||
type: "PUBLIC" | "PRIVATE" | "PUBLIC_JWT";
|
||||
} {
|
||||
const type = isPublicApiKey(apiKey)
|
||||
? "PUBLIC"
|
||||
: isSecretApiKey(apiKey)
|
||||
? "PRIVATE"
|
||||
: isPublicJWT(apiKey)
|
||||
? "PUBLIC_JWT"
|
||||
: "PRIVATE"; // Fallback to private key
|
||||
return { apiKey, type };
|
||||
}
|
||||
|
||||
export type AuthenticationResult =
|
||||
| {
|
||||
type: "personalAccessToken";
|
||||
result: PersonalAccessTokenAuthenticationResult;
|
||||
}
|
||||
| {
|
||||
type: "organizationAccessToken";
|
||||
result: OrganizationAccessTokenAuthenticationResult;
|
||||
}
|
||||
| {
|
||||
type: "apiKey";
|
||||
result: ApiAuthenticationResult;
|
||||
};
|
||||
|
||||
type AuthenticationMethod = "personalAccessToken" | "organizationAccessToken" | "apiKey";
|
||||
|
||||
type AllowedAuthenticationMethods = Record<AuthenticationMethod, boolean> &
|
||||
({ personalAccessToken: true } | { organizationAccessToken: true } | { apiKey: true });
|
||||
|
||||
const defaultAllowedAuthenticationMethods: AllowedAuthenticationMethods = {
|
||||
personalAccessToken: true,
|
||||
organizationAccessToken: true,
|
||||
apiKey: true,
|
||||
};
|
||||
|
||||
type FilteredAuthenticationResult<
|
||||
T extends AllowedAuthenticationMethods = AllowedAuthenticationMethods,
|
||||
> =
|
||||
| (T["personalAccessToken"] extends true
|
||||
? Extract<AuthenticationResult, { type: "personalAccessToken" }>
|
||||
: never)
|
||||
| (T["organizationAccessToken"] extends true
|
||||
? Extract<AuthenticationResult, { type: "organizationAccessToken" }>
|
||||
: never)
|
||||
| (T["apiKey"] extends true ? Extract<AuthenticationResult, { type: "apiKey" }> : never);
|
||||
|
||||
/**
|
||||
* Authenticates an incoming request by checking for various token types.
|
||||
*
|
||||
* Supports personal access tokens, organization access tokens, and API keys.
|
||||
* Returns the appropriate authentication result based on the token type found.
|
||||
*
|
||||
* This method currently only allows private keys for the `apiKey` authentication method.
|
||||
*
|
||||
* @template T - The allowed authentication methods configuration type
|
||||
* @param request - The incoming HTTP request containing authentication headers
|
||||
* @param allowedAuthenticationMethods - Configuration object specifying which authentication methods are allowed.
|
||||
* At least one method must be set to `true`. Defaults to allowing all methods.
|
||||
* @returns Authentication result with only the enabled auth method types, or undefined if no valid token found
|
||||
*
|
||||
* @example
|
||||
* ```typescript
|
||||
* // Only allow personal access tokens
|
||||
* const result = await authenticateRequest(request, {
|
||||
* personalAccessToken: true,
|
||||
* organizationAccessToken: false,
|
||||
* apiKey: false,
|
||||
* });
|
||||
* // result type: { type: "personalAccessToken"; result: PersonalAccessTokenAuthenticationResult } | undefined
|
||||
* ```
|
||||
*/
|
||||
export async function authenticateRequest<
|
||||
T extends AllowedAuthenticationMethods = AllowedAuthenticationMethods,
|
||||
>(
|
||||
request: Request,
|
||||
allowedAuthenticationMethods?: T
|
||||
): Promise<FilteredAuthenticationResult<T> | undefined> {
|
||||
const allowedMethods = allowedAuthenticationMethods ?? defaultAllowedAuthenticationMethods;
|
||||
|
||||
const { apiKey, branchName } = getApiKeyFromRequest(request);
|
||||
if (!apiKey) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (allowedMethods.personalAccessToken && isPersonalAccessToken(apiKey)) {
|
||||
const result = await authenticateApiRequestWithPersonalAccessToken(request);
|
||||
|
||||
if (!result) {
|
||||
return;
|
||||
}
|
||||
|
||||
return {
|
||||
type: "personalAccessToken",
|
||||
result,
|
||||
} satisfies Extract<
|
||||
AuthenticationResult,
|
||||
{ type: "personalAccessToken" }
|
||||
> as FilteredAuthenticationResult<T>;
|
||||
}
|
||||
|
||||
if (allowedMethods.organizationAccessToken && isOrganizationAccessToken(apiKey)) {
|
||||
const result = await authenticateApiRequestWithOrganizationAccessToken(request);
|
||||
|
||||
if (!result) {
|
||||
return;
|
||||
}
|
||||
|
||||
return {
|
||||
type: "organizationAccessToken",
|
||||
result,
|
||||
} satisfies Extract<
|
||||
AuthenticationResult,
|
||||
{ type: "organizationAccessToken" }
|
||||
> as FilteredAuthenticationResult<T>;
|
||||
}
|
||||
|
||||
if (allowedMethods.apiKey) {
|
||||
const result = await authenticateApiKey(apiKey, { allowPublicKey: false, branchName });
|
||||
|
||||
if (!result) {
|
||||
return;
|
||||
}
|
||||
|
||||
return {
|
||||
type: "apiKey",
|
||||
result,
|
||||
} satisfies Extract<
|
||||
AuthenticationResult,
|
||||
{ type: "apiKey" }
|
||||
> as FilteredAuthenticationResult<T>;
|
||||
}
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
export async function authenticatedEnvironmentForAuthentication(
|
||||
auth: AuthenticationResult,
|
||||
projectRef: string,
|
||||
slug: string,
|
||||
branch?: string
|
||||
): Promise<AuthenticatedEnvironment> {
|
||||
if (slug === "staging") {
|
||||
slug = "stg";
|
||||
}
|
||||
|
||||
// Normalize the requested branch once: sanitize it, then collapse the dev
|
||||
// `"default"` sentinel to "no branch" so it resolves to the root dev env
|
||||
// rather than a (non-existent) branch literally named "default".
|
||||
// TODO this slug check is brittle
|
||||
const sanitizedBranch = sanitizeBranchName(branch);
|
||||
const resolvedBranch =
|
||||
slug === "dev" && isDefaultDevBranch(sanitizedBranch) ? null : sanitizedBranch;
|
||||
|
||||
switch (auth.type) {
|
||||
case "apiKey": {
|
||||
if (!auth.result.ok) {
|
||||
throw json({ error: auth.result.error }, { status: 401 });
|
||||
}
|
||||
|
||||
if (auth.result.environment.project.externalRef !== projectRef) {
|
||||
throw json(
|
||||
{
|
||||
error:
|
||||
"Invalid project ref for this API key. Make sure you are using an API key associated with that project.",
|
||||
},
|
||||
{ status: 400 }
|
||||
);
|
||||
}
|
||||
|
||||
if (
|
||||
auth.result.environment.slug !== slug &&
|
||||
auth.result.environment.branchName !== resolvedBranch
|
||||
) {
|
||||
throw json(
|
||||
{
|
||||
error:
|
||||
"Invalid environment slug for this API key. Make sure you are using an API key associated with that environment.",
|
||||
},
|
||||
{ status: 400 }
|
||||
);
|
||||
}
|
||||
|
||||
return auth.result.environment;
|
||||
}
|
||||
case "personalAccessToken": {
|
||||
const user = await $replica.user.findUnique({
|
||||
where: {
|
||||
id: auth.result.userId,
|
||||
},
|
||||
});
|
||||
|
||||
if (!user) {
|
||||
throw json({ error: "Invalid or missing personal access token" }, { status: 401 });
|
||||
}
|
||||
|
||||
const project = await findProjectByRef(projectRef, user.id);
|
||||
|
||||
if (!project) {
|
||||
throw json({ error: "Project not found" }, { status: 404 });
|
||||
}
|
||||
|
||||
if (!resolvedBranch) {
|
||||
const environment = await $replica.runtimeEnvironment.findFirst({
|
||||
where: {
|
||||
projectId: project.id,
|
||||
slug: slug,
|
||||
...(slug === "dev"
|
||||
? {
|
||||
orgMember: {
|
||||
userId: user.id,
|
||||
},
|
||||
}
|
||||
: {}),
|
||||
},
|
||||
include: authIncludeBase,
|
||||
});
|
||||
|
||||
if (!environment) {
|
||||
throw json({ error: "Environment not found" }, { status: 404 });
|
||||
}
|
||||
|
||||
return toAuthenticated(environment);
|
||||
}
|
||||
|
||||
const environment = await $replica.runtimeEnvironment.findFirst({
|
||||
where: {
|
||||
projectId: project.id,
|
||||
type: slug === "dev" ? "DEVELOPMENT" : "PREVIEW",
|
||||
branchName: resolvedBranch,
|
||||
...(slug === "dev"
|
||||
? {
|
||||
orgMember: {
|
||||
userId: user.id,
|
||||
},
|
||||
}
|
||||
: {}),
|
||||
archivedAt: null,
|
||||
},
|
||||
include: authIncludeWithParent,
|
||||
});
|
||||
|
||||
if (!environment) {
|
||||
throw json({ error: "Branch not found" }, { status: 404 });
|
||||
}
|
||||
|
||||
if (!environment.parentEnvironment) {
|
||||
throw json({ error: "Branch not associated with a parent environment" }, { status: 400 });
|
||||
}
|
||||
|
||||
// PREVIEW envs (and DEVELOPMENT branches) reuse the parent's apiKey for downstream auth flows
|
||||
// (signed JWTs, internal-fetch helpers). Override before mapping so
|
||||
// the slim shape carries the parent's key.
|
||||
return toAuthenticated({
|
||||
...environment,
|
||||
apiKey: environment.parentEnvironment.apiKey,
|
||||
});
|
||||
}
|
||||
case "organizationAccessToken": {
|
||||
const organization = await $replica.organization.findUnique({
|
||||
where: {
|
||||
id: auth.result.organizationId,
|
||||
},
|
||||
});
|
||||
|
||||
if (!organization) {
|
||||
throw json({ error: "Invalid or missing organization access token" }, { status: 401 });
|
||||
}
|
||||
|
||||
const project = await $replica.project.findFirst({
|
||||
where: {
|
||||
organizationId: organization.id,
|
||||
externalRef: projectRef,
|
||||
},
|
||||
});
|
||||
|
||||
if (!project) {
|
||||
throw json({ error: "Project not found" }, { status: 404 });
|
||||
}
|
||||
|
||||
if (!resolvedBranch) {
|
||||
const environment = await $replica.runtimeEnvironment.findFirst({
|
||||
where: {
|
||||
projectId: project.id,
|
||||
slug: slug,
|
||||
},
|
||||
include: authIncludeBase,
|
||||
});
|
||||
|
||||
if (!environment) {
|
||||
throw json({ error: "Environment not found" }, { status: 404 });
|
||||
}
|
||||
|
||||
return toAuthenticated(environment);
|
||||
}
|
||||
|
||||
const environment = await $replica.runtimeEnvironment.findFirst({
|
||||
where: {
|
||||
projectId: project.id,
|
||||
// No Development branches for OAT
|
||||
type: "PREVIEW",
|
||||
branchName: resolvedBranch,
|
||||
archivedAt: null,
|
||||
},
|
||||
include: authIncludeWithParent,
|
||||
});
|
||||
|
||||
if (!environment) {
|
||||
throw json({ error: "Branch not found" }, { status: 404 });
|
||||
}
|
||||
|
||||
if (!environment.parentEnvironment) {
|
||||
throw json({ error: "Branch not associated with a preview environment" }, { status: 400 });
|
||||
}
|
||||
|
||||
return toAuthenticated({
|
||||
...environment,
|
||||
apiKey: environment.parentEnvironment.apiKey,
|
||||
});
|
||||
}
|
||||
default: {
|
||||
auth satisfies never;
|
||||
throw json({ error: "Invalid authentication result" }, { status: 401 });
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const JWT_SECRET = new TextEncoder().encode(env.SESSION_SECRET);
|
||||
const JWT_ALGORITHM = "HS256";
|
||||
const DEFAULT_JWT_EXPIRATION_IN_MS = 1000 * 60 * 60; // 1 hour
|
||||
|
||||
export async function generateJWTTokenForEnvironment(
|
||||
environment: RuntimeEnvironmentForEnvRepo,
|
||||
payload: Record<string, string>
|
||||
) {
|
||||
const jwt = await new SignJWT({
|
||||
environment_id: environment.id,
|
||||
org_id: environment.organizationId,
|
||||
project_id: environment.projectId,
|
||||
...payload,
|
||||
})
|
||||
.setProtectedHeader({ alg: JWT_ALGORITHM })
|
||||
.setIssuedAt()
|
||||
.setIssuer("https://id.trigger.dev")
|
||||
.setAudience("https://api.trigger.dev")
|
||||
.setExpirationTime(calculateJWTExpiration())
|
||||
.sign(JWT_SECRET);
|
||||
|
||||
return jwt;
|
||||
}
|
||||
|
||||
export async function validateJWTTokenAndRenew<T extends z.ZodTypeAny>(
|
||||
request: Request,
|
||||
payloadSchema: T
|
||||
): Promise<{ payload: z.infer<T>; jwt: string } | undefined> {
|
||||
try {
|
||||
const jwt = request.headers.get("x-trigger-jwt");
|
||||
|
||||
if (!jwt) {
|
||||
// Log a safe breadcrumb, not the raw headers (which carry the
|
||||
// caller's Authorization credential).
|
||||
logger.debug("Missing JWT token in request", missingJwtLogContext(request));
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
const { payload: rawPayload } = await jwtVerify(jwt, JWT_SECRET, {
|
||||
issuer: "https://id.trigger.dev",
|
||||
audience: "https://api.trigger.dev",
|
||||
});
|
||||
|
||||
const payload = payloadSchema.safeParse(rawPayload);
|
||||
|
||||
if (!payload.success) {
|
||||
logger.error("Failed to validate JWT", { payload: rawPayload, issues: payload.error.issues });
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
const renewedJwt = await renewJWTToken(payload.data);
|
||||
|
||||
return {
|
||||
payload: payload.data,
|
||||
jwt: renewedJwt,
|
||||
};
|
||||
} catch (error) {
|
||||
if (error instanceof errors.JWTExpired) {
|
||||
// Now we need to try and renew the token using the API key auth
|
||||
const authenticatedEnv = await authenticateApiRequest(request);
|
||||
|
||||
if (!authenticatedEnv) {
|
||||
logger.error("Failed to renew JWT token, missing or invalid Authorization header", {
|
||||
error: error.message,
|
||||
});
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
if (!authenticatedEnv.ok) {
|
||||
logger.error("Failed to renew JWT token, invalid API key", {
|
||||
error: error.message,
|
||||
});
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
const payload = payloadSchema.safeParse(error.payload);
|
||||
|
||||
if (!payload.success) {
|
||||
logger.error("Failed to parse jwt payload after expired", {
|
||||
payload: error.payload,
|
||||
issues: payload.error.issues,
|
||||
});
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
const renewedJwt = await generateJWTTokenForEnvironment(authenticatedEnv.environment, {
|
||||
...payload.data,
|
||||
});
|
||||
|
||||
// The environment carries secret material; log only non-secret fields.
|
||||
logger.debug("Renewed JWT token from Authorization header API Key", {
|
||||
environment: safeEnvironmentLogFields(authenticatedEnv.environment),
|
||||
payload: payload.data,
|
||||
});
|
||||
|
||||
return {
|
||||
payload: payload.data,
|
||||
jwt: renewedJwt,
|
||||
};
|
||||
}
|
||||
|
||||
logger.error("Failed to validate JWT token", { error });
|
||||
}
|
||||
}
|
||||
|
||||
async function renewJWTToken(payload: Record<string, string>) {
|
||||
const jwt = await new SignJWT(payload)
|
||||
.setProtectedHeader({ alg: JWT_ALGORITHM })
|
||||
.setIssuedAt()
|
||||
.setIssuer("https://id.trigger.dev")
|
||||
.setAudience("https://api.trigger.dev")
|
||||
.setExpirationTime(calculateJWTExpiration())
|
||||
.sign(JWT_SECRET);
|
||||
|
||||
return jwt;
|
||||
}
|
||||
|
||||
function calculateJWTExpiration() {
|
||||
if (env.PROD_USAGE_HEARTBEAT_INTERVAL_MS) {
|
||||
return (
|
||||
(Date.now() + Math.max(DEFAULT_JWT_EXPIRATION_IN_MS, env.PROD_USAGE_HEARTBEAT_INTERVAL_MS)) /
|
||||
1000
|
||||
);
|
||||
}
|
||||
|
||||
return (Date.now() + DEFAULT_JWT_EXPIRATION_IN_MS) / 1000;
|
||||
}
|
||||
|
||||
export async function getOneTimeUseToken(
|
||||
auth: ApiAuthenticationResultSuccess
|
||||
): Promise<string | undefined> {
|
||||
if (auth.type !== "PUBLIC_JWT") {
|
||||
return;
|
||||
}
|
||||
|
||||
if (!auth.oneTimeUse) {
|
||||
return;
|
||||
}
|
||||
|
||||
// Hash the API key to make it unique
|
||||
const hash = await crypto.subtle.digest("SHA-256", new TextEncoder().encode(auth.apiKey));
|
||||
|
||||
return Buffer.from(hash).toString("hex");
|
||||
}
|
||||
@@ -0,0 +1,82 @@
|
||||
import { env } from "~/env.server";
|
||||
import { authenticateAuthorizationHeader } from "./apiAuth.server";
|
||||
import { authorizationRateLimitMiddleware } from "./authorizationRateLimitMiddleware.server";
|
||||
import type { Duration } from "./rateLimiter.server";
|
||||
|
||||
export const apiRateLimiter = authorizationRateLimitMiddleware({
|
||||
redis: {
|
||||
port: env.RATE_LIMIT_REDIS_PORT,
|
||||
host: env.RATE_LIMIT_REDIS_HOST,
|
||||
username: env.RATE_LIMIT_REDIS_USERNAME,
|
||||
password: env.RATE_LIMIT_REDIS_PASSWORD,
|
||||
tlsDisabled: env.RATE_LIMIT_REDIS_TLS_DISABLED === "true",
|
||||
clusterMode: env.RATE_LIMIT_REDIS_CLUSTER_MODE_ENABLED === "1",
|
||||
},
|
||||
keyPrefix: "api",
|
||||
defaultLimiter: {
|
||||
type: "tokenBucket",
|
||||
refillRate: env.API_RATE_LIMIT_REFILL_RATE,
|
||||
interval: env.API_RATE_LIMIT_REFILL_INTERVAL as Duration,
|
||||
maxTokens: env.API_RATE_LIMIT_MAX,
|
||||
},
|
||||
limiterCache: {
|
||||
fresh: 60_000 * 10, // Data is fresh for 10 minutes
|
||||
stale: 60_000 * 20, // Date is stale after 20 minutes
|
||||
maxItems: 1000,
|
||||
},
|
||||
limiterConfigOverride: async (authorizationValue) => {
|
||||
const authenticatedEnv = await authenticateAuthorizationHeader(authorizationValue, {
|
||||
allowPublicKey: true,
|
||||
allowJWT: true,
|
||||
});
|
||||
|
||||
if (!authenticatedEnv || !authenticatedEnv.ok) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (authenticatedEnv.type === "PUBLIC_JWT") {
|
||||
return {
|
||||
type: "fixedWindow",
|
||||
window: env.API_RATE_LIMIT_JWT_WINDOW,
|
||||
tokens: env.API_RATE_LIMIT_JWT_TOKENS,
|
||||
};
|
||||
} else {
|
||||
return authenticatedEnv.environment.organization.apiRateLimiterConfig;
|
||||
}
|
||||
},
|
||||
pathMatchers: [/^\/api/],
|
||||
// Allow /api/v1/tasks/:id/callback/:secret
|
||||
pathWhiteList: [
|
||||
"/api/internal/stripe_webhooks",
|
||||
"/api/v1/authorization-code",
|
||||
"/api/v1/token",
|
||||
"/api/v1/usage/ingest",
|
||||
"/api/v1/plain/customer-cards",
|
||||
/^\/api\/v1\/tasks\/[^/]+\/callback\/[^/]+$/, // /api/v1/tasks/$id/callback/$secret
|
||||
/^\/api\/v1\/runs\/[^/]+\/tasks\/[^/]+\/callback\/[^/]+$/, // /api/v1/runs/$runId/tasks/$id/callback/$secret
|
||||
/^\/api\/v1\/http-endpoints\/[^/]+\/env\/[^/]+\/[^/]+$/, // /api/v1/http-endpoints/$httpEndpointId/env/$envType/$shortcode
|
||||
/^\/api\/v1\/sources\/http\/[^/]+$/, // /api/v1/sources/http/$id
|
||||
/^\/api\/v1\/endpoints\/[^/]+\/[^/]+\/index\/[^/]+$/, // /api/v1/endpoints/$environmentId/$endpointSlug/index/$indexHookIdentifier
|
||||
"/api/v1/timezones",
|
||||
"/api/v1/usage/ingest",
|
||||
"/api/v1/auth/jwt/claims",
|
||||
/^\/api\/v1\/runs\/[^/]+\/attempts$/, // /api/v1/runs/$runFriendlyId/attempts
|
||||
/^\/api\/v1\/waitpoints\/tokens\/[^/]+\/callback\/[^/]+$/, // /api/v1/waitpoints/tokens/$waitpointFriendlyId/callback/$hash
|
||||
/^\/api\/v\d+\/deployments/, // /api/v{1,2,3,n}/deployments/*
|
||||
// Internal SDK plumbing — packets are presigned-URL handshakes for
|
||||
// payload uploads (v2 PUT) and downloads (v1 GET), authenticated via
|
||||
// run-scoped JWT, called once per task/turn boundary by the runtime.
|
||||
// Same shape as `/api/v1/runs/$runFriendlyId/attempts` above; not a
|
||||
// customer-facing surface so customer rate limits shouldn't apply.
|
||||
/^\/api\/v1\/packets\//,
|
||||
/^\/api\/v2\/packets\//,
|
||||
/^\/api\/v1\/sessions\/[^/]+\/snapshot-url$/,
|
||||
],
|
||||
log: {
|
||||
rejections: env.API_RATE_LIMIT_REJECTION_LOGS_ENABLED === "1",
|
||||
requests: env.API_RATE_LIMIT_REQUEST_LOGS_ENABLED === "1",
|
||||
limiter: env.API_RATE_LIMIT_LIMITER_LOGS_ENABLED === "1",
|
||||
},
|
||||
});
|
||||
|
||||
export type RateLimitMiddleware = ReturnType<typeof authorizationRateLimitMiddleware>;
|
||||
@@ -0,0 +1,112 @@
|
||||
import { type PrismaClient } from "@trigger.dev/database";
|
||||
import { prisma } from "~/db.server";
|
||||
import { logger } from "./logger.server";
|
||||
import { nanoid } from "nanoid";
|
||||
import { controlPlaneResolver } from "~/v3/runOpsMigration/controlPlaneResolver.server";
|
||||
|
||||
export class ArchiveBranchService {
|
||||
#prismaClient: PrismaClient;
|
||||
|
||||
constructor(prismaClient: PrismaClient = prisma) {
|
||||
this.#prismaClient = prismaClient;
|
||||
}
|
||||
|
||||
public async call(
|
||||
// The orgFilter approach is not ideal but we need to keep it this way for now because of how the service is used in routes and api endpoints.
|
||||
// Currently authorization checks are spread across the controller/route layer and the service layer. Often we check in multiple places for org/project membership.
|
||||
// Ideally we would take care of both the authentication and authorization checks in the controllers and routes.
|
||||
// That would unify how we handle authorization and org/project membership checks. Also it would make the service layer queries simpler.
|
||||
orgFilter:
|
||||
| { type: "userMembership"; userId: string }
|
||||
| { type: "orgId"; organizationId: string },
|
||||
{
|
||||
environmentId,
|
||||
}: {
|
||||
environmentId: string;
|
||||
}
|
||||
) {
|
||||
try {
|
||||
const environment = await this.#prismaClient.runtimeEnvironment.findFirstOrThrow({
|
||||
where: {
|
||||
id: environmentId,
|
||||
organization:
|
||||
orgFilter.type === "userMembership"
|
||||
? {
|
||||
members: {
|
||||
some: {
|
||||
userId: orgFilter.userId,
|
||||
},
|
||||
},
|
||||
}
|
||||
: { id: orgFilter.organizationId },
|
||||
// Dev branches are per-org-member, so org membership alone isn't enough:
|
||||
// only the owner may archive their own dev branch. Non-dev branches (e.g.
|
||||
// preview) remain scoped by org membership only.
|
||||
...(orgFilter.type === "userMembership"
|
||||
? {
|
||||
OR: [
|
||||
{ type: { not: "DEVELOPMENT" as const } },
|
||||
{ orgMember: { userId: orgFilter.userId } },
|
||||
],
|
||||
}
|
||||
: {}),
|
||||
},
|
||||
include: {
|
||||
organization: {
|
||||
select: {
|
||||
id: true,
|
||||
slug: true,
|
||||
maximumConcurrencyLimit: true,
|
||||
},
|
||||
},
|
||||
project: {
|
||||
select: {
|
||||
id: true,
|
||||
slug: true,
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
// A branch is defined by having a parent; any root (dev/preview parent,
|
||||
// prod, staging) has none and can't be archived. For dev, that root is
|
||||
// the default branch, so give the clearer message.
|
||||
if (!environment.parentEnvironmentId) {
|
||||
return {
|
||||
success: false as const,
|
||||
error:
|
||||
environment.type === "DEVELOPMENT"
|
||||
? "The default development branch cannot be archived."
|
||||
: "This isn't a branch, and cannot be archived.",
|
||||
};
|
||||
}
|
||||
|
||||
// Branch archive is a SOFT update — do NOT hard-delete run-ops rows here (it would destroy a
|
||||
// retained branch's history). Any env hard-delete/purge belongs on a dedicated purge path
|
||||
// (owned by the cloud env-purge runbook), which has no site today.
|
||||
const slug = `${environment.slug}-${nanoid(6)}`;
|
||||
const shortcode = slug;
|
||||
|
||||
const updatedBranch = await this.#prismaClient.runtimeEnvironment.update({
|
||||
where: { id: environmentId },
|
||||
data: { archivedAt: new Date(), slug, shortcode },
|
||||
});
|
||||
|
||||
// archivedAt/slug/shortcode changed in the control-plane; drop any cached copy.
|
||||
controlPlaneResolver.invalidateEnvironment(environmentId);
|
||||
|
||||
return {
|
||||
success: true as const,
|
||||
branch: updatedBranch,
|
||||
organization: environment.organization,
|
||||
project: environment.project,
|
||||
};
|
||||
} catch (e) {
|
||||
logger.error("ArchiveBranchService error", { environmentId, error: e });
|
||||
return {
|
||||
success: false as const,
|
||||
error: "Failed to archive branch",
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,153 @@
|
||||
import { z } from "zod";
|
||||
import { prisma } from "~/db.server";
|
||||
import { env } from "~/env.server";
|
||||
import { logger } from "./logger.server";
|
||||
|
||||
// Syncs new orgs/users into Attio (workspaces/users objects) at signup, via the
|
||||
// common worker so a slow Attio never blocks signup. Ongoing field updates are
|
||||
// handled by the scheduled sync, not here. No-op without ATTIO_API_KEY.
|
||||
|
||||
const ATTIO_API = "https://api.attio.com/v2";
|
||||
const IS_TEST = env.APP_ENV !== "production";
|
||||
|
||||
export const AttioWorkspaceSyncSchema = z.object({
|
||||
orgId: z.string(),
|
||||
title: z.string(),
|
||||
slug: z.string(),
|
||||
companySize: z.string().nullish(),
|
||||
createdAt: z.coerce.date(),
|
||||
adminUserId: z.string(),
|
||||
});
|
||||
export type AttioWorkspaceSync = z.infer<typeof AttioWorkspaceSyncSchema>;
|
||||
|
||||
export const AttioUserSyncSchema = z.object({
|
||||
userId: z.string(),
|
||||
email: z.string(),
|
||||
referralSource: z.string().nullish(),
|
||||
marketingEmails: z.boolean(),
|
||||
createdAt: z.coerce.date(),
|
||||
});
|
||||
export type AttioUserSync = z.infer<typeof AttioUserSyncSchema>;
|
||||
|
||||
class AttioClient {
|
||||
constructor(private readonly apiKey: string) {}
|
||||
|
||||
// Create-or-update by unique attribute; returns the record id. Throws on failure so the worker retries.
|
||||
async #assert(
|
||||
object: string,
|
||||
matchingAttribute: string,
|
||||
values: Record<string, unknown>
|
||||
): Promise<string> {
|
||||
const url = `${ATTIO_API}/objects/${object}/records?matching_attribute=${matchingAttribute}`;
|
||||
const response = await fetch(url, {
|
||||
method: "PUT",
|
||||
headers: { Authorization: `Bearer ${this.apiKey}`, "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ data: { values } }),
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
const body = await response.text();
|
||||
logger.error("Attio assert failed", {
|
||||
object,
|
||||
matchingAttribute,
|
||||
status: response.status,
|
||||
body,
|
||||
});
|
||||
throw new Error(`Attio assert ${object} failed with status ${response.status}`);
|
||||
}
|
||||
|
||||
const recordId = ((await response.json()) as any).data?.id?.record_id;
|
||||
if (typeof recordId !== "string") {
|
||||
throw new Error(`Attio assert ${object}: response missing data.id.record_id`);
|
||||
}
|
||||
return recordId;
|
||||
}
|
||||
|
||||
async upsertWorkspace(payload: AttioWorkspaceSync, emailDomain?: string) {
|
||||
// The creating user is an admin of the new org — set their role and link them to the workspace.
|
||||
const adminRecordId = await this.#assert("users", "user_id", {
|
||||
user_id: payload.adminUserId,
|
||||
role: "Admin",
|
||||
is_test: IS_TEST,
|
||||
});
|
||||
|
||||
await this.#assert("workspaces", "workspace_id", {
|
||||
workspace_id: payload.orgId,
|
||||
name: payload.title,
|
||||
org_slug: payload.slug,
|
||||
company_size: payload.companySize ?? undefined,
|
||||
email_domain: emailDomain,
|
||||
signup_date: toDate(payload.createdAt),
|
||||
plan: "Free",
|
||||
account_status: "Active",
|
||||
is_test: IS_TEST,
|
||||
users: [{ target_object: "users", target_record_id: adminRecordId }],
|
||||
});
|
||||
}
|
||||
|
||||
async upsertUser(payload: AttioUserSync) {
|
||||
await this.#assert("users", "user_id", {
|
||||
user_id: payload.userId,
|
||||
primary_email_address: payload.email,
|
||||
marketing_opt_in: payload.marketingEmails,
|
||||
referral_source: payload.referralSource ?? undefined,
|
||||
signup_date: toDate(payload.createdAt),
|
||||
is_test: IS_TEST,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Attio `date` attributes want a bare YYYY-MM-DD value.
|
||||
function toDate(date: Date): string {
|
||||
return date.toISOString().slice(0, 10);
|
||||
}
|
||||
|
||||
// Domain from an email; the cloud-side matcher normalizes it further.
|
||||
function domainFromEmail(email: string | undefined): string | undefined {
|
||||
return email?.split("@")[1]?.toLowerCase().trim() || undefined;
|
||||
}
|
||||
|
||||
export const attioClient = env.ATTIO_API_KEY ? new AttioClient(env.ATTIO_API_KEY) : null;
|
||||
|
||||
export async function enqueueAttioWorkspaceSync(payload: AttioWorkspaceSync) {
|
||||
if (!attioClient) return;
|
||||
try {
|
||||
// Lazy import to avoid a circular dependency with commonWorker (which imports this module's schemas).
|
||||
const { commonWorker } = await import("~/v3/commonWorker.server");
|
||||
await commonWorker.enqueue({
|
||||
id: `attio:workspace:${payload.orgId}`,
|
||||
job: "attio.syncWorkspace",
|
||||
payload,
|
||||
});
|
||||
} catch (error) {
|
||||
logger.error("Failed to enqueue Attio workspace sync", { orgId: payload.orgId, error });
|
||||
}
|
||||
}
|
||||
|
||||
export async function enqueueAttioUserSync(payload: AttioUserSync) {
|
||||
if (!attioClient) return;
|
||||
try {
|
||||
const { commonWorker } = await import("~/v3/commonWorker.server");
|
||||
await commonWorker.enqueue({
|
||||
id: `attio:user:${payload.userId}`,
|
||||
job: "attio.syncUser",
|
||||
payload,
|
||||
});
|
||||
} catch (error) {
|
||||
logger.error("Failed to enqueue Attio user sync", { userId: payload.userId, error });
|
||||
}
|
||||
}
|
||||
|
||||
export async function runAttioWorkspaceSync(payload: AttioWorkspaceSync) {
|
||||
if (!attioClient) return;
|
||||
const admin = await prisma.user.findFirst({
|
||||
where: { id: payload.adminUserId },
|
||||
select: { email: true },
|
||||
});
|
||||
await attioClient.upsertWorkspace(payload, domainFromEmail(admin?.email));
|
||||
}
|
||||
|
||||
export async function runAttioUserSync(payload: AttioUserSync) {
|
||||
if (!attioClient) return;
|
||||
await attioClient.upsertUser(payload);
|
||||
}
|
||||
@@ -0,0 +1,33 @@
|
||||
import { Authenticator } from "remix-auth";
|
||||
import type { AuthUser } from "./authUser";
|
||||
import { addEmailLinkStrategy } from "./emailAuth.server";
|
||||
import { addGitHubStrategy } from "./gitHubAuth.server";
|
||||
import { addGoogleStrategy } from "./googleAuth.server";
|
||||
import { sessionStorage } from "./sessionStorage.server";
|
||||
import { addSsoStrategy } from "./ssoAuth.server";
|
||||
import { env } from "~/env.server";
|
||||
|
||||
// Create an instance of the authenticator, pass a generic with what
|
||||
// strategies will return and will store in the session
|
||||
const authenticator = new Authenticator<AuthUser>(sessionStorage);
|
||||
|
||||
const isGithubAuthSupported =
|
||||
typeof env.AUTH_GITHUB_CLIENT_ID === "string" &&
|
||||
typeof env.AUTH_GITHUB_CLIENT_SECRET === "string";
|
||||
|
||||
const isGoogleAuthSupported =
|
||||
typeof env.AUTH_GOOGLE_CLIENT_ID === "string" &&
|
||||
typeof env.AUTH_GOOGLE_CLIENT_SECRET === "string";
|
||||
|
||||
if (env.AUTH_GITHUB_CLIENT_ID && env.AUTH_GITHUB_CLIENT_SECRET) {
|
||||
addGitHubStrategy(authenticator, env.AUTH_GITHUB_CLIENT_ID, env.AUTH_GITHUB_CLIENT_SECRET);
|
||||
}
|
||||
|
||||
if (env.AUTH_GOOGLE_CLIENT_ID && env.AUTH_GOOGLE_CLIENT_SECRET) {
|
||||
addGoogleStrategy(authenticator, env.AUTH_GOOGLE_CLIENT_ID, env.AUTH_GOOGLE_CLIENT_SECRET);
|
||||
}
|
||||
|
||||
addEmailLinkStrategy(authenticator);
|
||||
addSsoStrategy(authenticator);
|
||||
|
||||
export { authenticator, isGithubAuthSupported, isGoogleAuthSupported };
|
||||
@@ -0,0 +1,11 @@
|
||||
export type AuthUser = {
|
||||
userId: string;
|
||||
// Present only when the session was established via SSO. Carries the
|
||||
// minimum the periodic re-validation hook needs to ask the IdP whether
|
||||
// the session is still valid. Signed into the session cookie, so it's
|
||||
// tamper-proof. Absent ⇒ non-SSO session ⇒ never revalidated.
|
||||
sso?: {
|
||||
idpOrgId: string;
|
||||
connectionId: string;
|
||||
};
|
||||
};
|
||||
@@ -0,0 +1,307 @@
|
||||
import type { Cache as UnkeyCache } from "@unkey/cache";
|
||||
import { createCache, DefaultStatefulContext, Namespace } from "@unkey/cache";
|
||||
import { createLRUMemoryStore } from "@internal/cache";
|
||||
import { Ratelimit } from "@upstash/ratelimit";
|
||||
import type { Request as ExpressRequest, Response as ExpressResponse, NextFunction } from "express";
|
||||
import { createHash } from "node:crypto";
|
||||
import { z } from "zod";
|
||||
import { env } from "~/env.server";
|
||||
import type { RedisWithClusterOptions } from "~/redis.server";
|
||||
import { logger } from "./logger.server";
|
||||
import type { Duration, Limiter } from "./rateLimiter.server";
|
||||
import { createRedisRateLimitClient, RateLimiter } from "./rateLimiter.server";
|
||||
import { RedisCacheStore } from "./unkey/redisCacheStore.server";
|
||||
|
||||
const DurationSchema = z.custom<Duration>((value) => {
|
||||
if (typeof value !== "string") {
|
||||
throw new Error("Duration must be a string");
|
||||
}
|
||||
|
||||
return value as Duration;
|
||||
});
|
||||
|
||||
export const RateLimitFixedWindowConfig = z.object({
|
||||
type: z.literal("fixedWindow"),
|
||||
window: DurationSchema,
|
||||
tokens: z.number(),
|
||||
});
|
||||
|
||||
export type RateLimitFixedWindowConfig = z.infer<typeof RateLimitFixedWindowConfig>;
|
||||
|
||||
export const RateLimitSlidingWindowConfig = z.object({
|
||||
type: z.literal("slidingWindow"),
|
||||
window: DurationSchema,
|
||||
tokens: z.number(),
|
||||
});
|
||||
|
||||
export type RateLimitSlidingWindowConfig = z.infer<typeof RateLimitSlidingWindowConfig>;
|
||||
|
||||
export const RateLimitTokenBucketConfig = z.object({
|
||||
type: z.literal("tokenBucket"),
|
||||
refillRate: z.number(),
|
||||
interval: DurationSchema,
|
||||
maxTokens: z.number(),
|
||||
});
|
||||
|
||||
export type RateLimitTokenBucketConfig = z.infer<typeof RateLimitTokenBucketConfig>;
|
||||
|
||||
export const RateLimiterConfig = z.discriminatedUnion("type", [
|
||||
RateLimitFixedWindowConfig,
|
||||
RateLimitSlidingWindowConfig,
|
||||
RateLimitTokenBucketConfig,
|
||||
]);
|
||||
|
||||
export type RateLimiterConfig = z.infer<typeof RateLimiterConfig>;
|
||||
|
||||
type LimitConfigOverrideFunction = (authorizationValue: string) => Promise<unknown>;
|
||||
|
||||
type Options = {
|
||||
redis?: RedisWithClusterOptions;
|
||||
keyPrefix: string;
|
||||
pathMatchers: (RegExp | string)[];
|
||||
pathWhiteList?: (RegExp | string)[];
|
||||
defaultLimiter: RateLimiterConfig;
|
||||
limiterConfigOverride?: LimitConfigOverrideFunction;
|
||||
limiterCache?: {
|
||||
fresh: number;
|
||||
stale: number;
|
||||
maxItems: number;
|
||||
};
|
||||
log?: {
|
||||
requests?: boolean;
|
||||
rejections?: boolean;
|
||||
limiter?: boolean;
|
||||
};
|
||||
};
|
||||
|
||||
async function resolveLimitConfig(
|
||||
authorizationValue: string,
|
||||
hashedAuthorizationValue: string,
|
||||
defaultLimiter: RateLimiterConfig,
|
||||
cache: UnkeyCache<{ limiter: RateLimiterConfig }>,
|
||||
logsEnabled: boolean,
|
||||
limiterConfigOverride?: LimitConfigOverrideFunction
|
||||
): Promise<RateLimiterConfig> {
|
||||
if (!limiterConfigOverride) {
|
||||
return defaultLimiter;
|
||||
}
|
||||
|
||||
if (logsEnabled) {
|
||||
logger.info("RateLimiter: checking for override", {
|
||||
authorizationValue: hashedAuthorizationValue,
|
||||
defaultLimiter,
|
||||
});
|
||||
}
|
||||
|
||||
const cacheResult = await cache.limiter.swr(hashedAuthorizationValue, async (key) => {
|
||||
const override = await limiterConfigOverride(authorizationValue);
|
||||
|
||||
if (!override) {
|
||||
if (logsEnabled) {
|
||||
logger.info("RateLimiter: no override found", {
|
||||
authorizationValue,
|
||||
defaultLimiter,
|
||||
});
|
||||
}
|
||||
|
||||
return defaultLimiter;
|
||||
}
|
||||
|
||||
const parsedOverride = RateLimiterConfig.safeParse(override);
|
||||
|
||||
if (!parsedOverride.success) {
|
||||
logger.error("Error parsing rate limiter override", {
|
||||
override,
|
||||
errors: parsedOverride.error.errors,
|
||||
});
|
||||
|
||||
return defaultLimiter;
|
||||
}
|
||||
|
||||
if (logsEnabled && parsedOverride.data) {
|
||||
logger.info("RateLimiter: override found", {
|
||||
authorizationValue,
|
||||
defaultLimiter,
|
||||
override: parsedOverride.data,
|
||||
});
|
||||
}
|
||||
|
||||
return parsedOverride.data;
|
||||
});
|
||||
|
||||
return cacheResult.val ?? defaultLimiter;
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates a Ratelimit limiter from a RateLimiterConfig.
|
||||
* This function is shared across the codebase to ensure consistent limiter creation.
|
||||
*/
|
||||
export function createLimiterFromConfig(config: RateLimiterConfig): Limiter {
|
||||
return config.type === "fixedWindow"
|
||||
? Ratelimit.fixedWindow(config.tokens, config.window)
|
||||
: config.type === "tokenBucket"
|
||||
? Ratelimit.tokenBucket(config.refillRate, config.interval, config.maxTokens)
|
||||
: Ratelimit.slidingWindow(config.tokens, config.window);
|
||||
}
|
||||
|
||||
//returns an Express middleware that rate limits using the Bearer token in the Authorization header
|
||||
export function authorizationRateLimitMiddleware({
|
||||
redis,
|
||||
keyPrefix,
|
||||
defaultLimiter,
|
||||
pathMatchers,
|
||||
pathWhiteList = [],
|
||||
log = {
|
||||
rejections: true,
|
||||
requests: true,
|
||||
},
|
||||
limiterCache,
|
||||
limiterConfigOverride,
|
||||
}: Options) {
|
||||
const ctx = new DefaultStatefulContext();
|
||||
const memory = createLRUMemoryStore(limiterCache?.maxItems ?? 1000);
|
||||
const redisCacheStore = new RedisCacheStore({
|
||||
connection: {
|
||||
keyPrefix: `cache:${keyPrefix}:rate-limit-cache:`,
|
||||
...redis,
|
||||
},
|
||||
});
|
||||
|
||||
// This cache holds the rate limit configuration for each org, so we don't have to fetch it every request
|
||||
const cache = createCache({
|
||||
limiter: new Namespace<RateLimiterConfig>(ctx, {
|
||||
stores: [memory, redisCacheStore],
|
||||
fresh: limiterCache?.fresh ?? 30_000,
|
||||
stale: limiterCache?.stale ?? 60_000,
|
||||
}),
|
||||
});
|
||||
|
||||
const redisClient = createRedisRateLimitClient(
|
||||
redis ?? {
|
||||
port: env.RATE_LIMIT_REDIS_PORT,
|
||||
host: env.RATE_LIMIT_REDIS_HOST,
|
||||
username: env.RATE_LIMIT_REDIS_USERNAME,
|
||||
password: env.RATE_LIMIT_REDIS_PASSWORD,
|
||||
tlsDisabled: env.RATE_LIMIT_REDIS_TLS_DISABLED === "true",
|
||||
clusterMode: env.RATE_LIMIT_REDIS_CLUSTER_MODE_ENABLED === "1",
|
||||
}
|
||||
);
|
||||
|
||||
return async (req: ExpressRequest, res: ExpressResponse, next: NextFunction) => {
|
||||
if (log.requests) {
|
||||
logger.info(`RateLimiter (${keyPrefix}): request to ${req.path}`);
|
||||
}
|
||||
|
||||
// allow OPTIONS requests
|
||||
if (req.method.toUpperCase() === "OPTIONS") {
|
||||
return next();
|
||||
}
|
||||
|
||||
//first check if any of the pathMatchers match the request path
|
||||
const path = req.path;
|
||||
if (
|
||||
!pathMatchers.some((matcher) =>
|
||||
matcher instanceof RegExp ? matcher.test(path) : path === matcher
|
||||
)
|
||||
) {
|
||||
if (log.requests) {
|
||||
logger.info(`RateLimiter (${keyPrefix}): didn't match ${req.path}`);
|
||||
}
|
||||
return next();
|
||||
}
|
||||
|
||||
// Check if the path matches any of the whitelisted paths
|
||||
if (
|
||||
pathWhiteList.some((matcher) =>
|
||||
matcher instanceof RegExp ? matcher.test(path) : path === matcher
|
||||
)
|
||||
) {
|
||||
if (log.requests) {
|
||||
logger.info(`RateLimiter (${keyPrefix}): whitelisted ${req.path}`);
|
||||
}
|
||||
return next();
|
||||
}
|
||||
|
||||
if (log.requests) {
|
||||
logger.info(`RateLimiter (${keyPrefix}): matched ${req.path}`);
|
||||
}
|
||||
|
||||
const authorizationValue = req.headers.authorization;
|
||||
if (!authorizationValue) {
|
||||
if (log.requests) {
|
||||
logger.info(`RateLimiter (${keyPrefix}): no key`, { headers: req.headers, url: req.url });
|
||||
}
|
||||
res.setHeader("Content-Type", "application/problem+json");
|
||||
return res.status(401).send(
|
||||
JSON.stringify(
|
||||
{
|
||||
title: "Unauthorized",
|
||||
status: 401,
|
||||
type: "https://developer.mozilla.org/en-US/docs/Web/HTTP/Status/401",
|
||||
detail: "No authorization header provided",
|
||||
error: "No authorization header provided",
|
||||
},
|
||||
null,
|
||||
2
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
const hash = createHash("sha256");
|
||||
hash.update(authorizationValue);
|
||||
const hashedAuthorizationValue = hash.digest("hex");
|
||||
|
||||
const limiterConfig = await resolveLimitConfig(
|
||||
authorizationValue,
|
||||
hashedAuthorizationValue,
|
||||
defaultLimiter,
|
||||
cache,
|
||||
typeof log.limiter === "boolean" ? log.limiter : false,
|
||||
limiterConfigOverride
|
||||
);
|
||||
|
||||
const limiter = createLimiterFromConfig(limiterConfig);
|
||||
|
||||
const rateLimiter = new RateLimiter({
|
||||
redisClient,
|
||||
keyPrefix,
|
||||
limiter,
|
||||
logSuccess: log.requests,
|
||||
logFailure: log.rejections,
|
||||
});
|
||||
|
||||
const { success, limit, reset, remaining } = await rateLimiter.limit(hashedAuthorizationValue);
|
||||
|
||||
const $remaining = Math.max(0, remaining); // remaining can be negative if the user has exceeded the limit, so clamp it to 0
|
||||
|
||||
res.set("x-ratelimit-limit", limit.toString());
|
||||
res.set("x-ratelimit-remaining", $remaining.toString());
|
||||
res.set("x-ratelimit-reset", reset.toString());
|
||||
|
||||
if (success) {
|
||||
return next();
|
||||
}
|
||||
|
||||
res.setHeader("Content-Type", "application/problem+json");
|
||||
const secondsUntilReset = Math.max(0, (reset - new Date().getTime()) / 1000);
|
||||
return res.status(429).send(
|
||||
JSON.stringify(
|
||||
{
|
||||
title: "Rate Limit Exceeded",
|
||||
status: 429,
|
||||
type: "https://developer.mozilla.org/en-US/docs/Web/HTTP/Status/429",
|
||||
detail: `Rate limit exceeded ${$remaining}/${limit} requests remaining. Retry in ${secondsUntilReset} seconds.`,
|
||||
reset,
|
||||
limit,
|
||||
remaining,
|
||||
secondsUntilReset,
|
||||
error: `Rate limit exceeded ${$remaining}/${limit} requests remaining. Retry in ${secondsUntilReset} seconds.`,
|
||||
},
|
||||
null,
|
||||
2
|
||||
)
|
||||
);
|
||||
};
|
||||
}
|
||||
|
||||
export type RateLimitMiddleware = ReturnType<typeof authorizationRateLimitMiddleware>;
|
||||
@@ -0,0 +1,86 @@
|
||||
import type { RedisOptions } from "ioredis";
|
||||
import Redis from "ioredis";
|
||||
import { defaultReconnectOnError } from "@internal/redis";
|
||||
import type { PrismaClientOrTransaction, PrismaTransactionOptions } from "~/db.server";
|
||||
import { Prisma, prisma } from "~/db.server";
|
||||
import { env } from "~/env.server";
|
||||
import { singleton } from "~/utils/singleton";
|
||||
|
||||
export type AutoIncrementCounterOptions = {
|
||||
redis: RedisOptions;
|
||||
};
|
||||
|
||||
export class AutoIncrementCounter {
|
||||
private _redis: Redis;
|
||||
|
||||
constructor(private options: AutoIncrementCounterOptions) {
|
||||
this._redis = new Redis({ reconnectOnError: defaultReconnectOnError, ...options.redis });
|
||||
}
|
||||
|
||||
async incrementInTransaction<T>(
|
||||
key: string,
|
||||
callback: (num: number, tx: PrismaClientOrTransaction) => Promise<T>,
|
||||
backfiller?: (key: string, db: PrismaClientOrTransaction) => Promise<number | undefined>,
|
||||
client: PrismaClientOrTransaction = prisma,
|
||||
transactionOptions?: PrismaTransactionOptions
|
||||
): Promise<T | undefined> {
|
||||
let performedIncrement = false;
|
||||
let performedBackfill = false;
|
||||
|
||||
try {
|
||||
let newNumber = await this.#increment(key);
|
||||
|
||||
performedIncrement = true;
|
||||
|
||||
if (newNumber === 1 && backfiller) {
|
||||
const backfilledNumber = await backfiller(key, client);
|
||||
|
||||
if (backfilledNumber && backfilledNumber > 1) {
|
||||
newNumber = backfilledNumber + 1;
|
||||
await this._redis.set(key, newNumber);
|
||||
performedBackfill = true;
|
||||
}
|
||||
}
|
||||
|
||||
return await callback(newNumber, client);
|
||||
} catch (e) {
|
||||
if (
|
||||
e instanceof Prisma.PrismaClientKnownRequestError ||
|
||||
e instanceof Prisma.PrismaClientUnknownRequestError ||
|
||||
e instanceof Prisma.PrismaClientValidationError
|
||||
) {
|
||||
if (performedIncrement && !performedBackfill) {
|
||||
await this._redis.decr(key);
|
||||
}
|
||||
}
|
||||
|
||||
throw e;
|
||||
}
|
||||
}
|
||||
|
||||
async #increment(key: string): Promise<number> {
|
||||
return await this._redis.incr(key);
|
||||
}
|
||||
}
|
||||
|
||||
export const autoIncrementCounter = singleton("auto-increment-counter", getAutoIncrementCounter);
|
||||
|
||||
function getAutoIncrementCounter() {
|
||||
if (!env.REDIS_HOST || !env.REDIS_PORT) {
|
||||
throw new Error(
|
||||
"Could not initialize auto-increment counter because process.env.REDIS_HOST and process.env.REDIS_PORT are required to be set. "
|
||||
);
|
||||
}
|
||||
|
||||
return new AutoIncrementCounter({
|
||||
redis: {
|
||||
keyPrefix: "auto-counter:",
|
||||
port: env.REDIS_PORT,
|
||||
host: env.REDIS_HOST,
|
||||
username: env.REDIS_USERNAME,
|
||||
password: env.REDIS_PASSWORD,
|
||||
enableAutoPipelining: true,
|
||||
...(env.REDIS_TLS_DISABLED === "true" ? {} : { tls: {} }),
|
||||
},
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,169 @@
|
||||
import { wrapZodFetch } from "@trigger.dev/core/v3/zodfetch";
|
||||
import { createCache, DefaultStatefulContext, Namespace } from "@unkey/cache";
|
||||
import { createLRUMemoryStore } from "@internal/cache";
|
||||
import { z } from "zod";
|
||||
import { env } from "~/env.server";
|
||||
|
||||
const StatusPageSchema = z.object({
|
||||
data: z.object({
|
||||
id: z.string(),
|
||||
type: z.string(),
|
||||
attributes: z.object({
|
||||
aggregate_state: z.enum(["operational", "degraded", "downtime"]),
|
||||
}),
|
||||
}),
|
||||
});
|
||||
|
||||
const StatusReportsSchema = z.object({
|
||||
data: z.array(
|
||||
z.object({
|
||||
id: z.string(),
|
||||
type: z.literal("status_report"),
|
||||
attributes: z.object({
|
||||
title: z.string().nullable(),
|
||||
starts_at: z.string().nullable(),
|
||||
ends_at: z.string().nullable(),
|
||||
aggregate_state: z.string().nullable(),
|
||||
}),
|
||||
})
|
||||
),
|
||||
pagination: z.object({
|
||||
first: z.string().nullable(),
|
||||
last: z.string().nullable(),
|
||||
prev: z.string().nullable(),
|
||||
next: z.string().nullable(),
|
||||
}),
|
||||
});
|
||||
|
||||
export type AggregateState = "operational" | "degraded" | "downtime";
|
||||
|
||||
export type IncidentStatus = {
|
||||
status: AggregateState;
|
||||
title: string | null;
|
||||
};
|
||||
|
||||
type CachedResult = { success: true; data: IncidentStatus } | { success: false; error: unknown };
|
||||
|
||||
const ctx = new DefaultStatefulContext();
|
||||
const memory = createLRUMemoryStore(100);
|
||||
|
||||
const cache = createCache({
|
||||
query: new Namespace<CachedResult>(ctx, {
|
||||
stores: [memory],
|
||||
fresh: 15_000,
|
||||
stale: 30_000,
|
||||
}),
|
||||
});
|
||||
|
||||
export class BetterStackClient {
|
||||
private readonly baseUrl = "https://uptime.betterstack.com/api/v2";
|
||||
|
||||
async getIncidentStatus(): Promise<CachedResult> {
|
||||
const apiKey = env.BETTERSTACK_API_KEY;
|
||||
const statusPageId = env.BETTERSTACK_STATUS_PAGE_ID;
|
||||
|
||||
if (!apiKey || !statusPageId) {
|
||||
return { success: false, error: "Missing BetterStack configuration" };
|
||||
}
|
||||
|
||||
const cachedResult = await cache.query.swr("betterstack-incident-status", () =>
|
||||
this.fetchIncidentStatus(apiKey, statusPageId)
|
||||
);
|
||||
|
||||
if (cachedResult.err || !cachedResult.val) {
|
||||
return { success: false, error: cachedResult.err ?? "No result from cache" };
|
||||
}
|
||||
|
||||
return cachedResult.val;
|
||||
}
|
||||
|
||||
private async fetchIncidentStatus(apiKey: string, statusPageId: string): Promise<CachedResult> {
|
||||
const headers = {
|
||||
Authorization: `Bearer ${apiKey}`,
|
||||
"Content-Type": "application/json",
|
||||
};
|
||||
const retryConfig = {
|
||||
retry: { maxAttempts: 3, minTimeoutInMs: 1000, maxTimeoutInMs: 5000 },
|
||||
};
|
||||
|
||||
try {
|
||||
// Fetch the status page to get aggregate state
|
||||
const statusPageResult = await wrapZodFetch(
|
||||
StatusPageSchema,
|
||||
`${this.baseUrl}/status-pages/${statusPageId}`,
|
||||
{ headers },
|
||||
retryConfig
|
||||
);
|
||||
|
||||
if (!statusPageResult.success) {
|
||||
return { success: false, error: statusPageResult.error };
|
||||
}
|
||||
|
||||
const status = statusPageResult.data.data.attributes.aggregate_state;
|
||||
|
||||
// If operational, no need to fetch reports
|
||||
if (status === "operational") {
|
||||
return { success: true, data: { status, title: null } };
|
||||
}
|
||||
|
||||
// Fetch status reports to get the incident title
|
||||
const title = await this.fetchActiveReportTitle(apiKey, statusPageId, headers, retryConfig);
|
||||
|
||||
return { success: true, data: { status, title } };
|
||||
} catch (error) {
|
||||
console.error("Failed to fetch incident status from BetterStack:", error);
|
||||
return {
|
||||
success: false,
|
||||
error: error instanceof Error ? error.message : "Unknown error",
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
private async fetchActiveReportTitle(
|
||||
apiKey: string,
|
||||
statusPageId: string,
|
||||
headers: Record<string, string>,
|
||||
retryConfig: { retry: { maxAttempts: number; minTimeoutInMs: number; maxTimeoutInMs: number } }
|
||||
): Promise<string | null> {
|
||||
const reportsUrl = `${this.baseUrl}/status-pages/${statusPageId}/status-reports`;
|
||||
|
||||
let reportsResult = await wrapZodFetch(
|
||||
StatusReportsSchema,
|
||||
reportsUrl,
|
||||
{ headers },
|
||||
retryConfig
|
||||
);
|
||||
|
||||
if (!reportsResult.success) {
|
||||
return null;
|
||||
}
|
||||
|
||||
// Fetch last page if there are multiple pages (most recent reports are at the end)
|
||||
const { first, last } = reportsResult.data.pagination;
|
||||
if (last && last !== first) {
|
||||
const lastPageResult = await wrapZodFetch(
|
||||
StatusReportsSchema,
|
||||
last,
|
||||
{ headers },
|
||||
retryConfig
|
||||
);
|
||||
if (lastPageResult.success) {
|
||||
reportsResult = lastPageResult;
|
||||
}
|
||||
}
|
||||
|
||||
// Find active reports (not resolved, not ended)
|
||||
const activeReports = reportsResult.data.data.filter(
|
||||
(report) =>
|
||||
report.attributes.aggregate_state !== "resolved" && report.attributes.ends_at === null
|
||||
);
|
||||
|
||||
if (activeReports.length === 0) {
|
||||
return null;
|
||||
}
|
||||
|
||||
// Return the title from the most recent active report
|
||||
const mostRecent = activeReports[activeReports.length - 1];
|
||||
return mostRecent.attributes.title;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,181 @@
|
||||
import type { BillingClient } from "@trigger.dev/platform";
|
||||
import { z } from "zod";
|
||||
|
||||
/**
|
||||
* Billing limit API schemas for the billing platform service.
|
||||
*
|
||||
* These mirror the planned @trigger.dev/platform types and are used via
|
||||
* BillingClient.fetch until the platform package is published with native
|
||||
* BillingClient methods.
|
||||
*/
|
||||
|
||||
export const BillingLimitStateSchema = z.discriminatedUnion("status", [
|
||||
z.object({
|
||||
status: z.literal("ok"),
|
||||
}),
|
||||
z.object({
|
||||
status: z.literal("grace"),
|
||||
hitAt: z.string().datetime({ offset: true }),
|
||||
graceEndsAt: z.string().datetime({ offset: true }),
|
||||
}),
|
||||
z.object({
|
||||
status: z.literal("rejected"),
|
||||
hitAt: z.string().datetime({ offset: true }),
|
||||
graceEndsAt: z.string().datetime({ offset: true }),
|
||||
}),
|
||||
]);
|
||||
|
||||
export type BillingLimitState = z.infer<typeof BillingLimitStateSchema>;
|
||||
|
||||
export const BillingLimitConfigSchema = z.discriminatedUnion("mode", [
|
||||
z.object({
|
||||
mode: z.literal("none"),
|
||||
}),
|
||||
z.object({
|
||||
mode: z.literal("plan"),
|
||||
}),
|
||||
z.object({
|
||||
mode: z.literal("custom"),
|
||||
amountCents: z.number().int().positive(),
|
||||
}),
|
||||
]);
|
||||
|
||||
export type BillingLimitConfig = z.infer<typeof BillingLimitConfigSchema>;
|
||||
|
||||
export const BillingLimitUnconfiguredSchema = z.object({
|
||||
isConfigured: z.literal(false),
|
||||
gracePeriodMs: z.number().int().nonnegative(),
|
||||
});
|
||||
|
||||
const billingLimitConfiguredFields = {
|
||||
isConfigured: z.literal(true),
|
||||
cancelInProgressRuns: z.boolean(),
|
||||
limitState: BillingLimitStateSchema,
|
||||
effectiveAmountCents: z.number().int().nonnegative().nullable(),
|
||||
gracePeriodMs: z.number().int().nonnegative(),
|
||||
};
|
||||
|
||||
export const BillingLimitConfiguredNoneSchema = z.object({
|
||||
...billingLimitConfiguredFields,
|
||||
mode: z.literal("none"),
|
||||
});
|
||||
|
||||
export const BillingLimitConfiguredPlanSchema = z.object({
|
||||
...billingLimitConfiguredFields,
|
||||
mode: z.literal("plan"),
|
||||
});
|
||||
|
||||
export const BillingLimitConfiguredCustomSchema = z.object({
|
||||
...billingLimitConfiguredFields,
|
||||
mode: z.literal("custom"),
|
||||
amountCents: z.number().int().positive(),
|
||||
});
|
||||
|
||||
export const BillingLimitConfiguredSchema = z.discriminatedUnion("mode", [
|
||||
BillingLimitConfiguredNoneSchema,
|
||||
BillingLimitConfiguredPlanSchema,
|
||||
BillingLimitConfiguredCustomSchema,
|
||||
]);
|
||||
|
||||
export const BillingLimitResultSchema = z.union([
|
||||
BillingLimitUnconfiguredSchema,
|
||||
BillingLimitConfiguredNoneSchema,
|
||||
BillingLimitConfiguredPlanSchema,
|
||||
BillingLimitConfiguredCustomSchema,
|
||||
]);
|
||||
|
||||
export type BillingLimitResult = z.infer<typeof BillingLimitResultSchema>;
|
||||
|
||||
export const UpdateBillingLimitRequestSchema = z.discriminatedUnion("mode", [
|
||||
z.object({
|
||||
mode: z.literal("none"),
|
||||
cancelInProgressRuns: z.boolean(),
|
||||
}),
|
||||
z.object({
|
||||
mode: z.literal("plan"),
|
||||
cancelInProgressRuns: z.boolean(),
|
||||
}),
|
||||
z.object({
|
||||
mode: z.literal("custom"),
|
||||
amountCents: z.number().int().positive(),
|
||||
cancelInProgressRuns: z.boolean(),
|
||||
}),
|
||||
]);
|
||||
|
||||
export type UpdateBillingLimitRequest = z.infer<typeof UpdateBillingLimitRequestSchema>;
|
||||
|
||||
export const ResolveBillingLimitRequestSchema = z.discriminatedUnion("action", [
|
||||
z.object({
|
||||
action: z.literal("increase"),
|
||||
newAmountCents: z.number().int().positive(),
|
||||
resumeMode: z.enum(["queue", "new_only"]),
|
||||
}),
|
||||
z.object({
|
||||
action: z.literal("remove"),
|
||||
resumeMode: z.enum(["queue", "new_only"]),
|
||||
}),
|
||||
]);
|
||||
|
||||
export type ResolveBillingLimitRequest = z.infer<typeof ResolveBillingLimitRequestSchema>;
|
||||
|
||||
export const BillingLimitActiveOrgSchema = z.object({
|
||||
orgId: z.string(),
|
||||
limitState: z.enum(["grace", "rejected"]),
|
||||
});
|
||||
|
||||
export const BillingLimitsActiveResultSchema = z.object({
|
||||
orgs: z.array(BillingLimitActiveOrgSchema),
|
||||
});
|
||||
|
||||
export type BillingLimitsActiveResult = z.infer<typeof BillingLimitsActiveResultSchema>;
|
||||
|
||||
export const BillingLimitPendingResolveOrgSchema = z.object({
|
||||
organizationId: z.string(),
|
||||
resumeMode: z.enum(["queue", "new_only"]),
|
||||
resolvedAt: z.string().datetime({ offset: true }),
|
||||
});
|
||||
|
||||
export const BillingLimitsPendingResolvesResultSchema = z.object({
|
||||
orgs: z.array(BillingLimitPendingResolveOrgSchema),
|
||||
});
|
||||
|
||||
export type BillingLimitsPendingResolvesResult = z.infer<
|
||||
typeof BillingLimitsPendingResolvesResultSchema
|
||||
>;
|
||||
|
||||
export const BillingLimitHitWebhookBodySchema = z.object({
|
||||
hitAt: z.string().datetime({ offset: true }),
|
||||
cancelInProgressRuns: z.boolean(),
|
||||
limitState: z.literal("grace"),
|
||||
});
|
||||
|
||||
export type BillingLimitHitWebhookBody = z.infer<typeof BillingLimitHitWebhookBodySchema>;
|
||||
|
||||
/** Entitlement response — mirrors ReportUsageResult with billing limit fields until platform ships native types. */
|
||||
export const EntitlementResultSchema = z.object({
|
||||
hasAccess: z.boolean(),
|
||||
balance: z.number().optional(),
|
||||
usage: z.number().optional(),
|
||||
overage: z.number().optional(),
|
||||
plan: z
|
||||
.object({
|
||||
type: z.string(),
|
||||
code: z.string(),
|
||||
isPaying: z.boolean(),
|
||||
})
|
||||
.optional(),
|
||||
limitState: z.literal("grace").optional(),
|
||||
reason: z.enum(["free_tier_exceeded", "billing_limit"]).optional(),
|
||||
});
|
||||
|
||||
export type EntitlementResult = z.infer<typeof EntitlementResultSchema>;
|
||||
|
||||
export type BillingLimitPageData = BillingLimitResult & {
|
||||
queuedRunCount: number;
|
||||
currentSpendCents: number;
|
||||
};
|
||||
|
||||
/** Bridge webapp Zod schemas to BillingClient.fetch (separate Zod type instances). */
|
||||
export function asPlatformSchema(schema: z.ZodTypeAny) {
|
||||
return schema as unknown as Parameters<BillingClient["fetch"]>[1];
|
||||
}
|
||||
@@ -0,0 +1,599 @@
|
||||
import { ClickHouse } from "@internal/clickhouse";
|
||||
import { createHash } from "crypto";
|
||||
import { ClickhouseEventRepository } from "~/v3/eventRepository/clickhouseEventRepository.server";
|
||||
import { env } from "~/env.server";
|
||||
import { clampToEmergencySpanCap } from "~/v3/eventRepository/emergencySpanCap.server";
|
||||
import { singleton } from "~/utils/singleton";
|
||||
import type { OrganizationDataStoresRegistry } from "~/services/dataStores/organizationDataStoresRegistry.server";
|
||||
import { type IEventRepository } from "~/v3/eventRepository/eventRepository.types";
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Default clients (singleton per process)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
const defaultClickhouseClient = singleton("clickhouseClient", initializeClickhouseClient);
|
||||
|
||||
function initializeClickhouseClient() {
|
||||
const url = new URL(env.CLICKHOUSE_URL);
|
||||
url.searchParams.delete("secure");
|
||||
|
||||
console.log(`🗃️ Clickhouse service enabled to host ${url.host}`);
|
||||
|
||||
return new ClickHouse({
|
||||
url: url.toString(),
|
||||
name: "clickhouse-instance",
|
||||
keepAlive: {
|
||||
enabled: env.CLICKHOUSE_KEEP_ALIVE_ENABLED === "1",
|
||||
idleSocketTtl: env.CLICKHOUSE_KEEP_ALIVE_IDLE_SOCKET_TTL_MS,
|
||||
},
|
||||
logLevel: env.CLICKHOUSE_LOG_LEVEL,
|
||||
compression: { request: true },
|
||||
maxOpenConnections: env.CLICKHOUSE_MAX_OPEN_CONNECTIONS,
|
||||
});
|
||||
}
|
||||
|
||||
const defaultLogsClickhouseClient = singleton(
|
||||
"logsClickhouseClient",
|
||||
initializeLogsClickhouseClient
|
||||
);
|
||||
|
||||
function getLogsListClickhouseSettings() {
|
||||
return {
|
||||
max_memory_usage: env.CLICKHOUSE_LOGS_LIST_MAX_MEMORY_USAGE.toString(),
|
||||
max_bytes_before_external_sort:
|
||||
env.CLICKHOUSE_LOGS_LIST_MAX_BYTES_BEFORE_EXTERNAL_SORT.toString(),
|
||||
max_threads: env.CLICKHOUSE_LOGS_LIST_MAX_THREADS,
|
||||
// Cap per-part read buffers so read-in-order memory stays bounded. These exist everywhere.
|
||||
prefetch_buffer_size: env.CLICKHOUSE_LOGS_LIST_PREFETCH_BUFFER_SIZE.toString(),
|
||||
max_read_buffer_size: env.CLICKHOUSE_LOGS_LIST_MAX_READ_BUFFER_SIZE.toString(),
|
||||
// Object-storage only and newer than the buffers above, so only send it when configured to
|
||||
// avoid UNKNOWN_SETTING failures against older self-hosted ClickHouse that lack it.
|
||||
...(env.CLICKHOUSE_LOGS_LIST_FILESYSTEM_CACHE_PREFER_BIGGER_BUFFER_SIZE !== undefined && {
|
||||
filesystem_cache_prefer_bigger_buffer_size:
|
||||
env.CLICKHOUSE_LOGS_LIST_FILESYSTEM_CACHE_PREFER_BIGGER_BUFFER_SIZE,
|
||||
}),
|
||||
...(env.CLICKHOUSE_LOGS_LIST_MAX_ROWS_TO_READ && {
|
||||
max_rows_to_read: env.CLICKHOUSE_LOGS_LIST_MAX_ROWS_TO_READ.toString(),
|
||||
}),
|
||||
...(env.CLICKHOUSE_LOGS_LIST_MAX_EXECUTION_TIME && {
|
||||
max_execution_time: env.CLICKHOUSE_LOGS_LIST_MAX_EXECUTION_TIME,
|
||||
}),
|
||||
};
|
||||
}
|
||||
|
||||
function initializeLogsClickhouseClient() {
|
||||
if (!env.LOGS_CLICKHOUSE_URL) {
|
||||
throw new Error("LOGS_CLICKHOUSE_URL is not set");
|
||||
}
|
||||
|
||||
const url = new URL(env.LOGS_CLICKHOUSE_URL);
|
||||
url.searchParams.delete("secure");
|
||||
|
||||
return new ClickHouse({
|
||||
url: url.toString(),
|
||||
name: "logs-clickhouse",
|
||||
keepAlive: {
|
||||
enabled: env.CLICKHOUSE_KEEP_ALIVE_ENABLED === "1",
|
||||
idleSocketTtl: env.CLICKHOUSE_KEEP_ALIVE_IDLE_SOCKET_TTL_MS,
|
||||
},
|
||||
logLevel: env.CLICKHOUSE_LOG_LEVEL,
|
||||
compression: { request: true },
|
||||
maxOpenConnections: env.CLICKHOUSE_MAX_OPEN_CONNECTIONS,
|
||||
clickhouseSettings: getLogsListClickhouseSettings(),
|
||||
});
|
||||
}
|
||||
|
||||
const defaultAdminClickhouseClient = singleton(
|
||||
"adminClickhouseClient",
|
||||
initializeAdminClickhouseClient
|
||||
);
|
||||
|
||||
function initializeAdminClickhouseClient() {
|
||||
if (!env.ADMIN_CLICKHOUSE_URL) {
|
||||
throw new Error("ADMIN_CLICKHOUSE_URL is not set");
|
||||
}
|
||||
|
||||
const url = new URL(env.ADMIN_CLICKHOUSE_URL);
|
||||
url.searchParams.delete("secure");
|
||||
|
||||
return new ClickHouse({
|
||||
url: url.toString(),
|
||||
name: "admin-clickhouse",
|
||||
keepAlive: {
|
||||
enabled: env.CLICKHOUSE_KEEP_ALIVE_ENABLED === "1",
|
||||
idleSocketTtl: env.CLICKHOUSE_KEEP_ALIVE_IDLE_SOCKET_TTL_MS,
|
||||
},
|
||||
logLevel: env.CLICKHOUSE_LOG_LEVEL,
|
||||
compression: { request: true },
|
||||
maxOpenConnections: env.CLICKHOUSE_MAX_OPEN_CONNECTIONS,
|
||||
});
|
||||
}
|
||||
|
||||
const defaultQueryClickhouseClient = singleton(
|
||||
"queryClickhouseClient",
|
||||
initializeQueryClickhouseClient
|
||||
);
|
||||
|
||||
function initializeQueryClickhouseClient() {
|
||||
if (!env.QUERY_CLICKHOUSE_URL) {
|
||||
throw new Error("QUERY_CLICKHOUSE_URL is not set");
|
||||
}
|
||||
|
||||
const url = new URL(env.QUERY_CLICKHOUSE_URL);
|
||||
url.searchParams.delete("secure");
|
||||
|
||||
return new ClickHouse({
|
||||
url: url.toString(),
|
||||
name: "query-clickhouse",
|
||||
keepAlive: {
|
||||
enabled: env.CLICKHOUSE_KEEP_ALIVE_ENABLED === "1",
|
||||
idleSocketTtl: env.CLICKHOUSE_KEEP_ALIVE_IDLE_SOCKET_TTL_MS,
|
||||
},
|
||||
logLevel: env.CLICKHOUSE_LOG_LEVEL,
|
||||
compression: { request: true },
|
||||
maxOpenConnections: env.CLICKHOUSE_MAX_OPEN_CONNECTIONS,
|
||||
});
|
||||
}
|
||||
|
||||
/** TaskRun replication to ClickHouse (`RUN_REPLICATION_CLICKHOUSE_URL`); not exported. */
|
||||
const defaultRunsReplicationClickhouseClient = singleton(
|
||||
"runsReplicationClickhouseClient",
|
||||
initializeRunsReplicationClickhouseClient
|
||||
);
|
||||
|
||||
function initializeRunsReplicationClickhouseClient(): ClickHouse {
|
||||
if (!env.RUN_REPLICATION_CLICKHOUSE_URL) {
|
||||
// Runs replication worker gates on this URL; factory may still resolve "replication" for tests.
|
||||
return defaultClickhouseClient;
|
||||
}
|
||||
|
||||
const url = new URL(env.RUN_REPLICATION_CLICKHOUSE_URL);
|
||||
url.searchParams.delete("secure");
|
||||
|
||||
return new ClickHouse({
|
||||
url: url.toString(),
|
||||
name: "runs-replication",
|
||||
keepAlive: {
|
||||
enabled: env.RUN_REPLICATION_KEEP_ALIVE_ENABLED === "1",
|
||||
idleSocketTtl: env.RUN_REPLICATION_KEEP_ALIVE_IDLE_SOCKET_TTL_MS,
|
||||
},
|
||||
logLevel: env.RUN_REPLICATION_CLICKHOUSE_LOG_LEVEL,
|
||||
compression: { request: true },
|
||||
maxOpenConnections: env.RUN_REPLICATION_MAX_OPEN_CONNECTIONS,
|
||||
});
|
||||
}
|
||||
|
||||
/** Session replication to ClickHouse (`SESSION_REPLICATION_CLICKHOUSE_URL`); not exported. */
|
||||
const defaultSessionsReplicationClickhouseClient = singleton(
|
||||
"sessionsReplicationClickhouseClient",
|
||||
initializeSessionsReplicationClickhouseClient
|
||||
);
|
||||
|
||||
function initializeSessionsReplicationClickhouseClient(): ClickHouse {
|
||||
if (!env.SESSION_REPLICATION_CLICKHOUSE_URL) {
|
||||
// Sessions replication worker gates on this URL; factory may still resolve "sessions_replication" for tests.
|
||||
return defaultClickhouseClient;
|
||||
}
|
||||
|
||||
const url = new URL(env.SESSION_REPLICATION_CLICKHOUSE_URL);
|
||||
url.searchParams.delete("secure");
|
||||
|
||||
return new ClickHouse({
|
||||
url: url.toString(),
|
||||
name: "sessions-replication",
|
||||
keepAlive: {
|
||||
enabled: env.SESSION_REPLICATION_KEEP_ALIVE_ENABLED === "1",
|
||||
idleSocketTtl: env.SESSION_REPLICATION_KEEP_ALIVE_IDLE_SOCKET_TTL_MS,
|
||||
},
|
||||
logLevel: env.SESSION_REPLICATION_CLICKHOUSE_LOG_LEVEL,
|
||||
compression: { request: true },
|
||||
maxOpenConnections: env.SESSION_REPLICATION_MAX_OPEN_CONNECTIONS,
|
||||
});
|
||||
}
|
||||
|
||||
/** Run-engine PendingVersionSystem lookup (`RUN_ENGINE_CLICKHOUSE_URL`);
|
||||
* falls back to the default client if unset. */
|
||||
const defaultRunEngineClickhouseClient = singleton(
|
||||
"runEngineClickhouseClient",
|
||||
initializeRunEngineClickhouseClient
|
||||
);
|
||||
|
||||
function initializeRunEngineClickhouseClient(): ClickHouse {
|
||||
if (!env.RUN_ENGINE_CLICKHOUSE_URL) {
|
||||
return defaultClickhouseClient;
|
||||
}
|
||||
|
||||
const url = new URL(env.RUN_ENGINE_CLICKHOUSE_URL);
|
||||
url.searchParams.delete("secure");
|
||||
|
||||
return new ClickHouse({
|
||||
url: url.toString(),
|
||||
name: "run-engine-clickhouse",
|
||||
keepAlive: {
|
||||
enabled: env.RUN_ENGINE_CLICKHOUSE_KEEP_ALIVE_ENABLED === "1",
|
||||
idleSocketTtl: env.RUN_ENGINE_CLICKHOUSE_KEEP_ALIVE_IDLE_SOCKET_TTL_MS,
|
||||
},
|
||||
logLevel: env.RUN_ENGINE_CLICKHOUSE_LOG_LEVEL,
|
||||
compression: {
|
||||
request: env.RUN_ENGINE_CLICKHOUSE_COMPRESSION_REQUEST === "1",
|
||||
},
|
||||
maxOpenConnections: env.RUN_ENGINE_CLICKHOUSE_MAX_OPEN_CONNECTIONS,
|
||||
});
|
||||
}
|
||||
|
||||
/** Realtime runs feed tag/batch id resolution (`REALTIME_BACKEND_NATIVE_CLICKHOUSE_URL`);
|
||||
* falls back to the default client if unset. */
|
||||
const defaultRealtimeClickhouseClient = singleton(
|
||||
"realtimeClickhouseClient",
|
||||
initializeRealtimeClickhouseClient
|
||||
);
|
||||
|
||||
function initializeRealtimeClickhouseClient(): ClickHouse {
|
||||
if (!env.REALTIME_BACKEND_NATIVE_CLICKHOUSE_URL) {
|
||||
return defaultClickhouseClient;
|
||||
}
|
||||
|
||||
const url = new URL(env.REALTIME_BACKEND_NATIVE_CLICKHOUSE_URL);
|
||||
url.searchParams.delete("secure");
|
||||
|
||||
return new ClickHouse({
|
||||
url: url.toString(),
|
||||
name: "realtime-runs-clickhouse",
|
||||
keepAlive: {
|
||||
enabled: env.REALTIME_BACKEND_NATIVE_CLICKHOUSE_KEEP_ALIVE_ENABLED === "1",
|
||||
idleSocketTtl: env.REALTIME_BACKEND_NATIVE_CLICKHOUSE_KEEP_ALIVE_IDLE_SOCKET_TTL_MS,
|
||||
},
|
||||
logLevel: env.REALTIME_BACKEND_NATIVE_CLICKHOUSE_LOG_LEVEL,
|
||||
compression: {
|
||||
request: env.REALTIME_BACKEND_NATIVE_CLICKHOUSE_COMPRESSION_REQUEST === "1",
|
||||
},
|
||||
maxOpenConnections: env.REALTIME_BACKEND_NATIVE_CLICKHOUSE_MAX_OPEN_CONNECTIONS,
|
||||
});
|
||||
}
|
||||
|
||||
/** Task events (`EVENTS_CLICKHOUSE_URL`); not exported — accessed via factory. */
|
||||
const defaultEventsClickhouseClient = singleton(
|
||||
"eventsClickhouseClient",
|
||||
initializeEventsClickhouseClient
|
||||
);
|
||||
|
||||
function initializeEventsClickhouseClient(): ClickHouse {
|
||||
if (!env.EVENTS_CLICKHOUSE_URL) {
|
||||
throw new Error("EVENTS_CLICKHOUSE_URL is not set");
|
||||
}
|
||||
|
||||
const url = new URL(env.EVENTS_CLICKHOUSE_URL);
|
||||
url.searchParams.delete("secure");
|
||||
|
||||
return new ClickHouse({
|
||||
url: url.toString(),
|
||||
name: "task-events",
|
||||
keepAlive: {
|
||||
enabled: env.EVENTS_CLICKHOUSE_KEEP_ALIVE_ENABLED === "1",
|
||||
idleSocketTtl: env.EVENTS_CLICKHOUSE_KEEP_ALIVE_IDLE_SOCKET_TTL_MS,
|
||||
},
|
||||
logLevel: env.EVENTS_CLICKHOUSE_LOG_LEVEL,
|
||||
compression: {
|
||||
request: env.EVENTS_CLICKHOUSE_COMPRESSION_REQUEST === "1",
|
||||
},
|
||||
maxOpenConnections: env.EVENTS_CLICKHOUSE_MAX_OPEN_CONNECTIONS,
|
||||
});
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
function hashHostname(url: string): string {
|
||||
const parsed = new URL(url);
|
||||
return createHash("sha256").update(parsed.hostname).digest("hex");
|
||||
}
|
||||
|
||||
export type ClientType =
|
||||
| "standard"
|
||||
| "events"
|
||||
| "replication"
|
||||
| "sessions_replication"
|
||||
| "logs"
|
||||
| "query"
|
||||
| "admin"
|
||||
| "engine"
|
||||
| "realtime";
|
||||
|
||||
function buildOrgClickhouseClient(url: string, clientType: ClientType): ClickHouse {
|
||||
const parsed = new URL(url);
|
||||
parsed.searchParams.delete("secure");
|
||||
const name = `org-clickhouse-${clientType}`;
|
||||
|
||||
switch (clientType) {
|
||||
case "events":
|
||||
return new ClickHouse({
|
||||
url: parsed.toString(),
|
||||
name,
|
||||
keepAlive: {
|
||||
enabled: env.EVENTS_CLICKHOUSE_KEEP_ALIVE_ENABLED === "1",
|
||||
idleSocketTtl: env.EVENTS_CLICKHOUSE_KEEP_ALIVE_IDLE_SOCKET_TTL_MS,
|
||||
},
|
||||
logLevel: env.EVENTS_CLICKHOUSE_LOG_LEVEL,
|
||||
compression: {
|
||||
request: env.EVENTS_CLICKHOUSE_COMPRESSION_REQUEST === "1",
|
||||
},
|
||||
maxOpenConnections: env.EVENTS_CLICKHOUSE_MAX_OPEN_CONNECTIONS,
|
||||
});
|
||||
case "replication":
|
||||
return new ClickHouse({
|
||||
url: parsed.toString(),
|
||||
name,
|
||||
keepAlive: {
|
||||
enabled: env.RUN_REPLICATION_KEEP_ALIVE_ENABLED === "1",
|
||||
idleSocketTtl: env.RUN_REPLICATION_KEEP_ALIVE_IDLE_SOCKET_TTL_MS,
|
||||
},
|
||||
logLevel: env.RUN_REPLICATION_CLICKHOUSE_LOG_LEVEL,
|
||||
compression: { request: true },
|
||||
maxOpenConnections: env.RUN_REPLICATION_MAX_OPEN_CONNECTIONS,
|
||||
});
|
||||
case "sessions_replication":
|
||||
return new ClickHouse({
|
||||
url: parsed.toString(),
|
||||
name,
|
||||
keepAlive: {
|
||||
enabled: env.SESSION_REPLICATION_KEEP_ALIVE_ENABLED === "1",
|
||||
idleSocketTtl: env.SESSION_REPLICATION_KEEP_ALIVE_IDLE_SOCKET_TTL_MS,
|
||||
},
|
||||
logLevel: env.SESSION_REPLICATION_CLICKHOUSE_LOG_LEVEL,
|
||||
compression: { request: true },
|
||||
maxOpenConnections: env.SESSION_REPLICATION_MAX_OPEN_CONNECTIONS,
|
||||
});
|
||||
case "logs":
|
||||
return new ClickHouse({
|
||||
url: parsed.toString(),
|
||||
name,
|
||||
keepAlive: {
|
||||
enabled: env.CLICKHOUSE_KEEP_ALIVE_ENABLED === "1",
|
||||
idleSocketTtl: env.CLICKHOUSE_KEEP_ALIVE_IDLE_SOCKET_TTL_MS,
|
||||
},
|
||||
logLevel: env.CLICKHOUSE_LOG_LEVEL,
|
||||
compression: { request: true },
|
||||
maxOpenConnections: env.CLICKHOUSE_MAX_OPEN_CONNECTIONS,
|
||||
clickhouseSettings: getLogsListClickhouseSettings(),
|
||||
});
|
||||
case "engine":
|
||||
return new ClickHouse({
|
||||
url: parsed.toString(),
|
||||
name,
|
||||
keepAlive: {
|
||||
enabled: env.RUN_ENGINE_CLICKHOUSE_KEEP_ALIVE_ENABLED === "1",
|
||||
idleSocketTtl: env.RUN_ENGINE_CLICKHOUSE_KEEP_ALIVE_IDLE_SOCKET_TTL_MS,
|
||||
},
|
||||
logLevel: env.RUN_ENGINE_CLICKHOUSE_LOG_LEVEL,
|
||||
compression: {
|
||||
request: env.RUN_ENGINE_CLICKHOUSE_COMPRESSION_REQUEST === "1",
|
||||
},
|
||||
maxOpenConnections: env.RUN_ENGINE_CLICKHOUSE_MAX_OPEN_CONNECTIONS,
|
||||
});
|
||||
case "realtime":
|
||||
return new ClickHouse({
|
||||
url: parsed.toString(),
|
||||
name,
|
||||
keepAlive: {
|
||||
enabled: env.REALTIME_BACKEND_NATIVE_CLICKHOUSE_KEEP_ALIVE_ENABLED === "1",
|
||||
idleSocketTtl: env.REALTIME_BACKEND_NATIVE_CLICKHOUSE_KEEP_ALIVE_IDLE_SOCKET_TTL_MS,
|
||||
},
|
||||
logLevel: env.REALTIME_BACKEND_NATIVE_CLICKHOUSE_LOG_LEVEL,
|
||||
compression: {
|
||||
request: env.REALTIME_BACKEND_NATIVE_CLICKHOUSE_COMPRESSION_REQUEST === "1",
|
||||
},
|
||||
maxOpenConnections: env.REALTIME_BACKEND_NATIVE_CLICKHOUSE_MAX_OPEN_CONNECTIONS,
|
||||
});
|
||||
case "standard":
|
||||
case "query":
|
||||
case "admin":
|
||||
return new ClickHouse({
|
||||
url: parsed.toString(),
|
||||
name,
|
||||
keepAlive: {
|
||||
enabled: env.CLICKHOUSE_KEEP_ALIVE_ENABLED === "1",
|
||||
idleSocketTtl: env.CLICKHOUSE_KEEP_ALIVE_IDLE_SOCKET_TTL_MS,
|
||||
},
|
||||
logLevel: env.CLICKHOUSE_LOG_LEVEL,
|
||||
compression: { request: true },
|
||||
maxOpenConnections: env.CLICKHOUSE_MAX_OPEN_CONNECTIONS,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Factory class (injectable for testing)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
export class ClickhouseFactory {
|
||||
/** ClickHouse clients keyed by hostname hash + clientType. */
|
||||
private readonly _clientCache = new Map<string, ClickHouse>();
|
||||
/** Event repositories keyed by hostname hash (stateful, must be reused). */
|
||||
private readonly _eventRepositoryCache = new Map<string, ClickhouseEventRepository>();
|
||||
|
||||
constructor(private readonly _registry: OrganizationDataStoresRegistry) {}
|
||||
|
||||
async isReady(): Promise<boolean> {
|
||||
if (!this._registry.isLoaded) {
|
||||
await this._registry.isReady;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
async getClickhouseForOrganization(
|
||||
organizationId: string,
|
||||
clientType: ClientType
|
||||
): Promise<ClickHouse> {
|
||||
if (!this._registry.isLoaded) {
|
||||
await this._registry.isReady;
|
||||
}
|
||||
|
||||
return this.getClickhouseForOrganizationSync(organizationId, clientType);
|
||||
}
|
||||
|
||||
getClickhouseForOrganizationSync(organizationId: string, clientType: ClientType): ClickHouse {
|
||||
const dataStore = this._registry.get(organizationId, "CLICKHOUSE");
|
||||
|
||||
if (!dataStore) {
|
||||
switch (clientType) {
|
||||
case "standard":
|
||||
return defaultClickhouseClient;
|
||||
case "events":
|
||||
return defaultEventsClickhouseClient;
|
||||
case "replication":
|
||||
return defaultRunsReplicationClickhouseClient;
|
||||
case "sessions_replication":
|
||||
return defaultSessionsReplicationClickhouseClient;
|
||||
case "logs":
|
||||
return defaultLogsClickhouseClient;
|
||||
case "query":
|
||||
return defaultQueryClickhouseClient;
|
||||
case "admin":
|
||||
return defaultAdminClickhouseClient;
|
||||
case "engine":
|
||||
return defaultRunEngineClickhouseClient;
|
||||
case "realtime":
|
||||
return defaultRealtimeClickhouseClient;
|
||||
}
|
||||
}
|
||||
|
||||
const hostnameHash = hashHostname(dataStore.url);
|
||||
const cacheKey = `${hostnameHash}:${clientType}`;
|
||||
let client = this._clientCache.get(cacheKey);
|
||||
|
||||
if (!client) {
|
||||
client = buildOrgClickhouseClient(dataStore.url, clientType);
|
||||
this._clientCache.set(cacheKey, client);
|
||||
}
|
||||
|
||||
return client;
|
||||
}
|
||||
|
||||
async getEventRepositoryForOrganization(
|
||||
store: string,
|
||||
organizationId: string
|
||||
): Promise<{ key: string; repository: IEventRepository }> {
|
||||
if (!this._registry.isLoaded) {
|
||||
await this._registry.isReady;
|
||||
}
|
||||
|
||||
return this.getEventRepositoryForOrganizationSync(store, organizationId);
|
||||
}
|
||||
|
||||
getEventRepositoryForOrganizationSync(
|
||||
store: string,
|
||||
organizationId: string
|
||||
): { key: string; repository: IEventRepository } {
|
||||
const dataStore = this._registry.get(organizationId, "CLICKHOUSE");
|
||||
|
||||
if (!dataStore) {
|
||||
const defaultKey = `default:events:${store}`;
|
||||
let defaultRepo = this._eventRepositoryCache.get(defaultKey);
|
||||
if (!defaultRepo) {
|
||||
const eventsClickhouse = getEventsClickhouseClient();
|
||||
defaultRepo = buildEventRepository(store, eventsClickhouse);
|
||||
this._eventRepositoryCache.set(defaultKey, defaultRepo);
|
||||
}
|
||||
return { key: defaultKey, repository: defaultRepo };
|
||||
}
|
||||
|
||||
const hostnameHash = hashHostname(dataStore.url);
|
||||
const cacheKey = `${hostnameHash}:events:${store}`;
|
||||
let repository = this._eventRepositoryCache.get(cacheKey);
|
||||
|
||||
if (!repository) {
|
||||
const client = this.getClickhouseForOrganizationSync(organizationId, "events");
|
||||
repository = buildEventRepository(store, client);
|
||||
this._eventRepositoryCache.set(cacheKey, repository);
|
||||
}
|
||||
|
||||
return { key: cacheKey, repository: repository };
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Get admin ClickHouse client for cross-organization queries.
|
||||
* Only use for admin tools and analytics that need to query across all orgs.
|
||||
*/
|
||||
export function getAdminClickhouse(): ClickHouse {
|
||||
return defaultAdminClickhouseClient;
|
||||
}
|
||||
|
||||
export function getDefaultClickhouseClient(): ClickHouse {
|
||||
return defaultClickhouseClient;
|
||||
}
|
||||
|
||||
export function getDefaultLogsClickhouseClient(): ClickHouse {
|
||||
return defaultLogsClickhouseClient;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Private helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
function getEventsClickhouseClient(): ClickHouse {
|
||||
return defaultEventsClickhouseClient;
|
||||
}
|
||||
|
||||
function buildEventRepository(store: string, clickhouse: ClickHouse): ClickhouseEventRepository {
|
||||
switch (store) {
|
||||
case "clickhouse": {
|
||||
return new ClickhouseEventRepository({
|
||||
clickhouse,
|
||||
batchSize: env.EVENTS_CLICKHOUSE_BATCH_SIZE,
|
||||
flushInterval: env.EVENTS_CLICKHOUSE_FLUSH_INTERVAL_MS,
|
||||
maximumTraceSummaryViewCount: clampToEmergencySpanCap(
|
||||
env.EVENTS_CLICKHOUSE_MAX_TRACE_SUMMARY_VIEW_COUNT
|
||||
),
|
||||
maximumTraceDetailedSummaryViewCount: clampToEmergencySpanCap(
|
||||
env.EVENTS_CLICKHOUSE_MAX_TRACE_DETAILED_SUMMARY_VIEW_COUNT
|
||||
),
|
||||
maximumLiveReloadingSetting: env.EVENTS_CLICKHOUSE_MAX_LIVE_RELOADING_SETTING,
|
||||
insertStrategy: env.EVENTS_CLICKHOUSE_INSERT_STRATEGY,
|
||||
waitForAsyncInsert: env.EVENTS_CLICKHOUSE_WAIT_FOR_ASYNC_INSERT === "1",
|
||||
asyncInsertMaxDataSize: env.EVENTS_CLICKHOUSE_ASYNC_INSERT_MAX_DATA_SIZE,
|
||||
asyncInsertBusyTimeoutMs: env.EVENTS_CLICKHOUSE_ASYNC_INSERT_BUSY_TIMEOUT_MS,
|
||||
startTimeMaxAgeMs: env.EVENTS_CLICKHOUSE_START_TIME_MAX_AGE_MS,
|
||||
llmMetricsBatchSize: env.LLM_METRICS_BATCH_SIZE,
|
||||
llmMetricsFlushInterval: env.LLM_METRICS_FLUSH_INTERVAL_MS,
|
||||
llmMetricsMaxBatchSize: env.LLM_METRICS_MAX_BATCH_SIZE,
|
||||
llmMetricsMaxConcurrency: env.LLM_METRICS_MAX_CONCURRENCY,
|
||||
otlpMetricsBatchSize: env.METRICS_CLICKHOUSE_BATCH_SIZE,
|
||||
otlpMetricsFlushInterval: env.METRICS_CLICKHOUSE_FLUSH_INTERVAL_MS,
|
||||
otlpMetricsMaxConcurrency: env.METRICS_CLICKHOUSE_MAX_CONCURRENCY,
|
||||
version: "v1",
|
||||
});
|
||||
}
|
||||
case "clickhouse_v2": {
|
||||
return new ClickhouseEventRepository({
|
||||
clickhouse: clickhouse,
|
||||
batchSize: env.EVENTS_CLICKHOUSE_BATCH_SIZE,
|
||||
flushInterval: env.EVENTS_CLICKHOUSE_FLUSH_INTERVAL_MS,
|
||||
maximumTraceSummaryViewCount: clampToEmergencySpanCap(
|
||||
env.EVENTS_CLICKHOUSE_MAX_TRACE_SUMMARY_VIEW_COUNT
|
||||
),
|
||||
maximumTraceDetailedSummaryViewCount: clampToEmergencySpanCap(
|
||||
env.EVENTS_CLICKHOUSE_MAX_TRACE_DETAILED_SUMMARY_VIEW_COUNT
|
||||
),
|
||||
maximumLiveReloadingSetting: env.EVENTS_CLICKHOUSE_MAX_LIVE_RELOADING_SETTING,
|
||||
insertStrategy: env.EVENTS_CLICKHOUSE_INSERT_STRATEGY,
|
||||
waitForAsyncInsert: env.EVENTS_CLICKHOUSE_WAIT_FOR_ASYNC_INSERT === "1",
|
||||
asyncInsertMaxDataSize: env.EVENTS_CLICKHOUSE_ASYNC_INSERT_MAX_DATA_SIZE,
|
||||
asyncInsertBusyTimeoutMs: env.EVENTS_CLICKHOUSE_ASYNC_INSERT_BUSY_TIMEOUT_MS,
|
||||
startTimeMaxAgeMs: env.EVENTS_CLICKHOUSE_START_TIME_MAX_AGE_MS,
|
||||
llmMetricsBatchSize: env.LLM_METRICS_BATCH_SIZE,
|
||||
llmMetricsFlushInterval: env.LLM_METRICS_FLUSH_INTERVAL_MS,
|
||||
llmMetricsMaxBatchSize: env.LLM_METRICS_MAX_BATCH_SIZE,
|
||||
llmMetricsMaxConcurrency: env.LLM_METRICS_MAX_CONCURRENCY,
|
||||
otlpMetricsBatchSize: env.METRICS_CLICKHOUSE_BATCH_SIZE,
|
||||
otlpMetricsFlushInterval: env.METRICS_CLICKHOUSE_FLUSH_INTERVAL_MS,
|
||||
otlpMetricsMaxConcurrency: env.METRICS_CLICKHOUSE_MAX_CONCURRENCY,
|
||||
version: "v2",
|
||||
});
|
||||
}
|
||||
default: {
|
||||
throw new Error(`Unknown ClickHouse event repository store: ${store}`);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,13 @@
|
||||
import { organizationDataStoresRegistry } from "~/services/dataStores/organizationDataStoresRegistryInstance.server";
|
||||
import { singleton } from "~/utils/singleton";
|
||||
import { ClickhouseFactory } from "./clickhouseFactory.server";
|
||||
|
||||
/**
|
||||
* Production singleton wired to the global organization data-stores registry.
|
||||
* Import this only from app/runtime code — not from tests that construct a
|
||||
* {@link ClickhouseFactory} with a stub registry (see `clickhouseFactory.server.ts`).
|
||||
*/
|
||||
export const clickhouseFactory = singleton(
|
||||
"clickhouseFactory",
|
||||
() => new ClickhouseFactory(organizationDataStoresRegistry)
|
||||
);
|
||||
@@ -0,0 +1,11 @@
|
||||
import { z } from "zod";
|
||||
|
||||
export const ClickhouseConnectionSchema = z.object({
|
||||
url: z.string().url(),
|
||||
});
|
||||
|
||||
export type ClickhouseConnection = z.infer<typeof ClickhouseConnectionSchema>;
|
||||
|
||||
export function getClickhouseSecretKey(orgId: string, clientType: string): string {
|
||||
return `org:${orgId}:clickhouse:${clientType}`;
|
||||
}
|
||||
@@ -0,0 +1,229 @@
|
||||
import { signUserActorToken } from "@trigger.dev/rbac";
|
||||
import { TriggerClient } from "@trigger.dev/sdk";
|
||||
import { chat } from "@trigger.dev/sdk/ai";
|
||||
import { prisma } from "~/db.server";
|
||||
import { env } from "~/env.server";
|
||||
import { runStore } from "~/v3/runStore.server";
|
||||
import { githubApp } from "./gitHub.server";
|
||||
import { logger } from "./logger.server";
|
||||
|
||||
const TASK_ID = "dashboard-agent";
|
||||
|
||||
// Read-only cap on the agent's delegated user-actor token. `read:apiKeys` is
|
||||
// what lets it exchange the token for an env JWT (the gate on the exchange
|
||||
// route); the rest scope the actual reads. No write/admin scopes, so even a
|
||||
// leaked token can't mutate anything.
|
||||
const DASHBOARD_AGENT_UAT_CAP = [
|
||||
"read:apiKeys",
|
||||
"read:runs",
|
||||
"read:deployments",
|
||||
"read:environments",
|
||||
"read:errors",
|
||||
"read:query",
|
||||
];
|
||||
|
||||
// Minted fresh on every turn (the `in` proxy injects it), so the lifetime only
|
||||
// has to cover a single turn's tool calls. Short by design — a stale token in
|
||||
// the agent's run payload expires quickly.
|
||||
const DASHBOARD_AGENT_UAT_TTL_SECONDS = 10 * 60;
|
||||
|
||||
// The Trigger instance this webapp runs against — the same origin the agent
|
||||
// task calls back to (as the user) for its read tools.
|
||||
export function dashboardAgentApiOrigin(): string {
|
||||
return env.API_ORIGIN ?? env.APP_ORIGIN;
|
||||
}
|
||||
|
||||
// Mint a short-lived, read-only delegated token for the signed-in user. Self
|
||||
// service from the dashboard session (never a PAT), so a user can only ever
|
||||
// mint a token for themselves. The `in` proxy injects this into the turn's
|
||||
// metadata so the token reaches the agent without ever touching the browser.
|
||||
export function mintDashboardAgentUserActorToken(userId: string): Promise<string> {
|
||||
return signUserActorToken(env.SESSION_SECRET, {
|
||||
userId,
|
||||
client: "dashboard-agent",
|
||||
cap: DASHBOARD_AGENT_UAT_CAP,
|
||||
expirationTime: Math.floor(Date.now() / 1000) + DASHBOARD_AGENT_UAT_TTL_SECONDS,
|
||||
});
|
||||
}
|
||||
|
||||
// The session is created in whatever env DASHBOARD_AGENT_SECRET_KEY belongs to.
|
||||
// baseURL is the Trigger instance this webapp runs against (its own API origin).
|
||||
function dashboardAgentConfig() {
|
||||
const accessToken = env.DASHBOARD_AGENT_SECRET_KEY;
|
||||
if (!accessToken) return null;
|
||||
return { baseURL: dashboardAgentApiOrigin(), accessToken };
|
||||
}
|
||||
|
||||
export function isDashboardAgentConfigured(): boolean {
|
||||
return Boolean(env.DASHBOARD_AGENT_SECRET_KEY);
|
||||
}
|
||||
|
||||
// Pins every agent session (and its continuation runs) to a deployed version
|
||||
// when DASHBOARD_AGENT_VERSION is set; unset runs on the env's current version.
|
||||
export function dashboardAgentTriggerConfig(): { lockToVersion: string } | undefined {
|
||||
return env.DASHBOARD_AGENT_VERSION ? { lockToVersion: env.DASHBOARD_AGENT_VERSION } : undefined;
|
||||
}
|
||||
|
||||
export async function startDashboardAgentSession(params: {
|
||||
chatId: string;
|
||||
clientData?: Record<string, unknown>;
|
||||
}): Promise<{ publicAccessToken: string }> {
|
||||
const config = dashboardAgentConfig();
|
||||
if (!config) throw new Error("DASHBOARD_AGENT_SECRET_KEY is not set");
|
||||
const startSession = chat.createStartSessionAction(TASK_ID, {
|
||||
apiClient: config,
|
||||
triggerConfig: dashboardAgentTriggerConfig(),
|
||||
});
|
||||
return startSession({ chatId: params.chatId, clientData: params.clientData });
|
||||
}
|
||||
|
||||
export async function mintDashboardAgentToken(chatId: string): Promise<string> {
|
||||
const config = dashboardAgentConfig();
|
||||
if (!config) throw new Error("DASHBOARD_AGENT_SECRET_KEY is not set");
|
||||
const client = new TriggerClient(config);
|
||||
return client.auth.createPublicToken({
|
||||
scopes: { read: { sessions: chatId }, write: { sessions: chatId } },
|
||||
expirationTime: "1h",
|
||||
});
|
||||
}
|
||||
|
||||
// A signed, short-lived pointer to the project's connected repo at a commit. Only
|
||||
// the URL crosses to the agent; the GitHub token stays here. The agent's code
|
||||
// tools download + extract it on their own filesystem (see @internal/dashboard-agent).
|
||||
export type DashboardAgentRepoSnapshot = {
|
||||
tarballUrl: string;
|
||||
owner: string;
|
||||
repo: string;
|
||||
sha: string;
|
||||
defaultBranch?: string;
|
||||
};
|
||||
|
||||
// The GitHub archive redirect URL is valid for a few minutes; cache the resolved
|
||||
// pointer briefly so multi-turn chats don't re-mint a token + re-resolve on every
|
||||
// message. Keyed by project + ref.
|
||||
const repoSnapshotCache = new Map<
|
||||
string,
|
||||
{ snapshot: DashboardAgentRepoSnapshot; expiresAt: number }
|
||||
>();
|
||||
const REPO_SNAPSHOT_TTL_MS = 60_000;
|
||||
const REPO_SNAPSHOT_MAX_ENTRIES = 1_000;
|
||||
|
||||
// Drop expired entries (key cardinality grows with each unique project + pinned
|
||||
// SHA), then evict oldest-first if still over the cap, so the cache can't grow
|
||||
// unbounded over a process lifetime.
|
||||
function pruneRepoSnapshotCache(now = Date.now()) {
|
||||
for (const [key, value] of repoSnapshotCache) {
|
||||
if (value.expiresAt <= now) repoSnapshotCache.delete(key);
|
||||
}
|
||||
let overflow = repoSnapshotCache.size - REPO_SNAPSHOT_MAX_ENTRIES;
|
||||
if (overflow <= 0) return;
|
||||
for (const key of repoSnapshotCache.keys()) {
|
||||
repoSnapshotCache.delete(key);
|
||||
if (--overflow <= 0) break;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Resolve the code-mode repo snapshot for a project, or null when the GitHub App
|
||||
* is disabled / no repo is connected (which keeps the agent in assistant mode).
|
||||
*
|
||||
* Mints a `contents:read` installation token scoped to the one repo, resolves the
|
||||
* signed archive URL, and returns just that URL. The token never leaves the
|
||||
* server. `opts.ref` pins a specific commit (run-SHA pinning); without it, the
|
||||
* tracked prod branch (or the repo default) head is used.
|
||||
*/
|
||||
export async function resolveDashboardAgentRepoSnapshot(
|
||||
projectId: string,
|
||||
opts: { ref?: string } = {}
|
||||
): Promise<DashboardAgentRepoSnapshot | null> {
|
||||
if (!githubApp) return null;
|
||||
|
||||
// Cache per project + ref so HEAD and each pinned commit are cached separately.
|
||||
const cacheKey = `${projectId}:${opts.ref ?? "HEAD"}`;
|
||||
const cached = repoSnapshotCache.get(cacheKey);
|
||||
if (cached && cached.expiresAt > Date.now()) return cached.snapshot;
|
||||
if (cached) repoSnapshotCache.delete(cacheKey);
|
||||
|
||||
const connected = await prisma.connectedGithubRepository.findFirst({
|
||||
where: { projectId },
|
||||
select: {
|
||||
branchTracking: true,
|
||||
repository: {
|
||||
select: {
|
||||
fullName: true,
|
||||
defaultBranch: true,
|
||||
installation: { select: { appInstallationId: true } },
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
if (!connected) return null;
|
||||
|
||||
const [owner, repo] = connected.repository.fullName.split("/");
|
||||
if (!owner || !repo) return null;
|
||||
const installationId = Number(connected.repository.installation.appInstallationId);
|
||||
const defaultBranch = connected.repository.defaultBranch;
|
||||
const tracking = connected.branchTracking as { prod?: { branch?: string } } | null;
|
||||
// An explicit 40-char commit SHA is used directly (run-SHA pinning); otherwise
|
||||
// resolve the requested branch, the tracked prod branch, or the repo default.
|
||||
const requested = opts.ref;
|
||||
const isSha = !!requested && /^[0-9a-f]{40}$/i.test(requested);
|
||||
const branchRef = requested && !isSha ? requested : tracking?.prod?.branch || defaultBranch;
|
||||
|
||||
try {
|
||||
const octokit = await githubApp.getInstallationOctokit(installationId);
|
||||
const sha = isSha
|
||||
? requested!
|
||||
: (await octokit.rest.repos.getBranch({ owner, repo, branch: branchRef })).data.commit.sha;
|
||||
|
||||
const token = await githubApp.octokit.rest.apps.createInstallationAccessToken({
|
||||
installation_id: installationId,
|
||||
repositories: [repo],
|
||||
permissions: { contents: "read" },
|
||||
});
|
||||
|
||||
// Resolve the signed archive URL without downloading the bytes server-side.
|
||||
const redirect = await fetch(`https://api.github.com/repos/${owner}/${repo}/tarball/${sha}`, {
|
||||
headers: {
|
||||
Authorization: `Bearer ${token.data.token}`,
|
||||
Accept: "application/vnd.github+json",
|
||||
"User-Agent": "trigger-dashboard-agent",
|
||||
},
|
||||
redirect: "manual",
|
||||
});
|
||||
const tarballUrl = redirect.headers.get("location");
|
||||
if (!tarballUrl) return null;
|
||||
|
||||
const snapshot: DashboardAgentRepoSnapshot = { tarballUrl, owner, repo, sha, defaultBranch };
|
||||
pruneRepoSnapshotCache();
|
||||
repoSnapshotCache.set(cacheKey, { snapshot, expiresAt: Date.now() + REPO_SNAPSHOT_TTL_MS });
|
||||
return snapshot;
|
||||
} catch (error) {
|
||||
logger.error("Failed to resolve dashboard agent repo snapshot", { error, projectId });
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
// Map a run (by friendly id) to the commit its deployed version came from, for
|
||||
// run-SHA pinning. A run locks to a BackgroundWorker (`lockedToVersionId`), whose
|
||||
// WorkerDeployment carries the commit. Null for runs with no deployed version
|
||||
// (e.g. dev runs), so the agent falls back to the branch head.
|
||||
export async function resolveRunCommit(
|
||||
environmentId: string,
|
||||
runFriendlyId: string
|
||||
): Promise<{ sha: string; version: string; dirty: boolean } | null> {
|
||||
const run = await runStore.findRun(
|
||||
{ friendlyId: runFriendlyId, runtimeEnvironmentId: environmentId },
|
||||
{ select: { lockedToVersionId: true } }
|
||||
);
|
||||
if (!run?.lockedToVersionId) return null;
|
||||
|
||||
const deployment = await prisma.workerDeployment.findFirst({
|
||||
where: { workerId: run.lockedToVersionId },
|
||||
select: { commitSHA: true, version: true, git: true },
|
||||
});
|
||||
if (!deployment?.commitSHA) return null;
|
||||
|
||||
const dirty = (deployment.git as { dirty?: boolean } | null)?.dirty ?? false;
|
||||
return { sha: deployment.commitSHA, version: deployment.version, dirty };
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
import { createDashboardAgentDb, type DashboardAgentDb } from "@internal/dashboard-agent-db";
|
||||
import { env } from "~/env.server";
|
||||
import { singleton } from "~/utils/singleton";
|
||||
|
||||
/**
|
||||
* The webapp's connection to the dashboard-agent conversation store (the History
|
||||
* tab + the chat panel's create / rename / delete / resume actions). Same Drizzle
|
||||
* client the agent task uses, pointed at the same database.
|
||||
*
|
||||
* This is the agent's OWN datastore — NOT the main Prisma database, which the
|
||||
* agent has no access to. Cloud uses the dedicated PlanetScale database; OSS
|
||||
* falls back to DATABASE_URL with tables isolated in the `trigger_dashboard_agent`
|
||||
* schema.
|
||||
*/
|
||||
export const dashboardAgentDb: DashboardAgentDb = singleton("dashboardAgentDb", () => {
|
||||
const connectionString = env.DASHBOARD_AGENT_DATABASE_URL ?? env.DATABASE_URL;
|
||||
return createDashboardAgentDb(connectionString, {
|
||||
max: env.DATABASE_CONNECTION_LIMIT,
|
||||
}).db;
|
||||
});
|
||||
@@ -0,0 +1,72 @@
|
||||
import { createAnthropic } from "@ai-sdk/anthropic";
|
||||
import {
|
||||
DASHBOARD_AGENT_CODE_SYSTEM_PROMPT,
|
||||
DASHBOARD_AGENT_MODEL,
|
||||
DASHBOARD_AGENT_SYSTEM_PROMPT,
|
||||
dashboardAgentCodeToolSchemas,
|
||||
dashboardAgentToolSchemas,
|
||||
} from "@internal/dashboard-agent/tool-schemas";
|
||||
import { chat as chatServer } from "@trigger.dev/sdk/chat-server";
|
||||
import { streamText, type UIMessage } from "ai";
|
||||
import { env } from "~/env.server";
|
||||
import {
|
||||
dashboardAgentApiOrigin,
|
||||
dashboardAgentTriggerConfig,
|
||||
} from "~/services/dashboardAgent.server";
|
||||
import { logger } from "~/services/logger.server";
|
||||
|
||||
const TASK_ID = "dashboard-agent";
|
||||
|
||||
const anthropic = createAnthropic({ apiKey: env.ANTHROPIC_API_KEY });
|
||||
|
||||
/**
|
||||
* Server-owned head start. The webapp generates the chatId and owns the chat
|
||||
* record, then kicks off step 1 here via `chat.startHeadStart` (the detached
|
||||
* flow): it creates the session (externalId = chatId), triggers the
|
||||
* handover-prepare run, and streams step 1 into `session.out` in the background.
|
||||
* The browser resumes that stream rather than streaming step 1 inline. Step 1
|
||||
* runs the agent's SCHEMA-ONLY tools + the shared model/prompt for the mode the
|
||||
* agent run will be in; the agent run picks up tool execution and step 2+.
|
||||
*
|
||||
* `metadata` (the delegated UAT + context) is merged into the run's wire payload
|
||||
* server-side, so it reaches the agent without touching the browser.
|
||||
*/
|
||||
export async function startDashboardAgentHeadStart(params: {
|
||||
chatId: string;
|
||||
messages: UIMessage[];
|
||||
mode: "assistant" | "code";
|
||||
metadata: Record<string, unknown>;
|
||||
}): Promise<void> {
|
||||
const tools = params.mode === "code" ? dashboardAgentCodeToolSchemas : dashboardAgentToolSchemas;
|
||||
const system =
|
||||
params.mode === "code" ? DASHBOARD_AGENT_CODE_SYSTEM_PROMPT : DASHBOARD_AGENT_SYSTEM_PROMPT;
|
||||
|
||||
const { completion } = await chatServer.startHeadStart({
|
||||
agentId: TASK_ID,
|
||||
chatId: params.chatId,
|
||||
messages: params.messages,
|
||||
metadata: params.metadata,
|
||||
triggerConfig: dashboardAgentTriggerConfig(),
|
||||
// Scope session creation + the agent trigger to the agent's project/env. The
|
||||
// Anthropic key here only powers the warm step-1 call.
|
||||
apiClient: {
|
||||
baseURL: dashboardAgentApiOrigin(),
|
||||
accessToken: env.DASHBOARD_AGENT_SECRET_KEY,
|
||||
},
|
||||
run: async ({ chat: helper }) =>
|
||||
streamText({
|
||||
...helper.toStreamTextOptions({ tools }),
|
||||
model: anthropic(DASHBOARD_AGENT_MODEL),
|
||||
system,
|
||||
}),
|
||||
});
|
||||
|
||||
// The webapp is long-lived, so step 1's drain + the handover dispatch run in
|
||||
// the background after this resolves (createSession + trigger have completed).
|
||||
// Log a warm-step failure for observability: startHeadStart has already fired
|
||||
// handover-skip so the agent run exits cleanly, but the client (mounted as
|
||||
// streaming) then resumes an empty session.out, so the turn looks lost.
|
||||
completion.catch((error) => {
|
||||
logger.error("Dashboard agent head start failed", { chatId: params.chatId, error });
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,232 @@
|
||||
import { z } from "zod";
|
||||
import { prisma } from "~/db.server";
|
||||
import { logger } from "./logger.server";
|
||||
import { type UserFromSession } from "./session.server";
|
||||
|
||||
const SideMenuPreferences = z.object({
|
||||
isCollapsed: z.boolean().default(false),
|
||||
// Map for section collapsed states - keys are section identifiers
|
||||
collapsedSections: z.record(z.string(), z.boolean()).optional(),
|
||||
/** Organization-specific settings */
|
||||
organizations: z
|
||||
.record(
|
||||
z.string(),
|
||||
z.object({
|
||||
orderedItems: z.record(z.string(), z.array(z.string())),
|
||||
})
|
||||
)
|
||||
.optional(),
|
||||
});
|
||||
|
||||
export type SideMenuPreferences = z.infer<typeof SideMenuPreferences>;
|
||||
|
||||
import { type SideMenuSectionId } from "~/components/navigation/sideMenuTypes";
|
||||
export type { SideMenuSectionId };
|
||||
|
||||
const DashboardPreferences = z.object({
|
||||
version: z.literal("1"),
|
||||
currentProjectId: z.string().optional(),
|
||||
projects: z.record(
|
||||
z.string(),
|
||||
z.object({
|
||||
currentEnvironment: z.object({ id: z.string() }),
|
||||
})
|
||||
),
|
||||
sideMenu: SideMenuPreferences.optional(),
|
||||
});
|
||||
|
||||
export type DashboardPreferences = z.infer<typeof DashboardPreferences>;
|
||||
|
||||
export function getDashboardPreferences(data?: any | null): DashboardPreferences {
|
||||
if (!data) {
|
||||
return {
|
||||
version: "1",
|
||||
projects: {},
|
||||
};
|
||||
}
|
||||
|
||||
const result = DashboardPreferences.safeParse(data);
|
||||
if (!result.success) {
|
||||
logger.error("Failed to parse DashboardPreferences", { data, error: result.error });
|
||||
return {
|
||||
version: "1",
|
||||
projects: {},
|
||||
};
|
||||
}
|
||||
|
||||
return result.data;
|
||||
}
|
||||
|
||||
export async function updateCurrentProjectEnvironmentId({
|
||||
user,
|
||||
projectId,
|
||||
environmentId,
|
||||
}: {
|
||||
user: UserFromSession;
|
||||
projectId: string;
|
||||
environmentId: string;
|
||||
}) {
|
||||
if (user.isImpersonating) {
|
||||
return;
|
||||
}
|
||||
|
||||
//only update if the existing preferences are different
|
||||
if (
|
||||
user.dashboardPreferences.currentProjectId === projectId &&
|
||||
user.dashboardPreferences.projects[projectId]?.currentEnvironment?.id === environmentId
|
||||
) {
|
||||
return;
|
||||
}
|
||||
|
||||
//ok we need to update the preferences
|
||||
const updatedPreferences: DashboardPreferences = {
|
||||
...user.dashboardPreferences,
|
||||
currentProjectId: projectId,
|
||||
projects: {
|
||||
...user.dashboardPreferences.projects,
|
||||
[projectId]: {
|
||||
...user.dashboardPreferences.projects[projectId],
|
||||
currentEnvironment: { id: environmentId },
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
return prisma.user.update({
|
||||
where: {
|
||||
id: user.id,
|
||||
},
|
||||
data: {
|
||||
dashboardPreferences: updatedPreferences,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
export async function clearCurrentProject({ user }: { user: UserFromSession }) {
|
||||
if (user.isImpersonating) {
|
||||
return;
|
||||
}
|
||||
|
||||
const updatedPreferences: DashboardPreferences = {
|
||||
...user.dashboardPreferences,
|
||||
currentProjectId: undefined,
|
||||
};
|
||||
|
||||
return prisma.user.update({
|
||||
where: {
|
||||
id: user.id,
|
||||
},
|
||||
data: {
|
||||
dashboardPreferences: updatedPreferences,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
export async function updateSideMenuPreferences({
|
||||
user,
|
||||
isCollapsed,
|
||||
sectionCollapsed,
|
||||
}: {
|
||||
user: UserFromSession;
|
||||
isCollapsed?: boolean;
|
||||
/** Update a specific section's collapsed state */
|
||||
sectionCollapsed?: { sectionId: SideMenuSectionId; collapsed: boolean };
|
||||
}) {
|
||||
if (user.isImpersonating) {
|
||||
return;
|
||||
}
|
||||
|
||||
// Parse with schema to apply defaults, then overlay any new values
|
||||
const currentSideMenu = SideMenuPreferences.parse(user.dashboardPreferences.sideMenu ?? {});
|
||||
|
||||
// Build the updated collapsedSections map
|
||||
let updatedCollapsedSections = { ...currentSideMenu.collapsedSections };
|
||||
|
||||
if (sectionCollapsed) {
|
||||
updatedCollapsedSections[sectionCollapsed.sectionId] = sectionCollapsed.collapsed;
|
||||
}
|
||||
|
||||
const updatedSideMenu = SideMenuPreferences.parse({
|
||||
...currentSideMenu,
|
||||
...(isCollapsed !== undefined && { isCollapsed }),
|
||||
collapsedSections: updatedCollapsedSections,
|
||||
});
|
||||
|
||||
// Only update if something changed
|
||||
const hasCollapsedSectionsChanged =
|
||||
JSON.stringify(updatedSideMenu.collapsedSections) !==
|
||||
JSON.stringify(currentSideMenu.collapsedSections);
|
||||
|
||||
if (updatedSideMenu.isCollapsed === currentSideMenu.isCollapsed && !hasCollapsedSectionsChanged) {
|
||||
return;
|
||||
}
|
||||
|
||||
const updatedPreferences: DashboardPreferences = {
|
||||
...user.dashboardPreferences,
|
||||
sideMenu: updatedSideMenu,
|
||||
};
|
||||
|
||||
return prisma.user.update({
|
||||
where: {
|
||||
id: user.id,
|
||||
},
|
||||
data: {
|
||||
dashboardPreferences: updatedPreferences,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
/** Get the stored item order for a specific list within an organization */
|
||||
export function getItemOrder(
|
||||
sideMenu: SideMenuPreferences | undefined,
|
||||
organizationId: string,
|
||||
listId: string
|
||||
): string[] | undefined {
|
||||
return sideMenu?.organizations?.[organizationId]?.orderedItems?.[listId];
|
||||
}
|
||||
|
||||
export async function updateItemOrder({
|
||||
user,
|
||||
organizationId,
|
||||
listId,
|
||||
order,
|
||||
}: {
|
||||
user: UserFromSession;
|
||||
organizationId: string;
|
||||
listId: string;
|
||||
order: string[];
|
||||
}) {
|
||||
if (user.isImpersonating) {
|
||||
return;
|
||||
}
|
||||
|
||||
const currentSideMenu = SideMenuPreferences.parse(user.dashboardPreferences.sideMenu ?? {});
|
||||
const currentOrg = currentSideMenu.organizations?.[organizationId];
|
||||
|
||||
const updatedSideMenu = SideMenuPreferences.parse({
|
||||
...currentSideMenu,
|
||||
organizations: {
|
||||
...currentSideMenu.organizations,
|
||||
[organizationId]: {
|
||||
...currentOrg,
|
||||
orderedItems: {
|
||||
...currentOrg?.orderedItems,
|
||||
[listId]: order,
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
const updatedPreferences: DashboardPreferences = {
|
||||
...user.dashboardPreferences,
|
||||
sideMenu: updatedSideMenu,
|
||||
};
|
||||
|
||||
return prisma.user.update({
|
||||
where: {
|
||||
id: user.id,
|
||||
},
|
||||
data: {
|
||||
dashboardPreferences: updatedPreferences,
|
||||
},
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,39 @@
|
||||
import { z } from "zod";
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// ClickHouse config (kind = CLICKHOUSE)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/** V1: single secret-store key that supplies the ClickHouse connection URL. */
|
||||
export const ClickhouseDataStoreConfigV1 = z.object({
|
||||
version: z.literal(1),
|
||||
data: z.object({
|
||||
/** Key into the SecretStore that resolves to a ClickhouseConnection ({url}). */
|
||||
secretKey: z.string(),
|
||||
}),
|
||||
});
|
||||
|
||||
export type ClickhouseDataStoreConfigV1 = z.infer<typeof ClickhouseDataStoreConfigV1>;
|
||||
|
||||
/** Discriminated union over version — extend by adding new literals here. */
|
||||
export const ClickhouseDataStoreConfig = z.discriminatedUnion("version", [
|
||||
ClickhouseDataStoreConfigV1,
|
||||
]);
|
||||
|
||||
export type ClickhouseDataStoreConfig = z.infer<typeof ClickhouseDataStoreConfig>;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Top-level per-kind union
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/**
|
||||
* Secrets are resolved to URLs at registry load time so the factory never
|
||||
* needs to touch the secret store on the hot path.
|
||||
*/
|
||||
export type ParsedClickhouseDataStore = {
|
||||
kind: "CLICKHOUSE";
|
||||
url: string;
|
||||
};
|
||||
|
||||
/** Union of all parsed data store types. Extend as new DataStoreKind values are added. */
|
||||
export type ParsedDataStore = ParsedClickhouseDataStore;
|
||||
@@ -0,0 +1,201 @@
|
||||
import type { DataStoreKind, PrismaClient, PrismaReplicaClient } from "@trigger.dev/database";
|
||||
import {
|
||||
ClickhouseDataStoreConfig,
|
||||
type ParsedDataStore,
|
||||
} from "./organizationDataStoreConfigSchemas.server";
|
||||
import { getSecretStore } from "../secrets/secretStore.server";
|
||||
import { ClickhouseConnectionSchema } from "../clickhouse/clickhouseSecretSchemas.server";
|
||||
|
||||
export class OrganizationDataStoresRegistry {
|
||||
/**
|
||||
* Writer client — used by every method that mutates state
|
||||
* (`addDataStore` / `updateDataStore` / `deleteDataStore` and their backing
|
||||
* SecretStore writes). Must be the primary connection; replica-targeted
|
||||
* writes are rejected by Postgres with code 25006 (read-only transaction).
|
||||
*/
|
||||
private _writer: PrismaClient;
|
||||
/**
|
||||
* Read client used by the polling `loadFromDatabase()` (and its
|
||||
* `SecretStore.getSecret` lookups). Can be a replica — these are
|
||||
* cache-fillers, not on hot user-facing paths.
|
||||
*/
|
||||
private _replica: PrismaClient | PrismaReplicaClient;
|
||||
/** Keyed by `${organizationId}:${kind}` */
|
||||
private _lookup: Map<string, ParsedDataStore> = new Map();
|
||||
private _loaded = false;
|
||||
private _readyResolve!: () => void;
|
||||
|
||||
/**
|
||||
* Resolves once the initial `loadFromDatabase()` completes successfully.
|
||||
* At process startup the singleton loads the registry with unbounded retries
|
||||
* (exponential backoff, capped delay) until Postgres is reachable; until then
|
||||
* this promise stays pending and callers that await readiness will block.
|
||||
*/
|
||||
readonly isReady: Promise<void>;
|
||||
|
||||
constructor(writer: PrismaClient, replica: PrismaClient | PrismaReplicaClient) {
|
||||
this._writer = writer;
|
||||
this._replica = replica;
|
||||
this.isReady = new Promise<void>((resolve) => {
|
||||
this._readyResolve = resolve;
|
||||
});
|
||||
}
|
||||
|
||||
get isLoaded(): boolean {
|
||||
return this._loaded;
|
||||
}
|
||||
|
||||
async loadFromDatabase(): Promise<void> {
|
||||
// Sort by `key` (unique, immutable) to ensure a deterministic winner when the
|
||||
// same `${orgId}:${kind}` appears in multiple rows. The registry must never
|
||||
// throw on overlap — failing the load would break every customer, not just the
|
||||
// misconfigured orgs — so we keep the first entry and log an error instead.
|
||||
const rows = await this._replica.organizationDataStore.findMany({
|
||||
orderBy: { key: "asc" },
|
||||
});
|
||||
const secretStore = getSecretStore("DATABASE", { prismaClient: this._replica });
|
||||
|
||||
const lookup = new Map<string, ParsedDataStore>();
|
||||
/** Tracks which row's `key` already owns each `${orgId}:${kind}` so we can log conflicts. */
|
||||
const winnerByLookupKey = new Map<string, string>();
|
||||
|
||||
for (const row of rows) {
|
||||
let parsed: ParsedDataStore | null = null;
|
||||
|
||||
switch (row.kind) {
|
||||
case "CLICKHOUSE": {
|
||||
const result = ClickhouseDataStoreConfig.safeParse(row.config);
|
||||
if (!result.success) {
|
||||
console.warn(
|
||||
`[OrganizationDataStoresRegistry] Invalid config for OrganizationDataStore "${row.key}" (kind=CLICKHOUSE): ${result.error.message}`
|
||||
);
|
||||
continue;
|
||||
}
|
||||
|
||||
const connection = await secretStore.getSecret(
|
||||
ClickhouseConnectionSchema,
|
||||
result.data.data.secretKey
|
||||
);
|
||||
|
||||
if (!connection) {
|
||||
console.warn(
|
||||
`[OrganizationDataStoresRegistry] Secret "${result.data.data.secretKey}" not found for OrganizationDataStore "${row.key}" — skipping`
|
||||
);
|
||||
continue;
|
||||
}
|
||||
|
||||
parsed = { kind: "CLICKHOUSE", url: connection.url };
|
||||
break;
|
||||
}
|
||||
default: {
|
||||
console.warn(
|
||||
`[OrganizationDataStoresRegistry] Unknown kind "${row.kind}" for OrganizationDataStore "${row.key}" — skipping`
|
||||
);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
for (const orgId of row.organizationIds) {
|
||||
const lookupKey = `${orgId}:${row.kind}`;
|
||||
const existingWinner = winnerByLookupKey.get(lookupKey);
|
||||
if (existingWinner) {
|
||||
console.error(
|
||||
`[OrganizationDataStoresRegistry] Overlapping OrganizationDataStore assignment for orgId="${orgId}" kind=${row.kind}: already routed to "${existingWinner}", ignoring "${row.key}". Pick one store per (org, kind) to resolve.`
|
||||
);
|
||||
continue;
|
||||
}
|
||||
winnerByLookupKey.set(lookupKey, row.key);
|
||||
lookup.set(lookupKey, parsed);
|
||||
}
|
||||
}
|
||||
|
||||
this._lookup = lookup;
|
||||
|
||||
if (!this._loaded) {
|
||||
this._loaded = true;
|
||||
this._readyResolve();
|
||||
}
|
||||
}
|
||||
|
||||
async reload(): Promise<void> {
|
||||
await this.loadFromDatabase();
|
||||
}
|
||||
|
||||
#secretKey(key: string, kind: DataStoreKind) {
|
||||
return `data-store:${key}:${kind.toLocaleLowerCase()}`;
|
||||
}
|
||||
|
||||
async addDataStore({
|
||||
key,
|
||||
kind,
|
||||
organizationIds,
|
||||
config,
|
||||
}: {
|
||||
key: string;
|
||||
kind: DataStoreKind;
|
||||
organizationIds: string[];
|
||||
config: any;
|
||||
}) {
|
||||
const secretKey = this.#secretKey(key, kind);
|
||||
|
||||
const secretStore = getSecretStore("DATABASE", { prismaClient: this._writer });
|
||||
await secretStore.setSecret(secretKey, config);
|
||||
|
||||
return this._writer.organizationDataStore.create({
|
||||
data: {
|
||||
key,
|
||||
organizationIds,
|
||||
kind,
|
||||
config: { version: 1, data: { secretKey } },
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
async updateDataStore({
|
||||
key,
|
||||
kind,
|
||||
organizationIds,
|
||||
config,
|
||||
}: {
|
||||
key: string;
|
||||
kind: DataStoreKind;
|
||||
organizationIds: string[];
|
||||
config?: any;
|
||||
}) {
|
||||
const secretKey = this.#secretKey(key, kind);
|
||||
|
||||
if (config) {
|
||||
const secretStore = getSecretStore("DATABASE", { prismaClient: this._writer });
|
||||
await secretStore.setSecret(secretKey, config);
|
||||
}
|
||||
|
||||
return this._writer.organizationDataStore.update({
|
||||
where: {
|
||||
key,
|
||||
},
|
||||
data: {
|
||||
organizationIds,
|
||||
kind: "CLICKHOUSE",
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
async deleteDataStore({ key, kind }: { key: string; kind: DataStoreKind }) {
|
||||
const secretKey = this.#secretKey(key, kind);
|
||||
const secretStore = getSecretStore("DATABASE", { prismaClient: this._writer });
|
||||
await secretStore.deleteSecret(secretKey).catch(() => {
|
||||
// Secret may not exist — proceed with deletion
|
||||
});
|
||||
|
||||
await this._writer.organizationDataStore.delete({ where: { key } });
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns the parsed data store config for the given organization and kind,
|
||||
* or `null` if no override is configured (caller should use the default).
|
||||
*/
|
||||
get(organizationId: string, kind: DataStoreKind): ParsedDataStore | null {
|
||||
if (!this._loaded) return null;
|
||||
return this._lookup.get(`${organizationId}:${kind}`) ?? null;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,42 @@
|
||||
import pRetry from "p-retry";
|
||||
import { $replica, prisma } from "~/db.server";
|
||||
import { env } from "~/env.server";
|
||||
import { logger } from "~/services/logger.server";
|
||||
import { signalsEmitter } from "~/services/signals.server";
|
||||
import { singleton } from "~/utils/singleton";
|
||||
import { OrganizationDataStoresRegistry } from "./organizationDataStoresRegistry.server";
|
||||
|
||||
export const organizationDataStoresRegistry = singleton("organizationDataStoresRegistry", () => {
|
||||
const registry = new OrganizationDataStoresRegistry(prisma, $replica);
|
||||
|
||||
// Runs as soon as this singleton is created (first import of this module). The
|
||||
// registry’s `isReady` promise resolves when this eventually succeeds.
|
||||
const startupLoadPromise = pRetry(() => registry.loadFromDatabase(), {
|
||||
forever: true,
|
||||
retries: 10,
|
||||
minTimeout: 1_000,
|
||||
maxTimeout: 60_000,
|
||||
factor: 2,
|
||||
onFailedAttempt: (error) => {
|
||||
logger.warn("[OrganizationDataStoresRegistry] Startup load failed, retrying", {
|
||||
attemptNumber: error.attemptNumber,
|
||||
retriesLeft: error.retriesLeft,
|
||||
error: error.message,
|
||||
});
|
||||
},
|
||||
});
|
||||
startupLoadPromise.catch((err) => {
|
||||
console.error("[OrganizationDataStoresRegistry] Unexpected startup load failure", err);
|
||||
});
|
||||
|
||||
const interval = setInterval(() => {
|
||||
registry.reload().catch((err) => {
|
||||
console.error("[OrganizationDataStoresRegistry] Failed to reload", err);
|
||||
});
|
||||
}, env.ORGANIZATION_DATA_STORES_RELOAD_INTERVAL_MS);
|
||||
|
||||
signalsEmitter.on("SIGTERM", () => clearInterval(interval));
|
||||
signalsEmitter.on("SIGINT", () => clearInterval(interval));
|
||||
|
||||
return registry;
|
||||
});
|
||||
@@ -0,0 +1,96 @@
|
||||
import { runMigrations } from "graphile-worker";
|
||||
import type { PrismaClient } from "~/db.server";
|
||||
import { prisma } from "~/db.server";
|
||||
import { env } from "~/env.server";
|
||||
import { logger } from "~/services/logger.server";
|
||||
import { PgNotifyService } from "./pgNotify.server";
|
||||
import { z } from "zod";
|
||||
|
||||
export class GraphileMigrationHelperService {
|
||||
#prismaClient: PrismaClient;
|
||||
|
||||
constructor(prismaClient: PrismaClient = prisma) {
|
||||
this.#prismaClient = prismaClient;
|
||||
}
|
||||
|
||||
public async call() {
|
||||
this.#logDebug("GraphileMigrationHelperService.call");
|
||||
|
||||
await this.#detectAndPrepareForMigrations();
|
||||
|
||||
await runMigrations({
|
||||
connectionString: env.DATABASE_URL,
|
||||
schema: env.WORKER_SCHEMA,
|
||||
});
|
||||
}
|
||||
|
||||
#logDebug(message: string, args?: any) {
|
||||
logger.debug(`[migrationHelper] ${message}`, args);
|
||||
}
|
||||
|
||||
async #getLatestMigration() {
|
||||
const migrationQueryResult = await this.#prismaClient.$queryRawUnsafe(`
|
||||
SELECT id FROM ${env.WORKER_SCHEMA}.migrations
|
||||
ORDER BY id DESC LIMIT 1
|
||||
`);
|
||||
|
||||
const MigrationQueryResultSchema = z.array(z.object({ id: z.number() }));
|
||||
|
||||
const migrationResults = MigrationQueryResultSchema.parse(migrationQueryResult);
|
||||
|
||||
if (!migrationResults.length) {
|
||||
// no migrations applied yet
|
||||
return -1;
|
||||
}
|
||||
|
||||
return migrationResults[0].id;
|
||||
}
|
||||
|
||||
async #graphileSchemaExists() {
|
||||
const schemaCount = await this.#prismaClient.$executeRaw`
|
||||
SELECT schema_name FROM information_schema.schemata
|
||||
WHERE schema_name = ${env.WORKER_SCHEMA}
|
||||
`;
|
||||
|
||||
return schemaCount === 1;
|
||||
}
|
||||
|
||||
/** Helper for graphile-worker v0.14.0 migration. No-op if already migrated. */
|
||||
async #detectAndPrepareForMigrations() {
|
||||
if (!(await this.#graphileSchemaExists())) {
|
||||
// no schema yet, likely first start
|
||||
return;
|
||||
}
|
||||
|
||||
const latestMigration = await this.#getLatestMigration();
|
||||
|
||||
if (latestMigration < 0) {
|
||||
// no migrations found
|
||||
return;
|
||||
}
|
||||
|
||||
// the first v0.14.0 migration has ID 11
|
||||
if (latestMigration > 10) {
|
||||
// already migrated
|
||||
return;
|
||||
}
|
||||
|
||||
// add 15s to graceful shutdown timeout, just to be safe
|
||||
const migrationDelayInMs = env.GRACEFUL_SHUTDOWN_TIMEOUT + 15000;
|
||||
|
||||
this.#logDebug("Delaying worker startup due to pending migration", {
|
||||
latestMigration,
|
||||
migrationDelayInMs,
|
||||
});
|
||||
|
||||
console.log(`⚠️ detected pending graphile migration`);
|
||||
console.log(`⚠️ notifying running workers`);
|
||||
|
||||
const pgNotify = new PgNotifyService();
|
||||
await pgNotify.call("trigger:graphile:migrate", { latestMigration });
|
||||
|
||||
console.log(`⚠️ delaying worker startup by ${migrationDelayInMs}ms`);
|
||||
|
||||
await new Promise((resolve) => setTimeout(resolve, migrationDelayInMs));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
import type { z } from "zod";
|
||||
import type { PrismaClient } from "~/db.server";
|
||||
import { prisma } from "~/db.server";
|
||||
import { logger } from "~/services/logger.server";
|
||||
import type { NotificationCatalog, NotificationChannel } from "./types";
|
||||
|
||||
export class PgNotifyService {
|
||||
#prismaClient: PrismaClient;
|
||||
|
||||
constructor(prismaClient: PrismaClient = prisma) {
|
||||
this.#prismaClient = prismaClient;
|
||||
}
|
||||
|
||||
public async call<TChannel extends NotificationChannel>(
|
||||
channelName: TChannel,
|
||||
payload: z.infer<NotificationCatalog[TChannel]>
|
||||
) {
|
||||
this.#logDebug("Sending notification", { channelName, notifyPayload: payload });
|
||||
|
||||
await this.#prismaClient.$executeRaw`
|
||||
SELECT pg_notify(${channelName}, ${JSON.stringify(payload)})
|
||||
`;
|
||||
}
|
||||
|
||||
#logDebug(message: string, args?: any) {
|
||||
logger.debug(`[pgNotify] ${message}`, args);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
import { z } from "zod";
|
||||
|
||||
export const notificationCatalog = {
|
||||
"trigger:graphile:migrate": z.object({
|
||||
latestMigration: z.number(),
|
||||
}),
|
||||
};
|
||||
|
||||
export type NotificationCatalog = typeof notificationCatalog;
|
||||
|
||||
export type NotificationChannel = keyof NotificationCatalog;
|
||||
@@ -0,0 +1,90 @@
|
||||
import { DateFormatter } from "@internationalized/date";
|
||||
import type { PrismaClient } from "@trigger.dev/database";
|
||||
import { prisma } from "~/db.server";
|
||||
import { featuresForRequest } from "~/features.server";
|
||||
import { DeleteProjectService } from "./deleteProject.server";
|
||||
import { getCurrentPlan } from "./platform.v3.server";
|
||||
import { controlPlaneResolver } from "~/v3/runOpsMigration/controlPlaneResolver.server";
|
||||
|
||||
export class DeleteOrganizationService {
|
||||
#prismaClient: PrismaClient;
|
||||
|
||||
constructor(prismaClient: PrismaClient = prisma) {
|
||||
this.#prismaClient = prismaClient;
|
||||
}
|
||||
|
||||
public async call({
|
||||
organizationSlug,
|
||||
userId,
|
||||
request,
|
||||
}: {
|
||||
organizationSlug: string;
|
||||
userId: string;
|
||||
request: Request;
|
||||
}) {
|
||||
const organization = await this.#prismaClient.organization.findFirst({
|
||||
include: {
|
||||
projects: true,
|
||||
members: true,
|
||||
},
|
||||
where: {
|
||||
slug: organizationSlug,
|
||||
members: { some: { userId: userId } },
|
||||
},
|
||||
});
|
||||
|
||||
if (!organization) {
|
||||
throw new Error("Organization not found");
|
||||
}
|
||||
|
||||
if (organization.deletedAt) {
|
||||
throw new Error("Organization already deleted");
|
||||
}
|
||||
|
||||
//Check if they have an active subscription
|
||||
const { isManagedCloud } = featuresForRequest(request);
|
||||
const currentPlan = isManagedCloud ? await getCurrentPlan(organization.id) : undefined;
|
||||
|
||||
if (currentPlan && currentPlan.v3Subscription && currentPlan.v3Subscription.isPaying) {
|
||||
//they've cancelled and that date hasn't passed yet
|
||||
if (
|
||||
currentPlan.v3Subscription.canceledAt &&
|
||||
new Date(currentPlan.v3Subscription.canceledAt) > new Date()
|
||||
) {
|
||||
//a dateformatter that produces results like "Jan 1 2024"
|
||||
const dateFormatter = new DateFormatter("en-us", {
|
||||
year: "numeric",
|
||||
month: "short",
|
||||
day: "numeric",
|
||||
});
|
||||
throw new Error(
|
||||
`This Organization has a canceled subscription. You can delete it when the cancelation date (${dateFormatter.format(
|
||||
new Date(currentPlan.v3Subscription.canceledAt)
|
||||
)}) is in the past.`
|
||||
);
|
||||
}
|
||||
|
||||
throw new Error("You can't delete an Organization that has an active subscription");
|
||||
}
|
||||
|
||||
// loop through the projects and delete them
|
||||
const projectDeleteService = new DeleteProjectService();
|
||||
for (const project of organization.projects) {
|
||||
await projectDeleteService.call({ projectId: project.id, userId });
|
||||
}
|
||||
|
||||
//mark the organization as deleted
|
||||
await this.#prismaClient.organization.update({
|
||||
where: {
|
||||
id: organization.id,
|
||||
},
|
||||
data: {
|
||||
runsEnabled: false,
|
||||
deletedAt: new Date(),
|
||||
},
|
||||
});
|
||||
|
||||
// runsEnabled + the org's projects (project.deletedAt) changed; drop all cached env rows.
|
||||
controlPlaneResolver.invalidateOrganization(organization.id);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,89 @@
|
||||
import type { PrismaClient } from "@trigger.dev/database";
|
||||
import { prisma } from "~/db.server";
|
||||
import { marqs } from "~/v3/marqs/index.server";
|
||||
import { engine } from "~/v3/runEngine.server";
|
||||
import { controlPlaneResolver } from "~/v3/runOpsMigration/controlPlaneResolver.server";
|
||||
|
||||
type Options = ({ projectId: string } | { projectSlug: string }) & {
|
||||
userId: string;
|
||||
};
|
||||
|
||||
export class DeleteProjectService {
|
||||
#prismaClient: PrismaClient;
|
||||
|
||||
constructor(prismaClient: PrismaClient = prisma) {
|
||||
this.#prismaClient = prismaClient;
|
||||
}
|
||||
|
||||
public async call(options: Options) {
|
||||
const projectId = await this.#getProjectId(options);
|
||||
const project = await this.#prismaClient.project.findFirst({
|
||||
include: {
|
||||
environments: true,
|
||||
organization: true,
|
||||
},
|
||||
where: {
|
||||
id: projectId,
|
||||
organization: { members: { some: { userId: options.userId } } },
|
||||
},
|
||||
});
|
||||
|
||||
if (!project) {
|
||||
throw new Error("Project not found");
|
||||
}
|
||||
|
||||
if (project.deletedAt) {
|
||||
return;
|
||||
}
|
||||
|
||||
// Remove queues from MARQS
|
||||
for (const environment of project.environments) {
|
||||
await marqs?.removeEnvironmentQueuesFromMasterQueue(project.organization.id, environment.id);
|
||||
}
|
||||
|
||||
// Delete all queues from the RunEngine 2 prod master queues
|
||||
for (const environment of project.environments) {
|
||||
await engine.removeEnvironmentQueuesFromMasterQueue({
|
||||
runtimeEnvironmentId: environment.id,
|
||||
organizationId: project.organization.id,
|
||||
projectId: project.id,
|
||||
});
|
||||
}
|
||||
|
||||
// Soft delete only: run-ops rows are intentionally retained (no hard-delete cascade here).
|
||||
|
||||
// Mark the project as deleted (do this last because it makes it impossible to try again)
|
||||
// - This disables all API keys
|
||||
// - This disables all schedules from being scheduled
|
||||
await this.#prismaClient.project.update({
|
||||
where: {
|
||||
id: project.id,
|
||||
},
|
||||
data: {
|
||||
deletedAt: new Date(),
|
||||
},
|
||||
});
|
||||
|
||||
// project.deletedAt (which gates env resolution) changed; drop every cached env of this project.
|
||||
for (const environment of project.environments) {
|
||||
controlPlaneResolver.invalidateEnvironment(environment.id);
|
||||
}
|
||||
}
|
||||
|
||||
async #getProjectId(options: Options) {
|
||||
if ("projectId" in options) {
|
||||
return options.projectId;
|
||||
}
|
||||
|
||||
const { id } = await this.#prismaClient.project.findFirstOrThrow({
|
||||
select: {
|
||||
id: true,
|
||||
},
|
||||
where: {
|
||||
slug: options.projectSlug,
|
||||
},
|
||||
});
|
||||
|
||||
return id;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,152 @@
|
||||
import type { DirectorySyncEffect } from "@trigger.dev/plugins";
|
||||
import { prisma } from "~/db.server";
|
||||
import { logger } from "~/services/logger.server";
|
||||
import { rbac } from "~/services/rbac.server";
|
||||
import {
|
||||
ensureOrgMember,
|
||||
ensureUserForDirectory,
|
||||
removeOrgMemberForDirectory,
|
||||
} from "~/models/orgMember.server";
|
||||
import { createPlatformNotification } from "~/services/platformNotifications.server";
|
||||
|
||||
const LAST_OWNER_NOTIFICATION_TITLE = "Directory Sync: last Owner protected";
|
||||
|
||||
// Effects are idempotent and the worker retries, so a single failed attempt is
|
||||
// transient (usually a serializable-conflict retry), not an alert. `logLevel`
|
||||
// makes the worker log it at warn instead of paging.
|
||||
function retryableEffectError(message: string): Error {
|
||||
return Object.assign(new Error(message), { logLevel: "warn" as const });
|
||||
}
|
||||
|
||||
// Deduped notification when the directory tried to remove the org's last Owner:
|
||||
// we keep the member, and one undismissed notification is enough (no retry spam).
|
||||
async function notifyLastOwnerProtected(userId: string, organizationId: string): Promise<void> {
|
||||
const existing = await prisma.platformNotification.findFirst({
|
||||
where: {
|
||||
scope: "USER",
|
||||
userId,
|
||||
surface: "WEBAPP",
|
||||
title: LAST_OWNER_NOTIFICATION_TITLE,
|
||||
archivedAt: null,
|
||||
endsAt: { gt: new Date() },
|
||||
},
|
||||
select: { id: true, interactions: { where: { userId }, select: { webappDismissedAt: true } } },
|
||||
});
|
||||
if (existing && !existing.interactions[0]?.webappDismissedAt) {
|
||||
return;
|
||||
}
|
||||
|
||||
const endsAt = new Date();
|
||||
endsAt.setFullYear(endsAt.getFullYear() + 1);
|
||||
|
||||
const result = await createPlatformNotification({
|
||||
title: LAST_OWNER_NOTIFICATION_TITLE,
|
||||
surface: "WEBAPP",
|
||||
scope: "USER",
|
||||
userId,
|
||||
endsAt: endsAt.toISOString(),
|
||||
priority: 10,
|
||||
payload: {
|
||||
version: "1",
|
||||
data: {
|
||||
type: "card",
|
||||
title: "Directory Sync kept your Owner access",
|
||||
description:
|
||||
"Your identity provider tried to remove you from this organization, but you are its only Owner. " +
|
||||
"We kept your membership to prevent a lockout. Assign another Owner, then the directory change will apply.",
|
||||
},
|
||||
},
|
||||
});
|
||||
if (result.isErr()) {
|
||||
logger.warn("directorySync: failed to create last-owner notification", {
|
||||
userId,
|
||||
organizationId,
|
||||
error: result.error,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
async function applyEffect(effect: DirectorySyncEffect): Promise<void> {
|
||||
switch (effect.kind) {
|
||||
case "provision": {
|
||||
const userId =
|
||||
effect.userId ??
|
||||
(
|
||||
await ensureUserForDirectory({
|
||||
email: effect.email,
|
||||
firstName: effect.firstName,
|
||||
lastName: effect.lastName,
|
||||
})
|
||||
).userId;
|
||||
|
||||
await ensureOrgMember({
|
||||
userId,
|
||||
organizationId: effect.organizationId,
|
||||
roleId: effect.roleId,
|
||||
source: "directory_sync",
|
||||
});
|
||||
|
||||
// Directory owns the role: overwrite even an existing member
|
||||
// (ensureOrgMember only sets it on create).
|
||||
if (effect.roleId) {
|
||||
const result = await rbac.setUserRole({
|
||||
userId,
|
||||
organizationId: effect.organizationId,
|
||||
roleId: effect.roleId,
|
||||
});
|
||||
if (!result.ok) {
|
||||
// The org must keep one Owner: skip the role overwrite for the last
|
||||
// Owner (they keep Owner) instead of failing the whole batch. Applies
|
||||
// to a directory burst and to a dashboard group remap alike.
|
||||
if (result.code === "last_owner") {
|
||||
logger.info("directorySync: kept last Owner, skipped provision role overwrite", {
|
||||
userId,
|
||||
organizationId: effect.organizationId,
|
||||
});
|
||||
} else {
|
||||
throw retryableEffectError(
|
||||
`directorySync provision setUserRole failed: ${result.error}`
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
case "set_role": {
|
||||
const result = await rbac.setUserRole({
|
||||
userId: effect.userId,
|
||||
organizationId: effect.organizationId,
|
||||
roleId: effect.roleId,
|
||||
});
|
||||
if (!result.ok) {
|
||||
// Keeping the org's last Owner is expected, not a failure — skip this
|
||||
// one member and let the rest of the remap apply (no server error).
|
||||
if (result.code === "last_owner") {
|
||||
logger.info("directorySync: kept last Owner, skipped set_role", {
|
||||
userId: effect.userId,
|
||||
organizationId: effect.organizationId,
|
||||
});
|
||||
return;
|
||||
}
|
||||
throw retryableEffectError(`directorySync set_role failed: ${result.error}`);
|
||||
}
|
||||
return;
|
||||
}
|
||||
case "deprovision": {
|
||||
const outcome = await removeOrgMemberForDirectory({
|
||||
userId: effect.userId,
|
||||
organizationId: effect.organizationId,
|
||||
});
|
||||
if (!outcome.removed && outcome.reason === "last_owner_protected") {
|
||||
await notifyLastOwnerProtected(effect.userId, effect.organizationId);
|
||||
}
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
export async function applyDirectorySyncEffects(effects: DirectorySyncEffect[]): Promise<void> {
|
||||
for (const effect of effects) {
|
||||
await applyEffect(effect);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,100 @@
|
||||
import type { DeliverEmail, SendPlainTextOptions, MailTransportOptions } from "emails";
|
||||
import { EmailClient } from "emails";
|
||||
import type { SendEmailOptions } from "remix-auth-email-link";
|
||||
import { redirect } from "remix-typedjson";
|
||||
import { env } from "~/env.server";
|
||||
import type { AuthUser } from "./authUser";
|
||||
import { logger } from "./logger.server";
|
||||
import { singleton } from "~/utils/singleton";
|
||||
import { assertEmailAllowed } from "~/utils/email";
|
||||
|
||||
const client = singleton(
|
||||
"email-client",
|
||||
() =>
|
||||
new EmailClient({
|
||||
transport: buildTransportOptions(),
|
||||
imagesBaseUrl: env.APP_ORIGIN,
|
||||
from: env.FROM_EMAIL ?? "team@email.trigger.dev",
|
||||
replyTo: env.REPLY_TO_EMAIL ?? "help@email.trigger.dev",
|
||||
})
|
||||
);
|
||||
|
||||
const alertsClient = singleton(
|
||||
"alerts-email-client",
|
||||
() =>
|
||||
new EmailClient({
|
||||
transport: buildTransportOptions(true),
|
||||
imagesBaseUrl: env.APP_ORIGIN,
|
||||
from: env.ALERT_FROM_EMAIL ?? "noreply@alerts.trigger.dev",
|
||||
// Fallback to `REPLY_TO_EMAIL` for backwards compat
|
||||
replyTo: env.ALERT_REPLY_TO_EMAIL ?? env.REPLY_TO_EMAIL ?? "help@email.trigger.dev",
|
||||
})
|
||||
);
|
||||
|
||||
function buildTransportOptions(alerts?: boolean): MailTransportOptions {
|
||||
const transportType = alerts ? env.ALERT_EMAIL_TRANSPORT : env.EMAIL_TRANSPORT;
|
||||
logger.debug(
|
||||
`Constructing email transport '${transportType}' for usage '${alerts ? "alerts" : "general"}'`
|
||||
);
|
||||
|
||||
switch (transportType) {
|
||||
case "aws-ses":
|
||||
return { type: "aws-ses" };
|
||||
case "resend":
|
||||
return {
|
||||
type: "resend",
|
||||
config: {
|
||||
apiKey: alerts ? env.ALERT_RESEND_API_KEY : env.RESEND_API_KEY,
|
||||
},
|
||||
};
|
||||
case "smtp":
|
||||
return {
|
||||
type: "smtp",
|
||||
config: {
|
||||
host: alerts ? env.ALERT_SMTP_HOST : env.SMTP_HOST,
|
||||
port: alerts ? env.ALERT_SMTP_PORT : env.SMTP_PORT,
|
||||
secure: alerts ? env.ALERT_SMTP_SECURE : env.SMTP_SECURE,
|
||||
auth: {
|
||||
user: alerts ? env.ALERT_SMTP_USER : env.SMTP_USER,
|
||||
pass: alerts ? env.ALERT_SMTP_PASSWORD : env.SMTP_PASSWORD,
|
||||
},
|
||||
},
|
||||
};
|
||||
default:
|
||||
return { type: undefined };
|
||||
}
|
||||
}
|
||||
|
||||
export async function sendMagicLinkEmail(options: SendEmailOptions<AuthUser>): Promise<void> {
|
||||
assertEmailAllowed(options.emailAddress);
|
||||
|
||||
// Auto redirect when in development mode
|
||||
if (env.NODE_ENV === "development") {
|
||||
throw redirect(options.magicLink);
|
||||
}
|
||||
|
||||
logger.debug("Sending magic link email", { emailAddress: options.emailAddress });
|
||||
|
||||
try {
|
||||
return await client.send({
|
||||
email: "magic_link",
|
||||
to: options.emailAddress,
|
||||
magicLink: options.magicLink,
|
||||
});
|
||||
} catch (error) {
|
||||
logger.error("Error sending magic link email", { error: JSON.stringify(error) });
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
export async function sendPlainTextEmail(options: SendPlainTextOptions) {
|
||||
return client.sendPlainText(options);
|
||||
}
|
||||
|
||||
export async function sendEmail(data: DeliverEmail) {
|
||||
return client.send(data);
|
||||
}
|
||||
|
||||
export async function sendAlertEmail(data: DeliverEmail) {
|
||||
return alertsClient.send(data);
|
||||
}
|
||||
@@ -0,0 +1,61 @@
|
||||
import type { Authenticator } from "remix-auth";
|
||||
import { EmailLinkStrategy } from "remix-auth-email-link";
|
||||
import { env } from "~/env.server";
|
||||
import { findOrCreateUser } from "~/models/user.server";
|
||||
import { sendMagicLinkEmail } from "~/services/email.server";
|
||||
import type { AuthUser } from "./authUser";
|
||||
import { logger } from "./logger.server";
|
||||
|
||||
import { postAuthentication } from "./postAuth.server";
|
||||
import { SsoRequiredError, ssoRedirectForEmail } from "./ssoAutoDiscovery.server";
|
||||
|
||||
let secret = env.MAGIC_LINK_SECRET;
|
||||
if (!secret) throw new Error("Missing MAGIC_LINK_SECRET env variable.");
|
||||
|
||||
const emailStrategy = new EmailLinkStrategy(
|
||||
{
|
||||
sendEmail: sendMagicLinkEmail,
|
||||
secret,
|
||||
callbackURL: "/magic",
|
||||
sessionMagicLinkKey: "triggerdotdev:magiclink",
|
||||
},
|
||||
async ({
|
||||
email,
|
||||
form,
|
||||
magicLinkVerify,
|
||||
}: {
|
||||
email: string;
|
||||
form: FormData;
|
||||
magicLinkVerify: boolean;
|
||||
}) => {
|
||||
logger.info("Magic link user authenticated", { email, magicLinkVerify });
|
||||
|
||||
// Gate the link CLICK, not just the send: a magic link issued before
|
||||
// SSO enforcement flipped on (or replayed within its validity
|
||||
// window) must not mint a session for an enforced domain.
|
||||
if (magicLinkVerify) {
|
||||
const ssoRedirect = await ssoRedirectForEmail(email, "domain_policy");
|
||||
if (ssoRedirect) {
|
||||
throw new SsoRequiredError(ssoRedirect);
|
||||
}
|
||||
}
|
||||
|
||||
try {
|
||||
const { user, isNewUser } = await findOrCreateUser({
|
||||
email,
|
||||
authenticationMethod: "MAGIC_LINK",
|
||||
});
|
||||
|
||||
await postAuthentication({ user, isNewUser, loginMethod: "MAGIC_LINK" });
|
||||
|
||||
return { userId: user.id };
|
||||
} catch (error) {
|
||||
logger.debug("Magic link user failed to authenticate", { error: JSON.stringify(error) });
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
);
|
||||
|
||||
export function addEmailLinkStrategy(authenticator: Authenticator<AuthUser>) {
|
||||
authenticator.use(emailStrategy);
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
import { env } from "~/env.server";
|
||||
import { authorizationRateLimitMiddleware } from "./authorizationRateLimitMiddleware.server";
|
||||
import type { Duration } from "./rateLimiter.server";
|
||||
|
||||
export const engineRateLimiter = authorizationRateLimitMiddleware({
|
||||
redis: {
|
||||
port: env.RATE_LIMIT_REDIS_PORT,
|
||||
host: env.RATE_LIMIT_REDIS_HOST,
|
||||
username: env.RATE_LIMIT_REDIS_USERNAME,
|
||||
password: env.RATE_LIMIT_REDIS_PASSWORD,
|
||||
tlsDisabled: env.RATE_LIMIT_REDIS_TLS_DISABLED === "true",
|
||||
clusterMode: env.RATE_LIMIT_REDIS_CLUSTER_MODE_ENABLED === "1",
|
||||
},
|
||||
keyPrefix: "engine",
|
||||
defaultLimiter: {
|
||||
type: "tokenBucket",
|
||||
refillRate: env.RUN_ENGINE_RATE_LIMIT_REFILL_RATE,
|
||||
interval: env.RUN_ENGINE_RATE_LIMIT_REFILL_INTERVAL as Duration,
|
||||
maxTokens: env.RUN_ENGINE_RATE_LIMIT_MAX,
|
||||
},
|
||||
limiterCache: {
|
||||
fresh: 60_000 * 10, // Data is fresh for 10 minutes
|
||||
stale: 60_000 * 20, // Date is stale after 20 minutes
|
||||
maxItems: 1000,
|
||||
},
|
||||
pathMatchers: [/^\/engine/],
|
||||
// Regex allow any path starting with /engine/v1/worker-actions/
|
||||
pathWhiteList: [/^\/engine\/v1\/worker-actions\/.*/],
|
||||
log: {
|
||||
rejections: env.RUN_ENGINE_RATE_LIMIT_REJECTION_LOGS_ENABLED === "1",
|
||||
requests: env.RUN_ENGINE_RATE_LIMIT_REQUEST_LOGS_ENABLED === "1",
|
||||
limiter: env.RUN_ENGINE_RATE_LIMIT_LIMITER_LOGS_ENABLED === "1",
|
||||
},
|
||||
});
|
||||
@@ -0,0 +1,233 @@
|
||||
import { type ClickHouse } from "@internal/clickhouse";
|
||||
import type { TaskRunStatus } from "@trigger.dev/database";
|
||||
import { QUEUED_STATUSES } from "~/components/runs/v3/TaskRunStatus";
|
||||
|
||||
export type DailyTaskActivity = Record<string, ({ day: string } & Record<TaskRunStatus, number>)[]>;
|
||||
export type CurrentRunningStats = Record<string, { queued: number; running: number }>;
|
||||
export type AverageDurations = Record<string, number>;
|
||||
|
||||
export interface EnvironmentMetricsRepository {
|
||||
getDailyTaskActivity(options: {
|
||||
organizationId: string;
|
||||
projectId: string;
|
||||
environmentId: string;
|
||||
days: number;
|
||||
tasks: string[];
|
||||
}): Promise<DailyTaskActivity>;
|
||||
|
||||
getCurrentRunningStats(options: {
|
||||
organizationId: string;
|
||||
projectId: string;
|
||||
environmentId: string;
|
||||
days: number;
|
||||
tasks: string[];
|
||||
}): Promise<CurrentRunningStats>;
|
||||
|
||||
getAverageDurations(options: {
|
||||
organizationId: string;
|
||||
projectId: string;
|
||||
environmentId: string;
|
||||
days: number;
|
||||
tasks: string[];
|
||||
}): Promise<AverageDurations>;
|
||||
}
|
||||
|
||||
export type ClickHouseEnvironmentMetricsRepositoryOptions = {
|
||||
clickhouse: ClickHouse;
|
||||
};
|
||||
|
||||
export class ClickHouseEnvironmentMetricsRepository implements EnvironmentMetricsRepository {
|
||||
constructor(private readonly options: ClickHouseEnvironmentMetricsRepositoryOptions) {}
|
||||
|
||||
public async getDailyTaskActivity({
|
||||
organizationId,
|
||||
projectId,
|
||||
environmentId,
|
||||
days,
|
||||
tasks,
|
||||
}: {
|
||||
organizationId: string;
|
||||
projectId: string;
|
||||
environmentId: string;
|
||||
days: number;
|
||||
tasks: string[];
|
||||
}): Promise<DailyTaskActivity> {
|
||||
if (tasks.length === 0) {
|
||||
return {};
|
||||
}
|
||||
|
||||
const [queryError, activity] = await this.options.clickhouse.taskRuns.getTaskActivity({
|
||||
organizationId,
|
||||
projectId,
|
||||
environmentId,
|
||||
days,
|
||||
});
|
||||
|
||||
if (queryError) {
|
||||
throw queryError;
|
||||
}
|
||||
|
||||
return fillInDailyTaskActivity(
|
||||
activity.map((a) => ({
|
||||
taskIdentifier: a.task_identifier,
|
||||
status: a.status as TaskRunStatus,
|
||||
day: new Date(a.day),
|
||||
count: BigInt(a.count),
|
||||
})),
|
||||
days
|
||||
);
|
||||
}
|
||||
|
||||
public async getCurrentRunningStats({
|
||||
organizationId,
|
||||
projectId,
|
||||
environmentId,
|
||||
days,
|
||||
tasks,
|
||||
}: {
|
||||
organizationId: string;
|
||||
projectId: string;
|
||||
environmentId: string;
|
||||
days: number;
|
||||
tasks: string[];
|
||||
}): Promise<CurrentRunningStats> {
|
||||
if (tasks.length === 0) {
|
||||
return {};
|
||||
}
|
||||
|
||||
const [queryError, stats] = await this.options.clickhouse.taskRuns.getCurrentRunningStats({
|
||||
organizationId,
|
||||
projectId,
|
||||
environmentId,
|
||||
days,
|
||||
});
|
||||
|
||||
if (queryError) {
|
||||
throw queryError;
|
||||
}
|
||||
|
||||
return fillInCurrentRunningStats(
|
||||
stats.map((s) => ({
|
||||
taskIdentifier: s.task_identifier,
|
||||
status: s.status as TaskRunStatus,
|
||||
count: BigInt(s.count),
|
||||
})),
|
||||
tasks
|
||||
);
|
||||
}
|
||||
|
||||
public async getAverageDurations({
|
||||
organizationId,
|
||||
projectId,
|
||||
environmentId,
|
||||
days,
|
||||
tasks,
|
||||
}: {
|
||||
organizationId: string;
|
||||
projectId: string;
|
||||
environmentId: string;
|
||||
days: number;
|
||||
tasks: string[];
|
||||
}): Promise<AverageDurations> {
|
||||
if (tasks.length === 0) {
|
||||
return {};
|
||||
}
|
||||
|
||||
const [queryError, durations] = await this.options.clickhouse.taskRuns.getAverageDurations({
|
||||
organizationId,
|
||||
projectId,
|
||||
environmentId,
|
||||
days,
|
||||
});
|
||||
|
||||
if (queryError) {
|
||||
throw queryError;
|
||||
}
|
||||
|
||||
return Object.fromEntries(durations.map((d) => [d.task_identifier, Number(d.duration)]));
|
||||
}
|
||||
}
|
||||
|
||||
type TaskActivityResults = Array<{
|
||||
taskIdentifier: string;
|
||||
status: TaskRunStatus;
|
||||
day: Date;
|
||||
count: bigint;
|
||||
}>;
|
||||
|
||||
function fillInDailyTaskActivity(activity: TaskActivityResults, days: number): DailyTaskActivity {
|
||||
//today with no time
|
||||
const today = new Date();
|
||||
today.setUTCHours(0, 0, 0, 0);
|
||||
|
||||
return activity.reduce((acc, a) => {
|
||||
let existingTask = acc[a.taskIdentifier];
|
||||
|
||||
if (!existingTask) {
|
||||
existingTask = [];
|
||||
//populate the array with the past 7 days
|
||||
for (let i = days; i >= 0; i--) {
|
||||
const day = new Date(today);
|
||||
day.setUTCDate(today.getDate() - i);
|
||||
day.setUTCHours(0, 0, 0, 0);
|
||||
|
||||
existingTask.push({
|
||||
day: day.toISOString(),
|
||||
["COMPLETED_SUCCESSFULLY"]: 0,
|
||||
} as { day: string } & Record<TaskRunStatus, number>);
|
||||
}
|
||||
|
||||
acc[a.taskIdentifier] = existingTask;
|
||||
}
|
||||
|
||||
const dayString = a.day.toISOString();
|
||||
const day = existingTask.find((d) => d.day === dayString);
|
||||
|
||||
if (!day) {
|
||||
return acc;
|
||||
}
|
||||
|
||||
day[a.status] = Number(a.count);
|
||||
|
||||
return acc;
|
||||
}, {} as DailyTaskActivity);
|
||||
}
|
||||
|
||||
type CurrentRunningStatsResults = Array<{
|
||||
taskIdentifier: string;
|
||||
status: TaskRunStatus;
|
||||
count: bigint;
|
||||
}>;
|
||||
|
||||
function fillInCurrentRunningStats(
|
||||
stats: CurrentRunningStatsResults,
|
||||
tasks: string[]
|
||||
): CurrentRunningStats {
|
||||
//create an object combining the queued and concurrency counts
|
||||
const result: Record<string, { queued: number; running: number }> = {};
|
||||
for (const task of tasks) {
|
||||
const queued = stats.filter(
|
||||
(q) => q.taskIdentifier === task && QUEUED_STATUSES.includes(q.status)
|
||||
);
|
||||
const queuedCount =
|
||||
queued.length === 0
|
||||
? 0
|
||||
: queued.reduce((acc, q) => {
|
||||
return acc + Number(q.count);
|
||||
}, 0);
|
||||
|
||||
const running = stats.filter((r) => r.taskIdentifier === task && r.status === "EXECUTING");
|
||||
const runningCount =
|
||||
running.length === 0
|
||||
? 0
|
||||
: running.reduce((acc, r) => {
|
||||
return acc + Number(r.count);
|
||||
}, 0);
|
||||
|
||||
result[task] = {
|
||||
queued: queuedCount,
|
||||
running: runningCount,
|
||||
};
|
||||
}
|
||||
return result;
|
||||
}
|
||||
@@ -0,0 +1,87 @@
|
||||
import { json } from "@remix-run/server-runtime";
|
||||
import type { RuntimeEnvironmentType } from "@trigger.dev/database";
|
||||
import { isUserActorToken } from "@trigger.dev/rbac";
|
||||
import { rbac } from "~/services/rbac.server";
|
||||
|
||||
type EnvironmentScopedResource = "envvars" | "apiKeys";
|
||||
|
||||
const RESOURCE_LABELS: Record<EnvironmentScopedResource, string> = {
|
||||
envvars: "environment variables",
|
||||
apiKeys: "API keys",
|
||||
};
|
||||
|
||||
/**
|
||||
* Env-tier RBAC for environment-scoped API routes (env vars, and the endpoints
|
||||
* that hand out an environment's secret credentials).
|
||||
*
|
||||
* Machine credentials (an environment's secret/public API key) are already
|
||||
* scoped to a single environment, so they pass through unchanged. A personal
|
||||
* access token (or a delegated user-actor token) carries a user, so enforce
|
||||
* that user's role for the targeted environment tier — e.g. a Developer can't
|
||||
* read deployed env vars or API keys via the API, matching the dashboard
|
||||
* restriction. Blocking the credential read for deployed tiers is also what
|
||||
* stops a restricted role deploying via the CLI (deploy needs the
|
||||
* environment's secret key).
|
||||
*
|
||||
* Returns a `Response` to short-circuit with when access is denied, or
|
||||
* `undefined` when the request may proceed.
|
||||
*/
|
||||
export async function authorizePatEnvironmentAccess({
|
||||
request,
|
||||
authType,
|
||||
organizationId,
|
||||
projectId,
|
||||
envType,
|
||||
resource,
|
||||
action,
|
||||
}: {
|
||||
request: Request;
|
||||
authType: "personalAccessToken" | "organizationAccessToken" | "apiKey";
|
||||
organizationId: string;
|
||||
projectId: string;
|
||||
envType: RuntimeEnvironmentType;
|
||||
resource: EnvironmentScopedResource;
|
||||
action: "read" | "write";
|
||||
}): Promise<Response | undefined> {
|
||||
const bearer = request.headers
|
||||
.get("Authorization")
|
||||
?.replace(/^Bearer /, "")
|
||||
.trim();
|
||||
const isUat = !!bearer && isUserActorToken(bearer);
|
||||
|
||||
// Machine creds (apiKey) and org tokens carry no user role to enforce. A
|
||||
// user-actor token carries a user just like a PAT, so it's gated too.
|
||||
if (authType !== "personalAccessToken" && !isUat) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
const userAuth = isUat
|
||||
? await rbac.authenticateUserActor(request, { organizationId, projectId })
|
||||
: await rbac.authenticatePat(request, { organizationId, projectId });
|
||||
if (!userAuth.ok) {
|
||||
return json({ error: userAuth.error }, { status: userAuth.status });
|
||||
}
|
||||
|
||||
if (!userAuth.ability.can(action, { type: resource, envType })) {
|
||||
return json(
|
||||
{
|
||||
error: `You don't have permission to access this environment's ${RESOURCE_LABELS[resource]}.`,
|
||||
},
|
||||
{ status: 403 }
|
||||
);
|
||||
}
|
||||
|
||||
return undefined;
|
||||
}
|
||||
|
||||
/** Env-tier env var access for the env var API routes. */
|
||||
export function authorizeEnvVarApiRequest(opts: {
|
||||
request: Request;
|
||||
authType: "personalAccessToken" | "organizationAccessToken" | "apiKey";
|
||||
organizationId: string;
|
||||
projectId: string;
|
||||
envType: RuntimeEnvironmentType;
|
||||
action: "read" | "write";
|
||||
}): Promise<Response | undefined> {
|
||||
return authorizePatEnvironmentAccess({ ...opts, resource: "envvars" });
|
||||
}
|
||||
@@ -0,0 +1,190 @@
|
||||
import { App, type Octokit } from "octokit";
|
||||
import { env } from "../env.server";
|
||||
import { prisma } from "~/db.server";
|
||||
import { logger } from "./logger.server";
|
||||
import { errAsync, fromPromise, okAsync, type ResultAsync } from "neverthrow";
|
||||
|
||||
export const githubApp =
|
||||
env.GITHUB_APP_ENABLED === "1"
|
||||
? new App({
|
||||
appId: env.GITHUB_APP_ID,
|
||||
privateKey: env.GITHUB_APP_PRIVATE_KEY,
|
||||
webhooks: {
|
||||
secret: env.GITHUB_APP_WEBHOOK_SECRET,
|
||||
},
|
||||
})
|
||||
: null;
|
||||
|
||||
/**
|
||||
* Links a GitHub App installation to a Trigger organization
|
||||
*/
|
||||
export async function linkGitHubAppInstallation(
|
||||
installationId: number,
|
||||
organizationId: string
|
||||
): Promise<void> {
|
||||
if (!githubApp) {
|
||||
throw new Error("GitHub App is not enabled");
|
||||
}
|
||||
|
||||
const octokit = await githubApp.getInstallationOctokit(installationId);
|
||||
const { data: installation } = await octokit.rest.apps.getInstallation({
|
||||
installation_id: installationId,
|
||||
});
|
||||
|
||||
const repositories = await fetchInstallationRepositories(octokit, installationId);
|
||||
|
||||
const repositorySelection = installation.repository_selection === "all" ? "ALL" : "SELECTED";
|
||||
|
||||
await prisma.githubAppInstallation.create({
|
||||
data: {
|
||||
appInstallationId: installationId,
|
||||
organizationId,
|
||||
targetId: installation.target_id,
|
||||
targetType: installation.target_type,
|
||||
accountHandle: installation.account
|
||||
? "login" in installation.account
|
||||
? installation.account.login
|
||||
: "slug" in installation.account
|
||||
? installation.account.slug
|
||||
: "-"
|
||||
: "-",
|
||||
permissions: installation.permissions,
|
||||
repositorySelection,
|
||||
repositories: {
|
||||
create: repositories,
|
||||
},
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Links a GitHub App installation to a Trigger organization
|
||||
*/
|
||||
export async function updateGitHubAppInstallation(installationId: number): Promise<void> {
|
||||
if (!githubApp) {
|
||||
throw new Error("GitHub App is not enabled");
|
||||
}
|
||||
|
||||
const octokit = await githubApp.getInstallationOctokit(installationId);
|
||||
const { data: installation } = await octokit.rest.apps.getInstallation({
|
||||
installation_id: installationId,
|
||||
});
|
||||
|
||||
const existingInstallation = await prisma.githubAppInstallation.findFirst({
|
||||
where: { appInstallationId: installationId },
|
||||
});
|
||||
|
||||
if (!existingInstallation) {
|
||||
throw new Error("GitHub App installation not found");
|
||||
}
|
||||
|
||||
const repositorySelection = installation.repository_selection === "all" ? "ALL" : "SELECTED";
|
||||
|
||||
// repos are updated asynchronously via webhook events
|
||||
await prisma.githubAppInstallation.update({
|
||||
where: { id: existingInstallation?.id },
|
||||
data: {
|
||||
appInstallationId: installationId,
|
||||
targetId: installation.target_id,
|
||||
targetType: installation.target_type,
|
||||
accountHandle: installation.account
|
||||
? "login" in installation.account
|
||||
? installation.account.login
|
||||
: "slug" in installation.account
|
||||
? installation.account.slug
|
||||
: "-"
|
||||
: "-",
|
||||
permissions: installation.permissions,
|
||||
suspendedAt: existingInstallation?.suspendedAt,
|
||||
repositorySelection,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
async function fetchInstallationRepositories(octokit: Octokit, installationId: number) {
|
||||
const iterator = octokit.paginate.iterator(octokit.rest.apps.listReposAccessibleToInstallation, {
|
||||
installation_id: installationId,
|
||||
per_page: 100,
|
||||
});
|
||||
|
||||
const allRepos = [];
|
||||
const maxPages = 3;
|
||||
let pageCount = 0;
|
||||
|
||||
for await (const { data } of iterator) {
|
||||
pageCount++;
|
||||
allRepos.push(...data);
|
||||
|
||||
if (maxPages && pageCount >= maxPages) {
|
||||
logger.warn("GitHub installation repository fetch truncated", {
|
||||
installationId,
|
||||
maxPages,
|
||||
totalReposFetched: allRepos.length,
|
||||
});
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
return allRepos.map((repo) => ({
|
||||
githubId: repo.id,
|
||||
name: repo.name,
|
||||
fullName: repo.full_name,
|
||||
htmlUrl: repo.html_url,
|
||||
private: repo.private,
|
||||
defaultBranch: repo.default_branch,
|
||||
}));
|
||||
}
|
||||
|
||||
/**
|
||||
* Checks if a branch exists in a GitHub repository
|
||||
*/
|
||||
export function checkGitHubBranchExists(
|
||||
installationId: number,
|
||||
fullRepoName: string,
|
||||
branch: string
|
||||
): ResultAsync<boolean, { type: "other" | "github_app_not_enabled"; cause?: unknown }> {
|
||||
if (!githubApp) {
|
||||
return errAsync({ type: "github_app_not_enabled" as const });
|
||||
}
|
||||
|
||||
if (!branch || branch.trim() === "") {
|
||||
return okAsync(false);
|
||||
}
|
||||
|
||||
const [owner, repo] = fullRepoName.split("/");
|
||||
|
||||
const getOctokit = () =>
|
||||
fromPromise(githubApp.getInstallationOctokit(installationId), (error) => ({
|
||||
type: "other" as const,
|
||||
cause: error,
|
||||
}));
|
||||
|
||||
const getBranch = (octokit: Octokit) =>
|
||||
fromPromise(
|
||||
octokit.rest.repos.getBranch({
|
||||
owner,
|
||||
repo,
|
||||
branch,
|
||||
}),
|
||||
(error) => ({
|
||||
type: "other" as const,
|
||||
cause: error,
|
||||
})
|
||||
);
|
||||
|
||||
return getOctokit()
|
||||
.andThen((octokit) => getBranch(octokit))
|
||||
.map(() => true)
|
||||
.orElse((error) => {
|
||||
if (
|
||||
error.cause &&
|
||||
error.cause instanceof Error &&
|
||||
"status" in error.cause &&
|
||||
error.cause.status === 404
|
||||
) {
|
||||
return okAsync(false);
|
||||
}
|
||||
|
||||
return errAsync(error);
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
import type { Authenticator } from "remix-auth";
|
||||
import { GitHubStrategy } from "remix-auth-github";
|
||||
import { env } from "~/env.server";
|
||||
import { findOrCreateUser } from "~/models/user.server";
|
||||
import type { AuthUser } from "./authUser";
|
||||
import { logger } from "./logger.server";
|
||||
import { postAuthentication } from "./postAuth.server";
|
||||
import { SsoRequiredError, ssoRedirectForEmail } from "./ssoAutoDiscovery.server";
|
||||
|
||||
export function addGitHubStrategy(
|
||||
authenticator: Authenticator<AuthUser>,
|
||||
clientID: string,
|
||||
clientSecret: string
|
||||
) {
|
||||
const gitHubStrategy = new GitHubStrategy(
|
||||
{
|
||||
clientID,
|
||||
clientSecret,
|
||||
callbackURL: `${env.LOGIN_ORIGIN}/auth/github/callback`,
|
||||
},
|
||||
async ({ extraParams, profile }) => {
|
||||
const emails = profile.emails;
|
||||
|
||||
if (!emails?.length) {
|
||||
throw new Error("GitHub login requires an email address");
|
||||
}
|
||||
|
||||
const email = emails[0].value;
|
||||
|
||||
// SSO auto-discovery gate — BEFORE findOrCreateUser, so an
|
||||
// SSO-enforced domain never gets this GitHub identity linked onto
|
||||
// an existing account.
|
||||
const ssoRedirect = await ssoRedirectForEmail(email, "oauth_blocked");
|
||||
if (ssoRedirect) {
|
||||
throw new SsoRequiredError(ssoRedirect);
|
||||
}
|
||||
|
||||
try {
|
||||
logger.debug("GitHub login", {
|
||||
emails,
|
||||
profile,
|
||||
extraParams,
|
||||
});
|
||||
|
||||
const { user, isNewUser } = await findOrCreateUser({
|
||||
email,
|
||||
authenticationMethod: "GITHUB",
|
||||
authenticationProfile: profile,
|
||||
authenticationExtraParams: extraParams,
|
||||
});
|
||||
|
||||
await postAuthentication({ user, isNewUser, loginMethod: "GITHUB" });
|
||||
|
||||
return {
|
||||
userId: user.id,
|
||||
};
|
||||
} catch (error) {
|
||||
logger.error("GitHub login failed", { error: JSON.stringify(error) });
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
);
|
||||
|
||||
authenticator.use(gitHubStrategy);
|
||||
}
|
||||
@@ -0,0 +1,124 @@
|
||||
import { createCookieSessionStorage } from "@remix-run/node";
|
||||
import { randomBytes } from "crypto";
|
||||
import { env } from "../env.server";
|
||||
import { logger } from "./logger.server";
|
||||
|
||||
const sessionStorage = createCookieSessionStorage({
|
||||
cookie: {
|
||||
name: "__github_app_install",
|
||||
httpOnly: true,
|
||||
maxAge: 60 * 60, // 1 hour
|
||||
path: "/",
|
||||
sameSite: "lax",
|
||||
secrets: [env.SESSION_SECRET],
|
||||
secure: env.NODE_ENV === "production",
|
||||
},
|
||||
});
|
||||
|
||||
/**
|
||||
* Creates a secure session for GitHub App installation with organization tracking
|
||||
*/
|
||||
export async function createGitHubAppInstallSession(
|
||||
organizationId: string,
|
||||
redirectTo: string
|
||||
): Promise<{ url: string; cookieHeader: string }> {
|
||||
if (env.GITHUB_APP_ENABLED !== "1") {
|
||||
throw new Error("GitHub App is not enabled");
|
||||
}
|
||||
|
||||
const state = randomBytes(32).toString("hex");
|
||||
|
||||
const session = await sessionStorage.getSession();
|
||||
session.set("organizationId", organizationId);
|
||||
session.set("redirectTo", redirectTo);
|
||||
session.set("state", state);
|
||||
session.set("createdAt", Date.now());
|
||||
|
||||
const githubAppSlug = env.GITHUB_APP_SLUG;
|
||||
|
||||
// the state query param gets passed through to the installation callback
|
||||
const url = `https://github.com/apps/${githubAppSlug}/installations/new?state=${state}`;
|
||||
|
||||
const cookieHeader = await sessionStorage.commitSession(session);
|
||||
|
||||
return { url, cookieHeader };
|
||||
}
|
||||
|
||||
/**
|
||||
* Validates and retrieves the GitHub App installation session
|
||||
*/
|
||||
export async function validateGitHubAppInstallSession(
|
||||
cookieHeader: string | null,
|
||||
state: string
|
||||
): Promise<
|
||||
{ valid: true; organizationId: string; redirectTo: string } | { valid: false; error?: string }
|
||||
> {
|
||||
if (!cookieHeader) {
|
||||
return {
|
||||
valid: false,
|
||||
error: "No installation session cookie found",
|
||||
};
|
||||
}
|
||||
|
||||
const session = await sessionStorage.getSession(cookieHeader);
|
||||
|
||||
const sessionState = session.get("state");
|
||||
const organizationId = session.get("organizationId");
|
||||
const redirectTo = session.get("redirectTo");
|
||||
const createdAt = session.get("createdAt");
|
||||
|
||||
if (!sessionState || !organizationId || !createdAt || !redirectTo) {
|
||||
logger.warn("GitHub App installation session missing required fields", {
|
||||
hasState: !!sessionState,
|
||||
hasOrgId: !!organizationId,
|
||||
hasCreatedAt: !!createdAt,
|
||||
hasRedirectTo: !!redirectTo,
|
||||
});
|
||||
|
||||
return {
|
||||
valid: false,
|
||||
error: "invalid_session_data",
|
||||
};
|
||||
}
|
||||
|
||||
if (sessionState !== state) {
|
||||
logger.warn("GitHub App installation state mismatch", {
|
||||
expectedState: sessionState,
|
||||
receivedState: state,
|
||||
});
|
||||
return {
|
||||
valid: false,
|
||||
error: "state_mismatch",
|
||||
};
|
||||
}
|
||||
|
||||
const expirationTime = createdAt + 60 * 60 * 1000;
|
||||
if (Date.now() > expirationTime) {
|
||||
logger.warn("GitHub App installation session expired", {
|
||||
createdAt: new Date(createdAt),
|
||||
now: new Date(),
|
||||
});
|
||||
return {
|
||||
valid: false,
|
||||
error: "session_expired",
|
||||
};
|
||||
}
|
||||
|
||||
return {
|
||||
valid: true,
|
||||
organizationId,
|
||||
redirectTo,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Destroys the GitHub App installation cookie session
|
||||
*/
|
||||
export async function destroyGitHubAppInstallSession(cookieHeader: string | null): Promise<string> {
|
||||
if (!cookieHeader) {
|
||||
return "";
|
||||
}
|
||||
|
||||
const session = await sessionStorage.getSession(cookieHeader);
|
||||
return await sessionStorage.destroySession(session);
|
||||
}
|
||||
@@ -0,0 +1,74 @@
|
||||
import type { Authenticator } from "remix-auth";
|
||||
import { GoogleStrategy } from "remix-auth-google";
|
||||
import { env } from "~/env.server";
|
||||
import { findOrCreateUser } from "~/models/user.server";
|
||||
import type { AuthUser } from "./authUser";
|
||||
import { isGoogleEmailVerified } from "./googleEmailVerification";
|
||||
import { logger } from "./logger.server";
|
||||
import { postAuthentication } from "./postAuth.server";
|
||||
import { SsoRequiredError, ssoRedirectForEmail } from "./ssoAutoDiscovery.server";
|
||||
|
||||
export function addGoogleStrategy(
|
||||
authenticator: Authenticator<AuthUser>,
|
||||
clientID: string,
|
||||
clientSecret: string
|
||||
) {
|
||||
const googleStrategy = new GoogleStrategy(
|
||||
{
|
||||
clientID,
|
||||
clientSecret,
|
||||
callbackURL: `${env.LOGIN_ORIGIN}/auth/google/callback`,
|
||||
},
|
||||
async ({ extraParams, profile }) => {
|
||||
const emails = profile.emails;
|
||||
|
||||
if (!emails?.length) {
|
||||
throw new Error("Google login requires an email address");
|
||||
}
|
||||
|
||||
const email = emails[0].value;
|
||||
|
||||
// Only trust the email if Google asserts it's verified, since account
|
||||
// linking keys off it. See isGoogleEmailVerified.
|
||||
if (!isGoogleEmailVerified(profile)) {
|
||||
throw new Error(
|
||||
"Google login refused: the Google account's email is not verified. Sign in with an account whose email Google has verified, or use magic-link / GitHub."
|
||||
);
|
||||
}
|
||||
|
||||
// SSO auto-discovery gate — BEFORE findOrCreateUser, so an
|
||||
// SSO-enforced domain never gets this Google identity linked onto
|
||||
// an existing account.
|
||||
const ssoRedirect = await ssoRedirectForEmail(email, "oauth_blocked");
|
||||
if (ssoRedirect) {
|
||||
throw new SsoRequiredError(ssoRedirect);
|
||||
}
|
||||
|
||||
try {
|
||||
logger.debug("Google login", {
|
||||
emails,
|
||||
profile,
|
||||
extraParams,
|
||||
});
|
||||
|
||||
const { user, isNewUser } = await findOrCreateUser({
|
||||
email,
|
||||
authenticationMethod: "GOOGLE",
|
||||
authenticationProfile: profile,
|
||||
authenticationExtraParams: extraParams,
|
||||
});
|
||||
|
||||
await postAuthentication({ user, isNewUser, loginMethod: "GOOGLE" });
|
||||
|
||||
return {
|
||||
userId: user.id,
|
||||
};
|
||||
} catch (error) {
|
||||
logger.error("Google login failed", { error: JSON.stringify(error) });
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
);
|
||||
|
||||
authenticator.use(googleStrategy);
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
import type { GoogleProfile } from "remix-auth-google";
|
||||
|
||||
/**
|
||||
* Whether Google has asserted that the profile's email is verified. A
|
||||
* successful OAuth flow proves control of the Google account, not ownership of
|
||||
* the email it carries, and account linking keys off the email.
|
||||
*
|
||||
* Strict by design: only a real boolean `true` counts. A missing claim, missing
|
||||
* `_json`, the string `"true"`, or a truthy `1` are all treated as unverified.
|
||||
*/
|
||||
export function isGoogleEmailVerified(profile: GoogleProfile): boolean {
|
||||
const emailVerified = (profile as { _json?: { email_verified?: unknown } })?._json
|
||||
?.email_verified;
|
||||
return emailVerified === true;
|
||||
}
|
||||
@@ -0,0 +1,33 @@
|
||||
import { AsyncLocalStorage } from "node:async_hooks";
|
||||
|
||||
export type HttpLocalStorage = {
|
||||
requestId: string;
|
||||
path: string;
|
||||
host: string;
|
||||
method: string;
|
||||
abortController: AbortController;
|
||||
};
|
||||
|
||||
const httpLocalStorage = new AsyncLocalStorage<HttpLocalStorage>();
|
||||
|
||||
export type RunWithHttpContextFunction = <T>(context: HttpLocalStorage, fn: () => T) => T;
|
||||
|
||||
export function runWithHttpContext<T>(context: HttpLocalStorage, fn: () => T): T {
|
||||
return httpLocalStorage.run(context, fn);
|
||||
}
|
||||
|
||||
export function getHttpContext(): HttpLocalStorage | undefined {
|
||||
return httpLocalStorage.getStore();
|
||||
}
|
||||
|
||||
// Fallback signal that is never aborted, safe for tests and non-Express contexts.
|
||||
const neverAbortedSignal = new AbortController().signal;
|
||||
|
||||
/**
|
||||
* Returns an AbortSignal wired to the Express response's "close" event.
|
||||
* This bypasses the broken request.signal chain in @remix-run/express
|
||||
* (caused by Node.js undici GC bug nodejs/node#55428).
|
||||
*/
|
||||
export function getRequestAbortSignal(): AbortSignal {
|
||||
return httpLocalStorage.getStore()?.abortController.signal ?? neverAbortedSignal;
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
import { WaitpointId } from "@trigger.dev/core/v3/isomorphic";
|
||||
import nodeCrypto from "node:crypto";
|
||||
import { env } from "~/env.server";
|
||||
|
||||
export function generateHttpCallbackUrl(waitpointId: string, apiKey: string) {
|
||||
const hash = generateHttpCallbackHash(waitpointId, apiKey);
|
||||
|
||||
return `${env.API_ORIGIN ?? env.APP_ORIGIN}/api/v1/waitpoints/tokens/${WaitpointId.toFriendlyId(
|
||||
waitpointId
|
||||
)}/callback/${hash}`;
|
||||
}
|
||||
|
||||
function generateHttpCallbackHash(waitpointId: string, apiKey: string) {
|
||||
const hmac = nodeCrypto.createHmac("sha256", apiKey);
|
||||
hmac.update(waitpointId);
|
||||
return hmac.digest("hex");
|
||||
}
|
||||
|
||||
export function verifyHttpCallbackHash(waitpointId: string, hash: string, apiKey: string) {
|
||||
const expectedHash = generateHttpCallbackHash(waitpointId, apiKey);
|
||||
|
||||
if (
|
||||
hash.length === expectedHash.length &&
|
||||
nodeCrypto.timingSafeEqual(Buffer.from(hash, "hex"), Buffer.from(expectedHash, "hex"))
|
||||
) {
|
||||
return true;
|
||||
}
|
||||
|
||||
return false;
|
||||
}
|
||||
@@ -0,0 +1,154 @@
|
||||
import { createCookieSessionStorage, type Session } from "@remix-run/node";
|
||||
import { SignJWT, jwtVerify, errors } from "jose";
|
||||
import { randomUUID } from "node:crypto";
|
||||
import { singleton } from "~/utils/singleton";
|
||||
import { createRedisClient, type RedisClient } from "~/redis.server";
|
||||
import { env } from "~/env.server";
|
||||
import { logger } from "~/services/logger.server";
|
||||
|
||||
export const impersonationSessionStorage = createCookieSessionStorage({
|
||||
cookie: {
|
||||
name: "__impersonate", // use any name you want here
|
||||
sameSite: "lax", // this helps with CSRF
|
||||
path: "/", // remember to add this so the cookie will work in all routes
|
||||
httpOnly: true, // for security reasons, make this cookie http only
|
||||
secrets: [env.SESSION_SECRET],
|
||||
secure: env.NODE_ENV === "production", // enable this in prod only
|
||||
maxAge: 60 * 60 * 24, // 1 day
|
||||
},
|
||||
});
|
||||
|
||||
export function getImpersonationSession(request: Request) {
|
||||
return impersonationSessionStorage.getSession(request.headers.get("Cookie"));
|
||||
}
|
||||
|
||||
export function commitImpersonationSession(session: Session) {
|
||||
return impersonationSessionStorage.commitSession(session);
|
||||
}
|
||||
|
||||
export async function getImpersonationId(request: Request) {
|
||||
const session = await getImpersonationSession(request);
|
||||
|
||||
return session.get("impersonatedUserId") as string | undefined;
|
||||
}
|
||||
|
||||
export async function setImpersonationId(userId: string, request: Request) {
|
||||
const session = await getImpersonationSession(request);
|
||||
|
||||
session.set("impersonatedUserId", userId);
|
||||
|
||||
return session;
|
||||
}
|
||||
|
||||
export async function clearImpersonationId(request: Request) {
|
||||
const session = await getImpersonationSession(request);
|
||||
|
||||
session.unset("impersonatedUserId");
|
||||
|
||||
return session;
|
||||
}
|
||||
|
||||
// Impersonation token utilities for CSRF protection
|
||||
const IMPERSONATION_TOKEN_EXPIRY_SECONDS = 5 * 60; // 5 minutes
|
||||
|
||||
function getImpersonationTokenSecret(): Uint8Array {
|
||||
return new TextEncoder().encode(env.SESSION_SECRET);
|
||||
}
|
||||
|
||||
function getImpersonationTokenRedisClient(): RedisClient {
|
||||
return singleton("impersonationTokenRedis", () =>
|
||||
createRedisClient("impersonation:token", {
|
||||
host: env.CACHE_REDIS_HOST,
|
||||
port: env.CACHE_REDIS_PORT,
|
||||
username: env.CACHE_REDIS_USERNAME,
|
||||
password: env.CACHE_REDIS_PASSWORD,
|
||||
tlsDisabled: env.CACHE_REDIS_TLS_DISABLED === "true",
|
||||
clusterMode: env.CACHE_REDIS_CLUSTER_MODE_ENABLED === "1",
|
||||
keyPrefix: "impersonation:token:",
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Generate a signed one-time impersonation token for a user
|
||||
*/
|
||||
export async function generateImpersonationToken(userId: string): Promise<string> {
|
||||
const secret = getImpersonationTokenSecret();
|
||||
const now = Math.floor(Date.now() / 1000);
|
||||
|
||||
const token = await new SignJWT({ userId })
|
||||
.setProtectedHeader({ alg: "HS256" })
|
||||
.setIssuedAt(now)
|
||||
.setExpirationTime(now + IMPERSONATION_TOKEN_EXPIRY_SECONDS)
|
||||
.setIssuer("https://trigger.dev")
|
||||
.setAudience("https://trigger.dev/admin")
|
||||
.setJti(randomUUID())
|
||||
.sign(secret);
|
||||
|
||||
return token;
|
||||
}
|
||||
|
||||
/**
|
||||
* Validate and consume an impersonation token (prevents replay attacks)
|
||||
*/
|
||||
export async function validateAndConsumeImpersonationToken(
|
||||
token: string
|
||||
): Promise<string | undefined> {
|
||||
try {
|
||||
const secret = getImpersonationTokenSecret();
|
||||
|
||||
// Verify the token signature and expiration
|
||||
const { payload } = await jwtVerify(token, secret, {
|
||||
issuer: "https://trigger.dev",
|
||||
audience: "https://trigger.dev/admin",
|
||||
});
|
||||
|
||||
const userId = payload.userId as string | undefined;
|
||||
if (!userId || typeof userId !== "string") {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
// Atomically mark the token as used to prevent replay attacks.
|
||||
// Use the jti (a short UUID) as the Redis key rather than the full JWT string.
|
||||
// If Redis is unavailable, fall back to JWT expiry as the only protection.
|
||||
if (env.CACHE_REDIS_HOST) {
|
||||
// Defensively reject tokens without a jti (e.g. issued before this change)
|
||||
if (!payload.jti) {
|
||||
logger.warn("Impersonation token missing jti claim, rejecting for safety");
|
||||
return undefined;
|
||||
}
|
||||
|
||||
try {
|
||||
const redis = getImpersonationTokenRedisClient();
|
||||
const result = await redis.set(
|
||||
payload.jti,
|
||||
"1",
|
||||
"EX",
|
||||
IMPERSONATION_TOKEN_EXPIRY_SECONDS,
|
||||
"NX"
|
||||
);
|
||||
if (result !== "OK") {
|
||||
// Token was already used
|
||||
return undefined;
|
||||
}
|
||||
} catch (redisError) {
|
||||
logger.warn(
|
||||
"Redis unavailable for impersonation token tracking, relying on JWT expiry only",
|
||||
{
|
||||
error: redisError instanceof Error ? redisError.message : String(redisError),
|
||||
}
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
return userId;
|
||||
} catch (error) {
|
||||
if (error instanceof errors.JWTExpired || error instanceof errors.JWTInvalid) {
|
||||
return undefined;
|
||||
}
|
||||
logger.error("Error validating impersonation token", {
|
||||
error: error instanceof Error ? error.message : String(error),
|
||||
});
|
||||
return undefined;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,102 @@
|
||||
import { Redis } from "ioredis";
|
||||
import { defaultReconnectOnError } from "@internal/redis";
|
||||
import { env } from "~/env.server";
|
||||
import { singleton } from "~/utils/singleton";
|
||||
import { logger } from "./logger.server";
|
||||
|
||||
const KEY_PREFIX = "isw:";
|
||||
const DEFAULT_TTL_MS = 7 * 24 * 60 * 60 * 1000; // 7 days
|
||||
|
||||
function buildKey(runFriendlyId: string, streamId: string): string {
|
||||
return `${KEY_PREFIX}${runFriendlyId}:${streamId}`;
|
||||
}
|
||||
|
||||
function initializeRedis(): Redis | undefined {
|
||||
const host = env.CACHE_REDIS_HOST;
|
||||
if (!host) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
return new Redis({
|
||||
connectionName: "inputStreamWaitpointCache",
|
||||
host,
|
||||
port: env.CACHE_REDIS_PORT,
|
||||
username: env.CACHE_REDIS_USERNAME,
|
||||
password: env.CACHE_REDIS_PASSWORD,
|
||||
keyPrefix: "tr:",
|
||||
enableAutoPipelining: true,
|
||||
reconnectOnError: defaultReconnectOnError,
|
||||
...(env.CACHE_REDIS_TLS_DISABLED === "true" ? {} : { tls: {} }),
|
||||
});
|
||||
}
|
||||
|
||||
const redis = singleton("inputStreamWaitpointCache", initializeRedis);
|
||||
|
||||
/**
|
||||
* Store a mapping from input stream to waitpoint ID in Redis.
|
||||
* Called when `.wait()` creates a new waitpoint.
|
||||
*/
|
||||
export async function setInputStreamWaitpoint(
|
||||
runFriendlyId: string,
|
||||
streamId: string,
|
||||
waitpointId: string,
|
||||
ttlMs?: number
|
||||
): Promise<void> {
|
||||
if (!redis) return;
|
||||
|
||||
try {
|
||||
const key = buildKey(runFriendlyId, streamId);
|
||||
await redis.set(key, waitpointId, "PX", ttlMs ?? DEFAULT_TTL_MS);
|
||||
} catch (error) {
|
||||
logger.error("Failed to set input stream waitpoint cache", {
|
||||
runFriendlyId,
|
||||
streamId,
|
||||
error,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the waitpoint ID for an input stream without deleting it.
|
||||
* Called from the `.send()` route before completing the waitpoint.
|
||||
*/
|
||||
export async function getInputStreamWaitpoint(
|
||||
runFriendlyId: string,
|
||||
streamId: string
|
||||
): Promise<string | null> {
|
||||
if (!redis) return null;
|
||||
|
||||
try {
|
||||
const key = buildKey(runFriendlyId, streamId);
|
||||
return await redis.get(key);
|
||||
} catch (error) {
|
||||
logger.error("Failed to get input stream waitpoint cache", {
|
||||
runFriendlyId,
|
||||
streamId,
|
||||
error,
|
||||
});
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Delete the cache entry for an input stream waitpoint.
|
||||
* Called when a waitpoint is completed or timed out.
|
||||
*/
|
||||
export async function deleteInputStreamWaitpoint(
|
||||
runFriendlyId: string,
|
||||
streamId: string
|
||||
): Promise<void> {
|
||||
if (!redis) return;
|
||||
|
||||
try {
|
||||
const key = buildKey(runFriendlyId, streamId);
|
||||
await redis.del(key);
|
||||
} catch (error) {
|
||||
logger.error("Failed to delete input stream waitpoint cache", {
|
||||
runFriendlyId,
|
||||
streamId,
|
||||
error,
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
import { createCookie } from "@remix-run/node";
|
||||
import { env } from "~/env.server";
|
||||
|
||||
export type LastAuthMethod = "github" | "google" | "email" | "sso";
|
||||
|
||||
// Cookie that persists for 1 year to remember the user's last login method
|
||||
export const lastAuthMethodCookie = createCookie("last-auth-method", {
|
||||
maxAge: 60 * 60 * 24 * 365, // 1 year
|
||||
httpOnly: true,
|
||||
sameSite: "lax",
|
||||
secure: env.NODE_ENV === "production",
|
||||
});
|
||||
|
||||
export async function getLastAuthMethod(request: Request): Promise<LastAuthMethod | null> {
|
||||
const cookie = request.headers.get("Cookie");
|
||||
const value = await lastAuthMethodCookie.parse(cookie);
|
||||
if (value === "github" || value === "google" || value === "email" || value === "sso") {
|
||||
return value;
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
export async function setLastAuthMethodHeader(method: LastAuthMethod): Promise<string> {
|
||||
return lastAuthMethodCookie.serialize(method);
|
||||
}
|
||||
@@ -0,0 +1,22 @@
|
||||
import { type Params } from "@remix-run/react";
|
||||
import { findProjectBySlug } from "~/models/project.server";
|
||||
import { findEnvironmentBySlug } from "~/models/runtimeEnvironment.server";
|
||||
import { requireUserId } from "~/services/session.server";
|
||||
import { EnvironmentParamSchema } from "~/utils/pathBuilder";
|
||||
|
||||
export async function loadProjectEnvironmentFromRequest(request: Request, params: Params<string>) {
|
||||
const userId = await requireUserId(request);
|
||||
const { organizationSlug, projectParam, envParam } = EnvironmentParamSchema.parse(params);
|
||||
|
||||
const project = await findProjectBySlug(organizationSlug, projectParam, userId);
|
||||
if (!project) {
|
||||
throw new Response(undefined, { status: 404, statusText: "Project not found" });
|
||||
}
|
||||
|
||||
const environment = await findEnvironmentBySlug(project.id, envParam, userId);
|
||||
if (!environment) {
|
||||
throw new Response(undefined, { status: 404, statusText: "Environment not found" });
|
||||
}
|
||||
|
||||
return { userId, project, environment };
|
||||
}
|
||||
@@ -0,0 +1,98 @@
|
||||
import type { LogLevel } from "@trigger.dev/core/logger";
|
||||
import { Logger } from "@trigger.dev/core/logger";
|
||||
import { patchConsoleToTelnet, startTelnetLogServer } from "@trigger.dev/core/v3/telnetLogServer";
|
||||
import { sensitiveDataReplacer } from "./sensitiveDataReplacer";
|
||||
import { AsyncLocalStorage } from "async_hooks";
|
||||
import { getHttpContext } from "./httpAsyncStorage.server";
|
||||
import { captureException, captureMessage } from "@sentry/remix";
|
||||
|
||||
const currentFieldsStore = new AsyncLocalStorage<Record<string, unknown>>();
|
||||
|
||||
export function trace<T>(fields: Record<string, unknown>, fn: () => T): T {
|
||||
return currentFieldsStore.run(fields, fn);
|
||||
}
|
||||
|
||||
Logger.onError = (message, ...args) => {
|
||||
const error = extractErrorFromArgs(args);
|
||||
|
||||
if (error) {
|
||||
captureException(error, {
|
||||
extra: {
|
||||
message,
|
||||
...flattenArgs(args),
|
||||
},
|
||||
});
|
||||
} else {
|
||||
captureMessage(message, {
|
||||
level: "error",
|
||||
extra: flattenArgs(args),
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
function extractErrorFromArgs(args: Array<Record<string, unknown> | undefined>) {
|
||||
for (const arg of args) {
|
||||
if (arg && "error" in arg && arg.error instanceof Error) {
|
||||
return arg.error;
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
function flattenArgs(args: Array<Record<string, unknown> | undefined>) {
|
||||
return args.reduce((acc, arg) => {
|
||||
if (arg) {
|
||||
return { ...acc, ...arg };
|
||||
}
|
||||
return acc;
|
||||
}, {});
|
||||
}
|
||||
export const logger = new Logger(
|
||||
"webapp",
|
||||
(process.env.APP_LOG_LEVEL ?? "info") as LogLevel,
|
||||
["examples", "output", "connectionString", "payload"],
|
||||
sensitiveDataReplacer,
|
||||
() => {
|
||||
const fields = currentFieldsStore.getStore();
|
||||
const httpContext = getHttpContext();
|
||||
return { ...fields, http: httpContext };
|
||||
}
|
||||
);
|
||||
|
||||
export const workerLogger = new Logger(
|
||||
"worker",
|
||||
(process.env.APP_LOG_LEVEL ?? "info") as LogLevel,
|
||||
["examples", "output", "connectionString"],
|
||||
sensitiveDataReplacer,
|
||||
() => {
|
||||
const fields = currentFieldsStore.getStore();
|
||||
return fields ? { ...fields } : {};
|
||||
}
|
||||
);
|
||||
|
||||
export const socketLogger = new Logger(
|
||||
"socket",
|
||||
(process.env.APP_LOG_LEVEL ?? "info") as LogLevel,
|
||||
[],
|
||||
sensitiveDataReplacer,
|
||||
() => {
|
||||
const fields = currentFieldsStore.getStore();
|
||||
return fields ? { ...fields } : {};
|
||||
}
|
||||
);
|
||||
|
||||
// Opt-in, dev-only: mirror this process's stdout to a local telnet/TCP stream.
|
||||
// We patch console (rather than the static Logger.onLog sink) so the stream also captures logs
|
||||
// from separate/bundled copies of the Logger — e.g. the enterprise SSO plugin, which bundles its
|
||||
// own @trigger.dev/core and logs via its own console.log, invisible to the webapp's onLog hook.
|
||||
const telnetLogsPort = process.env.WEBAPP_TELNET_LOGS_PORT
|
||||
? Number(process.env.WEBAPP_TELNET_LOGS_PORT)
|
||||
: undefined;
|
||||
if (telnetLogsPort && Number.isFinite(telnetLogsPort) && telnetLogsPort > 0) {
|
||||
const telnetGlobal = globalThis as typeof globalThis & { __webappTelnetLogs?: boolean };
|
||||
if (!telnetGlobal.__webappTelnetLogs) {
|
||||
telnetGlobal.__webappTelnetLogs = true;
|
||||
const telnetLogServer = startTelnetLogServer({ port: telnetLogsPort, name: "webapp" });
|
||||
patchConsoleToTelnet(telnetLogServer, { pretty: true });
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,82 @@
|
||||
import { env } from "~/env.server";
|
||||
import { logger } from "./logger.server";
|
||||
|
||||
class LoopsClient {
|
||||
constructor(private readonly apiKey: string) {}
|
||||
|
||||
async userCreated({
|
||||
userId,
|
||||
email,
|
||||
name,
|
||||
}: {
|
||||
userId: string;
|
||||
email: string;
|
||||
name: string | null;
|
||||
}) {
|
||||
logger.info(`Loops send "sign-up" event`, { userId, email, name });
|
||||
return this.#sendEvent({
|
||||
email,
|
||||
userId,
|
||||
firstName: name?.split(" ").at(0),
|
||||
eventName: "sign-up",
|
||||
});
|
||||
}
|
||||
|
||||
async #sendEvent({
|
||||
email,
|
||||
userId,
|
||||
firstName,
|
||||
eventName,
|
||||
eventProperties,
|
||||
}: {
|
||||
email: string;
|
||||
userId: string;
|
||||
firstName?: string;
|
||||
eventName: string;
|
||||
eventProperties?: Record<string, string | number | boolean>;
|
||||
}) {
|
||||
const options = {
|
||||
method: "POST",
|
||||
headers: { Authorization: `Bearer ${this.apiKey}`, "Content-Type": "application/json" },
|
||||
body: JSON.stringify({
|
||||
email,
|
||||
userId,
|
||||
firstName,
|
||||
eventName,
|
||||
eventProperties,
|
||||
}),
|
||||
};
|
||||
|
||||
try {
|
||||
const response = await fetch("https://app.loops.so/api/v1/events/send", options);
|
||||
|
||||
if (!response.ok) {
|
||||
logger.error(`Loops sendEvent ${eventName} bad status`, {
|
||||
status: response.status,
|
||||
email,
|
||||
userId,
|
||||
firstName,
|
||||
eventProperties,
|
||||
eventName,
|
||||
});
|
||||
return false;
|
||||
}
|
||||
|
||||
const responseBody = (await response.json()) as any;
|
||||
|
||||
if (!responseBody.success) {
|
||||
logger.error(`Loops sendEvent ${eventName} failed response`, {
|
||||
message: responseBody.message,
|
||||
});
|
||||
return false;
|
||||
}
|
||||
|
||||
return true;
|
||||
} catch (error) {
|
||||
logger.error(`Loops sendEvent ${eventName} failed`, { error });
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
export const loopsClient = env.LOOPS_API_KEY ? new LoopsClient(env.LOOPS_API_KEY) : null;
|
||||
@@ -0,0 +1,96 @@
|
||||
import { Ratelimit } from "@upstash/ratelimit";
|
||||
import { env } from "~/env.server";
|
||||
import { createRedisRateLimitClient, RateLimiter } from "~/services/rateLimiter.server";
|
||||
import { singleton } from "~/utils/singleton";
|
||||
|
||||
export class MagicLinkRateLimitError extends Error {
|
||||
public readonly retryAfter: number;
|
||||
|
||||
constructor(retryAfter: number) {
|
||||
super("Magic link request rate limit exceeded.");
|
||||
this.retryAfter = retryAfter;
|
||||
}
|
||||
}
|
||||
|
||||
function getRedisClient() {
|
||||
return createRedisRateLimitClient({
|
||||
port: env.RATE_LIMIT_REDIS_PORT,
|
||||
host: env.RATE_LIMIT_REDIS_HOST,
|
||||
username: env.RATE_LIMIT_REDIS_USERNAME,
|
||||
password: env.RATE_LIMIT_REDIS_PASSWORD,
|
||||
tlsDisabled: env.RATE_LIMIT_REDIS_TLS_DISABLED === "true",
|
||||
clusterMode: env.RATE_LIMIT_REDIS_CLUSTER_MODE_ENABLED === "1",
|
||||
});
|
||||
}
|
||||
|
||||
const magicLinkEmailRateLimiter = singleton(
|
||||
"magicLinkEmailRateLimiter",
|
||||
initializeMagicLinkEmailRateLimiter
|
||||
);
|
||||
|
||||
function initializeMagicLinkEmailRateLimiter() {
|
||||
return new RateLimiter({
|
||||
redisClient: getRedisClient(),
|
||||
keyPrefix: "auth:magiclink:email",
|
||||
limiter: Ratelimit.slidingWindow(3, "1 m"), // 3 requests per minute per email
|
||||
logSuccess: false,
|
||||
logFailure: true,
|
||||
});
|
||||
}
|
||||
|
||||
const magicLinkEmailDailyRateLimiter = singleton(
|
||||
"magicLinkEmailDailyRateLimiter",
|
||||
initializeMagicLinkEmailDailyRateLimiter
|
||||
);
|
||||
|
||||
function initializeMagicLinkEmailDailyRateLimiter() {
|
||||
return new RateLimiter({
|
||||
redisClient: getRedisClient(),
|
||||
keyPrefix: "auth:magiclink:email:daily",
|
||||
limiter: Ratelimit.slidingWindow(30, "1 d"), // 30 requests per day per email
|
||||
logSuccess: false,
|
||||
logFailure: true,
|
||||
});
|
||||
}
|
||||
|
||||
const magicLinkIpRateLimiter = singleton(
|
||||
"magicLinkIpRateLimiter",
|
||||
initializeMagicLinkIpRateLimiter
|
||||
);
|
||||
|
||||
function initializeMagicLinkIpRateLimiter() {
|
||||
return new RateLimiter({
|
||||
redisClient: getRedisClient(),
|
||||
keyPrefix: "auth:magiclink:ip",
|
||||
limiter: Ratelimit.slidingWindow(10, "1 m"), // 10 requests per minute per IP
|
||||
logSuccess: false,
|
||||
logFailure: true,
|
||||
});
|
||||
}
|
||||
|
||||
export async function checkMagicLinkEmailRateLimit(identifier: string): Promise<void> {
|
||||
const result = await magicLinkEmailRateLimiter.limit(identifier);
|
||||
|
||||
if (!result.success) {
|
||||
const retryAfter = new Date(result.reset).getTime() - Date.now();
|
||||
throw new MagicLinkRateLimitError(retryAfter);
|
||||
}
|
||||
}
|
||||
|
||||
export async function checkMagicLinkEmailDailyRateLimit(identifier: string): Promise<void> {
|
||||
const result = await magicLinkEmailDailyRateLimiter.limit(identifier);
|
||||
|
||||
if (!result.success) {
|
||||
const retryAfter = new Date(result.reset).getTime() - Date.now();
|
||||
throw new MagicLinkRateLimitError(retryAfter);
|
||||
}
|
||||
}
|
||||
|
||||
export async function checkMagicLinkIpRateLimit(ip: string): Promise<void> {
|
||||
const result = await magicLinkIpRateLimiter.limit(ip);
|
||||
|
||||
if (!result.success) {
|
||||
const retryAfter = new Date(result.reset).getTime() - Date.now();
|
||||
throw new MagicLinkRateLimitError(retryAfter);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,628 @@
|
||||
import type {
|
||||
IOPacket,
|
||||
RunMetadataChangeOperation,
|
||||
UpdateMetadataRequestBody,
|
||||
} from "@trigger.dev/core/v3";
|
||||
import { applyMetadataOperations, parsePacket } from "@trigger.dev/core/v3";
|
||||
import type { PrismaClientOrTransaction } from "~/db.server";
|
||||
import type { AuthenticatedEnvironment } from "~/services/apiAuth.server";
|
||||
import { handleMetadataPacket, MetadataTooLargeError } from "~/utils/packets";
|
||||
import { ServiceValidationError } from "~/v3/services/common.server";
|
||||
import { Effect, Schedule, Duration, Fiber } from "effect";
|
||||
import { type RuntimeFiber } from "effect/Fiber";
|
||||
import { setTimeout } from "timers/promises";
|
||||
import type { LogLevel } from "@trigger.dev/core/logger";
|
||||
import { Logger } from "@trigger.dev/core/logger";
|
||||
import type { RunStore } from "@internal/run-store";
|
||||
|
||||
const RUN_UPDATABLE_WINDOW_MS = 60 * 60 * 1000; // 1 hour
|
||||
|
||||
type BufferedRunMetadataChangeOperation = {
|
||||
runId: string;
|
||||
timestamp: number;
|
||||
operation: RunMetadataChangeOperation;
|
||||
};
|
||||
|
||||
export type UpdateMetadataServiceOptions = {
|
||||
prisma: PrismaClientOrTransaction;
|
||||
runStore: RunStore;
|
||||
flushIntervalMs?: number;
|
||||
flushEnabled?: boolean;
|
||||
flushLoggingEnabled?: boolean;
|
||||
maximumSize?: number;
|
||||
logger?: Logger;
|
||||
logLevel?: LogLevel;
|
||||
/** Called after the batched flusher writes a run's buffered operations, with everything
|
||||
* a realtime change record needs — buffered (parent/root) updates otherwise never wake
|
||||
* live feeds. */
|
||||
onRunFlushed?: (run: {
|
||||
runId: string;
|
||||
environmentId: string;
|
||||
tags: string[];
|
||||
batchId: string | null;
|
||||
updatedAtMs: number;
|
||||
}) => void;
|
||||
// Testing hooks
|
||||
onBeforeUpdate?: (runId: string) => Promise<void>;
|
||||
onAfterRead?: (runId: string, metadataVersion: number) => Promise<void>;
|
||||
};
|
||||
|
||||
export class UpdateMetadataService {
|
||||
private _bufferedOperations: Map<string, BufferedRunMetadataChangeOperation[]> = new Map();
|
||||
private _flushFiber: RuntimeFiber<void> | null = null;
|
||||
private readonly _prisma: PrismaClientOrTransaction;
|
||||
private readonly _runStore: RunStore;
|
||||
private readonly flushIntervalMs: number;
|
||||
private readonly flushEnabled: boolean;
|
||||
private readonly flushLoggingEnabled: boolean;
|
||||
private readonly maximumSize: number;
|
||||
private readonly logger: Logger;
|
||||
|
||||
constructor(private readonly options: UpdateMetadataServiceOptions) {
|
||||
this._prisma = options.prisma;
|
||||
this._runStore = options.runStore;
|
||||
this.flushIntervalMs = options.flushIntervalMs ?? 5000;
|
||||
this.flushEnabled = options.flushEnabled ?? true;
|
||||
this.flushLoggingEnabled = options.flushLoggingEnabled ?? true;
|
||||
this.maximumSize = options.maximumSize ?? 1024 * 1024 * 1; // 1MB
|
||||
this.logger = options.logger ?? new Logger("UpdateMetadataService", options.logLevel ?? "info");
|
||||
|
||||
this._startFlushing();
|
||||
}
|
||||
|
||||
// Start a loop that periodically flushes buffered operations
|
||||
private _startFlushing() {
|
||||
if (!this.flushEnabled) {
|
||||
this.logger.info("[UpdateMetadataService] 🚽 Flushing disabled");
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
this.logger.info("[UpdateMetadataService] 🚽 Flushing started");
|
||||
|
||||
// Create a program that sleeps, then processes buffered ops
|
||||
const program = Effect.gen(this, function* (_) {
|
||||
while (true) {
|
||||
// Wait for flushIntervalMs before flushing again
|
||||
yield* _(Effect.sleep(Duration.millis(this.flushIntervalMs)));
|
||||
|
||||
// Atomically get and clear current operations
|
||||
const currentOperations = new Map(this._bufferedOperations);
|
||||
this._bufferedOperations.clear();
|
||||
|
||||
yield* Effect.sync(() => {
|
||||
if (this.flushLoggingEnabled) {
|
||||
this.logger.debug(`[UpdateMetadataService] Flushing operations`, {
|
||||
operations: Object.fromEntries(currentOperations),
|
||||
});
|
||||
}
|
||||
});
|
||||
|
||||
// If we have operations, process them
|
||||
if (currentOperations.size > 0) {
|
||||
yield* _(this._processBufferedOperations(currentOperations));
|
||||
}
|
||||
}
|
||||
}).pipe(
|
||||
// Handle any unexpected errors, ensuring program does not fail
|
||||
Effect.catchAll((error) =>
|
||||
Effect.sync(() => {
|
||||
this.logger.error("Error in flushing program:", { error });
|
||||
})
|
||||
)
|
||||
);
|
||||
|
||||
// Fork the program so it runs in the background
|
||||
this._flushFiber = Effect.runFork(program as Effect.Effect<void, never, never>);
|
||||
}
|
||||
|
||||
stopFlushing() {
|
||||
if (this._flushFiber) {
|
||||
Effect.runFork(Fiber.interrupt(this._flushFiber));
|
||||
}
|
||||
}
|
||||
|
||||
private _processBufferedOperations = (
|
||||
operations: Map<string, BufferedRunMetadataChangeOperation[]>
|
||||
) => {
|
||||
return Effect.gen(this, function* (_) {
|
||||
for (const [runId, ops] of operations) {
|
||||
// Process and cull operations
|
||||
const processedOps = this._cullOperations(ops);
|
||||
|
||||
// If there are no operations to process, skip
|
||||
if (processedOps.length === 0) {
|
||||
continue;
|
||||
}
|
||||
|
||||
yield* Effect.sync(() => {
|
||||
if (this.flushLoggingEnabled) {
|
||||
this.logger.debug(`[UpdateMetadataService] Processing operations for run`, {
|
||||
runId,
|
||||
operationsCount: processedOps.length,
|
||||
});
|
||||
}
|
||||
});
|
||||
|
||||
// Update run with retry
|
||||
yield* _(
|
||||
this._updateRunWithOperations(runId, processedOps).pipe(
|
||||
// Catch MetadataTooLargeError before retry logic
|
||||
Effect.catchIf(
|
||||
(error) => error instanceof MetadataTooLargeError,
|
||||
(error) =>
|
||||
Effect.sync(() => {
|
||||
// Log the error but don't return operations to buffer
|
||||
console.error(
|
||||
`[UpdateMetadataService] Dropping operations for run ${runId} due to metadata size limit:`,
|
||||
error.message
|
||||
);
|
||||
if (this.flushLoggingEnabled) {
|
||||
this.logger.warn(`[UpdateMetadataService] Metadata too large for run`, {
|
||||
runId,
|
||||
operationsCount: processedOps.length,
|
||||
error: error.message,
|
||||
});
|
||||
}
|
||||
})
|
||||
),
|
||||
Effect.retry(Schedule.exponential(Duration.millis(100), 1.4)),
|
||||
Effect.catchAll((error) =>
|
||||
Effect.sync(() => {
|
||||
// On complete failure, return ops to buffer
|
||||
const existingOps = this._bufferedOperations.get(runId) ?? [];
|
||||
this._bufferedOperations.set(runId, [...existingOps, ...ops]);
|
||||
console.error(`Failed to process run ${runId}:`, error);
|
||||
})
|
||||
)
|
||||
)
|
||||
);
|
||||
}
|
||||
});
|
||||
};
|
||||
|
||||
private _updateRunWithOperations = (
|
||||
runId: string,
|
||||
operations: BufferedRunMetadataChangeOperation[]
|
||||
) => {
|
||||
return Effect.gen(this, function* (_) {
|
||||
// Fetch current run (+ the realtime membership keys, so a flush can publish)
|
||||
const run = yield* _(
|
||||
Effect.tryPromise(() =>
|
||||
this._runStore.findRun(
|
||||
{ id: runId },
|
||||
{
|
||||
select: {
|
||||
id: true,
|
||||
metadata: true,
|
||||
metadataType: true,
|
||||
metadataVersion: true,
|
||||
runtimeEnvironmentId: true,
|
||||
runTags: true,
|
||||
batchId: true,
|
||||
},
|
||||
},
|
||||
this._prisma
|
||||
)
|
||||
)
|
||||
);
|
||||
|
||||
if (!run) {
|
||||
return yield* _(Effect.fail(new Error(`Run ${runId} not found`)));
|
||||
}
|
||||
|
||||
// Testing hook after read
|
||||
if (this.options.onAfterRead) {
|
||||
yield* _(Effect.tryPromise(() => this.options.onAfterRead!(runId, run.metadataVersion)));
|
||||
}
|
||||
|
||||
const metadata = yield* _(
|
||||
Effect.tryPromise(() =>
|
||||
run.metadata
|
||||
? parsePacket({ data: run.metadata, dataType: run.metadataType })
|
||||
: Promise.resolve({})
|
||||
)
|
||||
);
|
||||
|
||||
// Apply operations and update
|
||||
const applyResult = applyMetadataOperations(
|
||||
metadata,
|
||||
operations.map((op) => op.operation)
|
||||
);
|
||||
|
||||
if (applyResult.unappliedOperations.length === operations.length) {
|
||||
this.logger.warn(`No operations applied for run ${runId}`);
|
||||
// If no operations were applied, return
|
||||
return;
|
||||
}
|
||||
|
||||
// Stringify the metadata
|
||||
const newMetadataPacket = yield* _(
|
||||
Effect.try(() =>
|
||||
handleMetadataPacket(applyResult.newMetadata, run.metadataType, this.maximumSize)
|
||||
).pipe(
|
||||
Effect.mapError((error) => {
|
||||
// Preserve the original error if it's MetadataTooLargeError
|
||||
if ("cause" in error && error.cause instanceof MetadataTooLargeError) {
|
||||
return error.cause;
|
||||
}
|
||||
return error;
|
||||
})
|
||||
)
|
||||
);
|
||||
|
||||
if (!newMetadataPacket) {
|
||||
// Log and skip if metadata is invalid
|
||||
this.logger.warn(`Invalid metadata after operations, skipping update`);
|
||||
return;
|
||||
}
|
||||
|
||||
// Testing hook before update
|
||||
if (this.options.onBeforeUpdate) {
|
||||
yield* _(Effect.tryPromise(() => this.options.onBeforeUpdate!(runId)));
|
||||
}
|
||||
|
||||
// Stamp updatedAt explicitly so the realtime publish can carry the exact committed
|
||||
// value without a follow-up read (updateMany can't RETURNING).
|
||||
const writeTime = new Date();
|
||||
const result = yield* _(
|
||||
Effect.tryPromise(() =>
|
||||
this._runStore.updateMetadata(
|
||||
runId,
|
||||
{
|
||||
metadata: newMetadataPacket.data!,
|
||||
metadataVersion: { increment: 1 },
|
||||
updatedAt: writeTime,
|
||||
},
|
||||
{ expectedMetadataVersion: run.metadataVersion },
|
||||
this._prisma
|
||||
)
|
||||
)
|
||||
);
|
||||
|
||||
if (result.count === 0) {
|
||||
yield* Effect.sync(() => {
|
||||
this.logger.warn(`Optimistic lock failed for run ${runId}`, {
|
||||
metadataVersion: run.metadataVersion,
|
||||
});
|
||||
});
|
||||
|
||||
return yield* _(Effect.fail(new Error("Optimistic lock failed")));
|
||||
}
|
||||
|
||||
yield* Effect.sync(() => {
|
||||
this.options.onRunFlushed?.({
|
||||
runId,
|
||||
environmentId: run.runtimeEnvironmentId,
|
||||
tags: run.runTags,
|
||||
batchId: run.batchId,
|
||||
updatedAtMs: writeTime.getTime(),
|
||||
});
|
||||
});
|
||||
|
||||
return result;
|
||||
});
|
||||
};
|
||||
|
||||
private _cullOperations(
|
||||
operations: BufferedRunMetadataChangeOperation[]
|
||||
): BufferedRunMetadataChangeOperation[] {
|
||||
// Sort by timestamp
|
||||
const sortedOps = [...operations].sort((a, b) => a.timestamp - b.timestamp);
|
||||
|
||||
// Track latest set operations by key
|
||||
const latestSetOps = new Map<string, BufferedRunMetadataChangeOperation>();
|
||||
const resultOps: BufferedRunMetadataChangeOperation[] = [];
|
||||
|
||||
for (const op of sortedOps) {
|
||||
if (op.operation.type === "set") {
|
||||
latestSetOps.set(op.operation.key, op);
|
||||
} else {
|
||||
resultOps.push(op);
|
||||
}
|
||||
}
|
||||
|
||||
// Add winning set operations
|
||||
resultOps.push(...latestSetOps.values());
|
||||
|
||||
return resultOps;
|
||||
}
|
||||
|
||||
public async call(
|
||||
runId: string,
|
||||
body: UpdateMetadataRequestBody,
|
||||
environment?: AuthenticatedEnvironment
|
||||
) {
|
||||
const runIdType = runId.startsWith("run_") ? "friendly" : "internal";
|
||||
|
||||
const taskRun = await this._runStore.findRun(
|
||||
environment
|
||||
? {
|
||||
runtimeEnvironmentId: environment.id,
|
||||
...(runIdType === "internal" ? { id: runId } : { friendlyId: runId }),
|
||||
}
|
||||
: {
|
||||
...(runIdType === "internal" ? { id: runId } : { friendlyId: runId }),
|
||||
},
|
||||
{
|
||||
select: {
|
||||
id: true,
|
||||
batchId: true,
|
||||
runTags: true,
|
||||
completedAt: true,
|
||||
status: true,
|
||||
metadata: true,
|
||||
metadataType: true,
|
||||
metadataVersion: true,
|
||||
parentTaskRun: {
|
||||
select: {
|
||||
id: true,
|
||||
status: true,
|
||||
},
|
||||
},
|
||||
rootTaskRun: {
|
||||
select: {
|
||||
id: true,
|
||||
status: true,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
this._prisma
|
||||
);
|
||||
|
||||
if (!taskRun) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (!this.#isRunUpdatable(taskRun)) {
|
||||
throw new ServiceValidationError("Cannot update metadata for a completed run");
|
||||
}
|
||||
|
||||
if (body.parentOperations && body.parentOperations.length > 0) {
|
||||
this.#ingestRunOperations(taskRun.parentTaskRun?.id ?? taskRun.id, body.parentOperations);
|
||||
}
|
||||
|
||||
if (body.rootOperations && body.rootOperations.length > 0) {
|
||||
this.#ingestRunOperations(taskRun.rootTaskRun?.id ?? taskRun.id, body.rootOperations);
|
||||
}
|
||||
|
||||
const result = await this.#updateRunMetadata({
|
||||
runId: taskRun.id,
|
||||
body,
|
||||
existingMetadata: {
|
||||
data: taskRun.metadata ?? undefined,
|
||||
dataType: taskRun.metadataType,
|
||||
},
|
||||
});
|
||||
|
||||
return {
|
||||
metadata: result?.metadata,
|
||||
// Internal id + membership keys, so callers can publish full realtime records the router routes by index.
|
||||
runId: taskRun.id,
|
||||
batchId: taskRun.batchId,
|
||||
runTags: taskRun.runTags,
|
||||
// The committed row's updatedAt — the realtime watermark. Undefined when nothing was
|
||||
// written here (no-op, or buffered for the flusher, which publishes itself).
|
||||
updatedAtMs: result?.updatedAtMs,
|
||||
};
|
||||
}
|
||||
|
||||
async #updateRunMetadata({
|
||||
runId,
|
||||
body,
|
||||
existingMetadata,
|
||||
}: {
|
||||
runId: string;
|
||||
body: UpdateMetadataRequestBody;
|
||||
existingMetadata: IOPacket;
|
||||
}): Promise<{ metadata: Record<string, unknown> | undefined; updatedAtMs?: number }> {
|
||||
if (Array.isArray(body.operations)) {
|
||||
return this.#updateRunMetadataWithOperations(runId, body.operations);
|
||||
} else {
|
||||
return this.#updateRunMetadataDirectly(runId, body, existingMetadata);
|
||||
}
|
||||
}
|
||||
|
||||
async #updateRunMetadataWithOperations(
|
||||
runId: string,
|
||||
operations: RunMetadataChangeOperation[]
|
||||
): Promise<{ metadata: Record<string, unknown> | undefined; updatedAtMs?: number }> {
|
||||
const MAX_RETRIES = 3;
|
||||
let attempts = 0;
|
||||
|
||||
while (attempts <= MAX_RETRIES) {
|
||||
// Fetch the latest run data
|
||||
const run = await this._runStore.findRun(
|
||||
{ id: runId },
|
||||
{
|
||||
select: { metadata: true, metadataType: true, metadataVersion: true },
|
||||
},
|
||||
this._prisma
|
||||
);
|
||||
|
||||
if (!run) {
|
||||
throw new Error(`Run ${runId} not found`);
|
||||
}
|
||||
|
||||
// Testing hook after read
|
||||
if (this.options.onAfterRead) {
|
||||
await this.options.onAfterRead(runId, run.metadataVersion);
|
||||
}
|
||||
|
||||
// Parse the current metadata
|
||||
const currentMetadata = await (run.metadata
|
||||
? parsePacket({ data: run.metadata, dataType: run.metadataType })
|
||||
: Promise.resolve({}));
|
||||
|
||||
// Apply operations to the current metadata
|
||||
const applyResults = applyMetadataOperations(currentMetadata, operations);
|
||||
|
||||
// If no operations were applied, return the current metadata (nothing written)
|
||||
if (applyResults.unappliedOperations.length === operations.length) {
|
||||
return { metadata: currentMetadata };
|
||||
}
|
||||
|
||||
const newMetadataPacket = handleMetadataPacket(
|
||||
applyResults.newMetadata,
|
||||
run.metadataType,
|
||||
this.maximumSize
|
||||
);
|
||||
|
||||
if (!newMetadataPacket) {
|
||||
throw new ServiceValidationError("Unable to update metadata");
|
||||
}
|
||||
|
||||
// Testing hook before update
|
||||
if (this.options.onBeforeUpdate) {
|
||||
await this.options.onBeforeUpdate(runId);
|
||||
}
|
||||
|
||||
// Update with optimistic locking; updatedAt stamped explicitly so the caller can
|
||||
// publish the exact committed watermark without a follow-up read.
|
||||
const writeTime = new Date();
|
||||
const result = await this._runStore.updateMetadata(
|
||||
runId,
|
||||
{
|
||||
metadata: newMetadataPacket.data!,
|
||||
metadataType: newMetadataPacket.dataType,
|
||||
metadataVersion: {
|
||||
increment: 1,
|
||||
},
|
||||
updatedAt: writeTime,
|
||||
},
|
||||
{ expectedMetadataVersion: run.metadataVersion },
|
||||
this._prisma
|
||||
);
|
||||
|
||||
if (result.count === 0) {
|
||||
if (this.flushLoggingEnabled) {
|
||||
this.logger.debug(
|
||||
`[updateRunMetadataWithOperations] Optimistic lock failed for run ${runId}`,
|
||||
{
|
||||
metadataVersion: run.metadataVersion,
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
// If this was our last attempt, buffer the operations and return optimistically
|
||||
// (no watermark — the flusher writes later and publishes itself).
|
||||
if (attempts === MAX_RETRIES) {
|
||||
this.#ingestRunOperations(runId, operations);
|
||||
return { metadata: applyResults.newMetadata };
|
||||
}
|
||||
|
||||
// Otherwise sleep and try again
|
||||
await setTimeout(100 * Math.pow(1.4, attempts));
|
||||
attempts++;
|
||||
continue;
|
||||
}
|
||||
|
||||
if (this.flushLoggingEnabled) {
|
||||
this.logger.debug(`[updateRunMetadataWithOperations] Updated metadata for run`, {
|
||||
metadata: applyResults.newMetadata,
|
||||
operations: operations,
|
||||
runId,
|
||||
});
|
||||
}
|
||||
|
||||
// Success! Return the new metadata
|
||||
return { metadata: applyResults.newMetadata, updatedAtMs: writeTime.getTime() };
|
||||
}
|
||||
|
||||
return { metadata: undefined };
|
||||
}
|
||||
|
||||
// Checks to see if a run is updatable
|
||||
// if there is no completedAt, the run is updatable
|
||||
// if the run is completed, but the completedAt is within the last 10 minutes, the run is updatable
|
||||
#isRunUpdatable(run: { completedAt: Date | null }) {
|
||||
if (!run.completedAt) {
|
||||
return true;
|
||||
}
|
||||
|
||||
return run.completedAt.getTime() > Date.now() - RUN_UPDATABLE_WINDOW_MS;
|
||||
}
|
||||
|
||||
async #updateRunMetadataDirectly(
|
||||
runId: string,
|
||||
body: UpdateMetadataRequestBody,
|
||||
existingMetadata: IOPacket
|
||||
): Promise<{ metadata: Record<string, unknown> | undefined; updatedAtMs?: number }> {
|
||||
const metadataPacket = handleMetadataPacket(
|
||||
body.metadata,
|
||||
"application/json",
|
||||
this.maximumSize
|
||||
);
|
||||
|
||||
if (!metadataPacket) {
|
||||
return { metadata: {} };
|
||||
}
|
||||
|
||||
let updatedAtMs: number | undefined;
|
||||
|
||||
if (
|
||||
metadataPacket.data !== "{}" ||
|
||||
(existingMetadata.data && metadataPacket.data !== existingMetadata.data)
|
||||
) {
|
||||
if (this.flushLoggingEnabled) {
|
||||
this.logger.debug(`[updateRunMetadataDirectly] Updating metadata directly for run`, {
|
||||
metadata: metadataPacket.data,
|
||||
runId,
|
||||
});
|
||||
}
|
||||
|
||||
// Update the metadata without version check; updatedAt stamped explicitly so the
|
||||
// caller can publish the exact committed watermark.
|
||||
const writeTime = new Date();
|
||||
await this._runStore.updateMetadata(
|
||||
runId,
|
||||
{
|
||||
metadata: metadataPacket?.data!,
|
||||
metadataType: metadataPacket?.dataType,
|
||||
metadataVersion: {
|
||||
increment: 1,
|
||||
},
|
||||
updatedAt: writeTime,
|
||||
},
|
||||
{},
|
||||
this._prisma
|
||||
);
|
||||
updatedAtMs = writeTime.getTime();
|
||||
}
|
||||
|
||||
const newMetadata = await parsePacket(metadataPacket);
|
||||
return { metadata: newMetadata, updatedAtMs };
|
||||
}
|
||||
|
||||
#ingestRunOperations(runId: string, operations: RunMetadataChangeOperation[]) {
|
||||
const bufferedOperations: BufferedRunMetadataChangeOperation[] = operations.map((operation) => {
|
||||
return {
|
||||
runId,
|
||||
timestamp: Date.now(),
|
||||
operation,
|
||||
};
|
||||
});
|
||||
|
||||
if (this.flushLoggingEnabled) {
|
||||
this.logger.debug(`[ingestRunOperations] Ingesting operations for run`, {
|
||||
runId,
|
||||
bufferedOperations,
|
||||
});
|
||||
}
|
||||
|
||||
const existingBufferedOperations = this._bufferedOperations.get(runId) ?? [];
|
||||
|
||||
this._bufferedOperations.set(runId, [...existingBufferedOperations, ...bufferedOperations]);
|
||||
}
|
||||
|
||||
// Testing method to manually trigger flush
|
||||
async flushOperations() {
|
||||
const currentOperations = new Map(this._bufferedOperations);
|
||||
this._bufferedOperations.clear();
|
||||
|
||||
if (currentOperations.size > 0) {
|
||||
await Effect.runPromise(this._processBufferedOperations(currentOperations));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
import { singleton } from "~/utils/singleton";
|
||||
import { env } from "~/env.server";
|
||||
import { UpdateMetadataService } from "./updateMetadata.server";
|
||||
import { prisma } from "~/db.server";
|
||||
import { runStore } from "~/v3/runStore.server";
|
||||
import { publishChangeRecord } from "~/services/realtime/runChangeNotifierInstance.server";
|
||||
|
||||
export const updateMetadataService = singleton(
|
||||
"update-metadata-service",
|
||||
() =>
|
||||
new UpdateMetadataService({
|
||||
prisma,
|
||||
runStore,
|
||||
flushIntervalMs: env.BATCH_METADATA_OPERATIONS_FLUSH_INTERVAL_MS,
|
||||
flushEnabled: env.BATCH_METADATA_OPERATIONS_FLUSH_ENABLED === "1",
|
||||
flushLoggingEnabled: env.BATCH_METADATA_OPERATIONS_FLUSH_LOGGING_ENABLED === "1",
|
||||
maximumSize: env.TASK_RUN_METADATA_MAXIMUM_SIZE,
|
||||
logLevel: env.BATCH_METADATA_OPERATIONS_FLUSH_LOGGING_ENABLED === "1" ? "debug" : "info",
|
||||
// Buffered (parent/root) operations land via the flusher, not the caller's request —
|
||||
// publish here so those changes wake live feeds too (no-op when the backend is off).
|
||||
onRunFlushed: (run) => {
|
||||
publishChangeRecord({
|
||||
runId: run.runId,
|
||||
envId: run.environmentId,
|
||||
tags: run.tags,
|
||||
batchId: run.batchId,
|
||||
updatedAtMs: run.updatedAtMs,
|
||||
});
|
||||
},
|
||||
})
|
||||
);
|
||||
@@ -0,0 +1,79 @@
|
||||
import { Ratelimit } from "@upstash/ratelimit";
|
||||
import { type RedisWithClusterOptions } from "~/redis.server";
|
||||
import { createRedisRateLimitClient, RateLimiter } from "~/services/rateLimiterCore.server";
|
||||
|
||||
// MFA rate limiting: two sliding windows in series, both of which must pass.
|
||||
// A per-minute window covers interactive retries; a cumulative daily window
|
||||
// caps total attempts per pending-MFA session.
|
||||
//
|
||||
// Free of `env.server` so it can be tested directly against a container Redis;
|
||||
// the env-derived production singletons live in `mfaRateLimiterGlobal.server.ts`.
|
||||
|
||||
// Production policy. Exported so tests assert against the real numbers.
|
||||
export const MFA_PER_MINUTE_ATTEMPTS = 5;
|
||||
export const MFA_DAILY_ATTEMPTS = 30;
|
||||
|
||||
export type MfaRateLimiters = {
|
||||
perMinute: Pick<RateLimiter, "limit">;
|
||||
daily: Pick<RateLimiter, "limit">;
|
||||
};
|
||||
|
||||
/**
|
||||
* Build the pair of MFA rate limiters. Production passes the env-derived
|
||||
* Redis connection and the default policy; tests inject a container
|
||||
* Redis (and may override the attempt caps to isolate one window).
|
||||
*/
|
||||
export function createMfaRateLimiters(options: {
|
||||
redisOptions: RedisWithClusterOptions;
|
||||
perMinuteAttempts?: number;
|
||||
dailyAttempts?: number;
|
||||
}): { perMinute: RateLimiter; daily: RateLimiter } {
|
||||
const redisClient = createRedisRateLimitClient(options.redisOptions);
|
||||
|
||||
return {
|
||||
perMinute: new RateLimiter({
|
||||
redisClient,
|
||||
keyPrefix: "mfa:validation",
|
||||
limiter: Ratelimit.slidingWindow(options.perMinuteAttempts ?? MFA_PER_MINUTE_ATTEMPTS, "1 m"),
|
||||
logSuccess: false,
|
||||
logFailure: true,
|
||||
}),
|
||||
daily: new RateLimiter({
|
||||
redisClient,
|
||||
keyPrefix: "mfa:validation:daily",
|
||||
limiter: Ratelimit.slidingWindow(options.dailyAttempts ?? MFA_DAILY_ATTEMPTS, "24 h"),
|
||||
logSuccess: false,
|
||||
logFailure: true,
|
||||
}),
|
||||
};
|
||||
}
|
||||
|
||||
export class MfaRateLimitError extends Error {
|
||||
public readonly retryAfter: number;
|
||||
|
||||
constructor(retryAfter: number) {
|
||||
super(`MFA validation rate limit exceeded.`);
|
||||
this.retryAfter = retryAfter;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Check whether the user can attempt MFA validation, enforcing both the
|
||||
* per-minute and the daily cap. The daily cap is checked first.
|
||||
* @param userId - The user ID to rate limit
|
||||
* @param limiters - The limiter pair (production singletons or test-injected)
|
||||
* @throws {MfaRateLimitError} If either rate limit is exceeded
|
||||
*/
|
||||
export async function checkMfaRateLimit(userId: string, limiters: MfaRateLimiters): Promise<void> {
|
||||
const dailyResult = await limiters.daily.limit(userId);
|
||||
if (!dailyResult.success) {
|
||||
const retryAfter = new Date(dailyResult.reset).getTime() - Date.now();
|
||||
throw new MfaRateLimitError(retryAfter);
|
||||
}
|
||||
|
||||
const result = await limiters.perMinute.limit(userId);
|
||||
if (!result.success) {
|
||||
const retryAfter = new Date(result.reset).getTime() - Date.now();
|
||||
throw new MfaRateLimitError(retryAfter);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
import { env } from "~/env.server";
|
||||
import { singleton } from "~/utils/singleton";
|
||||
import {
|
||||
checkMfaRateLimit as checkMfaRateLimitWith,
|
||||
createMfaRateLimiters,
|
||||
type MfaRateLimiters,
|
||||
} from "./mfaRateLimiter.server";
|
||||
|
||||
// Production singletons, wired to the env-derived rate-limit Redis.
|
||||
// Kept out of `mfaRateLimiter.server.ts` so that module stays free of
|
||||
// `env.server` and remains testable in isolation (see that file).
|
||||
const mfaRateLimiters = singleton("mfaRateLimiters", () =>
|
||||
createMfaRateLimiters({
|
||||
redisOptions: {
|
||||
port: env.RATE_LIMIT_REDIS_PORT,
|
||||
host: env.RATE_LIMIT_REDIS_HOST,
|
||||
username: env.RATE_LIMIT_REDIS_USERNAME,
|
||||
password: env.RATE_LIMIT_REDIS_PASSWORD,
|
||||
tlsDisabled: env.RATE_LIMIT_REDIS_TLS_DISABLED === "true",
|
||||
clusterMode: env.RATE_LIMIT_REDIS_CLUSTER_MODE_ENABLED === "1",
|
||||
},
|
||||
})
|
||||
);
|
||||
|
||||
export const mfaRateLimiter = mfaRateLimiters.perMinute;
|
||||
export const mfaDailyRateLimiter = mfaRateLimiters.daily;
|
||||
|
||||
/**
|
||||
* Production entrypoint: rate-limit an MFA validation attempt for `userId`
|
||||
* against the env-configured limiter pair. Throws `MfaRateLimitError` when
|
||||
* either the per-minute or the cumulative daily cap is exceeded.
|
||||
*/
|
||||
export function checkMfaRateLimit(userId: string, limiters: MfaRateLimiters = mfaRateLimiters) {
|
||||
return checkMfaRateLimitWith(userId, limiters);
|
||||
}
|
||||
|
||||
export { MfaRateLimitError } from "./mfaRateLimiter.server";
|
||||
@@ -0,0 +1,378 @@
|
||||
import { type SecretReference, type User, type PrismaClient } from "@trigger.dev/database";
|
||||
import { prisma } from "~/db.server";
|
||||
import { ServiceValidationError } from "~/v3/services/baseService.server";
|
||||
import { createRandomStringGenerator } from "@better-auth/utils/random";
|
||||
import { getSecretStore } from "~/services/secrets/secretStore.server";
|
||||
import { createHash } from "@better-auth/utils/hash";
|
||||
import { createOTP } from "@better-auth/utils/otp";
|
||||
import { base32 } from "@better-auth/utils/base32";
|
||||
import { z } from "zod";
|
||||
import { scheduleEmail } from "../scheduleEmail.server";
|
||||
|
||||
const generateRandomString = createRandomStringGenerator("A-Z", "0-9");
|
||||
|
||||
const SecretSchema = z.object({
|
||||
secret: z.string(),
|
||||
});
|
||||
|
||||
export class MultiFactorAuthenticationService {
|
||||
#prismaClient: PrismaClient;
|
||||
|
||||
constructor(prismaClient: PrismaClient = prisma) {
|
||||
this.#prismaClient = prismaClient;
|
||||
}
|
||||
|
||||
public async disableTotp(userId: string, params: { totpCode?: string; recoveryCode?: string }) {
|
||||
const user = await this.#prismaClient.user.findFirst({
|
||||
where: { id: userId },
|
||||
include: {
|
||||
mfaSecretReference: true,
|
||||
},
|
||||
});
|
||||
|
||||
if (!user) {
|
||||
return {
|
||||
success: false,
|
||||
};
|
||||
}
|
||||
|
||||
if (!user.mfaEnabledAt) {
|
||||
return {
|
||||
success: false,
|
||||
};
|
||||
}
|
||||
|
||||
if (!user.mfaSecretReference) {
|
||||
return {
|
||||
success: false,
|
||||
};
|
||||
}
|
||||
|
||||
// validate the TOTP code
|
||||
const secretStore = getSecretStore(user.mfaSecretReference.provider);
|
||||
const secretResult = await secretStore.getSecret(SecretSchema, user.mfaSecretReference.key);
|
||||
|
||||
if (!secretResult) {
|
||||
return {
|
||||
success: false,
|
||||
};
|
||||
}
|
||||
|
||||
const isValid = await this.#verifyTotpCodeOrRecoveryCode(
|
||||
user,
|
||||
user.mfaSecretReference,
|
||||
params.totpCode,
|
||||
params.recoveryCode
|
||||
);
|
||||
|
||||
if (!isValid) {
|
||||
return {
|
||||
success: false,
|
||||
};
|
||||
}
|
||||
|
||||
// Delete the MFA secret
|
||||
await secretStore.deleteSecret(user.mfaSecretReference.key);
|
||||
|
||||
// Delete the MFA backup codes
|
||||
await this.#prismaClient.mfaBackupCode.deleteMany({
|
||||
where: {
|
||||
userId,
|
||||
},
|
||||
});
|
||||
|
||||
await this.#prismaClient.user.update({
|
||||
where: { id: userId },
|
||||
data: {
|
||||
mfaEnabledAt: null,
|
||||
mfaSecretReference: {
|
||||
delete: true,
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
await scheduleEmail({
|
||||
email: "mfa-disabled",
|
||||
to: user.email,
|
||||
userEmail: user.email,
|
||||
});
|
||||
|
||||
return {
|
||||
success: true,
|
||||
};
|
||||
}
|
||||
|
||||
public async enableTotp(userId: string) {
|
||||
const user = await this.#prismaClient.user.findFirst({
|
||||
where: { id: userId },
|
||||
});
|
||||
|
||||
if (!user) {
|
||||
throw new ServiceValidationError("User not found");
|
||||
}
|
||||
|
||||
const secretStore = getSecretStore("DATABASE");
|
||||
|
||||
// Generate a new secret
|
||||
const secret = generateRandomString(24);
|
||||
const secretKey = `mfa:${userId}:${generateRandomString(8)}`;
|
||||
|
||||
// Store the secret in the SecretStore
|
||||
await secretStore.setSecret(secretKey, {
|
||||
secret,
|
||||
});
|
||||
|
||||
// Update the user's secret reference to the secret store
|
||||
await this.#prismaClient.user.update({
|
||||
where: { id: userId },
|
||||
data: {
|
||||
mfaSecretReference: {
|
||||
create: {
|
||||
provider: "DATABASE",
|
||||
key: secretKey,
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
// Return the secret and the recovery codes
|
||||
const otpAuthUrl = createOTP(secret).url("trigger.dev", user.email);
|
||||
|
||||
const displaySecret = base32.encode(secret, {
|
||||
padding: false,
|
||||
});
|
||||
|
||||
return {
|
||||
secret: displaySecret,
|
||||
otpAuthUrl,
|
||||
};
|
||||
}
|
||||
|
||||
public async validateTotpSetup(userId: string, totpCode: string) {
|
||||
const user = await this.#prismaClient.user.findFirst({
|
||||
where: { id: userId },
|
||||
include: {
|
||||
mfaSecretReference: true,
|
||||
},
|
||||
});
|
||||
|
||||
if (!user) {
|
||||
throw new ServiceValidationError("User not found");
|
||||
}
|
||||
|
||||
if (!user.mfaSecretReference) {
|
||||
throw new ServiceValidationError("User has not enabled MFA");
|
||||
}
|
||||
|
||||
const secretStore = getSecretStore(user.mfaSecretReference.provider);
|
||||
const secretResult = await secretStore.getSecret(SecretSchema, user.mfaSecretReference.key);
|
||||
|
||||
if (!secretResult) {
|
||||
throw new ServiceValidationError("User has not enabled MFA");
|
||||
}
|
||||
|
||||
const secret = secretResult.secret;
|
||||
|
||||
const otp = createOTP(secret, {
|
||||
digits: 6,
|
||||
period: 30,
|
||||
});
|
||||
|
||||
const isValid = await otp.verify(totpCode);
|
||||
|
||||
if (!isValid) {
|
||||
// Return the secret and the recovery codes
|
||||
const otpAuthUrl = createOTP(secret).url("trigger.dev", user.email);
|
||||
|
||||
const displaySecret = base32.encode(secret, {
|
||||
padding: false,
|
||||
});
|
||||
|
||||
return {
|
||||
success: false,
|
||||
otpAuthUrl,
|
||||
secret: displaySecret,
|
||||
};
|
||||
}
|
||||
|
||||
// Now that we've validated the TOTP code, we can enable MFA for the user
|
||||
await this.#prismaClient.user.update({
|
||||
where: { id: userId },
|
||||
data: {
|
||||
mfaEnabledAt: new Date(),
|
||||
},
|
||||
});
|
||||
|
||||
// Generate a new set of recovery codes
|
||||
const recoveryCodes = Array.from({ length: 9 }, () => generateRandomString(16, "a-z", "0-9"));
|
||||
|
||||
// Delete any existing recovery codes
|
||||
await this.#prismaClient.mfaBackupCode.deleteMany({
|
||||
where: {
|
||||
userId,
|
||||
},
|
||||
});
|
||||
|
||||
// Hash and store the recovery codes
|
||||
for (const code of recoveryCodes) {
|
||||
const hashedCode = await createHash("SHA-512", "hex").digest(code);
|
||||
await this.#prismaClient.mfaBackupCode.create({
|
||||
data: {
|
||||
userId,
|
||||
code: hashedCode,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
await scheduleEmail({
|
||||
email: "mfa-enabled",
|
||||
to: user.email,
|
||||
userEmail: user.email,
|
||||
});
|
||||
|
||||
return {
|
||||
success: true,
|
||||
recoveryCodes,
|
||||
};
|
||||
}
|
||||
|
||||
async #verifyTotpCodeOrRecoveryCode(
|
||||
user: User,
|
||||
secretReference: SecretReference,
|
||||
totpCode?: string,
|
||||
recoveryCode?: string
|
||||
) {
|
||||
if (!totpCode && !recoveryCode) {
|
||||
return false;
|
||||
}
|
||||
|
||||
if (typeof totpCode === "string" && totpCode.length === 6) {
|
||||
return this.#verifyTotpCode(user, secretReference, totpCode);
|
||||
}
|
||||
|
||||
if (typeof recoveryCode === "string") {
|
||||
return this.#verifyRecoveryCode(user, recoveryCode);
|
||||
}
|
||||
|
||||
return false;
|
||||
}
|
||||
|
||||
async #verifyTotpCode(user: User, secretReference: SecretReference, totpCode: string) {
|
||||
const secretStore = getSecretStore(secretReference.provider);
|
||||
const secretResult = await secretStore.getSecret(SecretSchema, secretReference.key);
|
||||
|
||||
if (!secretResult) {
|
||||
return false;
|
||||
}
|
||||
|
||||
const secret = secretResult.secret;
|
||||
|
||||
const isValid = await createOTP(secret, {
|
||||
digits: 6,
|
||||
period: 30,
|
||||
}).verify(totpCode);
|
||||
|
||||
return isValid;
|
||||
}
|
||||
|
||||
async #verifyRecoveryCode(user: User, recoveryCode: string) {
|
||||
const hashedCode = await createHash("SHA-512", "hex").digest(recoveryCode);
|
||||
|
||||
const backupCode = await this.#prismaClient.mfaBackupCode.findFirst({
|
||||
where: { userId: user.id, code: hashedCode, usedAt: null },
|
||||
});
|
||||
|
||||
return !!backupCode;
|
||||
}
|
||||
|
||||
// Public methods for login flow with security measures
|
||||
public async verifyTotpForLogin(userId: string, totpCode: string) {
|
||||
const user = await this.#prismaClient.user.findFirst({
|
||||
where: { id: userId },
|
||||
include: {
|
||||
mfaSecretReference: true,
|
||||
},
|
||||
});
|
||||
|
||||
if (!user || !user.mfaEnabledAt || !user.mfaSecretReference) {
|
||||
return {
|
||||
success: false,
|
||||
error: "Invalid authentication code",
|
||||
};
|
||||
}
|
||||
|
||||
// Check for replay attack - if this code was already used
|
||||
const hashedCode = await createHash("SHA-512", "hex").digest(totpCode);
|
||||
if (user.mfaLastUsedCode === hashedCode) {
|
||||
return {
|
||||
success: false,
|
||||
error: "Invalid authentication code",
|
||||
};
|
||||
}
|
||||
|
||||
// Verify the TOTP code
|
||||
const isValid = await this.#verifyTotpCode(user, user.mfaSecretReference, totpCode);
|
||||
|
||||
if (!isValid) {
|
||||
return {
|
||||
success: false,
|
||||
error: "Invalid authentication code",
|
||||
};
|
||||
}
|
||||
|
||||
// Mark this code as used to prevent replay
|
||||
await this.#prismaClient.user.update({
|
||||
where: { id: userId },
|
||||
data: {
|
||||
mfaLastUsedCode: hashedCode,
|
||||
},
|
||||
});
|
||||
|
||||
return {
|
||||
success: true,
|
||||
};
|
||||
}
|
||||
|
||||
public async verifyRecoveryCodeForLogin(userId: string, recoveryCode: string) {
|
||||
const user = await this.#prismaClient.user.findFirst({
|
||||
where: { id: userId },
|
||||
});
|
||||
|
||||
if (!user || !user.mfaEnabledAt) {
|
||||
return {
|
||||
success: false,
|
||||
error: "Invalid authentication code",
|
||||
};
|
||||
}
|
||||
|
||||
const hashedCode = await createHash("SHA-512", "hex").digest(recoveryCode);
|
||||
|
||||
// Find an unused recovery code
|
||||
const backupCode = await this.#prismaClient.mfaBackupCode.findFirst({
|
||||
where: {
|
||||
userId: user.id,
|
||||
code: hashedCode,
|
||||
usedAt: null,
|
||||
},
|
||||
});
|
||||
|
||||
if (!backupCode) {
|
||||
return {
|
||||
success: false,
|
||||
error: "Invalid authentication code",
|
||||
};
|
||||
}
|
||||
|
||||
// Mark this recovery code as used
|
||||
await this.#prismaClient.mfaBackupCode.update({
|
||||
where: { id: backupCode.id },
|
||||
data: {
|
||||
usedAt: new Date(),
|
||||
},
|
||||
});
|
||||
|
||||
return {
|
||||
success: true,
|
||||
};
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,49 @@
|
||||
import type { Session } from "@remix-run/node";
|
||||
import { createCookieSessionStorage } from "@remix-run/node";
|
||||
import { env } from "~/env.server";
|
||||
|
||||
export const onboardingSessionStorage = createCookieSessionStorage({
|
||||
cookie: {
|
||||
name: "__onboarding", // use any name you want here
|
||||
sameSite: "lax", // this helps with CSRF
|
||||
path: "/", // remember to add this so the cookie will work in all routes
|
||||
httpOnly: true, // for security reasons, make this cookie http only
|
||||
secrets: [env.SESSION_SECRET],
|
||||
secure: env.NODE_ENV === "production", // enable this in prod only
|
||||
maxAge: 60 * 60 * 24, // 1 day
|
||||
},
|
||||
});
|
||||
|
||||
export function getOnboardingSession(request: Request) {
|
||||
return onboardingSessionStorage.getSession(request.headers.get("Cookie"));
|
||||
}
|
||||
|
||||
export function commitOnboardingSession(session: Session) {
|
||||
return onboardingSessionStorage.commitSession(session);
|
||||
}
|
||||
|
||||
export async function getWorkflowDate(request: Request) {
|
||||
const session = await getOnboardingSession(request);
|
||||
|
||||
const rawWorkflowDate = session.get("workflowDate");
|
||||
|
||||
if (rawWorkflowDate) {
|
||||
return new Date(rawWorkflowDate);
|
||||
}
|
||||
}
|
||||
|
||||
export async function setWorkflowDate(date: Date, request: Request) {
|
||||
const session = await getOnboardingSession(request);
|
||||
|
||||
session.set("workflowDate", date.toISOString());
|
||||
|
||||
return session;
|
||||
}
|
||||
|
||||
export async function clearWorkflowDate(request: Request) {
|
||||
const session = await getOnboardingSession(request);
|
||||
|
||||
session.unset("workflowDate");
|
||||
|
||||
return session;
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
import { prisma } from "~/db.server";
|
||||
import { requireUserId } from "./session.server";
|
||||
|
||||
export async function requireOrganization(request: Request, organizationSlug: string) {
|
||||
const userId = await requireUserId(request);
|
||||
|
||||
const organization = await prisma.organization.findFirst({
|
||||
where: {
|
||||
slug: organizationSlug,
|
||||
members: { some: { userId } },
|
||||
deletedAt: null,
|
||||
},
|
||||
});
|
||||
|
||||
if (!organization) {
|
||||
throw new Response("Organization not found", { status: 404 });
|
||||
}
|
||||
|
||||
return { organization, userId };
|
||||
}
|
||||
@@ -0,0 +1,175 @@
|
||||
import { customAlphabet } from "nanoid";
|
||||
import { z } from "zod";
|
||||
import { prisma } from "~/db.server";
|
||||
import { logger } from "./logger.server";
|
||||
import { hashToken } from "~/utils/tokens.server";
|
||||
|
||||
const tokenValueLength = 40;
|
||||
//lowercase only, removed 0 and l to avoid confusion
|
||||
const tokenGenerator = customAlphabet("123456789abcdefghijkmnopqrstuvwxyz", tokenValueLength);
|
||||
|
||||
// Skip the lastAccessedAt write if the existing value is already within this
|
||||
// window. Eliminates per-auth UPDATE churn on a small narrow hot table; the
|
||||
// settings UI reads this field at human granularity so a few-minute
|
||||
// staleness is fine.
|
||||
export const OAT_LAST_ACCESSED_THROTTLE_MS = 5 * 60 * 1000;
|
||||
|
||||
type CreateOrganizationAccessTokenOptions = {
|
||||
name: string;
|
||||
organizationId: string;
|
||||
expiresAt?: Date;
|
||||
};
|
||||
|
||||
export async function getValidOrganizationAccessTokens(organizationId: string) {
|
||||
const organizationAccessTokens = await prisma.organizationAccessToken.findMany({
|
||||
select: {
|
||||
id: true,
|
||||
name: true,
|
||||
createdAt: true,
|
||||
lastAccessedAt: true,
|
||||
expiresAt: true,
|
||||
},
|
||||
where: {
|
||||
organizationId,
|
||||
revokedAt: null,
|
||||
OR: [{ expiresAt: null }, { expiresAt: { gte: new Date() } }],
|
||||
},
|
||||
});
|
||||
|
||||
return organizationAccessTokens.map((oat) => ({
|
||||
id: oat.id,
|
||||
name: oat.name,
|
||||
createdAt: oat.createdAt,
|
||||
lastAccessedAt: oat.lastAccessedAt,
|
||||
expiresAt: oat.expiresAt,
|
||||
}));
|
||||
}
|
||||
|
||||
export type ObfuscatedOrganizationAccessToken = Awaited<
|
||||
ReturnType<typeof getValidOrganizationAccessTokens>
|
||||
>[number];
|
||||
|
||||
export async function revokeOrganizationAccessToken(tokenId: string) {
|
||||
await prisma.organizationAccessToken.update({
|
||||
where: {
|
||||
id: tokenId,
|
||||
},
|
||||
data: {
|
||||
revokedAt: new Date(),
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
export type OrganizationAccessTokenAuthenticationResult = {
|
||||
organizationId: string;
|
||||
};
|
||||
|
||||
const AuthorizationHeaderSchema = z.string().regex(/^Bearer .+$/);
|
||||
|
||||
export async function authenticateApiRequestWithOrganizationAccessToken(
|
||||
request: Request
|
||||
): Promise<OrganizationAccessTokenAuthenticationResult | undefined> {
|
||||
const token = getOrganizationAccessTokenFromRequest(request);
|
||||
if (!token) {
|
||||
return;
|
||||
}
|
||||
|
||||
return authenticateOrganizationAccessToken(token);
|
||||
}
|
||||
|
||||
function getOrganizationAccessTokenFromRequest(request: Request) {
|
||||
const rawAuthorization = request.headers.get("Authorization");
|
||||
|
||||
const authorization = AuthorizationHeaderSchema.safeParse(rawAuthorization);
|
||||
if (!authorization.success) {
|
||||
return;
|
||||
}
|
||||
|
||||
const organizationAccessToken = authorization.data.replace(/^Bearer /, "");
|
||||
return organizationAccessToken;
|
||||
}
|
||||
|
||||
export async function authenticateOrganizationAccessToken(
|
||||
token: string
|
||||
): Promise<OrganizationAccessTokenAuthenticationResult | undefined> {
|
||||
if (!token.startsWith(tokenPrefix)) {
|
||||
logger.warn(`OAT doesn't start with ${tokenPrefix}`);
|
||||
return;
|
||||
}
|
||||
|
||||
const hashedToken = hashToken(token);
|
||||
|
||||
const organizationAccessToken = await prisma.organizationAccessToken.findFirst({
|
||||
where: {
|
||||
hashedToken,
|
||||
revokedAt: null,
|
||||
OR: [{ expiresAt: null }, { expiresAt: { gte: new Date() } }],
|
||||
},
|
||||
});
|
||||
|
||||
if (!organizationAccessToken) {
|
||||
return;
|
||||
}
|
||||
|
||||
// Conditional updateMany — only writes if the existing lastAccessedAt is
|
||||
// null or older than the throttle window. The WHERE runs inside the UPDATE
|
||||
// so concurrent auths don't race into a double-write. `revokedAt: null`
|
||||
// matches the findFirst guard above so a token revoked between the read
|
||||
// and write doesn't get a stale lastAccessedAt update.
|
||||
await prisma.organizationAccessToken.updateMany({
|
||||
where: {
|
||||
id: organizationAccessToken.id,
|
||||
revokedAt: null,
|
||||
OR: [
|
||||
{ lastAccessedAt: null },
|
||||
{ lastAccessedAt: { lt: new Date(Date.now() - OAT_LAST_ACCESSED_THROTTLE_MS) } },
|
||||
],
|
||||
},
|
||||
data: {
|
||||
lastAccessedAt: new Date(),
|
||||
},
|
||||
});
|
||||
|
||||
return {
|
||||
organizationId: organizationAccessToken.organizationId,
|
||||
};
|
||||
}
|
||||
|
||||
export function isOrganizationAccessToken(token: string) {
|
||||
return token.startsWith(tokenPrefix);
|
||||
}
|
||||
|
||||
export async function createOrganizationAccessToken({
|
||||
name,
|
||||
organizationId,
|
||||
expiresAt,
|
||||
}: CreateOrganizationAccessTokenOptions) {
|
||||
const token = createToken();
|
||||
|
||||
const organizationAccessToken = await prisma.organizationAccessToken.create({
|
||||
data: {
|
||||
name,
|
||||
organizationId,
|
||||
hashedToken: hashToken(token),
|
||||
expiresAt,
|
||||
},
|
||||
});
|
||||
|
||||
return {
|
||||
id: organizationAccessToken.id,
|
||||
name,
|
||||
organizationId,
|
||||
token,
|
||||
expiresAt: organizationAccessToken.expiresAt,
|
||||
};
|
||||
}
|
||||
|
||||
export type CreatedOrganizationAccessToken = Awaited<
|
||||
ReturnType<typeof createOrganizationAccessToken>
|
||||
>;
|
||||
|
||||
const tokenPrefix = "tr_oat_";
|
||||
|
||||
function createToken() {
|
||||
return `${tokenPrefix}${tokenGenerator()}`;
|
||||
}
|
||||
@@ -0,0 +1,467 @@
|
||||
import { type PersonalAccessToken, type User } from "@trigger.dev/database";
|
||||
import { customAlphabet, nanoid } from "nanoid";
|
||||
import { z } from "zod";
|
||||
import { prisma } from "~/db.server";
|
||||
import { logger } from "./logger.server";
|
||||
import { rbac } from "./rbac.server";
|
||||
import { decryptToken, encryptToken, hashToken } from "~/utils/tokens.server";
|
||||
import { env } from "~/env.server";
|
||||
import { isUserActorToken } from "@trigger.dev/rbac";
|
||||
|
||||
const tokenValueLength = 40;
|
||||
//lowercase only, removed 0 and l to avoid confusion
|
||||
const tokenGenerator = customAlphabet("123456789abcdefghijkmnopqrstuvwxyz", tokenValueLength);
|
||||
|
||||
// Skip the lastAccessedAt write if the existing value is already within this
|
||||
// window. Eliminates per-auth UPDATE churn on a small narrow hot table; the
|
||||
// /account/tokens UI reads this field at human granularity so a few-minute
|
||||
// staleness is fine.
|
||||
export const PAT_LAST_ACCESSED_THROTTLE_MS = 5 * 60 * 1000;
|
||||
|
||||
type CreatePersonalAccessTokenOptions = {
|
||||
name: string;
|
||||
userId: string;
|
||||
// Optional: when provided, persist a TokenRole row alongside the PAT
|
||||
// so PAT-authenticated requests pick up that role's permissions
|
||||
// (TRI-8749). The dashboard tokens page passes a chosen system role;
|
||||
// the CLI auth-code path doesn't pass one (legacy behaviour
|
||||
// preserved — those PATs run with no explicit role).
|
||||
roleId?: string;
|
||||
};
|
||||
|
||||
/** Returns obfuscated access tokens that aren't revoked */
|
||||
export async function getValidPersonalAccessTokens(userId: string) {
|
||||
const personalAccessTokens = await prisma.personalAccessToken.findMany({
|
||||
select: {
|
||||
id: true,
|
||||
name: true,
|
||||
obfuscatedToken: true,
|
||||
createdAt: true,
|
||||
lastAccessedAt: true,
|
||||
},
|
||||
where: {
|
||||
userId,
|
||||
revokedAt: null,
|
||||
},
|
||||
});
|
||||
|
||||
return personalAccessTokens.map((pat) => ({
|
||||
id: pat.id,
|
||||
name: pat.name,
|
||||
obfuscatedToken: pat.obfuscatedToken,
|
||||
createdAt: pat.createdAt,
|
||||
lastAccessedAt: pat.lastAccessedAt,
|
||||
}));
|
||||
}
|
||||
|
||||
export type ObfuscatedPersonalAccessToken = Awaited<
|
||||
ReturnType<typeof getValidPersonalAccessTokens>
|
||||
>[number];
|
||||
|
||||
/** Gets a PersonalAccessToken from an Auth Code, this only works within 10 mins of the auth code being created */
|
||||
export async function getPersonalAccessTokenFromAuthorizationCode(authorizationCode: string) {
|
||||
//only allow authorization codes that were created less than 10 mins ago
|
||||
const tenMinutesAgo = new Date(Date.now() - 10 * 60 * 1000);
|
||||
const code = await prisma.authorizationCode.findUnique({
|
||||
select: {
|
||||
personalAccessToken: true,
|
||||
},
|
||||
where: {
|
||||
code: authorizationCode,
|
||||
createdAt: {
|
||||
gte: tenMinutesAgo,
|
||||
},
|
||||
},
|
||||
});
|
||||
if (!code) {
|
||||
throw new Error("Invalid authorization code, or code expired");
|
||||
}
|
||||
|
||||
//there's no PersonalAccessToken associated with this code
|
||||
if (!code.personalAccessToken) {
|
||||
return {
|
||||
token: null,
|
||||
};
|
||||
}
|
||||
|
||||
const decryptedToken = decryptPersonalAccessToken(code.personalAccessToken);
|
||||
return {
|
||||
token: {
|
||||
token: decryptedToken,
|
||||
obfuscatedToken: code.personalAccessToken.obfuscatedToken,
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
export async function revokePersonalAccessToken(tokenId: string, userId: string) {
|
||||
const result = await prisma.personalAccessToken.updateMany({
|
||||
where: {
|
||||
id: tokenId,
|
||||
userId,
|
||||
},
|
||||
data: {
|
||||
revokedAt: new Date(),
|
||||
},
|
||||
});
|
||||
|
||||
if (result.count === 0) {
|
||||
throw new Error("PAT not found or already revoked");
|
||||
}
|
||||
}
|
||||
|
||||
export type PersonalAccessTokenAuthenticationResult = {
|
||||
userId: string;
|
||||
};
|
||||
|
||||
/**
|
||||
* Smart-skip the `lastAccessedAt` write when the cached value is already
|
||||
* within the throttle window. Saves one DB roundtrip per "fresh" auth.
|
||||
*
|
||||
* Two layers of throttling: JS-side (`Date.now() - lastAccessedAt.getTime()`)
|
||||
* elides the SQL entirely when the caller already has a recent timestamp;
|
||||
* SQL-side `WHERE` clause inside the `updateMany` guards against concurrent
|
||||
* auths racing to a double-write when the JS check decides to fire.
|
||||
*
|
||||
* Called from both the legacy PAT flow (`authenticatePersonalAccessToken`)
|
||||
* and the apiBuilder's RBAC-routed PAT flow (which gets `lastAccessedAt`
|
||||
* from `rbac.authenticatePat`'s result). Keeps the throttle policy in one
|
||||
* place rather than expecting every plugin implementer to re-implement it.
|
||||
*/
|
||||
export async function updateLastAccessedAtIfStale(
|
||||
tokenId: string,
|
||||
lastAccessedAt: Date | null
|
||||
): Promise<void> {
|
||||
if (lastAccessedAt && Date.now() - lastAccessedAt.getTime() <= PAT_LAST_ACCESSED_THROTTLE_MS) {
|
||||
return; // fresh — no roundtrip
|
||||
}
|
||||
await prisma.personalAccessToken.updateMany({
|
||||
where: {
|
||||
id: tokenId,
|
||||
revokedAt: null,
|
||||
OR: [
|
||||
{ lastAccessedAt: null },
|
||||
{ lastAccessedAt: { lt: new Date(Date.now() - PAT_LAST_ACCESSED_THROTTLE_MS) } },
|
||||
],
|
||||
},
|
||||
data: { lastAccessedAt: new Date() },
|
||||
});
|
||||
}
|
||||
|
||||
const EncryptedSecretValueSchema = z.object({
|
||||
nonce: z.string(),
|
||||
ciphertext: z.string(),
|
||||
tag: z.string(),
|
||||
});
|
||||
|
||||
const AuthorizationHeaderSchema = z.string().regex(/^Bearer .+$/);
|
||||
|
||||
export async function authenticateApiRequestWithPersonalAccessToken(
|
||||
request: Request
|
||||
): Promise<PersonalAccessTokenAuthenticationResult | undefined> {
|
||||
const token = getPersonalAccessTokenFromRequest(request);
|
||||
if (!token) {
|
||||
return;
|
||||
}
|
||||
|
||||
// A user-actor token authenticates as the user wherever a PAT does.
|
||||
// The plugin verifies it (identity path → no org context to floor against).
|
||||
if (isUserActorToken(token)) {
|
||||
const result = await rbac.authenticateUserActor(request, {});
|
||||
return result.ok ? { userId: result.userId } : undefined;
|
||||
}
|
||||
|
||||
return authenticatePersonalAccessToken(token);
|
||||
}
|
||||
|
||||
export type AdminAuthenticationResult =
|
||||
| { ok: true; user: User }
|
||||
| { ok: false; status: 401 | 403; message: string };
|
||||
|
||||
/**
|
||||
* Authenticates a request via personal access token and checks the user is
|
||||
* an admin. Returns a discriminated result so callers can shape the failure
|
||||
* (throw a Response, wrap in neverthrow, return JSON, etc.) to fit their
|
||||
* context. See `requireAdminApiRequest` for the Remix loader/action wrapper.
|
||||
*/
|
||||
export async function authenticateAdminRequest(
|
||||
request: Request
|
||||
): Promise<AdminAuthenticationResult> {
|
||||
const authResult = await authenticateApiRequestWithPersonalAccessToken(request);
|
||||
|
||||
if (!authResult) {
|
||||
return { ok: false, status: 401, message: "Invalid or Missing API key" };
|
||||
}
|
||||
|
||||
const user = await prisma.user.findFirst({
|
||||
where: { id: authResult.userId },
|
||||
});
|
||||
|
||||
if (!user) {
|
||||
return { ok: false, status: 401, message: "Invalid or Missing API key" };
|
||||
}
|
||||
|
||||
if (!user.admin) {
|
||||
return { ok: false, status: 403, message: "You must be an admin to perform this action" };
|
||||
}
|
||||
|
||||
return { ok: true, user };
|
||||
}
|
||||
|
||||
/**
|
||||
* Remix loader/action wrapper around `authenticateAdminRequest` that throws
|
||||
* a Response on failure so routes can `await` without handling the error
|
||||
* branch. Uses `new Response` directly to avoid coupling this module to
|
||||
* `@remix-run/server-runtime`.
|
||||
*/
|
||||
export async function requireAdminApiRequest(request: Request): Promise<User> {
|
||||
const result = await authenticateAdminRequest(request);
|
||||
|
||||
if (!result.ok) {
|
||||
throw new Response(JSON.stringify({ error: result.message }), {
|
||||
status: result.status,
|
||||
headers: { "Content-Type": "application/json" },
|
||||
});
|
||||
}
|
||||
|
||||
return result.user;
|
||||
}
|
||||
|
||||
function getPersonalAccessTokenFromRequest(request: Request) {
|
||||
const rawAuthorization = request.headers.get("Authorization");
|
||||
|
||||
const authorization = AuthorizationHeaderSchema.safeParse(rawAuthorization);
|
||||
if (!authorization.success) {
|
||||
return;
|
||||
}
|
||||
|
||||
const personalAccessToken = authorization.data.replace(/^Bearer /, "");
|
||||
return personalAccessToken;
|
||||
}
|
||||
|
||||
export async function authenticatePersonalAccessToken(
|
||||
token: string
|
||||
): Promise<PersonalAccessTokenAuthenticationResult | undefined> {
|
||||
if (!token.startsWith(tokenPrefix)) {
|
||||
logger.warn(`PAT doesn't start with ${tokenPrefix}`);
|
||||
return;
|
||||
}
|
||||
|
||||
const hashedToken = hashToken(token);
|
||||
|
||||
const personalAccessToken = await prisma.personalAccessToken.findFirst({
|
||||
where: {
|
||||
hashedToken,
|
||||
revokedAt: null,
|
||||
},
|
||||
});
|
||||
|
||||
if (!personalAccessToken) {
|
||||
// The token may have been revoked or is entirely invalid
|
||||
return;
|
||||
}
|
||||
|
||||
await updateLastAccessedAtIfStale(personalAccessToken.id, personalAccessToken.lastAccessedAt);
|
||||
|
||||
const decryptedToken = decryptPersonalAccessToken(personalAccessToken);
|
||||
|
||||
if (decryptedToken !== token) {
|
||||
logger.error(
|
||||
`PersonalAccessToken with id: ${personalAccessToken.id} was found in the database with hash ${hashedToken}, but the decrypted token did not match the provided token.`
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
return {
|
||||
userId: personalAccessToken.userId,
|
||||
};
|
||||
}
|
||||
|
||||
export function isPersonalAccessToken(token: string) {
|
||||
return token.startsWith(tokenPrefix);
|
||||
}
|
||||
|
||||
export function createAuthorizationCode() {
|
||||
return prisma.authorizationCode.create({
|
||||
data: {
|
||||
code: nanoid(64),
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
/** Creates a PersonalAccessToken from an Auth Code, and return the token. We only ever return the unencrypted token once. */
|
||||
export async function createPersonalAccessTokenFromAuthorizationCode(
|
||||
authorizationCode: string,
|
||||
userId: string
|
||||
) {
|
||||
//only allow authorization codes that were created less than 10 mins ago
|
||||
const tenMinutesAgo = new Date(Date.now() - 10 * 60 * 1000);
|
||||
const code = await prisma.authorizationCode.findUnique({
|
||||
where: {
|
||||
code: authorizationCode,
|
||||
personalAccessTokenId: null,
|
||||
createdAt: {
|
||||
gte: tenMinutesAgo,
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
if (!code) {
|
||||
throw new Error("Invalid authorization code, code already used, or code expired");
|
||||
}
|
||||
|
||||
const existingCliPersonalAccessToken = await prisma.personalAccessToken.findFirst({
|
||||
where: {
|
||||
userId,
|
||||
name: "cli",
|
||||
},
|
||||
});
|
||||
|
||||
//we only allow you to have one CLI PAT at a time, so return this
|
||||
if (existingCliPersonalAccessToken) {
|
||||
//associate this authorization code with the existing personal access token
|
||||
await prisma.authorizationCode.update({
|
||||
where: {
|
||||
code: authorizationCode,
|
||||
},
|
||||
data: {
|
||||
personalAccessTokenId: existingCliPersonalAccessToken.id,
|
||||
},
|
||||
});
|
||||
|
||||
if (existingCliPersonalAccessToken.revokedAt) {
|
||||
// re-activate revoked CLI PAT so we can use it again
|
||||
await prisma.personalAccessToken.update({
|
||||
where: {
|
||||
id: existingCliPersonalAccessToken.id,
|
||||
},
|
||||
data: {
|
||||
revokedAt: null,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
//we don't return the decrypted token
|
||||
return {
|
||||
id: existingCliPersonalAccessToken.id,
|
||||
name: existingCliPersonalAccessToken.name,
|
||||
userId: existingCliPersonalAccessToken.userId,
|
||||
obfuscateToken: existingCliPersonalAccessToken.obfuscatedToken,
|
||||
};
|
||||
}
|
||||
|
||||
const token = await createPersonalAccessToken({
|
||||
name: "cli",
|
||||
userId,
|
||||
});
|
||||
|
||||
await prisma.authorizationCode.update({
|
||||
where: {
|
||||
code: authorizationCode,
|
||||
},
|
||||
data: {
|
||||
personalAccessTokenId: token.id,
|
||||
},
|
||||
});
|
||||
|
||||
return token;
|
||||
}
|
||||
|
||||
/** Created a new PersonalAccessToken, and return the token. We only ever return the unencrypted token once. */
|
||||
export async function createPersonalAccessToken({
|
||||
name,
|
||||
userId,
|
||||
roleId,
|
||||
}: CreatePersonalAccessTokenOptions) {
|
||||
const token = createToken();
|
||||
const encryptedToken = encryptToken(token, env.ENCRYPTION_KEY);
|
||||
|
||||
const personalAccessToken = await prisma.personalAccessToken.create({
|
||||
data: {
|
||||
name,
|
||||
userId,
|
||||
encryptedToken,
|
||||
obfuscatedToken: obfuscateToken(token),
|
||||
hashedToken: hashToken(token),
|
||||
},
|
||||
});
|
||||
|
||||
// Persist the role choice via the RBAC plugin's setTokenRole. The
|
||||
// plugin may store this in a separate datastore from Prisma (e.g.
|
||||
// Drizzle on a different schema), so co-transactional inserts are
|
||||
// awkward — we use a compensating-delete pattern instead: if
|
||||
// setTokenRole fails, roll back the PAT row by deleting it. The auth
|
||||
// path treats "no role" as permissive (matches the default fallback)
|
||||
// so a brief orphan window between the two writes is harmless. The
|
||||
// compensating delete narrows that window from "until manual cleanup"
|
||||
// to "until the request returns".
|
||||
//
|
||||
// Skip the call entirely when no RBAC plugin is loaded — the OSS
|
||||
// fallback has no TokenRole table to write to. Gating on
|
||||
// `rbac.isUsingPlugin()` (rather than parsing the fallback's error
|
||||
// string) keeps the OSS-vs-cloud branch explicit and decoupled from
|
||||
// any specific error message.
|
||||
if (roleId && (await rbac.isUsingPlugin())) {
|
||||
const roleResult = await rbac.setTokenRole({
|
||||
tokenId: personalAccessToken.id,
|
||||
roleId,
|
||||
});
|
||||
if (!roleResult.ok) {
|
||||
await prisma.personalAccessToken
|
||||
.delete({ where: { id: personalAccessToken.id } })
|
||||
.catch((err) => {
|
||||
logger.error("Failed to compensating-delete PAT after TokenRole insert failed", {
|
||||
patId: personalAccessToken.id,
|
||||
roleResultError: roleResult.error,
|
||||
deleteError: err instanceof Error ? err.message : String(err),
|
||||
});
|
||||
});
|
||||
throw new Error(`Failed to assign role to access token: ${roleResult.error}`);
|
||||
}
|
||||
} else if (roleId) {
|
||||
logger.debug("createPersonalAccessToken: no RBAC plugin, skipping role assignment", {
|
||||
patId: personalAccessToken.id,
|
||||
userId,
|
||||
});
|
||||
}
|
||||
|
||||
return {
|
||||
id: personalAccessToken.id,
|
||||
name,
|
||||
userId,
|
||||
token,
|
||||
obfuscatedToken: personalAccessToken.obfuscatedToken,
|
||||
};
|
||||
}
|
||||
|
||||
export type CreatedPersonalAccessToken = Awaited<ReturnType<typeof createPersonalAccessToken>>;
|
||||
|
||||
const tokenPrefix = "tr_pat_";
|
||||
|
||||
/** Creates a PersonalAccessToken that starts with tr_pat_ */
|
||||
function createToken() {
|
||||
return `${tokenPrefix}${tokenGenerator()}`;
|
||||
}
|
||||
|
||||
/** Obfuscates all but the first and last 4 characters of the token, so it looks like tr_pat_bhbd•••••••••••••••••••fd4a */
|
||||
function obfuscateToken(token: string) {
|
||||
const withoutPrefix = token.replace(tokenPrefix, "");
|
||||
const obfuscated = `${withoutPrefix.slice(0, 4)}${"•".repeat(18)}${withoutPrefix.slice(-4)}`;
|
||||
return `${tokenPrefix}${obfuscated}`;
|
||||
}
|
||||
|
||||
function decryptPersonalAccessToken(personalAccessToken: PersonalAccessToken) {
|
||||
const encryptedData = EncryptedSecretValueSchema.safeParse(personalAccessToken.encryptedToken);
|
||||
if (!encryptedData.success) {
|
||||
throw new Error(
|
||||
`Unable to parse encrypted PersonalAccessToken with id: ${personalAccessToken.id}: ${encryptedData.error.message}`
|
||||
);
|
||||
}
|
||||
|
||||
const decryptedToken = decryptToken(
|
||||
encryptedData.data.nonce,
|
||||
encryptedData.data.ciphertext,
|
||||
encryptedData.data.tag,
|
||||
env.ENCRYPTION_KEY
|
||||
);
|
||||
return decryptedToken;
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,47 @@
|
||||
import { Redis } from "ioredis";
|
||||
import { defaultReconnectOnError } from "@internal/redis";
|
||||
import { env } from "~/env.server";
|
||||
import { singleton } from "~/utils/singleton";
|
||||
import { logger } from "./logger.server";
|
||||
|
||||
const KEY_PREFIX = "cli-notif-ctr:";
|
||||
const MAX_COUNTER = 1000;
|
||||
|
||||
function initializeRedis(): Redis | undefined {
|
||||
const host = env.CACHE_REDIS_HOST;
|
||||
if (!host) return undefined;
|
||||
|
||||
return new Redis({
|
||||
connectionName: "platformNotificationCounter",
|
||||
host,
|
||||
port: env.CACHE_REDIS_PORT,
|
||||
username: env.CACHE_REDIS_USERNAME,
|
||||
password: env.CACHE_REDIS_PASSWORD,
|
||||
keyPrefix: "tr:",
|
||||
enableAutoPipelining: true,
|
||||
reconnectOnError: defaultReconnectOnError,
|
||||
...(env.CACHE_REDIS_TLS_DISABLED === "true" ? {} : { tls: {} }),
|
||||
});
|
||||
}
|
||||
|
||||
const redis = singleton("platformNotificationCounter", initializeRedis);
|
||||
|
||||
/** Increment and return the user's CLI request counter (0-based, wraps at 1000→0). */
|
||||
export async function incrementCliRequestCounter(userId: string): Promise<number> {
|
||||
if (!redis) return 0;
|
||||
|
||||
try {
|
||||
const key = `${KEY_PREFIX}${userId}`;
|
||||
const value = await redis.incr(key);
|
||||
|
||||
if (value > MAX_COUNTER) {
|
||||
await redis.set(key, "0");
|
||||
return 0;
|
||||
}
|
||||
|
||||
return value;
|
||||
} catch (error) {
|
||||
logger.error("Failed to increment CLI notification counter", { userId, error });
|
||||
return 0;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,815 @@
|
||||
import { z } from "zod";
|
||||
import { errAsync, fromPromise, type ResultAsync } from "neverthrow";
|
||||
import { prisma } from "~/db.server";
|
||||
import {
|
||||
type PlatformNotificationScope,
|
||||
type PlatformNotificationSurface,
|
||||
} from "@trigger.dev/database";
|
||||
import { incrementCliRequestCounter } from "./platformNotificationCounter.server";
|
||||
|
||||
// --- Payload schema (spec v1) ---
|
||||
|
||||
const DiscoverySchema = z.object({
|
||||
filePatterns: z.array(z.string().min(1)).min(1),
|
||||
contentPattern: z
|
||||
.string()
|
||||
.max(200)
|
||||
.optional()
|
||||
.refine(
|
||||
(val) => {
|
||||
if (!val) return true;
|
||||
try {
|
||||
new RegExp(val);
|
||||
return true;
|
||||
} catch {
|
||||
return false;
|
||||
}
|
||||
},
|
||||
{ message: "contentPattern must be a valid regular expression" }
|
||||
),
|
||||
matchBehavior: z.enum(["show-if-found", "show-if-not-found"]),
|
||||
});
|
||||
|
||||
// Constrain URL fields to http/https; `.url()` alone accepts other schemes
|
||||
// that would be unsafe to render into an `<a href>`.
|
||||
const httpUrl = z
|
||||
.string()
|
||||
.url()
|
||||
.refine(
|
||||
(v) => {
|
||||
try {
|
||||
const proto = new URL(v).protocol;
|
||||
return proto === "http:" || proto === "https:";
|
||||
} catch {
|
||||
return false;
|
||||
}
|
||||
},
|
||||
{ message: "URL must use http or https" }
|
||||
);
|
||||
|
||||
const CardDataV1Schema = z.object({
|
||||
type: z.enum(["card", "info", "warn", "error", "success", "changelog"]),
|
||||
title: z.string(),
|
||||
description: z.string(),
|
||||
image: httpUrl.optional(),
|
||||
actionLabel: z.string().optional(),
|
||||
actionUrl: httpUrl.optional(),
|
||||
dismissOnAction: z.boolean().optional(),
|
||||
discovery: DiscoverySchema.optional(),
|
||||
});
|
||||
|
||||
const PayloadV1Schema = z.object({
|
||||
version: z.literal("1"),
|
||||
data: CardDataV1Schema,
|
||||
});
|
||||
|
||||
export type PayloadV1 = z.infer<typeof PayloadV1Schema>;
|
||||
|
||||
export type PlatformNotificationWithPayload = {
|
||||
id: string;
|
||||
friendlyId: string;
|
||||
scope: string;
|
||||
priority: number;
|
||||
payload: PayloadV1;
|
||||
isRead: boolean;
|
||||
};
|
||||
|
||||
// --- Read: admin list with interaction stats ---
|
||||
|
||||
export async function getAdminNotificationsList({
|
||||
page = 1,
|
||||
pageSize = 20,
|
||||
hideInactive = false,
|
||||
}: {
|
||||
page?: number;
|
||||
pageSize?: number;
|
||||
hideInactive?: boolean;
|
||||
}) {
|
||||
const where = hideInactive ? { archivedAt: null, endsAt: { gt: new Date() } } : {};
|
||||
|
||||
const [notifications, total] = await Promise.all([
|
||||
prisma.platformNotification.findMany({
|
||||
where,
|
||||
orderBy: [{ createdAt: "desc" }],
|
||||
skip: (page - 1) * pageSize,
|
||||
take: pageSize,
|
||||
include: {
|
||||
_count: {
|
||||
select: { interactions: true },
|
||||
},
|
||||
interactions: {
|
||||
select: {
|
||||
webappDismissedAt: true,
|
||||
webappClickedAt: true,
|
||||
cliDismissedAt: true,
|
||||
},
|
||||
},
|
||||
},
|
||||
}),
|
||||
prisma.platformNotification.count({ where }),
|
||||
]);
|
||||
|
||||
return {
|
||||
notifications: notifications.map((n) => {
|
||||
const parsed = PayloadV1Schema.safeParse(n.payload);
|
||||
return {
|
||||
id: n.id,
|
||||
friendlyId: n.friendlyId,
|
||||
title: n.title,
|
||||
surface: n.surface,
|
||||
scope: n.scope,
|
||||
userId: n.userId,
|
||||
organizationId: n.organizationId,
|
||||
projectId: n.projectId,
|
||||
priority: n.priority,
|
||||
startsAt: n.startsAt,
|
||||
endsAt: n.endsAt,
|
||||
archivedAt: n.archivedAt,
|
||||
createdAt: n.createdAt,
|
||||
payload: n.payload,
|
||||
payloadTitle: parsed.success ? parsed.data.data.title : null,
|
||||
payloadType: parsed.success ? parsed.data.data.type : null,
|
||||
payloadDescription: parsed.success ? parsed.data.data.description : null,
|
||||
payloadActionUrl: parsed.success ? parsed.data.data.actionUrl : null,
|
||||
payloadImage: parsed.success ? parsed.data.data.image : null,
|
||||
payloadDismissOnAction: parsed.success
|
||||
? (parsed.data.data.dismissOnAction ?? false)
|
||||
: false,
|
||||
payloadDiscovery: parsed.success ? (parsed.data.data.discovery ?? null) : null,
|
||||
cliMaxShowCount: n.cliMaxShowCount,
|
||||
cliMaxDaysAfterFirstSeen: n.cliMaxDaysAfterFirstSeen,
|
||||
cliShowEvery: n.cliShowEvery,
|
||||
stats: {
|
||||
seen: n._count.interactions,
|
||||
clicked: n.interactions.filter((i) => i.webappClickedAt !== null).length,
|
||||
dismissed: n.interactions.filter(
|
||||
(i) => i.webappDismissedAt !== null || i.cliDismissedAt !== null
|
||||
).length,
|
||||
},
|
||||
};
|
||||
}),
|
||||
total,
|
||||
page,
|
||||
pageCount: Math.ceil(total / pageSize),
|
||||
};
|
||||
}
|
||||
|
||||
// --- Read: active notifications for webapp ---
|
||||
|
||||
export async function getActivePlatformNotifications({
|
||||
userId,
|
||||
organizationId,
|
||||
projectId,
|
||||
}: {
|
||||
userId: string;
|
||||
organizationId: string;
|
||||
projectId?: string;
|
||||
}) {
|
||||
const now = new Date();
|
||||
|
||||
const notifications = await prisma.platformNotification.findMany({
|
||||
where: {
|
||||
surface: "WEBAPP",
|
||||
archivedAt: null,
|
||||
startsAt: { lte: now },
|
||||
endsAt: { gt: now },
|
||||
AND: [
|
||||
{
|
||||
OR: [
|
||||
{ scope: "GLOBAL" },
|
||||
{ scope: "ORGANIZATION", organizationId },
|
||||
...(projectId ? [{ scope: "PROJECT" as const, projectId }] : []),
|
||||
{ scope: "USER", userId },
|
||||
],
|
||||
},
|
||||
],
|
||||
},
|
||||
include: {
|
||||
interactions: {
|
||||
where: { userId },
|
||||
},
|
||||
},
|
||||
orderBy: [{ priority: "desc" }, { createdAt: "desc" }],
|
||||
});
|
||||
|
||||
type InternalNotification = PlatformNotificationWithPayload & { createdAt: Date };
|
||||
const result: InternalNotification[] = [];
|
||||
|
||||
for (const n of notifications) {
|
||||
const interaction = n.interactions[0] ?? null;
|
||||
|
||||
if (interaction?.webappDismissedAt) continue;
|
||||
|
||||
const parsed = PayloadV1Schema.safeParse(n.payload);
|
||||
if (!parsed.success) continue;
|
||||
|
||||
result.push({
|
||||
id: n.id,
|
||||
friendlyId: n.friendlyId,
|
||||
scope: n.scope,
|
||||
priority: n.priority,
|
||||
createdAt: n.createdAt,
|
||||
payload: parsed.data,
|
||||
isRead: !!interaction,
|
||||
});
|
||||
}
|
||||
|
||||
result.sort(compareNotifications);
|
||||
|
||||
const unreadCount = result.filter((n) => !n.isRead).length;
|
||||
const notifications_out: PlatformNotificationWithPayload[] = result.map(
|
||||
({ createdAt: _, ...rest }) => rest
|
||||
);
|
||||
|
||||
return { notifications: notifications_out, unreadCount };
|
||||
}
|
||||
|
||||
function compareNotifications(
|
||||
a: { priority: number; createdAt: Date },
|
||||
b: { priority: number; createdAt: Date }
|
||||
) {
|
||||
const priorityDiff = b.priority - a.priority;
|
||||
if (priorityDiff !== 0) return priorityDiff;
|
||||
|
||||
return b.createdAt.getTime() - a.createdAt.getTime();
|
||||
}
|
||||
|
||||
// --- Write: upsert interaction ---
|
||||
|
||||
async function upsertInteraction({
|
||||
notificationId,
|
||||
userId,
|
||||
onUpdate,
|
||||
onCreate,
|
||||
}: {
|
||||
notificationId: string;
|
||||
userId: string;
|
||||
onUpdate: Record<string, unknown>;
|
||||
onCreate: Record<string, unknown>;
|
||||
}) {
|
||||
await prisma.platformNotificationInteraction.upsert({
|
||||
where: { notificationId_userId: { notificationId, userId } },
|
||||
update: onUpdate,
|
||||
create: {
|
||||
notificationId,
|
||||
userId,
|
||||
firstSeenAt: new Date(),
|
||||
showCount: 1,
|
||||
...onCreate,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
export async function recordNotificationSeen({
|
||||
notificationId,
|
||||
userId,
|
||||
}: {
|
||||
notificationId: string;
|
||||
userId: string;
|
||||
}) {
|
||||
return upsertInteraction({
|
||||
notificationId,
|
||||
userId,
|
||||
onUpdate: { showCount: { increment: 1 } },
|
||||
onCreate: {},
|
||||
});
|
||||
}
|
||||
|
||||
export async function dismissNotification({
|
||||
notificationId,
|
||||
userId,
|
||||
}: {
|
||||
notificationId: string;
|
||||
userId: string;
|
||||
}) {
|
||||
const now = new Date();
|
||||
return upsertInteraction({
|
||||
notificationId,
|
||||
userId,
|
||||
onUpdate: { webappDismissedAt: now },
|
||||
onCreate: { webappDismissedAt: now },
|
||||
});
|
||||
}
|
||||
|
||||
export async function recordNotificationClicked({
|
||||
notificationId,
|
||||
userId,
|
||||
}: {
|
||||
notificationId: string;
|
||||
userId: string;
|
||||
}) {
|
||||
const now = new Date();
|
||||
return upsertInteraction({
|
||||
notificationId,
|
||||
userId,
|
||||
onUpdate: { webappClickedAt: now },
|
||||
onCreate: { webappClickedAt: now },
|
||||
});
|
||||
}
|
||||
|
||||
// --- Membership verification ---
|
||||
|
||||
export async function verifyOrgMembership({
|
||||
userId,
|
||||
organizationId,
|
||||
projectId,
|
||||
}: {
|
||||
userId: string;
|
||||
organizationId?: string;
|
||||
projectId?: string;
|
||||
}): Promise<{ organizationId?: string; projectId?: string }> {
|
||||
if (!organizationId) return {};
|
||||
|
||||
const membership = await prisma.orgMember.findFirst({
|
||||
where: { userId, organizationId },
|
||||
select: { organizationId: true },
|
||||
});
|
||||
|
||||
if (!membership) return {};
|
||||
|
||||
if (projectId) {
|
||||
const project = await prisma.project.findFirst({
|
||||
where: { id: projectId, organizationId, deletedAt: null },
|
||||
select: { id: true },
|
||||
});
|
||||
if (!project) return { organizationId };
|
||||
}
|
||||
|
||||
return { organizationId, projectId };
|
||||
}
|
||||
|
||||
// --- Read: recent changelogs (for Help & Feedback) ---
|
||||
|
||||
export async function getRecentChangelogs({
|
||||
userId,
|
||||
organizationId,
|
||||
projectId,
|
||||
limit = 2,
|
||||
}: {
|
||||
userId: string;
|
||||
organizationId?: string;
|
||||
projectId?: string;
|
||||
limit?: number;
|
||||
}) {
|
||||
// NOTE: Intentionally not filtering by archivedAt or endsAt.
|
||||
// We want to show archived and expired changelogs in the "What's new" section
|
||||
// so users can still find recent release notes.
|
||||
// We DO filter by scope (to prevent user-scoped changelogs leaking to others)
|
||||
// and by startsAt (to hide changelogs scheduled for the future).
|
||||
const notifications = await prisma.platformNotification.findMany({
|
||||
where: {
|
||||
surface: "WEBAPP",
|
||||
payload: { path: ["data", "type"], equals: "changelog" },
|
||||
startsAt: { lte: new Date() },
|
||||
OR: [
|
||||
{ scope: "GLOBAL" },
|
||||
{ scope: "USER", userId },
|
||||
...(organizationId ? [{ scope: "ORGANIZATION" as const, organizationId }] : []),
|
||||
...(projectId ? [{ scope: "PROJECT" as const, projectId }] : []),
|
||||
],
|
||||
},
|
||||
orderBy: [{ createdAt: "desc" }],
|
||||
take: limit,
|
||||
});
|
||||
|
||||
return notifications
|
||||
.map((n) => {
|
||||
const parsed = PayloadV1Schema.safeParse(n.payload);
|
||||
if (!parsed.success) return null;
|
||||
return { id: n.id, title: parsed.data.data.title, actionUrl: parsed.data.data.actionUrl };
|
||||
})
|
||||
.filter(Boolean) as Array<{ id: string; title: string; actionUrl?: string }>;
|
||||
}
|
||||
|
||||
// --- CLI: next notification for CLI surface ---
|
||||
|
||||
function isCliNotificationExpired(
|
||||
interaction: {
|
||||
userId: string;
|
||||
firstSeenAt: Date;
|
||||
showCount: number;
|
||||
cliDismissedAt: Date | null;
|
||||
} | null,
|
||||
notification: {
|
||||
id: string;
|
||||
cliMaxDaysAfterFirstSeen: number | null;
|
||||
cliMaxShowCount: number | null;
|
||||
}
|
||||
): boolean {
|
||||
if (!interaction) return false;
|
||||
|
||||
let expired = false;
|
||||
|
||||
if (
|
||||
notification.cliMaxShowCount !== null &&
|
||||
interaction.showCount >= notification.cliMaxShowCount
|
||||
) {
|
||||
expired = true;
|
||||
}
|
||||
|
||||
if (!expired && notification.cliMaxDaysAfterFirstSeen !== null) {
|
||||
const daysSinceFirstSeen =
|
||||
(Date.now() - interaction.firstSeenAt.getTime()) / (1000 * 60 * 60 * 24);
|
||||
if (daysSinceFirstSeen > notification.cliMaxDaysAfterFirstSeen) {
|
||||
expired = true;
|
||||
}
|
||||
}
|
||||
|
||||
// For time-based expiration, persist the dismiss on the next request
|
||||
// (showCount-based dismissal is handled inline at display time)
|
||||
if (expired && !interaction.cliDismissedAt) {
|
||||
void prisma.platformNotificationInteraction.update({
|
||||
where: {
|
||||
notificationId_userId: {
|
||||
notificationId: notification.id,
|
||||
userId: interaction.userId,
|
||||
},
|
||||
},
|
||||
data: { cliDismissedAt: new Date() },
|
||||
});
|
||||
}
|
||||
|
||||
return expired;
|
||||
}
|
||||
|
||||
export async function getNextCliNotification({
|
||||
userId,
|
||||
projectRef,
|
||||
}: {
|
||||
userId: string;
|
||||
projectRef?: string;
|
||||
}): Promise<{
|
||||
id: string;
|
||||
payload: PayloadV1;
|
||||
showCount: number;
|
||||
firstSeenAt: string;
|
||||
} | null> {
|
||||
const now = new Date();
|
||||
|
||||
// Resolve organizationId and projectId from projectRef if provided
|
||||
let organizationId: string | undefined;
|
||||
let projectId: string | undefined;
|
||||
|
||||
if (projectRef) {
|
||||
const project = await prisma.project.findFirst({
|
||||
where: {
|
||||
externalRef: projectRef,
|
||||
deletedAt: null,
|
||||
organization: {
|
||||
deletedAt: null,
|
||||
members: { some: { userId } },
|
||||
},
|
||||
},
|
||||
select: { id: true, organizationId: true },
|
||||
});
|
||||
|
||||
if (project) {
|
||||
projectId = project.id;
|
||||
organizationId = project.organizationId;
|
||||
}
|
||||
}
|
||||
|
||||
// If no projectRef or project not found, get org from membership
|
||||
if (!organizationId) {
|
||||
const membership = await prisma.orgMember.findFirst({
|
||||
where: { userId },
|
||||
select: { organizationId: true },
|
||||
});
|
||||
if (membership) {
|
||||
organizationId = membership.organizationId;
|
||||
}
|
||||
}
|
||||
|
||||
const scopeFilter: Array<Record<string, unknown>> = [
|
||||
{ scope: "GLOBAL" },
|
||||
{ scope: "USER", userId },
|
||||
];
|
||||
|
||||
if (organizationId) {
|
||||
scopeFilter.push({ scope: "ORGANIZATION", organizationId });
|
||||
}
|
||||
|
||||
if (projectId) {
|
||||
scopeFilter.push({ scope: "PROJECT", projectId });
|
||||
}
|
||||
|
||||
const notifications = await prisma.platformNotification.findMany({
|
||||
where: {
|
||||
surface: "CLI",
|
||||
archivedAt: null,
|
||||
startsAt: { lte: now },
|
||||
endsAt: { gt: now },
|
||||
AND: [{ OR: scopeFilter }],
|
||||
},
|
||||
include: {
|
||||
interactions: {
|
||||
where: { userId },
|
||||
},
|
||||
},
|
||||
orderBy: [{ priority: "desc" }, { createdAt: "desc" }],
|
||||
});
|
||||
|
||||
const sorted = [...notifications].sort(compareNotifications);
|
||||
|
||||
// Global per-user request counter stored in Redis, used for cliShowEvery modulo.
|
||||
// This is independent of per-notification showCount so that cliMaxShowCount
|
||||
// correctly tracks actual displays, not API encounters.
|
||||
const requestCounter = await incrementCliRequestCounter(userId);
|
||||
|
||||
for (const n of sorted) {
|
||||
const interaction = n.interactions[0] ?? null;
|
||||
|
||||
if (interaction?.cliDismissedAt) continue;
|
||||
if (isCliNotificationExpired(interaction, n)) continue;
|
||||
|
||||
const parsed = PayloadV1Schema.safeParse(n.payload);
|
||||
if (!parsed.success) continue;
|
||||
|
||||
// Check cliShowEvery using the global request counter
|
||||
if (n.cliShowEvery !== null && requestCounter % n.cliShowEvery !== 0) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Only increment showCount when the notification will actually be displayed.
|
||||
// If this display reaches cliMaxShowCount, also set cliDismissedAt now
|
||||
// so it's recorded immediately rather than waiting for a future request.
|
||||
const reachedMaxShows =
|
||||
n.cliMaxShowCount !== null && (interaction?.showCount ?? 0) + 1 >= n.cliMaxShowCount;
|
||||
|
||||
const updated = await prisma.platformNotificationInteraction.upsert({
|
||||
where: { notificationId_userId: { notificationId: n.id, userId } },
|
||||
update: {
|
||||
showCount: { increment: 1 },
|
||||
...(reachedMaxShows ? { cliDismissedAt: now } : {}),
|
||||
},
|
||||
create: {
|
||||
notificationId: n.id,
|
||||
userId,
|
||||
firstSeenAt: now,
|
||||
showCount: 1,
|
||||
...(reachedMaxShows ? { cliDismissedAt: now } : {}),
|
||||
},
|
||||
});
|
||||
|
||||
return {
|
||||
id: n.id,
|
||||
payload: parsed.data,
|
||||
showCount: updated.showCount,
|
||||
firstSeenAt: updated.firstSeenAt.toISOString(),
|
||||
};
|
||||
}
|
||||
|
||||
return null;
|
||||
}
|
||||
|
||||
// --- Create: admin endpoint support ---
|
||||
|
||||
const SCOPE_REQUIRED_FK: Record<string, "userId" | "organizationId" | "projectId"> = {
|
||||
USER: "userId",
|
||||
ORGANIZATION: "organizationId",
|
||||
PROJECT: "projectId",
|
||||
};
|
||||
|
||||
const ALL_FK_FIELDS = ["userId", "organizationId", "projectId"] as const;
|
||||
const CLI_ONLY_FIELDS = ["cliMaxDaysAfterFirstSeen", "cliMaxShowCount", "cliShowEvery"] as const;
|
||||
|
||||
const NotificationBaseFields = {
|
||||
title: z.string().min(1),
|
||||
payload: PayloadV1Schema,
|
||||
surface: z.enum(["WEBAPP", "CLI"]),
|
||||
scope: z.enum(["USER", "PROJECT", "ORGANIZATION", "GLOBAL"]),
|
||||
userId: z.string().optional(),
|
||||
organizationId: z.string().optional(),
|
||||
projectId: z.string().optional(),
|
||||
endsAt: z
|
||||
.string()
|
||||
.datetime()
|
||||
.transform((s) => new Date(s)),
|
||||
priority: z.number().int().default(0),
|
||||
cliMaxDaysAfterFirstSeen: z.number().int().positive().optional(),
|
||||
cliMaxShowCount: z.number().int().positive().optional(),
|
||||
cliShowEvery: z.number().int().min(2).optional(),
|
||||
};
|
||||
|
||||
export const CreatePlatformNotificationSchema = z
|
||||
.object({
|
||||
...NotificationBaseFields,
|
||||
startsAt: z
|
||||
.string()
|
||||
.datetime()
|
||||
.transform((s) => new Date(s))
|
||||
.optional(),
|
||||
})
|
||||
.superRefine((data, ctx) => {
|
||||
validateScopeForeignKeys(data, ctx);
|
||||
validateSurfaceFields(data, ctx);
|
||||
validatePayloadTypeForSurface(data, ctx);
|
||||
validateStartsAt(data, ctx);
|
||||
validateEndsAt(data, ctx);
|
||||
});
|
||||
|
||||
function validateScopeForeignKeys(
|
||||
data: { scope: string; userId?: string; organizationId?: string; projectId?: string },
|
||||
ctx: z.RefinementCtx
|
||||
) {
|
||||
const requiredFk = SCOPE_REQUIRED_FK[data.scope];
|
||||
|
||||
if (requiredFk && !data[requiredFk]) {
|
||||
ctx.addIssue({
|
||||
code: "custom",
|
||||
message: `${requiredFk} is required when scope is ${data.scope}`,
|
||||
path: [requiredFk],
|
||||
});
|
||||
}
|
||||
|
||||
const forbiddenFks = ALL_FK_FIELDS.filter((fk) => fk !== requiredFk);
|
||||
for (const fk of forbiddenFks) {
|
||||
if (data[fk]) {
|
||||
ctx.addIssue({
|
||||
code: "custom",
|
||||
message: `${fk} must not be set when scope is ${data.scope}`,
|
||||
path: [fk],
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
function validateSurfaceFields(
|
||||
data: {
|
||||
surface: string;
|
||||
cliMaxDaysAfterFirstSeen?: number;
|
||||
cliMaxShowCount?: number;
|
||||
cliShowEvery?: number;
|
||||
},
|
||||
ctx: z.RefinementCtx
|
||||
) {
|
||||
if (data.surface !== "WEBAPP") return;
|
||||
|
||||
for (const field of CLI_ONLY_FIELDS) {
|
||||
if (data[field] !== undefined) {
|
||||
ctx.addIssue({
|
||||
code: "custom",
|
||||
message: `${field} is not allowed for WEBAPP surface`,
|
||||
path: [field],
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
function validateStartsAt(data: { startsAt?: Date }, ctx: z.RefinementCtx) {
|
||||
if (!data.startsAt) return;
|
||||
|
||||
const oneHourAgo = new Date(Date.now() - 60 * 60 * 1000);
|
||||
if (data.startsAt < oneHourAgo) {
|
||||
ctx.addIssue({
|
||||
code: "custom",
|
||||
message: "startsAt must be within the last hour or in the future",
|
||||
path: ["startsAt"],
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
const CLI_TYPES = new Set(["info", "warn", "error", "success"]);
|
||||
const WEBAPP_TYPES = new Set(["card", "changelog"]);
|
||||
|
||||
function validatePayloadTypeForSurface(
|
||||
data: { surface: string; payload: PayloadV1 },
|
||||
ctx: z.RefinementCtx
|
||||
) {
|
||||
const allowedTypes = data.surface === "CLI" ? CLI_TYPES : WEBAPP_TYPES;
|
||||
if (!allowedTypes.has(data.payload.data.type)) {
|
||||
ctx.addIssue({
|
||||
code: "custom",
|
||||
message: `payload.data.type "${data.payload.data.type}" is not allowed for ${data.surface} surface`,
|
||||
path: ["payload", "data", "type"],
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
function validateEndsAt(data: { startsAt?: Date; endsAt: Date }, ctx: z.RefinementCtx) {
|
||||
const effectiveStart = data.startsAt ?? new Date();
|
||||
if (data.endsAt <= effectiveStart) {
|
||||
ctx.addIssue({
|
||||
code: "custom",
|
||||
message: "endsAt must be after startsAt",
|
||||
path: ["endsAt"],
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
export type CreatePlatformNotificationInput = z.input<typeof CreatePlatformNotificationSchema>;
|
||||
|
||||
// --- Update: admin endpoint support ---
|
||||
|
||||
export const UpdatePlatformNotificationSchema = z
|
||||
.object({
|
||||
...NotificationBaseFields,
|
||||
id: z.string().min(1),
|
||||
startsAt: z
|
||||
.string()
|
||||
.datetime()
|
||||
.transform((s) => new Date(s)),
|
||||
})
|
||||
.superRefine((data, ctx) => {
|
||||
validateScopeForeignKeys(data, ctx);
|
||||
validateSurfaceFields(data, ctx);
|
||||
validatePayloadTypeForSurface(data, ctx);
|
||||
// NOTE: No validateStartsAt — existing notifications may have past startsAt
|
||||
validateEndsAt(data, ctx);
|
||||
});
|
||||
|
||||
type CreateError = { type: "validation"; issues: z.ZodIssue[] } | { type: "db"; message: string };
|
||||
|
||||
export function createPlatformNotification(
|
||||
input: CreatePlatformNotificationInput
|
||||
): ResultAsync<{ id: string; friendlyId: string }, CreateError> {
|
||||
const parseResult = CreatePlatformNotificationSchema.safeParse(input);
|
||||
|
||||
if (!parseResult.success) {
|
||||
return errAsync({ type: "validation", issues: parseResult.error.issues });
|
||||
}
|
||||
|
||||
const data = parseResult.data;
|
||||
|
||||
return fromPromise(
|
||||
prisma.platformNotification.create({
|
||||
data: {
|
||||
title: data.title,
|
||||
payload: data.payload,
|
||||
surface: data.surface as PlatformNotificationSurface,
|
||||
scope: data.scope as PlatformNotificationScope,
|
||||
userId: data.userId,
|
||||
organizationId: data.organizationId,
|
||||
projectId: data.projectId,
|
||||
startsAt: data.startsAt ?? new Date(),
|
||||
endsAt: data.endsAt,
|
||||
priority: data.priority,
|
||||
cliMaxDaysAfterFirstSeen: data.cliMaxDaysAfterFirstSeen,
|
||||
cliMaxShowCount: data.cliMaxShowCount,
|
||||
cliShowEvery: data.cliShowEvery,
|
||||
},
|
||||
select: { id: true, friendlyId: true },
|
||||
}),
|
||||
(e): CreateError => ({
|
||||
type: "db",
|
||||
message: e instanceof Error ? e.message : String(e),
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
export function updatePlatformNotification(
|
||||
input: z.input<typeof UpdatePlatformNotificationSchema>
|
||||
): ResultAsync<{ id: string; friendlyId: string }, CreateError> {
|
||||
const parseResult = UpdatePlatformNotificationSchema.safeParse(input);
|
||||
|
||||
if (!parseResult.success) {
|
||||
return errAsync({ type: "validation", issues: parseResult.error.issues });
|
||||
}
|
||||
|
||||
const data = parseResult.data;
|
||||
|
||||
return fromPromise(
|
||||
prisma.platformNotification.update({
|
||||
where: { id: data.id },
|
||||
data: {
|
||||
title: data.title,
|
||||
payload: data.payload,
|
||||
surface: data.surface as PlatformNotificationSurface,
|
||||
scope: data.scope as PlatformNotificationScope,
|
||||
userId: data.scope === "USER" ? data.userId : null,
|
||||
organizationId: data.scope === "ORGANIZATION" ? data.organizationId : null,
|
||||
projectId: data.scope === "PROJECT" ? data.projectId : null,
|
||||
startsAt: data.startsAt,
|
||||
endsAt: data.endsAt,
|
||||
priority: data.priority,
|
||||
cliMaxDaysAfterFirstSeen:
|
||||
data.surface === "CLI" ? (data.cliMaxDaysAfterFirstSeen ?? null) : null,
|
||||
cliMaxShowCount: data.surface === "CLI" ? (data.cliMaxShowCount ?? null) : null,
|
||||
cliShowEvery: data.surface === "CLI" ? (data.cliShowEvery ?? null) : null,
|
||||
},
|
||||
select: { id: true, friendlyId: true },
|
||||
}),
|
||||
(e): CreateError => ({
|
||||
type: "db",
|
||||
message: e instanceof Error ? e.message : String(e),
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
export async function deletePlatformNotification(id: string): Promise<void> {
|
||||
await prisma.platformNotification.delete({ where: { id } });
|
||||
}
|
||||
|
||||
export async function publishNowPlatformNotification(id: string): Promise<void> {
|
||||
await prisma.platformNotification.update({
|
||||
where: { id },
|
||||
data: { startsAt: new Date() },
|
||||
});
|
||||
}
|
||||
|
||||
export async function archivePlatformNotification(id: string): Promise<void> {
|
||||
await prisma.platformNotification.update({
|
||||
where: { id },
|
||||
data: { archivedAt: new Date() },
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
import type { User } from "~/models/user.server";
|
||||
import { telemetry } from "./telemetry.server";
|
||||
|
||||
export async function postAuthentication({
|
||||
user,
|
||||
loginMethod,
|
||||
isNewUser,
|
||||
}: {
|
||||
user: User;
|
||||
loginMethod: User["authenticationMethod"];
|
||||
isNewUser: boolean;
|
||||
}) {
|
||||
telemetry.user.identify({
|
||||
user,
|
||||
isNewUser,
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
import { createCookieSessionStorage } from "@remix-run/node";
|
||||
import { env } from "~/env.server";
|
||||
|
||||
export const uiPreferencesStorage = createCookieSessionStorage({
|
||||
cookie: {
|
||||
name: "__ui_prefs",
|
||||
sameSite: "lax",
|
||||
path: "/",
|
||||
httpOnly: true,
|
||||
secrets: [env.SESSION_SECRET],
|
||||
secure: env.NODE_ENV === "production",
|
||||
maxAge: 60 * 60 * 24 * 365, // 1 year
|
||||
},
|
||||
});
|
||||
|
||||
export function getUiPreferencesSession(request: Request) {
|
||||
return uiPreferencesStorage.getSession(request.headers.get("Cookie"));
|
||||
}
|
||||
|
||||
export async function getUsefulLinksPreference(request: Request): Promise<boolean | undefined> {
|
||||
const session = await getUiPreferencesSession(request);
|
||||
return session.get("showUsefulLinks");
|
||||
}
|
||||
|
||||
export async function setUsefulLinksPreference(show: boolean, request: Request) {
|
||||
const session = await getUiPreferencesSession(request);
|
||||
session.set("showUsefulLinks", show);
|
||||
return session;
|
||||
}
|
||||
|
||||
export async function getRootOnlyFilterPreference(request: Request): Promise<boolean> {
|
||||
const session = await getUiPreferencesSession(request);
|
||||
const rootOnly = session.get("rootOnly");
|
||||
if (rootOnly === undefined) {
|
||||
return false;
|
||||
}
|
||||
return rootOnly;
|
||||
}
|
||||
|
||||
export async function setRootOnlyFilterPreference(rootOnly: boolean, request: Request) {
|
||||
const session = await getUiPreferencesSession(request);
|
||||
session.set("rootOnly", rootOnly);
|
||||
return session;
|
||||
}
|
||||
|
||||
export async function getTimezonePreference(request: Request): Promise<string> {
|
||||
const session = await getUiPreferencesSession(request);
|
||||
const timezone = session.get("timezone");
|
||||
return typeof timezone === "string" ? timezone : "UTC";
|
||||
}
|
||||
|
||||
export async function setTimezonePreference(timezone: string, request: Request) {
|
||||
const session = await getUiPreferencesSession(request);
|
||||
session.set("timezone", timezone);
|
||||
return session;
|
||||
}
|
||||
@@ -0,0 +1,35 @@
|
||||
import type { Organization, Project } from "@trigger.dev/database";
|
||||
import { createEnvironment } from "~/models/organization.server";
|
||||
import { getCurrentPlan, isCloud } from "~/services/platform.v3.server";
|
||||
|
||||
// Extracted from platform.v3.server.ts to break a circular import:
|
||||
// platform.v3.server ↔ models/organization.server (via createEnvironment).
|
||||
// The cycle caused the bundled __esm wrappers to re-enter and short-circuit
|
||||
// the platform.v3.server init, leaving `defaultMachine` and `machines`
|
||||
// undefined in `singleton("machinePresets", ...)` — the boot crash at
|
||||
// `allMachines()` traced to TRI-8731.
|
||||
export async function projectCreated(
|
||||
organization: Pick<Organization, "id" | "maximumConcurrencyLimit">,
|
||||
project: Project
|
||||
) {
|
||||
if (!isCloud()) {
|
||||
await createEnvironment({ organization, project, type: "STAGING" });
|
||||
await createEnvironment({
|
||||
organization,
|
||||
project,
|
||||
type: "PREVIEW",
|
||||
isBranchableEnvironment: true,
|
||||
});
|
||||
} else {
|
||||
const plan = await getCurrentPlan(organization.id);
|
||||
if (plan?.v3Subscription?.plan?.limits?.hasStagingEnvironment) {
|
||||
await createEnvironment({ organization, project, type: "STAGING" });
|
||||
await createEnvironment({
|
||||
organization,
|
||||
project,
|
||||
type: "PREVIEW",
|
||||
isBranchableEnvironment: true,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,329 @@
|
||||
import { type PrismaClient } from "@trigger.dev/database";
|
||||
import { prisma } from "~/db.server";
|
||||
import { DeleteProjectService } from "~/services/deleteProject.server";
|
||||
import { BranchTrackingConfigSchema, type BranchTrackingConfig } from "~/v3/github";
|
||||
import { checkGitHubBranchExists } from "~/services/gitHub.server";
|
||||
import { errAsync, fromPromise, okAsync, ResultAsync } from "neverthrow";
|
||||
import { type BuildSettings } from "~/v3/buildSettings";
|
||||
|
||||
export class ProjectSettingsService {
|
||||
#prismaClient: PrismaClient;
|
||||
|
||||
constructor(prismaClient: PrismaClient = prisma) {
|
||||
this.#prismaClient = prismaClient;
|
||||
}
|
||||
|
||||
renameProject(projectId: string, newName: string) {
|
||||
return fromPromise(
|
||||
this.#prismaClient.project.update({
|
||||
where: {
|
||||
id: projectId,
|
||||
},
|
||||
data: {
|
||||
name: newName,
|
||||
},
|
||||
}),
|
||||
(error) => ({
|
||||
type: "other" as const,
|
||||
cause: error,
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
deleteProject(projectId: string, userId: string) {
|
||||
const deleteProjectService = new DeleteProjectService(this.#prismaClient);
|
||||
|
||||
return fromPromise(deleteProjectService.call({ projectId, userId }), (error) => ({
|
||||
type: "other" as const,
|
||||
cause: error,
|
||||
}));
|
||||
}
|
||||
|
||||
connectGitHubRepo(
|
||||
projectId: string,
|
||||
organizationId: string,
|
||||
repositoryId: string,
|
||||
installationId: string
|
||||
) {
|
||||
const getRepository = () =>
|
||||
fromPromise(
|
||||
this.#prismaClient.githubRepository.findFirst({
|
||||
where: {
|
||||
id: repositoryId,
|
||||
installationId,
|
||||
installation: {
|
||||
organizationId: organizationId,
|
||||
},
|
||||
},
|
||||
select: {
|
||||
id: true,
|
||||
name: true,
|
||||
defaultBranch: true,
|
||||
},
|
||||
}),
|
||||
(error) => ({
|
||||
type: "other" as const,
|
||||
cause: error,
|
||||
})
|
||||
).andThen((repository) => {
|
||||
if (!repository) {
|
||||
return errAsync({ type: "gh_repository_not_found" as const });
|
||||
}
|
||||
return okAsync(repository);
|
||||
});
|
||||
|
||||
const findExistingConnection = () =>
|
||||
fromPromise(
|
||||
this.#prismaClient.connectedGithubRepository.findFirst({
|
||||
where: {
|
||||
projectId: projectId,
|
||||
},
|
||||
}),
|
||||
(error) => ({ type: "other" as const, cause: error })
|
||||
);
|
||||
|
||||
const createConnectedRepo = (defaultBranch: string, previewDeploymentsEnabled: boolean) =>
|
||||
fromPromise(
|
||||
this.#prismaClient.connectedGithubRepository.create({
|
||||
data: {
|
||||
projectId: projectId,
|
||||
repositoryId: repositoryId,
|
||||
branchTracking: {
|
||||
prod: { branch: defaultBranch },
|
||||
staging: {},
|
||||
} satisfies BranchTrackingConfig,
|
||||
previewDeploymentsEnabled,
|
||||
},
|
||||
}),
|
||||
(error) => ({ type: "other" as const, cause: error })
|
||||
);
|
||||
|
||||
return ResultAsync.combine([
|
||||
getRepository(),
|
||||
findExistingConnection(),
|
||||
this.isPreviewEnvironmentEnabled(projectId),
|
||||
]).andThen(([repository, existingConnection, previewEnvironmentEnabled]) => {
|
||||
if (existingConnection) {
|
||||
return errAsync({ type: "project_already_has_connected_repository" as const });
|
||||
}
|
||||
|
||||
return createConnectedRepo(repository.defaultBranch, previewEnvironmentEnabled);
|
||||
});
|
||||
}
|
||||
|
||||
disconnectGitHubRepo(projectId: string) {
|
||||
return fromPromise(
|
||||
this.#prismaClient.connectedGithubRepository.delete({
|
||||
where: {
|
||||
projectId: projectId,
|
||||
},
|
||||
}),
|
||||
(error) => ({ type: "other" as const, cause: error })
|
||||
);
|
||||
}
|
||||
|
||||
updateGitSettings(
|
||||
projectId: string,
|
||||
productionBranch?: string,
|
||||
stagingBranch?: string,
|
||||
previewDeploymentsEnabled?: boolean
|
||||
) {
|
||||
const getExistingConnectedRepo = () =>
|
||||
fromPromise(
|
||||
this.#prismaClient.connectedGithubRepository.findFirst({
|
||||
where: {
|
||||
projectId: projectId,
|
||||
},
|
||||
include: {
|
||||
repository: {
|
||||
include: {
|
||||
installation: true,
|
||||
},
|
||||
},
|
||||
},
|
||||
}),
|
||||
(error) => ({ type: "other" as const, cause: error })
|
||||
)
|
||||
.andThen((connectedRepo) => {
|
||||
if (!connectedRepo) {
|
||||
return errAsync({ type: "connected_gh_repository_not_found" as const });
|
||||
}
|
||||
return okAsync(connectedRepo);
|
||||
})
|
||||
.map((connectedRepo) => {
|
||||
const branchTrackingOrFailure = BranchTrackingConfigSchema.safeParse(
|
||||
connectedRepo.branchTracking
|
||||
);
|
||||
const branchTracking = branchTrackingOrFailure.success
|
||||
? branchTrackingOrFailure.data
|
||||
: undefined;
|
||||
|
||||
return {
|
||||
...connectedRepo,
|
||||
branchTracking,
|
||||
};
|
||||
});
|
||||
|
||||
const validateProductionBranch = ({
|
||||
installationId,
|
||||
fullRepoName,
|
||||
oldProductionBranch,
|
||||
}: {
|
||||
installationId: number;
|
||||
fullRepoName: string;
|
||||
oldProductionBranch?: string;
|
||||
}) => {
|
||||
if (productionBranch && oldProductionBranch !== productionBranch) {
|
||||
return checkGitHubBranchExists(installationId, fullRepoName, productionBranch).andThen(
|
||||
(exists) => {
|
||||
if (!exists) {
|
||||
return errAsync({ type: "production_tracking_branch_not_found" as const });
|
||||
}
|
||||
return okAsync(productionBranch);
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
return okAsync(productionBranch);
|
||||
};
|
||||
|
||||
const validateStagingBranch = ({
|
||||
installationId,
|
||||
fullRepoName,
|
||||
oldStagingBranch,
|
||||
}: {
|
||||
installationId: number;
|
||||
fullRepoName: string;
|
||||
oldStagingBranch?: string;
|
||||
}) => {
|
||||
if (stagingBranch && oldStagingBranch !== stagingBranch) {
|
||||
return checkGitHubBranchExists(installationId, fullRepoName, stagingBranch).andThen(
|
||||
(exists) => {
|
||||
if (!exists) {
|
||||
return errAsync({ type: "staging_tracking_branch_not_found" as const });
|
||||
}
|
||||
return okAsync(stagingBranch);
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
return okAsync(stagingBranch);
|
||||
};
|
||||
|
||||
const updateConnectedRepo = (data: {
|
||||
productionBranch: string | undefined;
|
||||
stagingBranch: string | undefined;
|
||||
previewDeploymentsEnabled: boolean | undefined;
|
||||
}) =>
|
||||
fromPromise(
|
||||
this.#prismaClient.connectedGithubRepository.update({
|
||||
where: {
|
||||
projectId: projectId,
|
||||
},
|
||||
data: {
|
||||
branchTracking: {
|
||||
prod: data.productionBranch ? { branch: data.productionBranch } : {},
|
||||
staging: data.stagingBranch ? { branch: data.stagingBranch } : {},
|
||||
} satisfies BranchTrackingConfig,
|
||||
previewDeploymentsEnabled: data.previewDeploymentsEnabled,
|
||||
},
|
||||
}),
|
||||
(error) => ({ type: "other" as const, cause: error })
|
||||
);
|
||||
|
||||
return getExistingConnectedRepo()
|
||||
.andThen((connectedRepo) => {
|
||||
const installationId = Number(connectedRepo.repository.installation.appInstallationId);
|
||||
|
||||
return ResultAsync.combine([
|
||||
validateProductionBranch({
|
||||
installationId,
|
||||
fullRepoName: connectedRepo.repository.fullName,
|
||||
oldProductionBranch: connectedRepo.branchTracking?.prod?.branch,
|
||||
}),
|
||||
validateStagingBranch({
|
||||
installationId,
|
||||
fullRepoName: connectedRepo.repository.fullName,
|
||||
oldStagingBranch: connectedRepo.branchTracking?.staging?.branch,
|
||||
}),
|
||||
this.isPreviewEnvironmentEnabled(projectId),
|
||||
]);
|
||||
})
|
||||
.map(([productionBranch, stagingBranch, previewEnvironmentEnabled]) => ({
|
||||
productionBranch,
|
||||
stagingBranch,
|
||||
previewDeploymentsEnabled: previewDeploymentsEnabled && previewEnvironmentEnabled,
|
||||
}))
|
||||
.andThen(updateConnectedRepo);
|
||||
}
|
||||
|
||||
updateBuildSettings(projectId: string, buildSettings: BuildSettings) {
|
||||
return fromPromise(
|
||||
this.#prismaClient.project.update({
|
||||
where: {
|
||||
id: projectId,
|
||||
},
|
||||
data: {
|
||||
buildSettings: buildSettings,
|
||||
},
|
||||
}),
|
||||
(error) => ({
|
||||
type: "other" as const,
|
||||
cause: error,
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
verifyProjectMembership(organizationSlug: string, projectSlug: string, userId: string) {
|
||||
const findProject = () =>
|
||||
fromPromise(
|
||||
this.#prismaClient.project.findFirst({
|
||||
where: {
|
||||
slug: projectSlug,
|
||||
organization: {
|
||||
slug: organizationSlug,
|
||||
members: {
|
||||
some: {
|
||||
userId,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
select: {
|
||||
id: true,
|
||||
organizationId: true,
|
||||
},
|
||||
}),
|
||||
(error) => ({ type: "other" as const, cause: error })
|
||||
);
|
||||
|
||||
return findProject().andThen((project) => {
|
||||
if (!project) {
|
||||
return errAsync({ type: "user_not_in_project" as const });
|
||||
}
|
||||
|
||||
return okAsync({
|
||||
projectId: project.id,
|
||||
organizationId: project.organizationId,
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
private isPreviewEnvironmentEnabled(projectId: string) {
|
||||
return fromPromise(
|
||||
this.#prismaClient.runtimeEnvironment.findFirst({
|
||||
select: {
|
||||
id: true,
|
||||
},
|
||||
where: {
|
||||
projectId: projectId,
|
||||
slug: "preview",
|
||||
},
|
||||
}),
|
||||
(error) => ({
|
||||
type: "other" as const,
|
||||
cause: error,
|
||||
})
|
||||
).map((previewEnvironment) => previewEnvironment !== null);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,172 @@
|
||||
import { type PrismaClient } from "@trigger.dev/database";
|
||||
import { prisma } from "~/db.server";
|
||||
import { BranchTrackingConfigSchema } from "~/v3/github";
|
||||
import { env } from "~/env.server";
|
||||
import { findProjectBySlug } from "~/models/project.server";
|
||||
import { err, fromPromise, ok, ResultAsync } from "neverthrow";
|
||||
import { BuildSettingsSchema } from "~/v3/buildSettings";
|
||||
|
||||
export class ProjectSettingsPresenter {
|
||||
#prismaClient: PrismaClient;
|
||||
|
||||
constructor(prismaClient: PrismaClient = prisma) {
|
||||
this.#prismaClient = prismaClient;
|
||||
}
|
||||
|
||||
getProjectSettings(organizationSlug: string, projectSlug: string, userId: string) {
|
||||
const githubAppEnabled = env.GITHUB_APP_ENABLED === "1";
|
||||
|
||||
const getProject = () =>
|
||||
fromPromise(findProjectBySlug(organizationSlug, projectSlug, userId), (error) => ({
|
||||
type: "other" as const,
|
||||
cause: error,
|
||||
}))
|
||||
.andThen((project) => {
|
||||
if (!project) {
|
||||
return err({ type: "project_not_found" as const });
|
||||
}
|
||||
return ok(project);
|
||||
})
|
||||
.map((project) => {
|
||||
const buildSettingsOrFailure = BuildSettingsSchema.safeParse(project.buildSettings);
|
||||
const buildSettings = buildSettingsOrFailure.success
|
||||
? buildSettingsOrFailure.data
|
||||
: undefined;
|
||||
return { ...project, buildSettings };
|
||||
});
|
||||
|
||||
if (!githubAppEnabled) {
|
||||
return getProject().map(({ buildSettings }) => ({
|
||||
gitHubApp: {
|
||||
enabled: false,
|
||||
connectedRepository: undefined,
|
||||
installations: undefined,
|
||||
isPreviewEnvironmentEnabled: undefined,
|
||||
},
|
||||
buildSettings,
|
||||
}));
|
||||
}
|
||||
|
||||
const findConnectedGithubRepository = (projectId: string) =>
|
||||
fromPromise(
|
||||
this.#prismaClient.connectedGithubRepository.findFirst({
|
||||
where: {
|
||||
projectId,
|
||||
repository: {
|
||||
installation: {
|
||||
deletedAt: null,
|
||||
suspendedAt: null,
|
||||
},
|
||||
},
|
||||
},
|
||||
select: {
|
||||
branchTracking: true,
|
||||
previewDeploymentsEnabled: true,
|
||||
createdAt: true,
|
||||
repository: {
|
||||
select: {
|
||||
id: true,
|
||||
name: true,
|
||||
fullName: true,
|
||||
htmlUrl: true,
|
||||
private: true,
|
||||
},
|
||||
},
|
||||
},
|
||||
}),
|
||||
(error) => ({
|
||||
type: "other" as const,
|
||||
cause: error,
|
||||
})
|
||||
).map((connectedGithubRepository) => {
|
||||
if (!connectedGithubRepository) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
const branchTrackingOrFailure = BranchTrackingConfigSchema.safeParse(
|
||||
connectedGithubRepository.branchTracking
|
||||
);
|
||||
const branchTracking = branchTrackingOrFailure.success
|
||||
? branchTrackingOrFailure.data
|
||||
: undefined;
|
||||
|
||||
return {
|
||||
...connectedGithubRepository,
|
||||
branchTracking,
|
||||
};
|
||||
});
|
||||
|
||||
const listGithubAppInstallations = (organizationId: string) =>
|
||||
fromPromise(
|
||||
this.#prismaClient.githubAppInstallation.findMany({
|
||||
where: {
|
||||
organizationId,
|
||||
deletedAt: null,
|
||||
suspendedAt: null,
|
||||
},
|
||||
select: {
|
||||
id: true,
|
||||
accountHandle: true,
|
||||
targetType: true,
|
||||
appInstallationId: true,
|
||||
repositories: {
|
||||
select: {
|
||||
id: true,
|
||||
name: true,
|
||||
fullName: true,
|
||||
htmlUrl: true,
|
||||
private: true,
|
||||
},
|
||||
// Most installations will only have a couple of repos so loading them here should be fine.
|
||||
// However, there might be outlier organizations so it's best to expose the installation repos
|
||||
// via a resource endpoint and filter on user input.
|
||||
take: 200,
|
||||
},
|
||||
},
|
||||
take: 20,
|
||||
orderBy: {
|
||||
createdAt: "desc",
|
||||
},
|
||||
}),
|
||||
(error) => ({
|
||||
type: "other" as const,
|
||||
cause: error,
|
||||
})
|
||||
);
|
||||
|
||||
const isPreviewEnvironmentEnabled = (projectId: string) =>
|
||||
fromPromise(
|
||||
this.#prismaClient.runtimeEnvironment.findFirst({
|
||||
select: {
|
||||
id: true,
|
||||
},
|
||||
where: {
|
||||
projectId: projectId,
|
||||
slug: "preview",
|
||||
},
|
||||
}),
|
||||
(error) => ({
|
||||
type: "other" as const,
|
||||
cause: error,
|
||||
})
|
||||
).map((previewEnvironment) => previewEnvironment !== null);
|
||||
|
||||
return getProject().andThen((project) =>
|
||||
ResultAsync.combine([
|
||||
isPreviewEnvironmentEnabled(project.id),
|
||||
findConnectedGithubRepository(project.id),
|
||||
listGithubAppInstallations(project.organizationId),
|
||||
]).map(
|
||||
([isPreviewEnvironmentEnabled, connectedGithubRepository, githubAppInstallations]) => ({
|
||||
gitHubApp: {
|
||||
enabled: true,
|
||||
connectedRepository: connectedGithubRepository,
|
||||
installations: githubAppInstallations,
|
||||
isPreviewEnvironmentEnabled,
|
||||
},
|
||||
buildSettings: project.buildSettings,
|
||||
})
|
||||
)
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,26 @@
|
||||
import { createCookie } from "@remix-run/node";
|
||||
import { env } from "~/env.server";
|
||||
|
||||
// Carries a promo code from the landing page through signup to first-org
|
||||
// creation. httpOnly + sameSite=lax so it survives the OAuth round-trip,
|
||||
// matching the existing redirect-to cookie.
|
||||
export const promoCodeCookie = createCookie("promo-code", {
|
||||
maxAge: 60 * 60, // 1 hour — enough to complete signup
|
||||
httpOnly: true,
|
||||
sameSite: "lax",
|
||||
secure: env.NODE_ENV === "production",
|
||||
path: "/",
|
||||
});
|
||||
|
||||
export async function setPromoCodeCookie(code: string): Promise<string> {
|
||||
return await promoCodeCookie.serialize(code);
|
||||
}
|
||||
|
||||
export async function getPromoCodeFromCookie(request: Request): Promise<string | null> {
|
||||
const value = await promoCodeCookie.parse(request.headers.get("Cookie"));
|
||||
return typeof value === "string" && value.length > 0 ? value : null;
|
||||
}
|
||||
|
||||
export async function clearPromoCodeCookie(): Promise<string> {
|
||||
return await promoCodeCookie.serialize("", { maxAge: 0 });
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
import { env } from "~/env.server";
|
||||
import { singleton } from "~/utils/singleton";
|
||||
import { RedisConcurrencyLimiter } from "./redisConcurrencyLimiter.server";
|
||||
|
||||
function initializeQueryConcurrencyLimiter() {
|
||||
return new RedisConcurrencyLimiter({
|
||||
keyPrefix: "query:concurrency",
|
||||
redis: {
|
||||
port: env.RATE_LIMIT_REDIS_PORT,
|
||||
host: env.RATE_LIMIT_REDIS_HOST,
|
||||
username: env.RATE_LIMIT_REDIS_USERNAME,
|
||||
password: env.RATE_LIMIT_REDIS_PASSWORD,
|
||||
tlsDisabled: env.RATE_LIMIT_REDIS_TLS_DISABLED === "true",
|
||||
clusterMode: env.RATE_LIMIT_REDIS_CLUSTER_MODE_ENABLED === "1",
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
export const queryConcurrencyLimiter = singleton(
|
||||
"queryConcurrencyLimiter",
|
||||
initializeQueryConcurrencyLimiter
|
||||
);
|
||||
|
||||
/** Default per-org concurrency limit from environment */
|
||||
export const DEFAULT_ORG_CONCURRENCY_LIMIT = env.QUERY_DEFAULT_ORG_CONCURRENCY_LIMIT;
|
||||
|
||||
/** Global concurrency limit from environment */
|
||||
export const GLOBAL_CONCURRENCY_LIMIT = env.QUERY_GLOBAL_CONCURRENCY_LIMIT;
|
||||
@@ -0,0 +1,382 @@
|
||||
import {
|
||||
executeTSQL,
|
||||
QueryError,
|
||||
type ClickHouseSettings,
|
||||
type ExecuteTSQLOptions,
|
||||
type FieldMappings,
|
||||
type TSQLQueryResult,
|
||||
} from "@internal/clickhouse";
|
||||
import type { CustomerQuerySource } from "@trigger.dev/database";
|
||||
import type { TableSchema, WhereClauseCondition } from "@internal/tsql";
|
||||
import { z } from "zod";
|
||||
import { prisma } from "~/db.server";
|
||||
import { env } from "~/env.server";
|
||||
import { clickhouseFactory } from "./clickhouse/clickhouseFactoryInstance.server";
|
||||
import {
|
||||
queryConcurrencyLimiter,
|
||||
DEFAULT_ORG_CONCURRENCY_LIMIT,
|
||||
GLOBAL_CONCURRENCY_LIMIT,
|
||||
} from "./queryConcurrencyLimiter.server";
|
||||
import { getLimit } from "./platform.v3.server";
|
||||
import { timeFilters, timeFilterFromTo } from "~/components/runs/v3/SharedFilters";
|
||||
import parse from "parse-duration";
|
||||
import { querySchemas, QueryScopeSchema, type QueryScope } from "~/v3/querySchemas";
|
||||
|
||||
export { QueryScopeSchema };
|
||||
export type { TableSchema, TSQLQueryResult, QueryScope };
|
||||
|
||||
const scopeToEnum = {
|
||||
organization: "ORGANIZATION",
|
||||
project: "PROJECT",
|
||||
environment: "ENVIRONMENT",
|
||||
} as const;
|
||||
|
||||
/**
|
||||
* Default ClickHouse settings for query protection
|
||||
* Based on PostHog's HogQL settings to prevent expensive queries
|
||||
*/
|
||||
function getDefaultClickhouseSettings(): ClickHouseSettings {
|
||||
return {
|
||||
// Query execution limits
|
||||
max_execution_time: env.QUERY_CLICKHOUSE_MAX_EXECUTION_TIME,
|
||||
timeout_overflow_mode: "throw",
|
||||
max_memory_usage: String(env.QUERY_CLICKHOUSE_MAX_MEMORY_USAGE),
|
||||
|
||||
// AST complexity limits to prevent extremely complex queries
|
||||
max_ast_elements: String(env.QUERY_CLICKHOUSE_MAX_AST_ELEMENTS),
|
||||
max_expanded_ast_elements: String(env.QUERY_CLICKHOUSE_MAX_EXPANDED_AST_ELEMENTS),
|
||||
|
||||
// Memory management for GROUP BY operations
|
||||
max_bytes_before_external_group_by: String(
|
||||
env.QUERY_CLICKHOUSE_MAX_BYTES_BEFORE_EXTERNAL_GROUP_BY
|
||||
),
|
||||
|
||||
// Safety settings
|
||||
format_csv_allow_double_quotes: 0,
|
||||
readonly: "1", // Ensure queries are read-only
|
||||
};
|
||||
}
|
||||
|
||||
export type ExecuteQueryOptions<TOut extends z.ZodSchema> = Omit<
|
||||
ExecuteTSQLOptions<TOut>,
|
||||
"tableSchema" | "fieldMappings" | "enforcedWhereClause" | "whereClauseFallback" | "schema"
|
||||
> & {
|
||||
organizationId: string;
|
||||
projectId: string;
|
||||
environmentId: string;
|
||||
/** The scope of the query - determines tenant isolation */
|
||||
scope: QueryScope;
|
||||
period?: string | null;
|
||||
from?: string | null;
|
||||
to?: string | null;
|
||||
/** Filter to specific task identifiers */
|
||||
taskIdentifiers?: string[];
|
||||
/** Filter to specific queues */
|
||||
queues?: string[];
|
||||
/** Filter to specific response models */
|
||||
responseModels?: string[];
|
||||
/** Filter to specific prompt slugs */
|
||||
promptSlugs?: string[];
|
||||
/** Filter to specific prompt versions */
|
||||
promptVersions?: number[];
|
||||
/** Filter to specific operations (e.g. ai.generateText.doGenerate) */
|
||||
operations?: string[];
|
||||
/** Filter to specific providers (e.g. openai.responses) */
|
||||
providers?: string[];
|
||||
/** History options for saving query to billing/audit */
|
||||
history?: {
|
||||
/** Where the query originated from */
|
||||
source: CustomerQuerySource;
|
||||
/** User ID (optional, null for API calls) */
|
||||
userId?: string | null;
|
||||
/** Skip saving to history (e.g., when impersonating) */
|
||||
skip?: boolean;
|
||||
};
|
||||
/** Custom per-org concurrency limit (overrides default) */
|
||||
customOrgConcurrencyLimit?: number;
|
||||
};
|
||||
|
||||
/**
|
||||
* Extended result type that includes the optional queryId when saved to history
|
||||
*/
|
||||
export type ExecuteQueryResult<T> =
|
||||
| {
|
||||
success: true;
|
||||
result: T;
|
||||
queryId: string | null;
|
||||
periodClipped: number | null;
|
||||
maxQueryPeriod: number;
|
||||
timeRange: { from: Date; to: Date };
|
||||
}
|
||||
| { success: false; error: Error };
|
||||
|
||||
export async function getDefaultPeriod(organizationId: string): Promise<string> {
|
||||
const idealDefaultPeriodDays = 7;
|
||||
const maxQueryPeriod = await getLimit(organizationId, "queryPeriodDays", 30);
|
||||
if (maxQueryPeriod < idealDefaultPeriodDays) {
|
||||
return `${maxQueryPeriod}d`;
|
||||
}
|
||||
return `${idealDefaultPeriodDays}d`;
|
||||
}
|
||||
|
||||
/**
|
||||
* Execute a TSQL query against ClickHouse with tenant isolation
|
||||
* Handles building tenant options, field mappings, and optionally saves to history
|
||||
* Returns [error, result, queryId] where queryId is the CustomerQuery ID if saved to history
|
||||
*/
|
||||
export async function executeQuery<TOut extends z.ZodSchema>(
|
||||
options: ExecuteQueryOptions<TOut>
|
||||
): Promise<ExecuteQueryResult<Exclude<TSQLQueryResult<z.output<TOut>>[1], null>>> {
|
||||
const {
|
||||
period,
|
||||
from,
|
||||
to,
|
||||
scope,
|
||||
organizationId,
|
||||
projectId,
|
||||
environmentId,
|
||||
taskIdentifiers,
|
||||
queues,
|
||||
responseModels,
|
||||
promptSlugs,
|
||||
promptVersions,
|
||||
operations,
|
||||
providers,
|
||||
history,
|
||||
customOrgConcurrencyLimit,
|
||||
...baseOptions
|
||||
} = options;
|
||||
|
||||
// Generate unique request ID for concurrency tracking
|
||||
const requestId = crypto.randomUUID();
|
||||
const orgLimit = customOrgConcurrencyLimit ?? DEFAULT_ORG_CONCURRENCY_LIMIT;
|
||||
|
||||
// Acquire concurrency slot
|
||||
const acquireResult = await queryConcurrencyLimiter.acquire({
|
||||
key: projectId,
|
||||
requestId,
|
||||
keyLimit: orgLimit,
|
||||
globalLimit: GLOBAL_CONCURRENCY_LIMIT,
|
||||
});
|
||||
|
||||
if (!acquireResult.success) {
|
||||
const errorMessage =
|
||||
acquireResult.reason === "key_limit"
|
||||
? `You've exceeded your query concurrency of ${orgLimit} for this project. Please try again later.`
|
||||
: "We're experiencing a lot of queries at the moment. Please try again later.";
|
||||
return { success: false, error: new QueryError(errorMessage, { query: options.query }) };
|
||||
}
|
||||
|
||||
// Detect which table the query targets to determine the time column
|
||||
// Each table schema declares its primary time column via timeConstraint
|
||||
const matchedSchema = querySchemas.find((s) =>
|
||||
new RegExp(`\\bFROM\\s+${s.name}\\b`, "i").test(options.query)
|
||||
);
|
||||
const timeColumn = matchedSchema?.timeConstraint ?? "triggered_at";
|
||||
|
||||
// Build time filter fallback for the table's time column
|
||||
const defaultPeriod = await getDefaultPeriod(organizationId);
|
||||
const timeFilter = timeFilters({
|
||||
period: period ?? undefined,
|
||||
from: from ?? undefined,
|
||||
to: to ?? undefined,
|
||||
defaultPeriod,
|
||||
});
|
||||
|
||||
// Calculate the effective "from" date the user is requesting (for period clipping check)
|
||||
// This is null only when the user specifies just a "to" date (rare case)
|
||||
let requestedFromDate: Date | null = null;
|
||||
if (timeFilter.from) {
|
||||
requestedFromDate = new Date(timeFilter.from);
|
||||
} else if (!timeFilter.to) {
|
||||
// Period specified (or default) - calculate from now
|
||||
const periodMs = parse(timeFilter.period ?? defaultPeriod) ?? 7 * 24 * 60 * 60 * 1000;
|
||||
requestedFromDate = new Date(Date.now() - periodMs);
|
||||
}
|
||||
|
||||
// Build the fallback WHERE condition based on what the user specified
|
||||
let timeFallback: WhereClauseCondition;
|
||||
if (timeFilter.from && timeFilter.to) {
|
||||
timeFallback = { op: "between", low: timeFilter.from, high: timeFilter.to };
|
||||
} else if (timeFilter.from) {
|
||||
timeFallback = { op: "gte", value: timeFilter.from };
|
||||
} else if (timeFilter.to) {
|
||||
timeFallback = { op: "lte", value: timeFilter.to };
|
||||
} else {
|
||||
timeFallback = { op: "gte", value: requestedFromDate! };
|
||||
}
|
||||
|
||||
const maxQueryPeriod = await getLimit(organizationId, "queryPeriodDays", 30);
|
||||
const maxQueryPeriodDate = new Date(Date.now() - maxQueryPeriod * 24 * 60 * 60 * 1000);
|
||||
|
||||
// Check if the requested time period exceeds the plan limit
|
||||
const periodClipped = requestedFromDate !== null && requestedFromDate < maxQueryPeriodDate;
|
||||
|
||||
// Force tenant isolation and time period limits
|
||||
// Global tables (no tenantColumns) skip tenant isolation — they contain anonymized cross-tenant data
|
||||
const isGlobalTable = matchedSchema != null && !matchedSchema.tenantColumns;
|
||||
const enforcedWhereClause = {
|
||||
...(isGlobalTable
|
||||
? {}
|
||||
: {
|
||||
organization_id: { op: "eq", value: organizationId },
|
||||
project_id:
|
||||
scope === "project" || scope === "environment"
|
||||
? { op: "eq", value: projectId }
|
||||
: undefined,
|
||||
environment_id: scope === "environment" ? { op: "eq", value: environmentId } : undefined,
|
||||
}),
|
||||
[timeColumn]: { op: "gte", value: maxQueryPeriodDate },
|
||||
// Optional filters for tasks and queues
|
||||
task_identifier:
|
||||
taskIdentifiers && taskIdentifiers.length > 0
|
||||
? { op: "in", values: taskIdentifiers }
|
||||
: undefined,
|
||||
queue: queues && queues.length > 0 ? { op: "in", values: queues } : undefined,
|
||||
response_model:
|
||||
responseModels && responseModels.length > 0
|
||||
? { op: "in", values: responseModels }
|
||||
: undefined,
|
||||
prompt_slug:
|
||||
promptSlugs && promptSlugs.length > 0 ? { op: "in", values: promptSlugs } : undefined,
|
||||
prompt_version:
|
||||
promptVersions && promptVersions.length > 0
|
||||
? { op: "in", values: promptVersions }
|
||||
: undefined,
|
||||
operation_id:
|
||||
operations && operations.length > 0 ? { op: "in", values: operations } : undefined,
|
||||
gen_ai_system: providers && providers.length > 0 ? { op: "in", values: providers } : undefined,
|
||||
} satisfies Record<string, WhereClauseCondition | undefined>;
|
||||
|
||||
// Compute the effective time range for timeBucket() interval calculation
|
||||
const timeRange = timeFilterFromTo({
|
||||
period: period ?? undefined,
|
||||
from: from ?? undefined,
|
||||
to: to ?? undefined,
|
||||
defaultPeriod,
|
||||
});
|
||||
|
||||
try {
|
||||
// Build field mappings for project_ref → project_id and environment_id → slug translation
|
||||
const projects = await prisma.project.findMany({
|
||||
where: { organizationId },
|
||||
select: { id: true, externalRef: true },
|
||||
});
|
||||
|
||||
const environments = await prisma.runtimeEnvironment.findMany({
|
||||
where: { project: { organizationId } },
|
||||
select: { id: true, slug: true },
|
||||
});
|
||||
|
||||
const fieldMappings: FieldMappings = {
|
||||
project: Object.fromEntries(projects.map((p) => [p.id, p.externalRef])),
|
||||
environment: Object.fromEntries(environments.map((e) => [e.id, e.slug])),
|
||||
};
|
||||
|
||||
const queryClickhouse = await clickhouseFactory.getClickhouseForOrganization(
|
||||
organizationId,
|
||||
"query"
|
||||
);
|
||||
const result = await executeTSQL(queryClickhouse.reader, {
|
||||
...baseOptions,
|
||||
schema: z.record(z.any()),
|
||||
tableSchema: querySchemas,
|
||||
transformValues: true,
|
||||
enforcedWhereClause,
|
||||
fieldMappings,
|
||||
whereClauseFallback: {
|
||||
[timeColumn]: timeFallback,
|
||||
},
|
||||
timeRange,
|
||||
clickhouseSettings: {
|
||||
...getDefaultClickhouseSettings(),
|
||||
...baseOptions.clickhouseSettings, // Allow caller overrides if needed
|
||||
},
|
||||
querySettings: {
|
||||
maxRows: env.QUERY_CLICKHOUSE_MAX_RETURNED_ROWS,
|
||||
...baseOptions.querySettings, // Allow caller overrides if needed
|
||||
},
|
||||
});
|
||||
|
||||
// If query failed, return early with no queryId
|
||||
if (result[0] !== null) {
|
||||
return { success: false, error: result[0] };
|
||||
}
|
||||
|
||||
let queryId: string | null = null;
|
||||
|
||||
// If query succeeded and history options provided, save to history
|
||||
// Skip history for EXPLAIN queries (admin debugging) and when explicitly skipped (e.g., impersonating)
|
||||
if (history && !history.skip && !baseOptions.explain) {
|
||||
// Check if this query is the same as the last one saved (avoid duplicate history entries)
|
||||
const lastQuery = await prisma.customerQuery.findFirst({
|
||||
where: {
|
||||
organizationId,
|
||||
source: history.source,
|
||||
userId: history.userId ?? null,
|
||||
},
|
||||
orderBy: { createdAt: "desc" },
|
||||
select: {
|
||||
id: true,
|
||||
query: true,
|
||||
scope: true,
|
||||
filterPeriod: true,
|
||||
filterFrom: true,
|
||||
filterTo: true,
|
||||
},
|
||||
});
|
||||
|
||||
// Save the effective period used for the query (timeFilters() handles defaults)
|
||||
// Only save period if no custom from/to range was specified
|
||||
const historyTimeFilter = {
|
||||
period: timeFilter.from || timeFilter.to ? undefined : timeFilter.period,
|
||||
from: timeFilter.from,
|
||||
to: timeFilter.to,
|
||||
};
|
||||
const isDuplicate =
|
||||
lastQuery &&
|
||||
lastQuery.query === options.query &&
|
||||
lastQuery.scope === scopeToEnum[scope] &&
|
||||
lastQuery.filterPeriod === (timeFilter?.period ?? null) &&
|
||||
lastQuery.filterFrom?.getTime() === (timeFilter?.from?.getTime() ?? undefined) &&
|
||||
lastQuery.filterTo?.getTime() === (timeFilter?.to?.getTime() ?? undefined);
|
||||
|
||||
if (isDuplicate && lastQuery) {
|
||||
// Return the existing query's ID for duplicate queries
|
||||
queryId = lastQuery.id;
|
||||
} else {
|
||||
const created = await prisma.customerQuery.create({
|
||||
data: {
|
||||
query: options.query,
|
||||
scope: scopeToEnum[scope],
|
||||
stats: { ...result[1].stats },
|
||||
source: history.source,
|
||||
organizationId,
|
||||
projectId: scope === "project" || scope === "environment" ? projectId : null,
|
||||
environmentId: scope === "environment" ? environmentId : null,
|
||||
userId: history.userId ?? null,
|
||||
filterPeriod: historyTimeFilter?.period ?? null,
|
||||
filterFrom: historyTimeFilter?.from ?? null,
|
||||
filterTo: historyTimeFilter?.to ?? null,
|
||||
},
|
||||
});
|
||||
queryId = created.id;
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
success: true,
|
||||
result: result[1],
|
||||
queryId,
|
||||
periodClipped: periodClipped ? maxQueryPeriod : null,
|
||||
maxQueryPeriod,
|
||||
timeRange,
|
||||
};
|
||||
} finally {
|
||||
// Always release the concurrency slot
|
||||
await queryConcurrencyLimiter.release({
|
||||
key: projectId,
|
||||
requestId,
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,40 @@
|
||||
import { env } from "~/env.server";
|
||||
import type { RedisWithClusterOptions } from "~/redis.server";
|
||||
import {
|
||||
RateLimiter as CoreRateLimiter,
|
||||
type Limiter,
|
||||
type RateLimiterRedisClient,
|
||||
} from "./rateLimiterCore.server";
|
||||
|
||||
export {
|
||||
createRedisRateLimitClient,
|
||||
type Duration,
|
||||
type Limiter,
|
||||
type RateLimitResponse,
|
||||
type RateLimiterRedisClient,
|
||||
} from "./rateLimiterCore.server";
|
||||
|
||||
type Options = {
|
||||
redis?: RedisWithClusterOptions;
|
||||
redisClient?: RateLimiterRedisClient;
|
||||
keyPrefix: string;
|
||||
limiter: Limiter;
|
||||
logSuccess?: boolean;
|
||||
logFailure?: boolean;
|
||||
};
|
||||
|
||||
export class RateLimiter extends CoreRateLimiter {
|
||||
constructor(options: Options) {
|
||||
super({
|
||||
...options,
|
||||
redis: options.redis ?? {
|
||||
port: env.RATE_LIMIT_REDIS_PORT,
|
||||
host: env.RATE_LIMIT_REDIS_HOST,
|
||||
username: env.RATE_LIMIT_REDIS_USERNAME,
|
||||
password: env.RATE_LIMIT_REDIS_PASSWORD,
|
||||
tlsDisabled: env.RATE_LIMIT_REDIS_TLS_DISABLED === "true",
|
||||
clusterMode: env.RATE_LIMIT_REDIS_CLUSTER_MODE_ENABLED === "1",
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,100 @@
|
||||
import { Ratelimit } from "@upstash/ratelimit";
|
||||
import type { RedisWithClusterOptions } from "~/redis.server";
|
||||
import { createRedisClient } from "~/redis.server";
|
||||
import { logger } from "./logger.server";
|
||||
|
||||
type Options = {
|
||||
redis?: RedisWithClusterOptions;
|
||||
redisClient?: RateLimiterRedisClient;
|
||||
keyPrefix: string;
|
||||
limiter: Limiter;
|
||||
logSuccess?: boolean;
|
||||
logFailure?: boolean;
|
||||
};
|
||||
|
||||
export type Limiter = ConstructorParameters<typeof Ratelimit>[0]["limiter"];
|
||||
export type Duration = Parameters<typeof Ratelimit.slidingWindow>[1];
|
||||
export type RateLimitResponse = Awaited<ReturnType<Ratelimit["limit"]>>;
|
||||
export type RateLimiterRedisClient = ConstructorParameters<typeof Ratelimit>[0]["redis"];
|
||||
|
||||
export class RateLimiter {
|
||||
#ratelimit: Ratelimit;
|
||||
|
||||
constructor(private readonly options: Options) {
|
||||
const { redis, redisClient, keyPrefix, limiter } = options;
|
||||
const prefix = `ratelimit:${keyPrefix}`;
|
||||
const resolvedRedisClient =
|
||||
redisClient ?? (redis ? createRedisRateLimitClient(redis) : undefined);
|
||||
|
||||
if (!resolvedRedisClient) {
|
||||
throw new Error("RateLimiter requires either redis or redisClient options");
|
||||
}
|
||||
|
||||
this.#ratelimit = new Ratelimit({
|
||||
redis: resolvedRedisClient,
|
||||
limiter,
|
||||
ephemeralCache: new Map(),
|
||||
analytics: false,
|
||||
prefix,
|
||||
});
|
||||
}
|
||||
|
||||
async limit(identifier: string, rate = 1): Promise<RateLimitResponse> {
|
||||
const result = this.#ratelimit.limit(identifier, { rate });
|
||||
const { success, limit, reset, remaining } = await result;
|
||||
|
||||
if (success && this.options.logSuccess) {
|
||||
logger.info(`RateLimiter (${this.options.keyPrefix}): under rate limit`, {
|
||||
limit,
|
||||
reset,
|
||||
remaining,
|
||||
identifier,
|
||||
});
|
||||
}
|
||||
|
||||
//log these by default
|
||||
if (!success && this.options.logFailure !== false) {
|
||||
logger.info(`RateLimiter (${this.options.keyPrefix}): rate limit exceeded`, {
|
||||
limit,
|
||||
reset,
|
||||
remaining,
|
||||
identifier,
|
||||
});
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
}
|
||||
|
||||
export function createRedisRateLimitClient(
|
||||
redisOptions: RedisWithClusterOptions
|
||||
): RateLimiterRedisClient {
|
||||
const redis = createRedisClient("trigger:rateLimiter", redisOptions);
|
||||
|
||||
return {
|
||||
sadd: async <TData>(key: string, ...members: TData[]): Promise<number> => {
|
||||
return redis.sadd(key, members as (string | number | Buffer)[]);
|
||||
},
|
||||
hset: <TValue>(
|
||||
key: string,
|
||||
obj: {
|
||||
[key: string]: TValue;
|
||||
}
|
||||
): Promise<number> => {
|
||||
return redis.hset(key, obj);
|
||||
},
|
||||
eval: <TArgs extends unknown[], TData = unknown>(
|
||||
...args: [script: string, keys: string[], args: TArgs]
|
||||
): Promise<TData> => {
|
||||
const script = args[0];
|
||||
const keys = args[1];
|
||||
const argsArray = args[2];
|
||||
return redis.eval(
|
||||
script,
|
||||
keys.length,
|
||||
...keys,
|
||||
...(argsArray as (string | Buffer | number)[])
|
||||
) as Promise<TData>;
|
||||
},
|
||||
};
|
||||
}
|
||||
@@ -0,0 +1,44 @@
|
||||
import { $replica, prisma } from "~/db.server";
|
||||
import type { PrismaClient } from "@trigger.dev/database";
|
||||
import plugin from "@trigger.dev/rbac";
|
||||
import { env } from "~/env.server";
|
||||
|
||||
// plugin.create() is synchronous — returns a lazy controller that resolves
|
||||
// any installed RBAC plugin on first call. Top-level await is not used
|
||||
// because CJS output format does not support it.
|
||||
//
|
||||
// Auth-path reads run on every request — pass the replica explicitly so
|
||||
// they don't pile up on the primary. Writes (role mutations) still go
|
||||
// through the primary. Same separation findEnvironmentByApiKey used
|
||||
// before this PR moved bearer auth into the RBAC plugin.
|
||||
//
|
||||
// Session-cookie userId resolution lives at the call site (see
|
||||
// dashboardBuilder.server.ts), not here. Statically importing
|
||||
// `~/services/session.server` from this module dragged the entire
|
||||
// remix-auth pipeline (auth.server → emailAuth/gitHubAuth/googleAuth,
|
||||
// each validating their secret at module load) into anything that
|
||||
// transitively imported `rbac` — including PAT auth callers that have
|
||||
// no session-cookie path at all. Passing userId through the
|
||||
// `authenticateSession` context decouples the plugin host from the
|
||||
// host's session implementation.
|
||||
export const rbac = plugin.create(
|
||||
// $replica is structurally a PrismaClient minus `$transaction` — the
|
||||
// RBAC fallback only uses `findFirst` on it, so the cast is safe.
|
||||
{ primary: prisma, replica: $replica as PrismaClient },
|
||||
// SESSION_SECRET signs delegated user-actor tokens; the plugin verifies
|
||||
// them with it in authenticateUserActor.
|
||||
{
|
||||
forceFallback: env.RBAC_FORCE_FALLBACK,
|
||||
userActorSecret: env.SESSION_SECRET,
|
||||
// A plugin that owns its own database client gets the same
|
||||
// writer/replica topology the webapp's Prisma clients use (see
|
||||
// getClient/getReplicaClient in db.server.ts): control-plane URLs win,
|
||||
// and with no replica configured reads share the writer.
|
||||
database: {
|
||||
writerUrl: env.CONTROL_PLANE_DATABASE_URL ?? env.DATABASE_URL,
|
||||
readerUrl: env.CONTROL_PLANE_DATABASE_READ_REPLICA_URL ?? env.DATABASE_READ_REPLICA_URL,
|
||||
writerConnectionLimit: env.RBAC_DATABASE_WRITER_CONNECTION_LIMIT,
|
||||
readerConnectionLimit: env.RBAC_DATABASE_READER_CONNECTION_LIMIT,
|
||||
},
|
||||
}
|
||||
);
|
||||
@@ -0,0 +1,50 @@
|
||||
/**
|
||||
* Tiny in-process bounded TTL cache shared by the realtime feeds: entries expire after `ttlMs` (evicted on read),
|
||||
* and at-capacity writes sweep expired entries then drop the oldest. A stored `undefined` is indistinguishable from a miss (use `null` for absence).
|
||||
*/
|
||||
export class BoundedTtlCache<V> {
|
||||
readonly #entries = new Map<string, { value: V; expiresAt: number }>();
|
||||
|
||||
constructor(
|
||||
private readonly ttlMs: number,
|
||||
private readonly maxEntries: number
|
||||
) {}
|
||||
|
||||
get(key: string): V | undefined {
|
||||
const entry = this.#entries.get(key);
|
||||
if (!entry) {
|
||||
return undefined;
|
||||
}
|
||||
if (entry.expiresAt > Date.now()) {
|
||||
return entry.value;
|
||||
}
|
||||
// Evict on read so expired entries don't linger until the next at-capacity
|
||||
// sweep — important for read-heavy / low-churn caches (per-handle working sets).
|
||||
this.#entries.delete(key);
|
||||
return undefined;
|
||||
}
|
||||
|
||||
set(key: string, value: V): void {
|
||||
// Only run capacity eviction when inserting a NEW key — updating an existing key
|
||||
// doesn't grow the map, so it must never drop an unrelated live entry.
|
||||
if (!this.#entries.has(key) && this.#entries.size >= this.maxEntries) {
|
||||
const now = Date.now();
|
||||
for (const [key, entry] of this.#entries) {
|
||||
if (entry.expiresAt <= now) {
|
||||
this.#entries.delete(key);
|
||||
}
|
||||
}
|
||||
if (this.#entries.size >= this.maxEntries) {
|
||||
const oldest = this.#entries.keys().next().value;
|
||||
if (oldest !== undefined) {
|
||||
this.#entries.delete(oldest);
|
||||
}
|
||||
}
|
||||
}
|
||||
this.#entries.set(key, { value, expiresAt: Date.now() + this.ttlMs });
|
||||
}
|
||||
|
||||
get size(): number {
|
||||
return this.#entries.size;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
import { env } from "~/env.server";
|
||||
|
||||
/**
|
||||
* Canonical storage URI for a session's chat.agent snapshot. Stamped on
|
||||
* `Session.chatSnapshotStoragePath` at row creation so PUT/GET presigns
|
||||
* resolve to the same store even if `OBJECT_STORE_DEFAULT_PROTOCOL`
|
||||
* changes later.
|
||||
*/
|
||||
export function chatSnapshotStoragePathForSession(friendlyId: string): string {
|
||||
const path = `sessions/${friendlyId}/snapshot.json`;
|
||||
const protocol = env.OBJECT_STORE_DEFAULT_PROTOCOL;
|
||||
return protocol ? `${protocol}://${path}` : path;
|
||||
}
|
||||
|
||||
/**
|
||||
* Resolve the storage key/URI a session's chat snapshot is written to and read
|
||||
* from. Single source of truth shared by every reader/writer so they all hit
|
||||
* the same object store:
|
||||
* - the SDK write + boot read (via the `snapshot-url` presign route), and
|
||||
* - the dashboard `SessionPresenter` (Agent/Session view).
|
||||
*
|
||||
* Prefers `chatSnapshotStoragePath` stamped at row creation (already
|
||||
* protocol-qualified, e.g. `s3://sessions/{id}/snapshot.json`), falling back to
|
||||
* recomputing it for sessions created before the column existed. Using a bare,
|
||||
* unqualified key here is the bug this guards against: the object store applies
|
||||
* `OBJECT_STORE_DEFAULT_PROTOCOL` to unprefixed keys on PUT but not on GET, so a
|
||||
* bare key can write to one store and read from another.
|
||||
*/
|
||||
export function chatSnapshotStorageKey(session: {
|
||||
friendlyId: string;
|
||||
chatSnapshotStoragePath: string | null;
|
||||
}): string {
|
||||
return session.chatSnapshotStoragePath ?? chatSnapshotStoragePathForSession(session.friendlyId);
|
||||
}
|
||||
@@ -0,0 +1,39 @@
|
||||
import { type ClickHouse } from "@internal/clickhouse";
|
||||
import { type PrismaClientOrTransaction } from "~/db.server";
|
||||
import { RunsRepository } from "~/services/runsRepository/runsRepository.server";
|
||||
import { type RunListFilter, type RunListResolver } from "./runReader.server";
|
||||
|
||||
export type ClickHouseRunListResolverOptions = {
|
||||
/** Resolves the per-organization ClickHouse client (multi-tenant routing). */
|
||||
getClickhouse: (organizationId: string) => Promise<ClickHouse>;
|
||||
prisma: PrismaClientOrTransaction;
|
||||
};
|
||||
|
||||
/**
|
||||
* Resolves the realtime tag/list filter into matching run ids via ClickHouse `listRunIds` (filter-only;
|
||||
* rows hydrated from Postgres by id afterward). Tag matching is contains-ALL, byte-matching Electric's
|
||||
* `runTags @> ARRAY[...]` shape.
|
||||
*/
|
||||
export class ClickHouseRunListResolver implements RunListResolver {
|
||||
constructor(private readonly options: ClickHouseRunListResolverOptions) {}
|
||||
|
||||
async resolveMatchingRunIds(filter: RunListFilter): Promise<string[]> {
|
||||
const clickhouse = await this.options.getClickhouse(filter.organizationId);
|
||||
const repository = new RunsRepository({ clickhouse, prisma: this.options.prisma });
|
||||
|
||||
const { runIds } = await repository.listRunIds({
|
||||
organizationId: filter.organizationId,
|
||||
projectId: filter.projectId,
|
||||
environmentId: filter.environmentId,
|
||||
tags: filter.tags && filter.tags.length > 0 ? filter.tags : undefined,
|
||||
// Contains-ALL, matching the Electric shape's `runTags @> ARRAY[...]` semantics.
|
||||
tagsMatch: "all",
|
||||
batchId: filter.batchId,
|
||||
from: filter.createdAtAfter?.getTime(),
|
||||
page: { size: filter.limit },
|
||||
});
|
||||
|
||||
// listRunIds is keyset-paginated; runIds is already capped to page.size (= limit).
|
||||
return runIds;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,48 @@
|
||||
/**
|
||||
* Duration string parsing for stream-basin retention / delete-on-empty
|
||||
* configuration. Used by `streamBasinProvisioner` (to convert to S2's
|
||||
* integer-seconds wire format) and by `env.server.ts` (to validate
|
||||
* duration-shaped env vars at boot rather than at first use).
|
||||
*
|
||||
* Accepts the short forms (`7d`, `30d`, `365d`, `1h`, `90m`, `45s`,
|
||||
* `2w`, `1y`) and the human forms (`7days`, `1week`, `1year`).
|
||||
*/
|
||||
|
||||
const PATTERN =
|
||||
/^(\d+)\s*(s|sec|secs|seconds?|m|min|mins|minutes?|h|hour|hours?|d|day|days?|w|week|weeks?|y|year|years?)$/;
|
||||
|
||||
export function isValidDuration(input: string): boolean {
|
||||
return PATTERN.test(input.trim().toLowerCase());
|
||||
}
|
||||
|
||||
/**
|
||||
* Parse a duration string into seconds. Throws on garbage so a
|
||||
* misconfigured env var fails loudly. Use {@link isValidDuration}
|
||||
* for non-throwing validation (e.g. inside a Zod `.refine()`).
|
||||
*/
|
||||
export function parseDuration(input: string): number {
|
||||
const trimmed = input.trim().toLowerCase();
|
||||
const match = trimmed.match(PATTERN);
|
||||
if (!match) {
|
||||
throw new Error(`Invalid duration string: ${input}`);
|
||||
}
|
||||
const value = parseInt(match[1]!, 10);
|
||||
const unit = match[2]!;
|
||||
const multiplier = /^s/.test(unit)
|
||||
? 1
|
||||
: /^m(?:in|ins|inute|inutes)?$/.test(unit)
|
||||
? 60
|
||||
: /^h/.test(unit)
|
||||
? 3600
|
||||
: /^d/.test(unit)
|
||||
? 86400
|
||||
: /^w/.test(unit)
|
||||
? 604800
|
||||
: /^y/.test(unit)
|
||||
? 31_536_000
|
||||
: NaN;
|
||||
if (!Number.isFinite(multiplier)) {
|
||||
throw new Error(`Invalid duration unit: ${unit}`);
|
||||
}
|
||||
return value * multiplier;
|
||||
}
|
||||
@@ -0,0 +1,281 @@
|
||||
/**
|
||||
* Pure (no DB/Redis/env) Electric HTTP shape-stream wire serializer, byte-faithful to what the
|
||||
* deployed `@electric-sql/client` (1.0.14 + 0.4.0) and the SDK's `SubscribeRunRawShape` expect.
|
||||
* Each column value is wire-encoded as a string (or null) decoded via the `electric-schema` header;
|
||||
* `up-to-date` is the only control message that makes the client emit, and re-sending a full row is idempotent.
|
||||
*/
|
||||
|
||||
export type ElectricColumnType =
|
||||
| "text"
|
||||
| "timestamp"
|
||||
| "int4"
|
||||
| "int8"
|
||||
| "float8"
|
||||
| "bool"
|
||||
| "jsonb";
|
||||
|
||||
type ElectricColumn = {
|
||||
name: string;
|
||||
type: ElectricColumnType;
|
||||
/** Array dimensionality. 1 => `type[]` (Postgres `{a,b}` literal). */
|
||||
dims?: number;
|
||||
/** Array columns only: true when the column has no SQL default, so an empty value emits `null` (not `{}`). Prisma erases this distinction, so we re-derive it here. */
|
||||
emptyArrayAsNull?: boolean;
|
||||
};
|
||||
|
||||
/** Columns the realtime run feed exposes; keep in sync with `DEFAULT_ELECTRIC_COLUMNS`. `type`/`dims` drive the schema header and value encoding. */
|
||||
export const RUN_ELECTRIC_COLUMNS: ReadonlyArray<ElectricColumn> = [
|
||||
{ name: "id", type: "text" },
|
||||
{ name: "taskIdentifier", type: "text" },
|
||||
{ name: "createdAt", type: "timestamp" },
|
||||
{ name: "updatedAt", type: "timestamp" },
|
||||
{ name: "startedAt", type: "timestamp" },
|
||||
{ name: "delayUntil", type: "timestamp" },
|
||||
{ name: "queuedAt", type: "timestamp" },
|
||||
{ name: "expiredAt", type: "timestamp" },
|
||||
{ name: "completedAt", type: "timestamp" },
|
||||
{ name: "friendlyId", type: "text" },
|
||||
{ name: "number", type: "int4" },
|
||||
{ name: "isTest", type: "bool" },
|
||||
{ name: "status", type: "text" },
|
||||
{ name: "usageDurationMs", type: "int4" },
|
||||
{ name: "costInCents", type: "float8" },
|
||||
{ name: "baseCostInCents", type: "float8" },
|
||||
{ name: "ttl", type: "text" },
|
||||
{ name: "payload", type: "text" },
|
||||
{ name: "payloadType", type: "text" },
|
||||
{ name: "metadata", type: "text" },
|
||||
{ name: "metadataType", type: "text" },
|
||||
{ name: "output", type: "text" },
|
||||
{ name: "outputType", type: "text" },
|
||||
{ name: "runTags", type: "text", dims: 1, emptyArrayAsNull: true },
|
||||
{ name: "error", type: "jsonb" },
|
||||
{ name: "realtimeStreams", type: "text", dims: 1 },
|
||||
];
|
||||
|
||||
/** Columns that can never be skipped via `skipColumns` (mirrors realtimeClient). */
|
||||
export const RESERVED_COLUMNS = ["id", "taskIdentifier", "friendlyId", "status", "createdAt"];
|
||||
|
||||
/** A single run hydrated for the realtime feed; structurally compatible with the `RunHydrator` Prisma `TaskRun` projection. */
|
||||
export type RealtimeRunRow = {
|
||||
id: string;
|
||||
taskIdentifier: string;
|
||||
createdAt: Date;
|
||||
updatedAt: Date;
|
||||
startedAt: Date | null;
|
||||
delayUntil: Date | null;
|
||||
queuedAt: Date | null;
|
||||
expiredAt: Date | null;
|
||||
completedAt: Date | null;
|
||||
friendlyId: string;
|
||||
number: number;
|
||||
isTest: boolean;
|
||||
status: string;
|
||||
usageDurationMs: number;
|
||||
costInCents: number;
|
||||
baseCostInCents: number;
|
||||
ttl: string | null;
|
||||
payload: string;
|
||||
payloadType: string;
|
||||
metadata: string | null;
|
||||
metadataType: string;
|
||||
output: string | null;
|
||||
outputType: string;
|
||||
runTags: string[];
|
||||
error: unknown;
|
||||
realtimeStreams: string[];
|
||||
};
|
||||
|
||||
type Operation = "insert" | "update" | "delete";
|
||||
|
||||
type ChangeMessage = {
|
||||
key: string;
|
||||
value: Record<string, string | null>;
|
||||
headers: { operation: Operation };
|
||||
};
|
||||
|
||||
type ControlMessage = {
|
||||
headers: { control: "up-to-date" | "must-refetch" };
|
||||
};
|
||||
|
||||
type ShapeMessage = ChangeMessage | ControlMessage;
|
||||
|
||||
const UP_TO_DATE: ControlMessage = { headers: { control: "up-to-date" } };
|
||||
|
||||
function effectiveSkipColumns(skipColumns: string[]): Set<string> {
|
||||
return new Set(skipColumns.filter((c) => c !== "" && !RESERVED_COLUMNS.includes(c)));
|
||||
}
|
||||
|
||||
function quoteArrayElement(value: string): string {
|
||||
return `"${value.replace(/\\/g, "\\\\").replace(/"/g, '\\"')}"`;
|
||||
}
|
||||
|
||||
function pgArrayLiteral(values: unknown[]): string {
|
||||
if (values.length === 0) {
|
||||
return "{}";
|
||||
}
|
||||
return `{${values.map((v) => quoteArrayElement(String(v))).join(",")}}`;
|
||||
}
|
||||
|
||||
function serializeValue(value: unknown, column: ElectricColumn): string | null {
|
||||
if (value === null || value === undefined) {
|
||||
return null;
|
||||
}
|
||||
|
||||
if (column.dims && column.dims > 0) {
|
||||
if (!Array.isArray(value)) {
|
||||
return null;
|
||||
}
|
||||
// A no-default array column stores NULL when empty, so Electric emits `null`
|
||||
// (not `{}`); match that here since Prisma handed us `[]` for the NULL value.
|
||||
if (value.length === 0 && column.emptyArrayAsNull) {
|
||||
return null;
|
||||
}
|
||||
return pgArrayLiteral(value);
|
||||
}
|
||||
|
||||
switch (column.type) {
|
||||
case "bool":
|
||||
// Postgres text representation; the client's parseBool accepts "t"/"f".
|
||||
return value ? "t" : "f";
|
||||
case "timestamp":
|
||||
// The SDK's RawShapeDate appends "Z" before parsing, so we emit the ISO
|
||||
// string WITHOUT the trailing "Z".
|
||||
return value instanceof Date ? value.toISOString().slice(0, -1) : String(value);
|
||||
case "jsonb":
|
||||
return JSON.stringify(value);
|
||||
case "int4":
|
||||
case "int8":
|
||||
case "float8":
|
||||
case "text":
|
||||
default:
|
||||
return String(value);
|
||||
}
|
||||
}
|
||||
|
||||
/** The merge key the client uses to reassemble a row across insert/update cycles. */
|
||||
export function runShapeKey(runId: string): string {
|
||||
return `"public"."TaskRun"/"${runId}"`;
|
||||
}
|
||||
|
||||
/** Encode a single run row into the wire `value` object (column -> string|null). */
|
||||
export function serializeRunRow(
|
||||
row: RealtimeRunRow,
|
||||
skipColumns: string[] = []
|
||||
): Record<string, string | null> {
|
||||
const skip = effectiveSkipColumns(skipColumns);
|
||||
const value: Record<string, string | null> = {};
|
||||
|
||||
for (const column of RUN_ELECTRIC_COLUMNS) {
|
||||
if (skip.has(column.name)) {
|
||||
continue;
|
||||
}
|
||||
value[column.name] = serializeValue((row as Record<string, unknown>)[column.name], column);
|
||||
}
|
||||
|
||||
return value;
|
||||
}
|
||||
|
||||
/** The `electric-schema` response header value for the (optionally trimmed) column set. */
|
||||
export function buildElectricSchemaHeader(skipColumns: string[] = []): string {
|
||||
const skip = effectiveSkipColumns(skipColumns);
|
||||
const schema: Record<string, { type: string; dims?: number }> = {};
|
||||
|
||||
for (const column of RUN_ELECTRIC_COLUMNS) {
|
||||
if (skip.has(column.name)) {
|
||||
continue;
|
||||
}
|
||||
schema[column.name] = column.dims
|
||||
? { type: column.type, dims: column.dims }
|
||||
: { type: column.type };
|
||||
}
|
||||
|
||||
return JSON.stringify(schema);
|
||||
}
|
||||
|
||||
/** Initial snapshot body: an `insert` for the row (if present) then `up-to-date`; an absent row emits a bare `up-to-date` (empty shape). */
|
||||
export function buildSnapshotBody(row: RealtimeRunRow | null, skipColumns: string[] = []): string {
|
||||
const messages: ShapeMessage[] = [];
|
||||
if (row) {
|
||||
messages.push({
|
||||
key: runShapeKey(row.id),
|
||||
value: serializeRunRow(row, skipColumns),
|
||||
headers: { operation: "insert" },
|
||||
});
|
||||
}
|
||||
messages.push(UP_TO_DATE);
|
||||
return JSON.stringify(messages);
|
||||
}
|
||||
|
||||
/** Live body when the row advanced: a full-row `update` followed by `up-to-date`. */
|
||||
export function buildUpdateBody(row: RealtimeRunRow, skipColumns: string[] = []): string {
|
||||
const messages: ShapeMessage[] = [
|
||||
{
|
||||
key: runShapeKey(row.id),
|
||||
value: serializeRunRow(row, skipColumns),
|
||||
headers: { operation: "update" },
|
||||
},
|
||||
UP_TO_DATE,
|
||||
];
|
||||
return JSON.stringify(messages);
|
||||
}
|
||||
|
||||
/** Live body when nothing advanced: a bare `up-to-date` (no row emission). */
|
||||
export function buildUpToDateBody(): string {
|
||||
return JSON.stringify([UP_TO_DATE]);
|
||||
}
|
||||
|
||||
export type RowChange = { row: RealtimeRunRow; operation: "insert" | "update" };
|
||||
|
||||
/** Multi-row body for the tag-list feed: one change message per row then `up-to-date` (empty `changes` emits a bare `up-to-date`). */
|
||||
export function buildRowsBody(changes: RowChange[], skipColumns: string[] = []): string {
|
||||
const messages: ShapeMessage[] = changes.map((change) => ({
|
||||
key: runShapeKey(change.row.id),
|
||||
value: serializeRunRow(change.row, skipColumns),
|
||||
headers: { operation: change.operation },
|
||||
}));
|
||||
messages.push(UP_TO_DATE);
|
||||
return JSON.stringify(messages);
|
||||
}
|
||||
|
||||
/** A row change whose wire `value` was already serialized (once, shared across feeds by
|
||||
* the EnvChangeRouter); the per-feed `operation` is applied here. */
|
||||
export type SerializedRowChange = {
|
||||
runId: string;
|
||||
value: Record<string, string | null>;
|
||||
operation: "insert" | "update";
|
||||
};
|
||||
|
||||
/** Like `buildRowsBody`, but from values serialized once per (runId, columnSet) upstream,
|
||||
* so a run matching many feeds is serialized once and reused across their bodies. */
|
||||
export function buildRowsBodyFromSerialized(changes: SerializedRowChange[]): string {
|
||||
const messages: ShapeMessage[] = changes.map((change) => ({
|
||||
key: runShapeKey(change.runId),
|
||||
value: change.value,
|
||||
headers: { operation: change.operation },
|
||||
}));
|
||||
messages.push(UP_TO_DATE);
|
||||
return JSON.stringify(messages);
|
||||
}
|
||||
|
||||
export const INITIAL_OFFSET = "-1";
|
||||
|
||||
/** Opaque `<updatedAtMs>_<seq>` offset token (client `${number}_${number}` type); the first segment lets a live request detect whether the row advanced. */
|
||||
export function encodeOffset(updatedAtMs: number, seq: number): string {
|
||||
return `${Math.trunc(updatedAtMs)}_${Math.trunc(seq)}`;
|
||||
}
|
||||
|
||||
/** Extract the `updatedAt` epoch-ms a client last saw from its echoed offset. */
|
||||
export function parseOffsetUpdatedAtMs(offset: string | null | undefined): number {
|
||||
if (!offset) {
|
||||
return 0;
|
||||
}
|
||||
const [first] = offset.split("_");
|
||||
const value = Number(first);
|
||||
return Number.isFinite(value) && value > 0 ? value : 0;
|
||||
}
|
||||
|
||||
/** Mirror of realtimeClient's DEQUEUED->EXECUTING rewrite for non-current API versions. */
|
||||
export function rewriteBodyForLegacyApiVersion(body: string): string {
|
||||
return body.replace(/"status":"DEQUEUED"/g, '"status":"EXECUTING"');
|
||||
}
|
||||
@@ -0,0 +1,695 @@
|
||||
import { type ChangeRecord } from "./runChangeNotifier.server";
|
||||
import { type RealtimeRunRow, serializeRunRow } from "./electricStreamProtocol.server";
|
||||
import { logger } from "~/services/logger.server";
|
||||
|
||||
/**
|
||||
* EnvChangeRouter — per-instance routing layer that fans one env's change stream out to the feeds it
|
||||
* matches. Owns one subscription per env (over the RunChangeNotifier) plus an inverted index of held
|
||||
* feeds, then per batch: routes via the index, batch-hydrates matched runs once per column set,
|
||||
* serializes each row's wire value once, and resolves each matched feed's pending wait. Stateless across reconnects.
|
||||
*/
|
||||
|
||||
export type WakeReason = "notify" | "timeout" | "abort";
|
||||
|
||||
/** A feed's membership predicate over the env stream. */
|
||||
export type FeedFilter =
|
||||
| { kind: "run"; runId: string }
|
||||
| { kind: "tag"; tags: string[]; createdAtFloorMs?: number }
|
||||
| { kind: "batch"; batchId: string };
|
||||
|
||||
/** A matched run handed to a feed: the hydrated row (for the feed's working-set diff) and
|
||||
* its wire `value` serialized once for this feed's column set (shared across feeds). */
|
||||
export type MatchedRow = { row: RealtimeRunRow; value: Record<string, string | null> };
|
||||
|
||||
export type WaitResult = { reason: WakeReason; rows: MatchedRow[] };
|
||||
|
||||
/** Minimal deps so the router is unit-testable without Redis/Postgres. */
|
||||
export interface EnvChangeSource {
|
||||
subscribeToEnv(environmentId: string, onBatch: (records: ChangeRecord[]) => void): () => void;
|
||||
}
|
||||
export interface RowHydrator {
|
||||
hydrateByIds(
|
||||
environmentId: string,
|
||||
ids: string[],
|
||||
skipColumns: string[]
|
||||
): Promise<RealtimeRunRow[]>;
|
||||
}
|
||||
|
||||
export type EnvChangeRouterOptions = {
|
||||
source: EnvChangeSource;
|
||||
hydrator: RowHydrator;
|
||||
/** Observability: a hydrate-by-id batch ran (count = runs hydrated this tick). */
|
||||
onHydrate?: (runCount: number) => void;
|
||||
/** How far back (ms) a newly-armed feed replays buffered records. 0 disables replay. */
|
||||
replayWindowMs?: number;
|
||||
/** Cap on buffered recent records per env (latest record per run). */
|
||||
replayMaxRunsPerEnv?: number;
|
||||
/** How long (ms) to keep an env subscribed + buffering after its last feed closes. 0 disables. */
|
||||
unsubscribeLingerMs?: number;
|
||||
/** Observability: a replay scan found candidates and delivered rows (or none survived). */
|
||||
onReplay?: (result: "delivered" | "empty") => void;
|
||||
/** Observability: a buffered record was evicted. `cap` evictions mean the env churns more
|
||||
* runs inside the window than the buffer holds (the replay guarantee is degrading). */
|
||||
onReplayEviction?: (reason: "cap" | "window") => void;
|
||||
/** Read-your-writes gate over the replica: delays wake-path hydrates until the replica
|
||||
* should have applied the change (record.updatedAtMs + lag + margin), and re-hydrates
|
||||
* rows the tripwire still finds stale. Omit to hydrate immediately (legacy behavior). */
|
||||
replicaLag?: ReplicaLagGate;
|
||||
};
|
||||
|
||||
export type ReplicaLagGate = {
|
||||
/** Current replica-lag estimate (ms). */
|
||||
getLagMs(): number;
|
||||
/** Feedback: a hydrate provably read at least this far behind the primary. */
|
||||
noteObservedLagMs(lagMs: number): void;
|
||||
/** Safety margin added on top of the estimate (clock skew + scheduling). */
|
||||
marginMs: number;
|
||||
/** Hard cap on any single gate delay — a sick replica degrades freshness, never liveness. */
|
||||
maxDelayMs: number;
|
||||
/** Re-hydrate attempts for rows that still read stale after the delay. */
|
||||
staleRetries: number;
|
||||
/** Observability: stale rows recovered by a retry, or delivered stale after exhausting them. */
|
||||
onStaleHydrate?: (outcome: "recovered" | "gave_up", runCount: number) => void;
|
||||
};
|
||||
|
||||
const DEFAULT_REPLAY_WINDOW_MS = 2_000;
|
||||
const DEFAULT_REPLAY_MAX_RUNS_PER_ENV = 512;
|
||||
const DEFAULT_UNSUBSCRIBE_LINGER_MS = 5_000;
|
||||
|
||||
/** Handle a feed holds for the duration of one long-poll. */
|
||||
export type FeedRegistration = {
|
||||
/** Wait for the next batch matching this feed (or timeout/abort), with the matched runs
|
||||
* hydrated + serialized for this feed's columns. One wait active at a time. */
|
||||
waitForMatch(signal: AbortSignal | undefined, timeoutMs: number): Promise<WaitResult>;
|
||||
/** Deregister from the index; unsubscribes the env when the last feed leaves. */
|
||||
close(): void;
|
||||
/** False when this instance's env subscription is younger than the replay window, so a
|
||||
* change in the caller's inter-poll gap may have been missed (hop/cold start) — the
|
||||
* caller should resolve once instead of holding blind. */
|
||||
gapCovered: boolean;
|
||||
};
|
||||
|
||||
type Feed = {
|
||||
filter: FeedFilter;
|
||||
skipColumns: string[];
|
||||
columnSig: string;
|
||||
/** The currently-waiting poll's resolver (null between polls). */
|
||||
resolve: ((result: WaitResult) => void) | null;
|
||||
/** Buffered records at or before this timestamp have been replayed (or predate this feed). */
|
||||
replayCursorMs: number;
|
||||
};
|
||||
|
||||
type EnvState = {
|
||||
unsubscribe: () => void;
|
||||
feeds: Set<Feed>;
|
||||
byRunId: Map<string, Set<Feed>>;
|
||||
byTag: Map<string, Set<Feed>>;
|
||||
byBatchId: Map<string, Set<Feed>>;
|
||||
/** All tag feeds, for routing partial records (no tags) as hydrate-to-classify candidates. */
|
||||
tagFeeds: Set<Feed>;
|
||||
/** Tag feeds with no tag filter — they match every record but are unreachable via byTag. */
|
||||
unfilteredTagFeeds: Set<Feed>;
|
||||
/** When this env's channel subscription started (for the gap-coverage check). */
|
||||
subscribedAtMs: number;
|
||||
/** Latest record per run, insertion-ordered, for replaying inter-poll gaps to newly-armed feeds. */
|
||||
recent: Map<string, { record: ChangeRecord; receivedAtMs: number }>;
|
||||
/** Pending teardown while the env lingers with zero feeds. */
|
||||
lingerTimer?: ReturnType<typeof setTimeout>;
|
||||
};
|
||||
|
||||
function sleepMs(ms: number): Promise<void> {
|
||||
return new Promise((resolve) => {
|
||||
const timer = setTimeout(resolve, ms);
|
||||
timer.unref?.();
|
||||
});
|
||||
}
|
||||
|
||||
function addToIndex(index: Map<string, Set<Feed>>, key: string, feed: Feed) {
|
||||
let set = index.get(key);
|
||||
if (!set) {
|
||||
set = new Set();
|
||||
index.set(key, set);
|
||||
}
|
||||
set.add(feed);
|
||||
}
|
||||
|
||||
function removeFromIndex(index: Map<string, Set<Feed>>, key: string, feed: Feed) {
|
||||
const set = index.get(key);
|
||||
if (set) {
|
||||
set.delete(feed);
|
||||
if (set.size === 0) {
|
||||
index.delete(key);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
export class EnvChangeRouter {
|
||||
readonly #envs = new Map<string, EnvState>();
|
||||
|
||||
constructor(private readonly options: EnvChangeRouterOptions) {}
|
||||
|
||||
register(
|
||||
environmentId: string,
|
||||
filter: FeedFilter,
|
||||
skipColumns: string[],
|
||||
opts?: {
|
||||
/** When the caller last received data for this connection. Bounds the replay to the
|
||||
* true inter-poll gap; older than the window can't be proven covered. */
|
||||
replaySinceMs?: number;
|
||||
}
|
||||
): FeedRegistration {
|
||||
const env = this.#ensureEnv(environmentId);
|
||||
const replayWindowMs = this.options.replayWindowMs ?? DEFAULT_REPLAY_WINDOW_MS;
|
||||
const now = Date.now();
|
||||
const windowFloorMs = now - replayWindowMs;
|
||||
const sinceMs = opts?.replaySinceMs ?? windowFloorMs;
|
||||
const feed: Feed = {
|
||||
filter,
|
||||
skipColumns,
|
||||
columnSig: skipColumns.length > 0 ? [...skipColumns].sort().join(",") : "",
|
||||
resolve: null,
|
||||
// First arm replays the caller's inter-poll gap; later arms only what arrived since.
|
||||
// The buffer only spans the window, so never rewind past it.
|
||||
replayCursorMs: Math.max(sinceMs, windowFloorMs),
|
||||
};
|
||||
|
||||
env.feeds.add(feed);
|
||||
this.#indexFeed(env, feed);
|
||||
|
||||
const waitForMatch = (signal: AbortSignal | undefined, timeoutMs: number) =>
|
||||
new Promise<WaitResult>((resolve) => {
|
||||
if (signal?.aborted) {
|
||||
resolve({ reason: "abort", rows: [] });
|
||||
return;
|
||||
}
|
||||
let settled = false;
|
||||
let timer: ReturnType<typeof setTimeout> | undefined;
|
||||
let onAbort: (() => void) | undefined;
|
||||
const settle = (result: WaitResult) => {
|
||||
if (settled) return;
|
||||
settled = true;
|
||||
feed.resolve = null;
|
||||
if (timer) clearTimeout(timer);
|
||||
if (signal && onAbort) signal.removeEventListener("abort", onAbort);
|
||||
resolve(result);
|
||||
};
|
||||
feed.resolve = settle;
|
||||
timer = setTimeout(() => settle({ reason: "timeout", rows: [] }), timeoutMs);
|
||||
timer.unref?.();
|
||||
if (signal) {
|
||||
onAbort = () => settle({ reason: "abort", rows: [] });
|
||||
signal.addEventListener("abort", onAbort, { once: true });
|
||||
}
|
||||
// Deliver any buffered records this feed hasn't seen (catches changes that
|
||||
// landed while the caller was between polls).
|
||||
if (replayWindowMs > 0 && env.recent.size > 0) {
|
||||
this.#replayRecent(environmentId, env, feed).catch((error) => {
|
||||
logger.error("[envChangeRouter] failed to replay buffered records", {
|
||||
environmentId,
|
||||
error,
|
||||
});
|
||||
});
|
||||
}
|
||||
});
|
||||
|
||||
const close = () => {
|
||||
if (!env.feeds.has(feed)) {
|
||||
return;
|
||||
}
|
||||
env.feeds.delete(feed);
|
||||
this.#deindexFeed(env, feed);
|
||||
// Resolve any in-flight wait so the poll doesn't hang.
|
||||
feed.resolve?.({ reason: "abort", rows: [] });
|
||||
feed.resolve = null;
|
||||
if (env.feeds.size === 0) {
|
||||
this.#scheduleEnvTeardown(environmentId, env);
|
||||
}
|
||||
};
|
||||
|
||||
return {
|
||||
waitForMatch,
|
||||
close,
|
||||
// Covered when this instance was already subscribed (and buffering) at the gap's
|
||||
// start, and the gap fits inside the buffer's window.
|
||||
gapCovered:
|
||||
replayWindowMs <= 0 || (env.subscribedAtMs <= sinceMs && sinceMs >= windowFloorMs),
|
||||
};
|
||||
}
|
||||
|
||||
/** Distinct environments currently routed (for metrics). */
|
||||
get activeEnvCount(): number {
|
||||
return this.#envs.size;
|
||||
}
|
||||
|
||||
/** Currently-held feeds by kind (for metrics) — the system's capacity unit. */
|
||||
get heldFeedCounts(): { run: number; tag: number; batch: number } {
|
||||
const counts = { run: 0, tag: 0, batch: 0 };
|
||||
for (const env of this.#envs.values()) {
|
||||
for (const feed of env.feeds) {
|
||||
counts[feed.filter.kind]++;
|
||||
}
|
||||
}
|
||||
return counts;
|
||||
}
|
||||
|
||||
#ensureEnv(environmentId: string): EnvState {
|
||||
const existing = this.#envs.get(environmentId);
|
||||
if (existing) {
|
||||
// A pending teardown is cancelled by new interest; the buffer survives the gap.
|
||||
if (existing.lingerTimer) {
|
||||
clearTimeout(existing.lingerTimer);
|
||||
existing.lingerTimer = undefined;
|
||||
}
|
||||
return existing;
|
||||
}
|
||||
const env: EnvState = {
|
||||
unsubscribe: () => {},
|
||||
feeds: new Set(),
|
||||
byRunId: new Map(),
|
||||
byTag: new Map(),
|
||||
byBatchId: new Map(),
|
||||
tagFeeds: new Set(),
|
||||
unfilteredTagFeeds: new Set(),
|
||||
subscribedAtMs: Date.now(),
|
||||
recent: new Map(),
|
||||
};
|
||||
this.#envs.set(environmentId, env);
|
||||
env.unsubscribe = this.options.source.subscribeToEnv(environmentId, (records) => {
|
||||
this.#bufferRecent(env, records);
|
||||
// Fire-and-forget; catch hydrate failures here (unhandled rejection exits the process) — waiters time out into the backstop.
|
||||
this.#onBatch(environmentId, env, records).catch((error) => {
|
||||
logger.error("[envChangeRouter] failed to route a change batch", {
|
||||
environmentId,
|
||||
error,
|
||||
});
|
||||
});
|
||||
});
|
||||
return env;
|
||||
}
|
||||
|
||||
/** Keep the env subscribed + buffering for a linger after its last feed closes, so a
|
||||
* client's next poll (or another instance hop landing back here) can replay the gap. */
|
||||
#scheduleEnvTeardown(environmentId: string, env: EnvState) {
|
||||
const lingerMs = this.options.unsubscribeLingerMs ?? DEFAULT_UNSUBSCRIBE_LINGER_MS;
|
||||
if (lingerMs <= 0) {
|
||||
this.#envs.delete(environmentId);
|
||||
env.unsubscribe();
|
||||
return;
|
||||
}
|
||||
if (env.lingerTimer) {
|
||||
clearTimeout(env.lingerTimer);
|
||||
}
|
||||
env.lingerTimer = setTimeout(() => {
|
||||
if (env.feeds.size === 0) {
|
||||
this.#envs.delete(environmentId);
|
||||
env.unsubscribe();
|
||||
}
|
||||
}, lingerMs);
|
||||
env.lingerTimer.unref?.();
|
||||
}
|
||||
|
||||
/** Upsert the latest record per run (insertion-ordered) and prune to the window + cap. */
|
||||
#bufferRecent(env: EnvState, records: ChangeRecord[]) {
|
||||
const windowMs = this.options.replayWindowMs ?? DEFAULT_REPLAY_WINDOW_MS;
|
||||
if (windowMs <= 0) {
|
||||
return;
|
||||
}
|
||||
const maxRuns = this.options.replayMaxRunsPerEnv ?? DEFAULT_REPLAY_MAX_RUNS_PER_ENV;
|
||||
const now = Date.now();
|
||||
for (const record of records) {
|
||||
env.recent.delete(record.runId);
|
||||
env.recent.set(record.runId, { record, receivedAtMs: now });
|
||||
}
|
||||
const cutoff = now - windowMs;
|
||||
for (const [runId, entry] of env.recent) {
|
||||
if (entry.receivedAtMs >= cutoff && env.recent.size <= maxRuns) {
|
||||
break;
|
||||
}
|
||||
this.options.onReplayEviction?.(entry.receivedAtMs < cutoff ? "window" : "cap");
|
||||
env.recent.delete(runId);
|
||||
}
|
||||
}
|
||||
|
||||
/** Whether a buffered record matches a feed's predicate (mirrors #onBatch's routing). */
|
||||
#recordMatchesFeed(record: ChangeRecord, feed: Feed): boolean {
|
||||
switch (feed.filter.kind) {
|
||||
case "run":
|
||||
return record.runId === feed.filter.runId;
|
||||
case "batch":
|
||||
return record.batchId != null && record.batchId === feed.filter.batchId;
|
||||
case "tag": {
|
||||
const tags = feed.filter.tags;
|
||||
// Unfiltered feed matches everything; partial record (no tags) = hydrate-to-classify.
|
||||
if (tags.length === 0 || record.tags === undefined) {
|
||||
return true;
|
||||
}
|
||||
return record.tags.some((tag) => tags.includes(tag));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/** How long to wait before hydrating so the replica has applied every change in the
|
||||
* batch: each record is safe at updatedAtMs + lag + margin (records without a watermark
|
||||
* anchor at now, degrading to a plain lag-sized delay). Capped — see ReplicaLagGate. */
|
||||
#gateDelayMs(records: ChangeRecord[]): number {
|
||||
const gate = this.options.replicaLag;
|
||||
if (!gate || records.length === 0) {
|
||||
return 0;
|
||||
}
|
||||
const now = Date.now();
|
||||
const lagMs = gate.getLagMs();
|
||||
let safeAtMs = 0;
|
||||
for (const record of records) {
|
||||
const anchorMs = record.updatedAtMs ?? now;
|
||||
safeAtMs = Math.max(safeAtMs, anchorMs + lagMs + gate.marginMs);
|
||||
}
|
||||
return Math.max(0, Math.min(safeAtMs - now, gate.maxDelayMs));
|
||||
}
|
||||
|
||||
/** Deliver buffered records newer than the feed's cursor through the normal
|
||||
* hydrate -> serialize -> settle pipeline. Already-seen rows diff to nothing downstream. */
|
||||
async #replayRecent(environmentId: string, env: EnvState, feed: Feed) {
|
||||
const cursor = feed.replayCursorMs;
|
||||
feed.replayCursorMs = Date.now();
|
||||
|
||||
const runIds: string[] = [];
|
||||
const candidateRecords: ChangeRecord[] = [];
|
||||
for (const [runId, entry] of env.recent) {
|
||||
if (entry.receivedAtMs > cursor && this.#recordMatchesFeed(entry.record, feed)) {
|
||||
runIds.push(runId);
|
||||
candidateRecords.push(entry.record);
|
||||
}
|
||||
}
|
||||
if (runIds.length === 0 || !feed.resolve) {
|
||||
return;
|
||||
}
|
||||
|
||||
// Replayed records are usually past the lag window already (delay computes to 0); a
|
||||
// just-buffered one gets the same read-your-writes gate as the live path. No tripwire
|
||||
// here — a stale replay diffs to a re-emission on the next wake or backstop.
|
||||
const replayDelayMs = this.#gateDelayMs(candidateRecords);
|
||||
if (replayDelayMs > 0) {
|
||||
await sleepMs(replayDelayMs);
|
||||
if (!feed.resolve) {
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
const hydrated = await this.options.hydrator.hydrateByIds(
|
||||
environmentId,
|
||||
runIds,
|
||||
feed.skipColumns
|
||||
);
|
||||
this.options.onHydrate?.(hydrated.length);
|
||||
|
||||
const rows: MatchedRow[] = [];
|
||||
for (const row of hydrated) {
|
||||
if (feed.filter.kind === "tag" && !this.#tagRowMatches(row, feed.filter)) {
|
||||
continue;
|
||||
}
|
||||
rows.push({ row, value: serializeRunRow(row, feed.skipColumns) });
|
||||
}
|
||||
|
||||
if (rows.length > 0 && feed.resolve) {
|
||||
this.options.onReplay?.("delivered");
|
||||
feed.resolve({ reason: "notify", rows });
|
||||
} else {
|
||||
this.options.onReplay?.("empty");
|
||||
}
|
||||
}
|
||||
|
||||
#indexFeed(env: EnvState, feed: Feed) {
|
||||
switch (feed.filter.kind) {
|
||||
case "run":
|
||||
addToIndex(env.byRunId, feed.filter.runId, feed);
|
||||
break;
|
||||
case "batch":
|
||||
addToIndex(env.byBatchId, feed.filter.batchId, feed);
|
||||
break;
|
||||
case "tag":
|
||||
env.tagFeeds.add(feed);
|
||||
if (feed.filter.tags.length === 0) {
|
||||
env.unfilteredTagFeeds.add(feed);
|
||||
}
|
||||
for (const tag of feed.filter.tags) {
|
||||
addToIndex(env.byTag, tag, feed);
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
#deindexFeed(env: EnvState, feed: Feed) {
|
||||
switch (feed.filter.kind) {
|
||||
case "run":
|
||||
removeFromIndex(env.byRunId, feed.filter.runId, feed);
|
||||
break;
|
||||
case "batch":
|
||||
removeFromIndex(env.byBatchId, feed.filter.batchId, feed);
|
||||
break;
|
||||
case "tag":
|
||||
env.tagFeeds.delete(feed);
|
||||
env.unfilteredTagFeeds.delete(feed);
|
||||
for (const tag of feed.filter.tags) {
|
||||
removeFromIndex(env.byTag, tag, feed);
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
async #onBatch(environmentId: string, env: EnvState, records: ChangeRecord[], attempt = 0) {
|
||||
// 0. Read-your-writes gate: wait out the replica's apply lag before hydrating, so the
|
||||
// rows we read contain the changes the records announce. Retry attempts were
|
||||
// scheduled with their own delay, so only the first pass gates here.
|
||||
if (attempt === 0) {
|
||||
const delayMs = this.#gateDelayMs(records);
|
||||
if (delayMs > 0) {
|
||||
await sleepMs(delayMs);
|
||||
}
|
||||
}
|
||||
|
||||
// 1. Route each record to the held feeds it matches; collect matched runIds per feed.
|
||||
const matchedRunIdsByFeed = new Map<Feed, Set<string>>();
|
||||
const addMatch = (feed: Feed, runId: string) => {
|
||||
if (!feed.resolve) {
|
||||
// Feed isn't currently waiting (between polls). Drop — its backstop catches gaps.
|
||||
return;
|
||||
}
|
||||
let set = matchedRunIdsByFeed.get(feed);
|
||||
if (!set) {
|
||||
set = new Set();
|
||||
matchedRunIdsByFeed.set(feed, set);
|
||||
}
|
||||
set.add(runId);
|
||||
};
|
||||
|
||||
for (const record of records) {
|
||||
// run feeds: exact runId match.
|
||||
const runFeeds = env.byRunId.get(record.runId);
|
||||
if (runFeeds) {
|
||||
for (const feed of runFeeds) addMatch(feed, record.runId);
|
||||
}
|
||||
|
||||
// batch feeds: exact batchId match (only when the record carries one).
|
||||
if (record.batchId) {
|
||||
const batchFeeds = env.byBatchId.get(record.batchId);
|
||||
if (batchFeeds) {
|
||||
for (const feed of batchFeeds) addMatch(feed, record.runId);
|
||||
}
|
||||
}
|
||||
|
||||
// tag feeds.
|
||||
if (record.tags !== undefined) {
|
||||
// Full record: prune via the tag index; only feeds whose filter intersects match.
|
||||
const seen = new Set<Feed>();
|
||||
for (const tag of record.tags) {
|
||||
const tagFeeds = env.byTag.get(tag);
|
||||
if (!tagFeeds) continue;
|
||||
for (const feed of tagFeeds) {
|
||||
if (seen.has(feed)) continue;
|
||||
seen.add(feed);
|
||||
addMatch(feed, record.runId);
|
||||
}
|
||||
}
|
||||
// Unfiltered tag feeds match every record but live outside the index.
|
||||
for (const feed of env.unfilteredTagFeeds) addMatch(feed, record.runId);
|
||||
} else {
|
||||
// Partial record (no membership data): route to every tag feed as a candidate to
|
||||
// hydrate-and-classify (rare; the publish side emits full records in practice).
|
||||
for (const feed of env.tagFeeds) addMatch(feed, record.runId);
|
||||
}
|
||||
}
|
||||
|
||||
if (matchedRunIdsByFeed.size === 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
// 2. Batch-hydrate ONCE per column set, then 3. serialize ONCE per (runId, column set).
|
||||
const runIdsByColumnSig = new Map<string, { skipColumns: string[]; runIds: Set<string> }>();
|
||||
for (const [feed, runIds] of matchedRunIdsByFeed) {
|
||||
let group = runIdsByColumnSig.get(feed.columnSig);
|
||||
if (!group) {
|
||||
group = { skipColumns: feed.skipColumns, runIds: new Set() };
|
||||
runIdsByColumnSig.set(feed.columnSig, group);
|
||||
}
|
||||
for (const id of runIds) group.runIds.add(id);
|
||||
}
|
||||
|
||||
const hydratedByColumnSig = new Map<string, Map<string, MatchedRow>>();
|
||||
await Promise.all(
|
||||
[...runIdsByColumnSig.entries()].map(async ([columnSig, group]) => {
|
||||
const ids = [...group.runIds];
|
||||
const rows = await this.options.hydrator.hydrateByIds(
|
||||
environmentId,
|
||||
ids,
|
||||
group.skipColumns
|
||||
);
|
||||
this.options.onHydrate?.(rows.length);
|
||||
const map = new Map<string, MatchedRow>();
|
||||
for (const row of rows) {
|
||||
map.set(row.id, { row, value: serializeRunRow(row, group.skipColumns) });
|
||||
}
|
||||
hydratedByColumnSig.set(columnSig, map);
|
||||
})
|
||||
);
|
||||
|
||||
// 3.5 Stale tripwire: a watermarked record whose hydrated row is older (or missing —
|
||||
// the insert race) read a replica that hadn't applied the change. Withhold those
|
||||
// rows and re-hydrate shortly. Exhausting the retry budget delivers what we have
|
||||
// (liveness over freshness) — but a stale emission advances the feed's cursor, so
|
||||
// it ALSO schedules echo passes past the gate: re-hydrates flowing through normal
|
||||
// emission, where the working-set diff drops unchanged rows and emits the fresh
|
||||
// version once the replica catches up. The backstop stays the terminal net.
|
||||
// Each detection feeds the lag estimator.
|
||||
const gate = this.options.replicaLag;
|
||||
const isEchoPass = gate !== undefined && attempt > gate.staleRetries;
|
||||
const staleRunIds = gate
|
||||
? this.#detectStaleRuns(records, runIdsByColumnSig, hydratedByColumnSig)
|
||||
: new Set<string>();
|
||||
if (attempt > 0 && !isEchoPass) {
|
||||
const recovered = new Set(records.map((r) => r.runId)).size - staleRunIds.size;
|
||||
if (recovered > 0) {
|
||||
gate?.onStaleHydrate?.("recovered", recovered);
|
||||
}
|
||||
}
|
||||
if (staleRunIds.size > 0 && gate) {
|
||||
const staleRecords = records.filter((record) => staleRunIds.has(record.runId));
|
||||
// Re-buffer the withheld records so a feed that re-arms between now and the next
|
||||
// pass replays them instead of waiting for its backstop.
|
||||
this.#bufferRecent(env, staleRecords);
|
||||
if (attempt >= gate.staleRetries) {
|
||||
// Budget exhausted: deliver the stale rows below (liveness) — but a stale emission
|
||||
// advances the feed's cursor, so keep echoing re-hydrates through normal emission
|
||||
// (the working-set diff drops unchanged rows, emits the fresh version when the
|
||||
// replica catches up). Echoes stop once the change ages past the horizon; deeper
|
||||
// outages are the backstop's job.
|
||||
if (attempt === gate.staleRetries) {
|
||||
gate.onStaleHydrate?.("gave_up", staleRunIds.size);
|
||||
}
|
||||
staleRunIds.clear();
|
||||
}
|
||||
const echoHorizonMs = gate.maxDelayMs * 10;
|
||||
const newestWatermarkMs = Math.max(...staleRecords.map((record) => record.updatedAtMs ?? 0));
|
||||
const withinEchoHorizon = Date.now() - newestWatermarkMs < echoHorizonMs;
|
||||
if (attempt < gate.staleRetries || withinEchoHorizon) {
|
||||
const retryDelayMs = Math.max(
|
||||
25,
|
||||
Math.min(gate.getLagMs() + gate.marginMs, gate.maxDelayMs)
|
||||
);
|
||||
const timer = setTimeout(() => {
|
||||
this.#onBatch(environmentId, env, staleRecords, attempt + 1).catch((error) => {
|
||||
logger.error("[envChangeRouter] failed to re-hydrate stale rows", {
|
||||
environmentId,
|
||||
error,
|
||||
});
|
||||
});
|
||||
}, retryDelayMs);
|
||||
timer.unref?.();
|
||||
}
|
||||
}
|
||||
|
||||
// 4. Assemble each feed's matched rows (post-filtering tag feeds against the
|
||||
// authoritative hydrated row) and resolve its pending wait.
|
||||
for (const [feed, runIds] of matchedRunIdsByFeed) {
|
||||
if (!feed.resolve) {
|
||||
continue; // stopped waiting while we hydrated; its next poll/backstop covers it
|
||||
}
|
||||
const hydrated = hydratedByColumnSig.get(feed.columnSig);
|
||||
if (!hydrated) continue;
|
||||
|
||||
const rows: MatchedRow[] = [];
|
||||
for (const runId of runIds) {
|
||||
if (staleRunIds.has(runId)) {
|
||||
continue; // withheld; the scheduled re-hydrate delivers the fresh version
|
||||
}
|
||||
const matched = hydrated.get(runId);
|
||||
if (!matched) continue; // run not found / left the table
|
||||
if (feed.filter.kind === "tag" && !this.#tagRowMatches(matched.row, feed.filter)) {
|
||||
continue; // re-confirm tags + createdAt floor against the authoritative row
|
||||
}
|
||||
rows.push(matched);
|
||||
}
|
||||
|
||||
if (rows.length > 0) {
|
||||
feed.resolve({ reason: "notify", rows });
|
||||
}
|
||||
// No surviving rows (e.g. a partial-record candidate that didn't actually match):
|
||||
// leave the feed waiting; nothing relevant changed for it.
|
||||
}
|
||||
}
|
||||
|
||||
/** Runs whose hydrated row is provably behind its record's watermark (stale content),
|
||||
* or absent entirely despite a watermark (the insert hasn't applied). Records without
|
||||
* `updatedAtMs` can't be judged and always pass. */
|
||||
#detectStaleRuns(
|
||||
records: ChangeRecord[],
|
||||
runIdsByColumnSig: Map<string, { skipColumns: string[]; runIds: Set<string> }>,
|
||||
hydratedByColumnSig: Map<string, Map<string, MatchedRow>>
|
||||
): Set<string> {
|
||||
const gate = this.options.replicaLag;
|
||||
const stale = new Set<string>();
|
||||
if (!gate) {
|
||||
return stale;
|
||||
}
|
||||
const expectedByRunId = new Map<string, number>();
|
||||
for (const record of records) {
|
||||
if (record.updatedAtMs !== undefined) {
|
||||
const existing = expectedByRunId.get(record.runId);
|
||||
if (existing === undefined || record.updatedAtMs > existing) {
|
||||
expectedByRunId.set(record.runId, record.updatedAtMs);
|
||||
}
|
||||
}
|
||||
}
|
||||
if (expectedByRunId.size === 0) {
|
||||
return stale;
|
||||
}
|
||||
const now = Date.now();
|
||||
for (const [columnSig, group] of runIdsByColumnSig) {
|
||||
const hydrated = hydratedByColumnSig.get(columnSig);
|
||||
for (const runId of group.runIds) {
|
||||
const expected = expectedByRunId.get(runId);
|
||||
if (expected === undefined || stale.has(runId)) {
|
||||
continue;
|
||||
}
|
||||
const matched = hydrated?.get(runId);
|
||||
if (!matched || matched.row.updatedAt.getTime() < expected) {
|
||||
stale.add(runId);
|
||||
gate.noteObservedLagMs(now - expected);
|
||||
}
|
||||
}
|
||||
}
|
||||
return stale;
|
||||
}
|
||||
|
||||
/** Authoritative re-check for tag feeds: the hydrated row carries ALL the filter's tags
|
||||
* (Electric's `runTags @> ARRAY[...]` semantics) and its createdAt is within the window. */
|
||||
#tagRowMatches(row: RealtimeRunRow, filter: Extract<FeedFilter, { kind: "tag" }>): boolean {
|
||||
if (
|
||||
filter.createdAtFloorMs !== undefined &&
|
||||
row.createdAt.getTime() < filter.createdAtFloorMs
|
||||
) {
|
||||
return false;
|
||||
}
|
||||
const rowTags = row.runTags ?? [];
|
||||
return filter.tags.every((tag) => rowTags.includes(tag));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,164 @@
|
||||
import { validateJWT, type ValidationResult } from "@trigger.dev/core/v3/jwt";
|
||||
import { $replica } from "~/db.server";
|
||||
import { findEnvironmentById } from "~/models/runtimeEnvironment.server";
|
||||
import type { AuthenticatedEnvironment } from "../apiAuth.server";
|
||||
|
||||
export type ValidatePublicJwtKeySuccess = {
|
||||
ok: true;
|
||||
environment: AuthenticatedEnvironment;
|
||||
claims: Record<string, unknown>;
|
||||
};
|
||||
|
||||
export type ValidatePublicJwtKeyError = {
|
||||
ok: false;
|
||||
error: string;
|
||||
};
|
||||
|
||||
export type ValidatePublicJwtKeyResult = ValidatePublicJwtKeySuccess | ValidatePublicJwtKeyError;
|
||||
|
||||
export async function validatePublicJwtKey(token: string): Promise<ValidatePublicJwtKeyResult> {
|
||||
// Get the sub claim from the token
|
||||
// Use the sub claim to find the environment
|
||||
// Validate the token against the environment.apiKey
|
||||
// Once that's done, return the environment and the claims
|
||||
const sub = extractJWTSub(token);
|
||||
|
||||
if (!sub) {
|
||||
return { ok: false, error: "Invalid Public Access Token, missing subject." };
|
||||
}
|
||||
|
||||
const environment = await findEnvironmentById(sub);
|
||||
|
||||
if (!environment) {
|
||||
return { ok: false, error: "Invalid Public Access Token, environment not found." };
|
||||
}
|
||||
|
||||
let result = await validateJWT(
|
||||
token,
|
||||
environment.parentEnvironment?.apiKey ?? environment.apiKey
|
||||
);
|
||||
|
||||
// PATs are signed with the env's apiKey at mint time. If the env's apiKey
|
||||
// has since been rotated, signature verification fails against the current
|
||||
// key — fall back to any RevokedApiKey rows still in their grace window.
|
||||
// Only run this query on the failure path so the success path is unchanged.
|
||||
if (!result.ok) {
|
||||
result = await validateAgainstRevokedApiKeys(
|
||||
token,
|
||||
environment.parentEnvironment?.id ?? environment.id,
|
||||
result
|
||||
);
|
||||
}
|
||||
|
||||
if (!result.ok) {
|
||||
switch (result.code) {
|
||||
case "ERR_JWT_EXPIRED": {
|
||||
return {
|
||||
ok: false,
|
||||
error:
|
||||
"Public Access Token has expired. See https://trigger.dev/docs/frontend/overview#authentication for more information.",
|
||||
};
|
||||
}
|
||||
case "ERR_JWT_CLAIM_INVALID": {
|
||||
return {
|
||||
ok: false,
|
||||
error: `Public Access Token is invalid: ${result.error}. See https://trigger.dev/docs/frontend/overview#authentication for more information.`,
|
||||
};
|
||||
}
|
||||
default: {
|
||||
return {
|
||||
ok: false,
|
||||
error:
|
||||
"Public Access Token is invalid. See https://trigger.dev/docs/frontend/overview#authentication for more information.",
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
ok: true,
|
||||
environment,
|
||||
claims: result.payload,
|
||||
};
|
||||
}
|
||||
|
||||
async function validateAgainstRevokedApiKeys(
|
||||
token: string,
|
||||
signingEnvironmentId: string,
|
||||
primaryResult: ValidationResult
|
||||
): Promise<ValidationResult> {
|
||||
const revokedApiKeys = await $replica.revokedApiKey.findMany({
|
||||
where: {
|
||||
runtimeEnvironmentId: signingEnvironmentId,
|
||||
expiresAt: { gt: new Date() },
|
||||
},
|
||||
select: { apiKey: true },
|
||||
});
|
||||
|
||||
for (const { apiKey } of revokedApiKeys) {
|
||||
const fallbackResult = await validateJWT(token, apiKey);
|
||||
if (fallbackResult.ok) {
|
||||
return fallbackResult;
|
||||
}
|
||||
}
|
||||
|
||||
return primaryResult;
|
||||
}
|
||||
|
||||
export function isPublicJWT(token: string): boolean {
|
||||
// Split the token
|
||||
const parts = token.split(".");
|
||||
if (parts.length !== 3) return false;
|
||||
|
||||
try {
|
||||
// Decode the payload (second part)
|
||||
const payload = JSON.parse(decodeBase64Url(parts[1]));
|
||||
|
||||
if (payload === null || typeof payload !== "object") return false;
|
||||
|
||||
// Check for the pub: true claim
|
||||
return "pub" in payload && payload.pub === true;
|
||||
} catch (_error) {
|
||||
// If there's any error in decoding or parsing, it's not a valid JWT
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
export function extractJwtSigningSecretKey(environment: AuthenticatedEnvironment) {
|
||||
return environment.parentEnvironment?.apiKey ?? environment.apiKey;
|
||||
}
|
||||
|
||||
function extractJWTSub(token: string): string | undefined {
|
||||
// Split the token
|
||||
const parts = token.split(".");
|
||||
if (parts.length !== 3) return;
|
||||
|
||||
try {
|
||||
// Decode the payload (second part)
|
||||
const payload = JSON.parse(decodeBase64Url(parts[1]));
|
||||
|
||||
if (payload === null || typeof payload !== "object") return;
|
||||
|
||||
// Check for the pub: true claim
|
||||
return "sub" in payload && typeof payload.sub === "string" ? payload.sub : undefined;
|
||||
} catch (_error) {
|
||||
// If there's any error in decoding or parsing, it's not a valid JWT
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
function decodeBase64Url(str: string): string {
|
||||
// Replace URL-safe characters and add padding
|
||||
str = str.replace(/-/g, "+").replace(/_/g, "/");
|
||||
switch (str.length % 4) {
|
||||
case 2:
|
||||
str += "==";
|
||||
break;
|
||||
case 3:
|
||||
str += "=";
|
||||
break;
|
||||
}
|
||||
|
||||
// Decode using Node.js Buffer
|
||||
return Buffer.from(str, "base64").toString("utf8");
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
import { generateJWT as internal_generateJWT } from "@trigger.dev/core/v3";
|
||||
import { extractJwtSigningSecretKey } from "./jwtAuth.server";
|
||||
|
||||
type Environment = Parameters<typeof extractJwtSigningSecretKey>[0];
|
||||
|
||||
export type MintRunTokenOptions = {
|
||||
/** Include the input-stream write scope (needed for steering messages from the playground). */
|
||||
includeInputStreamWrite?: boolean;
|
||||
/** Token expiration. Defaults to "1h". */
|
||||
expirationTime?: string;
|
||||
};
|
||||
|
||||
/**
|
||||
* Mint a run-scoped public access token (JWT) for browser subscription to a
|
||||
* run's realtime streams.
|
||||
*
|
||||
* Used by:
|
||||
* - The playground action to give a freshly triggered chat session a token.
|
||||
* - The run details page to let the agent view subscribe to the chat stream
|
||||
* of an existing run (read-only).
|
||||
*/
|
||||
export async function mintRunToken(
|
||||
environment: Environment,
|
||||
runFriendlyId: string,
|
||||
options: MintRunTokenOptions = {}
|
||||
): Promise<string> {
|
||||
const scopes = [`read:runs:${runFriendlyId}`];
|
||||
if (options.includeInputStreamWrite) {
|
||||
scopes.push(`write:inputStreams:${runFriendlyId}`);
|
||||
}
|
||||
|
||||
return internal_generateJWT({
|
||||
secretKey: extractJwtSigningSecretKey(environment),
|
||||
payload: {
|
||||
sub: environment.id,
|
||||
pub: true,
|
||||
scopes,
|
||||
},
|
||||
expirationTime: options.expirationTime ?? "1h",
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,40 @@
|
||||
import { generateJWT as internal_generateJWT } from "@trigger.dev/core/v3";
|
||||
import { extractJwtSigningSecretKey } from "./jwtAuth.server";
|
||||
|
||||
type Environment = Parameters<typeof extractJwtSigningSecretKey>[0];
|
||||
|
||||
export type MintSessionTokenOptions = {
|
||||
/** Token expiration. Defaults to "1h". */
|
||||
expirationTime?: string;
|
||||
};
|
||||
|
||||
/**
|
||||
* Mint a session-scoped public access token (JWT) covering both `.in`
|
||||
* append and `.out` subscribe for a session's realtime channels.
|
||||
*
|
||||
* Returned by `POST /api/v1/sessions` so the browser holds a single
|
||||
* long-lived token that survives across runs (sessions outlive any
|
||||
* single run). Includes both read and write scopes since the transport
|
||||
* needs both: read for SSE subscribe on `.out`, write for `.in` appends
|
||||
* (`stop`, follow-up messages, action chunks).
|
||||
*/
|
||||
export async function mintSessionToken(
|
||||
environment: Environment,
|
||||
sessionAddressingKey: string,
|
||||
options: MintSessionTokenOptions = {}
|
||||
): Promise<string> {
|
||||
const scopes = [
|
||||
`read:sessions:${sessionAddressingKey}`,
|
||||
`write:sessions:${sessionAddressingKey}`,
|
||||
];
|
||||
|
||||
return internal_generateJWT({
|
||||
secretKey: extractJwtSigningSecretKey(environment),
|
||||
payload: {
|
||||
sub: environment.id,
|
||||
pub: true,
|
||||
scopes,
|
||||
},
|
||||
expirationTime: options.expirationTime ?? "1h",
|
||||
});
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,272 @@
|
||||
import { getMeter } from "@internal/tracing";
|
||||
import { $replica } from "~/db.server";
|
||||
import { runStore } from "~/v3/runStore.server";
|
||||
import { env } from "~/env.server";
|
||||
import { singleton } from "~/utils/singleton";
|
||||
import { getCachedLimit } from "../platform.v3.server";
|
||||
import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactoryInstance.server";
|
||||
import { ClickHouseRunListResolver } from "./clickHouseRunListResolver.server";
|
||||
import { EnvChangeRouter, type EnvChangeSource } from "./envChangeRouter.server";
|
||||
import { NativeRealtimeClient } from "./nativeRealtimeClient.server";
|
||||
import { RealtimeConcurrencyLimiter } from "./realtimeConcurrencyLimiter.server";
|
||||
import { getRunChangeNotifier } from "./runChangeNotifierInstance.server";
|
||||
import { RedisReplayCursorStore } from "./replayCursorStore.server";
|
||||
import { createPostgresReplicaLagSource, ReplicaLagEstimator } from "./replicaLagEstimator.server";
|
||||
import { RunHydrator } from "./runReader.server";
|
||||
|
||||
// Process-singleton wiring for the native realtime client; only constructed when a
|
||||
// request actually routes to it, so a disabled webapp never instantiates it.
|
||||
function initializeNativeRealtimeClient(): NativeRealtimeClient {
|
||||
const meter = getMeter("realtime-native");
|
||||
|
||||
const wakeups = meter.createCounter("realtime_native.wakeups", {
|
||||
description:
|
||||
"Live realtime wakeups by reason. A rising 'timeout' share suggests a write site is missing its publishChangeRecord delegate.",
|
||||
});
|
||||
|
||||
const runSetResolves = meter.createCounter("realtime_native.runset_resolves", {
|
||||
description:
|
||||
"Multi-run (tag-list/batch) resolve+hydrate outcomes. 'hit'/'coalesced' vs 'miss' shows how effectively concurrent same-filter feeds share a single ClickHouse + Postgres query.",
|
||||
});
|
||||
|
||||
const runSetQueryMs = meter.createHistogram("realtime_native.runset_query_ms", {
|
||||
description: "Latency of the multi-run resolve (ClickHouse) and hydrate (Postgres) stages.",
|
||||
unit: "ms",
|
||||
});
|
||||
|
||||
const livePollPaths = meter.createCounter("realtime_native.live_polls", {
|
||||
description:
|
||||
"How live polls resolved. 'fast-hydrate' = router wake with rows hydrated by id (no ClickHouse); 'full-resolve' = backstop; 'cold-resolve' = fresh env subscription probed once.",
|
||||
});
|
||||
|
||||
const routerHydrates = meter.createCounter("realtime_native.router_hydrated_runs", {
|
||||
description:
|
||||
"Runs hydrated by the EnvChangeRouter's batch-hydrate (one query per column set per wake, shared across all feeds matching the same run).",
|
||||
});
|
||||
|
||||
const resolveAdmissionWaits = meter.createCounter("realtime_native.resolve_admission_waits", {
|
||||
description:
|
||||
"Fresh ClickHouse resolves that had to queue for an admission permit. A rising count means a distinct-filter reconnect stampede is being throttled (the gate is doing its job).",
|
||||
});
|
||||
|
||||
const replays = meter.createCounter("realtime_native.replays", {
|
||||
description:
|
||||
"Buffered change records replayed to a newly-armed feed (inter-poll gap recovery). 'delivered' = rows reached the feed; 'empty' = candidates hydrated but none survived the filter/diff.",
|
||||
});
|
||||
|
||||
const replayEvictions = meter.createCounter("realtime_native.replay_evictions", {
|
||||
description:
|
||||
"Replay-buffer evictions. 'window' expiry is normal; 'cap' means an env churns more runs inside the window than the buffer holds (replay guarantee degrading — retune the knobs).",
|
||||
});
|
||||
|
||||
const deliveryLagMs = meter.createHistogram("realtime_native.delivery_lag_ms", {
|
||||
description:
|
||||
"Live emissions: now minus the newest emitted row's updatedAt (PG clock vs app clock, so approximate). The end-to-end delivery SLI — a p99 near the backstop hold means wakes are being missed.",
|
||||
unit: "ms",
|
||||
});
|
||||
|
||||
const emittedRows = meter.createHistogram("realtime_native.emitted_rows", {
|
||||
description:
|
||||
"Rows per live emission. Deltas should be small; a fat tail means working-set/offset-floor fallbacks are re-emitting full sets.",
|
||||
unit: "rows",
|
||||
});
|
||||
|
||||
const backstops = meter.createCounter("realtime_native.backstops", {
|
||||
description:
|
||||
"Backstop full resolves by outcome. 'empty' is normal idle behavior; sustained 'delivered' means the notify/replay path missed changes — alert on it.",
|
||||
});
|
||||
|
||||
const concurrencyRejections = meter.createCounter("realtime_native.concurrency_rejections", {
|
||||
description: "Polls rejected (429) by the per-env concurrency limiter.",
|
||||
});
|
||||
|
||||
const replayCursorOps = meter.createCounter("realtime_native.replay_cursor_ops", {
|
||||
description:
|
||||
"Shared replay-cursor store operations by outcome. Errors degrade hops to cold resolves (watch live_polls{path='cold-resolve'} rise with them), never failed polls.",
|
||||
});
|
||||
|
||||
const staleHydrates = meter.createCounter("realtime_native.stale_hydrates", {
|
||||
description:
|
||||
"Wake hydrates the read-your-writes tripwire caught reading behind the publish. 'recovered' = a retry delivered the fresh row; sustained 'gave_up' means replica lag is outrunning the retry budget.",
|
||||
});
|
||||
|
||||
const limiter = new RealtimeConcurrencyLimiter({
|
||||
keyPrefix: "tr:realtime:native:concurrency",
|
||||
redis: {
|
||||
port: env.RATE_LIMIT_REDIS_PORT,
|
||||
host: env.RATE_LIMIT_REDIS_HOST,
|
||||
username: env.RATE_LIMIT_REDIS_USERNAME,
|
||||
password: env.RATE_LIMIT_REDIS_PASSWORD,
|
||||
tlsDisabled: env.RATE_LIMIT_REDIS_TLS_DISABLED === "true",
|
||||
clusterMode: env.RATE_LIMIT_REDIS_CLUSTER_MODE_ENABLED === "1",
|
||||
},
|
||||
});
|
||||
|
||||
// Fleet-shared replay cursors (one timestamp per connection) on the same Redis as the
|
||||
// change channel, so a load-balancer hop reads the connection's true inter-poll gap.
|
||||
const replayCursorStore =
|
||||
env.REALTIME_BACKEND_NATIVE_SHARED_REPLAY_CURSORS === "1"
|
||||
? new RedisReplayCursorStore({
|
||||
redis: {
|
||||
host: env.REALTIME_BACKEND_NATIVE_PUBSUB_REDIS_HOST,
|
||||
port: env.REALTIME_BACKEND_NATIVE_PUBSUB_REDIS_PORT,
|
||||
username: env.REALTIME_BACKEND_NATIVE_PUBSUB_REDIS_USERNAME,
|
||||
password: env.REALTIME_BACKEND_NATIVE_PUBSUB_REDIS_PASSWORD,
|
||||
tlsDisabled: env.REALTIME_BACKEND_NATIVE_PUBSUB_REDIS_TLS_DISABLED === "true",
|
||||
clusterMode: env.REALTIME_BACKEND_NATIVE_PUBSUB_REDIS_CLUSTER_MODE_ENABLED === "1",
|
||||
},
|
||||
ttlMs: env.REALTIME_BACKEND_NATIVE_WORKING_SET_TTL_MS,
|
||||
onResult: (op, ok) => replayCursorOps.add(1, { op, result: ok ? "ok" : "error" }),
|
||||
})
|
||||
: undefined;
|
||||
|
||||
// One RunHydrator shared by the router and the client, so its single-flight + short-TTL cache covers both.
|
||||
const runReader = new RunHydrator({
|
||||
replica: $replica,
|
||||
runStore,
|
||||
cacheTtlMs: env.REALTIME_BACKEND_NATIVE_RUN_CACHE_TTL_MS,
|
||||
maxCacheEntries: env.REALTIME_BACKEND_NATIVE_RUN_CACHE_MAX_ENTRIES,
|
||||
});
|
||||
|
||||
// Read-your-writes gate: the estimator samples replica lag (reader-side only, paused
|
||||
// when idle) and the router delays wake hydrates by it, anchored to each record's
|
||||
// updatedAtMs — so a publish racing the replica's apply is waited out, not read stale.
|
||||
const lagEstimator =
|
||||
env.REALTIME_BACKEND_NATIVE_REPLICA_LAG_GATE_ENABLED === "1"
|
||||
? new ReplicaLagEstimator({
|
||||
source: createPostgresReplicaLagSource($replica),
|
||||
sampleIntervalMs: env.REALTIME_BACKEND_NATIVE_REPLICA_LAG_SAMPLE_INTERVAL_MS,
|
||||
idleAfterMs: env.REALTIME_BACKEND_NATIVE_REPLICA_LAG_IDLE_AFTER_MS,
|
||||
windowMs: env.REALTIME_BACKEND_NATIVE_REPLICA_LAG_WINDOW_MS,
|
||||
defaultLagMs: env.REALTIME_BACKEND_NATIVE_REPLICA_LAG_DEFAULT_MS,
|
||||
observedFloorTtlMs: env.REALTIME_BACKEND_NATIVE_REPLICA_LAG_OBSERVED_FLOOR_TTL_MS,
|
||||
})
|
||||
: undefined;
|
||||
|
||||
// The notifier wrapped so router activity keeps the lag sampler warm.
|
||||
const notifier = getRunChangeNotifier();
|
||||
const source: EnvChangeSource = lagEstimator
|
||||
? {
|
||||
subscribeToEnv(environmentId, onBatch) {
|
||||
lagEstimator.touch();
|
||||
return notifier.subscribeToEnv(environmentId, (records) => {
|
||||
lagEstimator.touch();
|
||||
onBatch(records);
|
||||
});
|
||||
},
|
||||
}
|
||||
: notifier;
|
||||
|
||||
const router = new EnvChangeRouter({
|
||||
source,
|
||||
hydrator: runReader,
|
||||
onHydrate: (runCount) => routerHydrates.add(runCount),
|
||||
replayWindowMs: env.REALTIME_BACKEND_NATIVE_REPLAY_WINDOW_MS,
|
||||
replayMaxRunsPerEnv: env.REALTIME_BACKEND_NATIVE_REPLAY_MAX_RUNS,
|
||||
unsubscribeLingerMs: env.REALTIME_BACKEND_NATIVE_UNSUBSCRIBE_LINGER_MS,
|
||||
onReplay: (result) => replays.add(1, { result }),
|
||||
onReplayEviction: (reason) => replayEvictions.add(1, { reason }),
|
||||
replicaLag: lagEstimator
|
||||
? {
|
||||
getLagMs: () => lagEstimator.getLagMs(),
|
||||
noteObservedLagMs: (lagMs) => lagEstimator.noteObservedLagMs(lagMs),
|
||||
marginMs: env.REALTIME_BACKEND_NATIVE_REPLICA_LAG_MARGIN_MS,
|
||||
maxDelayMs: env.REALTIME_BACKEND_NATIVE_REPLICA_LAG_MAX_DELAY_MS,
|
||||
staleRetries: env.REALTIME_BACKEND_NATIVE_STALE_HYDRATE_RETRIES,
|
||||
onStaleHydrate: (outcome, runCount) => staleHydrates.add(runCount, { outcome }),
|
||||
}
|
||||
: undefined,
|
||||
});
|
||||
|
||||
const client = new NativeRealtimeClient({
|
||||
runReader,
|
||||
runListResolver: new ClickHouseRunListResolver({
|
||||
getClickhouse: (organizationId) =>
|
||||
clickhouseFactory.getClickhouseForOrganization(organizationId, "realtime"),
|
||||
prisma: $replica,
|
||||
}),
|
||||
router,
|
||||
limiter,
|
||||
cachedLimitProvider: {
|
||||
async getCachedLimit(organizationId, defaultValue) {
|
||||
const result = await getCachedLimit(
|
||||
organizationId,
|
||||
"realtimeConcurrentConnections",
|
||||
defaultValue
|
||||
);
|
||||
return result.val;
|
||||
},
|
||||
},
|
||||
defaultConcurrencyLimit: env.REALTIME_BACKEND_NATIVE_DEFAULT_CONCURRENCY_LIMIT,
|
||||
livePollTimeoutMs: env.REALTIME_BACKEND_NATIVE_LIVE_POLL_TIMEOUT_MS,
|
||||
livePollJitterRatio: env.REALTIME_BACKEND_NATIVE_LIVE_POLL_JITTER_RATIO,
|
||||
maximumCreatedAtFilterAgeMs: env.REALTIME_MAXIMUM_CREATED_AT_FILTER_AGE_IN_MS,
|
||||
maxListResults: env.REALTIME_BACKEND_NATIVE_MAX_LIST_RESULTS,
|
||||
runSetResolveCacheTtlMs: env.REALTIME_BACKEND_NATIVE_RUNSET_CACHE_TTL_MS,
|
||||
runSetResolveCacheMaxEntries: env.REALTIME_BACKEND_NATIVE_RUNSET_CACHE_MAX_ENTRIES,
|
||||
listCacheMaxEntries: env.REALTIME_BACKEND_NATIVE_WORKING_SET_MAX_ENTRIES,
|
||||
workingSetCacheTtlMs: env.REALTIME_BACKEND_NATIVE_WORKING_SET_TTL_MS,
|
||||
runSetCreatedAtBucketMs: env.REALTIME_BACKEND_NATIVE_RUNSET_CREATED_AT_BUCKET_MS,
|
||||
holdOnEmpty: env.REALTIME_BACKEND_NATIVE_HOLD_ON_EMPTY === "1",
|
||||
resolveAdmissionLimit: env.REALTIME_BACKEND_NATIVE_RESOLVE_ADMISSION_LIMIT,
|
||||
replayCursorStore,
|
||||
onWakeup: (reason) => wakeups.add(1, { reason }),
|
||||
onLivePollPath: (path) => livePollPaths.add(1, { path }),
|
||||
onRunSetResolve: (result) => runSetResolves.add(1, { result }),
|
||||
onRunSetQuery: (stage, ms) => runSetQueryMs.record(ms, { stage }),
|
||||
onResolveAdmissionWait: () => resolveAdmissionWaits.add(1),
|
||||
onEmit: (path, lagMs, rowCount) => {
|
||||
deliveryLagMs.record(Math.max(lagMs, 0), { path });
|
||||
emittedRows.record(rowCount);
|
||||
},
|
||||
onBackstopResult: (result) => backstops.add(1, { result }),
|
||||
onConcurrencyRejected: () => concurrencyRejections.add(1),
|
||||
});
|
||||
|
||||
meter
|
||||
.createObservableGauge("realtime_native.working_set_size", {
|
||||
description:
|
||||
"Entries in the per-handle working-set cache (one per active multi-run feed session).",
|
||||
})
|
||||
.addCallback((result) => result.observe(client.workingSetCacheSize));
|
||||
|
||||
meter
|
||||
.createObservableGauge("realtime_native.resolve_admission_in_use", {
|
||||
description:
|
||||
"Fresh ClickHouse resolves currently holding an admission permit (live concurrency against the gate's limit).",
|
||||
})
|
||||
.addCallback((result) => result.observe(client.resolveAdmissionInUse));
|
||||
|
||||
meter
|
||||
.createObservableGauge("realtime_native.held_feeds", {
|
||||
description: "Long-polls currently held, by feed kind — the system's capacity unit.",
|
||||
})
|
||||
.addCallback((result) => {
|
||||
const counts = router.heldFeedCounts;
|
||||
result.observe(counts.run, { kind: "run" });
|
||||
result.observe(counts.tag, { kind: "tag" });
|
||||
result.observe(counts.batch, { kind: "batch" });
|
||||
});
|
||||
|
||||
meter
|
||||
.createObservableGauge("realtime_native.active_envs", {
|
||||
description:
|
||||
"Environments currently routed on this instance (held feeds + lingering subscriptions).",
|
||||
})
|
||||
.addCallback((result) => result.observe(router.activeEnvCount));
|
||||
|
||||
if (lagEstimator) {
|
||||
meter
|
||||
.createObservableGauge("realtime_native.replica_lag_estimate_ms", {
|
||||
description:
|
||||
"The read-your-writes gate's current replica-lag estimate (max sample in the window). Wake hydrates are delayed by roughly this much past each change's commit.",
|
||||
})
|
||||
.addCallback((result) => result.observe(lagEstimator.getLagMs()));
|
||||
}
|
||||
|
||||
return client;
|
||||
}
|
||||
|
||||
export function getNativeRealtimeClient(): NativeRealtimeClient {
|
||||
return singleton("nativeRealtimeClient", initializeNativeRealtimeClient);
|
||||
}
|
||||
@@ -0,0 +1,112 @@
|
||||
import type { Callback, Result } from "ioredis";
|
||||
import type { RedisClient, RedisWithClusterOptions } from "~/redis.server";
|
||||
import { createRedisClient } from "~/redis.server";
|
||||
import { logger } from "../logger.server";
|
||||
|
||||
export type RealtimeConcurrencyLimiterOptions = {
|
||||
redis: RedisWithClusterOptions;
|
||||
keyPrefix: string;
|
||||
/** How long a tracked request lives before it's swept as stale (seconds). */
|
||||
expiryTimeInSeconds?: number;
|
||||
connectionName?: string;
|
||||
};
|
||||
|
||||
/**
|
||||
* Per-environment concurrent-connection limiter for realtime long-polls; a standalone copy of the limiter in
|
||||
* `realtimeClient.server.ts` (identical Lua + key shape, different key prefix) so the native backend tracks independently.
|
||||
*/
|
||||
export class RealtimeConcurrencyLimiter {
|
||||
private redis: RedisClient;
|
||||
private expiryTimeInSeconds: number;
|
||||
|
||||
constructor(private options: RealtimeConcurrencyLimiterOptions) {
|
||||
this.redis = createRedisClient(
|
||||
options.connectionName ?? "trigger:realtime:native:concurrency",
|
||||
options.redis
|
||||
);
|
||||
this.expiryTimeInSeconds = options.expiryTimeInSeconds ?? 60 * 5;
|
||||
this.#registerCommands();
|
||||
}
|
||||
|
||||
async incrementAndCheck(
|
||||
environmentId: string,
|
||||
requestId: string,
|
||||
limit: number
|
||||
): Promise<boolean> {
|
||||
const key = this.#getKey(environmentId);
|
||||
const now = Date.now();
|
||||
|
||||
const result = await this.redis.incrementAndCheckRealtimeNativeConcurrency(
|
||||
key,
|
||||
now.toString(),
|
||||
requestId,
|
||||
this.expiryTimeInSeconds.toString(),
|
||||
(now - this.expiryTimeInSeconds * 1000).toString(),
|
||||
limit.toString()
|
||||
);
|
||||
|
||||
return result === 1;
|
||||
}
|
||||
|
||||
async decrement(environmentId: string, requestId: string): Promise<void> {
|
||||
const key = this.#getKey(environmentId);
|
||||
await this.redis.zrem(key, requestId);
|
||||
}
|
||||
|
||||
#getKey(environmentId: string): string {
|
||||
return `${this.options.keyPrefix}:${environmentId}`;
|
||||
}
|
||||
|
||||
#registerCommands() {
|
||||
this.redis.defineCommand("incrementAndCheckRealtimeNativeConcurrency", {
|
||||
numberOfKeys: 1,
|
||||
lua: /* lua */ `
|
||||
local concurrencyKey = KEYS[1]
|
||||
|
||||
local timestamp = tonumber(ARGV[1])
|
||||
local requestId = ARGV[2]
|
||||
local expiryTime = tonumber(ARGV[3])
|
||||
local cutoffTime = tonumber(ARGV[4])
|
||||
local limit = tonumber(ARGV[5])
|
||||
|
||||
-- Remove expired entries
|
||||
redis.call('ZREMRANGEBYSCORE', concurrencyKey, '-inf', cutoffTime)
|
||||
|
||||
-- Add the new request to the sorted set
|
||||
redis.call('ZADD', concurrencyKey, timestamp, requestId)
|
||||
|
||||
-- Set the expiry time on the key
|
||||
redis.call('EXPIRE', concurrencyKey, expiryTime)
|
||||
|
||||
-- Get the total number of concurrent requests
|
||||
local totalRequests = redis.call('ZCARD', concurrencyKey)
|
||||
|
||||
-- Check if the limit has been exceeded
|
||||
if totalRequests > limit then
|
||||
redis.call('ZREM', concurrencyKey, requestId)
|
||||
return 0
|
||||
end
|
||||
|
||||
return 1
|
||||
`,
|
||||
});
|
||||
|
||||
this.redis.on("error", (error) => {
|
||||
logger.error("[realtimeConcurrencyLimiter] redis error", { error });
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
declare module "ioredis" {
|
||||
interface RedisCommander<Context> {
|
||||
incrementAndCheckRealtimeNativeConcurrency(
|
||||
key: string,
|
||||
timestamp: string,
|
||||
requestId: string,
|
||||
expiryTime: string,
|
||||
cutoffTime: string,
|
||||
limit: string,
|
||||
callback?: Callback<number>
|
||||
): Result<number, Context>;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,479 @@
|
||||
import type { LogLevel } from "@trigger.dev/core/logger";
|
||||
import { Logger } from "@trigger.dev/core/logger";
|
||||
import type { RedisOptions } from "ioredis";
|
||||
import Redis from "ioredis";
|
||||
import { defaultReconnectOnError } from "@internal/redis";
|
||||
import { env } from "~/env.server";
|
||||
import type { StreamIngestor, StreamResponder, StreamResponseOptions } from "./types";
|
||||
|
||||
export type RealtimeStreamsOptions = {
|
||||
redis: RedisOptions | undefined;
|
||||
logger?: Logger;
|
||||
logLevel?: LogLevel;
|
||||
inactivityTimeoutMs?: number; // Close stream after this many ms of no new data (default: 15000)
|
||||
};
|
||||
|
||||
// Legacy constant for backward compatibility (no longer written, but still recognized when reading)
|
||||
const END_SENTINEL = "<<CLOSE_STREAM>>";
|
||||
|
||||
// Internal types for stream pipeline
|
||||
type StreamChunk =
|
||||
| { type: "ping" }
|
||||
| { type: "data"; redisId: string; data: string }
|
||||
| { type: "legacy-data"; redisId: string; data: string };
|
||||
|
||||
// Class implementing both interfaces
|
||||
export class RedisRealtimeStreams implements StreamIngestor, StreamResponder {
|
||||
private logger: Logger;
|
||||
private inactivityTimeoutMs: number;
|
||||
// Shared connection for short-lived non-blocking operations (XADD, XREVRANGE, EXPIRE).
|
||||
// Lazily created on first use so we don't open a connection if only streamResponse is called.
|
||||
private _sharedRedis: Redis | undefined;
|
||||
|
||||
constructor(private options: RealtimeStreamsOptions) {
|
||||
this.logger = options.logger ?? new Logger("RedisRealtimeStreams", options.logLevel ?? "info");
|
||||
this.inactivityTimeoutMs = options.inactivityTimeoutMs ?? 15000; // Default: 15 seconds
|
||||
}
|
||||
|
||||
private get sharedRedis(): Redis {
|
||||
if (!this._sharedRedis) {
|
||||
this._sharedRedis = new Redis({
|
||||
reconnectOnError: defaultReconnectOnError,
|
||||
...this.options.redis,
|
||||
connectionName: "realtime:shared",
|
||||
});
|
||||
}
|
||||
return this._sharedRedis;
|
||||
}
|
||||
|
||||
async initializeStream(
|
||||
runId: string,
|
||||
streamId: string
|
||||
): Promise<{ responseHeaders?: Record<string, string> }> {
|
||||
return {};
|
||||
}
|
||||
|
||||
async streamResponse(
|
||||
request: Request,
|
||||
runId: string,
|
||||
streamId: string,
|
||||
signal: AbortSignal,
|
||||
options?: StreamResponseOptions
|
||||
): Promise<Response> {
|
||||
const redis = new Redis({
|
||||
reconnectOnError: defaultReconnectOnError,
|
||||
...this.options.redis,
|
||||
connectionName: "realtime:streamResponse",
|
||||
});
|
||||
const streamKey = `stream:${runId}:${streamId}`;
|
||||
let isCleanedUp = false;
|
||||
|
||||
const stream = new ReadableStream<StreamChunk>({
|
||||
start: async (controller) => {
|
||||
// Start from lastEventId if provided, otherwise from beginning
|
||||
let lastId = options?.lastEventId ?? "0";
|
||||
let retryCount = 0;
|
||||
const maxRetries = 3;
|
||||
let lastDataTime = Date.now();
|
||||
let lastEnqueueTime = Date.now();
|
||||
const blockTimeMs = 5000;
|
||||
const pingIntervalMs = 10000; // 10 seconds
|
||||
|
||||
if (options?.lastEventId) {
|
||||
this.logger.debug("[RealtimeStreams][streamResponse] Resuming from lastEventId", {
|
||||
streamKey,
|
||||
lastEventId: options?.lastEventId,
|
||||
});
|
||||
}
|
||||
|
||||
try {
|
||||
while (!signal.aborted) {
|
||||
// Check if we need to send a ping
|
||||
const timeSinceLastEnqueue = Date.now() - lastEnqueueTime;
|
||||
if (timeSinceLastEnqueue >= pingIntervalMs) {
|
||||
controller.enqueue({ type: "ping" });
|
||||
lastEnqueueTime = Date.now();
|
||||
}
|
||||
|
||||
// Compute inactivity threshold once to use consistently in both branches
|
||||
const inactivityThresholdMs = options?.timeoutInSeconds
|
||||
? options.timeoutInSeconds * 1000
|
||||
: this.inactivityTimeoutMs;
|
||||
|
||||
try {
|
||||
const messages = await redis.xread(
|
||||
"COUNT",
|
||||
100,
|
||||
"BLOCK",
|
||||
blockTimeMs,
|
||||
"STREAMS",
|
||||
streamKey,
|
||||
lastId
|
||||
);
|
||||
|
||||
retryCount = 0;
|
||||
|
||||
if (messages && messages.length > 0) {
|
||||
const [_key, entries] = messages[0];
|
||||
let foundData = false;
|
||||
|
||||
for (let i = 0; i < entries.length; i++) {
|
||||
const [id, fields] = entries[i];
|
||||
lastId = id;
|
||||
|
||||
if (fields && fields.length >= 2) {
|
||||
// Extract the data field from the Redis entry
|
||||
// Fields format: ["field1", "value1", "field2", "value2", ...]
|
||||
let data: string | null = null;
|
||||
|
||||
for (let j = 0; j < fields.length; j += 2) {
|
||||
if (fields[j] === "data") {
|
||||
data = fields[j + 1];
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Handle legacy entries that don't have field names (just data at index 1)
|
||||
if (data === null && fields.length >= 2) {
|
||||
data = fields[1];
|
||||
}
|
||||
|
||||
if (data) {
|
||||
// Skip legacy END_SENTINEL entries (backward compatibility)
|
||||
if (data === END_SENTINEL) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Enqueue structured chunk with Redis stream ID
|
||||
controller.enqueue({
|
||||
type: "data",
|
||||
redisId: id,
|
||||
data,
|
||||
});
|
||||
|
||||
foundData = true;
|
||||
lastDataTime = Date.now();
|
||||
lastEnqueueTime = Date.now();
|
||||
|
||||
if (signal.aborted) {
|
||||
controller.close();
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// If we didn't find any data in this batch, might have only seen sentinels
|
||||
if (!foundData) {
|
||||
// Check for inactivity timeout
|
||||
const inactiveMs = Date.now() - lastDataTime;
|
||||
if (inactiveMs >= inactivityThresholdMs) {
|
||||
this.logger.debug(
|
||||
"[RealtimeStreams][streamResponse] Closing stream due to inactivity",
|
||||
{
|
||||
streamKey,
|
||||
inactiveMs,
|
||||
threshold: inactivityThresholdMs,
|
||||
}
|
||||
);
|
||||
controller.close();
|
||||
return;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// No messages received (timed out on BLOCK)
|
||||
// Check for inactivity timeout
|
||||
const inactiveMs = Date.now() - lastDataTime;
|
||||
if (inactiveMs >= inactivityThresholdMs) {
|
||||
this.logger.debug(
|
||||
"[RealtimeStreams][streamResponse] Closing stream due to inactivity",
|
||||
{
|
||||
streamKey,
|
||||
inactiveMs,
|
||||
threshold: inactivityThresholdMs,
|
||||
}
|
||||
);
|
||||
controller.close();
|
||||
return;
|
||||
}
|
||||
}
|
||||
} catch (error) {
|
||||
if (signal.aborted) break;
|
||||
|
||||
this.logger.error(
|
||||
"[RealtimeStreams][streamResponse] Error reading from Redis stream:",
|
||||
{
|
||||
error,
|
||||
}
|
||||
);
|
||||
retryCount++;
|
||||
if (retryCount >= maxRetries) throw error;
|
||||
await new Promise((resolve) => setTimeout(resolve, 1000 * retryCount));
|
||||
}
|
||||
}
|
||||
} catch (error) {
|
||||
this.logger.error("[RealtimeStreams][streamResponse] Fatal error in stream processing:", {
|
||||
error,
|
||||
});
|
||||
controller.error(error);
|
||||
} finally {
|
||||
await cleanup();
|
||||
}
|
||||
},
|
||||
cancel: async () => {
|
||||
await cleanup();
|
||||
},
|
||||
})
|
||||
.pipeThrough(
|
||||
// Transform 1: Buffer partial lines across Redis entries
|
||||
(() => {
|
||||
let buffer = "";
|
||||
let lastRedisId = "0";
|
||||
|
||||
return new TransformStream<StreamChunk, StreamChunk & { line: string }>({
|
||||
transform(chunk, controller) {
|
||||
if (chunk.type === "ping") {
|
||||
controller.enqueue(chunk as any);
|
||||
} else if (chunk.type === "data" || chunk.type === "legacy-data") {
|
||||
// Buffer partial lines: accumulate until we see newlines
|
||||
buffer += chunk.data;
|
||||
|
||||
// Split on newlines
|
||||
const lines = buffer.split("\n");
|
||||
|
||||
// The last element might be incomplete, hold it back in buffer
|
||||
buffer = lines.pop() || "";
|
||||
|
||||
// Emit complete lines with the Redis ID of the chunk that completed them
|
||||
for (const line of lines) {
|
||||
if (line.trim().length > 0) {
|
||||
controller.enqueue({
|
||||
...chunk,
|
||||
line,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Update last Redis ID for next iteration
|
||||
lastRedisId = chunk.redisId;
|
||||
}
|
||||
},
|
||||
flush(controller) {
|
||||
// On stream end, emit any leftover buffered text
|
||||
if (buffer.trim().length > 0) {
|
||||
controller.enqueue({
|
||||
type: "data",
|
||||
redisId: lastRedisId,
|
||||
data: "",
|
||||
line: buffer.trim(),
|
||||
});
|
||||
}
|
||||
},
|
||||
});
|
||||
})()
|
||||
)
|
||||
.pipeThrough(
|
||||
// Transform 2: Format as SSE
|
||||
new TransformStream<StreamChunk & { line?: string }, string>({
|
||||
transform(chunk, controller) {
|
||||
if (chunk.type === "ping") {
|
||||
controller.enqueue(`: ping\n\n`);
|
||||
} else if ((chunk.type === "data" || chunk.type === "legacy-data") && chunk.line) {
|
||||
// Use Redis stream ID as SSE event ID
|
||||
controller.enqueue(`id: ${chunk.redisId}\ndata: ${chunk.line}\n\n`);
|
||||
}
|
||||
},
|
||||
})
|
||||
)
|
||||
.pipeThrough(new TextEncoderStream());
|
||||
|
||||
async function cleanup() {
|
||||
if (isCleanedUp) return;
|
||||
isCleanedUp = true;
|
||||
// disconnect() tears down the TCP socket immediately, which causes any
|
||||
// pending XREAD BLOCK to reject right away instead of waiting for the
|
||||
// block timeout to elapse. quit() would queue behind the blocking command.
|
||||
redis.disconnect();
|
||||
}
|
||||
|
||||
signal.addEventListener("abort", cleanup, { once: true });
|
||||
|
||||
return new Response(stream, {
|
||||
headers: {
|
||||
"Content-Type": "text/event-stream",
|
||||
"Cache-Control": "no-cache",
|
||||
Connection: "keep-alive",
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
async ingestData(
|
||||
stream: ReadableStream<Uint8Array>,
|
||||
runId: string,
|
||||
streamId: string,
|
||||
clientId: string,
|
||||
resumeFromChunk?: number
|
||||
): Promise<Response> {
|
||||
const redis = this.sharedRedis;
|
||||
const streamKey = `stream:${runId}:${streamId}`;
|
||||
const startChunk = resumeFromChunk ?? 0;
|
||||
// Start counting from the resume point, not from 0
|
||||
let currentChunkIndex = startChunk;
|
||||
|
||||
try {
|
||||
const textStream = stream.pipeThrough(new TextDecoderStream());
|
||||
const reader = textStream.getReader();
|
||||
|
||||
while (true) {
|
||||
const { done, value } = await reader.read();
|
||||
|
||||
if (done || !value) {
|
||||
break;
|
||||
}
|
||||
|
||||
// Write each chunk with its index and clientId
|
||||
this.logger.debug("[RedisRealtimeStreams][ingestData] Writing chunk", {
|
||||
streamKey,
|
||||
runId,
|
||||
clientId,
|
||||
chunkIndex: currentChunkIndex,
|
||||
resumeFromChunk: startChunk,
|
||||
value,
|
||||
});
|
||||
|
||||
await redis.xadd(
|
||||
streamKey,
|
||||
"MAXLEN",
|
||||
"~",
|
||||
String(env.REALTIME_STREAM_MAX_LENGTH),
|
||||
"*",
|
||||
"clientId",
|
||||
clientId,
|
||||
"chunkIndex",
|
||||
currentChunkIndex.toString(),
|
||||
"data",
|
||||
value
|
||||
);
|
||||
|
||||
currentChunkIndex++;
|
||||
}
|
||||
|
||||
// Set TTL for cleanup when stream is done
|
||||
await redis.expire(streamKey, env.REALTIME_STREAM_TTL);
|
||||
|
||||
return new Response(null, { status: 200 });
|
||||
} catch (error) {
|
||||
if (error instanceof Error) {
|
||||
if ("code" in error && error.code === "ECONNRESET") {
|
||||
this.logger.info("[RealtimeStreams][ingestData] Connection reset during ingestData:", {
|
||||
error,
|
||||
});
|
||||
return new Response(null, { status: 500 });
|
||||
}
|
||||
}
|
||||
|
||||
this.logger.error("[RealtimeStreams][ingestData] Error in ingestData:", { error });
|
||||
|
||||
return new Response(null, { status: 500 });
|
||||
}
|
||||
}
|
||||
|
||||
async appendPart(part: string, partId: string, runId: string, streamId: string): Promise<void> {
|
||||
const redis = this.sharedRedis;
|
||||
const streamKey = `stream:${runId}:${streamId}`;
|
||||
|
||||
await redis.xadd(
|
||||
streamKey,
|
||||
"MAXLEN",
|
||||
"~",
|
||||
String(env.REALTIME_STREAM_MAX_LENGTH),
|
||||
"*",
|
||||
"clientId",
|
||||
"",
|
||||
"chunkIndex",
|
||||
"0",
|
||||
"data",
|
||||
JSON.stringify(part) + "\n"
|
||||
);
|
||||
|
||||
// Set TTL for cleanup when stream is done
|
||||
await redis.expire(streamKey, env.REALTIME_STREAM_TTL);
|
||||
}
|
||||
|
||||
async getLastChunkIndex(runId: string, streamId: string, clientId: string): Promise<number> {
|
||||
const redis = this.sharedRedis;
|
||||
const streamKey = `stream:${runId}:${streamId}`;
|
||||
|
||||
try {
|
||||
// Paginate through the stream from newest to oldest until we find this client's last chunk
|
||||
const batchSize = 100;
|
||||
let lastId = "+"; // Start from newest
|
||||
|
||||
while (true) {
|
||||
const entries = await redis.xrevrange(streamKey, lastId, "-", "COUNT", batchSize);
|
||||
|
||||
if (!entries || entries.length === 0) {
|
||||
// Reached the beginning of the stream, no chunks from this client
|
||||
this.logger.debug(
|
||||
"[RedisRealtimeStreams][getLastChunkIndex] No chunks found for client",
|
||||
{
|
||||
streamKey,
|
||||
clientId,
|
||||
}
|
||||
);
|
||||
return -1;
|
||||
}
|
||||
|
||||
// Search through this batch for the client's last chunk
|
||||
for (const [_id, fields] of entries) {
|
||||
let entryClientId: string | null = null;
|
||||
let chunkIndex: number | null = null;
|
||||
let data: string | null = null;
|
||||
|
||||
for (let i = 0; i < fields.length; i += 2) {
|
||||
if (fields[i] === "clientId") {
|
||||
entryClientId = fields[i + 1];
|
||||
}
|
||||
if (fields[i] === "chunkIndex") {
|
||||
chunkIndex = parseInt(fields[i + 1], 10);
|
||||
}
|
||||
if (fields[i] === "data") {
|
||||
data = fields[i + 1];
|
||||
}
|
||||
}
|
||||
|
||||
// Skip legacy END_SENTINEL entries (backward compatibility)
|
||||
if (data === END_SENTINEL) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Check if this entry is from our client and has a chunkIndex
|
||||
if (entryClientId === clientId && chunkIndex !== null) {
|
||||
this.logger.debug("[RedisRealtimeStreams][getLastChunkIndex] Found last chunk", {
|
||||
streamKey,
|
||||
clientId,
|
||||
chunkIndex,
|
||||
});
|
||||
return chunkIndex;
|
||||
}
|
||||
}
|
||||
|
||||
// Move to next batch (older entries)
|
||||
// Use the ID of the last entry in this batch as the new cursor
|
||||
lastId = `(${entries[entries.length - 1][0]}`; // Exclusive range with (
|
||||
}
|
||||
} catch (error) {
|
||||
this.logger.error("[RedisRealtimeStreams][getLastChunkIndex] Error getting last chunk:", {
|
||||
error,
|
||||
streamKey,
|
||||
clientId,
|
||||
});
|
||||
// Return -1 to indicate we don't know what the server has
|
||||
return -1;
|
||||
}
|
||||
}
|
||||
|
||||
async readRecords(): Promise<never> {
|
||||
throw new Error("readRecords is not implemented for Redis realtime streams");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,145 @@
|
||||
import { createRedisClient, type RedisClient, type RedisWithClusterOptions } from "~/redis.server";
|
||||
import { logger } from "../logger.server";
|
||||
import { BoundedTtlCache } from "./boundedTtlCache";
|
||||
|
||||
/**
|
||||
* Per-connection replay cursors ("when did this connection last receive data"), keyed by the
|
||||
* env-prefixed working-set key. Sharing them fleet-wide makes an instance hop look like a normal
|
||||
* inter-poll gap instead of an unknown one, so hops stop triggering cold resolves and full-window
|
||||
* replays. Values are single timestamps, so the shared store stays cheap.
|
||||
*/
|
||||
export interface ReplayCursorStore {
|
||||
/** The connection's last-response timestamp; undefined on miss OR error (the caller
|
||||
* degrades to a cold probe / full-window replay, never blocks the poll). */
|
||||
get(key: string): Promise<number | undefined>;
|
||||
/** Fire-and-forget stamp; must never throw. */
|
||||
set(key: string, ms: number): void;
|
||||
}
|
||||
|
||||
/** Per-instance fallback with the same shape (used when the shared store is disabled, and in tests). */
|
||||
export class InMemoryReplayCursorStore implements ReplayCursorStore {
|
||||
readonly #cache: BoundedTtlCache<number>;
|
||||
|
||||
constructor(ttlMs: number, maxEntries: number) {
|
||||
this.#cache = new BoundedTtlCache<number>(ttlMs, maxEntries);
|
||||
}
|
||||
|
||||
async get(key: string): Promise<number | undefined> {
|
||||
return this.#cache.get(key);
|
||||
}
|
||||
|
||||
set(key: string, ms: number): void {
|
||||
this.#cache.set(key, ms);
|
||||
}
|
||||
}
|
||||
|
||||
export type RedisReplayCursorStoreOptions = {
|
||||
redis: RedisWithClusterOptions;
|
||||
/** Entry TTL (ms); matches the working-set TTL so both views of a connection age out together. */
|
||||
ttlMs: number;
|
||||
/** Read deadline (ms): a slow or down Redis degrades the poll to a cold probe instead of stalling it. */
|
||||
getTimeoutMs?: number;
|
||||
keyPrefix?: string;
|
||||
connectionName?: string;
|
||||
/** Observability hook: a store op settled (errors are the degradation signal, not failures). */
|
||||
onResult?: (op: "get" | "set", ok: boolean) => void;
|
||||
};
|
||||
|
||||
const DEFAULT_KEY_PREFIX = "realtime:replay-cursor:";
|
||||
const DEFAULT_GET_TIMEOUT_MS = 250;
|
||||
const TIMED_OUT = Symbol("replay-cursor-get-timeout");
|
||||
|
||||
export class RedisReplayCursorStore implements ReplayCursorStore {
|
||||
#client: RedisClient | undefined;
|
||||
|
||||
constructor(private readonly options: RedisReplayCursorStoreOptions) {}
|
||||
|
||||
async get(key: string): Promise<number | undefined> {
|
||||
try {
|
||||
const raw = await this.#getWithDeadline(this.#key(key));
|
||||
if (raw === TIMED_OUT) {
|
||||
this.options.onResult?.("get", false);
|
||||
logger.warn("[replayCursorStore] replay-cursor read timed out", { key });
|
||||
return undefined;
|
||||
}
|
||||
this.options.onResult?.("get", true);
|
||||
if (raw === null) {
|
||||
return undefined;
|
||||
}
|
||||
const ms = Number(raw);
|
||||
return Number.isFinite(ms) && ms > 0 ? ms : undefined;
|
||||
} catch (error) {
|
||||
this.options.onResult?.("get", false);
|
||||
logger.error("[replayCursorStore] failed to read a replay cursor", { error, key });
|
||||
return undefined;
|
||||
}
|
||||
}
|
||||
|
||||
/** GET raced against the read deadline (ioredis queues commands while disconnected, which
|
||||
* would otherwise stall every poll start through an outage). */
|
||||
#getWithDeadline(key: string): Promise<string | null | typeof TIMED_OUT> {
|
||||
return new Promise((resolve, reject) => {
|
||||
const timer = setTimeout(
|
||||
() => resolve(TIMED_OUT),
|
||||
this.options.getTimeoutMs ?? DEFAULT_GET_TIMEOUT_MS
|
||||
);
|
||||
timer.unref?.();
|
||||
this.#ensureClient()
|
||||
.get(key)
|
||||
.then(
|
||||
(value) => {
|
||||
clearTimeout(timer);
|
||||
resolve(value);
|
||||
},
|
||||
(error) => {
|
||||
clearTimeout(timer);
|
||||
reject(error);
|
||||
}
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
set(key: string, ms: number): void {
|
||||
try {
|
||||
this.#ensureClient()
|
||||
.set(this.#key(key), String(ms), "PX", this.options.ttlMs)
|
||||
.then(
|
||||
() => this.options.onResult?.("set", true),
|
||||
(error) => {
|
||||
this.options.onResult?.("set", false);
|
||||
logger.error("[replayCursorStore] failed to write a replay cursor", { error, key });
|
||||
}
|
||||
);
|
||||
} catch (error) {
|
||||
this.options.onResult?.("set", false);
|
||||
logger.error("[replayCursorStore] failed to write a replay cursor", { error, key });
|
||||
}
|
||||
}
|
||||
|
||||
async quit(): Promise<void> {
|
||||
const client = this.#client;
|
||||
this.#client = undefined;
|
||||
if (!client) return;
|
||||
try {
|
||||
// Bounded graceful QUIT; cursor writes are best-effort, so force-close beyond it.
|
||||
await Promise.race([client.quit(), new Promise((resolve) => setTimeout(resolve, 500))]);
|
||||
} catch {
|
||||
// force-close below
|
||||
}
|
||||
client.disconnect();
|
||||
}
|
||||
|
||||
#key(key: string): string {
|
||||
return `${this.options.keyPrefix ?? DEFAULT_KEY_PREFIX}${key}`;
|
||||
}
|
||||
|
||||
#ensureClient(): RedisClient {
|
||||
if (!this.#client) {
|
||||
this.#client = createRedisClient(
|
||||
this.options.connectionName ?? "trigger:realtime:replay-cursors",
|
||||
this.options.redis
|
||||
);
|
||||
}
|
||||
return this.#client;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,257 @@
|
||||
import { logger } from "~/services/logger.server";
|
||||
|
||||
/**
|
||||
* ReplicaLagEstimator — tracks how far the read replica trails the primary so the
|
||||
* EnvChangeRouter can delay wake-path hydrates just long enough to read their own writes.
|
||||
* Two inputs: a ReplicaLagSource (active, reader-side only — never queries the primary)
|
||||
* sampled on an interval while the router is busy, and passive observations fed back by
|
||||
* the router's stale-hydrate tripwire. The estimate is the max over a short window —
|
||||
* floored by recent observations — so spikes widen the delay immediately and decay back
|
||||
* out as fresh samples land.
|
||||
*/
|
||||
|
||||
type RawQueryable = {
|
||||
$queryRawUnsafe<T = unknown>(query: string): Promise<T>;
|
||||
};
|
||||
|
||||
/** A dialect-specific reader-side lag measure. `sampleLagMs()` returns the current lag in
|
||||
* ms, or undefined when lag is genuinely unmeasurable right now (NOT an error — errors
|
||||
* throw, and the composing source uses them to rule the dialect out). */
|
||||
export interface ReplicaLagSource {
|
||||
readonly name: string;
|
||||
sampleLagMs(): Promise<number | undefined>;
|
||||
}
|
||||
|
||||
/** Aurora: replicas share the storage layer and reject every standard WAL function;
|
||||
* `aurora_replica_status()` is the only live lag source. Max across readers, since the
|
||||
* `$replica` pool balances over all of them. No reader rows = `$replica` is the writer =
|
||||
* no lag. Throws on non-Aurora (the function doesn't exist). */
|
||||
export class AuroraReplicaLagSource implements ReplicaLagSource {
|
||||
readonly name = "aurora";
|
||||
|
||||
constructor(private readonly db: RawQueryable) {}
|
||||
|
||||
async sampleLagMs(): Promise<number | undefined> {
|
||||
const rows = await this.db.$queryRawUnsafe<{ lag: number | null }[]>(
|
||||
`SELECT max(replica_lag_in_msec)::float8 AS lag FROM aurora_replica_status() WHERE session_id <> 'MASTER_SESSION_ID' AND replica_lag_in_msec IS NOT NULL`
|
||||
);
|
||||
const lag = rows[0]?.lag;
|
||||
return typeof lag === "number" && Number.isFinite(lag) ? Math.max(0, lag) : 0;
|
||||
}
|
||||
}
|
||||
|
||||
/** Vanilla PG streaming replication. A primary (not in recovery — no replica configured,
|
||||
* `$replica` is the writer) has no lag by definition; a caught-up replica (receive LSN ==
|
||||
* replay LSN) reports 0. Mid-apply there is NO honest reader-side timestamp measure —
|
||||
* `now() - pg_last_xact_replay_timestamp()` reads as the full inter-write gap on
|
||||
* low-traffic systems, which (measured locally) pins the estimate at the delay cap — so
|
||||
* mid-apply reports undefined and the tripwire's observed-staleness floor carries the
|
||||
* estimate instead. */
|
||||
export class VanillaPgReplicaLagSource implements ReplicaLagSource {
|
||||
readonly name = "vanilla-pg";
|
||||
|
||||
constructor(private readonly db: RawQueryable) {}
|
||||
|
||||
async sampleLagMs(): Promise<number | undefined> {
|
||||
const rows = await this.db.$queryRawUnsafe<{ caught_up: boolean | null }[]>(
|
||||
`SELECT CASE
|
||||
WHEN NOT pg_is_in_recovery() THEN true
|
||||
WHEN pg_last_wal_receive_lsn() IS NOT DISTINCT FROM pg_last_wal_replay_lsn() THEN true
|
||||
ELSE false
|
||||
END AS caught_up`
|
||||
);
|
||||
return rows[0]?.caught_up ? 0 : undefined;
|
||||
}
|
||||
}
|
||||
|
||||
/** Composes dialect sources: the first whose sample succeeds is selected and used from
|
||||
* then on; a database where none work degrades to never-measuring (the estimator then
|
||||
* runs on its default + tripwire observations). Selection is by thrown-vs-returned —
|
||||
* sources throw on unsupported dialects and return undefined for "can't measure now". */
|
||||
export class FirstSupportedReplicaLagSource implements ReplicaLagSource {
|
||||
/** undefined = not probed yet; null = no candidate works here. */
|
||||
#selected: ReplicaLagSource | null | undefined;
|
||||
|
||||
constructor(private readonly candidates: ReplicaLagSource[]) {}
|
||||
|
||||
get name(): string {
|
||||
return this.#selected ? this.#selected.name : "undetected";
|
||||
}
|
||||
|
||||
async sampleLagMs(): Promise<number | undefined> {
|
||||
if (this.#selected === null) {
|
||||
return undefined;
|
||||
}
|
||||
if (this.#selected) {
|
||||
// Transient errors don't unselect the dialect; the sample is just skipped.
|
||||
try {
|
||||
return await this.#selected.sampleLagMs();
|
||||
} catch {
|
||||
return undefined;
|
||||
}
|
||||
}
|
||||
for (const candidate of this.candidates) {
|
||||
try {
|
||||
const lag = await candidate.sampleLagMs();
|
||||
this.#selected = candidate;
|
||||
logger.info("[replicaLagEstimator] selected lag source", { source: candidate.name });
|
||||
return lag;
|
||||
} catch {
|
||||
// unsupported dialect; try the next
|
||||
}
|
||||
}
|
||||
this.#selected = null;
|
||||
logger.warn(
|
||||
"[replicaLagEstimator] no usable lag source; relying on default + tripwire observations"
|
||||
);
|
||||
return undefined;
|
||||
}
|
||||
}
|
||||
|
||||
/** The standard composition for a Prisma replica client. */
|
||||
export function createPostgresReplicaLagSource(replica: RawQueryable): ReplicaLagSource {
|
||||
return new FirstSupportedReplicaLagSource([
|
||||
new AuroraReplicaLagSource(replica),
|
||||
new VanillaPgReplicaLagSource(replica),
|
||||
]);
|
||||
}
|
||||
|
||||
export type ReplicaLagEstimatorOptions = {
|
||||
source: ReplicaLagSource;
|
||||
/** Sample cadence while active. */
|
||||
sampleIntervalMs?: number;
|
||||
/** Stop sampling this long after the last touch(); the next touch resumes. */
|
||||
idleAfterMs?: number;
|
||||
/** The estimate is the max sample inside this window. */
|
||||
windowMs?: number;
|
||||
/** Estimate before any sample lands (and the floor when sampling is unavailable). */
|
||||
defaultLagMs?: number;
|
||||
/** Ceiling on accepted samples — shields the estimate from a wild observation. */
|
||||
maxLagMs?: number;
|
||||
/** How long a tripwire observation floors the estimate. Sources that can't measure
|
||||
* mid-apply lag (vanilla PG) return nothing, so without this floor the estimate decays
|
||||
* to the caught-up zeros within windowMs and every ~window one wake pays a stale retry
|
||||
* to re-learn. */
|
||||
observedFloorTtlMs?: number;
|
||||
/** Observability: a sample (active or passive) was accepted. */
|
||||
onSample?: (lagMs: number, source: "probe" | "observed") => void;
|
||||
};
|
||||
|
||||
const DEFAULT_SAMPLE_INTERVAL_MS = 250;
|
||||
const DEFAULT_IDLE_AFTER_MS = 30_000;
|
||||
const DEFAULT_WINDOW_MS = 5_000;
|
||||
const DEFAULT_DEFAULT_LAG_MS = 30;
|
||||
const DEFAULT_MAX_LAG_MS = 60_000;
|
||||
const DEFAULT_OBSERVED_FLOOR_TTL_MS = 60_000;
|
||||
|
||||
export class ReplicaLagEstimator {
|
||||
readonly #sampleIntervalMs: number;
|
||||
readonly #idleAfterMs: number;
|
||||
readonly #windowMs: number;
|
||||
readonly #defaultLagMs: number;
|
||||
readonly #maxLagMs: number;
|
||||
readonly #observedFloorTtlMs: number;
|
||||
#samples: { atMs: number; lagMs: number }[] = [];
|
||||
#lastKnownLagMs: number | undefined;
|
||||
#observedFloorLagMs = 0;
|
||||
#observedFloorAtMs = 0;
|
||||
#lastTouchMs = 0;
|
||||
#timer: ReturnType<typeof setInterval> | undefined;
|
||||
#sampling = false;
|
||||
|
||||
constructor(private readonly options: ReplicaLagEstimatorOptions) {
|
||||
this.#sampleIntervalMs = options.sampleIntervalMs ?? DEFAULT_SAMPLE_INTERVAL_MS;
|
||||
this.#idleAfterMs = options.idleAfterMs ?? DEFAULT_IDLE_AFTER_MS;
|
||||
this.#windowMs = options.windowMs ?? DEFAULT_WINDOW_MS;
|
||||
this.#defaultLagMs = options.defaultLagMs ?? DEFAULT_DEFAULT_LAG_MS;
|
||||
this.#maxLagMs = options.maxLagMs ?? DEFAULT_MAX_LAG_MS;
|
||||
this.#observedFloorTtlMs = options.observedFloorTtlMs ?? DEFAULT_OBSERVED_FLOOR_TTL_MS;
|
||||
}
|
||||
|
||||
/** Mark router activity; starts (or keeps) the sampler running. */
|
||||
touch() {
|
||||
this.#lastTouchMs = Date.now();
|
||||
if (!this.#timer) {
|
||||
this.#timer = setInterval(() => this.#tick(), this.#sampleIntervalMs);
|
||||
this.#timer.unref?.();
|
||||
// Sample immediately so the first wake after idle doesn't run on a stale estimate.
|
||||
this.#tick();
|
||||
}
|
||||
}
|
||||
|
||||
/** Current lag estimate (ms): the max recent sample (else last known, else the default),
|
||||
* floored by the latest tripwire observation while it's fresh. Never throws. */
|
||||
getLagMs(): number {
|
||||
const now = Date.now();
|
||||
const cutoff = now - this.#windowMs;
|
||||
let max: number | undefined;
|
||||
for (const sample of this.#samples) {
|
||||
if (sample.atMs >= cutoff && (max === undefined || sample.lagMs > max)) {
|
||||
max = sample.lagMs;
|
||||
}
|
||||
}
|
||||
const base = max ?? this.#lastKnownLagMs ?? this.#defaultLagMs;
|
||||
const floor =
|
||||
now - this.#observedFloorAtMs < this.#observedFloorTtlMs ? this.#observedFloorLagMs : 0;
|
||||
return Math.max(base, floor);
|
||||
}
|
||||
|
||||
/** Feedback from the stale-hydrate tripwire: a read provably ran at least this far
|
||||
* behind the primary. Widens the estimate immediately AND floors it for a while —
|
||||
* sources that can't measure mid-apply lag would otherwise decay it straight back. */
|
||||
noteObservedLagMs(lagMs: number) {
|
||||
const clamped = Math.min(Math.max(0, lagMs), this.#maxLagMs);
|
||||
const floorExpired = Date.now() - this.#observedFloorAtMs >= this.#observedFloorTtlMs;
|
||||
if (clamped >= this.#observedFloorLagMs || floorExpired) {
|
||||
this.#observedFloorLagMs = clamped;
|
||||
this.#observedFloorAtMs = Date.now();
|
||||
}
|
||||
this.#accept(lagMs, "observed");
|
||||
}
|
||||
|
||||
stop() {
|
||||
if (this.#timer) {
|
||||
clearInterval(this.#timer);
|
||||
this.#timer = undefined;
|
||||
}
|
||||
}
|
||||
|
||||
#accept(lagMs: number, source: "probe" | "observed") {
|
||||
if (!Number.isFinite(lagMs)) {
|
||||
return;
|
||||
}
|
||||
const clamped = Math.min(Math.max(0, lagMs), this.#maxLagMs);
|
||||
const now = Date.now();
|
||||
this.#samples.push({ atMs: now, lagMs: clamped });
|
||||
this.#lastKnownLagMs = clamped;
|
||||
const cutoff = now - this.#windowMs;
|
||||
while (this.#samples.length > 0 && this.#samples[0].atMs < cutoff) {
|
||||
this.#samples.shift();
|
||||
}
|
||||
this.options.onSample?.(clamped, source);
|
||||
}
|
||||
|
||||
#tick() {
|
||||
if (Date.now() - this.#lastTouchMs > this.#idleAfterMs) {
|
||||
this.stop();
|
||||
return;
|
||||
}
|
||||
if (this.#sampling) {
|
||||
return; // a slow sample shouldn't stack
|
||||
}
|
||||
this.#sampling = true;
|
||||
this.options.source
|
||||
.sampleLagMs()
|
||||
.then((lagMs) => {
|
||||
if (lagMs !== undefined) {
|
||||
this.#accept(lagMs, "probe");
|
||||
}
|
||||
})
|
||||
.catch(() => {
|
||||
// sampling errors never propagate; the estimate just ages
|
||||
})
|
||||
.finally(() => {
|
||||
this.#sampling = false;
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,93 @@
|
||||
import { $replica } from "~/db.server";
|
||||
import { env } from "~/env.server";
|
||||
import { singleton } from "~/utils/singleton";
|
||||
import { FEATURE_FLAG } from "~/v3/featureFlags";
|
||||
import { makeFlag } from "~/v3/featureFlags.server";
|
||||
import { logger } from "../logger.server";
|
||||
import { type RealtimeEnvironment } from "../realtimeClient.server";
|
||||
import { realtimeClient } from "../realtimeClientGlobal.server";
|
||||
import { BoundedTtlCache } from "./boundedTtlCache";
|
||||
import { type RealtimeStreamClient } from "./nativeRealtimeClient.server";
|
||||
import { getNativeRealtimeClient } from "./nativeRealtimeClientInstance.server";
|
||||
import { getShadowRealtimeClient } from "./shadowRealtimeClientInstance.server";
|
||||
|
||||
type RealtimeBackend = "electric" | "native" | "shadow";
|
||||
|
||||
// Two gates: the env master switch, then the per-org `realtimeBackend` feature flag (cached so
|
||||
// long-polls don't hit the DB per request), falling back to REALTIME_BACKEND_DEFAULT.
|
||||
const nativeBackendEnabled = env.REALTIME_BACKEND_NATIVE_ENABLED === "1";
|
||||
|
||||
const flag = singleton("realtimeBackendFlag", () => makeFlag($replica));
|
||||
const backendCache = singleton(
|
||||
"realtimeBackendCache",
|
||||
() =>
|
||||
new BoundedTtlCache<RealtimeBackend>(
|
||||
env.REALTIME_BACKEND_FLAG_CACHE_TTL_MS,
|
||||
env.REALTIME_BACKEND_FLAG_CACHE_MAX_ENTRIES
|
||||
)
|
||||
);
|
||||
|
||||
export async function resolveRealtimeStreamClient(
|
||||
environment: RealtimeEnvironment & { organization?: { featureFlags?: unknown } }
|
||||
): Promise<RealtimeStreamClient> {
|
||||
if (!nativeBackendEnabled) {
|
||||
return realtimeClient;
|
||||
}
|
||||
|
||||
// The authenticated environment already carries the org's feature flags; pass them
|
||||
// through so a cache miss doesn't need an extra organization read.
|
||||
const orgFeatureFlags = environment.organization
|
||||
? (environment.organization.featureFlags ?? {})
|
||||
: undefined;
|
||||
|
||||
switch (await getRealtimeBackend(environment.organizationId, orgFeatureFlags)) {
|
||||
case "native":
|
||||
return getNativeRealtimeClient();
|
||||
case "shadow":
|
||||
// The client is still served Electric; the native path is diffed in the background.
|
||||
return getShadowRealtimeClient();
|
||||
case "electric":
|
||||
default:
|
||||
return realtimeClient;
|
||||
}
|
||||
}
|
||||
|
||||
async function getRealtimeBackend(
|
||||
organizationId: string,
|
||||
orgFeatureFlags: unknown | undefined
|
||||
): Promise<RealtimeBackend> {
|
||||
const cached = backendCache.get(organizationId);
|
||||
if (cached !== undefined) {
|
||||
return cached;
|
||||
}
|
||||
|
||||
let backend: RealtimeBackend = env.REALTIME_BACKEND_DEFAULT;
|
||||
|
||||
try {
|
||||
const overrides =
|
||||
orgFeatureFlags !== undefined
|
||||
? orgFeatureFlags
|
||||
: (
|
||||
await $replica.organization.findFirst({
|
||||
where: { id: organizationId },
|
||||
select: { featureFlags: true },
|
||||
})
|
||||
)?.featureFlags;
|
||||
|
||||
backend = await flag({
|
||||
key: FEATURE_FLAG.realtimeBackend,
|
||||
defaultValue: env.REALTIME_BACKEND_DEFAULT,
|
||||
overrides: (overrides as Record<string, unknown>) ?? {},
|
||||
});
|
||||
} catch (error) {
|
||||
// Never let a flag lookup failure break the realtime feed.
|
||||
logger.error("[resolveRealtimeStreamClient] failed to resolve realtimeBackend flag", {
|
||||
organizationId,
|
||||
error,
|
||||
});
|
||||
backend = env.REALTIME_BACKEND_DEFAULT;
|
||||
}
|
||||
|
||||
backendCache.set(organizationId, backend);
|
||||
return backend;
|
||||
}
|
||||
@@ -0,0 +1,332 @@
|
||||
import type { RedisClient, RedisWithClusterOptions } from "~/redis.server";
|
||||
import { createRedisClient } from "~/redis.server";
|
||||
import { logger } from "../logger.server";
|
||||
|
||||
export const CHANGE_RECORD_VERSION = 1;
|
||||
|
||||
/**
|
||||
* A self-describing run-change fact published once to the run's environment channel; row state is
|
||||
* never on the wire. `tags` present (even `[]`) marks a "full" record a feed can classify locally;
|
||||
* `tags` absent marks a "partial" record (envId+runId only) a tag feed must hydrate to classify.
|
||||
*/
|
||||
export type ChangeRecord = {
|
||||
v: number;
|
||||
runId: string;
|
||||
envId: string;
|
||||
tags?: string[];
|
||||
batchId?: string | null;
|
||||
createdAtMs?: number;
|
||||
updatedAtMs?: number;
|
||||
status?: string;
|
||||
};
|
||||
|
||||
/** What a publish site provides; the notifier stamps the version. */
|
||||
export type ChangeRecordInput = Omit<ChangeRecord, "v">;
|
||||
|
||||
export function encodeChangeRecord(record: ChangeRecord): string {
|
||||
return JSON.stringify(record);
|
||||
}
|
||||
|
||||
/** Decode a wire message into a ChangeRecord; a bare/malformed frame degrades to a partial record rather than throwing. */
|
||||
export function decodeChangeRecord(message: string): ChangeRecord {
|
||||
if (message.length === 0 || message[0] !== "{") {
|
||||
return { v: 0, runId: message, envId: "" };
|
||||
}
|
||||
try {
|
||||
const parsed = JSON.parse(message) as Partial<ChangeRecord>;
|
||||
if (parsed && typeof parsed.runId === "string") {
|
||||
return {
|
||||
v: parsed.v ?? 0,
|
||||
runId: parsed.runId,
|
||||
envId: parsed.envId ?? "",
|
||||
tags: parsed.tags,
|
||||
batchId: parsed.batchId,
|
||||
createdAtMs: parsed.createdAtMs,
|
||||
updatedAtMs: parsed.updatedAtMs,
|
||||
status: parsed.status,
|
||||
};
|
||||
}
|
||||
} catch {
|
||||
// fall through to the bare-runId fallback
|
||||
}
|
||||
return { v: 0, runId: message, envId: "" };
|
||||
}
|
||||
|
||||
export type RunChangeNotifierOptions = {
|
||||
redis: RedisWithClusterOptions;
|
||||
/** Channel name prefix; the envId is appended inside a hash-tag for slot locality. */
|
||||
channelPrefix?: string;
|
||||
connectionName?: string;
|
||||
/** Leading-edge throttle (ms) for the per-env channel, bounding the wake rate per env. Defaults to 100ms; 0 disables. */
|
||||
envWakeCoalesceWindowMs?: number;
|
||||
/** Use Redis sharded pub/sub (SSUBSCRIBE/SPUBLISH); cluster-only and requires `clusterOptions.shardedSubscribers`. Defaults to false (classic). */
|
||||
shardedPubSub?: boolean;
|
||||
/** Observability hook: a publish settled (ok) or failed (the leading degradation signal). */
|
||||
onPublishResult?: (ok: boolean) => void;
|
||||
/** Observability hook: a raw channel message arrived (pre-coalesce). */
|
||||
onMessageReceived?: () => void;
|
||||
/** Observability hook: a coalesced batch was delivered to listeners (records per batch). */
|
||||
onBatchDelivered?: (recordCount: number) => void;
|
||||
};
|
||||
|
||||
const DEFAULT_CHANNEL_PREFIX = "realtime:";
|
||||
const DEFAULT_ENV_WAKE_COALESCE_WINDOW_MS = 100;
|
||||
|
||||
/**
|
||||
* RunChangeNotifier — carries "run X changed" facts from write sites to the realtime feeds over ONE
|
||||
* per-environment channel (`<prefix>env:{<envId>}`, hash-tagged so an env stays on one cluster slot).
|
||||
* Uses one shared multiplexed subscriber per process (refcounted), created lazily, and a fire-and-forget
|
||||
* `publish` that never throws — a dropped publish only costs latency because the consumer has a backstop.
|
||||
*/
|
||||
export class RunChangeNotifier {
|
||||
#publisher: RedisClient | undefined;
|
||||
#subscriber: RedisClient | undefined;
|
||||
readonly #listeners = new Map<string, Set<(records: ChangeRecord[]) => void>>();
|
||||
/** Per-channel accumulator of records since the last delivery, deduped by runId (latest per run wins), so a coalesced wake carries every run that moved. */
|
||||
readonly #pending = new Map<string, Map<string, ChangeRecord>>();
|
||||
readonly #channelPrefix: string;
|
||||
readonly #connectionName: string;
|
||||
readonly #coalesceWindowMs: number;
|
||||
/** When true, use sharded pub/sub (SSUBSCRIBE/SPUBLISH/smessage) — see options. */
|
||||
readonly #sharded: boolean;
|
||||
/** Active coalescing windows per channel. */
|
||||
readonly #coalesceTimers = new Map<string, ReturnType<typeof setTimeout>>();
|
||||
/** Channels that received a message while their window was open (need a trailing wake). */
|
||||
readonly #coalesceDirty = new Set<string>();
|
||||
|
||||
constructor(private readonly options: RunChangeNotifierOptions) {
|
||||
this.#channelPrefix = options.channelPrefix ?? DEFAULT_CHANNEL_PREFIX;
|
||||
this.#connectionName = options.connectionName ?? "trigger:realtime:run-change-notifier";
|
||||
this.#coalesceWindowMs = options.envWakeCoalesceWindowMs ?? DEFAULT_ENV_WAKE_COALESCE_WINDOW_MS;
|
||||
this.#sharded = options.shardedPubSub ?? false;
|
||||
}
|
||||
|
||||
/** Fire-and-forget publish of a run-changed fact to the run's environment channel; never throws. */
|
||||
publish(input: ChangeRecordInput): void {
|
||||
const record: ChangeRecord = { v: CHANGE_RECORD_VERSION, ...input };
|
||||
this.#publishToChannel(this.#channelForEnv(record.envId), encodeChangeRecord(record));
|
||||
}
|
||||
|
||||
/** Fire-and-forget publish of many run-changed facts. Never throws. */
|
||||
publishMany(inputs: ChangeRecordInput[]): void {
|
||||
for (const input of inputs) {
|
||||
this.publish(input);
|
||||
}
|
||||
}
|
||||
|
||||
#publishToChannel(channel: string, payload: string): void {
|
||||
try {
|
||||
const publisher = this.#ensurePublisher();
|
||||
// Sharded pub/sub (SPUBLISH) routes to the channel's slot owner; classic PUBLISH
|
||||
// broadcasts cluster-wide. The channel is hash-tagged by envId.
|
||||
const result = this.#sharded
|
||||
? publisher.spublish(channel, payload)
|
||||
: publisher.publish(channel, payload);
|
||||
if (typeof (result as Promise<number>)?.then === "function") {
|
||||
(result as Promise<number>).then(
|
||||
() => this.options.onPublishResult?.(true),
|
||||
(error) => {
|
||||
this.options.onPublishResult?.(false);
|
||||
logger.error("[runChangeNotifier] Failed to publish run-changed notification", {
|
||||
error,
|
||||
channel,
|
||||
});
|
||||
}
|
||||
);
|
||||
} else {
|
||||
this.options.onPublishResult?.(true);
|
||||
}
|
||||
} catch (error) {
|
||||
this.options.onPublishResult?.(false);
|
||||
logger.error("[runChangeNotifier] Failed to publish run-changed notification", {
|
||||
error,
|
||||
channel,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/** Subscribe to an env's run-change stream; refcounted over the shared subscriber (first listener SUBSCRIBEs, last UNSUBSCRIBEs). */
|
||||
subscribeToEnv(environmentId: string, onBatch: (records: ChangeRecord[]) => void): () => void {
|
||||
const channel = this.#channelForEnv(environmentId);
|
||||
const subscriber = this.#ensureSubscriber();
|
||||
|
||||
let listeners = this.#listeners.get(channel);
|
||||
if (!listeners) {
|
||||
listeners = new Set();
|
||||
this.#listeners.set(channel, listeners);
|
||||
this.#subscribeChannel(subscriber, channel).catch((error) => {
|
||||
logger.error("[runChangeNotifier] Failed to subscribe to run-change channel", {
|
||||
error,
|
||||
channel,
|
||||
});
|
||||
});
|
||||
}
|
||||
listeners.add(onBatch);
|
||||
|
||||
let unsubscribed = false;
|
||||
return () => {
|
||||
if (unsubscribed) {
|
||||
return;
|
||||
}
|
||||
unsubscribed = true;
|
||||
|
||||
const current = this.#listeners.get(channel);
|
||||
if (!current) {
|
||||
return;
|
||||
}
|
||||
current.delete(onBatch);
|
||||
if (current.size === 0) {
|
||||
// Drop the channel from the map only after Redis confirms UNSUBSCRIBE and no new listener re-subscribed in the meantime.
|
||||
this.#unsubscribeChannel(subscriber, channel)
|
||||
.then(() => {
|
||||
const latest = this.#listeners.get(channel);
|
||||
if (!latest) {
|
||||
return;
|
||||
}
|
||||
if (latest.size === 0) {
|
||||
this.#listeners.delete(channel);
|
||||
} else {
|
||||
// A listener arrived during the in-flight UNSUBSCRIBE; re-subscribe so it keeps receiving (the backstop covers the gap).
|
||||
this.#subscribeChannel(subscriber, channel).catch((error) => {
|
||||
logger.error("[runChangeNotifier] Failed to re-subscribe to run-change channel", {
|
||||
error,
|
||||
channel,
|
||||
});
|
||||
});
|
||||
}
|
||||
})
|
||||
.catch((error) => {
|
||||
// UNSUBSCRIBE failed (likely still subscribed in Redis): keep the empty map entry so a future subscriber reuses it.
|
||||
logger.error("[runChangeNotifier] Failed to unsubscribe from run-change channel", {
|
||||
error,
|
||||
channel,
|
||||
});
|
||||
});
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
/** Number of distinct env channels currently subscribed (for metrics). */
|
||||
get activeSubscriptionCount(): number {
|
||||
return this.#listeners.size;
|
||||
}
|
||||
|
||||
async quit(): Promise<void> {
|
||||
for (const timer of this.#coalesceTimers.values()) {
|
||||
clearTimeout(timer);
|
||||
}
|
||||
this.#coalesceTimers.clear();
|
||||
this.#coalesceDirty.clear();
|
||||
this.#pending.clear();
|
||||
await Promise.allSettled([this.#subscriber?.quit(), this.#publisher?.quit()]);
|
||||
this.#subscriber = undefined;
|
||||
this.#publisher = undefined;
|
||||
this.#listeners.clear();
|
||||
}
|
||||
|
||||
#ensurePublisher(): RedisClient {
|
||||
if (!this.#publisher) {
|
||||
// Publishes are fire-and-forget with a consumer-side backstop, so a dropped publish is
|
||||
// latency-only. Cap retries (vs ioredis's default 20) so a pub/sub outage rejects publishes
|
||||
// after ~1 reconnect cycle instead of buffering them in memory across the fleet.
|
||||
this.#publisher = createRedisClient(`${this.#connectionName}:pub`, {
|
||||
...this.options.redis,
|
||||
maxRetriesPerRequest: 1,
|
||||
});
|
||||
}
|
||||
return this.#publisher;
|
||||
}
|
||||
|
||||
#ensureSubscriber(): RedisClient {
|
||||
if (!this.#subscriber) {
|
||||
const subscriber = createRedisClient(`${this.#connectionName}:sub`, this.options.redis);
|
||||
const onMessage = (channel: string, message: string) => this.#onMessage(channel, message);
|
||||
// Classic pub/sub delivers "message"; sharded pub/sub delivers "smessage". Register
|
||||
// both so the delivery path is identical regardless of mode.
|
||||
subscriber.on("message", onMessage);
|
||||
subscriber.on("smessage", onMessage);
|
||||
this.#subscriber = subscriber;
|
||||
}
|
||||
return this.#subscriber;
|
||||
}
|
||||
|
||||
/** SUBSCRIBE (classic) vs SSUBSCRIBE (sharded, cluster-only). */
|
||||
#subscribeChannel(subscriber: RedisClient, channel: string): Promise<unknown> {
|
||||
return this.#sharded ? subscriber.ssubscribe(channel) : subscriber.subscribe(channel);
|
||||
}
|
||||
|
||||
/** UNSUBSCRIBE (classic) vs SUNSUBSCRIBE (sharded, cluster-only). */
|
||||
#unsubscribeChannel(subscriber: RedisClient, channel: string): Promise<unknown> {
|
||||
return this.#sharded ? subscriber.sunsubscribe(channel) : subscriber.unsubscribe(channel);
|
||||
}
|
||||
|
||||
#onMessage(channel: string, message: string) {
|
||||
this.options.onMessageReceived?.();
|
||||
// Accumulate the decoded record (deduped by runId) before delivering, so a coalesced
|
||||
// wake carries every run that moved during the window.
|
||||
this.#addPending(channel, decodeChangeRecord(message));
|
||||
|
||||
if (this.#coalesceWindowMs > 0) {
|
||||
this.#deliverCoalesced(channel);
|
||||
return;
|
||||
}
|
||||
this.#deliver(channel);
|
||||
}
|
||||
|
||||
/** Accumulate a record into the channel's pending batch, deduped by runId (a later
|
||||
* record for the same run replaces the earlier one, keeping the freshest keys). */
|
||||
#addPending(channel: string, record: ChangeRecord) {
|
||||
let batch = this.#pending.get(channel);
|
||||
if (!batch) {
|
||||
batch = new Map();
|
||||
this.#pending.set(channel, batch);
|
||||
}
|
||||
batch.set(record.runId, record);
|
||||
}
|
||||
|
||||
#deliver(channel: string) {
|
||||
// Drain the accumulated batch (and clear it) so listeners woken now get every run that
|
||||
// changed since the last delivery, and a later message starts a fresh batch.
|
||||
const batchMap = this.#pending.get(channel);
|
||||
const batch = batchMap ? [...batchMap.values()] : [];
|
||||
this.#pending.delete(channel);
|
||||
|
||||
const listeners = this.#listeners.get(channel);
|
||||
if (!listeners || batch.length === 0) {
|
||||
return;
|
||||
}
|
||||
this.options.onBatchDelivered?.(batch.length);
|
||||
for (const onBatch of [...listeners]) {
|
||||
onBatch(batch);
|
||||
}
|
||||
}
|
||||
|
||||
/** Leading-edge throttle capping the wake rate to ~1/window: deliver the first wake immediately, then one trailing wake per window while activity continues. Lossless. */
|
||||
#deliverCoalesced(channel: string) {
|
||||
if (this.#coalesceTimers.has(channel)) {
|
||||
this.#coalesceDirty.add(channel);
|
||||
return;
|
||||
}
|
||||
this.#deliver(channel);
|
||||
this.#openCoalesceWindow(channel);
|
||||
}
|
||||
|
||||
#openCoalesceWindow(channel: string) {
|
||||
const timer = setTimeout(() => {
|
||||
this.#coalesceTimers.delete(channel);
|
||||
if (this.#coalesceDirty.delete(channel)) {
|
||||
this.#deliver(channel);
|
||||
this.#openCoalesceWindow(channel);
|
||||
}
|
||||
}, this.#coalesceWindowMs);
|
||||
// Don't let a pending coalescing window hold the process open at shutdown.
|
||||
timer.unref?.();
|
||||
this.#coalesceTimers.set(channel, timer);
|
||||
}
|
||||
|
||||
// Hash-tagged (`...{<envId>}`) so all of an env's traffic maps to one cluster slot (one
|
||||
// shard) under sharded pub/sub.
|
||||
#channelForEnv(environmentId: string): string {
|
||||
return `${this.#channelPrefix}env:{${environmentId}}`;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,87 @@
|
||||
import { env } from "~/env.server";
|
||||
import { engine } from "~/v3/runEngine.server";
|
||||
import { logger } from "../logger.server";
|
||||
import { publishChangeRecord } from "./runChangeNotifierInstance.server";
|
||||
|
||||
/**
|
||||
* Builds and publishes a self-describing `ChangeRecord` for the lifecycle events whose engine-bus payload
|
||||
* already carries env + tags + batchId. Terminal transitions, runAttemptFailed, and runMetadataUpdated publish
|
||||
* from `runEngineHandlers.server.ts` instead. Coverage isn't exhaustive — a dropped transition only adds latency
|
||||
* because the consumer has a periodic backstop full-resolve. The env master switch is `REALTIME_BACKEND_NATIVE_ENABLED`.
|
||||
*/
|
||||
export function registerRunChangeNotifierHandlers() {
|
||||
// Return truthy in every path so singleton() caches this factory and never re-runs it (re-running would attach duplicate engine-bus listeners on dev reload).
|
||||
if (env.REALTIME_BACKEND_NATIVE_ENABLED !== "1") {
|
||||
return true;
|
||||
}
|
||||
|
||||
// Run created: the first signal for a brand-new run (born QUEUED with no status transition), so it surfaces before ClickHouse ingests it.
|
||||
engine.eventBus.on("runCreated", ({ run, environment }) => {
|
||||
publishChangeRecord({
|
||||
runId: run.id,
|
||||
envId: environment.id,
|
||||
tags: run.runTags,
|
||||
batchId: run.batchId,
|
||||
});
|
||||
});
|
||||
|
||||
// Status transitions (checkpoint suspend/resume, pending version, dequeue).
|
||||
engine.eventBus.on("runStatusChanged", ({ run, environment }) => {
|
||||
publishChangeRecord({
|
||||
runId: run.id,
|
||||
envId: environment.id,
|
||||
tags: run.runTags,
|
||||
batchId: run.batchId,
|
||||
});
|
||||
});
|
||||
|
||||
// Dequeue/lock (sets startedAt) and attempt start (DEQUEUED -> EXECUTING) — the
|
||||
// most-watched "my run started" transitions.
|
||||
engine.eventBus.on("runLocked", ({ run, environment }) => {
|
||||
publishChangeRecord({
|
||||
runId: run.id,
|
||||
envId: environment.id,
|
||||
tags: run.runTags,
|
||||
batchId: run.batchId,
|
||||
});
|
||||
});
|
||||
engine.eventBus.on("runAttemptStarted", ({ run, environment }) => {
|
||||
publishChangeRecord({
|
||||
runId: run.id,
|
||||
envId: environment.id,
|
||||
tags: run.runTags,
|
||||
batchId: run.batchId,
|
||||
});
|
||||
});
|
||||
|
||||
engine.eventBus.on("runRetryScheduled", ({ run, environment }) => {
|
||||
publishChangeRecord({
|
||||
runId: run.id,
|
||||
envId: environment.id,
|
||||
tags: run.runTags,
|
||||
batchId: run.batchId,
|
||||
});
|
||||
});
|
||||
|
||||
// Delay lifecycle (delayUntil / queued-after-delay changes).
|
||||
engine.eventBus.on("runDelayRescheduled", ({ run, environment }) => {
|
||||
publishChangeRecord({
|
||||
runId: run.id,
|
||||
envId: environment.id,
|
||||
tags: run.runTags,
|
||||
batchId: run.batchId,
|
||||
});
|
||||
});
|
||||
engine.eventBus.on("runEnqueuedAfterDelay", ({ run, environment }) => {
|
||||
publishChangeRecord({
|
||||
runId: run.id,
|
||||
envId: environment.id,
|
||||
tags: run.runTags,
|
||||
batchId: run.batchId,
|
||||
});
|
||||
});
|
||||
|
||||
logger.info("[runChangeNotifier] realtime change-record builder registered");
|
||||
|
||||
return true;
|
||||
}
|
||||
@@ -0,0 +1,99 @@
|
||||
import { getMeter } from "@internal/tracing";
|
||||
import { env } from "~/env.server";
|
||||
import { singleton } from "~/utils/singleton";
|
||||
import { logger } from "../logger.server";
|
||||
import { RunChangeNotifier, type ChangeRecordInput } from "./runChangeNotifier.server";
|
||||
|
||||
/**
|
||||
* Process-singleton wiring for the RunChangeNotifier plus the gated convenience functions write sites
|
||||
* delegate to. The notifier is constructed lazily, so `REALTIME_BACKEND_NATIVE_ENABLED=0` (default) opens no Redis connections.
|
||||
*/
|
||||
const nativeBackendEnabled = env.REALTIME_BACKEND_NATIVE_ENABLED === "1";
|
||||
|
||||
function initializeRunChangeNotifier(): RunChangeNotifier {
|
||||
const clusterMode = env.REALTIME_BACKEND_NATIVE_PUBSUB_REDIS_CLUSTER_MODE_ENABLED === "1";
|
||||
// Sharded pub/sub only works against a cluster; classic pub/sub there would
|
||||
// broadcast every message to every node, so this is what actually shards load.
|
||||
const shardedPubSub =
|
||||
clusterMode && env.REALTIME_BACKEND_NATIVE_PUBSUB_REDIS_SHARDED_ENABLED === "1";
|
||||
|
||||
const meter = getMeter("realtime-notifier");
|
||||
|
||||
const publishes = meter.createCounter("realtime_notifier.publishes", {
|
||||
description:
|
||||
"Change-record publishes by outcome. Failures are the leading indicator that feeds are degrading to their backstops (pub/sub Redis trouble).",
|
||||
});
|
||||
|
||||
const received = meter.createCounter("realtime_notifier.messages_received", {
|
||||
description: "Raw channel messages received by this instance's subscriber, pre-coalesce.",
|
||||
});
|
||||
|
||||
const delivered = meter.createCounter("realtime_notifier.batches_delivered", {
|
||||
description:
|
||||
"Coalesced batches delivered to listeners. received/batches = the coalesce ratio (how hard a busy env is being collapsed).",
|
||||
});
|
||||
|
||||
const notifier = new RunChangeNotifier({
|
||||
redis: {
|
||||
host: env.REALTIME_BACKEND_NATIVE_PUBSUB_REDIS_HOST,
|
||||
port: env.REALTIME_BACKEND_NATIVE_PUBSUB_REDIS_PORT,
|
||||
username: env.REALTIME_BACKEND_NATIVE_PUBSUB_REDIS_USERNAME,
|
||||
password: env.REALTIME_BACKEND_NATIVE_PUBSUB_REDIS_PASSWORD,
|
||||
tlsDisabled: env.REALTIME_BACKEND_NATIVE_PUBSUB_REDIS_TLS_DISABLED === "true",
|
||||
clusterMode,
|
||||
// One subscriber connection per shard so SSUBSCRIBE routes to the slot owner.
|
||||
...(shardedPubSub ? { clusterOptions: { shardedSubscribers: true } } : {}),
|
||||
},
|
||||
envWakeCoalesceWindowMs: env.REALTIME_BACKEND_NATIVE_ENV_WAKE_COALESCE_WINDOW_MS,
|
||||
shardedPubSub,
|
||||
onPublishResult: (ok) => publishes.add(1, { result: ok ? "ok" : "error" }),
|
||||
onMessageReceived: () => received.add(1),
|
||||
onBatchDelivered: () => delivered.add(1),
|
||||
});
|
||||
|
||||
meter
|
||||
.createObservableGauge("realtime_notifier.active_subscriptions", {
|
||||
description: "Distinct env channels currently subscribed for realtime change notifications.",
|
||||
})
|
||||
.addCallback((result) => result.observe(notifier.activeSubscriptionCount));
|
||||
|
||||
return notifier;
|
||||
}
|
||||
|
||||
/** Lazily construct (and memoize) the notifier singleton. */
|
||||
export function getRunChangeNotifier(): RunChangeNotifier {
|
||||
return singleton("runChangeNotifier", initializeRunChangeNotifier);
|
||||
}
|
||||
|
||||
/** Whether the notifier subsystem is enabled for this process. */
|
||||
export function isRunChangeNotifierEnabled(): boolean {
|
||||
return nativeBackendEnabled;
|
||||
}
|
||||
|
||||
/** Fire-and-forget publish of a run-changed record. No-op (and no notifier construction)
|
||||
* when disabled, so publish sites can call it unconditionally. */
|
||||
export function publishChangeRecord(input: ChangeRecordInput): void {
|
||||
if (!nativeBackendEnabled) {
|
||||
return;
|
||||
}
|
||||
// Publish runs on the run-engine event bus / metadata flush loop; lazy init + encoding happen
|
||||
// before the notifier's own try/catch, so guard the whole call — it must never throw at its caller.
|
||||
try {
|
||||
getRunChangeNotifier().publish(input);
|
||||
} catch (error) {
|
||||
logger.error("[runChangeNotifier] publishChangeRecord threw; dropping notification", { error });
|
||||
}
|
||||
}
|
||||
|
||||
export function publishManyChangeRecords(inputs: ChangeRecordInput[]): void {
|
||||
if (!nativeBackendEnabled) {
|
||||
return;
|
||||
}
|
||||
try {
|
||||
getRunChangeNotifier().publishMany(inputs);
|
||||
} catch (error) {
|
||||
logger.error("[runChangeNotifier] publishManyChangeRecords threw; dropping notifications", {
|
||||
error,
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,174 @@
|
||||
import {
|
||||
type Prisma,
|
||||
type PrismaClient,
|
||||
type PrismaClientOrTransaction,
|
||||
} from "@trigger.dev/database";
|
||||
import type { RunStore } from "@internal/run-store";
|
||||
import { BoundedTtlCache } from "./boundedTtlCache";
|
||||
import { RESERVED_COLUMNS, type RealtimeRunRow } from "./electricStreamProtocol.server";
|
||||
|
||||
/**
|
||||
* RunReader — the pluggable read half of the native-backend realtime feed: ClickHouse is filter-only
|
||||
* (resolves ids), Postgres always hydrates row columns. Owns the `RunHydrator` (by-id) and the
|
||||
* `RunListResolver` interface (the tag/list filter -> id-set seam, implemented over ClickHouse).
|
||||
*/
|
||||
|
||||
/** The TaskRun columns the realtime feed projects (mirrors DEFAULT_ELECTRIC_COLUMNS). */
|
||||
export const RUN_HYDRATOR_SELECT = {
|
||||
id: true,
|
||||
taskIdentifier: true,
|
||||
createdAt: true,
|
||||
updatedAt: true,
|
||||
startedAt: true,
|
||||
delayUntil: true,
|
||||
queuedAt: true,
|
||||
expiredAt: true,
|
||||
completedAt: true,
|
||||
friendlyId: true,
|
||||
number: true,
|
||||
isTest: true,
|
||||
status: true,
|
||||
usageDurationMs: true,
|
||||
costInCents: true,
|
||||
baseCostInCents: true,
|
||||
ttl: true,
|
||||
payload: true,
|
||||
payloadType: true,
|
||||
metadata: true,
|
||||
metadataType: true,
|
||||
output: true,
|
||||
outputType: true,
|
||||
runTags: true,
|
||||
error: true,
|
||||
realtimeStreams: true,
|
||||
} satisfies Prisma.TaskRunSelect;
|
||||
|
||||
/** Columns hydrated regardless of `skipColumns`: `id` keys the row, `updatedAt` drives the offset and working-set diff. */
|
||||
const ALWAYS_HYDRATED_COLUMNS = new Set<string>(["id", "updatedAt", ...RESERVED_COLUMNS]);
|
||||
|
||||
/** Project `RUN_HYDRATOR_SELECT` down to the columns the client didn't skip (plus
|
||||
* the always-needed ones). An empty skip set returns the full select unchanged. */
|
||||
export function buildHydratorSelect(skipColumns: string[] = []): Prisma.TaskRunSelect {
|
||||
if (skipColumns.length === 0) {
|
||||
return RUN_HYDRATOR_SELECT;
|
||||
}
|
||||
const skip = new Set(skipColumns);
|
||||
const select: Record<string, boolean> = {};
|
||||
for (const column of Object.keys(RUN_HYDRATOR_SELECT)) {
|
||||
if (ALWAYS_HYDRATED_COLUMNS.has(column) || !skip.has(column)) {
|
||||
select[column] = true;
|
||||
}
|
||||
}
|
||||
return select as Prisma.TaskRunSelect;
|
||||
}
|
||||
|
||||
export type RunListFilter = {
|
||||
organizationId: string;
|
||||
projectId: string;
|
||||
environmentId: string;
|
||||
/** Contains-ANY tag match (OR). Omit/empty for non-tag feeds. */
|
||||
tags?: string[];
|
||||
/** Restrict to a single batch (internal batch id) — the batch feed. */
|
||||
batchId?: string;
|
||||
/** Lower bound on createdAt (the tag-list feed pins this; batch omits it). */
|
||||
createdAtAfter?: Date;
|
||||
/** Hard cap on the result set so a broad filter can't unbound the snapshot. */
|
||||
limit: number;
|
||||
};
|
||||
|
||||
/** Resolves a tag/list filter into the matching run id-set, filter-only (rows hydrated from Postgres by id afterward). ClickHouse impl in `clickHouseRunListResolver.server.ts`. */
|
||||
export interface RunListResolver {
|
||||
resolveMatchingRunIds(filter: RunListFilter): Promise<string[]>;
|
||||
}
|
||||
|
||||
export type RunHydratorOptions = {
|
||||
/** A read-replica Prisma client (`$replica`). Always Postgres. */
|
||||
replica: Pick<PrismaClient, "taskRun">;
|
||||
/** RunStore the reads are routed through; `replica` is passed as the read client. */
|
||||
runStore: RunStore;
|
||||
/** Read-through cache TTL (ms) collapsing duplicate refetches for the same run. Set 0 to disable. Defaults to 250ms. */
|
||||
cacheTtlMs?: number;
|
||||
/** Hard cap on cache entries before expired entries are swept. */
|
||||
maxCacheEntries?: number;
|
||||
};
|
||||
|
||||
const DEFAULT_CACHE_TTL_MS = 250;
|
||||
const DEFAULT_MAX_CACHE_ENTRIES = 5_000;
|
||||
|
||||
/** Hydrates runs by id through the runStore seam (split routing lives in the store, below this file), projected to the realtime columns; concurrent same-run refetches are single-flighted + short-TTL cached. */
|
||||
export class RunHydrator {
|
||||
readonly #inflight = new Map<string, Promise<RealtimeRunRow | null>>();
|
||||
readonly #cache: BoundedTtlCache<RealtimeRunRow | null>;
|
||||
readonly #cacheTtlMs: number;
|
||||
|
||||
constructor(private readonly options: RunHydratorOptions) {
|
||||
this.#cacheTtlMs = options.cacheTtlMs ?? DEFAULT_CACHE_TTL_MS;
|
||||
this.#cache = new BoundedTtlCache(
|
||||
this.#cacheTtlMs,
|
||||
options.maxCacheEntries ?? DEFAULT_MAX_CACHE_ENTRIES
|
||||
);
|
||||
}
|
||||
|
||||
async getRunById(environmentId: string, runId: string): Promise<RealtimeRunRow | null> {
|
||||
const key = `${environmentId}:${runId}`;
|
||||
|
||||
if (this.#cacheTtlMs > 0) {
|
||||
// A cached null is a valid "run not found" hit; only undefined is a miss.
|
||||
const cached = this.#cache.get(key);
|
||||
if (cached !== undefined) {
|
||||
return cached;
|
||||
}
|
||||
}
|
||||
|
||||
const existing = this.#inflight.get(key);
|
||||
if (existing) {
|
||||
return existing;
|
||||
}
|
||||
|
||||
const promise = this.#fetch(environmentId, runId).finally(() => this.#inflight.delete(key));
|
||||
this.#inflight.set(key, promise);
|
||||
|
||||
const row = await promise;
|
||||
|
||||
if (this.#cacheTtlMs > 0) {
|
||||
this.#cache.set(key, row);
|
||||
}
|
||||
|
||||
return row;
|
||||
}
|
||||
|
||||
/** Hydrate many runs by id in one query (order not guaranteed); `skipColumns` projects the SELECT so dropped columns aren't shipped. */
|
||||
async hydrateByIds(
|
||||
environmentId: string,
|
||||
ids: string[],
|
||||
skipColumns: string[] = []
|
||||
): Promise<RealtimeRunRow[]> {
|
||||
if (ids.length === 0) {
|
||||
return [];
|
||||
}
|
||||
const rows = await this.options.runStore.findRuns(
|
||||
{
|
||||
where: {
|
||||
runtimeEnvironmentId: environmentId,
|
||||
id: { in: ids },
|
||||
},
|
||||
select: buildHydratorSelect(skipColumns),
|
||||
},
|
||||
this.options.replica as PrismaClientOrTransaction
|
||||
);
|
||||
return rows as unknown as RealtimeRunRow[];
|
||||
}
|
||||
|
||||
async #fetch(environmentId: string, runId: string): Promise<RealtimeRunRow | null> {
|
||||
const run = await this.options.runStore.findRun(
|
||||
{
|
||||
id: runId,
|
||||
runtimeEnvironmentId: environmentId,
|
||||
},
|
||||
{ select: RUN_HYDRATOR_SELECT },
|
||||
this.options.replica as PrismaClientOrTransaction
|
||||
);
|
||||
|
||||
return (run ?? null) as RealtimeRunRow | null;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,742 @@
|
||||
// app/realtime/S2RealtimeStreams.ts
|
||||
import type { UnkeyCache } from "@internal/cache";
|
||||
import type { StreamIngestor, StreamRecord, StreamResponder, StreamResponseOptions } from "./types";
|
||||
import type { LogLevel } from "@trigger.dev/core/logger";
|
||||
import { Logger } from "@trigger.dev/core/logger";
|
||||
import { headerValue } from "@trigger.dev/core/v3";
|
||||
import { randomUUID } from "node:crypto";
|
||||
import { ServiceValidationError } from "~/v3/services/common.server";
|
||||
|
||||
// S2's per-record metered-size limit. Verified empirically against
|
||||
// cloud S2: append succeeds at metered=1048576 and 422s at 1048577
|
||||
// with `"record must have metered size less than 1 MiB"` (the "less
|
||||
// than" wording is slightly off — the boundary is inclusive).
|
||||
//
|
||||
// Metered size formula:
|
||||
// metered = 8 (record overhead) + 2*H + Σ(header name + value) + body
|
||||
// where `body` is the unescaped record body length in UTF-8 bytes and
|
||||
// `H` is the number of S2 record headers.
|
||||
//
|
||||
// We attach no record headers (H=0), so the budget reduces to:
|
||||
// 8 + body ≤ 1048576 → body ≤ 1048568
|
||||
export const S2_MAX_METERED_BYTES = 1024 * 1024; // 1 MiB
|
||||
export const S2_RECORD_BASE_OVERHEAD_BYTES = 8;
|
||||
|
||||
/**
|
||||
* Thrown when a record's metered size would exceed S2's hard per-record
|
||||
* limit. Caught by the route handler and surfaced as 413.
|
||||
*/
|
||||
export class S2RecordTooLargeError extends ServiceValidationError {
|
||||
constructor(public readonly meteredBytes: number) {
|
||||
super(
|
||||
`Record metered size ${meteredBytes} bytes exceeds the S2 per-record limit of ${S2_MAX_METERED_BYTES} bytes. Reduce tool-output size or split into smaller parts.`,
|
||||
413
|
||||
);
|
||||
this.name = "S2RecordTooLargeError";
|
||||
}
|
||||
}
|
||||
|
||||
export type S2RealtimeStreamsOptions = {
|
||||
// S2
|
||||
basin: string; // e.g., "my-basin"
|
||||
accessToken: string; // "Bearer" token issued in S2 console
|
||||
streamPrefix?: string; // defaults to ""
|
||||
|
||||
// Custom endpoint for s2-lite (self-hosted)
|
||||
endpoint?: string; // e.g., "http://localhost:4566/v1"
|
||||
|
||||
// Skip access token issuance (s2-lite doesn't support /access-tokens)
|
||||
skipAccessTokens?: boolean;
|
||||
|
||||
// Read behavior
|
||||
s2WaitSeconds?: number;
|
||||
|
||||
flushIntervalMs?: number; // how often to flush buffered chunks (default 200ms)
|
||||
maxRetries?: number; // max number of retries for failed flushes (default 10)
|
||||
|
||||
logger?: Logger;
|
||||
logLevel?: LogLevel;
|
||||
|
||||
accessTokenExpirationInMs?: number;
|
||||
|
||||
cache?: UnkeyCache<{
|
||||
accessToken: string;
|
||||
}>;
|
||||
};
|
||||
|
||||
// Ops the issued S2 access token is scoped to. `trim` is a distinct op
|
||||
// from `append` even though trim records are appended like any other —
|
||||
// without it, `AppendRecord.trim()` 403s with "Operation not permitted".
|
||||
// `chat.agent`'s per-turn trim chain depends on it.
|
||||
//
|
||||
// The fingerprint folds the ops list into the cache key, so any future
|
||||
// scope change auto-invalidates pre-deploy cached tokens.
|
||||
const S2_TOKEN_OPS = ["append", "create-stream", "trim"] as const;
|
||||
const S2_TOKEN_OPS_FINGERPRINT = [...S2_TOKEN_OPS].sort().join(",");
|
||||
|
||||
type S2IssueAccessTokenResponse = { access_token: string };
|
||||
type S2AppendInput = { records: { body: string }[] };
|
||||
type S2AppendAck = {
|
||||
start: { seq_num: number; timestamp: number };
|
||||
end: { seq_num: number; timestamp: number };
|
||||
tail: { seq_num: number; timestamp: number };
|
||||
};
|
||||
|
||||
export class S2RealtimeStreams implements StreamResponder, StreamIngestor {
|
||||
private readonly basin: string;
|
||||
private readonly baseUrl: string;
|
||||
private readonly accountUrl: string;
|
||||
private readonly endpoint?: string;
|
||||
private readonly token: string;
|
||||
private readonly streamPrefix: string;
|
||||
private readonly skipAccessTokens: boolean;
|
||||
|
||||
private readonly s2WaitSeconds: number;
|
||||
|
||||
private readonly flushIntervalMs: number;
|
||||
private readonly maxRetries: number;
|
||||
|
||||
private readonly logger: Logger;
|
||||
private readonly level: LogLevel;
|
||||
|
||||
private readonly accessTokenExpirationInMs: number;
|
||||
|
||||
private readonly cache?: UnkeyCache<{
|
||||
accessToken: string;
|
||||
}>;
|
||||
|
||||
constructor(opts: S2RealtimeStreamsOptions) {
|
||||
this.basin = opts.basin;
|
||||
this.baseUrl = opts.endpoint ?? `https://${this.basin}.b.aws.s2.dev/v1`;
|
||||
this.accountUrl = opts.endpoint ?? `https://aws.s2.dev/v1`;
|
||||
this.endpoint = opts.endpoint;
|
||||
this.token = opts.accessToken;
|
||||
this.streamPrefix = opts.streamPrefix ?? "";
|
||||
this.skipAccessTokens = opts.skipAccessTokens ?? false;
|
||||
|
||||
this.s2WaitSeconds = opts.s2WaitSeconds ?? 60;
|
||||
|
||||
this.flushIntervalMs = opts.flushIntervalMs ?? 200;
|
||||
this.maxRetries = opts.maxRetries ?? 10;
|
||||
|
||||
this.logger = opts.logger ?? new Logger("S2RealtimeStreams", opts.logLevel ?? "info");
|
||||
this.level = opts.logLevel ?? "info";
|
||||
|
||||
this.cache = opts.cache;
|
||||
this.accessTokenExpirationInMs = opts.accessTokenExpirationInMs ?? 60_000 * 60 * 24; // 1 day
|
||||
}
|
||||
|
||||
private toStreamName(runId: string, streamId: string): string {
|
||||
return `${this.streamPrefix}/runs/${runId}/${streamId}`;
|
||||
}
|
||||
|
||||
/**
|
||||
* Build an S2 stream name for a `Session`-primitive channel, addressed by
|
||||
* the session's `friendlyId` and the I/O direction. Used by the session
|
||||
* realtime routes to route traffic to `sessions/{friendlyId}/{out|in}`.
|
||||
*/
|
||||
public toSessionStreamName(friendlyId: string, io: "out" | "in"): string {
|
||||
return `${this.streamPrefix}/sessions/${friendlyId}/${io}`;
|
||||
}
|
||||
|
||||
async initializeStream(
|
||||
runId: string,
|
||||
streamId: string
|
||||
): Promise<{ responseHeaders?: Record<string, string> }> {
|
||||
return this.#initializeStreamByName(
|
||||
this.toStreamName(runId, streamId),
|
||||
`/runs/${runId}/${streamId}`
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Initialize an S2 stream by `(sessionFriendlyId, io)` — mirrors
|
||||
* {@link initializeStream} but addresses the new `sessions/*` key format.
|
||||
*/
|
||||
async initializeSessionStream(
|
||||
friendlyId: string,
|
||||
io: "out" | "in"
|
||||
): Promise<{ responseHeaders?: Record<string, string> }> {
|
||||
return this.#initializeStreamByName(
|
||||
this.toSessionStreamName(friendlyId, io),
|
||||
`/sessions/${friendlyId}/${io}`
|
||||
);
|
||||
}
|
||||
|
||||
async #initializeStreamByName(
|
||||
prefixedName: string,
|
||||
relativeName: string
|
||||
): Promise<{ responseHeaders?: Record<string, string> }> {
|
||||
const accessToken = this.skipAccessTokens
|
||||
? this.token
|
||||
: await this.getS2AccessToken(randomUUID());
|
||||
|
||||
return {
|
||||
responseHeaders: {
|
||||
"X-S2-Access-Token": accessToken,
|
||||
"X-S2-Stream-Name": this.skipAccessTokens ? prefixedName : relativeName,
|
||||
"X-S2-Basin": this.basin,
|
||||
"X-S2-Flush-Interval-Ms": this.flushIntervalMs.toString(),
|
||||
"X-S2-Max-Retries": this.maxRetries.toString(),
|
||||
...(this.endpoint ? { "X-S2-Endpoint": this.endpoint } : {}),
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
ingestData(
|
||||
stream: ReadableStream<Uint8Array>,
|
||||
runId: string,
|
||||
streamId: string,
|
||||
clientId: string,
|
||||
resumeFromChunk?: number
|
||||
): Promise<Response> {
|
||||
throw new Error("S2 streams are written to S2 via the client, not from the server");
|
||||
}
|
||||
|
||||
async appendPart(part: string, partId: string, runId: string, streamId: string): Promise<void> {
|
||||
await this.#appendPartByName(part, partId, this.toStreamName(runId, streamId));
|
||||
}
|
||||
|
||||
/** Append one record to a `Session` channel; returns its seq (same space as `session-in-event-id`). */
|
||||
async appendPartToSessionStream(
|
||||
part: string,
|
||||
partId: string,
|
||||
friendlyId: string,
|
||||
io: "out" | "in"
|
||||
): Promise<number> {
|
||||
return this.#appendPartByName(part, partId, this.toSessionStreamName(friendlyId, io));
|
||||
}
|
||||
|
||||
async #appendPartByName(part: string, partId: string, s2Stream: string): Promise<number> {
|
||||
this.logger.debug(`S2 appending to stream`, { part, stream: s2Stream });
|
||||
|
||||
const recordBody = JSON.stringify({ data: part, id: partId });
|
||||
const meteredBytes = Buffer.byteLength(recordBody, "utf8") + S2_RECORD_BASE_OVERHEAD_BYTES;
|
||||
if (meteredBytes > S2_MAX_METERED_BYTES) {
|
||||
throw new S2RecordTooLargeError(meteredBytes);
|
||||
}
|
||||
|
||||
const result = await this.s2Append(s2Stream, {
|
||||
records: [{ body: recordBody }],
|
||||
});
|
||||
|
||||
this.logger.debug(`S2 append result`, { result });
|
||||
|
||||
return result.start.seq_num;
|
||||
}
|
||||
|
||||
getLastChunkIndex(runId: string, streamId: string, clientId: string): Promise<number> {
|
||||
throw new Error("S2 streams are written to S2 via the client, not from the server");
|
||||
}
|
||||
|
||||
async readRecords(
|
||||
runId: string,
|
||||
streamId: string,
|
||||
afterSeqNum?: number
|
||||
): Promise<StreamRecord[]> {
|
||||
return this.#readRecordsByName(this.toStreamName(runId, streamId), afterSeqNum);
|
||||
}
|
||||
|
||||
/**
|
||||
* Read records from a `Session`-primitive channel starting after the
|
||||
* given sequence number. Used by the `.wait()` race-check path.
|
||||
*/
|
||||
async readSessionStreamRecords(
|
||||
friendlyId: string,
|
||||
io: "out" | "in",
|
||||
afterSeqNum?: number
|
||||
): Promise<StreamRecord[]> {
|
||||
return this.#readRecordsByName(this.toSessionStreamName(friendlyId, io), afterSeqNum);
|
||||
}
|
||||
|
||||
async #readRecordsByName(s2Stream: string, afterSeqNum?: number): Promise<StreamRecord[]> {
|
||||
const startSeq = afterSeqNum != null ? afterSeqNum + 1 : 0;
|
||||
|
||||
const qs = new URLSearchParams();
|
||||
qs.set("seq_num", String(startSeq));
|
||||
qs.set("clamp", "true");
|
||||
qs.set("wait", "0"); // Non-blocking: return immediately with existing records
|
||||
|
||||
const res = await fetch(
|
||||
`${this.baseUrl}/streams/${encodeURIComponent(s2Stream)}/records?${qs}`,
|
||||
{
|
||||
method: "GET",
|
||||
headers: {
|
||||
Authorization: `Bearer ${this.token}`,
|
||||
Accept: "text/event-stream",
|
||||
"S2-Format": "raw",
|
||||
"S2-Basin": this.basin,
|
||||
},
|
||||
}
|
||||
);
|
||||
|
||||
if (!res.ok) {
|
||||
// Stream may not exist yet (no data sent)
|
||||
if (res.status === 404) {
|
||||
return [];
|
||||
}
|
||||
const text = await res.text().catch(() => "");
|
||||
throw new Error(`S2 readRecords failed: ${res.status} ${res.statusText} ${text}`);
|
||||
}
|
||||
|
||||
// Parse the SSE response body to extract records
|
||||
const body = await res.text();
|
||||
return this.parseSSEBatchRecords(body);
|
||||
}
|
||||
|
||||
private parseSSEBatchRecords(sseText: string): StreamRecord[] {
|
||||
const records: StreamRecord[] = [];
|
||||
|
||||
// SSE events are separated by double newlines
|
||||
const events = sseText.split("\n\n").filter((e) => e.trim());
|
||||
|
||||
for (const event of events) {
|
||||
const lines = event.split("\n");
|
||||
let eventType: string | undefined;
|
||||
let data: string | undefined;
|
||||
|
||||
for (const line of lines) {
|
||||
if (line.startsWith("event:")) {
|
||||
eventType = line.slice(6).trim();
|
||||
} else if (line.startsWith("data:")) {
|
||||
data = line.slice(5).trim();
|
||||
}
|
||||
}
|
||||
|
||||
if (eventType === "batch" && data) {
|
||||
try {
|
||||
const parsed = JSON.parse(data) as {
|
||||
records: Array<{
|
||||
body: string;
|
||||
seq_num: number;
|
||||
timestamp: number;
|
||||
headers?: Array<[string, string]>;
|
||||
}>;
|
||||
};
|
||||
|
||||
for (const record of parsed.records) {
|
||||
// S2 command records (trim/fence) have a single header with
|
||||
// empty name. Skip — callers want only data + Trigger control
|
||||
// records.
|
||||
if (record.headers?.[0]?.[0] === "") {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Data records carry a JSON envelope; Trigger control records
|
||||
// have an empty body and route via headers. Tolerate non-JSON
|
||||
// bodies so a control record (or a malformed data record)
|
||||
// doesn't take the whole batch down with it.
|
||||
let parsedBody: { data: string; id: string } | undefined;
|
||||
try {
|
||||
parsedBody = JSON.parse(record.body) as { data: string; id: string };
|
||||
} catch {
|
||||
parsedBody = undefined;
|
||||
}
|
||||
records.push({
|
||||
data: parsedBody?.data ?? "",
|
||||
id: parsedBody?.id ?? "",
|
||||
seqNum: record.seq_num,
|
||||
headers: record.headers,
|
||||
});
|
||||
}
|
||||
} catch {
|
||||
// Skip malformed events
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return records;
|
||||
}
|
||||
|
||||
// ---------- Serve SSE from S2 ----------
|
||||
|
||||
async streamResponse(
|
||||
request: Request,
|
||||
runId: string,
|
||||
streamId: string,
|
||||
signal: AbortSignal,
|
||||
options?: StreamResponseOptions
|
||||
): Promise<Response> {
|
||||
return this.#streamResponseByName(this.toStreamName(runId, streamId), signal, options);
|
||||
}
|
||||
|
||||
/**
|
||||
* Serve SSE from a `Session`-primitive channel addressed by
|
||||
* `(friendlyId, io)`.
|
||||
*
|
||||
* For `io=out`, peek the tail of the stream. If the most recent
|
||||
* non-command record is a `turn-complete` control record (i.e. the
|
||||
* agent has finished a turn and is either idle-waiting on `.in` or
|
||||
* has exited), no more chunks will arrive without further user
|
||||
* action. We switch the downstream S2 read to `wait=0` (drain
|
||||
* whatever's left, close fast) and set `X-Session-Settled: true` so
|
||||
* the client knows this SSE close is terminal instead of the normal
|
||||
* 60s long-poll cycle.
|
||||
*
|
||||
* The actual tail is now usually an S2 `trim` command record (the
|
||||
* agent appends one after every turn-complete to keep `.out`
|
||||
* bounded). The peek reads two records and walks past the trim to
|
||||
* find the turn-complete underneath.
|
||||
*
|
||||
* Mid-turn tail (streaming UIMessageChunk) falls through to the
|
||||
* long-poll path; a crashed-mid-turn stream is indistinguishable
|
||||
* here and behaves like today (client sees wait=60 close, retries).
|
||||
*/
|
||||
async streamResponseFromSessionStream(
|
||||
request: Request,
|
||||
friendlyId: string,
|
||||
io: "out" | "in",
|
||||
signal: AbortSignal,
|
||||
options?: StreamResponseOptions
|
||||
): Promise<Response> {
|
||||
const s2Stream = this.toSessionStreamName(friendlyId, io);
|
||||
|
||||
let waitSeconds = options?.timeoutInSeconds ?? this.s2WaitSeconds;
|
||||
let settled = false;
|
||||
|
||||
// Only peek + settle when the client opts in via `options.peekSettled`.
|
||||
// Reconnect-on-reload paths (`TriggerChatTransport.reconnectToStream`)
|
||||
// set it; active send-a-message paths don't — otherwise the peek
|
||||
// races the newly-triggered turn's first chunk and the SSE closes
|
||||
// before records land.
|
||||
if (io === "out" && options?.peekSettled) {
|
||||
settled = await this.#peekIsSettled(s2Stream);
|
||||
if (settled) {
|
||||
waitSeconds = 0;
|
||||
}
|
||||
}
|
||||
|
||||
const s2Response = await this.#streamResponseByName(s2Stream, signal, {
|
||||
...options,
|
||||
timeoutInSeconds: waitSeconds,
|
||||
});
|
||||
|
||||
if (!settled) return s2Response;
|
||||
|
||||
const headers = new Headers(s2Response.headers);
|
||||
headers.set("X-Session-Settled", "true");
|
||||
return new Response(s2Response.body, {
|
||||
status: s2Response.status,
|
||||
statusText: s2Response.statusText,
|
||||
headers,
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Peek the tail of `.out` and return whether the stream is "settled" —
|
||||
* i.e. the most recent non-command record is a `turn-complete` control
|
||||
* record. The agent appends an S2 `trim` command record immediately
|
||||
* after every turn-complete to keep the stream bounded, so we read two
|
||||
* tail records and walk past any trim command to find the
|
||||
* turn-complete underneath.
|
||||
*/
|
||||
async #peekIsSettled(s2Stream: string): Promise<boolean> {
|
||||
const qs = new URLSearchParams();
|
||||
// `tail_offset=2` rewinds two seq positions; `count=2` caps it to
|
||||
// those two records. At steady state these are `[turn-complete, trim]`.
|
||||
// `wait=0` returns immediately with no long-poll.
|
||||
qs.set("tail_offset", "2");
|
||||
qs.set("count", "2");
|
||||
qs.set("wait", "0");
|
||||
|
||||
let res: Response;
|
||||
try {
|
||||
res = await fetch(`${this.baseUrl}/streams/${encodeURIComponent(s2Stream)}/records?${qs}`, {
|
||||
method: "GET",
|
||||
headers: {
|
||||
Authorization: `Bearer ${this.token}`,
|
||||
Accept: "application/json",
|
||||
"S2-Format": "raw",
|
||||
"S2-Basin": this.basin,
|
||||
},
|
||||
});
|
||||
} catch (err) {
|
||||
this.logger.warn("S2 peek last record: fetch failed", { err, stream: s2Stream });
|
||||
return false;
|
||||
}
|
||||
|
||||
if (!res.ok) {
|
||||
// 404: stream has never been written to. 416: range not
|
||||
// satisfiable (empty stream). Both mean "nothing to peek."
|
||||
if (res.status === 404 || res.status === 416) return false;
|
||||
const text = await res.text().catch(() => "");
|
||||
this.logger.warn("S2 peek last record failed", {
|
||||
status: res.status,
|
||||
statusText: res.statusText,
|
||||
text,
|
||||
stream: s2Stream,
|
||||
});
|
||||
return false;
|
||||
}
|
||||
|
||||
let records: Array<{
|
||||
body: string;
|
||||
seq_num: number;
|
||||
timestamp: number;
|
||||
headers?: Array<[string, string]>;
|
||||
}>;
|
||||
try {
|
||||
const json = (await res.json()) as {
|
||||
records?: Array<{
|
||||
body: string;
|
||||
seq_num: number;
|
||||
timestamp: number;
|
||||
headers?: Array<[string, string]>;
|
||||
}>;
|
||||
};
|
||||
records = json.records ?? [];
|
||||
} catch (err) {
|
||||
this.logger.warn("S2 peek last record: parse failed", { err, stream: s2Stream });
|
||||
return false;
|
||||
}
|
||||
|
||||
// Walk from most-recent backward, skipping S2 command records
|
||||
// (`headers[0][0] === ""`). The first non-command record is the
|
||||
// real tail — settled iff its `trigger-control` header is
|
||||
// `turn-complete`.
|
||||
for (let i = records.length - 1; i >= 0; i--) {
|
||||
const record = records[i]!;
|
||||
if (record.headers?.[0]?.[0] === "") {
|
||||
continue;
|
||||
}
|
||||
const controlValue = headerValue(record.headers, "trigger-control");
|
||||
return controlValue === "turn-complete";
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
async #streamResponseByName(
|
||||
s2Stream: string,
|
||||
signal: AbortSignal,
|
||||
options?: StreamResponseOptions
|
||||
): Promise<Response> {
|
||||
const startSeq = this.parseLastEventId(options?.lastEventId);
|
||||
|
||||
this.logger.info(`S2 streaming records from stream`, { stream: s2Stream, startSeq });
|
||||
|
||||
// Request SSE stream from S2 and return it directly
|
||||
const s2Response = await this.s2StreamRecords(s2Stream, {
|
||||
seq_num: startSeq ?? 0,
|
||||
clamp: true,
|
||||
wait: options?.timeoutInSeconds ?? this.s2WaitSeconds, // S2 will keep the connection open and stream new records
|
||||
signal, // Pass abort signal so S2 connection is cleaned up when client disconnects
|
||||
});
|
||||
|
||||
// Return S2's SSE response directly to the client
|
||||
return s2Response;
|
||||
}
|
||||
|
||||
// ---------- Internals: S2 REST ----------
|
||||
private async s2Append(stream: string, body: S2AppendInput): Promise<S2AppendAck> {
|
||||
// POST /v1/streams/{stream}/records (JSON).
|
||||
//
|
||||
// Retries transient failures (network errors and 5xx) up to 3 times with
|
||||
// exponential backoff. Undici's "fetch failed" errors observed locally
|
||||
// are pre-connection (DNS/TCP) so the request never reaches S2, making
|
||||
// retry safe — the alternative is a 500 surfacing to the SDK transport,
|
||||
// which then retries the whole `/in/append` round-trip and pollutes
|
||||
// logs. 4xx are not retried (genuine client errors).
|
||||
const url = `${this.baseUrl}/streams/${encodeURIComponent(stream)}/records`;
|
||||
const init: RequestInit = {
|
||||
method: "POST",
|
||||
headers: {
|
||||
Authorization: `Bearer ${this.token}`,
|
||||
"Content-Type": "application/json",
|
||||
"S2-Format": "raw",
|
||||
"S2-Basin": this.basin,
|
||||
},
|
||||
body: JSON.stringify(body),
|
||||
};
|
||||
|
||||
const maxAttempts = 3;
|
||||
const backoffsMs = [100, 250, 600];
|
||||
let lastError: unknown;
|
||||
|
||||
for (let attempt = 0; attempt < maxAttempts; attempt++) {
|
||||
// The `try` only wraps `fetch` — once we have a Response we handle status
|
||||
// outside the catch, so a 4xx throw can't be swallowed and retried.
|
||||
let res: Response | undefined;
|
||||
try {
|
||||
res = await fetch(url, init);
|
||||
} catch (err) {
|
||||
lastError = err;
|
||||
}
|
||||
|
||||
if (res) {
|
||||
if (res.ok) {
|
||||
return (await res.json()) as S2AppendAck;
|
||||
}
|
||||
const text = await res.text().catch(() => "");
|
||||
const httpError = new Error(`S2 append failed: ${res.status} ${res.statusText} ${text}`);
|
||||
if (res.status >= 400 && res.status < 500) {
|
||||
// 4xx — caller-side problem (auth, malformed body, closed stream).
|
||||
// Retrying won't help.
|
||||
throw httpError;
|
||||
}
|
||||
// 5xx — retryable.
|
||||
lastError = httpError;
|
||||
}
|
||||
|
||||
const isLastAttempt = attempt === maxAttempts - 1;
|
||||
const diagnostics = describeFetchError(lastError);
|
||||
if (isLastAttempt) {
|
||||
this.logger.error("S2 append failed after retries", {
|
||||
stream,
|
||||
attempts: maxAttempts,
|
||||
...diagnostics,
|
||||
});
|
||||
break;
|
||||
}
|
||||
|
||||
this.logger.warn("S2 append transient failure, retrying", {
|
||||
stream,
|
||||
attempt: attempt + 1,
|
||||
nextDelayMs: backoffsMs[attempt],
|
||||
...diagnostics,
|
||||
});
|
||||
await new Promise((resolve) => setTimeout(resolve, backoffsMs[attempt]));
|
||||
}
|
||||
|
||||
throw lastError instanceof Error ? lastError : new Error(String(lastError));
|
||||
}
|
||||
|
||||
private async getS2AccessToken(id: string): Promise<string> {
|
||||
if (!this.cache) {
|
||||
return this.s2IssueAccessToken(id);
|
||||
}
|
||||
|
||||
// Cache key includes basin so per-org basins never collide on
|
||||
// cached tokens, and the ops fingerprint so a scope change in code
|
||||
// (e.g. adding `trim` in #3644) auto-invalidates pre-deploy entries
|
||||
// instead of returning stale tokens for up to 24h.
|
||||
const cacheKey = `${this.basin}:${this.streamPrefix}:${S2_TOKEN_OPS_FINGERPRINT}`;
|
||||
const result = await this.cache.accessToken.swr(cacheKey, async () => {
|
||||
return this.s2IssueAccessToken(id);
|
||||
});
|
||||
|
||||
if (!result.val) {
|
||||
throw new Error("Failed to get S2 access token");
|
||||
}
|
||||
|
||||
return result.val;
|
||||
}
|
||||
|
||||
private async s2IssueAccessToken(id: string): Promise<string> {
|
||||
// POST /v1/access-tokens
|
||||
const res = await fetch(`${this.accountUrl}/access-tokens`, {
|
||||
method: "POST",
|
||||
headers: {
|
||||
Authorization: `Bearer ${this.token}`,
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
body: JSON.stringify({
|
||||
id,
|
||||
scope: {
|
||||
basins: {
|
||||
exact: this.basin,
|
||||
},
|
||||
ops: [...S2_TOKEN_OPS],
|
||||
streams: {
|
||||
prefix: this.streamPrefix,
|
||||
},
|
||||
},
|
||||
expires_at: new Date(Date.now() + this.accessTokenExpirationInMs).toISOString(),
|
||||
auto_prefix_streams: true,
|
||||
}),
|
||||
});
|
||||
|
||||
if (!res.ok) {
|
||||
const text = await res.text().catch(() => "");
|
||||
throw new Error(`S2 issue access token failed: ${res.status} ${res.statusText} ${text}`);
|
||||
}
|
||||
const data = (await res.json()) as S2IssueAccessTokenResponse;
|
||||
return data.access_token;
|
||||
}
|
||||
|
||||
private async s2StreamRecords(
|
||||
stream: string,
|
||||
opts: {
|
||||
seq_num?: number;
|
||||
clamp?: boolean;
|
||||
wait?: number;
|
||||
signal?: AbortSignal;
|
||||
}
|
||||
): Promise<Response> {
|
||||
// GET /v1/streams/{stream}/records with Accept: text/event-stream for SSE streaming
|
||||
const qs = new URLSearchParams();
|
||||
if (opts.seq_num != null) qs.set("seq_num", String(opts.seq_num));
|
||||
if (opts.clamp != null) qs.set("clamp", String(opts.clamp));
|
||||
if (opts.wait != null) qs.set("wait", String(opts.wait));
|
||||
|
||||
const res = await fetch(`${this.baseUrl}/streams/${encodeURIComponent(stream)}/records?${qs}`, {
|
||||
method: "GET",
|
||||
headers: {
|
||||
Authorization: `Bearer ${this.token}`,
|
||||
Accept: "text/event-stream",
|
||||
"S2-Format": "raw",
|
||||
"S2-Basin": this.basin,
|
||||
},
|
||||
signal: opts.signal,
|
||||
});
|
||||
|
||||
if (!res.ok) {
|
||||
const text = await res.text().catch(() => "");
|
||||
throw new Error(`S2 stream failed: ${res.status} ${res.statusText} ${text}`);
|
||||
}
|
||||
|
||||
const headers = new Headers(res.headers);
|
||||
headers.set("X-Stream-Version", "v2");
|
||||
headers.set("Access-Control-Expose-Headers", "*");
|
||||
|
||||
return new Response(res.body, {
|
||||
headers,
|
||||
status: res.status,
|
||||
statusText: res.statusText,
|
||||
});
|
||||
}
|
||||
|
||||
private parseLastEventId(lastEventId?: string): number | undefined {
|
||||
if (!lastEventId) return undefined;
|
||||
// tolerate formats like "1699999999999-5" (take leading digits)
|
||||
const digits = lastEventId.split("-")[0];
|
||||
const n = Number(digits);
|
||||
return Number.isFinite(n) && n >= 0 ? n + 1 : undefined;
|
||||
}
|
||||
}
|
||||
|
||||
// Pulls the underlying network error out of undici's generic "fetch failed".
|
||||
// undici sets `error.cause` to either a SystemError-shaped object with `code`
|
||||
// (e.g. `ECONNRESET`, `UND_ERR_SOCKET`, `ETIMEDOUT`), `errno`, and `syscall`,
|
||||
// or — for happy-eyeballs / multi-address connect attempts — an
|
||||
// `AggregateError` whose `errors[]` each carry their own code. Surfacing
|
||||
// those tells us whether failures are pre-connection (DNS / TCP), mid-stream
|
||||
// socket resets, or genuine S2 server errors.
|
||||
function describeFetchError(err: unknown): Record<string, unknown> {
|
||||
if (!(err instanceof Error)) {
|
||||
return { error: String(err) };
|
||||
}
|
||||
const out: Record<string, unknown> = {
|
||||
error: err.message,
|
||||
name: err.name,
|
||||
};
|
||||
const cause = (err as { cause?: unknown }).cause;
|
||||
if (cause && typeof cause === "object") {
|
||||
const c = cause as Record<string, unknown>;
|
||||
if (typeof c.code === "string") out.causeCode = c.code;
|
||||
if (typeof c.errno === "number" || typeof c.errno === "string") out.causeErrno = c.errno;
|
||||
if (typeof c.syscall === "string") out.causeSyscall = c.syscall;
|
||||
if (typeof c.message === "string") out.causeMessage = c.message;
|
||||
if (Array.isArray(c.errors)) {
|
||||
out.causeErrors = c.errors
|
||||
.filter((e: unknown): e is Error => e instanceof Error)
|
||||
.map((e) => ({
|
||||
message: e.message,
|
||||
code: (e as { code?: unknown }).code,
|
||||
syscall: (e as { syscall?: unknown }).syscall,
|
||||
address: (e as { address?: unknown }).address,
|
||||
port: (e as { port?: unknown }).port,
|
||||
}));
|
||||
}
|
||||
}
|
||||
return out;
|
||||
}
|
||||
@@ -0,0 +1,525 @@
|
||||
import type { Session, TaskRunStatus } from "@trigger.dev/database";
|
||||
import { SessionTriggerConfig as SessionTriggerConfigZod } from "@trigger.dev/core/v3";
|
||||
import type { z } from "zod";
|
||||
import { prisma, $replica } from "~/db.server";
|
||||
import { runStore } from "~/v3/runStore.server";
|
||||
import type { AuthenticatedEnvironment } from "~/services/apiAuth.server";
|
||||
import { logger } from "~/services/logger.server";
|
||||
import { CancelTaskRunService } from "~/v3/services/cancelTaskRun.server";
|
||||
import { TriggerTaskService } from "~/v3/services/triggerTask.server";
|
||||
import { isFinalRunStatus } from "~/v3/taskStatus";
|
||||
|
||||
/**
|
||||
* Schema for `Session.triggerConfig` (stored as JSONB). The wire-format
|
||||
* source of truth lives in `@trigger.dev/core/v3` as `SessionTriggerConfig`;
|
||||
* we re-export it here for the trigger machinery to validate on read.
|
||||
*
|
||||
* `basePayload` carries the customer's wire payload (for chat.agent:
|
||||
* `{ chatId, ...clientData, idleTimeoutInSeconds? }`). Runtime fields
|
||||
* specific to a particular trigger (e.g. `trigger: "trigger" | "preload"`,
|
||||
* an `isContinuation` flag) come in via the `payloadOverrides` argument
|
||||
* to `ensureRunForSession` and shallow-merge on top of `basePayload`.
|
||||
*/
|
||||
export const SessionTriggerConfigSchema = SessionTriggerConfigZod;
|
||||
|
||||
export type SessionTriggerConfig = z.infer<typeof SessionTriggerConfigSchema>;
|
||||
|
||||
export type EnsureRunReason = "initial" | "continuation" | "upgrade" | "manual";
|
||||
|
||||
/**
|
||||
* Hard cap on how many times `ensureRunForSession` will recurse on the
|
||||
* pathological "we lost the claim race AND the winner's run was already
|
||||
* terminal" path. In practice progress through the run engine bounds
|
||||
* this, but a misconfigured task that crashes before it can be dequeued
|
||||
* could otherwise loop without limit. After this many attempts we
|
||||
* surface `SessionRunManagerError` so the caller can 5xx instead of
|
||||
* blowing the stack.
|
||||
*/
|
||||
const ENSURE_RUN_FOR_SESSION_MAX_ATTEMPTS = 3;
|
||||
|
||||
type EnsureRunForSessionParams = {
|
||||
/**
|
||||
* Session row to operate on. Caller is responsible for the env match —
|
||||
* we don't re-check `runtimeEnvironmentId` against `environment.id`.
|
||||
*
|
||||
* `friendlyId` is used to pre-populate `payload.sessionId` on the new
|
||||
* run so the agent's `chat.agent` boot path can attach to `session.in/.out`
|
||||
* without a control-plane round-trip. `currentRunId` is also forwarded
|
||||
* as `payload.previousRunId` (with `continuation: true`) when the prior
|
||||
* run is dead, so the agent's boot gate triggers snapshot.read + replay
|
||||
* instead of treating the run as a fresh chat.
|
||||
*/
|
||||
session: Pick<
|
||||
Session,
|
||||
"id" | "friendlyId" | "taskIdentifier" | "triggerConfig" | "currentRunId" | "currentRunVersion"
|
||||
>;
|
||||
environment: AuthenticatedEnvironment;
|
||||
reason: EnsureRunReason;
|
||||
/**
|
||||
* Shallow-merged on top of `triggerConfig.basePayload`. Runtime fields
|
||||
* only — caller-controlled data that varies per trigger (`trigger:
|
||||
* "preload"` vs `"trigger"`, etc).
|
||||
*/
|
||||
payloadOverrides?: Record<string, unknown>;
|
||||
/**
|
||||
* @internal Recursion-guard counter for the lost-claim-race retry path.
|
||||
* Public callers should leave this unset; the function recurses with
|
||||
* an incremented value on the pathological "winner's run was already
|
||||
* terminal" branch and throws once it exceeds
|
||||
* {@link ENSURE_RUN_FOR_SESSION_MAX_ATTEMPTS}.
|
||||
*/
|
||||
_attempt?: number;
|
||||
};
|
||||
|
||||
export type EnsureRunResult = {
|
||||
runId: string;
|
||||
/** True if this call triggered a fresh run; false if it reused an alive existing one. */
|
||||
triggered: boolean;
|
||||
};
|
||||
|
||||
/**
|
||||
* Idempotently make sure the session has a live run.
|
||||
*
|
||||
* Algorithm:
|
||||
* 1. If `currentRunId` is set, probe its status. Alive → return as-is.
|
||||
* 2. Trigger a new run upfront (cheap to cancel if we lose the race).
|
||||
* 3. Atomic claim via `updateMany` keyed on `currentRunVersion`.
|
||||
* - Won: return new runId, record SessionRun audit row.
|
||||
* - Lost: cancel our triggered run, re-read session, reuse winner's
|
||||
* run if alive. If pathological (winner's run already terminal),
|
||||
* recurse.
|
||||
*
|
||||
* No DB lock is held across the trigger call. Wasted-trigger window is
|
||||
* the rare multi-tab race on a dead run; cancel cost is negligible and
|
||||
* the run-engine handles it gracefully.
|
||||
*/
|
||||
export async function ensureRunForSession(
|
||||
params: EnsureRunForSessionParams
|
||||
): Promise<EnsureRunResult> {
|
||||
const { session, environment, reason, payloadOverrides, _attempt = 1 } = params;
|
||||
|
||||
if (_attempt > ENSURE_RUN_FOR_SESSION_MAX_ATTEMPTS) {
|
||||
throw new SessionRunManagerError(
|
||||
`ensureRunForSession exceeded ${ENSURE_RUN_FOR_SESSION_MAX_ATTEMPTS} attempts for session ${session.id} — every triggered run reached a terminal state before claim could resolve`
|
||||
);
|
||||
}
|
||||
|
||||
// 1. Probe currentRunId.
|
||||
let priorDeadRunFriendlyId: string | undefined;
|
||||
if (session.currentRunId) {
|
||||
let probe = await getRunStatusAndFriendlyId(session.currentRunId);
|
||||
if (!probe) {
|
||||
// Replica miss on a row we just observed via `currentRunId` — the
|
||||
// run was likely triggered moments ago and hasn't replicated yet.
|
||||
// Re-probe the writer BEFORE deciding liveness: treating a lagging
|
||||
// replica as "row vanished" double-triggers the session (a fast
|
||||
// first append after session create races the replica apply delay
|
||||
// and spawns a second live run consuming the same `.in`).
|
||||
probe = await runStore.findRun(
|
||||
{ id: session.currentRunId },
|
||||
{ select: { status: true, friendlyId: true } },
|
||||
prisma
|
||||
);
|
||||
}
|
||||
if (probe && !isFinalRunStatus(probe.status)) {
|
||||
return { runId: session.currentRunId, triggered: false };
|
||||
}
|
||||
// Either the row vanished on the writer too (probe null) or its status
|
||||
// is final. Either way the prior run isn't going to consume new
|
||||
// appends — but the session may still hold conversation state on
|
||||
// `session.out` and an S3 snapshot keyed on `session.friendlyId`.
|
||||
// Forward the prior run's public-form id (friendlyId — same shape as
|
||||
// `ctx.run.id`) to the agent as `previousRunId` so its boot gate flips
|
||||
// `couldHavePriorState` and replays the persisted state instead of
|
||||
// treating this as a fresh chat. See `chat.agent`'s boot orchestration
|
||||
// in `packages/trigger-sdk/src/v3/ai.ts`.
|
||||
priorDeadRunFriendlyId = probe?.friendlyId ?? session.currentRunId;
|
||||
}
|
||||
|
||||
// 2. Validate config + trigger upfront. Continuation overrides
|
||||
// (`continuation`, `previousRunId`) are derived from session state above
|
||||
// and merged AFTER caller-supplied overrides — caller can't accidentally
|
||||
// unset them on a session that has had a prior run, but can still
|
||||
// override `trigger`/`metadata` etc. `sessionId` is always set so the
|
||||
// agent doesn't need a control-plane round-trip to look up the session
|
||||
// friendlyId from `payload.chatId`.
|
||||
// Continuation overrides strip the basePayload's first-run-only fields
|
||||
// so a continuation run doesn't inherit a stale boot payload. The Session
|
||||
// row's `triggerConfig.basePayload` is captured at create-time and used
|
||||
// verbatim for every Run we trigger; if the customer included `message`
|
||||
// / `messages` / `trigger: "submit-message"` to make the FIRST run boot
|
||||
// straight into a first turn (via `chat.createStartSessionAction`), those
|
||||
// values stick around and get replayed on every continuation. With
|
||||
// `continuation: true` and `message`/`messages` cleared, the SDK boot
|
||||
// path enters its continuation-wait branch and waits for the next
|
||||
// session.in record before running a turn.
|
||||
const continuationOverrides: Record<string, unknown> = {
|
||||
sessionId: session.friendlyId,
|
||||
...(priorDeadRunFriendlyId !== undefined
|
||||
? {
|
||||
continuation: true,
|
||||
previousRunId: priorDeadRunFriendlyId,
|
||||
// Clear sticky boot-payload fields so the new run waits for the
|
||||
// next session.in record instead of re-processing whatever was
|
||||
// in the original `createStartSessionAction({ basePayload })`.
|
||||
message: undefined,
|
||||
messages: undefined,
|
||||
trigger: undefined,
|
||||
}
|
||||
: {}),
|
||||
};
|
||||
const mergedPayloadOverrides: Record<string, unknown> = {
|
||||
...(payloadOverrides ?? {}),
|
||||
...continuationOverrides,
|
||||
};
|
||||
|
||||
const config = SessionTriggerConfigSchema.parse(session.triggerConfig);
|
||||
const triggered = await triggerSessionRun({
|
||||
session,
|
||||
config,
|
||||
environment,
|
||||
payloadOverrides: mergedPayloadOverrides,
|
||||
});
|
||||
|
||||
// 3. Try to claim the slot atomically.
|
||||
const claim = await prisma.session.updateMany({
|
||||
where: {
|
||||
id: session.id,
|
||||
currentRunVersion: session.currentRunVersion,
|
||||
},
|
||||
data: {
|
||||
currentRunId: triggered.id,
|
||||
currentRunVersion: { increment: 1 },
|
||||
},
|
||||
});
|
||||
|
||||
if (claim.count === 1) {
|
||||
// Won. Audit the SessionRun. Best-effort — failure here doesn't
|
||||
// invalidate the live run, just leaves a missing audit row.
|
||||
prisma.sessionRun
|
||||
.create({
|
||||
data: { sessionId: session.id, runId: triggered.id, reason },
|
||||
})
|
||||
.catch((error) => {
|
||||
logger.warn("Failed to record SessionRun audit row", {
|
||||
sessionId: session.id,
|
||||
runId: triggered.id,
|
||||
reason,
|
||||
error,
|
||||
});
|
||||
});
|
||||
|
||||
return { runId: triggered.id, triggered: true };
|
||||
}
|
||||
|
||||
// 4. Lost the race. Cancel our triggered run; reuse the winner's.
|
||||
cancelLostRaceRun(triggered.id, environment).catch((error) => {
|
||||
logger.warn("Failed to cancel lost-race session run", {
|
||||
sessionId: session.id,
|
||||
runId: triggered.id,
|
||||
error,
|
||||
});
|
||||
});
|
||||
|
||||
// Read-after-write: the winner just wrote `currentRunId` /
|
||||
// `currentRunVersion` on the writer. Reading from `$replica` could
|
||||
// return pre-race state and cause us to recurse with the same stale
|
||||
// version, losing the next claim, until we exhaust max attempts.
|
||||
const fresh = await prisma.session.findFirst({
|
||||
where: { id: session.id },
|
||||
select: {
|
||||
id: true,
|
||||
friendlyId: true,
|
||||
taskIdentifier: true,
|
||||
triggerConfig: true,
|
||||
currentRunId: true,
|
||||
currentRunVersion: true,
|
||||
},
|
||||
});
|
||||
|
||||
if (!fresh) {
|
||||
// Session vanished mid-flight. Surface as an error — caller decides
|
||||
// whether to 404 or retry.
|
||||
throw new SessionRunManagerError(`Session ${session.id} not found after lost claim race`);
|
||||
}
|
||||
|
||||
if (fresh.currentRunId) {
|
||||
// Same read-after-write reason as the `fresh` reload above: the winner
|
||||
// just wrote `currentRunId` on the writer, so probe the writer too —
|
||||
// the replica may not have the run row yet, and a missed probe forces
|
||||
// another trigger+recurse until `ENSURE_RUN_FOR_SESSION_MAX_ATTEMPTS`.
|
||||
const probe = await runStore.findRun(
|
||||
{ id: fresh.currentRunId },
|
||||
{ select: { status: true, friendlyId: true } },
|
||||
prisma
|
||||
);
|
||||
if (probe && !isFinalRunStatus(probe.status)) {
|
||||
return { runId: fresh.currentRunId, triggered: false };
|
||||
}
|
||||
}
|
||||
|
||||
// Pathological: winner's run already terminal. Recurse with the fresh
|
||||
// version. Bounded by `ENSURE_RUN_FOR_SESSION_MAX_ATTEMPTS` so a task
|
||||
// that always crashes before being dequeued surfaces as an error
|
||||
// instead of a stack overflow.
|
||||
return ensureRunForSession({
|
||||
session: fresh,
|
||||
environment,
|
||||
reason,
|
||||
payloadOverrides,
|
||||
_attempt: _attempt + 1,
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Trigger a single run for a session. Builds `TriggerTaskRequestBody`
|
||||
* by shallow-merging `payloadOverrides` over `config.basePayload` and
|
||||
* threading `config`'s machine/queue/tags through the trigger options.
|
||||
*/
|
||||
async function triggerSessionRun(params: {
|
||||
session: Pick<Session, "id" | "taskIdentifier">;
|
||||
config: SessionTriggerConfig;
|
||||
environment: AuthenticatedEnvironment;
|
||||
payloadOverrides?: Record<string, unknown>;
|
||||
}): Promise<{ id: string; friendlyId: string }> {
|
||||
const { session, config, environment, payloadOverrides } = params;
|
||||
|
||||
const payload = {
|
||||
...config.basePayload,
|
||||
...(config.idleTimeoutInSeconds !== undefined
|
||||
? { idleTimeoutInSeconds: config.idleTimeoutInSeconds }
|
||||
: {}),
|
||||
...(payloadOverrides ?? {}),
|
||||
};
|
||||
|
||||
const body = {
|
||||
payload,
|
||||
context: {},
|
||||
options: {
|
||||
...(config.machine ? { machine: config.machine as never } : {}),
|
||||
...(config.queue ? { queue: { name: config.queue } } : {}),
|
||||
...(config.tags ? { tags: config.tags } : {}),
|
||||
...(config.maxAttempts !== undefined ? { maxAttempts: config.maxAttempts } : {}),
|
||||
...(config.maxDuration !== undefined ? { maxDuration: config.maxDuration } : {}),
|
||||
...(config.lockToVersion ? { lockToVersion: config.lockToVersion } : {}),
|
||||
...(config.region ? { region: config.region } : {}),
|
||||
},
|
||||
};
|
||||
|
||||
const service = new TriggerTaskService();
|
||||
const result = await service.call(session.taskIdentifier, environment, body, {
|
||||
triggerSource: "session",
|
||||
triggerAction: "trigger",
|
||||
});
|
||||
|
||||
if (!result) {
|
||||
throw new SessionRunManagerError(
|
||||
`TriggerTaskService returned no result for taskIdentifier=${session.taskIdentifier}`
|
||||
);
|
||||
}
|
||||
|
||||
return { id: result.run.id, friendlyId: result.run.friendlyId };
|
||||
}
|
||||
|
||||
type SwapSessionRunParams = {
|
||||
/**
|
||||
* Session row to swap. `friendlyId` is forwarded as `payload.sessionId`
|
||||
* on the new run so the agent attaches to `session.in/.out` without a
|
||||
* control-plane round-trip (same convention as
|
||||
* {@link EnsureRunForSessionParams}).
|
||||
*/
|
||||
session: Pick<
|
||||
Session,
|
||||
"id" | "friendlyId" | "taskIdentifier" | "triggerConfig" | "currentRunId" | "currentRunVersion"
|
||||
>;
|
||||
/**
|
||||
* The run requesting the swap. Optimistic claim requires
|
||||
* `Session.currentRunId === callingRunId` so the swap can't clobber
|
||||
* a run triggered out-of-band (e.g. a parallel `.in/append` probe
|
||||
* that already replaced the dead run).
|
||||
*
|
||||
* Also forwarded as `payload.previousRunId` on the new run alongside
|
||||
* `continuation: true` — every swap is a continuation by construction
|
||||
* (`chat.requestUpgrade` / `chat.endRun` deliberately hand off prior
|
||||
* conversation state to a new run), so the agent's boot gate flips
|
||||
* `couldHavePriorState` and replays the snapshot + session.out tail.
|
||||
*/
|
||||
callingRunId: string;
|
||||
environment: AuthenticatedEnvironment;
|
||||
reason: EnsureRunReason;
|
||||
payloadOverrides?: Record<string, unknown>;
|
||||
};
|
||||
|
||||
export type SwapSessionRunResult = {
|
||||
/** runId of the newly-triggered run that has taken over the session. */
|
||||
runId: string;
|
||||
/**
|
||||
* False when the swap was preempted (currentRunId is no longer the
|
||||
* calling run). The caller should treat this as "someone else
|
||||
* already moved on" — exit cleanly without expecting to drive the
|
||||
* next run.
|
||||
*/
|
||||
swapped: boolean;
|
||||
};
|
||||
|
||||
/**
|
||||
* Force-swap the session to a freshly-triggered run, regardless of
|
||||
* whether the current run is alive. Called by `end-and-continue` when
|
||||
* the running agent wants a clean handoff (typically version upgrade).
|
||||
*
|
||||
* Differs from `ensureRunForSession`: never reuses the current run.
|
||||
* The optimistic claim is keyed on `currentRunId === callingRunId`, so
|
||||
* a parallel append-time probe that already swapped to a different
|
||||
* run wins the race and `swapped: false` is surfaced.
|
||||
*/
|
||||
export async function swapSessionRun(params: SwapSessionRunParams): Promise<SwapSessionRunResult> {
|
||||
const { session, callingRunId, environment, reason, payloadOverrides } = params;
|
||||
|
||||
// `callingRunId` is the internal cuid (`Session.currentRunId` stores
|
||||
// cuid; the route handler resolves the wire's friendlyId before passing
|
||||
// it here). The agent's `previousRunId` is customer-visible and must
|
||||
// match the public `run_*` form exposed via `ctx.run.id` — resolve
|
||||
// before forwarding.
|
||||
const callingRunFriendlyId = await resolveRunFriendlyId(callingRunId);
|
||||
|
||||
// Continuation overrides — unconditionally set on swap. Unlike
|
||||
// `ensureRunForSession`, there's no dead-run-detection branch here:
|
||||
// every swap is a deliberate handoff from `callingRunId` (which owned
|
||||
// prior conversation state) to a fresh run. Merged AFTER caller-supplied
|
||||
// overrides so a caller can't accidentally unset them.
|
||||
//
|
||||
// Sticky boot-payload fields (`message` / `messages` / `trigger`) are
|
||||
// cleared here for the same reason as in `ensureRunForSession`: the
|
||||
// Session's basePayload is captured at create-time and replays on every
|
||||
// continuation if not stripped. See the comment in `ensureRunForSession`.
|
||||
const mergedPayloadOverrides: Record<string, unknown> = {
|
||||
...(payloadOverrides ?? {}),
|
||||
sessionId: session.friendlyId,
|
||||
continuation: true,
|
||||
previousRunId: callingRunFriendlyId,
|
||||
message: undefined,
|
||||
messages: undefined,
|
||||
trigger: undefined,
|
||||
};
|
||||
|
||||
const config = SessionTriggerConfigSchema.parse(session.triggerConfig);
|
||||
const triggered = await triggerSessionRun({
|
||||
session,
|
||||
config,
|
||||
environment,
|
||||
payloadOverrides: mergedPayloadOverrides,
|
||||
});
|
||||
|
||||
const claim = await prisma.session.updateMany({
|
||||
where: {
|
||||
id: session.id,
|
||||
currentRunId: callingRunId,
|
||||
currentRunVersion: session.currentRunVersion,
|
||||
},
|
||||
data: {
|
||||
currentRunId: triggered.id,
|
||||
currentRunVersion: { increment: 1 },
|
||||
},
|
||||
});
|
||||
|
||||
if (claim.count === 1) {
|
||||
prisma.sessionRun
|
||||
.create({
|
||||
data: { sessionId: session.id, runId: triggered.id, reason },
|
||||
})
|
||||
.catch((error) => {
|
||||
logger.warn("Failed to record SessionRun audit row", {
|
||||
sessionId: session.id,
|
||||
runId: triggered.id,
|
||||
reason,
|
||||
error,
|
||||
});
|
||||
});
|
||||
return { runId: triggered.id, swapped: true };
|
||||
}
|
||||
|
||||
// Lost the race — someone else already swapped to a new run. Cancel
|
||||
// ours, surface the existing winner.
|
||||
cancelLostRaceRun(triggered.id, environment).catch((error) => {
|
||||
logger.warn("Failed to cancel preempted swap run", {
|
||||
sessionId: session.id,
|
||||
runId: triggered.id,
|
||||
error,
|
||||
});
|
||||
});
|
||||
|
||||
// Read-after-write: the winner's swap was just committed on the
|
||||
// writer. A replica read could return the pre-swap `currentRunId`
|
||||
// (often `callingRunId` itself), which would tell the caller it is
|
||||
// still the canonical run when in fact a different run has taken
|
||||
// over.
|
||||
const fresh = await prisma.session.findFirst({
|
||||
where: { id: session.id },
|
||||
select: { currentRunId: true },
|
||||
});
|
||||
|
||||
// Mirror `ensureRunForSession`'s "session vanished" branch: if we
|
||||
// can't find the row (or it has no current run) on the writer right
|
||||
// after losing the race, surface as an error rather than handing back
|
||||
// `callingRunId` with `swapped: false` — that would tell the caller
|
||||
// it's still the canonical run when in fact we don't know who is.
|
||||
if (!fresh?.currentRunId) {
|
||||
throw new SessionRunManagerError(
|
||||
`Session ${session.id} has no currentRunId after preempted swap`
|
||||
);
|
||||
}
|
||||
|
||||
return {
|
||||
runId: fresh.currentRunId,
|
||||
swapped: false,
|
||||
};
|
||||
}
|
||||
|
||||
async function getRunStatusAndFriendlyId(
|
||||
runId: string
|
||||
): Promise<{ status: TaskRunStatus; friendlyId: string } | null> {
|
||||
// Use the read replica — this is a hot-path probe and stale-by-ms is
|
||||
// fine. The append handler re-checks if it ends up reusing the runId.
|
||||
// `friendlyId` is fetched alongside `status` so the dead-run-detection
|
||||
// branch in `ensureRunForSession` can forward the public-form id as
|
||||
// `payload.previousRunId` without a second read. `Session.currentRunId`
|
||||
// stores the internal cuid; the agent's wire / customer hooks expose
|
||||
// the friendlyId via `ctx.run.id`, so consistency matters.
|
||||
const row = await runStore.findRun(
|
||||
{ id: runId },
|
||||
{ select: { status: true, friendlyId: true } },
|
||||
$replica
|
||||
);
|
||||
return row ?? null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Resolve a TaskRun cuid to its friendlyId. Used by `swapSessionRun` to
|
||||
* forward the calling run's public-form id as `payload.previousRunId` on
|
||||
* the new run. Falls back to the cuid on lookup miss so the swap doesn't
|
||||
* fail just because the read replica hasn't caught up — the agent only
|
||||
* uses `previousRunId` for customer-visible bookkeeping (e.g.
|
||||
* `runs.retrieve(previousRunId)`), so a stale-but-non-null value is
|
||||
* acceptable degraded behavior.
|
||||
*/
|
||||
async function resolveRunFriendlyId(runId: string): Promise<string> {
|
||||
const row = await runStore.findRun({ id: runId }, { select: { friendlyId: true } }, $replica);
|
||||
return row?.friendlyId ?? runId;
|
||||
}
|
||||
|
||||
async function cancelLostRaceRun(
|
||||
runId: string,
|
||||
environment: AuthenticatedEnvironment
|
||||
): Promise<void> {
|
||||
const service = new CancelTaskRunService();
|
||||
// Read-after-write: the run was just triggered on the writer, so go
|
||||
// through `prisma`. A `$replica` miss here would silently no-op the
|
||||
// cancel and leak an orphan run that no session is going to claim.
|
||||
const run = await runStore.findRun({ id: runId }, prisma);
|
||||
if (!run) return;
|
||||
await service.call(run, { reason: "Lost session-run claim race" });
|
||||
}
|
||||
|
||||
export class SessionRunManagerError extends Error {
|
||||
readonly name = "SessionRunManagerError";
|
||||
}
|
||||
@@ -0,0 +1,194 @@
|
||||
import type { PrismaClient, Session } from "@trigger.dev/database";
|
||||
import type { SessionItem } from "@trigger.dev/core/v3";
|
||||
import type { RunStore } from "@internal/run-store";
|
||||
import { $replica, prisma } from "~/db.server";
|
||||
import { runStore as defaultRunStore } from "~/v3/runStore.server";
|
||||
|
||||
/**
|
||||
* Prefix that {@link SessionId.generate} attaches to every Session friendlyId.
|
||||
* Used to distinguish friendlyId lookups (`session_abc...`) from externalId
|
||||
* lookups on the public `GET /api/v1/sessions/:session` route.
|
||||
*/
|
||||
const SESSION_FRIENDLY_ID_PREFIX = "session_";
|
||||
|
||||
/**
|
||||
* Resolve a session from a URL path parameter that may contain either a
|
||||
* friendlyId (`session_abc...`) or a user-supplied externalId.
|
||||
*
|
||||
* Disambiguated by prefix: values starting with `session_` are treated as
|
||||
* friendlyIds, anything else is looked up against `externalId` scoped to
|
||||
* the caller's environment.
|
||||
*/
|
||||
// CONTROL-PLANE: `Session` lives on the control-plane DB; these reads are NOT
|
||||
// routed to run-ops read-through — only the `TaskRun` currentRunId resolves in
|
||||
// this file are run-ops read-through routed.
|
||||
export async function resolveSessionByIdOrExternalId(
|
||||
prisma: Pick<PrismaClient, "session">,
|
||||
runtimeEnvironmentId: string,
|
||||
idOrExternalId: string
|
||||
): Promise<Session | null> {
|
||||
if (isSessionFriendlyIdForm(idOrExternalId)) {
|
||||
return prisma.session.findFirst({
|
||||
where: { friendlyId: idOrExternalId, runtimeEnvironmentId },
|
||||
});
|
||||
}
|
||||
|
||||
// `findFirst` rather than `findUnique` per the repo rule — `findUnique`'s
|
||||
// implicit DataLoader has open correctness bugs in Prisma 6.x that bite
|
||||
// hot-path lookups exactly like this one.
|
||||
return prisma.session.findFirst({
|
||||
where: { runtimeEnvironmentId, externalId: idOrExternalId },
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Replica-first session resolution with a writer fallback on miss. For the
|
||||
* hot realtime routes (append / SSE subscribe / end-and-continue): a fresh
|
||||
* session's first append or subscribe can arrive inside the replica's apply
|
||||
* window, and a bare replica miss there surfaces as a 404 (or a subscribe
|
||||
* that never finds its session) for a session that exists on the writer.
|
||||
*/
|
||||
export async function resolveSessionWithWriterFallback(
|
||||
runtimeEnvironmentId: string,
|
||||
idOrExternalId: string
|
||||
): Promise<Session | null> {
|
||||
const row = await resolveSessionByIdOrExternalId($replica, runtimeEnvironmentId, idOrExternalId);
|
||||
if (row) return row;
|
||||
return resolveSessionByIdOrExternalId(prisma, runtimeEnvironmentId, idOrExternalId);
|
||||
}
|
||||
|
||||
/** True for `session_*` friendlyId form, false for everything else. */
|
||||
export function isSessionFriendlyIdForm(value: string): boolean {
|
||||
return value.startsWith(SESSION_FRIENDLY_ID_PREFIX);
|
||||
}
|
||||
|
||||
/**
|
||||
* Canonicalise the addressing key used for everything stream-level: the
|
||||
* S2 stream path and the run-engine waitpoint cache key. `chat.agent`
|
||||
* and the rest of the operational surface always pass `externalId`, but
|
||||
* a public-API caller may legitimately address by `friendlyId` — and a
|
||||
* session created without an `externalId` only has a friendlyId at all.
|
||||
*
|
||||
* Rule:
|
||||
* - If we have a Session row, the canonical key is `externalId` if
|
||||
* set, else `friendlyId`. This way two callers addressing the same
|
||||
* row via different forms always converge to the same S2 stream.
|
||||
* - If we have no row (yet — chat.agent's transport may subscribe
|
||||
* before the agent's bind-time upsert lands), the canonical key is
|
||||
* whatever the URL had. Operationally that's always an externalId.
|
||||
* Friendlyid-form callers without a matching row are rejected by
|
||||
* the route handler before this is reached.
|
||||
*/
|
||||
export function canonicalSessionAddressingKey(row: Session | null, paramSession: string): string {
|
||||
if (row) {
|
||||
return row.externalId ?? row.friendlyId;
|
||||
}
|
||||
return paramSession;
|
||||
}
|
||||
|
||||
/**
|
||||
* Convert a Prisma `Session` row to the public {@link SessionItem} wire format.
|
||||
* Strips internal columns (project/environment/organization ids) and narrows
|
||||
* the `metadata` JSON to a record.
|
||||
*
|
||||
* Note: `currentRunId` is left as-is — Prisma stores the internal run id
|
||||
* (cuid), but `SessionItem.currentRunId` is the *friendly* form. Routes
|
||||
* that emit `SessionItem`s must translate: single-row endpoints via
|
||||
* {@link serializeSessionWithFriendlyRunId}, list endpoints via the
|
||||
* batched {@link serializeSessionsWithFriendlyRunIds}. Never put this
|
||||
* raw form on the wire directly.
|
||||
*/
|
||||
export function serializeSession(session: Session): SessionItem {
|
||||
return {
|
||||
id: session.friendlyId,
|
||||
externalId: session.externalId,
|
||||
type: session.type,
|
||||
taskIdentifier: session.taskIdentifier,
|
||||
triggerConfig: session.triggerConfig as SessionItem["triggerConfig"],
|
||||
currentRunId: session.currentRunId,
|
||||
tags: session.tags,
|
||||
metadata: (session.metadata ?? null) as SessionItem["metadata"],
|
||||
closedAt: session.closedAt,
|
||||
closedReason: session.closedReason,
|
||||
expiresAt: session.expiresAt,
|
||||
createdAt: session.createdAt,
|
||||
updatedAt: session.updatedAt,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Same as {@link serializeSession} but resolves `currentRunId` from the
|
||||
* internal cuid to the public `run_*` friendlyId via a TaskRun lookup.
|
||||
* Single-row endpoints (`POST/GET/PATCH/close /api/v1/sessions/:s`) use
|
||||
* this so the wire-side `currentRunId` is consistent with the rest of
|
||||
* the public API (which only accepts friendlyIds for run lookups).
|
||||
*
|
||||
* Skips the lookup when `currentRunId` is null.
|
||||
*
|
||||
* Resolves `currentRunId` -> `friendlyId` through `runStore.findRun` so a
|
||||
* run-ops id (NEW-DB) session run resolves from its owning store rather than the
|
||||
* control-plane replica. Mirrors `sessionRunManager.server.ts`.
|
||||
* Tenant-scoped because `Session.currentRunId` is a no-FK pointer.
|
||||
*/
|
||||
export async function serializeSessionWithFriendlyRunId(
|
||||
session: Session,
|
||||
runStore: RunStore = defaultRunStore
|
||||
): Promise<SessionItem> {
|
||||
const base = serializeSession(session);
|
||||
if (!session.currentRunId) return base;
|
||||
|
||||
const run = await runStore.findRun(
|
||||
{
|
||||
id: session.currentRunId,
|
||||
projectId: session.projectId,
|
||||
runtimeEnvironmentId: session.runtimeEnvironmentId,
|
||||
},
|
||||
{ select: { friendlyId: true } }
|
||||
);
|
||||
|
||||
return {
|
||||
...base,
|
||||
currentRunId: run?.friendlyId ?? null,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Batched form of {@link serializeSessionWithFriendlyRunId} for list
|
||||
* endpoints: one `IN` lookup per page instead of N+1. `currentRunId` on
|
||||
* the wire is always the public `run_*` friendlyId — the raw
|
||||
* {@link serializeSession} form leaks the internal cuid, which customers
|
||||
* can't use with `runs.retrieve(...)`.
|
||||
*/
|
||||
export async function serializeSessionsWithFriendlyRunIds(
|
||||
sessions: Session[],
|
||||
scope: { projectId: string; runtimeEnvironmentId: string },
|
||||
runStore: RunStore = defaultRunStore
|
||||
): Promise<SessionItem[]> {
|
||||
const runIds = [
|
||||
...new Set(sessions.map((s) => s.currentRunId).filter((id): id is string => !!id)),
|
||||
];
|
||||
|
||||
// `runStore.findRuns` fans out across both stores under split (NEW + LEGACY
|
||||
// replica merge) and is a plain `$replica` find when split is off. Tenant-
|
||||
// scoped: `Session.currentRunId` is a no-FK pointer, so a stale id must never
|
||||
// resolve a run in another env.
|
||||
const runs =
|
||||
runIds.length > 0
|
||||
? await runStore.findRuns({
|
||||
where: {
|
||||
id: { in: runIds },
|
||||
projectId: scope.projectId,
|
||||
runtimeEnvironmentId: scope.runtimeEnvironmentId,
|
||||
},
|
||||
select: { id: true, friendlyId: true },
|
||||
})
|
||||
: [];
|
||||
const friendlyIdByRunId = new Map(runs.map((run) => [run.id, run.friendlyId]));
|
||||
|
||||
return sessions.map((session) => ({
|
||||
...serializeSession(session),
|
||||
currentRunId: session.currentRunId
|
||||
? (friendlyIdByRunId.get(session.currentRunId) ?? null)
|
||||
: null,
|
||||
}));
|
||||
}
|
||||
@@ -0,0 +1,283 @@
|
||||
import {
|
||||
type ElectricColumnType,
|
||||
RUN_ELECTRIC_COLUMNS,
|
||||
serializeRunRow,
|
||||
} from "./electricStreamProtocol.server";
|
||||
import { type RunHydrator, type RunListFilter, type RunListResolver } from "./runReader.server";
|
||||
|
||||
/**
|
||||
* Dual-run shadow-compare: the client is always served the Electric response while this re-derives what
|
||||
* the native backend would emit and diffs the two, to prove parity on real traffic before cutover. Checks
|
||||
* serialization (semantic per-column compare, gated on same updatedAt so a changed row is "skew", not a
|
||||
* divergence) and membership (emitted id-set, only on tag/batch initial snapshots). Pure but for the injected deps.
|
||||
*/
|
||||
|
||||
export type ShadowFeed = "run" | "runs" | "batch";
|
||||
|
||||
type WireValue = Record<string, string | null>;
|
||||
|
||||
type ShapeMessage = {
|
||||
key?: string;
|
||||
value?: WireValue;
|
||||
headers: { operation?: string; control?: string };
|
||||
};
|
||||
|
||||
const COLUMN_BY_NAME = new Map(RUN_ELECTRIC_COLUMNS.map((column) => [column.name, column]));
|
||||
|
||||
export type ColumnDiff = {
|
||||
runId: string;
|
||||
column: string;
|
||||
electric: string | null;
|
||||
native: string | null;
|
||||
};
|
||||
|
||||
export type ShadowCompareOutcome = {
|
||||
feed: ShadowFeed;
|
||||
/** Runs whose every emitted column matched (same-version). */
|
||||
serializationMatched: number;
|
||||
/** Runs with at least one semantic column divergence (same-version). */
|
||||
serializationDiverged: number;
|
||||
/** Runs that changed between Electric's emit and our refetch (not a divergence). */
|
||||
serializationSkew: number;
|
||||
/** Per-column divergences (capped) for logging. */
|
||||
diffs: ColumnDiff[];
|
||||
/** Set membership (tag/batch initial snapshot only). undefined when not checked. */
|
||||
membershipMatch?: boolean;
|
||||
missingInNative?: string[];
|
||||
extraInNative?: string[];
|
||||
};
|
||||
|
||||
export type ShadowCompareInput = {
|
||||
feed: ShadowFeed;
|
||||
/** The served Electric response body (a JSON array of messages, or "" / "[]"). */
|
||||
electricBody: string;
|
||||
environment: { id: string };
|
||||
skipColumns: string[];
|
||||
/** True when this was an initial snapshot request (offset=-1); enables membership compare. */
|
||||
isInitialSnapshot: boolean;
|
||||
/** When set (tag/batch initial snapshot), compare the resolved id-set. */
|
||||
membershipFilter?: RunListFilter;
|
||||
};
|
||||
|
||||
const MAX_DIFFS = 20;
|
||||
|
||||
export class RealtimeShadowComparator {
|
||||
constructor(
|
||||
private readonly options: { runReader: RunHydrator; runListResolver: RunListResolver }
|
||||
) {}
|
||||
|
||||
async compare(input: ShadowCompareInput): Promise<ShadowCompareOutcome> {
|
||||
const messages = parseBody(input.electricBody);
|
||||
const changes = messages.filter(
|
||||
(m): m is ShapeMessage & { value: WireValue } =>
|
||||
typeof m.headers?.operation === "string" && !!m.value && m.headers.operation !== "delete"
|
||||
);
|
||||
|
||||
const outcome: ShadowCompareOutcome = {
|
||||
feed: input.feed,
|
||||
serializationMatched: 0,
|
||||
serializationDiverged: 0,
|
||||
serializationSkew: 0,
|
||||
diffs: [],
|
||||
};
|
||||
|
||||
// Bulk-hydrate every emitted run in one query rather than a per-message round
|
||||
// trip, so shadow mode doesn't inflate the very replica load it's measuring.
|
||||
const emittedIds = changes
|
||||
.map((m) => m.value.id)
|
||||
.filter((id): id is string => typeof id === "string");
|
||||
const hydrated = await this.options.runReader.hydrateByIds(input.environment.id, emittedIds);
|
||||
const rowsById = new Map(hydrated.map((row) => [row.id, row]));
|
||||
|
||||
for (const message of changes) {
|
||||
const runId = message.value.id ?? undefined;
|
||||
if (!runId) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const row = rowsById.get(runId);
|
||||
if (!row) {
|
||||
// Run no longer readable (deleted / replica miss). Not a serialization divergence.
|
||||
outcome.serializationSkew++;
|
||||
continue;
|
||||
}
|
||||
|
||||
const nativeValue = serializeRunRow(row, input.skipColumns);
|
||||
|
||||
// Only compare rows at the same version; otherwise the row advanced between
|
||||
// Electric's emit and our refetch (timing skew, not a divergence).
|
||||
if (!sameInstant(message.value.updatedAt, nativeValue.updatedAt)) {
|
||||
outcome.serializationSkew++;
|
||||
continue;
|
||||
}
|
||||
|
||||
let rowDiverged = false;
|
||||
for (const [column, electricRaw] of Object.entries(message.value)) {
|
||||
const meta = COLUMN_BY_NAME.get(column);
|
||||
if (!meta) {
|
||||
continue;
|
||||
}
|
||||
const nativeRaw = nativeValue[column] ?? null;
|
||||
if (!valuesEqual(electricRaw, nativeRaw, meta.type, meta.dims, column)) {
|
||||
rowDiverged = true;
|
||||
if (outcome.diffs.length < MAX_DIFFS) {
|
||||
outcome.diffs.push({ runId, column, electric: electricRaw, native: nativeRaw });
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (rowDiverged) {
|
||||
outcome.serializationDiverged++;
|
||||
} else {
|
||||
outcome.serializationMatched++;
|
||||
}
|
||||
}
|
||||
|
||||
if (input.isInitialSnapshot && input.membershipFilter) {
|
||||
const electricIds = new Set(
|
||||
changes.map((m) => m.value.id).filter((id): id is string => typeof id === "string")
|
||||
);
|
||||
const nativeIds = new Set(
|
||||
await this.options.runListResolver.resolveMatchingRunIds(input.membershipFilter)
|
||||
);
|
||||
|
||||
outcome.missingInNative = [...electricIds].filter((id) => !nativeIds.has(id));
|
||||
outcome.extraInNative = [...nativeIds].filter((id) => !electricIds.has(id));
|
||||
outcome.membershipMatch =
|
||||
outcome.missingInNative.length === 0 && outcome.extraInNative.length === 0;
|
||||
}
|
||||
|
||||
return outcome;
|
||||
}
|
||||
}
|
||||
|
||||
function parseBody(body: string): ShapeMessage[] {
|
||||
const text = body.trim();
|
||||
if (!text) {
|
||||
return [];
|
||||
}
|
||||
try {
|
||||
const parsed = JSON.parse(text);
|
||||
return Array.isArray(parsed) ? (parsed as ShapeMessage[]) : [];
|
||||
} catch {
|
||||
return [];
|
||||
}
|
||||
}
|
||||
|
||||
/** Status carries a known legacy rewrite (DEQUEUED -> EXECUTING) applied equally to
|
||||
* both paths for non-current API versions; treat them as equivalent. */
|
||||
function normalizeStatus(value: string): string {
|
||||
return value === "DEQUEUED" ? "EXECUTING" : value;
|
||||
}
|
||||
|
||||
function sameInstant(a: string | null | undefined, b: string | null | undefined): boolean {
|
||||
if (a == null || b == null) {
|
||||
return a == null && b == null;
|
||||
}
|
||||
// Mirror the SDK's RawShapeDate (`new Date(val + "Z")`).
|
||||
return new Date(`${a}Z`).getTime() === new Date(`${b}Z`).getTime();
|
||||
}
|
||||
|
||||
function valuesEqual(
|
||||
electricRaw: string | null,
|
||||
nativeRaw: string | null,
|
||||
type: ElectricColumnType,
|
||||
dims: number | undefined,
|
||||
column: string
|
||||
): boolean {
|
||||
if (electricRaw == null || nativeRaw == null) {
|
||||
return electricRaw == null && nativeRaw == null;
|
||||
}
|
||||
|
||||
if (dims && dims > 0) {
|
||||
return arraysEqual(parsePgTextArray(electricRaw), parsePgTextArray(nativeRaw));
|
||||
}
|
||||
|
||||
switch (type) {
|
||||
case "timestamp":
|
||||
return new Date(`${electricRaw}Z`).getTime() === new Date(`${nativeRaw}Z`).getTime();
|
||||
case "bool":
|
||||
return parseBool(electricRaw) === parseBool(nativeRaw);
|
||||
case "int4":
|
||||
case "int8":
|
||||
case "float8":
|
||||
return Number(electricRaw) === Number(nativeRaw);
|
||||
case "jsonb":
|
||||
return jsonEqual(electricRaw, nativeRaw);
|
||||
case "text":
|
||||
default:
|
||||
if (column === "status") {
|
||||
return normalizeStatus(electricRaw) === normalizeStatus(nativeRaw);
|
||||
}
|
||||
return electricRaw === nativeRaw;
|
||||
}
|
||||
}
|
||||
|
||||
function parseBool(value: string): boolean {
|
||||
return value === "t" || value === "true";
|
||||
}
|
||||
|
||||
function jsonEqual(a: string, b: string): boolean {
|
||||
try {
|
||||
return deepEqual(JSON.parse(a), JSON.parse(b));
|
||||
} catch {
|
||||
return a === b;
|
||||
}
|
||||
}
|
||||
|
||||
function deepEqual(a: unknown, b: unknown): boolean {
|
||||
if (a === b) return true;
|
||||
if (typeof a !== typeof b || a === null || b === null) return false;
|
||||
if (Array.isArray(a) && Array.isArray(b)) {
|
||||
return a.length === b.length && a.every((v, i) => deepEqual(v, b[i]));
|
||||
}
|
||||
if (typeof a === "object" && typeof b === "object") {
|
||||
const ak = Object.keys(a as object).sort();
|
||||
const bk = Object.keys(b as object).sort();
|
||||
return (
|
||||
ak.length === bk.length &&
|
||||
ak.every((k, i) => k === bk[i]) &&
|
||||
ak.every((k) => deepEqual((a as any)[k], (b as any)[k]))
|
||||
);
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
function arraysEqual(a: string[], b: string[]): boolean {
|
||||
return a.length === b.length && a.every((v, i) => v === b[i]);
|
||||
}
|
||||
|
||||
/** Parse a Postgres text-array literal (`{"a","b"}` / `{}`). Mirrors the client's pgArrayParser. */
|
||||
function parsePgTextArray(literal: string): string[] {
|
||||
if (literal === "{}" || literal === "") {
|
||||
return [];
|
||||
}
|
||||
const inner = literal.startsWith("{") && literal.endsWith("}") ? literal.slice(1, -1) : literal;
|
||||
const result: string[] = [];
|
||||
let i = 0;
|
||||
while (i < inner.length) {
|
||||
if (inner[i] === '"') {
|
||||
i++;
|
||||
let s = "";
|
||||
while (i < inner.length && inner[i] !== '"') {
|
||||
if (inner[i] === "\\") {
|
||||
i++;
|
||||
}
|
||||
s += inner[i];
|
||||
i++;
|
||||
}
|
||||
result.push(s);
|
||||
i++;
|
||||
if (inner[i] === ",") i++;
|
||||
} else {
|
||||
let s = "";
|
||||
while (i < inner.length && inner[i] !== ",") {
|
||||
s += inner[i];
|
||||
i++;
|
||||
}
|
||||
result.push(s);
|
||||
if (inner[i] === ",") i++;
|
||||
}
|
||||
}
|
||||
return result;
|
||||
}
|
||||
@@ -0,0 +1,188 @@
|
||||
import type { API_VERSIONS } from "~/api/versions";
|
||||
import { logger } from "../logger.server";
|
||||
import {
|
||||
type RealtimeEnvironment,
|
||||
type RealtimeRequestOptions,
|
||||
type RealtimeRunsParams,
|
||||
} from "../realtimeClient.server";
|
||||
import { RESERVED_COLUMNS } from "./electricStreamProtocol.server";
|
||||
import {
|
||||
type RealtimeListEnvironment,
|
||||
type RealtimeStreamClient,
|
||||
} from "./nativeRealtimeClient.server";
|
||||
import { type RunListFilter } from "./runReader.server";
|
||||
import {
|
||||
type RealtimeShadowComparator,
|
||||
type ShadowCompareOutcome,
|
||||
type ShadowFeed,
|
||||
} from "./shadowCompare.server";
|
||||
|
||||
export type ShadowRealtimeClientOptions = {
|
||||
/** The path actually served to the client (Electric). */
|
||||
electric: RealtimeStreamClient;
|
||||
comparator: RealtimeShadowComparator;
|
||||
/** createdAt window (ms) used to resolve tag-list membership for the compare. */
|
||||
maximumCreatedAtFilterAgeMs: number;
|
||||
/** Cap for the membership resolve. */
|
||||
maxListResults: number;
|
||||
/** Metrics sink for compare outcomes. */
|
||||
onOutcome?: (outcome: ShadowCompareOutcome) => void;
|
||||
};
|
||||
|
||||
/** Transparent wrapper that serves the Electric response unchanged and, in the background (fire-and-forget), diffs what the native backend would emit. */
|
||||
export class ShadowRealtimeClient implements RealtimeStreamClient {
|
||||
constructor(private readonly options: ShadowRealtimeClientOptions) {}
|
||||
|
||||
async streamRun(
|
||||
url: URL | string,
|
||||
environment: RealtimeEnvironment,
|
||||
runId: string,
|
||||
apiVersion: API_VERSIONS,
|
||||
requestOptions?: RealtimeRequestOptions,
|
||||
clientVersion?: string,
|
||||
signal?: AbortSignal
|
||||
): Promise<Response> {
|
||||
const response = await this.options.electric.streamRun(
|
||||
url,
|
||||
environment,
|
||||
runId,
|
||||
apiVersion,
|
||||
requestOptions,
|
||||
clientVersion,
|
||||
signal
|
||||
);
|
||||
this.#shadow("run", response, url, environment, requestOptions);
|
||||
return response;
|
||||
}
|
||||
|
||||
async streamRuns(
|
||||
url: URL | string,
|
||||
environment: RealtimeListEnvironment,
|
||||
params: RealtimeRunsParams,
|
||||
apiVersion: API_VERSIONS,
|
||||
requestOptions?: RealtimeRequestOptions,
|
||||
clientVersion?: string,
|
||||
signal?: AbortSignal
|
||||
): Promise<Response> {
|
||||
const response = await this.options.electric.streamRuns(
|
||||
url,
|
||||
environment,
|
||||
params,
|
||||
apiVersion,
|
||||
requestOptions,
|
||||
clientVersion,
|
||||
signal
|
||||
);
|
||||
this.#shadow("runs", response, url, environment, requestOptions, { tags: params.tags ?? [] });
|
||||
return response;
|
||||
}
|
||||
|
||||
async streamBatch(
|
||||
url: URL | string,
|
||||
environment: RealtimeListEnvironment,
|
||||
batchId: string,
|
||||
apiVersion: API_VERSIONS,
|
||||
requestOptions?: RealtimeRequestOptions,
|
||||
clientVersion?: string,
|
||||
signal?: AbortSignal
|
||||
): Promise<Response> {
|
||||
const response = await this.options.electric.streamBatch(
|
||||
url,
|
||||
environment,
|
||||
batchId,
|
||||
apiVersion,
|
||||
requestOptions,
|
||||
clientVersion,
|
||||
signal
|
||||
);
|
||||
this.#shadow("batch", response, url, environment, requestOptions, { batchId });
|
||||
return response;
|
||||
}
|
||||
|
||||
/** Fire-and-forget; never blocks the served response, never throws into the request. */
|
||||
#shadow(
|
||||
feed: ShadowFeed,
|
||||
electricResponse: Response,
|
||||
url: URL | string,
|
||||
environment: RealtimeEnvironment & { projectId?: string },
|
||||
requestOptions?: RealtimeRequestOptions,
|
||||
membership?: { tags?: string[]; batchId?: string }
|
||||
): void {
|
||||
// Clone synchronously before the client consumes the body.
|
||||
let bodyClone: Response;
|
||||
try {
|
||||
if (electricResponse.status !== 200) {
|
||||
return;
|
||||
}
|
||||
bodyClone = electricResponse.clone();
|
||||
} catch {
|
||||
return;
|
||||
}
|
||||
|
||||
void this.#runShadow(feed, bodyClone, url, environment, requestOptions, membership).catch(
|
||||
(error) => logger.debug("[shadowRealtime] compare failed", { feed, error })
|
||||
);
|
||||
}
|
||||
|
||||
async #runShadow(
|
||||
feed: ShadowFeed,
|
||||
bodyClone: Response,
|
||||
url: URL | string,
|
||||
environment: RealtimeEnvironment & { projectId?: string },
|
||||
requestOptions: RealtimeRequestOptions | undefined,
|
||||
membership: { tags?: string[]; batchId?: string } | undefined
|
||||
): Promise<void> {
|
||||
const $url = new URL(url.toString());
|
||||
const offset = $url.searchParams.get("offset") ?? "-1";
|
||||
const handle = $url.searchParams.get("handle") ?? $url.searchParams.get("shape_id");
|
||||
const isInitialSnapshot = offset === "-1" || !handle;
|
||||
const skipColumns = resolveSkipColumns($url, requestOptions);
|
||||
const electricBody = await bodyClone.text();
|
||||
|
||||
let membershipFilter: RunListFilter | undefined;
|
||||
if (isInitialSnapshot && membership && environment.projectId) {
|
||||
membershipFilter = {
|
||||
organizationId: environment.organizationId,
|
||||
projectId: environment.projectId,
|
||||
environmentId: environment.id,
|
||||
tags: membership.tags,
|
||||
batchId: membership.batchId,
|
||||
createdAtAfter: membership.batchId
|
||||
? undefined
|
||||
: new Date(Date.now() - this.options.maximumCreatedAtFilterAgeMs),
|
||||
limit: this.options.maxListResults,
|
||||
};
|
||||
}
|
||||
|
||||
const outcome = await this.options.comparator.compare({
|
||||
feed,
|
||||
electricBody,
|
||||
environment: { id: environment.id },
|
||||
skipColumns,
|
||||
isInitialSnapshot,
|
||||
membershipFilter,
|
||||
});
|
||||
|
||||
this.options.onOutcome?.(outcome);
|
||||
|
||||
if (outcome.serializationDiverged > 0 || outcome.membershipMatch === false) {
|
||||
logger.warn("[shadowRealtime] divergence detected", {
|
||||
feed,
|
||||
serializationDiverged: outcome.serializationDiverged,
|
||||
serializationMatched: outcome.serializationMatched,
|
||||
serializationSkew: outcome.serializationSkew,
|
||||
membershipMatch: outcome.membershipMatch,
|
||||
missingInNative: outcome.missingInNative?.slice(0, 20),
|
||||
extraInNative: outcome.extraInNative?.slice(0, 20),
|
||||
// Log only which run/column diverged, never the raw cell values — they can
|
||||
// include run payload/output/metadata and must not leak into logs.
|
||||
diffs: outcome.diffs.map(({ runId, column }) => ({ runId, column })),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
function resolveSkipColumns(url: URL, requestOptions?: RealtimeRequestOptions): string[] {
|
||||
const raw = requestOptions?.skipColumns ?? url.searchParams.get("skipColumns")?.split(",") ?? [];
|
||||
return raw.map((c) => c.trim()).filter((c) => c !== "" && !RESERVED_COLUMNS.includes(c));
|
||||
}
|
||||
@@ -0,0 +1,69 @@
|
||||
import { getMeter } from "@internal/tracing";
|
||||
import { $replica } from "~/db.server";
|
||||
import { runStore } from "~/v3/runStore.server";
|
||||
import { env } from "~/env.server";
|
||||
import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactoryInstance.server";
|
||||
import { singleton } from "~/utils/singleton";
|
||||
import { realtimeClient } from "../realtimeClientGlobal.server";
|
||||
import { ClickHouseRunListResolver } from "./clickHouseRunListResolver.server";
|
||||
import { RunHydrator } from "./runReader.server";
|
||||
import { RealtimeShadowComparator } from "./shadowCompare.server";
|
||||
import { ShadowRealtimeClient } from "./shadowRealtimeClient.server";
|
||||
|
||||
/**
|
||||
* Process-singleton wiring for the shadow-compare client. Only constructed
|
||||
* when an org's `realtimeBackend` flag is set to "shadow".
|
||||
*/
|
||||
function initializeShadowRealtimeClient(): ShadowRealtimeClient {
|
||||
const compares = getMeter("realtime-shadow").createCounter("realtime_shadow.compares", {
|
||||
description:
|
||||
"Dual-run shadow-compare outcomes (Electric vs native). kind=serialization|membership, result=match|diverge|skew.",
|
||||
});
|
||||
|
||||
const comparator = new RealtimeShadowComparator({
|
||||
runReader: new RunHydrator({ replica: $replica, runStore }),
|
||||
runListResolver: new ClickHouseRunListResolver({
|
||||
getClickhouse: (organizationId) =>
|
||||
clickhouseFactory.getClickhouseForOrganization(organizationId, "realtime"),
|
||||
prisma: $replica,
|
||||
}),
|
||||
});
|
||||
|
||||
return new ShadowRealtimeClient({
|
||||
electric: realtimeClient,
|
||||
comparator,
|
||||
maximumCreatedAtFilterAgeMs: env.REALTIME_MAXIMUM_CREATED_AT_FILTER_AGE_IN_MS,
|
||||
maxListResults: env.REALTIME_BACKEND_NATIVE_MAX_LIST_RESULTS,
|
||||
onOutcome: (outcome) => {
|
||||
const { feed } = outcome;
|
||||
if (outcome.serializationMatched) {
|
||||
compares.add(outcome.serializationMatched, {
|
||||
feed,
|
||||
kind: "serialization",
|
||||
result: "match",
|
||||
});
|
||||
}
|
||||
if (outcome.serializationDiverged) {
|
||||
compares.add(outcome.serializationDiverged, {
|
||||
feed,
|
||||
kind: "serialization",
|
||||
result: "diverge",
|
||||
});
|
||||
}
|
||||
if (outcome.serializationSkew) {
|
||||
compares.add(outcome.serializationSkew, { feed, kind: "serialization", result: "skew" });
|
||||
}
|
||||
if (outcome.membershipMatch !== undefined) {
|
||||
compares.add(1, {
|
||||
feed,
|
||||
kind: "membership",
|
||||
result: outcome.membershipMatch ? "match" : "diverge",
|
||||
});
|
||||
}
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
export function getShadowRealtimeClient(): ShadowRealtimeClient {
|
||||
return singleton("shadowRealtimeClient", initializeShadowRealtimeClient);
|
||||
}
|
||||
@@ -0,0 +1,246 @@
|
||||
/**
|
||||
* Per-org S2 basin provisioning. Gated by
|
||||
* `REALTIME_STREAMS_PER_ORG_BASINS_ENABLED`: when off, all orgs share
|
||||
* `REALTIME_STREAMS_S2_BASIN` and this module no-ops.
|
||||
*
|
||||
* Pure retention-string in / S2-call out. Plan vocabulary lives in the
|
||||
* cloud billing app, which calls into the admin sync route to drive
|
||||
* provisioning + reconfiguration.
|
||||
*/
|
||||
import type { PrismaClientOrTransaction } from "~/db.server";
|
||||
import { prisma } from "~/db.server";
|
||||
import { env } from "~/env.server";
|
||||
import { logger } from "~/services/logger.server";
|
||||
import { controlPlaneResolver } from "~/v3/runOpsMigration/controlPlaneResolver.server";
|
||||
import { parseDuration } from "./duration.server";
|
||||
|
||||
export function isPerOrgBasinsEnabled(): boolean {
|
||||
return env.REALTIME_STREAMS_PER_ORG_BASINS_ENABLED === "true";
|
||||
}
|
||||
|
||||
export function defaultRetention(): string {
|
||||
return env.REALTIME_STREAMS_BASIN_DEFAULT_RETENTION;
|
||||
}
|
||||
|
||||
// Org id is a cuid — fixed-length and stable, so the basin name is
|
||||
// collision-free without truncation. Slugs are user-editable and would
|
||||
// drift.
|
||||
export function basinNameForOrg(org: { id: string }): string {
|
||||
const prefix = env.REALTIME_STREAMS_BASIN_NAME_PREFIX;
|
||||
const envName = env.REALTIME_STREAMS_BASIN_NAME_ENV;
|
||||
return `${prefix}-${envName}-org-${org.id}`;
|
||||
}
|
||||
|
||||
type ProvisionInput = {
|
||||
id: string;
|
||||
retention?: string;
|
||||
streamBasinName: string | null | undefined;
|
||||
};
|
||||
|
||||
type ProvisionResult =
|
||||
| { kind: "skipped"; reason: "feature-disabled" | "already-provisioned"; basin: string | null }
|
||||
| { kind: "provisioned"; basin: string; retention: string };
|
||||
|
||||
// Idempotent. Treats S2 409 as success (race with another caller, or
|
||||
// previous run that crashed after S2 ack but before the column write).
|
||||
export async function provisionBasinForOrg(
|
||||
org: ProvisionInput,
|
||||
prismaClient: PrismaClientOrTransaction = prisma
|
||||
): Promise<ProvisionResult> {
|
||||
if (!isPerOrgBasinsEnabled()) {
|
||||
return { kind: "skipped", reason: "feature-disabled", basin: null };
|
||||
}
|
||||
|
||||
if (org.streamBasinName) {
|
||||
return { kind: "skipped", reason: "already-provisioned", basin: org.streamBasinName };
|
||||
}
|
||||
|
||||
const accessToken = env.REALTIME_STREAMS_S2_ACCESS_TOKEN;
|
||||
if (!accessToken) {
|
||||
throw new Error(
|
||||
"REALTIME_STREAMS_S2_ACCESS_TOKEN must be set when REALTIME_STREAMS_PER_ORG_BASINS_ENABLED=true"
|
||||
);
|
||||
}
|
||||
|
||||
const basin = basinNameForOrg(org);
|
||||
const retention = org.retention ?? defaultRetention();
|
||||
|
||||
await s2CreateBasin(basin, {
|
||||
accessToken,
|
||||
retentionPolicy: retention,
|
||||
storageClass: env.REALTIME_STREAMS_BASIN_STORAGE_CLASS,
|
||||
deleteOnEmptyMinAge: env.REALTIME_STREAMS_BASIN_DELETE_ON_EMPTY_MIN_AGE,
|
||||
});
|
||||
|
||||
await prismaClient.organization.update({
|
||||
where: { id: org.id },
|
||||
data: { streamBasinName: basin },
|
||||
});
|
||||
|
||||
// streamBasinName is embedded in every env of the org; drop all its cached env rows.
|
||||
controlPlaneResolver.invalidateOrganization(org.id);
|
||||
|
||||
logger.info("[streamBasinProvisioner] provisioned basin for org", {
|
||||
orgId: org.id,
|
||||
basin,
|
||||
retention,
|
||||
});
|
||||
|
||||
return { kind: "provisioned", basin, retention };
|
||||
}
|
||||
|
||||
export async function reconfigureBasinForOrg(orgId: string, retention: string): Promise<void> {
|
||||
if (!isPerOrgBasinsEnabled()) return;
|
||||
|
||||
const accessToken = env.REALTIME_STREAMS_S2_ACCESS_TOKEN;
|
||||
if (!accessToken) {
|
||||
throw new Error(
|
||||
"REALTIME_STREAMS_S2_ACCESS_TOKEN must be set when REALTIME_STREAMS_PER_ORG_BASINS_ENABLED=true"
|
||||
);
|
||||
}
|
||||
|
||||
const org = await prisma.organization.findFirst({
|
||||
where: { id: orgId },
|
||||
select: { id: true, streamBasinName: true },
|
||||
});
|
||||
if (!org?.streamBasinName) return;
|
||||
|
||||
await s2ReconfigureBasin(org.streamBasinName, { accessToken, retentionPolicy: retention });
|
||||
|
||||
logger.info("[streamBasinProvisioner] reconfigured basin retention", {
|
||||
orgId,
|
||||
basin: org.streamBasinName,
|
||||
retention,
|
||||
});
|
||||
}
|
||||
|
||||
type EnsureResult =
|
||||
| { kind: "skipped"; reason: "feature-disabled" | "org-not-found" }
|
||||
| { kind: "provisioned"; basin: string; retention: string }
|
||||
| { kind: "reconfigured"; basin: string; retention: string };
|
||||
|
||||
// Idempotent: provisions if the org has no basin, PATCHes retention if
|
||||
// it does. The single entrypoint the cloud billing app drives — both
|
||||
// for the live plan-change path and the bulk backfill.
|
||||
export async function ensureBasinForOrg(orgId: string, retention: string): Promise<EnsureResult> {
|
||||
if (!isPerOrgBasinsEnabled()) {
|
||||
return { kind: "skipped", reason: "feature-disabled" };
|
||||
}
|
||||
|
||||
const org = await prisma.organization.findFirst({
|
||||
where: { id: orgId },
|
||||
select: { id: true, streamBasinName: true },
|
||||
});
|
||||
if (!org) return { kind: "skipped", reason: "org-not-found" };
|
||||
|
||||
if (!org.streamBasinName) {
|
||||
const result = await provisionBasinForOrg({ id: org.id, streamBasinName: null, retention });
|
||||
if (result.kind === "provisioned") {
|
||||
return { kind: "provisioned", basin: result.basin, retention: result.retention };
|
||||
}
|
||||
return { kind: "skipped", reason: "feature-disabled" };
|
||||
}
|
||||
|
||||
await reconfigureBasinForOrg(org.id, retention);
|
||||
return { kind: "reconfigured", basin: org.streamBasinName, retention };
|
||||
}
|
||||
|
||||
// Inverse of ensureBasinForOrg: nulls the column so future runs/sessions
|
||||
// land in the shared global basin. The S2 basin lingers; existing streams
|
||||
// age out on their original retention.
|
||||
export async function deprovisionBasinForOrg(
|
||||
orgId: string
|
||||
): Promise<{ kind: "deprovisioned" } | { kind: "skipped"; reason: "no-basin" }> {
|
||||
const org = await prisma.organization.findFirst({
|
||||
where: { id: orgId },
|
||||
select: { id: true, streamBasinName: true },
|
||||
});
|
||||
if (!org?.streamBasinName) return { kind: "skipped", reason: "no-basin" };
|
||||
|
||||
await prisma.organization.update({
|
||||
where: { id: org.id },
|
||||
data: { streamBasinName: null },
|
||||
});
|
||||
|
||||
// streamBasinName is embedded in every env of the org; drop all its cached env rows.
|
||||
controlPlaneResolver.invalidateOrganization(org.id);
|
||||
|
||||
logger.info("[streamBasinProvisioner] deprovisioned basin for org", {
|
||||
orgId,
|
||||
previousBasin: org.streamBasinName,
|
||||
});
|
||||
|
||||
return { kind: "deprovisioned" };
|
||||
}
|
||||
|
||||
// S2 REST: POST /v1/basins to create, PATCH /v1/basins/{name} to
|
||||
// reconfigure. Wire shape takes integer seconds; we accept human strings
|
||||
// like "7d" / "1y" as env-var ergonomics and parse them here.
|
||||
|
||||
type CreateBasinOptions = {
|
||||
accessToken: string;
|
||||
retentionPolicy: string;
|
||||
storageClass: "express" | "standard";
|
||||
deleteOnEmptyMinAge: string;
|
||||
};
|
||||
|
||||
async function s2CreateBasin(name: string, opts: CreateBasinOptions): Promise<void> {
|
||||
const url = `https://aws.s2.dev/v1/basins`;
|
||||
const body = {
|
||||
basin: name,
|
||||
config: {
|
||||
create_stream_on_append: true,
|
||||
create_stream_on_read: true,
|
||||
default_stream_config: {
|
||||
storage_class: opts.storageClass,
|
||||
retention_policy: { age: parseDuration(opts.retentionPolicy) },
|
||||
delete_on_empty: { min_age_secs: parseDuration(opts.deleteOnEmptyMinAge) },
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
const res = await fetch(url, {
|
||||
signal: AbortSignal.timeout(10_000),
|
||||
method: "POST",
|
||||
headers: {
|
||||
Authorization: `Bearer ${opts.accessToken}`,
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
body: JSON.stringify(body),
|
||||
});
|
||||
|
||||
// 409 = basin already exists; treat as success (idempotent).
|
||||
if (res.ok || res.status === 409) return;
|
||||
|
||||
const text = await res.text().catch(() => "");
|
||||
throw new Error(`S2 createBasin failed: ${res.status} ${res.statusText} ${text}`);
|
||||
}
|
||||
|
||||
type ReconfigureBasinOptions = {
|
||||
accessToken: string;
|
||||
retentionPolicy: string;
|
||||
};
|
||||
|
||||
async function s2ReconfigureBasin(name: string, opts: ReconfigureBasinOptions): Promise<void> {
|
||||
const url = `https://aws.s2.dev/v1/basins/${encodeURIComponent(name)}`;
|
||||
const body = {
|
||||
default_stream_config: {
|
||||
retention_policy: { age: parseDuration(opts.retentionPolicy) },
|
||||
},
|
||||
};
|
||||
|
||||
const res = await fetch(url, {
|
||||
signal: AbortSignal.timeout(10_000),
|
||||
method: "PATCH",
|
||||
headers: {
|
||||
Authorization: `Bearer ${opts.accessToken}`,
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
body: JSON.stringify(body),
|
||||
});
|
||||
|
||||
if (res.ok) return;
|
||||
|
||||
const text = await res.text().catch(() => "");
|
||||
throw new Error(`S2 reconfigureBasin failed: ${res.status} ${res.statusText} ${text}`);
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
export type StreamRecord = {
|
||||
data: string;
|
||||
id: string;
|
||||
seqNum: number;
|
||||
/**
|
||||
* S2 record headers, when the underlying backend is the v2 (S2) shape.
|
||||
* Undefined or empty for run-scoped Redis streams. First-header empty-name
|
||||
* is an S2 command record (trim/fence); the parser strips those before
|
||||
* surfacing the record, so callers never see them.
|
||||
*/
|
||||
headers?: Array<[string, string]>;
|
||||
};
|
||||
|
||||
// Interface for stream ingestion
|
||||
export interface StreamIngestor {
|
||||
initializeStream(
|
||||
runId: string,
|
||||
streamId: string
|
||||
): Promise<{ responseHeaders?: Record<string, string> }>;
|
||||
|
||||
ingestData(
|
||||
stream: ReadableStream<Uint8Array>,
|
||||
runId: string,
|
||||
streamId: string,
|
||||
clientId: string,
|
||||
resumeFromChunk?: number
|
||||
): Promise<Response>;
|
||||
|
||||
appendPart(part: string, partId: string, runId: string, streamId: string): Promise<void>;
|
||||
|
||||
getLastChunkIndex(runId: string, streamId: string, clientId: string): Promise<number>;
|
||||
|
||||
readRecords(runId: string, streamId: string, afterSeqNum?: number): Promise<StreamRecord[]>;
|
||||
}
|
||||
|
||||
export type StreamResponseOptions = {
|
||||
timeoutInSeconds?: number;
|
||||
lastEventId?: string;
|
||||
/**
|
||||
* Session-stream-only. When `true`, the responder MAY peek the tail
|
||||
* of `.out` and short-circuit to `wait=0` + `X-Session-Settled: true`
|
||||
* if the last record is a terminal marker (a `trigger-control`
|
||||
* `turn-complete` control record, ignoring any trailing S2 trim
|
||||
* command record). Used by `TriggerChatTransport.reconnectToStream`
|
||||
* on page reload.
|
||||
*
|
||||
* When absent/false, the responder keeps the unconditional long-poll
|
||||
* behavior — required on the active send-a-message path where the
|
||||
* peek would race the newly-triggered turn's first chunk.
|
||||
*/
|
||||
peekSettled?: boolean;
|
||||
};
|
||||
|
||||
// Interface for stream response
|
||||
export interface StreamResponder {
|
||||
streamResponse(
|
||||
request: Request,
|
||||
runId: string,
|
||||
streamId: string,
|
||||
signal: AbortSignal,
|
||||
options?: StreamResponseOptions
|
||||
): Promise<Response>;
|
||||
}
|
||||
@@ -0,0 +1,33 @@
|
||||
export class LineTransformStream extends TransformStream<string, string[]> {
|
||||
private buffer = "";
|
||||
|
||||
constructor() {
|
||||
super({
|
||||
transform: (chunk, controller) => {
|
||||
// Append the chunk to the buffer
|
||||
this.buffer += chunk;
|
||||
|
||||
// Split on newlines
|
||||
const lines = this.buffer.split("\n");
|
||||
|
||||
// The last element might be incomplete, hold it back in buffer
|
||||
this.buffer = lines.pop() || "";
|
||||
|
||||
// Filter out empty or whitespace-only lines
|
||||
const fullLines = lines.filter((line) => line.trim().length > 0);
|
||||
|
||||
// If we got any complete lines, emit them as an array
|
||||
if (fullLines.length > 0) {
|
||||
controller.enqueue(fullLines);
|
||||
}
|
||||
},
|
||||
flush: (controller) => {
|
||||
// On stream end, if there's leftover text, emit it as a single-element array
|
||||
const trimmed = this.buffer.trim();
|
||||
if (trimmed.length > 0) {
|
||||
controller.enqueue([trimmed]);
|
||||
}
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,143 @@
|
||||
import {
|
||||
createCache,
|
||||
createLRUMemoryStore,
|
||||
DefaultStatefulContext,
|
||||
Namespace,
|
||||
RedisCacheStore,
|
||||
} from "@internal/cache";
|
||||
import { env } from "~/env.server";
|
||||
import { singleton } from "~/utils/singleton";
|
||||
import type { AuthenticatedEnvironment } from "../apiAuth.server";
|
||||
import { RedisRealtimeStreams } from "./redisRealtimeStreams.server";
|
||||
import { S2RealtimeStreams } from "./s2realtimeStreams.server";
|
||||
import type { StreamIngestor, StreamResponder } from "./types";
|
||||
|
||||
function initializeRedisRealtimeStreams() {
|
||||
return new RedisRealtimeStreams({
|
||||
redis: {
|
||||
port: env.REALTIME_STREAMS_REDIS_PORT,
|
||||
host: env.REALTIME_STREAMS_REDIS_HOST,
|
||||
username: env.REALTIME_STREAMS_REDIS_USERNAME,
|
||||
password: env.REALTIME_STREAMS_REDIS_PASSWORD,
|
||||
enableAutoPipelining: true,
|
||||
...(env.REALTIME_STREAMS_REDIS_TLS_DISABLED === "true" ? {} : { tls: {} }),
|
||||
keyPrefix: "tr:realtime:streams:",
|
||||
},
|
||||
inactivityTimeoutMs: env.REALTIME_STREAMS_INACTIVITY_TIMEOUT_MS,
|
||||
});
|
||||
}
|
||||
|
||||
export const v1RealtimeStreams = singleton("realtimeStreams", initializeRedisRealtimeStreams);
|
||||
|
||||
/**
|
||||
* Resolve a stream's basin. Precedence: run → session → org → global env.
|
||||
* Pre-migration rows have `streamBasinName: null` and fall through to
|
||||
* the global basin (where their streams actually live), so only pass
|
||||
* `organization` when no run/session row exists at all — otherwise a
|
||||
* null column would short-circuit to the org's *current* basin.
|
||||
*/
|
||||
export type StreamBasinContext = {
|
||||
run?: { streamBasinName: string | null } | null;
|
||||
session?: { streamBasinName: string | null } | null;
|
||||
organization?: { streamBasinName: string | null } | null;
|
||||
};
|
||||
|
||||
export function resolveStreamBasin(ctx: StreamBasinContext): string | undefined {
|
||||
return (
|
||||
ctx.run?.streamBasinName ??
|
||||
ctx.session?.streamBasinName ??
|
||||
ctx.organization?.streamBasinName ??
|
||||
env.REALTIME_STREAMS_S2_BASIN ??
|
||||
undefined
|
||||
);
|
||||
}
|
||||
|
||||
export function getRealtimeStreamInstance(
|
||||
environment: AuthenticatedEnvironment,
|
||||
streamVersion: string,
|
||||
basinContext?: StreamBasinContext
|
||||
): StreamIngestor & StreamResponder {
|
||||
if (streamVersion === "v1") {
|
||||
return v1RealtimeStreams;
|
||||
}
|
||||
|
||||
const resolvedBasin = resolveStreamBasin(basinContext ?? {});
|
||||
if (
|
||||
resolvedBasin &&
|
||||
(env.REALTIME_STREAMS_S2_ACCESS_TOKEN || env.REALTIME_STREAMS_S2_SKIP_ACCESS_TOKENS === "true")
|
||||
) {
|
||||
return new S2RealtimeStreams({
|
||||
basin: resolvedBasin,
|
||||
accessToken: env.REALTIME_STREAMS_S2_ACCESS_TOKEN ?? "",
|
||||
endpoint: env.REALTIME_STREAMS_S2_ENDPOINT,
|
||||
skipAccessTokens: env.REALTIME_STREAMS_S2_SKIP_ACCESS_TOKENS === "true",
|
||||
streamPrefix: streamPrefixFor(environment, resolvedBasin),
|
||||
logLevel: env.REALTIME_STREAMS_S2_LOG_LEVEL,
|
||||
flushIntervalMs: env.REALTIME_STREAMS_S2_FLUSH_INTERVAL_MS,
|
||||
maxRetries: env.REALTIME_STREAMS_S2_MAX_RETRIES,
|
||||
s2WaitSeconds: env.REALTIME_STREAMS_S2_WAIT_SECONDS,
|
||||
accessTokenExpirationInMs: env.REALTIME_STREAMS_S2_ACCESS_TOKEN_EXPIRATION_IN_MS,
|
||||
cache: s2RealtimeStreamsCache,
|
||||
});
|
||||
}
|
||||
|
||||
throw new Error("Realtime streams v2 is required for this run but S2 configuration is missing");
|
||||
}
|
||||
|
||||
// Shared basin needs `org/{orgId}` to namespace; per-org basin already
|
||||
// isolates so the segment drops.
|
||||
function streamPrefixFor(environment: AuthenticatedEnvironment, basin: string): string {
|
||||
const isPerOrgBasin = basin !== env.REALTIME_STREAMS_S2_BASIN;
|
||||
const segments = isPerOrgBasin
|
||||
? ["env", environment.slug, environment.id]
|
||||
: ["org", environment.organization.id, "env", environment.slug, environment.id];
|
||||
return segments.join("/");
|
||||
}
|
||||
|
||||
export function determineRealtimeStreamsVersion(streamVersion?: string): "v1" | "v2" {
|
||||
if (!streamVersion) {
|
||||
return env.REALTIME_STREAMS_DEFAULT_VERSION;
|
||||
}
|
||||
|
||||
if (
|
||||
streamVersion === "v2" &&
|
||||
env.REALTIME_STREAMS_S2_BASIN &&
|
||||
(env.REALTIME_STREAMS_S2_ACCESS_TOKEN || env.REALTIME_STREAMS_S2_SKIP_ACCESS_TOKENS === "true")
|
||||
) {
|
||||
return "v2";
|
||||
}
|
||||
|
||||
return "v1";
|
||||
}
|
||||
|
||||
const s2RealtimeStreamsCache = singleton(
|
||||
"s2RealtimeStreamsCache",
|
||||
initializeS2RealtimeStreamsCache
|
||||
);
|
||||
|
||||
function initializeS2RealtimeStreamsCache() {
|
||||
const ctx = new DefaultStatefulContext();
|
||||
const redisCacheStore = new RedisCacheStore({
|
||||
name: "s2-realtime-streams-cache",
|
||||
connection: {
|
||||
port: env.REALTIME_STREAMS_REDIS_PORT,
|
||||
host: env.REALTIME_STREAMS_REDIS_HOST,
|
||||
username: env.REALTIME_STREAMS_REDIS_USERNAME,
|
||||
password: env.REALTIME_STREAMS_REDIS_PASSWORD,
|
||||
enableAutoPipelining: true,
|
||||
...(env.REALTIME_STREAMS_REDIS_TLS_DISABLED === "true" ? {} : { tls: {} }),
|
||||
keyPrefix: "s2-realtime-streams-cache:",
|
||||
},
|
||||
useModernCacheKeyBuilder: true,
|
||||
});
|
||||
|
||||
const memoryStore = createLRUMemoryStore(5000);
|
||||
|
||||
return createCache({
|
||||
accessToken: new Namespace<string>(ctx, {
|
||||
stores: [memoryStore, redisCacheStore],
|
||||
fresh: Math.floor(env.REALTIME_STREAMS_S2_ACCESS_TOKEN_EXPIRATION_IN_MS / 2),
|
||||
stale: Math.floor(env.REALTIME_STREAMS_S2_ACCESS_TOKEN_EXPIRATION_IN_MS / 2 + 60_000),
|
||||
}),
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,614 @@
|
||||
import { json } from "@remix-run/server-runtime";
|
||||
import { tryCatch } from "@trigger.dev/core/utils";
|
||||
import { safeParseNaturalLanguageDurationAgo } from "@trigger.dev/core/v3/isomorphic";
|
||||
import type { Callback, Result } from "ioredis";
|
||||
import { randomUUID } from "node:crypto";
|
||||
import type { RedisClient, RedisWithClusterOptions } from "~/redis.server";
|
||||
import { createRedisClient } from "~/redis.server";
|
||||
import { longPollingFetch } from "~/utils/longPollingFetch";
|
||||
import { logger } from "./logger.server";
|
||||
import { jumpHash } from "@trigger.dev/core/v3/serverOnly";
|
||||
import type { Cache } from "@unkey/cache";
|
||||
import { createCache, DefaultStatefulContext, Namespace } from "@unkey/cache";
|
||||
import { createLRUMemoryStore } from "@internal/cache";
|
||||
import { RedisCacheStore } from "./unkey/redisCacheStore.server";
|
||||
import { env } from "~/env.server";
|
||||
import type { API_VERSIONS } from "~/api/versions";
|
||||
import { CURRENT_API_VERSION } from "~/api/versions";
|
||||
|
||||
export interface CachedLimitProvider {
|
||||
getCachedLimit: (organizationId: string, defaultValue: number) => Promise<number | undefined>;
|
||||
}
|
||||
|
||||
const DEFAULT_ELECTRIC_COLUMNS = [
|
||||
"id",
|
||||
"taskIdentifier",
|
||||
"createdAt",
|
||||
"updatedAt",
|
||||
"startedAt",
|
||||
"delayUntil",
|
||||
"queuedAt",
|
||||
"expiredAt",
|
||||
"completedAt",
|
||||
"friendlyId",
|
||||
"number",
|
||||
"isTest",
|
||||
"status",
|
||||
"usageDurationMs",
|
||||
"costInCents",
|
||||
"baseCostInCents",
|
||||
"ttl",
|
||||
"payload",
|
||||
"payloadType",
|
||||
"metadata",
|
||||
"metadataType",
|
||||
"output",
|
||||
"outputType",
|
||||
"runTags",
|
||||
"error",
|
||||
"realtimeStreams",
|
||||
];
|
||||
|
||||
const RESERVED_COLUMNS = ["id", "taskIdentifier", "friendlyId", "status", "createdAt"];
|
||||
const RESERVED_SEARCH_PARAMS = ["createdAt", "tags", "skipColumns"];
|
||||
|
||||
export type RealtimeClientOptions = {
|
||||
electricOrigin: string | string[];
|
||||
redis: RedisWithClusterOptions;
|
||||
cachedLimitProvider: CachedLimitProvider;
|
||||
keyPrefix: string;
|
||||
expiryTimeInSeconds?: number;
|
||||
};
|
||||
|
||||
export type RealtimeEnvironment = {
|
||||
id: string;
|
||||
organizationId: string;
|
||||
};
|
||||
|
||||
export type RealtimeRunsParams = {
|
||||
tags?: string[];
|
||||
createdAt?: string;
|
||||
};
|
||||
|
||||
export type RealtimeRequestOptions = {
|
||||
skipColumns?: string[];
|
||||
};
|
||||
|
||||
export class RealtimeClient {
|
||||
private redis: RedisClient;
|
||||
private expiryTimeInSeconds: number;
|
||||
private cachedLimitProvider: CachedLimitProvider;
|
||||
private cache: Cache<{ createdAtFilter: string }>;
|
||||
|
||||
constructor(private options: RealtimeClientOptions) {
|
||||
this.redis = createRedisClient("trigger:realtime", options.redis);
|
||||
this.expiryTimeInSeconds = options.expiryTimeInSeconds ?? 60 * 5; // default to 5 minutes
|
||||
this.cachedLimitProvider = options.cachedLimitProvider;
|
||||
this.#registerCommands();
|
||||
|
||||
const ctx = new DefaultStatefulContext();
|
||||
const memory = createLRUMemoryStore(1000);
|
||||
const redisCacheStore = new RedisCacheStore({
|
||||
connection: {
|
||||
keyPrefix: "tr:cache:realtime",
|
||||
port: options.redis.port,
|
||||
host: options.redis.host,
|
||||
username: options.redis.username,
|
||||
password: options.redis.password,
|
||||
tlsDisabled: options.redis.tlsDisabled,
|
||||
clusterMode: options.redis.clusterMode,
|
||||
},
|
||||
});
|
||||
|
||||
// This cache holds the limits fetched from the platform service
|
||||
const cache = createCache({
|
||||
createdAtFilter: new Namespace<string>(ctx, {
|
||||
stores: [memory, redisCacheStore],
|
||||
fresh: 60_000 * 60 * 24 * 7, // 1 week
|
||||
stale: 60_000 * 60 * 24 * 14, // 2 weeks
|
||||
}),
|
||||
});
|
||||
|
||||
this.cache = cache;
|
||||
}
|
||||
|
||||
async streamRun(
|
||||
url: URL | string,
|
||||
environment: RealtimeEnvironment,
|
||||
runId: string,
|
||||
apiVersion: API_VERSIONS,
|
||||
requestOptions?: RealtimeRequestOptions,
|
||||
clientVersion?: string,
|
||||
signal?: AbortSignal
|
||||
) {
|
||||
return this.#streamRunsWhere(
|
||||
url,
|
||||
environment,
|
||||
`id='${runId}'`,
|
||||
apiVersion,
|
||||
requestOptions,
|
||||
clientVersion,
|
||||
signal
|
||||
);
|
||||
}
|
||||
|
||||
async streamBatch(
|
||||
url: URL | string,
|
||||
environment: RealtimeEnvironment,
|
||||
batchId: string,
|
||||
apiVersion: API_VERSIONS,
|
||||
requestOptions?: RealtimeRequestOptions,
|
||||
clientVersion?: string,
|
||||
signal?: AbortSignal
|
||||
) {
|
||||
const whereClauses: string[] = [
|
||||
`"runtimeEnvironmentId"='${environment.id}'`,
|
||||
`"batchId"='${batchId}'`,
|
||||
];
|
||||
|
||||
const whereClause = whereClauses.join(" AND ");
|
||||
|
||||
return this.#streamRunsWhere(
|
||||
url,
|
||||
environment,
|
||||
whereClause,
|
||||
apiVersion,
|
||||
requestOptions,
|
||||
clientVersion,
|
||||
signal
|
||||
);
|
||||
}
|
||||
|
||||
async streamRuns(
|
||||
url: URL | string,
|
||||
environment: RealtimeEnvironment,
|
||||
params: RealtimeRunsParams,
|
||||
apiVersion: API_VERSIONS,
|
||||
requestOptions?: RealtimeRequestOptions,
|
||||
clientVersion?: string,
|
||||
signal?: AbortSignal
|
||||
) {
|
||||
const whereClauses: string[] = [`"runtimeEnvironmentId"='${environment.id}'`];
|
||||
|
||||
if (params.tags) {
|
||||
whereClauses.push(`"runTags" @> ARRAY[${params.tags.map((t) => `'${t}'`).join(",")}]`);
|
||||
}
|
||||
|
||||
const createdAtFilter = await this.#calculateCreatedAtFilter(url, params.createdAt);
|
||||
|
||||
if (createdAtFilter) {
|
||||
whereClauses.push(`"createdAt" > '${createdAtFilter.toISOString()}'`);
|
||||
}
|
||||
|
||||
const whereClause = whereClauses.join(" AND ");
|
||||
|
||||
const response = await this.#streamRunsWhere(
|
||||
url,
|
||||
environment,
|
||||
whereClause,
|
||||
apiVersion,
|
||||
requestOptions,
|
||||
clientVersion,
|
||||
signal
|
||||
);
|
||||
|
||||
if (createdAtFilter) {
|
||||
const [setCreatedAtFilterError] = await tryCatch(
|
||||
this.#setCreatedAtFilterFromResponse(response, createdAtFilter)
|
||||
);
|
||||
|
||||
if (setCreatedAtFilterError) {
|
||||
logger.error("[realtimeClient] Failed to set createdAt filter", {
|
||||
error: setCreatedAtFilterError,
|
||||
createdAtFilter,
|
||||
responseHeaders: Object.fromEntries(response.headers.entries()),
|
||||
responseStatus: response.status,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
return response;
|
||||
}
|
||||
|
||||
async #calculateCreatedAtFilter(url: URL | string, createdAt?: string) {
|
||||
const duration = createdAt ?? "24h";
|
||||
const $url = new URL(url.toString());
|
||||
const shapeId = extractShapeId($url);
|
||||
|
||||
if (!shapeId) {
|
||||
// This means we need to calculate the createdAt filter and store it in redis after we get back the response
|
||||
const createdAtFilter = safeParseNaturalLanguageDurationAgo(duration);
|
||||
|
||||
// Validate that the createdAt filter is in the past, and not more than the maximum age in the past.
|
||||
// if it's more than the maximum age in the past, just return the maximum age in the past Date
|
||||
if (
|
||||
createdAtFilter &&
|
||||
createdAtFilter < new Date(Date.now() - env.REALTIME_MAXIMUM_CREATED_AT_FILTER_AGE_IN_MS)
|
||||
) {
|
||||
return new Date(Date.now() - env.REALTIME_MAXIMUM_CREATED_AT_FILTER_AGE_IN_MS);
|
||||
}
|
||||
|
||||
return createdAtFilter;
|
||||
} else {
|
||||
// We need to get the createdAt filter value from redis, if there is none we need to return undefined
|
||||
const [createdAtFilterError, createdAtFilter] = await tryCatch(
|
||||
this.#getCreatedAtFilter(shapeId)
|
||||
);
|
||||
|
||||
if (createdAtFilterError) {
|
||||
logger.error("[realtimeClient] Failed to get createdAt filter", {
|
||||
shapeId,
|
||||
error: createdAtFilterError,
|
||||
});
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
return createdAtFilter;
|
||||
}
|
||||
}
|
||||
|
||||
async #getCreatedAtFilter(shapeId: string) {
|
||||
const createdAtFilterCacheResult = await this.cache.createdAtFilter.get(shapeId);
|
||||
|
||||
if (createdAtFilterCacheResult.err) {
|
||||
logger.error("[realtimeClient] Failed to get createdAt filter", {
|
||||
shapeId,
|
||||
error: createdAtFilterCacheResult.err,
|
||||
});
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
if (!createdAtFilterCacheResult.val) {
|
||||
return;
|
||||
}
|
||||
|
||||
return new Date(createdAtFilterCacheResult.val);
|
||||
}
|
||||
|
||||
async #setCreatedAtFilterFromResponse(response: Response, createdAtFilter: Date) {
|
||||
const shapeId = extractShapeIdFromResponse(response);
|
||||
|
||||
if (!shapeId) {
|
||||
return;
|
||||
}
|
||||
|
||||
await this.cache.createdAtFilter.set(shapeId, createdAtFilter.toISOString());
|
||||
}
|
||||
|
||||
async #streamRunsWhere(
|
||||
url: URL | string,
|
||||
environment: RealtimeEnvironment,
|
||||
whereClause: string,
|
||||
apiVersion: API_VERSIONS,
|
||||
requestOptions?: RealtimeRequestOptions,
|
||||
clientVersion?: string,
|
||||
signal?: AbortSignal
|
||||
) {
|
||||
const electricUrl = this.#constructRunsElectricUrl(
|
||||
url,
|
||||
environment,
|
||||
whereClause,
|
||||
requestOptions,
|
||||
clientVersion
|
||||
);
|
||||
|
||||
return this.#performElectricRequest(
|
||||
electricUrl,
|
||||
environment,
|
||||
apiVersion,
|
||||
signal,
|
||||
clientVersion
|
||||
);
|
||||
}
|
||||
|
||||
#constructRunsElectricUrl(
|
||||
url: URL | string,
|
||||
environment: RealtimeEnvironment,
|
||||
whereClause: string,
|
||||
requestOptions?: RealtimeRequestOptions,
|
||||
clientVersion?: string
|
||||
): URL {
|
||||
const $url = new URL(url.toString());
|
||||
|
||||
const electricOrigin = this.#resolveElectricOrigin($url, whereClause, environment.id);
|
||||
const electricUrl = new URL(`${electricOrigin}/v1/shape`);
|
||||
|
||||
// Copy over all the url search params to the electric url
|
||||
$url.searchParams.forEach((value, key) => {
|
||||
if (RESERVED_SEARCH_PARAMS.includes(key)) {
|
||||
return;
|
||||
}
|
||||
|
||||
electricUrl.searchParams.set(key, value);
|
||||
});
|
||||
|
||||
electricUrl.searchParams.set("where", whereClause);
|
||||
electricUrl.searchParams.set("table", 'public."TaskRun"');
|
||||
|
||||
if (!clientVersion) {
|
||||
// If the client version is not provided, that means we're using an older client
|
||||
// This means the client will be sending shape_id instead of handle
|
||||
electricUrl.searchParams.set("handle", electricUrl.searchParams.get("shape_id") ?? "");
|
||||
}
|
||||
|
||||
let skipColumns = getSkipColumns($url.searchParams, requestOptions);
|
||||
|
||||
if (skipColumns.length > 0) {
|
||||
skipColumns = skipColumns.filter((c) => c !== "" && !RESERVED_COLUMNS.includes(c));
|
||||
|
||||
electricUrl.searchParams.set(
|
||||
"columns",
|
||||
DEFAULT_ELECTRIC_COLUMNS.filter((c) => !skipColumns.includes(c))
|
||||
.map((c) => `"${c}"`)
|
||||
.join(",")
|
||||
);
|
||||
} else {
|
||||
electricUrl.searchParams.set(
|
||||
"columns",
|
||||
DEFAULT_ELECTRIC_COLUMNS.map((c) => `"${c}"`).join(",")
|
||||
);
|
||||
}
|
||||
|
||||
return electricUrl;
|
||||
}
|
||||
|
||||
async #performElectricRequest(
|
||||
url: URL,
|
||||
environment: RealtimeEnvironment,
|
||||
apiVersion: API_VERSIONS,
|
||||
signal?: AbortSignal,
|
||||
clientVersion?: string
|
||||
) {
|
||||
const shapeId = extractShapeId(url);
|
||||
|
||||
logger.debug("[realtimeClient] request", {
|
||||
url: url.toString(),
|
||||
});
|
||||
|
||||
const rewriteResponseHeaders: Record<string, string> = clientVersion
|
||||
? {}
|
||||
: { "electric-handle": "electric-shape-id", "electric-offset": "electric-chunk-last-offset" };
|
||||
|
||||
if (!shapeId) {
|
||||
// If the shapeId is not present, we're just getting the initial value
|
||||
return this.#doLongPollingFetch(url, apiVersion, signal, rewriteResponseHeaders);
|
||||
}
|
||||
|
||||
const isLive = isLiveRequestUrl(url);
|
||||
|
||||
if (!isLive) {
|
||||
return this.#doLongPollingFetch(url, apiVersion, signal, rewriteResponseHeaders);
|
||||
}
|
||||
|
||||
const requestId = randomUUID();
|
||||
|
||||
// We now need to wrap the longPollingFetch in a concurrency tracker
|
||||
const concurrencyLimit = await this.cachedLimitProvider.getCachedLimit(
|
||||
environment.organizationId,
|
||||
100_000
|
||||
);
|
||||
|
||||
if (!concurrencyLimit) {
|
||||
logger.error("Failed to get concurrency limit", {
|
||||
organizationId: environment.organizationId,
|
||||
});
|
||||
|
||||
return json({ error: "Failed to get concurrency limit" }, { status: 500 });
|
||||
}
|
||||
|
||||
logger.debug("[realtimeClient] increment and check", {
|
||||
concurrencyLimit,
|
||||
shapeId,
|
||||
requestId,
|
||||
environment: {
|
||||
id: environment.id,
|
||||
organizationId: environment.organizationId,
|
||||
},
|
||||
});
|
||||
|
||||
const canProceed = await this.#incrementAndCheck(environment.id, requestId, concurrencyLimit);
|
||||
|
||||
if (!canProceed) {
|
||||
logger.debug("[realtimeClient] too many concurrent requests", {
|
||||
requestId,
|
||||
environmentId: environment.id,
|
||||
});
|
||||
|
||||
return json({ error: "Too many concurrent requests" }, { status: 429 });
|
||||
}
|
||||
|
||||
try {
|
||||
// ... (rest of your existing code for the long polling request)
|
||||
const response = await this.#doLongPollingFetch(
|
||||
url,
|
||||
apiVersion,
|
||||
signal,
|
||||
rewriteResponseHeaders
|
||||
);
|
||||
|
||||
// If this is the initial request, the response.headers['electric-handle'] will be the shapeId
|
||||
// And we may need to set the "createdAt" filter timestamp keyed by the shapeId
|
||||
// Then in the next request, we will get the createdAt timestamp value via the shapeId and use it to filter the results
|
||||
|
||||
// Decrement the counter after the long polling request is complete
|
||||
await this.#decrementConcurrency(environment.id, requestId);
|
||||
|
||||
return response;
|
||||
} catch (error) {
|
||||
// Decrement the counter if the request fails
|
||||
await this.#decrementConcurrency(environment.id, requestId);
|
||||
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
async #doLongPollingFetch(
|
||||
url: URL,
|
||||
apiVersion: API_VERSIONS,
|
||||
signal?: AbortSignal,
|
||||
rewriteResponseHeaders?: Record<string, string>
|
||||
) {
|
||||
if (apiVersion === CURRENT_API_VERSION) {
|
||||
return longPollingFetch(url.toString(), { signal }, rewriteResponseHeaders);
|
||||
}
|
||||
|
||||
const response = await longPollingFetch(url.toString(), { signal }, rewriteResponseHeaders);
|
||||
|
||||
return this.#rewriteResponseForNoneApiVersion(response);
|
||||
}
|
||||
|
||||
async #rewriteResponseForNoneApiVersion(response: Response) {
|
||||
// Get the raw response body
|
||||
const responseBody = await response.text();
|
||||
|
||||
// Rewrite the response body
|
||||
const rewrittenResponseBody = this.#rewriteResponseBodyForNoneApiVersion(responseBody);
|
||||
|
||||
// Return the rewritten response
|
||||
return new Response(rewrittenResponseBody, {
|
||||
status: response.status,
|
||||
headers: response.headers,
|
||||
});
|
||||
}
|
||||
|
||||
// Rewrites "status":"DEQUEUED" to "status":"EXECUTING"
|
||||
#rewriteResponseBodyForNoneApiVersion(responseBody: string) {
|
||||
return responseBody.replace(/"status":"DEQUEUED"/g, '"status":"EXECUTING"');
|
||||
}
|
||||
|
||||
async #incrementAndCheck(environmentId: string, requestId: string, limit: number) {
|
||||
const key = this.#getKey(environmentId);
|
||||
const now = Date.now();
|
||||
|
||||
const result = await this.redis.incrementAndCheckConcurrency(
|
||||
key,
|
||||
now.toString(),
|
||||
requestId,
|
||||
this.expiryTimeInSeconds.toString(), // expiry time
|
||||
(now - this.expiryTimeInSeconds * 1000).toString(), // cutoff time
|
||||
limit.toString()
|
||||
);
|
||||
|
||||
return result === 1;
|
||||
}
|
||||
|
||||
async #decrementConcurrency(environmentId: string, requestId: string) {
|
||||
logger.debug("[realtimeClient] decrement", {
|
||||
requestId,
|
||||
environmentId,
|
||||
});
|
||||
|
||||
const key = this.#getKey(environmentId);
|
||||
|
||||
await this.redis.zrem(key, requestId);
|
||||
}
|
||||
|
||||
#getKey(environmentId: string): string {
|
||||
return `${this.options.keyPrefix}:${environmentId}`;
|
||||
}
|
||||
|
||||
#resolveElectricOrigin(url: URL, whereClause: string, environmentId: string) {
|
||||
if (typeof this.options.electricOrigin === "string") {
|
||||
return this.options.electricOrigin;
|
||||
}
|
||||
|
||||
const shardKey = this.#getShardKey(whereClause, environmentId);
|
||||
|
||||
const index = jumpHash(shardKey, this.options.electricOrigin.length);
|
||||
|
||||
const origin = this.options.electricOrigin[index] ?? this.options.electricOrigin[0];
|
||||
|
||||
logger.debug("[realtimeClient] resolveElectricOrigin", {
|
||||
whereClause,
|
||||
environmentId,
|
||||
shardKey,
|
||||
index,
|
||||
electricOrigin: origin,
|
||||
});
|
||||
|
||||
return origin;
|
||||
}
|
||||
|
||||
#getShardKey(whereClause: string, environmentId: string) {
|
||||
return [environmentId, whereClause].join(":");
|
||||
}
|
||||
|
||||
#registerCommands() {
|
||||
this.redis.defineCommand("incrementAndCheckConcurrency", {
|
||||
numberOfKeys: 1,
|
||||
lua: /* lua */ `
|
||||
local concurrencyKey = KEYS[1]
|
||||
|
||||
local timestamp = tonumber(ARGV[1])
|
||||
local requestId = ARGV[2]
|
||||
local expiryTime = tonumber(ARGV[3])
|
||||
local cutoffTime = tonumber(ARGV[4])
|
||||
local limit = tonumber(ARGV[5])
|
||||
|
||||
-- Remove expired entries
|
||||
redis.call('ZREMRANGEBYSCORE', concurrencyKey, '-inf', cutoffTime)
|
||||
|
||||
-- Add the new request to the sorted set
|
||||
redis.call('ZADD', concurrencyKey, timestamp, requestId)
|
||||
|
||||
-- Set the expiry time on the key
|
||||
redis.call('EXPIRE', concurrencyKey, expiryTime)
|
||||
|
||||
-- Get the total number of concurrent requests
|
||||
local totalRequests = redis.call('ZCARD', concurrencyKey)
|
||||
|
||||
-- Check if the limit has been exceeded
|
||||
if totalRequests > limit then
|
||||
-- Remove the request we just added
|
||||
redis.call('ZREM', concurrencyKey, requestId)
|
||||
return 0
|
||||
end
|
||||
|
||||
-- Return 1 to indicate success
|
||||
return 1
|
||||
`,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
function extractShapeId(url: URL) {
|
||||
return url.searchParams.get("handle") ?? url.searchParams.get("shape_id");
|
||||
}
|
||||
|
||||
function extractShapeIdFromResponse(response: Response) {
|
||||
return response.headers.get("electric-handle");
|
||||
}
|
||||
|
||||
function isLiveRequestUrl(url: URL) {
|
||||
return url.searchParams.has("live") && url.searchParams.get("live") === "true";
|
||||
}
|
||||
|
||||
declare module "ioredis" {
|
||||
interface RedisCommander<Context> {
|
||||
incrementAndCheckConcurrency(
|
||||
key: string,
|
||||
timestamp: string,
|
||||
requestId: string,
|
||||
expiryTime: string,
|
||||
cutoffTime: string,
|
||||
limit: string,
|
||||
callback?: Callback<number>
|
||||
): Result<number, Context>;
|
||||
}
|
||||
}
|
||||
|
||||
function getSkipColumns(searchParams: URLSearchParams, requestOptions?: RealtimeRequestOptions) {
|
||||
if (requestOptions?.skipColumns) {
|
||||
return requestOptions.skipColumns;
|
||||
}
|
||||
|
||||
const skipColumnsRaw = searchParams.get("skipColumns");
|
||||
|
||||
if (skipColumnsRaw) {
|
||||
return skipColumnsRaw.split(",").map((c) => c.trim());
|
||||
}
|
||||
|
||||
return [];
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user