From bcc6995bc7e890729f2b0f8909ce2fa057e163eb Mon Sep 17 00:00:00 2001 From: John Yeo Date: Fri, 24 Apr 2026 18:50:31 +0100 Subject: [PATCH] fix: improved otel tracing with tenant context --- AGENTS.md | 16 +- ai | 2 +- bun.lock | 1 + server/package.json | 1 + .../revenueCat/revenuecatWebhookRouter.ts | 2 + .../external/stripe/stripeWebhookRouter.ts | 3 + .../external/vercel/vercelWebhookRouter.ts | 2 + server/src/honoMiddlewares/baseMiddleware.ts | 1 + server/src/honoMiddlewares/traceMiddleware.ts | 62 ++++---- server/src/initHono.ts | 2 - server/src/instrumentation.ts | 29 +++- .../src/internal/checkouts/checkoutRouter.ts | 4 + server/src/internal/misc/trmnl/trmnlRouter.ts | 2 + server/src/queue/processMessage.ts | 12 +- server/src/routers/apiRouter.ts | 2 + server/src/routers/internalRouter.ts | 2 + server/src/utils/logging/loggerTypes.ts | 2 + .../src/utils/otel/TenantAttrSpanProcessor.ts | 33 ++++ server/src/utils/otel/tenantContext.ts | 37 +++++ server/src/utils/otel/withWorkerSpan.ts | 56 +++++++ server/src/workers.ts | 5 +- ...mediate-switch-entities-edge-cases.test.ts | 143 ++++++++++++++++++ 22 files changed, 381 insertions(+), 38 deletions(-) create mode 100644 server/src/utils/otel/TenantAttrSpanProcessor.ts create mode 100644 server/src/utils/otel/tenantContext.ts create mode 100644 server/src/utils/otel/withWorkerSpan.ts create mode 100644 server/tests/integration/billing/attach/immediate-switch/entities/immediate-switch-entities-edge-cases.test.ts diff --git a/AGENTS.md b/AGENTS.md index b837dc641..3311968f3 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -49,7 +49,17 @@ There is a legacy-compatibility exception in the adjust-balance flow: Projects maintain state in `.context//` folders across sessions. Tasks are optional parallel workstreams within a project. -### Reading context (session start) +### Default: NOT interacting with a project +**Unless the user explicitly mentions a project or task by name, assume the current conversation is NOT associated with any project.** Session-start hooks may surface a list of active projects as reference material — that alone is NOT a signal that the current work belongs to any of them. + +Do not: +- Read `.context/**` files proactively +- Write, update, or append to any project's STATUS.md / DECISIONS.md / sessions +- Assume a script, audit, or change is part of a project just because it touches files related to one (e.g. a user-of-framework script is not part of the framework's project) + +Only engage with `.context//` when the user explicitly references the project, opens a task in it, or asks for project-tracking actions. + +### Reading context (when a project IS referenced) When the user mentions a project or task name and `.context//` exists: 1. Read project STATUS.md first (20-30 line "resume card") 2. If the project has `tasks/`, list active tasks @@ -58,7 +68,7 @@ When the user mentions a project or task name and `.context//` exists: 5. Do NOT read everything upfront. Use progressive disclosure. ### Updating context (at breakpoints, NOT continuously) -Update at these moments ONLY: +Only update when the current work IS part of a project (see default-off rule above). When it is, update at these moments ONLY: - Phase or milestone completed - Architectural decision made (append to DECISIONS.md) - User says they're done or switching tasks @@ -67,6 +77,8 @@ Update at these moments ONLY: Do NOT update context during normal coding work. Work first, compact at breakpoints. +A STATUS.md entry should record changes to the project itself — not one-off work that merely uses the project (e.g. writing a consumer script of a framework is not a framework-project update). + ### Compaction quality STATUS.md must be: - Correct (reflects actual current state, not stale) diff --git a/ai b/ai index c598ec94e..e66bb8079 160000 --- a/ai +++ b/ai @@ -1 +1 @@ -Subproject commit c598ec94e5cb46ed10c6af4ba810f53cd564ef2c +Subproject commit e66bb80796bf8387e4fbcc550f9d520eaa182f6e diff --git a/bun.lock b/bun.lock index 5fa99fcec..9661b038e 100644 --- a/bun.lock +++ b/bun.lock @@ -349,6 +349,7 @@ "@opentelemetry/api": "^1.9.0", "@opentelemetry/exporter-trace-otlp-http": "^0.213.0", "@opentelemetry/sdk-node": "^0.202.0", + "@opentelemetry/sdk-trace-base": "^2.6.0", "@posthog/ai": "^7.4.2", "@puzzmo/revenue-cat-webhook-types": "^1.1.0", "@react-email/components": "^0.0.42", diff --git a/server/package.json b/server/package.json index c05cb9743..4249a9bbd 100644 --- a/server/package.json +++ b/server/package.json @@ -70,6 +70,7 @@ "@opentelemetry/api": "^1.9.0", "@opentelemetry/exporter-trace-otlp-http": "^0.213.0", "@opentelemetry/sdk-node": "^0.202.0", + "@opentelemetry/sdk-trace-base": "^2.6.0", "@posthog/ai": "^7.4.2", "@puzzmo/revenue-cat-webhook-types": "^1.1.0", "@react-email/components": "^0.0.42", diff --git a/server/src/external/revenueCat/revenuecatWebhookRouter.ts b/server/src/external/revenueCat/revenuecatWebhookRouter.ts index 274c75e19..d7f050767 100644 --- a/server/src/external/revenueCat/revenuecatWebhookRouter.ts +++ b/server/src/external/revenueCat/revenuecatWebhookRouter.ts @@ -9,6 +9,7 @@ import type { WebhookUnCancellation, } from "@puzzmo/revenue-cat-webhook-types"; import { type Context, Hono } from "hono"; +import { traceEnrichMiddleware } from "@/honoMiddlewares/traceMiddleware.js"; import { getRevenuecatWebhookSecret } from "./misc/getRevenuecatWebhookSecret"; import { revenuecatLogMiddleware, @@ -31,6 +32,7 @@ revenuecatWebhookRouter.post( revenuecatSeederMiddleware, revenuecatLogMiddleware, revenuecatWebhookRefreshMiddleware, + traceEnrichMiddleware, async (c: Context) => { const ctx = c.get("ctx"); const { logger, org, env } = ctx; diff --git a/server/src/external/stripe/stripeWebhookRouter.ts b/server/src/external/stripe/stripeWebhookRouter.ts index 9a07d41ce..434f6b232 100644 --- a/server/src/external/stripe/stripeWebhookRouter.ts +++ b/server/src/external/stripe/stripeWebhookRouter.ts @@ -1,5 +1,6 @@ import { Hono } from "hono"; import { stripeLoggerMiddleware } from "@/external/stripe/webhookMiddlewares/stripeLoggerMiddleware.js"; +import { traceEnrichMiddleware } from "@/honoMiddlewares/traceMiddleware.js"; import { handleStripeWebhookEvent } from "./handleStripeWebhookEvent.js"; import { stripeConnectSeederMiddleware } from "./webhookMiddlewares/stripeConnectSeederMiddleware.js"; import { stripeIdempotencyMiddleware } from "./webhookMiddlewares/stripeIdempotencyMiddleware.js"; @@ -19,6 +20,7 @@ stripeWebhookRouter.post( stripeSyncMiddleware, stripeToAutumnCustomerMiddleware, stripeLoggerMiddleware, + traceEnrichMiddleware, stripeIdempotencyMiddleware, handleStripeWebhookEvent, ); @@ -31,6 +33,7 @@ stripeWebhookRouter.post( stripeSyncMiddleware, stripeToAutumnCustomerMiddleware, stripeLoggerMiddleware, + traceEnrichMiddleware, stripeIdempotencyMiddleware, handleStripeWebhookEvent, ); diff --git a/server/src/external/vercel/vercelWebhookRouter.ts b/server/src/external/vercel/vercelWebhookRouter.ts index 97e525661..6cdd9582d 100644 --- a/server/src/external/vercel/vercelWebhookRouter.ts +++ b/server/src/external/vercel/vercelWebhookRouter.ts @@ -2,6 +2,7 @@ import { AppEnv } from "@autumn/shared"; import { Hono } from "hono"; import { handleRotateResourceSecret } from "@/external/vercel/handlers/resources/handleRotateResourceSecret.js"; import { analyticsMiddleware } from "@/honoMiddlewares/analyticsMiddleware.js"; +import { traceEnrichMiddleware } from "@/honoMiddlewares/traceMiddleware.js"; import type { HonoEnv } from "@/honoUtils/HonoEnv.js"; import { sendCustomSvixEvent } from "../svix/svixHelpers.js"; import { handleListBillingPlansPerInstall } from "./handlers/handleListBillingPlans.js"; @@ -40,6 +41,7 @@ vercelWebhookRouter.use( "/:orgId/:env/*", vercelSeederMiddleware, analyticsMiddleware, + traceEnrichMiddleware, ); // Product-level plans (no integrationConfigurationId in path) diff --git a/server/src/honoMiddlewares/baseMiddleware.ts b/server/src/honoMiddlewares/baseMiddleware.ts index 3f683ffc9..031b255e4 100644 --- a/server/src/honoMiddlewares/baseMiddleware.ts +++ b/server/src/honoMiddlewares/baseMiddleware.ts @@ -59,6 +59,7 @@ export const baseMiddleware = async (c: Context, next: Next) => { entity_id: entityId, user_agent: c.req.header("user-agent"), ip_address: c.req.header("x-forwarded-for"), + region: process.env.AWS_REGION, query: c.req.query(), body, diff --git a/server/src/honoMiddlewares/traceMiddleware.ts b/server/src/honoMiddlewares/traceMiddleware.ts index 80c7d4d36..61a840d0a 100644 --- a/server/src/honoMiddlewares/traceMiddleware.ts +++ b/server/src/honoMiddlewares/traceMiddleware.ts @@ -1,40 +1,46 @@ import { context, trace } from "@opentelemetry/api"; -import type { Context, Next } from "hono"; +import type { MiddlewareHandler } from "hono"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; import { isFullSubjectRolloutEnabled } from "@/internal/misc/rollouts/fullSubjectRolloutUtils.js"; -import type { HonoEnv } from "@/honoUtils/HonoEnv.js"; +import { + type TenantAttrs, + withTenantContext, +} from "@/utils/otel/tenantContext.js"; /** - * Enriches the active OTel span (created by @hono/otel) with - * request-specific context like org, env, and customer info. + * Stamps tenant attrs onto the root HTTP span, then wraps the rest of the + * chain in an OTel context carrying the same attrs so TenantAttrSpanProcessor + * propagates them to every child span. Must run AFTER auth/seeder middleware + * that populates ctx. Works on any router whose ctx extends AutumnContext + * (apiRouter, internalRouter, stripe/vercel/revenueCat webhook routers, etc). */ -export const traceEnrichMiddleware = async ( - c: Context, - next: Next, -) => { - await next(); +export const traceEnrichMiddleware: MiddlewareHandler = async (c, next) => { + const ctx = c.get("ctx") as AutumnContext | undefined; + if (!ctx) return next(); - const span = trace.getSpan(context.active()); - if (!span) return; - - const ctx = c.get("ctx"); - - const attrs: Record = { + const attrs: TenantAttrs = { req_id: ctx.id, + org_id: ctx.org?.id, + org_slug: ctx.org?.slug, + env: ctx.env, + customer_id: ctx.customerId, + entity_id: ctx.entityId, + user_id: ctx.userId || undefined, + auth_type: ctx.authType, + api_version: ctx.apiVersion?.semver, + region: process.env.AWS_REGION, + full_subject_rollout_enabled: ctx.org + ? isFullSubjectRolloutEnabled({ ctx }) + : undefined, }; - if (ctx.org) { - attrs.org_id = ctx.org.id; - attrs.org_slug = ctx.org.slug; + const rootSpan = trace.getSpan(context.active()); + if (rootSpan) { + for (const [key, value] of Object.entries(attrs)) { + if (value === undefined) continue; + rootSpan.setAttribute(key, value); + } } - if (ctx.env) { - attrs.env = ctx.env; - } - - if (ctx.customerId) { - attrs.customer_id = ctx.customerId; - attrs.full_subject_rollout_enabled = isFullSubjectRolloutEnabled({ ctx }); - } - - span.setAttributes(attrs); + return withTenantContext({ attrs, fn: () => next() }); }; diff --git a/server/src/initHono.ts b/server/src/initHono.ts index 1e5b36ec5..1aa062f61 100644 --- a/server/src/initHono.ts +++ b/server/src/initHono.ts @@ -14,7 +14,6 @@ import { vercelWebhookRouter } from "./external/vercel/vercelWebhookRouter.js"; import { baseMiddleware } from "./honoMiddlewares/baseMiddleware.js"; import { errorMiddleware } from "./honoMiddlewares/errorMiddleware.js"; import { replicaDbMiddleware } from "./honoMiddlewares/replicaDbMiddleware.js"; -import { traceEnrichMiddleware } from "./honoMiddlewares/traceMiddleware.js"; import type { HonoEnv } from "./honoUtils/HonoEnv.js"; import { handleHealthCheck } from "./honoUtils/handleHealthCheck.js"; import { handleReadyCheck } from "./honoUtils/handleReadyCheck.js"; @@ -95,7 +94,6 @@ export const createHonoApp = () => { ); app.use("*", baseMiddleware); app.use("*", replicaDbMiddleware); - app.use("*", traceEnrichMiddleware); // Public endpoint to get OAuth client name (for consent page) app.get("/oauth/client/:client_id", async (c) => { diff --git a/server/src/instrumentation.ts b/server/src/instrumentation.ts index 0c0ee0274..5e9774e91 100644 --- a/server/src/instrumentation.ts +++ b/server/src/instrumentation.ts @@ -2,8 +2,11 @@ import "dotenv/config"; import { DiagConsoleLogger, DiagLogLevel, diag } from "@opentelemetry/api"; import { OTLPTraceExporter } from "@opentelemetry/exporter-trace-otlp-http"; import { NodeSDK } from "@opentelemetry/sdk-node"; +import { BatchSpanProcessor } from "@opentelemetry/sdk-trace-base"; +import { TenantAttrSpanProcessor } from "./utils/otel/TenantAttrSpanProcessor.js"; -// Surface OTel internal warnings/errors so export failures are visible +// Surface OTel internal warnings/errors (export failures, auth issues, etc.) +// but not DEBUG-level span dumps. diag.setLogger(new DiagConsoleLogger(), DiagLogLevel.WARN); let sdk: NodeSDK | null = null; @@ -20,14 +23,34 @@ if (process.env.AXIOM_TOKEN) { }, }); + // Passing `spanProcessors` replaces the default pipeline — NodeSDK does NOT + // auto-add a BatchSpanProcessor for `traceExporter` when `spanProcessors` + // is set. We must wire the exporter processor explicitly. + // Dev: short 1s flush for fast feedback. Prod: default 5s for throughput. + const isDev = process.env.NODE_ENV !== "production"; + const exportProcessor = new BatchSpanProcessor(traceExporter, { + scheduledDelayMillis: isDev ? 1000 : 5000, + }); + // No auto-instrumentations — Bun doesn't support require-in-the-middle. // Stripe, Drizzle, and Redis are instrumented via manual patchers in utils/otel/. sdk = new NodeSDK({ - traceExporter, + spanProcessors: [new TenantAttrSpanProcessor(), exportProcessor], }); - console.log("Starting OpenTelemetry"); sdk.start(); + + // Flush spans on SIGTERM/SIGINT so dev restarts (nodemon) and prod rollouts + // don't swallow in-flight batches. + const shutdown = async () => { + try { + await sdk?.shutdown(); + } catch (err) { + console.error("[otel] shutdown error", err); + } + }; + process.once("SIGTERM", shutdown); + process.once("SIGINT", shutdown); } export { sdk as otelSdk }; diff --git a/server/src/internal/checkouts/checkoutRouter.ts b/server/src/internal/checkouts/checkoutRouter.ts index a30354c3a..b25f573d9 100644 --- a/server/src/internal/checkouts/checkoutRouter.ts +++ b/server/src/internal/checkouts/checkoutRouter.ts @@ -1,5 +1,6 @@ import { Hono } from "hono"; import { analyticsMiddleware } from "@/honoMiddlewares/analyticsMiddleware"; +import { traceEnrichMiddleware } from "@/honoMiddlewares/traceMiddleware"; import type { HonoEnv } from "@/honoUtils/HonoEnv"; import { handleConfirmCheckout } from "./handlers/handleConfirmCheckout"; import { handleGetCheckout } from "./handlers/handleGetCheckout"; @@ -24,6 +25,9 @@ publicCheckoutRouter.use("/:checkout_id/*", checkoutMiddleware); // publicCheckoutRouter.use("/:checkout_id", analyticsMiddleware); publicCheckoutRouter.use("/:checkout_id/*", analyticsMiddleware); +publicCheckoutRouter.use("/:checkout_id", traceEnrichMiddleware); +publicCheckoutRouter.use("/:checkout_id/*", traceEnrichMiddleware); + // Routes publicCheckoutRouter.get("/:checkout_id", ...handleGetCheckout); publicCheckoutRouter.post("/:checkout_id/preview", ...handlePreviewCheckout); diff --git a/server/src/internal/misc/trmnl/trmnlRouter.ts b/server/src/internal/misc/trmnl/trmnlRouter.ts index 42b53943a..3e0cf2df0 100644 --- a/server/src/internal/misc/trmnl/trmnlRouter.ts +++ b/server/src/internal/misc/trmnl/trmnlRouter.ts @@ -1,5 +1,6 @@ import { Hono } from "hono"; import { rateLimiter } from "hono-rate-limiter"; +import { traceEnrichMiddleware } from "@/honoMiddlewares/traceMiddleware.js"; import type { HonoEnv } from "@/honoUtils/HonoEnv.js"; import { handleGenerateTrmnlScreen } from "./handlers/handleGenerateTrmnlScreen.js"; import { handleGetTrmnlDeviceId } from "./handlers/handleGetTrmnlDeviceId.js"; @@ -28,5 +29,6 @@ publicTrmnlRouter.post( "/screen", trmnlScreenLimiter, trmnlAuthMiddleware, + traceEnrichMiddleware, ...handleGenerateTrmnlScreen, ); diff --git a/server/src/queue/processMessage.ts b/server/src/queue/processMessage.ts index 02f2cac57..229d68683 100644 --- a/server/src/queue/processMessage.ts +++ b/server/src/queue/processMessage.ts @@ -27,6 +27,7 @@ import { runTriggerCheckoutReward } from "@/internal/rewards/triggerCheckoutRewa import { generateId } from "@/utils/genUtils.js"; import { addWorkflowToLogs } from "@/utils/logging/addContextToLogs.js"; import { maskExtraLogs } from "@/utils/logging/maskExtraLogs.js"; +import { withWorkerSpan } from "@/utils/otel/withWorkerSpan.js"; import { setSentryTags } from "../external/sentry/sentryUtils.js"; import { createWorkerContext } from "./createWorkerContext.js"; import { JobName } from "./JobName.js"; @@ -285,7 +286,16 @@ export const processMessage = async ({ }; try { - await executeJob(); + await withWorkerSpan({ + workflowName: job.name, + workflowId: message.MessageId ?? generateId("job"), + tenantAttrs: { + org_id: job.data?.orgId, + env: job.data?.env, + customer_id: job.data?.customerId, + }, + fn: executeJob, + }); } catch (error) { const errorLogger = workerCtx?.logger ?? workerLogger; // Sync jobs: re-throw infrastructure errors so the message stays in SQS. diff --git a/server/src/routers/apiRouter.ts b/server/src/routers/apiRouter.ts index f6054506c..663b1f86d 100644 --- a/server/src/routers/apiRouter.ts +++ b/server/src/routers/apiRouter.ts @@ -17,6 +17,7 @@ import { requestBlockMiddleware } from "../honoMiddlewares/requestBlockMiddlewar import { responseFilterMiddleware } from "../honoMiddlewares/responseFilter/responseFilterMiddleware.js"; import { rolloutMiddleware } from "../honoMiddlewares/rolloutMiddleware.js"; import { secretKeyMiddleware } from "../honoMiddlewares/secretKeyMiddleware.js"; +import { traceEnrichMiddleware } from "../honoMiddlewares/traceMiddleware.js"; import type { HonoEnv } from "../honoUtils/HonoEnv.js"; import { redemptionRouter, @@ -48,6 +49,7 @@ apiRouter.use("*", requestBlockMiddleware); apiRouter.use("*", orgConfigMiddleware); apiRouter.use("*", rolloutMiddleware); apiRouter.use("*", apiVersionMiddleware); +apiRouter.use("*", traceEnrichMiddleware); apiRouter.use("*", refreshCacheMiddleware); apiRouter.use("*", refreshProductsCacheMiddleware); apiRouter.use("*", analyticsMiddleware); diff --git a/server/src/routers/internalRouter.ts b/server/src/routers/internalRouter.ts index 7bd7a359e..a792697d9 100644 --- a/server/src/routers/internalRouter.ts +++ b/server/src/routers/internalRouter.ts @@ -8,6 +8,7 @@ import { betterAuthMiddleware } from "../honoMiddlewares/betterAuthMiddleware"; import { orgConfigMiddleware } from "../honoMiddlewares/orgConfigMiddleware"; import { queryMiddleware } from "../honoMiddlewares/queryMiddleware"; import { refreshCacheMiddleware } from "../honoMiddlewares/refreshCacheMiddleware"; +import { traceEnrichMiddleware } from "../honoMiddlewares/traceMiddleware"; import type { HonoEnv } from "../honoUtils/HonoEnv"; import { honoAdminRouter } from "../internal/admin/adminRouter"; import { internalAnalyticsRouter } from "../internal/analytics/internalAnalyticsRouter"; @@ -26,6 +27,7 @@ export const internalRouter = new Hono(); internalRouter.use("*", betterAuthMiddleware); internalRouter.use("*", orgConfigMiddleware); internalRouter.use("*", apiVersionMiddleware); +internalRouter.use("*", traceEnrichMiddleware); internalRouter.use("*", analyticsMiddleware); internalRouter.use("*", refreshCacheMiddleware); internalRouter.use("*", queryMiddleware()); diff --git a/server/src/utils/logging/loggerTypes.ts b/server/src/utils/logging/loggerTypes.ts index 34e75431e..97b3ac25e 100644 --- a/server/src/utils/logging/loggerTypes.ts +++ b/server/src/utils/logging/loggerTypes.ts @@ -12,6 +12,8 @@ export type LogRequestContext = { user_agent?: string; ip_address?: string; + region?: string; + // New fields query: Record; body: unknown; diff --git a/server/src/utils/otel/TenantAttrSpanProcessor.ts b/server/src/utils/otel/TenantAttrSpanProcessor.ts new file mode 100644 index 000000000..745157950 --- /dev/null +++ b/server/src/utils/otel/TenantAttrSpanProcessor.ts @@ -0,0 +1,33 @@ +import type { Context as OtelContext } from "@opentelemetry/api"; +import type { + ReadableSpan, + Span, + SpanProcessor, +} from "@opentelemetry/sdk-trace-base"; +import { getTenantAttrs } from "./tenantContext.js"; + +/** + * Reads tenant attrs from the OTel parent context at span start and stamps + * them onto every child span. No-op when parent context has no tenant data. + */ +export class TenantAttrSpanProcessor implements SpanProcessor { + onStart(span: Span, parentContext: OtelContext): void { + const attrs = getTenantAttrs(parentContext); + if (!attrs) return; + + for (const [key, value] of Object.entries(attrs)) { + if (value === undefined) continue; + span.setAttribute(key, value); + } + } + + onEnd(_span: ReadableSpan): void {} + + shutdown(): Promise { + return Promise.resolve(); + } + + forceFlush(): Promise { + return Promise.resolve(); + } +} diff --git a/server/src/utils/otel/tenantContext.ts b/server/src/utils/otel/tenantContext.ts new file mode 100644 index 000000000..d2f255605 --- /dev/null +++ b/server/src/utils/otel/tenantContext.ts @@ -0,0 +1,37 @@ +import { + context, + createContextKey, + type Context as OtelContext, +} from "@opentelemetry/api"; + +export type TenantAttrs = { + req_id?: string; + org_id?: string; + org_slug?: string; + env?: string; + customer_id?: string; + entity_id?: string; + user_id?: string; + auth_type?: string; + api_version?: string; + region?: string; + full_subject_rollout_enabled?: boolean; +}; + +export const TENANT_CONTEXT_KEY = createContextKey("autumn.tenant"); + +export const withTenantContext = ({ + attrs, + fn, +}: { + attrs: TenantAttrs; + fn: () => T; +}): T => { + const activeContext = context.active(); + const nextContext = activeContext.setValue(TENANT_CONTEXT_KEY, attrs); + return context.with(nextContext, fn); +}; + +export const getTenantAttrs = (ctx: OtelContext): TenantAttrs | undefined => { + return ctx.getValue(TENANT_CONTEXT_KEY) as TenantAttrs | undefined; +}; diff --git a/server/src/utils/otel/withWorkerSpan.ts b/server/src/utils/otel/withWorkerSpan.ts new file mode 100644 index 000000000..a568e6ce8 --- /dev/null +++ b/server/src/utils/otel/withWorkerSpan.ts @@ -0,0 +1,56 @@ +import { SpanKind, SpanStatusCode, trace } from "@opentelemetry/api"; +import { type TenantAttrs, withTenantContext } from "./tenantContext.js"; + +const tracer = trace.getTracer("autumn.worker"); + +/** + * Root span wrapper for worker/cron jobs. Creates a span named after the + * workflow, stamps workflow + tenant attrs on it AND on the OTel context so + * every child span inherits them. Records exceptions and sets ERROR status on + * throw. Always re-throws so SQS retry semantics are preserved. + */ +export const withWorkerSpan = async ({ + workflowName, + workflowId, + tenantAttrs, + fn, +}: { + workflowName: string; + workflowId: string; + tenantAttrs?: TenantAttrs; + fn: () => Promise; +}): Promise => { + return tracer.startActiveSpan( + `worker.${workflowName}`, + { kind: SpanKind.CONSUMER }, + async (span) => { + span.setAttributes({ + workflow_id: workflowId, + workflow_name: workflowName, + }); + + if (tenantAttrs) { + for (const [key, value] of Object.entries(tenantAttrs)) { + if (value === undefined) continue; + span.setAttribute(key, value); + } + } + + try { + return await withTenantContext({ + attrs: tenantAttrs ?? {}, + fn, + }); + } catch (err) { + span.recordException(err as Error); + span.setStatus({ + code: SpanStatusCode.ERROR, + message: err instanceof Error ? err.message : String(err), + }); + throw err; + } finally { + span.end(); + } + }, + ); +}; diff --git a/server/src/workers.ts b/server/src/workers.ts index bac4ac87f..b86f0d7ee 100644 --- a/server/src/workers.ts +++ b/server/src/workers.ts @@ -90,7 +90,10 @@ if (cluster.isPrimary) { } }); } else { - // Worker process + // Worker process — start OTel SDK so child spans (Stripe/Redis/Drizzle/ + // withSpan/withWorkerSpan) export to Axiom. + await import("./instrumentation.js"); + const startupStartedAt = Date.now(); const queueImplementation = "SQS"; startMemoryMonitor("worker", 60_000); diff --git a/server/tests/integration/billing/attach/immediate-switch/entities/immediate-switch-entities-edge-cases.test.ts b/server/tests/integration/billing/attach/immediate-switch/entities/immediate-switch-entities-edge-cases.test.ts new file mode 100644 index 000000000..9374ba896 --- /dev/null +++ b/server/tests/integration/billing/attach/immediate-switch/entities/immediate-switch-entities-edge-cases.test.ts @@ -0,0 +1,143 @@ +/** + * Immediate Switch Entity Edge Cases (Attach V2) + * + * Regression tests for edge cases observed in production (Mintlify customer) + * where attaching products with different billing intervals to different + * entities after a time advance produced duplicate charges. + */ + +import { expect, test } from "bun:test"; +import type { ApiCustomerV3, ApiEntityV0 } from "@autumn/shared"; +import { expectCustomerInvoiceCorrect } from "@tests/integration/billing/utils/expectCustomerInvoiceCorrect"; +import { + expectCustomerProducts, + expectProductActive, +} from "@tests/integration/billing/utils/expectCustomerProductCorrect"; +import { expectStripeSubscriptionCorrect } from "@tests/integration/billing/utils/expectStripeSubCorrect"; +import { calculateProrationFromPeriod } from "@tests/integration/billing/utils/proration"; +import { getStripeSubscription } from "@tests/integration/billing/utils/stripeSubscriptionUtils"; +import { advanceTestClock } from "@tests/utils/stripeUtils"; +import { TestFeature } from "@tests/setup/v2Features"; +import { items } from "@tests/utils/fixtures/items"; +import { products } from "@tests/utils/fixtures/products"; +import { initScenario, s } from "@tests/utils/testInitUtils/initScenario"; +import chalk from "chalk"; + +// ═══════════════════════════════════════════════════════════════════════════════ +// TEST 1: Entity 1 pro annual, advance 3 weeks, attach pro monthly to entity 2 +// ═══════════════════════════════════════════════════════════════════════════════ + +/** + * Scenario: + * - Entity 1 attaches pro annual ($200/yr) + * - Advance test clock by 3 weeks + * - Entity 2 attaches pro monthly ($20/mo) + * + * Expected Result: + * - Entity 1 still has pro annual, entity 2 has pro monthly + * - Second attach invoice is exactly $20 for the monthly item only + * - No duplicate charge for the annual item that was already billed on entity 1 + */ +test.concurrent(`${chalk.yellowBright("immediate-switch-entities-edge-cases 1: entity 1 annual, advance 3 weeks, attach monthly to entity 2")}`, async () => { + const customerId = "imm-switch-ent-edge-annual-then-monthly"; + + const proAnnualMessages = items.monthlyMessages({ includedUsage: 500 }); + const proAnnual = products.proAnnual({ + id: "pro-annual", + items: [proAnnualMessages], + }); + + const proMonthlyMessages = items.monthlyMessages({ includedUsage: 500 }); + const proMonthly = products.pro({ + id: "pro-monthly", + items: [proMonthlyMessages], + }); + + const { autumnV1, ctx, entities, advancedTo, testClockId } = await initScenario({ + customerId, + setup: [ + s.customer({ paymentMethod: "success" }), + s.products({ list: [proAnnual, proMonthly] }), + s.entities({ count: 2, featureId: TestFeature.Users }), + ], + actions: [s.billing.attach({ productId: proAnnual.id, entityIndex: 0 })], + }); + + // Set the annual sub's billing_cycle_anchor to exactly one year from its + // start date — this mirrors the Mintlify production scenario. + const { + stripeCli, + subscription: annualSub, + } = await getStripeSubscription({ customerId }); + const annualStart = annualSub.start_date; + const oneYearFromStart = annualStart + 365 * 24 * 60 * 60; + await stripeCli.subscriptions.update(annualSub.id, { + trial_end: oneYearFromStart, + proration_behavior: "none", + }); + + // Now advance the clock 3 weeks into the annual cycle + await advanceTestClock({ + stripeCli, + testClockId: testClockId!, + numberOfWeeks: 3, + }); + const advancedToAfter = advancedTo + 3 * 7 * 24 * 60 * 60 * 1000; + + // 1. Preview attach pro monthly to entity 2 + const preview = await autumnV1.billing.previewAttach({ + customer_id: customerId, + product_id: proMonthly.id, + entity_id: entities[1].id, + }); + + // 2. Attach pro monthly to entity 2 + await autumnV1.billing.attach({ + customer_id: customerId, + product_id: proMonthly.id, + entity_id: entities[1].id, + redirect_mode: "if_required", + }); + + // 3. Compute the EXACT expected proration of $20 against the monthly sub's + // period via getStripeSubscription + calculateProrationFromPeriod + const { billingPeriod } = await getStripeSubscription({ customerId }); + const expectedMonthlyCharge = calculateProrationFromPeriod({ + billingPeriod, + advancedTo: advancedToAfter, + amount: 20, + }); + + // Preview must match the exact prorated amount — proves no duplicate annual + expect(preview.total).toBe(expectedMonthlyCharge); + + // 4. Invoice count: annual attach + monthly attach = 2 invoices; most recent + // invoice total must exactly equal the computed prorated amount + const customer = await autumnV1.customers.get(customerId); + expectCustomerInvoiceCorrect({ + customer, + count: 2, + latestTotal: expectedMonthlyCharge, + latestStatus: "paid", + }); + + // 4. Entity 1 still has pro annual, entity 2 has pro monthly + const entity1 = await autumnV1.entities.get( + customerId, + entities[0].id, + ); + const entity2 = await autumnV1.entities.get( + customerId, + entities[1].id, + ); + + await expectProductActive({ customer: entity1, productId: proAnnual.id }); + await expectCustomerProducts({ + customer: entity2, + active: [proMonthly.id], + notPresent: [proAnnual.id], + }); + + // 5. Stripe subscriptions reflect the two independent entity subs + await expectStripeSubscriptionCorrect({ ctx, customerId }); +});