Skip to content
Merged
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,13 @@ public class TeamMembership implements Serializable {
@Column(name = "updated_at", nullable = false)
private LocalDateTime updatedAt;

/**
* Optional per-member spend cap inside the team's wallet, in doc units. NULL means the member
* is bounded only by the team-wide cap.
*/
@Column(name = "cap_units")
private Long capUnits;

public boolean isLeader() {
return role == TeamRole.LEADER;
}
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,156 @@
package stirling.software.saas.payg.docs;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.file.Files;
import java.util.List;
import java.util.Objects;

import org.springframework.stereotype.Component;
import org.springframework.web.multipart.MultipartFile;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;

import stirling.software.common.util.TempFile;
import stirling.software.common.util.TempFileManager;
import stirling.software.jpdfium.PdfDocument;
import stirling.software.saas.payg.policy.PricingPolicy;

/**
* Reads pages via jpdfium for PDF inputs; treats every other content type as bytes-only.
*
* <p>For PDFs, units are the larger of {@code ceil(pages / docPagesPerUnit)} and {@code ceil(bytes
* / docBytesPerUnit)}. For non-PDFs, only the bytes axis contributes. A single file is clamped to
* {@code [1, policy.fileUnitCap]}; a multi-file group is clamped to {@code [1, policy.fileUnitCap *
* file_count]} applied to the sum of raw per-file units.
*
* <p>Malformed or encrypted PDFs fall back to bytes-only classification — the file still has a size
* we can charge against, and the caller decides whether to reject the upload on other grounds.
*/
@Slf4j
@Component
@RequiredArgsConstructor
public class DefaultDocumentClassifier implements DocumentClassifier {

private static final String PDF_CONTENT_TYPE = "application/pdf";
private static final String DEFAULT_CONTENT_TYPE = "application/octet-stream";

private final TempFileManager tempFileManager;

@Override
public DocumentMetrics classify(MultipartFile file, PricingPolicy policy) {
Objects.requireNonNull(file, "file");
Objects.requireNonNull(policy, "policy");

FileFacts facts = inspect(file);
long rawUnits = computeRawUnits(facts.pages, facts.bytes, policy);
int units = (int) Math.max(1, Math.min(policy.getFileUnitCap(), rawUnits));
return new DocumentMetrics(facts.pages, facts.bytes, facts.contentType, units);
}

@Override
public DocumentMetrics classify(List<MultipartFile> files, PricingPolicy policy) {
Objects.requireNonNull(files, "files");
Objects.requireNonNull(policy, "policy");
if (files.isEmpty()) {
throw new IllegalArgumentException("files must not be empty");
}

int totalPages = 0;
long totalBytes = 0;
long rawUnitsSum = 0;
String firstContentType = null;

for (MultipartFile file : files) {
FileFacts facts = inspect(file);
// Sum the *raw* (unclamped) per-file units so the group cap below can actually bind.
// Per-file clamping in this loop would make the group cap a no-op.
rawUnitsSum =
saturatedAdd(rawUnitsSum, computeRawUnits(facts.pages, facts.bytes, policy));
totalPages = saturatedAdd(totalPages, facts.pages);
totalBytes = saturatedAdd(totalBytes, facts.bytes);
if (firstContentType == null) {
firstContentType = facts.contentType;
}
}

long groupCap = (long) policy.getFileUnitCap() * files.size();
int totalUnits = (int) Math.max(1L, Math.min(groupCap, rawUnitsSum));

return new DocumentMetrics(
totalPages,
totalBytes,
firstContentType != null ? firstContentType : DEFAULT_CONTENT_TYPE,
totalUnits);
}

private FileFacts inspect(MultipartFile file) {
long bytes = file.getSize();
String contentType =
file.getContentType() != null ? file.getContentType() : DEFAULT_CONTENT_TYPE;
int pages = isPdf(contentType, file.getOriginalFilename()) ? readPageCount(file) : 0;
return new FileFacts(pages, bytes, contentType);
}

private static long computeRawUnits(int pages, long bytes, PricingPolicy policy) {
long pageUnits = pages > 0 ? ceilDiv(pages, policy.getDocPagesPerUnit()) : 0L;
long byteUnits = ceilDiv(bytes, policy.getDocBytesPerUnit());
return Math.max(pageUnits, byteUnits);
}

private static long ceilDiv(long numerator, long divisor) {
if (numerator <= 0) {
return 0;
}
return (numerator + divisor - 1) / divisor;
}

private static boolean isPdf(String contentType, String filename) {
if (PDF_CONTENT_TYPE.equalsIgnoreCase(contentType)) {
return true;
}
return filename != null && filename.toLowerCase().endsWith(".pdf");
}

/**
* Materialises the upload to a managed temp file and asks jpdfium for the page count. Returns 0
* if the file can't be parsed — the byte-derived axis still produces a charge.
*/
private int readPageCount(MultipartFile file) {
try (TempFile temp = tempFileManager.createManagedTempFile(".pdf")) {
try (InputStream in = file.getInputStream();
OutputStream out = Files.newOutputStream(temp.getPath())) {
in.transferTo(out);
}
try (PdfDocument doc = PdfDocument.open(temp.getPath())) {
return doc.pageCount();
}
} catch (IOException | RuntimeException e) {
log.debug(
"Could not read PDF page count for {} ({}); falling back to bytes-only units",
file.getOriginalFilename(),
e.getClass().getSimpleName());
return 0;
}
}

private static int saturatedAdd(int a, int b) {
long sum = (long) a + b;
if (sum > Integer.MAX_VALUE) {
return Integer.MAX_VALUE;
}
return (int) sum;
}

private static long saturatedAdd(long a, long b) {
try {
return Math.addExact(a, b);
} catch (ArithmeticException e) {
return Long.MAX_VALUE;
}
}

private record FileFacts(int pages, long bytes, String contentType) {}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
package stirling.software.saas.payg.docs;

import java.util.List;

import org.springframework.web.multipart.MultipartFile;

import stirling.software.saas.payg.policy.PricingPolicy;

/** Computes the doc-unit cost of an uploaded file (or multi-file input) under a given policy. */
public interface DocumentClassifier {

/** Classify a single uploaded file. Returns 1 unit minimum, {@code fileUnitCap} maximum. */
DocumentMetrics classify(MultipartFile file, PricingPolicy policy);

/**
* Classify a multi-file input (e.g. a merge or overlay). Returns the sum of each file's units,
* capped at {@code fileUnitCap × files.size()}.
*/
DocumentMetrics classify(List<MultipartFile> files, PricingPolicy policy);
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
package stirling.software.saas.payg.docs;

/**
* Output of {@link DocumentClassifier#classify}. {@code pages} is {@code 0} for non-PDF inputs.
*
* @param pages page count (0 for non-PDFs and for files whose page count couldn't be read)
* @param bytes raw byte length of the file
* @param contentType MIME type as reported by the upload, or {@code "application/octet-stream"}
* when unknown
* @param docUnits computed unit cost, clamped to the policy's {@code fileUnitCap}
*/
public record DocumentMetrics(int pages, long bytes, String contentType, int docUnits) {}
Original file line number Diff line number Diff line change
@@ -0,0 +1,110 @@
package stirling.software.saas.payg.entitlement;

import java.io.Serializable;
import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;

import org.hibernate.annotations.CreationTimestamp;
import org.hibernate.annotations.JdbcTypeCode;
import org.hibernate.type.SqlTypes;

import jakarta.persistence.Column;
import jakarta.persistence.Embeddable;
import jakarta.persistence.EmbeddedId;
import jakarta.persistence.Entity;
import jakarta.persistence.EnumType;
import jakarta.persistence.Enumerated;
import jakarta.persistence.Table;

import lombok.Getter;
import lombok.NoArgsConstructor;
import lombok.Setter;

import stirling.software.saas.payg.model.EntitlementState;
import stirling.software.saas.payg.model.FeatureGate;
import stirling.software.saas.payg.model.FeatureSet;

/**
* Cached entitlement state for the team (one row with {@code user_id = 0}, the team-wide sentinel)
* plus optional per-member rows when a member sub-cap is configured. Read on the hot path by the
* entitlement guard.
*
* <p>Composite PK {@code (team_id, user_id)} uses 0 as the team-wide sentinel because Postgres
* treats {@code NULL} as not-equal-to-NULL in unique constraints — 0 keeps the PK well-defined.
*/
@Entity
@Table(name = "wallet_entitlement_snapshot")
@NoArgsConstructor
@Getter
@Setter
public class WalletEntitlementSnapshot implements Serializable {

private static final long serialVersionUID = 1L;

public static final long TEAM_WIDE_USER_ID = 0L;

@EmbeddedId private WalletEntitlementSnapshotId id;

@Column(name = "period_start", nullable = false)
private LocalDateTime periodStart;

@Column(name = "period_end", nullable = false)
private LocalDateTime periodEnd;

@Column(name = "period_spend_units", nullable = false)
private Long periodSpendUnits = 0L;

@Column(name = "period_cap_units")
private Long periodCapUnits;

@Enumerated(EnumType.STRING)
@Column(name = "state", nullable = false, length = 16)
private EntitlementState state = EntitlementState.FULL;

@Enumerated(EnumType.STRING)
@Column(name = "feature_set", nullable = false, length = 32)
private FeatureSet featureSet = FeatureSet.FULL;

@JdbcTypeCode(SqlTypes.JSON)
@Column(name = "enabled_gates", columnDefinition = "jsonb", nullable = false)
private List<FeatureGate> enabledGates = new ArrayList<>();

@CreationTimestamp
@Column(name = "computed_at", nullable = false, updatable = false)
private LocalDateTime computedAt;

@Embeddable
@NoArgsConstructor
@Getter
@Setter
public static class WalletEntitlementSnapshotId implements Serializable {

private static final long serialVersionUID = 1L;

@Column(name = "team_id", nullable = false)
private Long teamId;

/** Use {@link #TEAM_WIDE_USER_ID} for the team-wide row. */
@Column(name = "user_id", nullable = false)
private Long userId;

public WalletEntitlementSnapshotId(Long teamId, Long userId) {
this.teamId = teamId;
this.userId = userId;
}

@Override
public boolean equals(Object o) {
if (this == o) return true;
if (!(o instanceof WalletEntitlementSnapshotId other)) return false;
return Objects.equals(teamId, other.teamId) && Objects.equals(userId, other.userId);
}

@Override
public int hashCode() {
return Objects.hash(teamId, userId);
}
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,81 @@
package stirling.software.saas.payg.job;

import java.io.Serializable;
import java.time.LocalDateTime;
import java.util.Objects;
import java.util.UUID;

import org.hibernate.annotations.CreationTimestamp;

import jakarta.persistence.Column;
import jakarta.persistence.Embeddable;
import jakarta.persistence.EmbeddedId;
import jakarta.persistence.Entity;
import jakarta.persistence.EnumType;
import jakarta.persistence.Enumerated;
import jakarta.persistence.Table;

import lombok.Getter;
import lombok.NoArgsConstructor;
import lombok.Setter;

import stirling.software.saas.payg.model.ArtifactKind;

/**
* Per-step input/output content hash. Used by the lineage detector to decide whether a tool call
* joins an open process (matching an earlier input or output) or opens a new one.
*/
@Entity
@Table(name = "job_artifact_hash")
@NoArgsConstructor
@Getter
@Setter
public class JobArtifactHash implements Serializable {

private static final long serialVersionUID = 1L;

@EmbeddedId private JobArtifactHashId id;

@CreationTimestamp
@Column(name = "created_at", nullable = false, updatable = false)
private LocalDateTime createdAt;

@Embeddable
@NoArgsConstructor
@Getter
@Setter
public static class JobArtifactHashId implements Serializable {

private static final long serialVersionUID = 1L;

@Column(name = "job_id", nullable = false)
private UUID jobId;

@Column(name = "content_hash", nullable = false, length = 64)
private String contentHash;

@Enumerated(EnumType.STRING)
@Column(name = "kind", nullable = false, length = 8)
private ArtifactKind kind;

public JobArtifactHashId(UUID jobId, String contentHash, ArtifactKind kind) {
this.jobId = jobId;
this.contentHash = contentHash;
this.kind = kind;
}

@Override
public boolean equals(Object o) {
if (this == o) return true;
if (!(o instanceof JobArtifactHashId other)) return false;
return Objects.equals(jobId, other.jobId)
&& Objects.equals(contentHash, other.contentHash)
&& kind == other.kind;
}

@Override
public int hashCode() {
return Objects.hash(jobId, contentHash, kind);
}
}
}
Loading
Loading