diff --git a/apps/api/src/lib/autumn-tracker.ts b/apps/api/src/lib/autumn-tracker.ts index 249f4db62..10eb219bf 100644 --- a/apps/api/src/lib/autumn-tracker.ts +++ b/apps/api/src/lib/autumn-tracker.ts @@ -12,7 +12,13 @@ export interface TrackTokenUsageOptions { readonly inputTokens: number readonly outputTokens: number readonly idempotencyKey: string - readonly source: "triage" + /** + * Namespaces the idempotency key per producer, so a triage run and a Slack + * bot turn can never collide on a shared id. `slack-agent` is the + * Railway-hosted bot (apps/slack-agent), which reports its own Workers AI + * usage through `POST /internal/slack/workspaces/:teamId/usage`. + */ + readonly source: "triage" | "slack-agent" } interface TrackEvent { @@ -21,12 +27,51 @@ interface TrackEvent { readonly idempotencyKey: string } +/** + * Why a call reported nothing. Callers that care about billing coverage (the + * Slack usage endpoint) turn anything other than `tracked` into a log line and + * a span annotation; callers that don't (triage, investigations) ignore it. + * + * `no-credentials` is the dangerous one: `AUTUMN_SECRET_KEY` is bound with + * `optionalSecret()`, so a key missing from an Infisical environment silently + * disables billing with no CI error at all. + */ +export type TrackTokenUsageOutcome = + | "tracked" + | "partial" + | "failed" + | "no-credentials" + | "excluded-org" + | "no-tokens" + +export interface TrackTokenUsageResult { + readonly outcome: TrackTokenUsageOutcome + /** Events posted to Autumn (0 when nothing was attempted). */ + readonly attempted: number + readonly succeeded: number + /** First failure reason, for the caller's log line. Absent when nothing failed. */ + readonly failureReason?: string +} + +/** + * Hard ceiling on one Autumn round-trip. `trackTokenUsage` is awaited inline by + * the Slack usage endpoint, and at one request per model step a stalled Autumn + * would otherwise tie up a Worker request per step indefinitely. Billing is + * best-effort by contract, so giving up beats hanging. + */ +const AUTUMN_TIMEOUT_MS = 5_000 + +interface PostResult { + readonly ok: boolean + readonly reason?: string +} + const postTrack = async ( apiUrl: string, secretKey: string, customerId: string, event: TrackEvent, -): Promise => { +): Promise => { try { const response = await fetch(`${apiUrl}/v1/track`, { method: "POST", @@ -40,28 +85,40 @@ const postTrack = async ( value: event.value, idempotency_key: event.idempotencyKey, }), + signal: AbortSignal.timeout(AUTUMN_TIMEOUT_MS), }) if (!response.ok) { const body = await response.text().catch(() => "") - console.warn( - `[autumn-tracker] track failed: ${response.status} feature=${event.featureId} body=${body}`, - ) + const reason = `${event.featureId}: HTTP ${response.status} ${body.slice(0, 200)}`.trim() + console.warn(`[autumn-tracker] track failed: ${reason}`) + return { ok: false, reason } } + return { ok: true } } catch (error) { - console.warn( - `[autumn-tracker] track error feature=${event.featureId}: ${error instanceof Error ? error.message : String(error)}`, - ) + const reason = `${event.featureId}: ${error instanceof Error ? error.message : String(error)}` + console.warn(`[autumn-tracker] track error ${reason}`) + return { ok: false, reason } } } +/** + * Never rejects — billing is best-effort and must not fail the work that + * produced the tokens. It DOES report what happened, so a caller can tell + * "billed" from "silently billed nothing for a week"; callers that predate the + * return value can keep ignoring it. + */ export const trackTokenUsage = async ( env: Record, { orgId, inputTokens, outputTokens, idempotencyKey, source }: TrackTokenUsageOptions, -): Promise => { +): Promise => { const secretKey = typeof env.AUTUMN_SECRET_KEY === "string" ? env.AUTUMN_SECRET_KEY : undefined - if (!secretKey) return - if (typeof env.MAPLE_DEFAULT_ORG_ID === "string" && orgId === env.MAPLE_DEFAULT_ORG_ID) return - if (inputTokens <= 0 && outputTokens <= 0) return + if (!secretKey) return { outcome: "no-credentials", attempted: 0, succeeded: 0 } + if (typeof env.MAPLE_DEFAULT_ORG_ID === "string" && orgId === env.MAPLE_DEFAULT_ORG_ID) { + return { outcome: "excluded-org", attempted: 0, succeeded: 0 } + } + if (inputTokens <= 0 && outputTokens <= 0) { + return { outcome: "no-tokens", attempted: 0, succeeded: 0 } + } const apiUrl = ( typeof env.AUTUMN_API_URL === "string" ? env.AUTUMN_API_URL : DEFAULT_AUTUMN_API_URL @@ -82,5 +139,23 @@ export const trackTokenUsage = async ( }) } - await Promise.allSettled(events.map((event) => postTrack(apiUrl, secretKey, orgId, event))) + // `postTrack` resolves with an outcome rather than throwing, so `allSettled` + // here is belt-and-braces against a future throw path — not the error handling. + const settled = await Promise.allSettled( + events.map((event) => postTrack(apiUrl, secretKey, orgId, event)), + ) + const results: PostResult[] = settled.map((entry) => + entry.status === "fulfilled" ? entry.value : { ok: false, reason: String(entry.reason) }, + ) + const succeeded = results.filter((result) => result.ok).length + const failureReason = results.find((result) => !result.ok)?.reason + const outcome: TrackTokenUsageOutcome = + succeeded === results.length ? "tracked" : succeeded === 0 ? "failed" : "partial" + + return { + outcome, + attempted: results.length, + succeeded, + ...(failureReason === undefined ? {} : { failureReason }), + } } diff --git a/apps/api/src/routes/slack-integration.http.test.ts b/apps/api/src/routes/slack-integration.http.test.ts index d3076c39f..c7e9da056 100644 --- a/apps/api/src/routes/slack-integration.http.test.ts +++ b/apps/api/src/routes/slack-integration.http.test.ts @@ -79,7 +79,16 @@ const insertWorkspace = async ( ) } -const makeConfig = (tokens: { slack?: string; shared?: string } = {}) => +/** Router-layer knobs: which internal bearers are configured, and Autumn wiring for usage tracking. */ +interface HarnessConfig { + slack?: string + shared?: string + autumnSecret?: string + autumnApiUrl?: string + defaultOrgId?: string +} + +const makeConfig = (tokens: HarnessConfig = {}) => ConfigProvider.layer( ConfigProvider.fromUnknown({ PORT: "3472", @@ -87,16 +96,18 @@ const makeConfig = (tokens: { slack?: string; shared?: string } = {}) => TINYBIRD_TOKEN: "test-token", MAPLE_AUTH_MODE: "self_hosted", MAPLE_ROOT_PASSWORD: "test-root-password", - MAPLE_DEFAULT_ORG_ID: "default", + MAPLE_DEFAULT_ORG_ID: tokens.defaultOrgId ?? "default", MAPLE_INGEST_KEY_ENCRYPTION_KEY: ENCRYPTION_KEY.toString("base64"), MAPLE_INGEST_KEY_LOOKUP_HMAC_KEY: "maple-test-lookup-secret", MAPLE_APP_BASE_URL: "https://web.localhost", ...(tokens.slack !== undefined ? { SLACK_INTERNAL_SERVICE_TOKEN: tokens.slack } : {}), ...(tokens.shared !== undefined ? { INTERNAL_SERVICE_TOKEN: tokens.shared } : {}), + ...(tokens.autumnSecret !== undefined ? { AUTUMN_SECRET_KEY: tokens.autumnSecret } : {}), + ...(tokens.autumnApiUrl !== undefined ? { AUTUMN_API_URL: tokens.autumnApiUrl } : {}), }), ) -const makeRouterLayer = (testDb: TestDb, tokens: { slack?: string; shared?: string } = {}) => +const makeRouterLayer = (testDb: TestDb, tokens: HarnessConfig = {}) => SlackInternalRouter.pipe( Layer.provide(SlackIntegrationService.layer), Layer.provide(Layer.mergeAll(ApiKeysService.layer, OAuthStateRepository.layer)), @@ -139,7 +150,7 @@ const postRevoke = ( /** Run `body` against a web handler for the internal router, disposing after. */ const withHandler = ( testDb: TestDb, - tokens: { slack?: string; shared?: string }, + tokens: HarnessConfig, body: (handler: (request: Request) => Promise) => Effect.Effect, ) => { const { handler, dispose } = HttpRouter.toWebHandler(makeRouterLayer(testDb, tokens), { @@ -495,3 +506,436 @@ describe("SlackInternalRouter (revoke)", () => { ) }) }) + +// --------------------------------------------------------------------------- +// Usage reporting. The Railway-hosted bot runs its own model (Cloudflare +// Workers AI), so its AI spend reaches billing only by being reported here and +// attributed to the org bound to the Slack team. +// --------------------------------------------------------------------------- + +const AUTUMN_URL = "https://autumn.test" + +interface AutumnCall { + readonly customerId: string + readonly featureId: string + readonly value: number + readonly idempotencyKey: string +} + +/** + * Stubs `globalThis.fetch` for the duration of `body`. `trackTokenUsage` issues + * a bare `fetch` (it is shared with Workers-runtime call sites and takes no + * injected client), so the global is the only seam. + */ +const withAutumnStub = ( + body: (calls: AutumnCall[]) => Effect.Effect, + respond: () => Response = () => new Response("{}", { status: 200 }), +) => + Effect.suspend(() => { + const realFetch = globalThis.fetch + const calls: AutumnCall[] = [] + globalThis.fetch = (async (input: string | URL | Request, init?: RequestInit) => { + const url = String(input instanceof Request ? input.url : input) + if (!url.startsWith(AUTUMN_URL)) return realFetch(input as RequestInfo, init) + const parsed = JSON.parse(String(init?.body)) as { + customer_id: string + feature_id: string + value: number + idempotency_key: string + } + calls.push({ + customerId: parsed.customer_id, + featureId: parsed.feature_id, + value: parsed.value, + idempotencyKey: parsed.idempotency_key, + }) + return respond() + }) as typeof fetch + return body(calls).pipe( + Effect.ensuring( + Effect.sync(() => { + globalThis.fetch = realFetch + }), + ), + ) + }) + +const postUsage = ( + handler: (request: Request) => Promise, + teamId: string, + body: unknown, + bearer?: string, +) => + Effect.promise(() => + handler( + new Request(`http://api.localhost${TEAM_PATH}/${teamId}/usage`, { + method: "POST", + headers: { + "content-type": "application/json", + ...(bearer !== undefined ? { authorization: bearer } : {}), + }, + body: typeof body === "string" ? body : JSON.stringify(body), + }), + ), + ) + +const USAGE_TOKENS: HarnessConfig = { + slack: "slack-secret-token", + autumnSecret: "autumn-secret", + autumnApiUrl: AUTUMN_URL, +} + +const USAGE_BEARER = "Bearer maple_svc_slack-secret-token" + +describe("SlackInternalRouter (usage)", () => { + it.effect("tracks input and output tokens against the team's org", () => { + const testDb = createTestDb(trackedDbs) + return Effect.gen(function* () { + yield* Effect.promise(() => + insertWorkspace(testDb, { + id: "sw_us1", + orgId: "org_us1", + teamId: "T-US1", + teamName: "UsOrg1", + botToken: "xoxb-us1", + apiKey: "maple_ak_us1", + }), + ) + yield* withAutumnStub((calls) => + withHandler( + testDb, + USAGE_TOKENS, + Effect.fnUntraced(function* (handler) { + const response = yield* postUsage( + handler, + "T-US1", + { + idempotencyKey: "sess_1:turn_1:0", + inputTokens: 120, + outputTokens: 34, + model: "@cf/zai-org/glm-5.2", + }, + USAGE_BEARER, + ) + assert.strictEqual(response.status, 202) + assert.deepStrictEqual(yield* Effect.promise(() => response.json()), { + tracked: true, + outcome: "tracked", + }) + // Billed to the org the team resolves to — never to the team id. + // The key is team-namespaced and source-tagged so a bot step and a + // triage run can never collide. + assert.deepStrictEqual( + [...calls].sort((a, b) => a.featureId.localeCompare(b.featureId)), + [ + { + customerId: "org_us1", + featureId: "ai_input_tokens", + value: 120, + idempotencyKey: "slack:T-US1:sess_1:turn_1:0:slack-agent:input", + }, + { + customerId: "org_us1", + featureId: "ai_output_tokens", + value: 34, + idempotencyKey: "slack:T-US1:sess_1:turn_1:0:slack-agent:output", + }, + ], + ) + }), + ), + ) + }).pipe(Effect.provide(testDb.layer)) + }) + + it.effect("skips a zero-token side rather than tracking a zero", () => { + const testDb = createTestDb(trackedDbs) + return Effect.gen(function* () { + yield* Effect.promise(() => + insertWorkspace(testDb, { + id: "sw_us2", + orgId: "org_us2", + teamId: "T-US2", + teamName: "UsOrg2", + botToken: "xoxb-us2", + apiKey: "maple_ak_us2", + }), + ) + yield* withAutumnStub((calls) => + withHandler( + testDb, + USAGE_TOKENS, + Effect.fnUntraced(function* (handler) { + const response = yield* postUsage( + handler, + "T-US2", + { idempotencyKey: "k", inputTokens: 7, outputTokens: 0 }, + USAGE_BEARER, + ) + assert.strictEqual(response.status, 202) + assert.strictEqual(calls.length, 1) + assert.strictEqual(calls[0]?.featureId, "ai_input_tokens") + }), + ), + ) + }).pipe(Effect.provide(testDb.layer)) + }) + + it.effect("answers 202 but reports NOT tracked when the billing provider fails", () => { + const testDb = createTestDb(trackedDbs) + return Effect.gen(function* () { + yield* Effect.promise(() => + insertWorkspace(testDb, { + id: "sw_us3", + orgId: "org_us3", + teamId: "T-US3", + teamName: "UsOrg3", + botToken: "xoxb-us3", + apiKey: "maple_ak_us3", + }), + ) + yield* withAutumnStub( + (calls) => + withHandler( + testDb, + USAGE_TOKENS, + Effect.fnUntraced(function* (handler) { + const response = yield* postUsage( + handler, + "T-US3", + { idempotencyKey: "k", inputTokens: 5, outputTokens: 5 }, + USAGE_BEARER, + ) + // Still 202 — the Slack reply already went out and a retry would + // only re-post something Autumn may have recorded. But the body + // must not claim success, or a week-long Autumn outage is + // indistinguishable from a week of billed traffic. + assert.strictEqual(response.status, 202) + assert.deepStrictEqual(yield* Effect.promise(() => response.json()), { + tracked: false, + outcome: "failed", + }) + // It really did attempt both writes — this asserts the failure + // path, not merely that nothing threw. + assert.strictEqual(calls.length, 2) + }), + ), + () => new Response("nope", { status: 500 }), + ) + }).pipe(Effect.provide(testDb.layer)) + }) + + it.effect("reports NOT tracked when Autumn is unconfigured — optionalSecret can drop the key", () => { + const testDb = createTestDb(trackedDbs) + return Effect.gen(function* () { + yield* Effect.promise(() => + insertWorkspace(testDb, { + id: "sw_us3b", + orgId: "org_us3b", + teamId: "T-US3B", + teamName: "UsOrg3b", + botToken: "xoxb-us3b", + apiKey: "maple_ak_us3b", + }), + ) + yield* withAutumnStub((calls) => + // No `autumnSecret`: exactly what a PR preview looks like when the key + // is missing from the Infisical environment `optionalSecret()` reads. + withHandler( + testDb, + { slack: "slack-secret-token", autumnApiUrl: AUTUMN_URL }, + Effect.fnUntraced(function* (handler) { + const response = yield* postUsage( + handler, + "T-US3B", + { idempotencyKey: "k", inputTokens: 5, outputTokens: 5 }, + USAGE_BEARER, + ) + assert.strictEqual(response.status, 202) + assert.deepStrictEqual(yield* Effect.promise(() => response.json()), { + tracked: false, + outcome: "no-credentials", + }) + assert.strictEqual(calls.length, 0) + }), + ), + ) + }).pipe(Effect.provide(testDb.layer)) + }) + + it.effect("does not bill the self-hosted default org", () => { + const testDb = createTestDb(trackedDbs) + return Effect.gen(function* () { + yield* Effect.promise(() => + insertWorkspace(testDb, { + id: "sw_us4", + orgId: "org_default", + teamId: "T-US4", + teamName: "UsOrg4", + botToken: "xoxb-us4", + apiKey: "maple_ak_us4", + }), + ) + yield* withAutumnStub((calls) => + withHandler( + testDb, + { ...USAGE_TOKENS, defaultOrgId: "org_default" }, + Effect.fnUntraced(function* (handler) { + const response = yield* postUsage( + handler, + "T-US4", + { idempotencyKey: "k", inputTokens: 9, outputTokens: 9 }, + USAGE_BEARER, + ) + assert.strictEqual(response.status, 202) + assert.strictEqual(calls.length, 0) + }), + ), + ) + }).pipe(Effect.provide(testDb.layer)) + }) + + it.effect("rejects an unknown or revoked team with 404", () => { + const testDb = createTestDb(trackedDbs) + return Effect.gen(function* () { + yield* Effect.promise(() => + insertWorkspace(testDb, { + id: "sw_us5", + orgId: "org_us5", + teamId: "T-US5", + teamName: "UsOrg5", + botToken: "xoxb-us5", + apiKey: "maple_ak_us5", + revoked: true, + }), + ) + yield* withAutumnStub((calls) => + withHandler( + testDb, + USAGE_TOKENS, + Effect.fnUntraced(function* (handler) { + const body = { idempotencyKey: "k", inputTokens: 1, outputTokens: 1 } + assert.strictEqual((yield* postUsage(handler, "T-US5", body, USAGE_BEARER)).status, 404) + assert.strictEqual((yield* postUsage(handler, "T-NOPE", body, USAGE_BEARER)).status, 404) + assert.strictEqual(calls.length, 0) + }), + ), + ) + }).pipe(Effect.provide(testDb.layer)) + }) + + it.effect("requires the dedicated slack bearer", () => { + const testDb = createTestDb(trackedDbs) + return withHandler( + testDb, + { ...USAGE_TOKENS, shared: "shared-secret-token" }, + Effect.fnUntraced(function* (handler) { + const body = { idempotencyKey: "k", inputTokens: 1, outputTokens: 1 } + assert.strictEqual((yield* postUsage(handler, "T-US1", body)).status, 401) + assert.strictEqual( + (yield* postUsage(handler, "T-US1", body, "Bearer maple_svc_wrong")).status, + 401, + ) + // The shared internal token is NOT accepted here, same as resolve/revoke. + assert.strictEqual( + (yield* postUsage(handler, "T-US1", body, "Bearer maple_svc_shared-secret-token")).status, + 401, + ) + }), + ) + }) + + it.effect("answers 401 when the slack internal token is unset", () => { + const testDb = createTestDb(trackedDbs) + return withHandler( + testDb, + { autumnSecret: "autumn-secret", autumnApiUrl: AUTUMN_URL }, + Effect.fnUntraced(function* (handler) { + const response = yield* postUsage( + handler, + "T-US1", + { idempotencyKey: "k", inputTokens: 1, outputTokens: 1 }, + USAGE_BEARER, + ) + assert.strictEqual(response.status, 401) + }), + ) + }) + + it.effect("rejects a malformed body with 400", () => { + const testDb = createTestDb(trackedDbs) + return withHandler( + testDb, + USAGE_TOKENS, + Effect.fnUntraced(function* (handler) { + const cases: unknown[] = [ + {}, + { idempotencyKey: "", inputTokens: 1, outputTokens: 1 }, + { idempotencyKey: "k", inputTokens: -1, outputTokens: 1 }, + { idempotencyKey: "k", inputTokens: 1.5, outputTokens: 1 }, + { idempotencyKey: "k", inputTokens: "1", outputTokens: 1 }, + { idempotencyKey: "k", inputTokens: 1 }, + "not json", + ] + for (const body of cases) { + const response = yield* postUsage(handler, "T-US1", body, USAGE_BEARER) + assert.strictEqual(response.status, 400, `expected 400 for ${JSON.stringify(body)}`) + } + }), + ) + }) + + it.effect("rejects an absurd token count — this is a financial write on a shared bearer", () => { + const testDb = createTestDb(trackedDbs) + return Effect.gen(function* () { + yield* Effect.promise(() => + insertWorkspace(testDb, { + id: "sw_us6", + orgId: "org_us6", + teamId: "T-US6", + teamName: "UsOrg6", + botToken: "xoxb-us6", + apiKey: "maple_ak_us6", + }), + ) + yield* withAutumnStub((calls) => + withHandler( + testDb, + USAGE_TOKENS, + Effect.fnUntraced(function* (handler) { + // The bearer is one shared service token that works against every + // org and there is no rate limiting, so a single POST of 1e15 would + // otherwise land on a customer's meter and blow their invoice. + const overCap = 8 * 262_144 + 1 + const cases: unknown[] = [ + { idempotencyKey: "k", inputTokens: overCap, outputTokens: 1 }, + { idempotencyKey: "k", inputTokens: 1, outputTokens: overCap }, + { idempotencyKey: "k", inputTokens: 1e15, outputTokens: 1 }, + { idempotencyKey: "k", inputTokens: Number.MAX_SAFE_INTEGER, outputTokens: 1 }, + ] + for (const body of cases) { + const response = yield* postUsage(handler, "T-US6", body, USAGE_BEARER) + assert.strictEqual( + response.status, + 400, + `expected 400 for ${JSON.stringify(body)}`, + ) + } + // Nothing reached the billing provider. + assert.strictEqual(calls.length, 0) + + // A plausible large step still goes through — the cap must not + // reject real traffic after a context-window bump. + const atCap = yield* postUsage( + handler, + "T-US6", + { idempotencyKey: "k2", inputTokens: 8 * 262_144, outputTokens: 0 }, + USAGE_BEARER, + ) + assert.strictEqual(atCap.status, 202) + }), + ), + ) + }).pipe(Effect.provide(testDb.layer)) + }) +}) diff --git a/apps/api/src/routes/slack-integration.http.ts b/apps/api/src/routes/slack-integration.http.ts index 08ad5643b..c3d03cdc8 100644 --- a/apps/api/src/routes/slack-integration.http.ts +++ b/apps/api/src/routes/slack-integration.http.ts @@ -1,7 +1,8 @@ import { HttpRouter, HttpServerRequest, HttpServerResponse } from "effect/unstable/http" -import { Effect, Option, Redacted, Schema } from "effect" +import { Cause, Effect, Option, Redacted, Schema } from "effect" import { timingSafeEqual } from "node:crypto" import { SlackBotResolutionResponseSchema } from "@maple/domain/http" +import { trackTokenUsage } from "../lib/autumn-tracker" import { Env } from "../lib/Env" import { SlackIntegrationService, @@ -147,22 +148,34 @@ export const SlackInternalRouter = HttpRouter.use((router) => onSome: Redacted.value, }) - const logAccess = Effect.fnUntraced(function* ( - teamId: string | undefined, - outcome: "found" | "not-found" | "invalid" | "unauthorized" | "unavailable", - status: number, - ) { - yield* Effect.annotateCurrentSpan({ - ...(teamId === undefined ? {} : { teamId }), - outcome, - "http.response.status_code": status, + /** + * Shared by all three internal endpoints, so the line has to name the one + * that produced it: a 401 on `usage` and a 401 on `resolve` mean different + * things, and labelling both "resolve" sends anyone triaging a billing + * anomaly to the wrong endpoint. + */ + const logAccess = (route: "resolve" | "revoke" | "usage") => + Effect.fnUntraced(function* ( + teamId: string | undefined, + outcome: "found" | "not-found" | "invalid" | "unauthorized" | "unavailable", + status: number, + ) { + yield* Effect.annotateCurrentSpan({ + ...(teamId === undefined ? {} : { teamId }), + outcome, + "http.response.status_code": status, + }) + // `resolve` hands out decrypted tokens and `usage` is a financial write + // primitive, so a rejected caller on either is a security signal rather + // than routine traffic. + yield* outcome === "unauthorized" + ? Effect.logWarning(`Slack internal ${route} rejected`, { teamId, route, outcome }) + : Effect.logInfo(`Slack internal ${route} access`, { teamId, route, outcome }) }) - // This endpoint hands out decrypted tokens — a rejected caller is a - // security signal, not routine traffic. - yield* outcome === "unauthorized" - ? Effect.logWarning("Slack internal resolve rejected", { teamId, outcome }) - : Effect.logInfo("Slack internal resolve access", { teamId, outcome }) - }) + + const logResolveAccess = logAccess("resolve") + const logRevokeAccess = logAccess("revoke") + const logUsageAccess = logAccess("usage") const handle = Effect.fn("SlackInternal.resolve")(function* ( req: HttpServerRequest.HttpServerRequest, @@ -170,11 +183,11 @@ export const SlackInternalRouter = HttpRouter.use((router) => // Auth first: everything below (including path-param decoding) must be // unreachable for an unauthenticated caller. if (!internalToken) { - yield* logAccess(undefined, "unauthorized", 401) + yield* logResolveAccess(undefined, "unauthorized", 401) return errorText("Slack internal service token is not configured", 401) } if (!isValidServiceBearer(req.headers.authorization, internalToken)) { - yield* logAccess(undefined, "unauthorized", 401) + yield* logResolveAccess(undefined, "unauthorized", 401) return errorText("Unauthorized", 401) } @@ -185,21 +198,21 @@ export const SlackInternalRouter = HttpRouter.use((router) => : undefined, ) if (Option.isNone(teamIdOption)) { - yield* logAccess(undefined, "invalid", 400) + yield* logResolveAccess(undefined, "invalid", 400) return errorText("Missing teamId", 400) } const teamId = teamIdOption.value return yield* slack.resolveForBot(teamId).pipe( Effect.flatMap((resolved) => - logAccess(teamId, "found", 200).pipe( + logResolveAccess(teamId, "found", 200).pipe( Effect.andThen(encodeBotResolution(resolved).pipe(Effect.orDie)), Effect.flatMap((encoded) => HttpServerResponse.json(encoded)), ), ), Effect.catchTags({ "@maple/http/errors/IntegrationsNotConnectedError": () => - logAccess(teamId, "not-found", 404).pipe( + logResolveAccess(teamId, "not-found", 404).pipe( Effect.as(errorText("No active Slack installation for this team", 404)), ), "@maple/http/errors/IntegrationsPersistenceError": (error) => @@ -207,7 +220,7 @@ export const SlackInternalRouter = HttpRouter.use((router) => teamId, message: error.message, }).pipe( - Effect.andThen(logAccess(teamId, "unavailable", 503)), + Effect.andThen(logResolveAccess(teamId, "unavailable", 503)), Effect.as(errorText("Slack workspace lookup unavailable", 503)), ), }), @@ -242,11 +255,11 @@ export const SlackInternalRouter = HttpRouter.use((router) => req: HttpServerRequest.HttpServerRequest, ) { if (!internalToken) { - yield* logAccess(undefined, "unauthorized", 401) + yield* logRevokeAccess(undefined, "unauthorized", 401) return errorText("Slack internal service token is not configured", 401) } if (!isValidServiceBearer(req.headers.authorization, internalToken)) { - yield* logAccess(undefined, "unauthorized", 401) + yield* logRevokeAccess(undefined, "unauthorized", 401) return errorText("Unauthorized", 401) } @@ -257,7 +270,7 @@ export const SlackInternalRouter = HttpRouter.use((router) => : undefined, ) if (Option.isNone(teamIdOption)) { - yield* logAccess(undefined, "invalid", 400) + yield* logRevokeAccess(undefined, "invalid", 400) return errorText("Missing teamId", 400) } const teamId = teamIdOption.value @@ -267,14 +280,14 @@ export const SlackInternalRouter = HttpRouter.use((router) => Option.liftThrowable(() => JSON.parse(body))().pipe(Option.flatMap(decodeRevokeBody)), ) if (Option.isNone(reasonOption)) { - yield* logAccess(teamId, "invalid", 400) + yield* logRevokeAccess(teamId, "invalid", 400) return errorText('Body must be JSON: { "reason": "app_uninstalled" | "tokens_revoked" }', 400) } const reason: SlackRevocationReason = reasonOption.value.reason return yield* slack.revokeByTeamId(teamId, reason).pipe( Effect.flatMap((result) => - logAccess(teamId, result.revoked ? "found" : "not-found", 200).pipe( + logRevokeAccess(teamId, result.revoked ? "found" : "not-found", 200).pipe( Effect.andThen(HttpServerResponse.json({ revoked: result.revoked })), ), ), @@ -284,14 +297,199 @@ export const SlackInternalRouter = HttpRouter.use((router) => reason, message: error.message, }).pipe( - Effect.andThen(logAccess(teamId, "unavailable", 503)), + Effect.andThen(logRevokeAccess(teamId, "unavailable", 503)), Effect.as(errorText("Slack workspace revoke unavailable", 503)), ), ), ) }) + // --------------------------------------------------------------------- + // AI usage reporting — the Slack bot runs its own model (Cloudflare + // Workers AI) on Railway, so its spend is invisible to billing unless it + // tells us. It reports one call per completed model step (eve's + // `step.completed`, which carries provider-reported token counts) and + // this endpoint attributes it to the team's bound org. + // + // Autumn credentials deliberately stay here rather than on the bot: + // `AUTUMN_SECRET_KEY` is org-agnostic and would let its holder write + // usage for ANY customer. The bot only ever proves it is the bot. + // --------------------------------------------------------------------- + + /** + * Per-step ceiling on a reported token count. This is a financial write — + * the value lands on the org's Autumn meter — and the bearer is a single + * shared service token that works against EVERY org, on a plain + * `HttpRouter` with no rate limiting. Without a ceiling one malformed (or + * malicious) POST of `1e15` blows a customer's plan limit and invoice. + * + * 8x the bot's 262,144-token context window (`modelContextWindowTokens` in + * apps/slack-agent/agent/agent.ts): comfortably above anything one model + * call can legitimately report even after a context-window bump, and far + * below the range where a bogus value does damage. + */ + const MAX_STEP_TOKENS = 8 * 262_144 + + const decodeUsageBody = Schema.decodeUnknownOption( + Schema.Struct({ + /** + * Producer-side de-dupe id, unique per reported model call. Guards + * against a duplicate DELIVERY of one call, not against two distinct + * calls — see `apps/slack-agent/agent/lib/token-usage.ts` for why the + * bot appends a per-process unique component. + */ + idempotencyKey: Schema.String.check(Schema.isMinLength(1), Schema.isTrimmed()), + inputTokens: Schema.Int.check( + Schema.isGreaterThanOrEqualTo(0), + Schema.isLessThanOrEqualTo(MAX_STEP_TOKENS), + ), + outputTokens: Schema.Int.check( + Schema.isGreaterThanOrEqualTo(0), + Schema.isLessThanOrEqualTo(MAX_STEP_TOKENS), + ), + /** Reported for attribution on the span/logs only — Autumn bills tokens, not models. */ + model: Schema.optionalKey(Schema.String), + }), + ) + + const autumnEnv = { + AUTUMN_SECRET_KEY: Option.match(env.AUTUMN_SECRET_KEY, { + onNone: () => undefined, + onSome: Redacted.value, + }), + AUTUMN_API_URL: env.AUTUMN_API_URL, + MAPLE_DEFAULT_ORG_ID: env.MAPLE_DEFAULT_ORG_ID, + } + + const handleUsage = Effect.fn("SlackInternal.usage")(function* ( + req: HttpServerRequest.HttpServerRequest, + ) { + if (!internalToken) { + yield* logUsageAccess(undefined, "unauthorized", 401) + return errorText("Slack internal service token is not configured", 401) + } + if (!isValidServiceBearer(req.headers.authorization, internalToken)) { + yield* logUsageAccess(undefined, "unauthorized", 401) + return errorText("Unauthorized", 401) + } + + const params = yield* HttpRouter.params + const teamIdOption = decodeTeamIdParam( + typeof params.teamId === "string" + ? Option.getOrUndefined(decodeUriComponentOption(params.teamId)) + : undefined, + ) + if (Option.isNone(teamIdOption)) { + yield* logUsageAccess(undefined, "invalid", 400) + return errorText("Missing teamId", 400) + } + const teamId = teamIdOption.value + + const bodyOption = yield* req.text.pipe(Effect.option) + const usageOption = Option.flatMap(bodyOption, (body) => + Option.liftThrowable(() => JSON.parse(body))().pipe(Option.flatMap(decodeUsageBody)), + ) + if (Option.isNone(usageOption)) { + yield* logUsageAccess(teamId, "invalid", 400) + return errorText( + `Body must be JSON: { "idempotencyKey": string, "inputTokens": int 0..${MAX_STEP_TOKENS}, "outputTokens": int 0..${MAX_STEP_TOKENS}, "model"?: string }`, + 400, + ) + } + const usage = usageOption.value + + return yield* slack.resolveOrgForTeam(teamId).pipe( + Effect.flatMap(({ orgId }) => + Effect.gen(function* () { + yield* Effect.annotateCurrentSpan({ + orgId, + "maple.ai.input_tokens": usage.inputTokens, + "maple.ai.output_tokens": usage.outputTokens, + ...(usage.model === undefined ? {} : { "maple.ai.model": usage.model }), + }) + // Namespaced by team so two workspaces can never collide on a + // key minted independently by the bot. + const result = yield* Effect.tryPromise(() => + trackTokenUsage(autumnEnv, { + orgId, + inputTokens: usage.inputTokens, + outputTokens: usage.outputTokens, + idempotencyKey: `slack:${teamId}:${usage.idempotencyKey}`, + source: "slack-agent", + }), + ).pipe( + // `trackTokenUsage` does not reject, so this only catches a defect + // in it; either way an unbilled write must not read as billed. + Effect.catchCause((cause) => + Effect.succeed({ + outcome: "failed" as const, + attempted: 0, + succeeded: 0, + failureReason: Cause.pretty(cause), + }), + ), + ) + + const tracked = result.outcome === "tracked" + yield* Effect.annotateCurrentSpan({ + "maple.billing.outcome": result.outcome, + "maple.billing.events_succeeded": result.succeeded, + "maple.billing.events_attempted": result.attempted, + }) + // A billing write that silently does nothing is the failure mode + // worth paging on: without this, a rotated AUTUMN_SECRET_KEY or a + // week-long Autumn outage looks exactly like success — every + // response 202 `{tracked:true}`, every span Ok, nothing billed. + // `excluded-org` / `no-tokens` are deliberate no-ops, not failures. + if (result.outcome === "failed" || result.outcome === "partial") { + yield* Effect.logError("Slack agent token usage tracking failed", { + teamId, + orgId, + outcome: result.outcome, + succeeded: result.succeeded, + attempted: result.attempted, + reason: result.failureReason, + }) + } else if (result.outcome === "no-credentials") { + // AUTUMN_SECRET_KEY is bound with `optionalSecret()`, so a key + // missing from an Infisical environment disables billing with no + // deploy-time error anywhere. + yield* Effect.logError("Slack agent token usage not billed: Autumn is unconfigured", { + teamId, + orgId, + }) + } + + yield* logUsageAccess(teamId, "found", 202) + // Still 202 regardless: the Slack reply already went out and cannot + // be held or undone, so a 5xx would only buy a retry of something + // Autumn may already have recorded. But `tracked` now reports what + // actually happened instead of asserting success unconditionally. + return yield* HttpServerResponse.json( + { tracked, outcome: result.outcome }, + { status: 202 }, + ) + }), + ), + Effect.catchTags({ + "@maple/http/errors/IntegrationsNotConnectedError": () => + logUsageAccess(teamId, "not-found", 404).pipe( + Effect.as(errorText("No active Slack installation for this team", 404)), + ), + "@maple/http/errors/IntegrationsPersistenceError": (error) => + Effect.logError("Slack internal usage lookup failed", { + teamId, + message: error.message, + }).pipe( + Effect.andThen(logUsageAccess(teamId, "unavailable", 503)), + Effect.as(errorText("Slack workspace lookup unavailable", 503)), + ), + }), + ) + }) + yield* router.add("GET", "/internal/slack/workspaces/:teamId", handle) yield* router.add("POST", "/internal/slack/workspaces/:teamId/revoke", handleRevoke) + yield* router.add("POST", "/internal/slack/workspaces/:teamId/usage", handleUsage) }), ) diff --git a/apps/api/src/routes/v2/v2-test-support.ts b/apps/api/src/routes/v2/v2-test-support.ts index 5aea85e2c..476d5de0b 100644 --- a/apps/api/src/routes/v2/v2-test-support.ts +++ b/apps/api/src/routes/v2/v2-test-support.ts @@ -215,6 +215,7 @@ export const SlackIntegrationServiceStubLayer = Layer.succeed( uninstall: die, listChannels: die, resolveForBot: die, + resolveOrgForTeam: die, revokeByTeamId: die, reconcileWorkspaces: die, }), diff --git a/apps/api/src/services/SlackIntegrationService.ts b/apps/api/src/services/SlackIntegrationService.ts index 6a898bad7..2ff448552 100644 --- a/apps/api/src/services/SlackIntegrationService.ts +++ b/apps/api/src/services/SlackIntegrationService.ts @@ -244,6 +244,19 @@ export interface SlackIntegrationServiceShape { readonly resolveForBot: ( teamId: string, ) => Effect.Effect + /** + * The org a Slack team's active install is bound to — and nothing else. + * + * Deliberately NOT {@link resolveForBot}: usage reporting only needs the + * billing subject, so it must not decrypt (or risk logging) the workspace's + * bot token and full-access Maple API key just to attribute a token count. + */ + readonly resolveOrgForTeam: ( + teamId: string, + ) => Effect.Effect< + { readonly orgId: OrgId }, + IntegrationsNotConnectedError | IntegrationsPersistenceError + > /** * Revoke a workspace binding by Slack team id without calling Slack's * `auth.revoke` — for the two cases where Slack has already told us (or we've @@ -1000,6 +1013,44 @@ const make: Effect.Effect< } satisfies SlackBotResolution }) + const resolveOrgForTeam = Effect.fn("SlackIntegrationService.resolveOrgForTeam")(function* ( + teamId: string, + ) { + yield* Effect.annotateCurrentSpan({ teamId }) + // Only the org column — the secret columns are never read here, so an + // install whose secrets were dropped by `revokeByTeamId` is still matched + // by the `revokedAt IS NULL` filter alone, exactly as elsewhere. + const rowOption: Option.Option<{ orgId: string }> = yield* database + .execute((db) => + db + .select({ orgId: slackWorkspaces.orgId }) + .from(slackWorkspaces) + .where(and(eq(slackWorkspaces.teamId, teamId), isNull(slackWorkspaces.revokedAt))) + .limit(1), + ) + .pipe( + Effect.mapError(toPersistenceError), + Effect.map((rows) => Option.fromNullishOr(rows[0])), + ) + if (Option.isNone(rowOption)) { + return yield* Effect.fail( + new IntegrationsNotConnectedError({ + message: "No active Slack installation for this team", + }), + ) + } + const orgId = yield* decodeOrgId(rowOption.value.orgId).pipe( + Effect.mapError( + (error) => + new IntegrationsPersistenceError({ + message: `Stored Slack workspace has an invalid orgId: ${error.message}`, + }), + ), + ) + yield* Effect.annotateCurrentSpan({ orgId }) + return { orgId } + }) + const revokeByTeamId = Effect.fn("SlackIntegrationService.revokeByTeamId")(function* ( teamId: string, reason: SlackRevocationReason, @@ -1119,6 +1170,7 @@ const make: Effect.Effect< uninstall, listChannels, resolveForBot, + resolveOrgForTeam, revokeByTeamId, reconcileWorkspaces, }) diff --git a/apps/slack-agent/README.md b/apps/slack-agent/README.md index c95e8bac6..2fd31a99f 100644 --- a/apps/slack-agent/README.md +++ b/apps/slack-agent/README.md @@ -404,6 +404,47 @@ Authorization: Bearer maple_svc_{MAPLE_INTERNAL_SERVICE_TOKEN} `agent/lib/maple.ts` wraps this in `resolveWorkspace(teamId)` with an in-memory TTL cache (5 min positive, 30 s negative, in-flight de-dupe). It never caches 5xx/network errors as "not installed". +**AI usage reporting (billing):** this agent runs its own model (Cloudflare Workers AI) on Railway, +outside every path the Maple API already meters — so without reporting, Slack-driven AI spend is +invisible to billing. `agent/hooks/token-usage.ts` subscribes to eve's `step.completed`, the only +hook event carrying provider-reported token counts (one per model call; `turn.completed` has none), +and posts them per step: + +``` +POST {MAPLE_API_BASE_URL}/internal/slack/workspaces/{teamId}/usage +Authorization: Bearer maple_svc_{MAPLE_INTERNAL_SERVICE_TOKEN} +{ "idempotencyKey": "{sessionId}:{turnId}:{stepIndex}:{bootId}-{n}", "inputTokens": 120, + "outputTokens": 34, "model": "@cf/..." } + +202 { "tracked": true|false, "outcome": "tracked"|"failed"|"no-credentials"|… } +404 → team not installed / revoked +400 → malformed body, or a token count above the per-step cap (8x the context window) +``` + +The key's `{bootId}-{n}` suffix is load-bearing: eve derives `turnId` from a `sequence` that a +*terminal* model failure never advances, so `{sessionId}:{turnId}:{stepIndex}` alone repeats across +a user's retry and the billing provider silently drops the second one. See +`agent/lib/token-usage.ts` for the full trade-off. + +The 202 is unconditional — the Slack reply has already gone out — but `tracked`/`outcome` report +whether the billing write actually landed, so an outage or a missing credential is visible in Maple's +logs and spans instead of looking exactly like success. + +The Maple API resolves the team's org and forwards it to the billing provider. The billing +credential deliberately stays on the API side: it is org-agnostic, so holding it would mean being +able to write usage against any customer — the bot only ever proves it is the bot. No new env vars: +this reuses `MAPLE_API_BASE_URL` + `MAPLE_INTERNAL_SERVICE_TOKEN`. + +Reported per step rather than accumulated per turn, because per-turn batching needs mutable state +keyed by turn id that leaks on a turn which never terminates and loses the whole turn's spend if +the process dies mid-turn. Steps with zero tokens (a truncated stream — `workers-ai-provider` +reports all-zeros rather than `undefined` when the provider never sends its usage chunk) and +sessions with no Slack team (local dev, no org to bill) are skipped client-side. Zero-token skips +emit a throttled `usage_report_zero_tokens` warning — if the provider never surfaces usage for the +configured model, every step is skipped and the feature bills exactly nothing, which must not be +silent. `agent/lib/token-usage.ts` never throws: a failed billing write is logged and dropped, never +costing the user their reply. + **Inbound (webhook verification):** `agent/channels/slack.ts` installs a custom `webhookVerifier`. Because a `webhookVerifier` is set, eve skips its built-in signing-secret check and this verifier owns verification: it HMAC-verifies the Slack **v0** signature (`v0:{timestamp}:{rawBody}`, SHA-256, diff --git a/apps/slack-agent/agent/agent.ts b/apps/slack-agent/agent/agent.ts index c53fd707b..3d0c15eaf 100644 --- a/apps/slack-agent/agent/agent.ts +++ b/apps/slack-agent/agent/agent.ts @@ -1,5 +1,6 @@ import { defineAgent } from "eve" import { createWorkersAI } from "workers-ai-provider" +import { workersAiModelId } from "#lib/model.js" /** * Cloudflare Workers AI over its REST API (no Workers runtime / AI binding @@ -22,11 +23,8 @@ if (missingModelEnv.length > 0 && !isEveBuildInvocation) { ) } -/** - * Must support tool calling **while streaming** — eve's harness is tool-driven and - * always streams. - */ -const modelId = process.env.WORKERS_AI_MODEL ?? "@cf/zai-org/glm-5.2" +/** Shared with AI-usage reporting — see `agent/lib/model.ts`. */ +const modelId = workersAiModelId() const contextWindowTokens = Number(process.env.WORKERS_AI_CONTEXT_WINDOW ?? 262_144) /** diff --git a/apps/slack-agent/agent/hooks/token-usage.ts b/apps/slack-agent/agent/hooks/token-usage.ts new file mode 100644 index 000000000..d1ffe84da --- /dev/null +++ b/apps/slack-agent/agent/hooks/token-usage.ts @@ -0,0 +1,38 @@ +import { defineHook, type HookContext } from "eve/hooks" +import { trackStepUsage } from "#lib/token-usage.js" + +/** + * Reports the bot's Workers AI token spend to Maple, which bills it to the org + * bound to the Slack workspace. + * + * `step.completed` is the only hook event eve gives usage on (one per model + * call, carrying the provider's own counts); `turn.completed` has none. See + * `agent/lib/token-usage.ts` for why this reports per step rather than + * accumulating per turn. + * + * Fired without awaiting, exactly like the uninstall forward in + * `agent/lib/uninstall-detection.ts`: hook handlers run inline in the turn, and + * a billing round-trip must not sit between two model calls of a live Slack + * reply. `trackStepUsage` never rejects, so the un-awaited promise cannot + * become an unhandled rejection. + */ + +function teamIdOf(ctx: HookContext): string | undefined { + const team = ctx.session.auth.current?.attributes?.team_id + return typeof team === "string" && team.length > 0 ? team : undefined +} + +export default defineHook({ + events: { + "step.completed"(event, ctx) { + void trackStepUsage({ + sessionId: ctx.session.id, + teamId: teamIdOf(ctx), + turnId: event.data.turnId, + stepIndex: event.data.stepIndex, + inputTokens: event.data.usage?.inputTokens, + outputTokens: event.data.usage?.outputTokens, + }) + }, + }, +}) diff --git a/apps/slack-agent/agent/lib/maple.test.ts b/apps/slack-agent/agent/lib/maple.test.ts index 947999435..7e1e0fd8f 100644 --- a/apps/slack-agent/agent/lib/maple.test.ts +++ b/apps/slack-agent/agent/lib/maple.test.ts @@ -11,6 +11,7 @@ delete process.env.SLACK_BOT_TOKEN import { installFetchStub } from "./fetch-stub.js" import { notifyMapleRevocation, + reportTokenUsage, resetWorkspaceCacheForTests, resolveBotToken, resolveWorkspace, @@ -310,6 +311,44 @@ describe("notifyMapleRevocation", () => { }) }) +// ── reportTokenUsage: AI spend attribution ───────────────────────────────── + +describe("reportTokenUsage", () => { + afterEach(() => { + globalThis.fetch = realFetch + }) + + const usage = { + idempotencyKey: "sess_1:turn_1:2", + inputTokens: 120, + outputTokens: 34, + model: "@cf/zai-org/glm-5.2", + } + + test("POSTs the usage to the team's usage endpoint with the internal bearer", async () => { + const stub = installFetchStub(() => new Response(null, { status: 202 })) + + await reportTokenUsage("T1", usage) + + expect(stub.calls.length).toBe(1) + expect(stub.calls[0]?.method).toBe("POST") + expect(stub.calls[0]?.url).toBe("https://maple-api.test/internal/slack/workspaces/T1/usage") + expect(stub.calls[0]?.headers.authorization).toBe("Bearer maple_svc_test-service-token") + expect(JSON.parse(String(stub.calls[0]?.body))).toEqual(usage) + }) + + test("percent-encodes the team id into the path", async () => { + const stub = installFetchStub(() => new Response(null, { status: 202 })) + await reportTokenUsage("T/1", usage) + expect(stub.calls[0]?.url).toBe("https://maple-api.test/internal/slack/workspaces/T%2F1/usage") + }) + + test("throws on a non-ok response", async () => { + installFetchStub(() => new Response(null, { status: 503 })) + await expect(reportTokenUsage("T1", usage)).rejects.toThrow(/HTTP 503/) + }) +}) + // ── resolveBotToken: patched credential context → env fallback ────────────── describe("resolveBotToken", () => { diff --git a/apps/slack-agent/agent/lib/maple.ts b/apps/slack-agent/agent/lib/maple.ts index 25ac128b2..de78ae869 100644 --- a/apps/slack-agent/agent/lib/maple.ts +++ b/apps/slack-agent/agent/lib/maple.ts @@ -171,6 +171,51 @@ export async function notifyMapleRevocation( } } +/** + * Same budget and the same reasoning as {@link NOTIFY_REVOCATION_TIMEOUT_MS}: + * usage reporting is fired without awaiting from a hook, so nothing in the + * Slack reply path is waiting on it. + */ +const REPORT_USAGE_TIMEOUT_MS = 10_000 + +/** One model call's provider-reported token usage, as sent to Maple. */ +export interface TokenUsageReport { + /** De-dupe id, unique per model call — Maple namespaces it by team. */ + readonly idempotencyKey: string + readonly inputTokens: number + readonly outputTokens: number + readonly model: string +} + +/** + * Reports one model call's token usage to Maple, which attributes it to the + * org bound to this Slack team and forwards it to the billing provider. + * + * The bot deliberately does NOT talk to the billing provider itself: that + * credential is org-agnostic, so holding it would mean being able to write + * usage against any customer. This endpoint lets the bot prove only that it is + * the bot, and Maple decides whose spend it is. + * + * Throws on transport/server errors (mirrors the other two internal calls) — + * the caller (`agent/lib/token-usage.ts`) catches and logs, because a billing + * write must never cost the user their Slack reply. + */ +export async function reportTokenUsage(teamId: string, usage: TokenUsageReport): Promise { + const url = `${mapleApiBaseUrl()}/internal/slack/workspaces/${encodeURIComponent(teamId)}/usage` + const res = await fetch(url, { + method: "POST", + headers: { + authorization: `Bearer maple_svc_${mapleServiceToken()}`, + "content-type": "application/json", + }, + body: JSON.stringify(usage), + signal: AbortSignal.timeout(REPORT_USAGE_TIMEOUT_MS), + }) + if (!res.ok) { + throw new Error(`Maple usage report failed for team ${teamId}: HTTP ${res.status}`) + } +} + async function fetchWorkspace(teamId: string): Promise { const url = `${mapleApiBaseUrl()}/internal/slack/workspaces/${encodeURIComponent(teamId)}` const res = await fetch(url, { diff --git a/apps/slack-agent/agent/lib/model.ts b/apps/slack-agent/agent/lib/model.ts new file mode 100644 index 000000000..4d1e1fd9d --- /dev/null +++ b/apps/slack-agent/agent/lib/model.ts @@ -0,0 +1,19 @@ +/** + * The Workers AI model the agent runs on. + * + * Lives here rather than inline in `agent/agent.ts` because two places need + * the same answer: the agent definition itself, and AI-usage reporting + * (`agent/lib/token-usage.ts`), which stamps the model on every usage report. + * eve's `step.completed` event carries token counts but no model id, so the + * reporter has to resolve it the same way the agent did — and a drift between + * the two would mis-attribute spend to a model that never ran. + */ + +/** Must support tool calling **while streaming** — eve's harness is tool-driven and always streams. */ +export const DEFAULT_WORKERS_AI_MODEL = "@cf/zai-org/glm-5.2" + +/** Read at call time, not module load, so a var set after import still counts. */ +export function workersAiModelId(): string { + const raw = process.env.WORKERS_AI_MODEL + return raw && raw.length > 0 ? raw : DEFAULT_WORKERS_AI_MODEL +} diff --git a/apps/slack-agent/agent/lib/token-usage.test.ts b/apps/slack-agent/agent/lib/token-usage.test.ts new file mode 100644 index 000000000..7683a0b5c --- /dev/null +++ b/apps/slack-agent/agent/lib/token-usage.test.ts @@ -0,0 +1,156 @@ +import { afterEach, describe, expect, test } from "bun:test" +import { DEFAULT_WORKERS_AI_MODEL } from "./model.js" +import { + prepareUsageReport, + resetZeroTokenDiagnosticsForTests, + trackStepUsage, + type StepUsageInput, +} from "./token-usage.js" + +const step = (overrides: Partial = {}): StepUsageInput => ({ + sessionId: "sess_1", + teamId: "T1", + turnId: "turn_1", + stepIndex: 0, + inputTokens: 120, + outputTokens: 34, + ...overrides, +}) + +const originalModel = process.env.WORKERS_AI_MODEL +afterEach(() => { + if (originalModel === undefined) delete process.env.WORKERS_AI_MODEL + else process.env.WORKERS_AI_MODEL = originalModel +}) + +describe("prepareUsageReport", () => { + test("builds a report prefixed by session, turn and step", () => { + const prepared = prepareUsageReport(step({ stepIndex: 3 })) + expect(prepared?.teamId).toBe("T1") + expect(prepared?.report).toMatchObject({ + inputTokens: 120, + outputTokens: 34, + model: DEFAULT_WORKERS_AI_MODEL, + }) + // Readable prefix for debugging, plus a per-process unique suffix. + expect(prepared?.report.idempotencyKey).toStartWith("sess_1:turn_1:3:") + }) + + test("stamps the configured model", () => { + process.env.WORKERS_AI_MODEL = "@cf/meta/llama-4" + expect(prepareUsageReport(step())?.report.model).toBe("@cf/meta/llama-4") + }) + + test("keys distinct steps of one turn distinctly", () => { + const first = prepareUsageReport(step({ stepIndex: 0 }))?.report.idempotencyKey + const second = prepareUsageReport(step({ stepIndex: 1 }))?.report.idempotencyKey + expect(first).not.toBe(second) + }) + + test("keys a re-used session/turn/step triple distinctly — eve reuses turn ids", () => { + // A terminal model failure does not advance eve's `sequence`, so the next + // turn is minted as the same `turn_${sequence}` with `stepIndex` back at 0. + // Identical input must NOT produce an identical key, or the billing + // provider silently drops the retry's spend. + const first = prepareUsageReport(step())?.report.idempotencyKey + const second = prepareUsageReport(step())?.report.idempotencyKey + expect(first).not.toBe(second) + expect(first).toStartWith("sess_1:turn_1:0:") + expect(second).toStartWith("sess_1:turn_1:0:") + }) + + test("skips a session with no Slack team — nothing maps it to a payer", () => { + expect(prepareUsageReport(step({ teamId: undefined }))).toBeNull() + expect(prepareUsageReport(step({ teamId: "" }))).toBeNull() + }) + + test("skips an all-zero step (truncated stream), but keeps a one-sided one", () => { + expect(prepareUsageReport(step({ inputTokens: 0, outputTokens: 0 }))).toBeNull() + expect(prepareUsageReport(step({ inputTokens: undefined, outputTokens: undefined }))).toBeNull() + expect(prepareUsageReport(step({ outputTokens: 0 }))?.report.outputTokens).toBe(0) + }) + + test("coerces provider counts to non-negative integers", () => { + const prepared = prepareUsageReport(step({ inputTokens: 12.7, outputTokens: -5 })) + expect(prepared?.report).toMatchObject({ inputTokens: 12, outputTokens: 0 }) + }) + + test("drops non-finite counts rather than sending NaN", () => { + expect(prepareUsageReport(step({ inputTokens: Number.NaN, outputTokens: 0 }))).toBeNull() + }) +}) + +describe("trackStepUsage", () => { + test("reports a billable step to Maple", async () => { + const calls: Array<{ teamId: string; idempotencyKey: string; inputTokens: number }> = [] + await trackStepUsage(step(), { + reportTokenUsage: async (teamId, usage) => { + calls.push({ teamId, idempotencyKey: usage.idempotencyKey, inputTokens: usage.inputTokens }) + }, + }) + expect(calls).toHaveLength(1) + expect(calls[0]).toMatchObject({ teamId: "T1", inputTokens: 120 }) + expect(calls[0]?.idempotencyKey).toStartWith("sess_1:turn_1:0:") + }) + + test("does not call Maple for a skipped step", async () => { + let called = false + await trackStepUsage(step({ teamId: undefined }), { + reportTokenUsage: async () => { + called = true + }, + }) + expect(called).toBe(false) + }) + + test("swallows a reporting failure — a billing write must not fail the turn", async () => { + const promise = trackStepUsage(step(), { + reportTokenUsage: async () => { + throw new Error("HTTP 503") + }, + }) + await expect(promise).resolves.toBeUndefined() + }) + + test("warns once per window when steps are skipped for zero tokens", async () => { + // If the provider never surfaces usage, EVERY step is skipped and nothing + // is ever billed — that has to be visible in the logs, not silent. + resetZeroTokenDiagnosticsForTests() + const warnings: string[] = [] + const originalWarn = console.warn + console.warn = (line: unknown) => { + warnings.push(String(line)) + } + try { + const zero = step({ inputTokens: 0, outputTokens: 0 }) + const deps = { reportTokenUsage: async () => {} } + await trackStepUsage(zero, deps) + await trackStepUsage(zero, deps) + await trackStepUsage(zero, deps) + } finally { + console.warn = originalWarn + resetZeroTokenDiagnosticsForTests() + } + + // Throttled to the first of the window, not one line per skipped step. + const zeroTokenLines = warnings.filter((line) => line.includes("usage_report_zero_tokens")) + expect(zeroTokenLines).toHaveLength(1) + expect(zeroTokenLines[0]).toContain("maple.ai.zero_token_steps") + }) + + test("does not warn about zero tokens when there is simply no team to bill", async () => { + resetZeroTokenDiagnosticsForTests() + const warnings: string[] = [] + const originalWarn = console.warn + console.warn = (line: unknown) => { + warnings.push(String(line)) + } + try { + await trackStepUsage(step({ teamId: undefined }), { reportTokenUsage: async () => {} }) + } finally { + console.warn = originalWarn + resetZeroTokenDiagnosticsForTests() + } + expect(warnings.filter((line) => line.includes("usage_report_zero_tokens"))).toHaveLength(0) + }) +}) diff --git a/apps/slack-agent/agent/lib/token-usage.ts b/apps/slack-agent/agent/lib/token-usage.ts new file mode 100644 index 000000000..0f7213a48 --- /dev/null +++ b/apps/slack-agent/agent/lib/token-usage.ts @@ -0,0 +1,186 @@ +// Relative, not `#lib/*.js`: bun's test runner does not rewrite the package +// `imports` map onto .ts sources, and this module is under test. +import { workersAiModelId } from "./model.js" +import { reportTokenUsage, type TokenUsageReport } from "./maple.js" +import { emitAgentLog } from "./telemetry-log.js" + +/** + * AI-usage capture for the Slack bot. + * + * The bot runs its own model (Cloudflare Workers AI on Railway) outside every + * path Maple's API already meters, so without this its spend is invisible to + * billing entirely. eve surfaces provider-reported token counts on exactly one + * hook event — `step.completed`, one per model call — so that is where usage is + * read; `turn.completed` carries only `{ sequence, turnId }`. + * + * **Reported per step, not accumulated per turn.** Per-turn batching would mean + * fewer HTTP calls but holding mutable state keyed by turn id, which leaks on + * any turn that never terminates and silently drops the whole turn's spend if + * the process dies mid-turn. Per-step is stateless. + */ + +import { randomUUID } from "node:crypto" + +/** + * Unique-per-process suffix source for the idempotency key. + * + * `::` is NOT unique per model call, which cost + * real revenue: eve derives `turnId` as `turn_${sequence}` and only advances + * `sequence` in `emitTurnEpilogue` / `emitRecoverableFailedTurn`. A **terminal** + * model failure (provider 400, context-length, model-config error) runs + * `emitFailedStep`, which advances nothing — so when the user re-asks, the next + * turn is minted with the SAME `turn_${sequence}` and `stepIndex` back at 0, and + * that retry's step 0 (usually the largest input step) reuses a key the billing + * provider has already seen. It is dropped silently: no log, no 4xx, just + * under-billing. eve's park-and-resume branches (`setPendingRuntimeActionBatch`, + * `setPendingAuthorization`) return emission state unchanged and have the same + * shape. + * + * The trade-off, taken deliberately: appending a per-process unique component + * means a genuine duplicate DELIVERY of one call would be billed twice. That is + * near-impossible here — `reportTokenUsage` issues a single `fetch` with no + * retry (its `AbortSignal.timeout` throws and is swallowed) — whereas the + * collision above is a live path. Under-billing silently is worse than a + * theoretical double-count, so uniqueness wins. The readable + * `::` prefix is kept so a key still says which + * call it came from. + */ +const REPORT_BOOT_ID = randomUUID().slice(0, 8) +let reportSequence = 0 + +/** The subset of eve's `step.completed` event this module needs. */ +export interface StepUsageInput { + readonly sessionId: string + /** + * Slack workspace the turn belongs to. Absent for a session with no Slack + * auth context (local dev against a single-workspace token), where there is + * no org to bill. + */ + readonly teamId: string | undefined + readonly turnId: string + readonly stepIndex: number + readonly inputTokens: number | undefined + readonly outputTokens: number | undefined +} + +/** A report ready to send, paired with the team it bills to. */ +export interface PreparedUsageReport { + readonly teamId: string + readonly report: TokenUsageReport +} + +/** Maple's endpoint takes non-negative integers; providers are trusted but not verified. */ +function toTokenCount(value: number | undefined): number { + if (typeof value !== "number" || !Number.isFinite(value)) return 0 + return Math.max(0, Math.trunc(value)) +} + +/** + * Why a step is not billable, or `"billable"` when it is. + * + * Two skip cases, both routine rather than exceptional: + * - **No team.** Nothing maps the session to an org, so the usage has no payer. + * - **Zero tokens.** `workers-ai-provider` reports all-zero usage when a stream + * is truncated before the provider sends its usage chunk, so zero-token steps + * are expected noise — and Maple would drop the report anyway. + */ +export type StepUsageClass = "billable" | "no-team" | "zero-tokens" + +export function classifyStep(input: StepUsageInput): StepUsageClass { + if (input.teamId === undefined || input.teamId.length === 0) return "no-team" + if (toTokenCount(input.inputTokens) === 0 && toTokenCount(input.outputTokens) === 0) { + return "zero-tokens" + } + return "billable" +} + +/** + * Turns one `step.completed` into a billable report, or `null` when there is + * nothing to bill (see {@link classifyStep}). + */ +export function prepareUsageReport(input: StepUsageInput): PreparedUsageReport | null { + if (classifyStep(input) !== "billable") return null + reportSequence += 1 + return { + // `classifyStep` already rejected an absent team. + teamId: input.teamId as string, + report: { + idempotencyKey: `${input.sessionId}:${input.turnId}:${input.stepIndex}:${REPORT_BOOT_ID}-${reportSequence}`, + inputTokens: toTokenCount(input.inputTokens), + outputTokens: toTokenCount(input.outputTokens), + model: workersAiModelId(), + }, + } +} + +/** + * Zero-token steps are skipped as truncation noise, which hides a total + * failure mode: if `workers-ai-provider` never surfaces usage for the + * configured model — it calls `/ai/run/{model}` with `{stream:true}` and never + * sets `stream_options.include_usage`, leaving its reducer's all-zero + * initializer intact when no chunk carries usage — then EVERY step is skipped + * and the feature bills exactly nothing, with no error anywhere. + * + * So the skip is counted and surfaced: the first one, then at most one line per + * window, carrying the cumulative count. Deliberately not a per-step log (a + * genuinely truncated stream is routine) and deliberately not "bill the zeros". + */ +const ZERO_TOKEN_LOG_INTERVAL_MS = 5 * 60_000 +let zeroTokenSkips = 0 +let zeroTokenLastLoggedAt: number | undefined + +/** Test-only: resets the zero-token skip throttle. */ +export function resetZeroTokenDiagnosticsForTests(): void { + zeroTokenSkips = 0 + zeroTokenLastLoggedAt = undefined +} + +function noteZeroTokenSkip(input: StepUsageInput): void { + zeroTokenSkips += 1 + const now = Date.now() + if (zeroTokenLastLoggedAt !== undefined && now - zeroTokenLastLoggedAt < ZERO_TOKEN_LOG_INTERVAL_MS) { + return + } + zeroTokenLastLoggedAt = now + emitAgentLog("warn", "usage_report_zero_tokens", { + "maple.agent.event": "usage_report_zero_tokens", + "session.id": input.sessionId, + "maple.ai.model": workersAiModelId(), + "maple.ai.zero_token_steps": zeroTokenSkips, + }) +} + +export interface TokenUsageDeps { + reportTokenUsage(teamId: string, usage: TokenUsageReport): Promise +} + +const defaultDeps: TokenUsageDeps = { reportTokenUsage } + +/** + * Reports one completed model step's usage to Maple. Never throws and never + * rejects: it is fired without awaiting from a hook, and a failed billing write + * must not surface as a failed turn. Failures are logged (queryable, unlike a + * bare console line) and dropped — there is no local retry, because the Slack + * reply has already gone out and cannot be held for it. + */ +export async function trackStepUsage( + input: StepUsageInput, + deps: TokenUsageDeps = defaultDeps, +): Promise { + if (classifyStep(input) === "zero-tokens") noteZeroTokenSkip(input) + const prepared = prepareUsageReport(input) + if (prepared === null) return + try { + await deps.reportTokenUsage(prepared.teamId, prepared.report) + } catch (error) { + emitAgentLog("warn", "usage_report_failed", { + "maple.agent.event": "usage_report_failed", + "session.id": input.sessionId, + "maple.slack.team_id": prepared.teamId, + "maple.ai.model": prepared.report.model, + "maple.ai.input_tokens": prepared.report.inputTokens, + "maple.ai.output_tokens": prepared.report.outputTokens, + "maple.agent.error_message": error instanceof Error ? error.message : String(error), + }) + } +}