Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions config.example.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -170,6 +170,11 @@ claim:
# Backoff after failed attempts (10 min default); quota_exhausted /
# already_claimed back off to the server-provided next window instead.
cooldownMs: 600000
# Optional daily claim window in UTC+8, "HH:MM-HH:MM" (end exclusive).
# Outside it the scheduler sends no upstream requests at all, and a
# successful claim holds until the next day's window opens (one claim per
# day). Empty = no gate (poll around the clock). Example: "19:00-20:30".
window: ""
# Optional: claim a specific plan_id. Empty = highest-priority preview.
planId: ""

Expand Down
3 changes: 2 additions & 1 deletion src/claim/runtime.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ import type { ClaimablePlan, ClaimOutcome } from "./types.js";
import type { Credential } from "../auth/types.js";
import { credentialString } from "../auth/types.js";
import { createClaimClient, ClaimPreviewError } from "./client.js";
import { ClaimScheduler } from "./scheduler.js";
import { ClaimScheduler, parseClaimWindow } from "./scheduler.js";
import { getCaptchaToken } from "../proxy/captcha.js";
import { loadCredential } from "../auth/store.js";
import { getDefaultEndpointRouting, requiresClaimCaptcha } from "../proxy/endpoint-routing.js";
Expand Down Expand Up @@ -79,6 +79,7 @@ export function startAutoClaim(config: ProxyConfig, auth: AuthManager): ClaimSch
planId: config.claim.planId || undefined,
pollIntervalMs: config.claim.pollIntervalMs,
cooldownMs: config.claim.cooldownMs,
window: parseClaimWindow(config.claim.window),
},
log: (message) => console.log(`[claim] ${message}`),
});
Expand Down
82 changes: 81 additions & 1 deletion src/claim/scheduler.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
* failure kind, credential loss → stop, and error backoff.
*/
import { describe, it, expect } from "bun:test";
import { ClaimScheduler, type ClaimSchedulerDeps } from "./scheduler.js";
import { ClaimScheduler, parseClaimWindow, minutesUntilWindow, type ClaimSchedulerDeps } from "./scheduler.js";
import { ClaimPreviewError } from "./client.js";
import type { ClaimOutcome, ClaimablePlan } from "./types.js";

Expand Down Expand Up @@ -201,6 +201,86 @@ describe("ClaimScheduler.tick", () => {
});
});

describe("ClaimScheduler — daily claim window (UTC+8)", () => {
// Window below is the campaign's usual daily slot, 19:00-20:30 UTC+8.
function windowHarness() {
const h = makeHarness();
h.config.window = { startMin: 19 * 60, endMin: 20 * 60 + 30 };
return h;
}

it("suppresses upstream requests outside the window and holds until the next open", async () => {
const h = windowHarness();
// 2026-11-11T12:31:00Z = 20:31 UTC+8 — outside (end is 20:30).
h.nowMs = Date.UTC(2026, 10, 11, 12, 31, 0);
const res = await h.scheduler.tick();
expect(res).toEqual({ action: "idle" });
// ~22h29m until tomorrow 19:00; hold must cover it (re-armed every poll by scheduleNext).
h.nowMs += 22 * 60 * 60_000 + 29 * 60_000 - 1;
expect(await h.scheduler.tick()).toEqual({ action: "skipped_hold" });
expect(h.claimCalls).toHaveLength(0);
});

it("polls normally inside the window", async () => {
const h = windowHarness();
h.nowMs = Date.UTC(2026, 10, 11, 11, 30, 0); // 19:30 UTC+8 — inside
const res = await h.scheduler.tick();
expect(res).toEqual({ action: "claimed", planId: "weekend-1" });
});

it("after a successful claim, the next tick waits for the NEXT day's window", async () => {
const h = windowHarness();
h.nowMs = Date.UTC(2026, 10, 11, 11, 30, 0); // 19:30 UTC+8 — inside
await h.scheduler.tick();
// holdUntil = now + 23.5h (minutes until tomorrow 19:00); endsAt undefined.
const expectedHold = 23 * 60 * 60_000 + 30 * 60_000;
h.nowMs += expectedHold - 60_000; // 18:59 next day — still held
expect(await h.scheduler.tick()).toEqual({ action: "skipped_hold" });
h.nowMs = Date.UTC(2026, 10, 12, 11, 30, 0); // 19:30 next day — window open again
expect((await h.scheduler.tick()).action).toBe("claimed");
});

it("window start is inclusive", async () => {
const h = windowHarness();
h.nowMs = Date.UTC(2026, 10, 11, 11, 0, 0); // exactly 19:00 UTC+8
expect((await h.scheduler.tick()).action).toBe("claimed");
});

it("window end is exclusive (20:30:00 is outside)", async () => {
const h = windowHarness();
h.nowMs = Date.UTC(2026, 10, 11, 12, 30, 0); // exactly 20:30 UTC+8
expect((await h.scheduler.tick()).action).toBe("idle");
expect(h.claimCalls).toHaveLength(0);
});

it("no window configured = legacy always-poll behavior", async () => {
const h = makeHarness(); // no window
h.nowMs = Date.UTC(2026, 10, 11, 12, 1, 0); // 20:01 UTC+8
expect((await h.scheduler.tick()).action).toBe("claimed");
});
});

describe("parseClaimWindow / minutesUntilWindow", () => {
it("parses HH:MM-HH:MM and rejects malformed input", () => {
expect(parseClaimWindow("19:00-20:30")).toEqual({ startMin: 1140, endMin: 1230 });
expect(parseClaimWindow(" 0:05-23:59 ")).toEqual({ startMin: 5, endMin: 1439 });
expect(parseClaimWindow("")).toBeUndefined();
expect(parseClaimWindow(undefined)).toBeUndefined();
expect(() => parseClaimWindow("19:00")).toThrow(/HH:MM-HH:MM/);
expect(() => parseClaimWindow("24:00-25:00")).toThrow(/out of range/);
expect(() => parseClaimWindow("19:00-19:00")).toThrow(/start equals end/);
});

it("minutesUntilWindow handles before/inside/after and midnight wrap", () => {
const win = { startMin: 1140, endMin: 1230 }; // 19:00-20:30
expect(minutesUntilWindow(1140, win)).toBe(0); // inside (start)
expect(minutesUntilWindow(1229, win)).toBe(0); // inside
expect(minutesUntilWindow(1230, win)).toBe(22 * 60 + 30); // 20:30 → next 19:00
expect(minutesUntilWindow(600, win)).toBe(540); // 10:00 → 19:00
expect(minutesUntilWindow(1439, win)).toBe(1141); // 23:59 → 19:00 next day (wrap)
});
});

describe("ClaimScheduler lifecycle", () => {
it("start/stop drives ticks on the timer without overlap", async () => {
const h = makeHarness({ pollIntervalMs: 5 });
Expand Down
68 changes: 66 additions & 2 deletions src/claim/scheduler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,12 @@ interface ClaimSchedulerConfig {
planId?: string;
pollIntervalMs: number;
cooldownMs: number;
/**
* Daily claim window in UTC+8, `"HH:MM-HH:MM"` (24h). Outside the window
* ticks are suppressed (no preview requests) and a successful claim holds
* until the NEXT day's window start. `undefined` = gate off (poll always).
*/
window?: { startMin: number; endMin: number };
}

export interface ClaimSchedulerDeps {
Expand All @@ -48,6 +54,43 @@ export type TickResult =
| { action: "failed"; outcome: Extract<ClaimOutcome, { ok: false }>; holdMs: number }
| { action: "error"; message: string; holdMs: number };

/** Daily claim window in UTC+8 ("HH:MM-HH:MM"): outside it, don't touch upstream. */
export function parseClaimWindow(raw: string | undefined): { startMin: number; endMin: number } | undefined {
const s = raw?.trim();
if (!s) return undefined;
const m = /^(\d{1,2}):(\d{2})-(\d{1,2}):(\d{2})$/.exec(s);
if (!m) throw new Error(`claim.window: expected "HH:MM-HH:MM", got "${s}"`);
const startMin = Number(m[1]) * 60 + Number(m[2]);
const endMin = Number(m[3]) * 60 + Number(m[4]);
if (Number(m[1]) > 23 || Number(m[3]) > 23 || Number(m[2]) > 59 || Number(m[4]) > 59) {
throw new Error(`claim.window: time out of range in "${s}"`);
}
if (startMin === endMin) throw new Error(`claim.window: start equals end in "${s}"`);
return { startMin, endMin };
}

/** Minutes since local midnight mapped into the UTC+8 wall clock. */
export function utcPlus8MinutesOfDay(nowMs: number): number {
const shifted = new Date(nowMs + 8 * 60 * 60 * 1000);
return shifted.getUTCHours() * 60 + shifted.getUTCMinutes();
}

function fmtMinute(min: number): string {
return `${String(Math.floor(min / 60)).padStart(2, "0")}:${String(min % 60).padStart(2, "0")}`;
}

/** Minutes until the window opens (startMin inclusive, endMin exclusive); 0 = inside. */
export function minutesUntilWindow(minutesOfDay: number, win: { startMin: number; endMin: number }): number {
if (minutesOfDay < win.startMin) return win.startMin - minutesOfDay;
if (minutesOfDay >= win.endMin) return 24 * 60 - minutesOfDay + win.startMin;
return 0; // inside the window
}

/** Minutes until the window NEXT opens (always ≥1 — inside the window counts to tomorrow's start). */
export function minutesUntilNextStart(minutesOfDay: number, win: { startMin: number; endMin: number }): number {
return minutesOfDay < win.startMin ? win.startMin - minutesOfDay : 24 * 60 - minutesOfDay + win.startMin;
}

export class ClaimScheduler {
private stopped = false;
private holdUntil = 0;
Expand Down Expand Up @@ -83,6 +126,19 @@ export class ClaimScheduler {
const nowMs = this.now();
if (nowMs < this.holdUntil) return { action: "skipped_hold" };

// Daily claim window (UTC+8): outside the window, don't touch upstream.
// Hold until the next window opens; scheduleNext clamps to pollIntervalMs,
// so the hold re-arms every ~5min without any request being sent.
const win = this.deps.config.window;
if (win) {
const waitMin = minutesUntilWindow(utcPlus8MinutesOfDay(nowMs), win);
if (waitMin > 0) {
this.holdUntil = nowMs + waitMin * 60_000;
this.log(`claim: outside claim window (UTC+8 ${fmtMinute(win.startMin)}-${fmtMinute(win.endMin)}) — next poll at window open (~${waitMin}m)`);
return { action: "idle" };
}
}

let jwt: string | undefined;
try {
jwt = await this.deps.getJwt();
Expand Down Expand Up @@ -148,8 +204,16 @@ export class ClaimScheduler {

if (outcome.ok) {
const endsAtMs = outcome.endsAt !== undefined ? outcome.endsAt * 1000 : undefined;
this.holdUntil = endsAtMs ?? nowMs + this.deps.config.pollIntervalMs;
this.log(`claim: claimed plan ${target.planId}${outcome.startsAt !== undefined ? ` (activates ${new Date(outcome.startsAt * 1000).toISOString()})` : ""}`);
// Claimed: stop polling until the NEXT window opens (success once per
// day; re-polling inside the same window would just burn daily quota).
const nextWindowMs = win ? minutesUntilNextStart(utcPlus8MinutesOfDay(nowMs), win) * 60_000 : undefined;
const holdCandidates: number[] = [];
if (endsAtMs !== undefined) holdCandidates.push(endsAtMs);
if (nextWindowMs !== undefined) holdCandidates.push(nowMs + nextWindowMs);
this.holdUntil = holdCandidates.length > 0
? Math.max(...holdCandidates)
: nowMs + this.deps.config.pollIntervalMs;
this.log(`claim: claimed plan ${target.planId}${outcome.startsAt !== undefined ? ` (activates ${new Date(outcome.startsAt * 1000).toISOString()})` : ""}${nextWindowMs !== undefined ? ` — next poll at next window open (~${Math.round(nextWindowMs / 60_000)}m)` : ""}`);
return { action: "claimed", planId: target.planId, startsAt: outcome.startsAt, endsAt: outcome.endsAt };
}

Expand Down
19 changes: 19 additions & 0 deletions src/config/loader.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -140,6 +140,7 @@ logging:
origin: "https://zcode.z.ai",
pollIntervalMs: 300000,
cooldownMs: 600000,
window: "",
planId: "",
});
});
Expand Down Expand Up @@ -304,10 +305,28 @@ claim:
origin: "https://zcode.z.ai",
pollIntervalMs: 60000,
cooldownMs: 120000,
window: "",
planId: "weekend-special",
});
});

it("claim: window parses and passes through to the scheduler config", () => {
const path = writeYaml(`
claim:
window: "19:00-20:30"
`);
const cfg = loadConfig(path);
expect(cfg.claim.window).toBe("19:00-20:30");
});

it("claim: malformed window throws at load time", () => {
const path = writeYaml(`
claim:
window: "25:00-99:99"
`);
expect(() => loadConfig(path)).toThrow(/claim\.window/);
});

it("claim: ZCODE_CLAIM_ENABLED / ZCODE_CLAIM_POLL_INTERVAL_MS env override", () => {
const path = writeYaml(`
claim:
Expand Down
7 changes: 7 additions & 0 deletions src/config/loader.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import { readFileSync, existsSync } from "node:fs";
import { parse } from "yaml";
import type { ClientIdentityConfig, ProxyConfig, ProviderEndpoints, ProxyIdentity, ResponsesConfig, McpConfig, AsyncConfig, EndpointRoutingConfig, ClientSigningConfig, HighspeedConfig, ClaimConfig } from "./types.js";
import { parseClaimWindow } from "../claim/scheduler.js";

/** Environment variable keys that override YAML values. */
const ENV = {
Expand Down Expand Up @@ -77,6 +78,7 @@ const DEFAULTS = {
CLAIM_ORIGIN: "https://zcode.z.ai",
CLAIM_POLL_INTERVAL_MS: 300000,
CLAIM_COOLDOWN_MS: 600000,
CLAIM_WINDOW: "",
CLAIM_PLAN_ID: "",
ENDPOINT_ROUTING_ENABLED: true,
ENDPOINT_ROUTING_ORIGIN: "https://zcode.z.ai",
Expand Down Expand Up @@ -401,12 +403,17 @@ function resolveClaimConfig(raw: unknown): ClaimConfig {
const origin = (originEnv ?? (typeof obj.origin === "string" ? obj.origin : DEFAULTS.CLAIM_ORIGIN)).trim() || DEFAULTS.CLAIM_ORIGIN;
validateOrigin(origin, "claim.origin");

// Validate the window format at load time (throws on malformed input).
const rawWindow = typeof obj.window === "string" ? obj.window : DEFAULTS.CLAIM_WINDOW;
const window = parseClaimWindow(rawWindow) ? rawWindow.trim() : DEFAULTS.CLAIM_WINDOW;

return {
enabled: enabledEnv !== undefined ? resolveBool(enabledEnv, DEFAULTS.CLAIM_ENABLED) : resolveBool(obj.enabled, DEFAULTS.CLAIM_ENABLED),
auto: autoEnv !== undefined ? resolveBool(autoEnv, DEFAULTS.CLAIM_AUTO) : resolveBool(obj.auto, DEFAULTS.CLAIM_AUTO),
origin,
pollIntervalMs: resolvePositiveInt(pollIntervalEnv ?? obj.pollIntervalMs ?? obj.poll_interval_ms, DEFAULTS.CLAIM_POLL_INTERVAL_MS, "claim.pollIntervalMs"),
cooldownMs: resolvePositiveInt(obj.cooldownMs ?? obj.cooldown_ms, DEFAULTS.CLAIM_COOLDOWN_MS, "claim.cooldownMs"),
window,
planId: typeof obj.planId === "string" ? obj.planId.trim() : DEFAULTS.CLAIM_PLAN_ID,
};
}
Expand Down
7 changes: 7 additions & 0 deletions src/config/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -169,6 +169,13 @@ export interface ClaimConfig {
pollIntervalMs: number;
/** Backoff after a failed claim attempt in ms. Default `600000` (10 min). */
cooldownMs: number;
/**
* Daily claim window in UTC+8, `"HH:MM-HH:MM"` (24h). Outside the window the
* scheduler sends no upstream requests, and a successful claim waits for the
* next day's window before polling again (one claim per day). Empty disables
* the gate and polls around the clock. Default `""`.
*/
window?: string;
/** Optional `plan_id` to claim; empty string claims the highest-priority preview. Default `""`. */
planId: string;
}
Expand Down
Loading