Skip to content
Sign in

flagon-io/g1t

public
0

Commit

runner: a run sent straight to AI Gateway (no model proxy) carries a session in its metadata and on its billing run, so billing settles it to the gateway's price instead of keeping the sandbox's figure

syntaqxcommitted Parent0d6fd03Browse files
3 files+37−50/3 viewed
+7−2
6868 canReachModel,
6969 changeSize,
7070 chooseTier,
71+ gatewaySession,
7172 lastAttemptFailed,
7273 modelEnv,
7374 parseRouting,
13931394 session = opened.value;
13941395 }
13951396 const own = session?.billedTo === "workspace";
1397+ // Straight to the gateway, without the proxy: the run still gets a
1398+ // session there, so billing settles it to what the gateway priced it
1399+ // at instead of leaving the sandbox's own figure.
1400+ const direct = !session && this.env.AI_GATEWAY_ID ? gatewaySession() : undefined;
13961401 // A workspace's own provider is not routed by tier: it runs the model
13971402 // its route names, or for an Anthropic provider, the large tier's.
13981403 const routed = routing.tiers[own ? "large" : tier];
14051410 task,
14061411 model: own ? `${modelName} (${session?.providerName ?? "own provider"})` : modelName,
14071412 billedTo: own ? "workspace" : "g1t",
1408− session: own ? null : (session?.id ?? null),
1413+ session: own ? null : (session?.id ?? direct ?? null),
14091414 tier: own ? null : tier,
14101415 });
14111416 if (!ticket.ok) return ticket;
14231428 // the harness's small tasks too.
14241429 ...(session.model ? { ANTHROPIC_SMALL_FAST_MODEL: session.model } : {}),
14251430 }
1426− : modelEnv(this.env, routing, task, tier, tags);
1431+ : modelEnv(this.env, routing, task, tier, direct ? { ...tags, session: direct } : tags);
14271432 if (ticket.value) {
14281433 // How the sandbox says what the run cost. Kept from the agent.
14291434 vars.BILLING_RUN = ticket.value.runId;
+12−0
99 canReachModel,
1010 changeSize,
1111 chooseTier,
12+ gatewaySession,
1213 lastAttemptFailed,
1314 modelEnv,
1415 parseRouting,
127128 });
128129 });
129130
131+test("a run straight to the gateway carries its session, so billing can settle it", () => {
132+ const session = gatewaySession();
133+ assert.match(session, /^rs_[0-9a-f]{24}$/);
134+ assert.notEqual(gatewaySession(), session);
135+ const vars = modelEnv({ ...direct, AI_GATEWAY_ID: "g1t" }, routes, "implement", "small", { ...tags, session });
136+ const metadata = JSON.parse(customHeaders(vars)["cf-aig-metadata"]);
137+ assert.equal(metadata.session, session);
138+ // The gateway keeps at most five metadata entries.
139+ assert.ok(Object.keys(metadata).length <= 5);
140+});
141+
130142 test("an authenticated gateway is sent its token", () => {
131143 const vars = modelEnv({ ...direct, AI_GATEWAY_ID: "g1t", AI_GATEWAY_TOKEN: "tok" }, routes, "implement", "large", tags);
132144 assert.equal(customHeaders(vars)["cf-aig-authorization"], "Bearer tok");
+18−3
123123 AI_GATEWAY_TOKEN?: string;
124124 };
125125
126−/** What a run is for, attached to each of its requests at the gateway. */
127−export type RunTags = { repo: string; pull: number };
126+/**
127+ * What a run is for, attached to each of its requests at the gateway.
128+ * `session` is the run's id there: billing finds the run's requests by it
129+ * and settles the run to what the gateway priced them at.
130+ */
131+export type RunTags = { repo: string; pull: number; session?: string };
128132
133+/**
134+ * A session id for a run that goes straight to the gateway (no model
135+ * proxy): `rs_` and 24 hex characters, which billing's log filter needs
136+ * no escaping for.
137+ */
138+export function gatewaySession(): string {
139+ const bytes = crypto.getRandomValues(new Uint8Array(12));
140+ return `rs_${Array.from(bytes, (b) => b.toString(16).padStart(2, "0")).join("")}`;
141+}
142+
129143 /** Whether there is a way to reach a model at all. */
130144 export function canReachModel(env: ModelRouting): boolean {
131145 return Boolean(env.ANTHROPIC_API_KEY || (env.AI_GATEWAY_ID && env.AI_GATEWAY_TOKEN));
160174
161175 vars.ANTHROPIC_BASE_URL = `https://gateway.ai.cloudflare.com/v1/${env.CLOUDFLARE_ACCOUNT_ID}/${env.AI_GATEWAY_ID}/anthropic`;
162176 // The gateway logs these with every request, so spend and failures can
163− // be read per kind of work, tier, repository and pull request.
177+ // be read per kind of work, tier, repository and pull request; and by
178+ // the run's session, which billing settles the run's charge by.
164179 const headers = [`cf-aig-metadata: ${JSON.stringify({ task, tier, ...tags })}`];
165180 if (env.AI_GATEWAY_TOKEN) {
166181 vars.AI_GATEWAY_TOKEN = env.AI_GATEWAY_TOKEN;