Skip to content

Commit 63d5c6d

Browse files
Yash-Singh1Yash-Singh1
authored
feat(opencode): read usage limits from the connected server (#34146)
1 parent 76e4cb5 commit 63d5c6d

10 files changed

Lines changed: 630 additions & 29 deletions

File tree

‎packages/contracts/src/providerUsageLimits.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,9 @@ export const ServerProviderUsageWindow = Schema.Struct({
1414
id: TrimmedNonEmptyString,
1515
kind: Schema.Literals(["session", "weekly", "monthly", "other"]),
1616
label: TrimmedNonEmptyString,
17+
/** When present, this limit applies only to these provider model IDs. */
18+
modelIds: Schema.optional(Schema.Array(TrimmedNonEmptyString)),
19+
providerIds: Schema.optional(Schema.Array(TrimmedNonEmptyString)),
1720
usedPercent: Schema.Number.check(Schema.isBetween({ minimum: 0, maximum: 100 })),
1821
resetsAt: Schema.optional(IsoDateTime),
1922
windowDurationMins: Schema.optional(NonNegativeInt),

‎packages/provider-core/src/server/usageLimits.test.ts‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,29 @@ describe("applyUsageLimitsUpdate", () => {
3333
},
3434
});
3535
expect(next).toBe(published);
36+
const scoped = { checkedAt, windows: [{ ...weekly, modelIds: ["opencode/a"] }] };
37+
const providerScoped = { checkedAt, windows: [{ ...weekly, providerIds: ["opencode"] }] };
38+
expect(
39+
applyUsageLimitsUpdate({
40+
previous: providerScoped,
41+
checkedAt,
42+
update: { windows: [{ ...weekly, providerIds: ["console-custom"] }] },
43+
})?.windows[0]?.providerIds,
44+
).toEqual(["console-custom"]);
45+
expect(
46+
applyUsageLimitsUpdate({
47+
previous: scoped,
48+
checkedAt,
49+
update: { windows: [{ ...weekly, modelIds: ["opencode/a"] }] },
50+
}),
51+
).toBe(scoped);
52+
expect(
53+
applyUsageLimitsUpdate({
54+
previous: scoped,
55+
checkedAt,
56+
update: { windows: [{ ...weekly, modelIds: ["opencode/b"] }] },
57+
})?.windows[0]?.modelIds,
58+
).toEqual(["opencode/b"]);
3659
});
3760

3861
it("upserts by id and keeps the reset a percent-only update omits", () => {

‎packages/provider-core/src/server/usageLimits.ts‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,16 @@ function usageWindowEquals(a: ServerProviderUsageWindow, b: ServerProviderUsageW
9494
a.kind === b.kind &&
9595
a.label === b.label &&
9696
a.usedPercent === b.usedPercent &&
97+
(a.modelIds === b.modelIds ||
98+
(a.modelIds !== undefined &&
99+
b.modelIds !== undefined &&
100+
a.modelIds.length === b.modelIds.length &&
101+
a.modelIds.every((id, index) => id === b.modelIds?.[index]))) &&
102+
(a.providerIds === b.providerIds ||
103+
(a.providerIds !== undefined &&
104+
b.providerIds !== undefined &&
105+
a.providerIds.length === b.providerIds.length &&
106+
a.providerIds.every((id, index) => id === b.providerIds?.[index]))) &&
97107
a.resetsAt === b.resetsAt &&
98108
a.windowDurationMins === b.windowDurationMins
99109
);

‎packages/provider-opencode/src/server/driver.test.ts‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -804,6 +804,81 @@ it.layer(layer)("OpenCodeDriver runtime selection", (it) => {
804804
}).pipe(Effect.scoped, Effect.provide(TestClock.layer())),
805805
);
806806

807+
it.effect.each([
808+
{ status: 200, data: null, reason: "unsupported" },
809+
{ status: 404, data: null, reason: "unsupported" },
810+
{ status: 503, data: null, reason: "probeFailed" },
811+
{
812+
status: 200,
813+
data: {
814+
workspace: {
815+
id: "remote-workspace",
816+
name: "Remote workspace",
817+
usageURL: "https://opencode.ai/console/remote-workspace",
818+
},
819+
budgets: [
820+
{
821+
id: "all",
822+
limit: 20,
823+
spent: 15,
824+
scope: "combined",
825+
per: "set",
826+
coverage: "all",
827+
models: [],
828+
shared: true,
829+
resetsAt: 1793491200000,
830+
},
831+
],
832+
},
833+
reason: undefined,
834+
},
835+
])(
836+
"reads external workspace usage without local credentials, HTTP $status, $reason",
837+
(scenario) =>
838+
Effect.gen(function* () {
839+
const requested: string[] = [];
840+
const replay = replayOpenCodeServer(OPENCODE_2_RESPONSES, "secret", requested);
841+
const server = HttpClient.make((request) => {
842+
assert.notInclude(request.url, "opencode.ai/zen");
843+
if (request.url.endsWith("/api/experimental/usage/limits")) {
844+
assert.strictEqual(request.headers.authorization, "Basic b3BlbmNvZGU6c2VjcmV0");
845+
return Effect.succeed(
846+
HttpClientResponse.fromWeb(
847+
request,
848+
Response.json(
849+
{
850+
location: { directory: "/work" },
851+
data: scenario.data,
852+
},
853+
{ status: scenario.status },
854+
),
855+
),
856+
);
857+
}
858+
return replay.execute(request);
859+
});
860+
const instance = yield* create(
861+
{ serverUrl: "https://remote.example", serverPassword: "secret" },
862+
server,
863+
);
864+
const [snapshot] = Array.from(
865+
yield* instance.snapshot.streamChanges.pipe(
866+
Stream.filter((snapshot) => snapshot.usageLimits !== undefined),
867+
Stream.take(1),
868+
Stream.runCollect,
869+
),
870+
);
871+
assert.strictEqual(snapshot?.usageLimits?.unavailable?.reason, scenario.reason);
872+
if (scenario.reason === undefined) {
873+
assert.strictEqual(snapshot?.usageLimits?.windows[0]?.usedPercent, 75);
874+
assert.strictEqual(snapshot?.usageLimits?.externalUsage?.label, "Remote workspace");
875+
assert.deepStrictEqual(snapshot?.usageLimits?.windows[0]?.providerIds, []);
876+
} else {
877+
assert.deepStrictEqual(snapshot?.usageLimits?.windows, []);
878+
}
879+
}).pipe(Effect.scoped),
880+
);
881+
807882
it.effect.each(
808883
[
809884
{ name: "recovers" },
@@ -843,6 +918,10 @@ it.layer(layer)("OpenCodeDriver runtime selection", (it) => {
843918
);
844919
}
845920
const path = new URL(request.url).pathname;
921+
if (path === "/api/experimental/usage/limits")
922+
return Effect.succeed(
923+
HttpClientResponse.fromWeb(request, new Response(null, { status: 404 })),
924+
);
846925
if (path === "/api/model") modelReads += 1;
847926
const body =
848927
path === "/api/info"

‎packages/provider-opencode/src/server/driver.ts‎

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import * as Digest from "effect/crypto/Digest";
88
import * as SecureRandom from "effect/crypto/SecureRandom";
99
import * as Deferred from "effect/Deferred";
1010
import * as Effect from "effect/Effect";
11+
import * as DateTime from "effect/DateTime";
1112
import * as Equal from "effect/Equal";
1213
import * as Option from "effect/Option";
1314
import * as FileSystem from "effect/FileSystem";
@@ -459,7 +460,7 @@ export const OpenCodeDriver: ProviderDriver<
459460
),
460461
);
461462

462-
const checkUsageLimits = readOpenCodeGoUsageLimits({
463+
const legacyUsageLimits = readOpenCodeGoUsageLimits({
463464
enabled: effectiveConfig.enabled,
464465
serverUrl: effectiveConfig.serverUrl,
465466
environment: processEnv,
@@ -469,6 +470,30 @@ export const OpenCodeDriver: ProviderDriver<
469470
Effect.provideService(HttpClient.HttpClient, httpClient),
470471
Effect.provideService(Digest.Digest, digest),
471472
);
473+
let unsupportedUsageVersion: string | undefined;
474+
const checkUsageLimits = Effect.gen(function* () {
475+
if (!effectiveConfig.enabled) return yield* legacyUsageLimits;
476+
const runtime = yield* runtimeProbe.get;
477+
if (runtime.generation === "v1" || runtime.version === unsupportedUsageVersion)
478+
return yield* legacyUsageLimits;
479+
const limits = yield* openCode2Server.withConnection((connection) =>
480+
connection.usageLimits(host.paths.cwd),
481+
);
482+
if (limits === undefined) unsupportedUsageVersion = runtime.version;
483+
return limits ?? (yield* legacyUsageLimits);
484+
}).pipe(
485+
Effect.timeout("10 seconds"),
486+
Effect.catch(() =>
487+
Effect.map(DateTime.now, (now) => ({
488+
checkedAt: DateTime.formatIso(now),
489+
windows: [],
490+
unavailable: {
491+
reason: "probeFailed" as const,
492+
message: "OpenCode could not read workspace usage.",
493+
},
494+
})),
495+
),
496+
);
472497
const checkProvider = Effect.all(
473498
{
474499
// A configured loopback server reports its running version. Probe the local

0 commit comments

Comments
 (0)