feat(ai): centralize durable usage
Add a content-free usage ledger with legacy backfill. Reserve Workspace and durable Strategy/CV work before execution so deleted history or duplicate admission cannot reset limits.
This commit is contained in:
@@ -13,7 +13,9 @@ public sealed record CreateUserOperation(
|
||||
string? SubjectId = null,
|
||||
int Priority = 0,
|
||||
int MaxAttempts = 3,
|
||||
DateTime? DeadlineAtUtc = null);
|
||||
DateTime? DeadlineAtUtc = null,
|
||||
int UsageInputCharacters = 0,
|
||||
int UsageReservedTokens = 0);
|
||||
|
||||
public sealed record UserOperationCreation(UserOperation Operation, bool Created);
|
||||
public sealed record UserOperationLease(Guid OperationId, string OwnerUserId, string LeaseToken, string TaskType, string PrivacyPolicy, string? SubjectType, string? SubjectId, int AttemptCount, DateTime? DeadlineAtUtc);
|
||||
@@ -73,6 +75,18 @@ public sealed class UserOperationStore(JobTrackerContext db, TimeProvider timePr
|
||||
DeadlineAtUtc = request.DeadlineAtUtc,
|
||||
};
|
||||
db.UserOperations.Add(operation);
|
||||
AiUsageRecord? usage = null;
|
||||
if (request.UsageReservedTokens > 0)
|
||||
{
|
||||
usage = AiUsageMeter.NewOperationRecord(
|
||||
owner,
|
||||
operation.Id,
|
||||
operation.TaskType,
|
||||
request.UsageInputCharacters,
|
||||
request.UsageReservedTokens,
|
||||
new DateTimeOffset(now));
|
||||
db.AiUsageRecords.Add(usage);
|
||||
}
|
||||
try
|
||||
{
|
||||
await db.SaveChangesAsync(cancellationToken);
|
||||
@@ -81,6 +95,7 @@ public sealed class UserOperationStore(JobTrackerContext db, TimeProvider timePr
|
||||
catch (DbUpdateException)
|
||||
{
|
||||
db.Entry(operation).State = EntityState.Detached;
|
||||
if (usage is not null) db.Entry(usage).State = EntityState.Detached;
|
||||
existing = await db.UserOperations.FirstOrDefaultAsync(
|
||||
item => item.TaskType == request.TaskType && item.IdempotencyKey == request.IdempotencyKey,
|
||||
cancellationToken);
|
||||
@@ -170,6 +185,18 @@ public sealed class UserOperationStore(JobTrackerContext db, TimeProvider timePr
|
||||
string? model,
|
||||
string? completionStage,
|
||||
CancellationToken cancellationToken)
|
||||
=> await CompleteAsync(operationId, leaseToken, resultReference, provider, model, completionStage, null, null, cancellationToken);
|
||||
|
||||
public async Task<int> CompleteAsync(
|
||||
Guid operationId,
|
||||
string leaseToken,
|
||||
string? resultReference,
|
||||
string? provider,
|
||||
string? model,
|
||||
string? completionStage,
|
||||
int? usageInputCharacters,
|
||||
int? usageOutputCharacters,
|
||||
CancellationToken cancellationToken)
|
||||
{
|
||||
EnsureOwnerScope();
|
||||
ValidateOptional(resultReference, 256, nameof(resultReference));
|
||||
@@ -196,6 +223,15 @@ public sealed class UserOperationStore(JobTrackerContext db, TimeProvider timePr
|
||||
cancellationToken);
|
||||
if (affected == 1)
|
||||
{
|
||||
if (usageInputCharacters is not null && usageOutputCharacters is not null)
|
||||
{
|
||||
var estimatedTokens = (usageInputCharacters.Value + usageOutputCharacters.Value + 3) / 4;
|
||||
await db.AiUsageRecords.Where(item => item.SourceType == "operation" && item.SourceId == operationId.ToString("D"))
|
||||
.ExecuteUpdateAsync(setters => setters
|
||||
.SetProperty(item => item.InputCharacterCount, usageInputCharacters.Value)
|
||||
.SetProperty(item => item.OutputCharacterCount, usageOutputCharacters.Value)
|
||||
.SetProperty(item => item.EstimatedTokenCount, estimatedTokens), cancellationToken);
|
||||
}
|
||||
db.UserNotifications.Add(CreateTerminalNotification(operation, OperationStatuses.Succeeded, now));
|
||||
await db.SaveChangesAsync(cancellationToken);
|
||||
if (transaction is not null) await transaction.CommitAsync(cancellationToken);
|
||||
|
||||
Reference in New Issue
Block a user