Skip to content

fix: refresh Codex quota without running Codex #6

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
mabyes1 wants to merge 1 commit into
base: master
Choose a base branch
from
Draft

fix: refresh Codex quota without running Codex #6

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
4 changes: 3 additions & 1 deletion src/PhoneMonitor.Host/Quotas/AiQuotaService.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,9 @@ private async Task<QuotaSnapshot> BuildSnapshotAsync(bool forceAgyRefresh, Cance
Providers = new List<AiQuotaStatus>()
};

snapshot.Providers.AddRange(CodexQuotaReader.ReadCodexQuotas());
snapshot.Providers.AddRange(forceAgyRefresh
? await CodexQuotaReader.RefreshCodexQuotasAsync(cancellationToken)
: CodexQuotaReader.ReadCodexQuotas());
snapshot.Providers.AddRange(await agy.ReadAgyQuotasAsync(forceAgyRefresh, cancellationToken));
return snapshot;
}
Expand Down
149 changes: 149 additions & 0 deletions src/PhoneMonitor.Host/Quotas/CodexQuotaReader.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,12 @@
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Net.Http;
using System.Net.Http.Headers;
using System.Text;
using System.Text.Json;
using System.Threading;
using System.Threading.Tasks;
using static PhoneMonitor.Host.Quotas.QuotaJsonHelpers;
using static PhoneMonitor.Host.Quotas.QuotaPaths;
using static PhoneMonitor.Host.Quotas.QuotaShared;
Expand All @@ -19,8 +23,23 @@ internal static class CodexQuotaReader
private const int MaxSessionFiles = 120;
private const int TailBytes = 768 * 1024;
private static readonly object CacheFileLock = new object();
private static readonly HttpClient UsageClient = new HttpClient
{
Timeout = TimeSpan.FromSeconds(15)
};

internal static IEnumerable<AiQuotaStatus> ReadCodexQuotas()
{
return ReadCodexQuotas(null);
}

internal static async Task<IEnumerable<AiQuotaStatus>> RefreshCodexQuotasAsync(CancellationToken cancellationToken)
{
var refreshed = await TryFetchCodexQuotaAsync(cancellationToken);
return ReadCodexQuotas(refreshed);
}

private static IEnumerable<AiQuotaStatus> ReadCodexQuotas(AiQuotaStatus refreshedQuota)
{
var codexHomes = ResolveCodexHomes().ToList();
var cacheDirectory = CodexQuotaCacheDirectory();
Expand Down Expand Up @@ -48,6 +67,10 @@ internal static IEnumerable<AiQuotaStatus> ReadCodexQuotas()
var activeQuota = !string.IsNullOrWhiteSpace(activeHome)
? ReadCodexQuota(activeHome)
: Unavailable("codex", "Codex", "Codex session directory was not found.", activeHome);
if (IsUsableQuota(refreshedQuota))
{
activeQuota = refreshedQuota;
}
if (!IsUsableQuota(activeQuota))
{
activeQuota = codexHomes
Expand Down Expand Up @@ -96,6 +119,100 @@ internal static IEnumerable<AiQuotaStatus> ReadCodexQuotas()
return new[] { activeQuota };
}

private static async Task<AiQuotaStatus> TryFetchCodexQuotaAsync(CancellationToken cancellationToken)
{
var authFile = ResolveCodexHomes()
.Select(home => Path.Combine(home, "auth.json"))
.FirstOrDefault(File.Exists);
if (string.IsNullOrWhiteSpace(authFile))
{
return null;
}

try
{
using var authDocument = JsonDocument.Parse(File.ReadAllText(authFile));
if (!TryGetProperty(authDocument.RootElement, "tokens", out var tokens) ||
tokens.ValueKind != JsonValueKind.Object)
{
return null;
}

var accessToken = TryGetString(tokens, "access_token");
if (string.IsNullOrWhiteSpace(accessToken))
{
return null;
}

using var request = new HttpRequestMessage(
HttpMethod.Get,
"https://chatgpt.com/backend-api/wham/usage");
request.Headers.Authorization = new AuthenticationHeaderValue("Bearer", accessToken);
request.Headers.UserAgent.ParseAdd("VibeDeck/1.0");
var accountId = TryGetString(tokens, "account_id");
if (!string.IsNullOrWhiteSpace(accountId))
{
request.Headers.TryAddWithoutValidation("ChatGPT-Account-ID", accountId);
}

using var response = await UsageClient.SendAsync(request, HttpCompletionOption.ResponseHeadersRead, cancellationToken);
if (!response.IsSuccessStatusCode)
{
return null;
}

var json = await response.Content.ReadAsStringAsync(cancellationToken);
return TryReadCodexQuotaFromUsageResponse(json);
}
catch (Exception ex) when (
ex is IOException ||
ex is UnauthorizedAccessException ||
ex is JsonException ||
ex is HttpRequestException ||
ex is TaskCanceledException ||
ex is FormatException)
{
return null;
}
}

internal static AiQuotaStatus TryReadCodexQuotaFromUsageResponse(string json)
{
if (string.IsNullOrWhiteSpace(json))
{
return null;
}

using var document = JsonDocument.Parse(json);
var root = document.RootElement;
if (!TryGetProperty(root, "rate_limit", out var limits) ||
limits.ValueKind != JsonValueKind.Object)
{
return null;
}

var credits = TryGetProperty(root, "credits", out var creditInfo) &&
creditInfo.ValueKind == JsonValueKind.Object
? creditInfo
: default;
return new AiQuotaStatus
{
Id = "codex",
Label = "Codex",
Family = "codex",
AccountId = "local",
State = "ok",
Source = "https://chatgpt.com/backend-api/wham/usage",
Detail = "Latest Codex quota fetched directly by VibeDeck.",
ObservedAt = DateTimeOffset.UtcNow,
AccountTier = TryGetString(root, "plan_type"),
CreditBalance = credits.ValueKind == JsonValueKind.Object ? TryGetDouble(credits, "balance") : null,
CreditUnlimited = credits.ValueKind == JsonValueKind.Object ? TryGetBoolean(credits, "unlimited") : null,
Primary = ReadUsageWindow(limits, "primary_window", "5h"),
Secondary = ReadUsageWindow(limits, "secondary_window", "Weekly")
};
}

private static IReadOnlyList<string> ResolveCodexHomes()
{
var homes = new List<string>();
Expand Down Expand Up @@ -336,6 +453,20 @@ private static void WriteCodexQuotaCache(string cacheDirectory, AiQuotaStatus st
Directory.CreateDirectory(cacheDirectory);
var accountKey = FirstNonEmpty(status.AccountId, status.AccountEmail, status.Id, "local");
var path = Path.Combine(cacheDirectory, $"{SafeFileName(accountKey)}.json");
if (File.Exists(path))
{
try
{
var cached = JsonSerializer.Deserialize<AiQuotaStatus>(File.ReadAllText(path), CacheJsonOptions);
if (cached?.ObservedAt > status.ObservedAt)
{
return;
}
}
catch (Exception ex) when (ex is IOException || ex is UnauthorizedAccessException || ex is JsonException)
{
}
}
File.WriteAllText(path, JsonSerializer.Serialize(status, CacheJsonOptions));
}
}
Expand Down Expand Up @@ -505,6 +636,24 @@ private static QuotaWindow ReadWindow(JsonElement limits, string propertyName, s
};
}

private static QuotaWindow ReadUsageWindow(JsonElement limits, string propertyName, string label)
{
if (!TryGetProperty(limits, propertyName, out var window) ||
window.ValueKind != JsonValueKind.Object)
{
return null;
}

var windowSeconds = TryGetInt(window, "limit_window_seconds");
return new QuotaWindow
{
Label = label,
UsedPercent = TryGetDouble(window, "used_percent"),
WindowMinutes = windowSeconds.HasValue ? windowSeconds.Value / 60 : (int?)null,
ResetsAt = TryGetUnixTime(window, "reset_at")
};
}

internal sealed class CodexAccountIdentity
{
public string AccountId { get; set; }
Expand Down
17 changes: 17 additions & 0 deletions tests/PhoneMonitor.Host.Tests/AiQuotaServiceTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,23 @@ public void ReadsCodexCreditBalanceFromLatestRateLimitEvent()
}
}

[Fact]
public void ReadsCodexQuotaFromLiveUsageResponse()
{
const string response = "{\"plan_type\":\"plus\",\"rate_limit\":{\"primary_window\":{\"used_percent\":23,\"limit_window_seconds\":18000,\"reset_at\":1780000000},\"secondary_window\":{\"used_percent\":41,\"limit_window_seconds\":604800,\"reset_at\":1780500000}},\"credits\":{\"balance\":\"12.5\",\"unlimited\":false}}";

var status = CodexQuotaReader.TryReadCodexQuotaFromUsageResponse(response);

Assert.Equal("ok", status.State);
Assert.Equal("plus", status.AccountTier);
Assert.Equal(23d, status.Primary.UsedPercent);
Assert.Equal(300, status.Primary.WindowMinutes);
Assert.Equal(41d, status.Secondary.UsedPercent);
Assert.Equal(10080, status.Secondary.WindowMinutes);
Assert.Equal(12.5d, status.CreditBalance);
Assert.False(status.CreditUnlimited);
}

[Fact]
public void AtomicallyReplacesCodexAuthFileFromSavedProfile()
{
Expand Down
Loading