From 4d09b3aaa360ca029298aa3c71fe223b66904637 Mon Sep 17 00:00:00 2001 From: Leejaewang03 <192717122+Leejaewang03@users.noreply.github.com> Date: Wed, 22 Jul 2026 10:56:06 +0900 Subject: [PATCH] =?UTF-8?q?=EC=9B=94=EB=B3=84=20budget=20window=EB=A5=BC?= =?UTF-8?q?=20Clock=20=EA=B8=B0=EB=B0=98=20key=EB=A1=9C=20=EA=B5=90?= =?UTF-8?q?=EC=A0=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- AGENTS.md | 17 +- .../TokenPilotAutoConfiguration.java | 15 +- .../autoconfigure/TokenPilotProperties.java | 75 ++++++++ .../TokenPilotAutoConfigurationTest.java | 79 +++++++- .../io/tokenpilot/budget/BudgetDecision.java | 7 +- .../java/io/tokenpilot/budget/BudgetKey.java | 33 ++++ .../io/tokenpilot/budget/BudgetPolicy.java | 51 +++++ .../tokenpilot/budget/BudgetStateStore.java | 15 +- .../io/tokenpilot/budget/BudgetWindow.java | 33 ++++ .../internal/DefaultBudgetEvaluator.java | 155 +++++++-------- .../internal/InMemoryBudgetStateStore.java | 60 ++++-- .../internal/LedgerBudgetComponents.java | 11 +- .../internal/DefaultBudgetEvaluatorTest.java | 180 ++++++++++++------ .../InMemoryBudgetStateStoreTest.java | 101 ++++++++++ .../notification/BudgetNotificationEvent.java | 6 +- .../BudgetNotificationService.java | 13 +- .../InMemoryNotificationStateStore.java | 24 +-- .../notification/NotificationStateStore.java | 11 +- .../BudgetNotificationServiceTest.java | 133 ++++++------- .../tokenpilot/sample/SampleController.java | 9 +- .../internal/DefaultLedgerAdvisor.java | 25 ++- .../internal/DefaultLedgerAdvisorTest.java | 54 +++++- 22 files changed, 797 insertions(+), 310 deletions(-) create mode 100644 token-pilot-budget/src/main/java/io/tokenpilot/budget/BudgetKey.java create mode 100644 token-pilot-budget/src/main/java/io/tokenpilot/budget/BudgetPolicy.java create mode 100644 token-pilot-budget/src/main/java/io/tokenpilot/budget/BudgetWindow.java create mode 100644 token-pilot-budget/src/test/java/io/tokenpilot/budget/internal/InMemoryBudgetStateStoreTest.java diff --git a/AGENTS.md b/AGENTS.md index 3503a64..656dc7b 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -4,7 +4,7 @@ Token Pilot is evolving from a Spring AI usage-tracking starter into a framework-independent Java LLM control and accounting core with optional framework and observability adapters. -Current truth: post-call usage normalization, cost calculation, ledger events, Micrometer publishing, basic non-atomic budget evaluation, Spring AI integration, and starter autoconfiguration are implemented. Preflight estimation, context admission, atomic reservation, and estimate/actual reconciliation are 30-day MVP targets, not current capabilities. +Current truth: post-call usage normalization, cost calculation, ledger events, Micrometer publishing, Clock-based monthly budget windows, basic non-atomic budget evaluation, Spring AI integration, and starter autoconfiguration are implemented. Preflight estimation, context admission, atomic reservation, and estimate/actual reconciliation are 30-day MVP targets, not current capabilities. Distribution direction: publish a framework-independent core and an optional Spring AI convenience starter from the same repository and release train. The existing starter artifact is `token-pilot-starter`; `token-pilot-spring-ai-starter` is only a target name until a compatibility ADR and module change land. @@ -74,7 +74,7 @@ Token Pilot의 제품 포지션은 framework-independent Java LLM control and ac | `token-pilot-core` | Basic implementation complete | Domain records, pricing, calculator, registry, ledger manager | | `token-pilot-spring-ai` | Basic implementation complete | `UsageExtractor`, `LedgerAdvisor`, response usage recording | | `token-pilot-micrometer` | Basic implementation complete | `MetricsOptions`, tag whitelist, and metric metadata exist; metric ownership must be narrowed | -| `token-pilot-budget` | Basic non-atomic implementation | Needs BLOCK enforcement, window semantics, reservation, idempotency, and reconciliation | +| `token-pilot-budget` | Basic non-atomic implementation | Typed monthly keys and Clock/ZoneId windows implemented; needs BLOCK enforcement, reservation, idempotency, and reconciliation | | `token-pilot-notification` | Basic implementation complete | Event API and deduplication exist; not yet connected to the full advisor/budget lifecycle | | `token-pilot-autoconfigure` | Basic implementation complete | Bean registration, property binding, pricing/budget/notification wiring, and ChatClient customizer implemented | | `token-pilot-starter` | Basic implementation complete | Thin final user entrypoint that brings runtime modules together | @@ -233,7 +233,7 @@ class MailBudgetNotificationHandler implements BudgetNotificationHandler { - `BudgetNotificationHandler` 빈이 없으면 `BudgetNotificationService`는 등록되지 않는다 (no-op). - `token-pilot.notification.enabled=true` 설정 시에만 notification 빈이 등록된다. -- 알림 중복 방지는 `(targetId, budgetWindow)` 조합으로 처리된다. +- 알림 중복 방지는 evaluator가 확정한 `BudgetKey(policyId, targetType, targetId, window)`로 처리된다. - 같은 window 안에서는 낮거나 같은 threshold 재발송이 방지된다. - 새 window에서는 50/80/100% 알림이 다시 가능하다. @@ -256,7 +256,13 @@ token-pilot: - model budget: enabled: false + policy-id: default-monthly + target-type: tenant + target-tag-key: tenant_id + # fallback-target-id: shared monthly-limit: 10.00 + currency: USD + zone-id: UTC notification: enabled: false ``` @@ -381,6 +387,11 @@ Stage and deploy a Central release: ## Update History +### 2026-07-22 + +- Added immutable monthly `BudgetKey`/`BudgetWindow` identity resolved from an injected `Clock` and configured `ZoneId`. +- Made budget policy identity, target fallback, limit, and currency explicit, and propagated the resolved key through state updates and notifications. + ### 2026-07-20 - Replaced generic cached input/cached output details with optional cache-read input, cache-creation input, and reasoning-output breakdowns; `null` now means unreported and `0` means reported zero. diff --git a/token-pilot-autoconfigure/src/main/java/io/tokenpilot/autoconfigure/TokenPilotAutoConfiguration.java b/token-pilot-autoconfigure/src/main/java/io/tokenpilot/autoconfigure/TokenPilotAutoConfiguration.java index 5101b19..620c2eb 100644 --- a/token-pilot-autoconfigure/src/main/java/io/tokenpilot/autoconfigure/TokenPilotAutoConfiguration.java +++ b/token-pilot-autoconfigure/src/main/java/io/tokenpilot/autoconfigure/TokenPilotAutoConfiguration.java @@ -28,6 +28,8 @@ import io.tokenpilot.notification.InMemoryNotificationStateStore; import io.tokenpilot.notification.NotificationStateStore; +import java.time.Clock; + /** * Token Pilot 라이브러리의 자동 설정을 담당하는 클래스. */ @@ -165,9 +167,16 @@ public BudgetStateStore budgetStateStore() { @ConditionalOnMissingBean @ConditionalOnClass(LedgerBudgetComponents.class) @ConditionalOnProperty(prefix = "token-pilot.budget", name = "enabled", havingValue = "true") - public BudgetEvaluator budgetEvaluator(BudgetStateStore budgetStateStore, TokenPilotProperties properties) { - return LedgerBudgetComponents.defaultBudgetEvaluator(budgetStateStore, properties.getBudget() - .getMonthlyLimit()); + public BudgetEvaluator budgetEvaluator( + BudgetStateStore budgetStateStore, + TokenPilotProperties properties, + ObjectProvider clock + ) { + return LedgerBudgetComponents.defaultBudgetEvaluator( + budgetStateStore, + properties.toBudgetPolicy(), + clock.getIfAvailable(Clock::systemUTC) + ); } /** diff --git a/token-pilot-autoconfigure/src/main/java/io/tokenpilot/autoconfigure/TokenPilotProperties.java b/token-pilot-autoconfigure/src/main/java/io/tokenpilot/autoconfigure/TokenPilotProperties.java index 6f787f6..8f9bd26 100644 --- a/token-pilot-autoconfigure/src/main/java/io/tokenpilot/autoconfigure/TokenPilotProperties.java +++ b/token-pilot-autoconfigure/src/main/java/io/tokenpilot/autoconfigure/TokenPilotProperties.java @@ -1,10 +1,13 @@ package io.tokenpilot.autoconfigure; +import io.tokenpilot.budget.BudgetPolicy; import io.tokenpilot.core.domain.PricingPlan; import org.springframework.boot.context.properties.ConfigurationProperties; import org.springframework.boot.context.properties.NestedConfigurationProperty; +import java.time.ZoneId; import java.util.ArrayList; +import java.util.Currency; import java.util.HashSet; import java.util.List; import java.util.Set; @@ -80,6 +83,18 @@ public List toPricingPlans() { .toList(); } + public BudgetPolicy toBudgetPolicy() { + return new BudgetPolicy( + budget.getPolicyId(), + budget.getTargetType(), + budget.getTargetTagKey(), + budget.getFallbackTargetId(), + budget.getMonthlyLimit(), + Currency.getInstance(budget.getCurrency()), + ZoneId.of(budget.getZoneId()) + ); + } + public static class PricingProperties { private List plans = new ArrayList<>(); @@ -116,6 +131,12 @@ public void setTagWhitelist(Set tagWhitelist) { public static class BudgetProperties { private boolean enabled = false; private java.math.BigDecimal monthlyLimit = new java.math.BigDecimal("10.00"); + private String policyId = "default-monthly"; + private String targetType = "tenant"; + private String targetTagKey = "tenant_id"; + private String fallbackTargetId; + private String currency = "USD"; + private String zoneId = "UTC"; public boolean isEnabled() { return enabled; @@ -132,6 +153,60 @@ public java.math.BigDecimal getMonthlyLimit() { public void setMonthlyLimit(java.math.BigDecimal monthlyLimit) { this.monthlyLimit = monthlyLimit; } + + public String getPolicyId() { + return policyId; + } + + public void setPolicyId(String policyId) { + this.policyId = policyId; + } + + public String getTargetType() { + return targetType; + } + + public void setTargetType(String targetType) { + this.targetType = targetType; + } + + public String getTargetTagKey() { + return targetTagKey; + } + + public void setTargetTagKey(String targetTagKey) { + this.targetTagKey = targetTagKey; + } + + /** + * 대상 tag가 없을 때 사용할 명시적 fallback입니다. 미설정 시 평가는 fail-closed 됩니다. + */ + public String getFallbackTargetId() { + return fallbackTargetId; + } + + public void setFallbackTargetId(String fallbackTargetId) { + this.fallbackTargetId = fallbackTargetId; + } + + public String getCurrency() { + return currency; + } + + public void setCurrency(String currency) { + this.currency = currency; + } + + /** + * 월별 budget window 경계를 계산하는 IANA ZoneId입니다. 기본값은 UTC입니다. + */ + public String getZoneId() { + return zoneId; + } + + public void setZoneId(String zoneId) { + this.zoneId = zoneId; + } } public static class NotificationProperties { diff --git a/token-pilot-autoconfigure/src/test/java/io/tokenpilot/autoconfigure/TokenPilotAutoConfigurationTest.java b/token-pilot-autoconfigure/src/test/java/io/tokenpilot/autoconfigure/TokenPilotAutoConfigurationTest.java index 0717c1c..c90e4be 100644 --- a/token-pilot-autoconfigure/src/test/java/io/tokenpilot/autoconfigure/TokenPilotAutoConfigurationTest.java +++ b/token-pilot-autoconfigure/src/test/java/io/tokenpilot/autoconfigure/TokenPilotAutoConfigurationTest.java @@ -4,9 +4,11 @@ import io.micrometer.core.instrument.simple.SimpleMeterRegistry; import io.tokenpilot.budget.BudgetDecision; import io.tokenpilot.budget.BudgetEvaluator; +import io.tokenpilot.budget.BudgetKey; import io.tokenpilot.budget.BudgetState; import io.tokenpilot.budget.BudgetStateStore; import io.tokenpilot.budget.BudgetThreshold; +import io.tokenpilot.budget.BudgetWindow; import io.tokenpilot.core.CostCalculator; import io.tokenpilot.core.LedgerManager; import io.tokenpilot.core.PricingProvider; @@ -27,12 +29,17 @@ import org.junit.jupiter.params.provider.MethodSource; import org.springframework.ai.chat.client.ChatClientRequest; import org.springframework.ai.chat.client.advisor.api.AdvisorChain; +import org.springframework.ai.chat.prompt.Prompt; import org.springframework.boot.autoconfigure.AutoConfigurations; import org.springframework.boot.test.context.runner.ApplicationContextRunner; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; import java.math.BigDecimal; +import java.time.Clock; +import java.time.Instant; +import java.time.ZoneOffset; +import java.util.Currency; import java.util.Map; import java.util.Optional; import java.util.stream.Stream; @@ -187,6 +194,49 @@ void shouldRegisterBudgetBeansWhenEnabled() { }); } + @Test + @DisplayName("Budget 빈은 사용자 Clock 기준의 월별 window를 공유해야 한다") + void shouldUseUserClockForMonthlyBudgetWindow() { + this.contextRunner + .withUserConfiguration(FixedClockConfiguration.class) + .withPropertyValues( + "token-pilot.budget.enabled=true", + "token-pilot.budget.monthly-limit=100.00", + "token-pilot.budget.zone-id=Asia/Seoul" + ) + .run(context -> { + BudgetStateStore store = context.getBean(BudgetStateStore.class); + BudgetEvaluator evaluator = context.getBean(BudgetEvaluator.class); + Map tags = Map.of("tenant_id", "tenant-clock"); + BudgetDecision initial = evaluator.evaluate(tags); + BudgetKey july = new BudgetKey( + initial.key().policyId(), + initial.key().targetType(), + initial.key().targetId(), + BudgetWindow.parse("2026-07") + ); + + store.addCost( + july, + initial.limit(), + initial.currency(), + new Cost(new BigDecimal("100.00"), initial.currency()) + ); + store.addCost( + initial.key(), + initial.limit(), + initial.currency(), + new Cost(new BigDecimal("50.00"), initial.currency()) + ); + + BudgetDecision decision = evaluator.evaluate(tags); + + assertThat(decision.key().window()).isEqualTo(BudgetWindow.parse("2026-08")); + assertThat(decision.threshold()).isEqualTo(BudgetThreshold.HALF); + assertThat(decision.currentUsage()).isEqualByComparingTo("50.00"); + }); + } + @Test @DisplayName("Budget가 활성화되면 LedgerAdvisor가 BudgetEvaluator를 사용해야 한다") void shouldWireBudgetEvaluatorIntoLedgerAdvisorWhenBudgetEnabled() { @@ -197,8 +247,10 @@ void shouldWireBudgetEvaluatorIntoLedgerAdvisorWhenBudgetEnabled() { LedgerAdvisor advisor = context.getBean(LedgerAdvisor.class); RecordingBudgetEvaluator evaluator = context.getBean(RecordingBudgetEvaluator.class); - ChatClientRequest request = mock(ChatClientRequest.class); - when(request.context()).thenReturn(Map.of("tenant_id", "tenant-abc")); + ChatClientRequest request = new ChatClientRequest( + new Prompt("test"), + Map.of("tenant_id", "tenant-abc") + ); advisor.before(request, mock(AdvisorChain.class)); @@ -338,6 +390,14 @@ public RecordingBudgetEvaluator budgetEvaluator() { } } + @Configuration(proxyBeanMethods = false) + static class FixedClockConfiguration { + @Bean + public Clock clock() { + return Clock.fixed(Instant.parse("2026-07-31T15:30:00Z"), ZoneOffset.UTC); + } + } + static class RecordingBudgetEvaluator implements BudgetEvaluator { private int evaluateCalls; private Map lastTags = Map.of(); @@ -346,7 +406,20 @@ static class RecordingBudgetEvaluator implements BudgetEvaluator { public BudgetDecision evaluate(Map tags) { this.evaluateCalls++; this.lastTags = tags; - return new BudgetDecision(BudgetState.ALLOW, BudgetThreshold.NONE, "allowed", BigDecimal.ZERO, BigDecimal.TEN); + return new BudgetDecision( + new BudgetKey( + "policy-a", + "tenant", + tags.get("tenant_id"), + BudgetWindow.parse("2026-07") + ), + BudgetState.ALLOW, + BudgetThreshold.NONE, + "allowed", + BigDecimal.ZERO, + BigDecimal.TEN, + Currency.getInstance("USD") + ); } @Override diff --git a/token-pilot-budget/src/main/java/io/tokenpilot/budget/BudgetDecision.java b/token-pilot-budget/src/main/java/io/tokenpilot/budget/BudgetDecision.java index 1b6cd01..6962d62 100644 --- a/token-pilot-budget/src/main/java/io/tokenpilot/budget/BudgetDecision.java +++ b/token-pilot-budget/src/main/java/io/tokenpilot/budget/BudgetDecision.java @@ -1,6 +1,7 @@ package io.tokenpilot.budget; import java.math.BigDecimal; +import java.util.Currency; /** * 예산 평가 결과를 나타내는 객체 @@ -10,11 +11,15 @@ * - reason: 상태 설명 * - currentUsage: 현재 사용량 * - limit: 총 예산 + * - key: 평가 시점에 확정된 예산 bucket 식별자 + * - currency: 예산 통화 */ public record BudgetDecision( + BudgetKey key, BudgetState state, BudgetThreshold threshold, String reason, BigDecimal currentUsage, - BigDecimal limit + BigDecimal limit, + Currency currency ) {} diff --git a/token-pilot-budget/src/main/java/io/tokenpilot/budget/BudgetKey.java b/token-pilot-budget/src/main/java/io/tokenpilot/budget/BudgetKey.java new file mode 100644 index 0000000..4c1ecd5 --- /dev/null +++ b/token-pilot-budget/src/main/java/io/tokenpilot/budget/BudgetKey.java @@ -0,0 +1,33 @@ +package io.tokenpilot.budget; + +import java.util.Objects; + +/** + * 월별 예산 bucket을 식별하는 불변 key입니다. + * + * @param policyId 예산 정책 식별자 + * @param targetType 예산 대상 종류 + * @param targetId 예산 대상 식별자 + * @param window 월별 예산 기간 + */ +public record BudgetKey( + String policyId, + String targetType, + String targetId, + BudgetWindow window +) { + + public BudgetKey { + policyId = requireText(policyId, "policyId"); + targetType = requireText(targetType, "targetType"); + targetId = requireText(targetId, "targetId"); + Objects.requireNonNull(window, "window must not be null"); + } + + private static String requireText(String value, String name) { + if (value == null || value.isBlank()) { + throw new IllegalArgumentException(name + " must not be blank"); + } + return value; + } +} diff --git a/token-pilot-budget/src/main/java/io/tokenpilot/budget/BudgetPolicy.java b/token-pilot-budget/src/main/java/io/tokenpilot/budget/BudgetPolicy.java new file mode 100644 index 0000000..31bd0d0 --- /dev/null +++ b/token-pilot-budget/src/main/java/io/tokenpilot/budget/BudgetPolicy.java @@ -0,0 +1,51 @@ +package io.tokenpilot.budget; + +import java.math.BigDecimal; +import java.time.ZoneId; +import java.util.Currency; +import java.util.Objects; + +/** + * 월별 예산 key 생성과 금액 검증에 사용하는 정책 snapshot입니다. + * 대상 tag가 없고 {@code fallbackTargetId}도 설정되지 않으면 평가는 fail-closed 됩니다. + * + * @param id 정책 식별자 + * @param targetType 예산 대상 종류 + * @param targetTagKey 대상 식별자를 읽을 tag key + * @param fallbackTargetId 누락된 대상에 사용할 명시적 fallback, 미설정 시 {@code null} + * @param monthlyLimit 월별 한도 + * @param currency 예산 통화 + * @param zoneId 월 경계를 계산할 시간대 + */ +public record BudgetPolicy( + String id, + String targetType, + String targetTagKey, + String fallbackTargetId, + BigDecimal monthlyLimit, + Currency currency, + ZoneId zoneId +) { + + public BudgetPolicy { + id = requireText(id, "id"); + targetType = requireText(targetType, "targetType"); + targetTagKey = requireText(targetTagKey, "targetTagKey"); + if (fallbackTargetId != null && fallbackTargetId.isBlank()) { + throw new IllegalArgumentException("fallbackTargetId must not be blank"); + } + Objects.requireNonNull(monthlyLimit, "monthlyLimit must not be null"); + if (monthlyLimit.compareTo(BigDecimal.ZERO) <= 0) { + throw new IllegalArgumentException("monthlyLimit must be greater than zero"); + } + Objects.requireNonNull(currency, "currency must not be null"); + Objects.requireNonNull(zoneId, "zoneId must not be null"); + } + + private static String requireText(String value, String name) { + if (value == null || value.isBlank()) { + throw new IllegalArgumentException(name + " must not be blank"); + } + return value; + } +} diff --git a/token-pilot-budget/src/main/java/io/tokenpilot/budget/BudgetStateStore.java b/token-pilot-budget/src/main/java/io/tokenpilot/budget/BudgetStateStore.java index 2fc475f..a4e8e96 100644 --- a/token-pilot-budget/src/main/java/io/tokenpilot/budget/BudgetStateStore.java +++ b/token-pilot-budget/src/main/java/io/tokenpilot/budget/BudgetStateStore.java @@ -1,19 +1,16 @@ package io.tokenpilot.budget; -import java.math.BigDecimal; -import java.util.Map; +import io.tokenpilot.core.domain.Cost; +import java.math.BigDecimal; +import java.util.Currency; /** - * 예산 판단을 위해 비용 누적 상태를 관리하는 저장소 인터페이스입니다. - *

- * 예산 식별자별(예: tenant)로 비용을 조회하고 - * 누적하는 책임을 가집니다. + * resolved {@link BudgetKey}별 누적 비용과 limit/currency snapshot을 관리합니다. */ - public interface BudgetStateStore { - BigDecimal getAccumulatedCost(Map tags); + BigDecimal getAccumulatedCost(BudgetKey key, BigDecimal limit, Currency currency); - void addCost(Map tags, BigDecimal amount); + void addCost(BudgetKey key, BigDecimal limit, Currency currency, Cost amount); } diff --git a/token-pilot-budget/src/main/java/io/tokenpilot/budget/BudgetWindow.java b/token-pilot-budget/src/main/java/io/tokenpilot/budget/BudgetWindow.java new file mode 100644 index 0000000..a5d13b3 --- /dev/null +++ b/token-pilot-budget/src/main/java/io/tokenpilot/budget/BudgetWindow.java @@ -0,0 +1,33 @@ +package io.tokenpilot.budget; + +import java.time.Clock; +import java.time.YearMonth; +import java.time.ZoneId; +import java.util.Objects; + +/** + * 설정된 시간대 기준의 월별 예산 기간입니다. + * + * @param value 예산이 적용되는 연월 + */ +public record BudgetWindow(YearMonth value) { + + public BudgetWindow { + Objects.requireNonNull(value, "value must not be null"); + } + + public static BudgetWindow resolve(Clock clock, ZoneId zoneId) { + Objects.requireNonNull(clock, "clock must not be null"); + Objects.requireNonNull(zoneId, "zoneId must not be null"); + return new BudgetWindow(YearMonth.now(clock.withZone(zoneId))); + } + + public static BudgetWindow parse(String value) { + return new BudgetWindow(YearMonth.parse(value)); + } + + @Override + public String toString() { + return value.toString(); + } +} diff --git a/token-pilot-budget/src/main/java/io/tokenpilot/budget/internal/DefaultBudgetEvaluator.java b/token-pilot-budget/src/main/java/io/tokenpilot/budget/internal/DefaultBudgetEvaluator.java index f6659f3..78ccf50 100644 --- a/token-pilot-budget/src/main/java/io/tokenpilot/budget/internal/DefaultBudgetEvaluator.java +++ b/token-pilot-budget/src/main/java/io/tokenpilot/budget/internal/DefaultBudgetEvaluator.java @@ -2,126 +2,107 @@ import io.tokenpilot.budget.BudgetDecision; import io.tokenpilot.budget.BudgetEvaluator; +import io.tokenpilot.budget.BudgetKey; +import io.tokenpilot.budget.BudgetPolicy; import io.tokenpilot.budget.BudgetState; import io.tokenpilot.budget.BudgetStateStore; import io.tokenpilot.budget.BudgetThreshold; +import io.tokenpilot.budget.BudgetWindow; import io.tokenpilot.budget.exception.BudgetExceededException; import java.math.BigDecimal; +import java.time.Clock; import java.util.Map; +import java.util.Objects; public class DefaultBudgetEvaluator implements BudgetEvaluator { private final BudgetStateStore store; - private final BigDecimal monthlyLimit; + private final BudgetPolicy policy; + private final Clock clock; - public DefaultBudgetEvaluator( - BudgetStateStore store, - BigDecimal monthlyLimit - ) { - this.store = store; - this.monthlyLimit = monthlyLimit; + public DefaultBudgetEvaluator(BudgetStateStore store, BudgetPolicy policy, Clock clock) { + this.store = Objects.requireNonNull(store, "store must not be null"); + this.policy = Objects.requireNonNull(policy, "policy must not be null"); + this.clock = Objects.requireNonNull(clock, "clock must not be null"); } @Override - public BudgetDecision evaluate( - Map tags, - BigDecimal costAmount - ) { - - BigDecimal currentUsage = store.getAccumulatedCost(tags); - BigDecimal newUsage = currentUsage.add(costAmount); - - BigDecimal halfThreshold = monthlyLimit.multiply(BigDecimal.valueOf(0.5)); - BigDecimal warningThreshold = monthlyLimit.multiply(BigDecimal.valueOf(0.8)); - - if (newUsage.compareTo(monthlyLimit) >= 0) { - BudgetDecision decision = new BudgetDecision( - BudgetState.BLOCK, - BudgetThreshold.EXCEEDED, - "월 예산을 초과했습니다", - newUsage, - monthlyLimit - ); + public BudgetDecision evaluate(Map tags, BigDecimal costAmount) { + BudgetKey key = resolveKey(tags); + BigDecimal usage = store.getAccumulatedCost( + key, + policy.monthlyLimit(), + policy.currency() + ).add(costAmount); + BudgetDecision decision = decide(key, usage); + if (decision.state() == BudgetState.BLOCK) { throw new BudgetExceededException(decision); } + return decision; + } - // ✅ 80% 이상 - if (newUsage.compareTo(warningThreshold) >= 0) { - return new BudgetDecision( - BudgetState.WARN, - BudgetThreshold.WARNING, - "월 예산의 80% 이상 사용", - newUsage, - monthlyLimit - ); - } + @Override + public BudgetDecision evaluate(Map tags) { + BudgetKey key = resolveKey(tags); + BigDecimal usage = store.getAccumulatedCost( + key, + policy.monthlyLimit(), + policy.currency() + ); + return decide(key, usage); + } - // ✅ 50% 이상 - if (newUsage.compareTo(halfThreshold) >= 0) { - return new BudgetDecision( - BudgetState.ALLOW, - BudgetThreshold.HALF, - "월 예산의 50% 이상 사용", - newUsage, - monthlyLimit + private BudgetKey resolveKey(Map tags) { + Objects.requireNonNull(tags, "tags must not be null"); + String targetId = tags.get(policy.targetTagKey()); + if (targetId == null || targetId.isBlank()) { + targetId = policy.fallbackTargetId(); + } + if (targetId == null) { + throw new IllegalArgumentException( + "Missing required budget target tag: " + policy.targetTagKey() ); } - - // ✅ 정상 - return new BudgetDecision( - BudgetState.ALLOW, - BudgetThreshold.NONE, - "정상 범위 사용", - newUsage, - monthlyLimit + return new BudgetKey( + policy.id(), + policy.targetType(), + targetId, + BudgetWindow.resolve(clock, policy.zoneId()) ); } - @Override - public BudgetDecision evaluate(Map tags) { + private BudgetDecision decide(BudgetKey key, BigDecimal usage) { + BigDecimal halfThreshold = policy.monthlyLimit().multiply(new BigDecimal("0.5")); + BigDecimal warningThreshold = policy.monthlyLimit().multiply(new BigDecimal("0.8")); - BigDecimal usage = store.getAccumulatedCost(tags); - - BigDecimal halfThreshold = monthlyLimit.multiply(BigDecimal.valueOf(0.5)); - BigDecimal warningThreshold = monthlyLimit.multiply(BigDecimal.valueOf(0.8)); - - if (usage.compareTo(monthlyLimit) >= 0) { - return new BudgetDecision( - BudgetState.BLOCK, - BudgetThreshold.EXCEEDED, - "월 예산을 초과했습니다", - usage, - monthlyLimit - ); + if (usage.compareTo(policy.monthlyLimit()) >= 0) { + return decision(key, BudgetState.BLOCK, BudgetThreshold.EXCEEDED, "월 예산을 초과했습니다", usage); } - if (usage.compareTo(warningThreshold) >= 0) { - return new BudgetDecision( - BudgetState.WARN, - BudgetThreshold.WARNING, - "월 예산의 80% 이상 사용", - usage, - monthlyLimit - ); + return decision(key, BudgetState.WARN, BudgetThreshold.WARNING, "월 예산의 80% 이상 사용", usage); } - if (usage.compareTo(halfThreshold) >= 0) { - return new BudgetDecision( - BudgetState.ALLOW, - BudgetThreshold.HALF, - "월 예산의 50% 이상 사용", - usage, - monthlyLimit - ); + return decision(key, BudgetState.ALLOW, BudgetThreshold.HALF, "월 예산의 50% 이상 사용", usage); } + return decision(key, BudgetState.ALLOW, BudgetThreshold.NONE, "정상 범위 사용", usage); + } + private BudgetDecision decision( + BudgetKey key, + BudgetState state, + BudgetThreshold threshold, + String reason, + BigDecimal usage + ) { return new BudgetDecision( - BudgetState.ALLOW, - BudgetThreshold.NONE, - "정상 범위 사용", + key, + state, + threshold, + reason, usage, - monthlyLimit + policy.monthlyLimit(), + policy.currency() ); } } diff --git a/token-pilot-budget/src/main/java/io/tokenpilot/budget/internal/InMemoryBudgetStateStore.java b/token-pilot-budget/src/main/java/io/tokenpilot/budget/internal/InMemoryBudgetStateStore.java index ea68db9..5754e7b 100644 --- a/token-pilot-budget/src/main/java/io/tokenpilot/budget/internal/InMemoryBudgetStateStore.java +++ b/token-pilot-budget/src/main/java/io/tokenpilot/budget/internal/InMemoryBudgetStateStore.java @@ -1,36 +1,62 @@ package io.tokenpilot.budget.internal; +import io.tokenpilot.budget.BudgetKey; import io.tokenpilot.budget.BudgetStateStore; +import io.tokenpilot.core.domain.Cost; import java.math.BigDecimal; -import java.util.Map; +import java.util.Currency; +import java.util.Objects; import java.util.concurrent.ConcurrentHashMap; - +import java.util.concurrent.ConcurrentMap; /** - * BudgetStateStore의 인메모리 기반 구현체입니다. - * - * 예산 사용량을 메모리 내에서 누적 관리하며, - * 테스트 및 간단한 실행 환경을 위한 구현입니다. + * resolved {@link BudgetKey}별 비용을 관리하는 인메모리 저장소입니다. */ - public class InMemoryBudgetStateStore implements BudgetStateStore { - private final Map store = new ConcurrentHashMap<>(); + private final ConcurrentMap store = new ConcurrentHashMap<>(); - private String key(Map tags) { - return tags.getOrDefault("tenant_id", "default"); + @Override + public BigDecimal getAccumulatedCost(BudgetKey key, BigDecimal limit, Currency currency) { + validateArguments(key, limit, currency); + Bucket bucket = store.computeIfAbsent( + key, + ignored -> new Bucket(limit, currency, BigDecimal.ZERO) + ); + bucket.validate(limit, currency); + return bucket.accumulatedCost(); } @Override - public BigDecimal getAccumulatedCost(Map tags) { - // 아직 사용 기록이 없으면 0원 - return store.getOrDefault(key(tags), BigDecimal.ZERO); + public void addCost(BudgetKey key, BigDecimal limit, Currency currency, Cost amount) { + validateArguments(key, limit, currency); + Objects.requireNonNull(amount, "amount must not be null"); + if (!currency.equals(amount.currency())) { + throw new IllegalArgumentException("Budget currency does not match cost currency"); + } + + store.compute(key, (ignored, bucket) -> { + if (bucket == null) { + return new Bucket(limit, currency, amount.value()); + } + bucket.validate(limit, currency); + return new Bucket(limit, currency, bucket.accumulatedCost().add(amount.value())); + }); } - @Override - public void addCost(Map tags, BigDecimal amount) { - // 기존 값에 amount를 더함 - store.merge(key(tags), amount, BigDecimal::add); + private void validateArguments(BudgetKey key, BigDecimal limit, Currency currency) { + Objects.requireNonNull(key, "key must not be null"); + Objects.requireNonNull(limit, "limit must not be null"); + Objects.requireNonNull(currency, "currency must not be null"); + } + + private record Bucket(BigDecimal limit, Currency currency, BigDecimal accumulatedCost) { + + private void validate(BigDecimal expectedLimit, Currency expectedCurrency) { + if (limit.compareTo(expectedLimit) != 0 || !currency.equals(expectedCurrency)) { + throw new IllegalArgumentException("Budget policy snapshot changed for an existing key"); + } + } } } diff --git a/token-pilot-budget/src/main/java/io/tokenpilot/budget/internal/LedgerBudgetComponents.java b/token-pilot-budget/src/main/java/io/tokenpilot/budget/internal/LedgerBudgetComponents.java index 2cb5bbc..996b590 100644 --- a/token-pilot-budget/src/main/java/io/tokenpilot/budget/internal/LedgerBudgetComponents.java +++ b/token-pilot-budget/src/main/java/io/tokenpilot/budget/internal/LedgerBudgetComponents.java @@ -1,8 +1,11 @@ package io.tokenpilot.budget.internal; import io.tokenpilot.budget.BudgetEvaluator; +import io.tokenpilot.budget.BudgetPolicy; import io.tokenpilot.budget.BudgetStateStore; +import java.time.Clock; + /** * 예산 제어 컴포넌트 생성을 위한 팩토리 클래스입니다. */ @@ -15,7 +18,11 @@ public static BudgetStateStore inMemoryBudgetStateStore() { return new InMemoryBudgetStateStore(); } - public static BudgetEvaluator defaultBudgetEvaluator(BudgetStateStore store, java.math.BigDecimal monthlyLimit) { - return new DefaultBudgetEvaluator(store, monthlyLimit); + public static BudgetEvaluator defaultBudgetEvaluator( + BudgetStateStore store, + BudgetPolicy policy, + Clock clock + ) { + return new DefaultBudgetEvaluator(store, policy, clock); } } diff --git a/token-pilot-budget/src/test/java/io/tokenpilot/budget/internal/DefaultBudgetEvaluatorTest.java b/token-pilot-budget/src/test/java/io/tokenpilot/budget/internal/DefaultBudgetEvaluatorTest.java index 1a1cd90..ff5740e 100644 --- a/token-pilot-budget/src/test/java/io/tokenpilot/budget/internal/DefaultBudgetEvaluatorTest.java +++ b/token-pilot-budget/src/test/java/io/tokenpilot/budget/internal/DefaultBudgetEvaluatorTest.java @@ -1,117 +1,179 @@ package io.tokenpilot.budget.internal; import io.tokenpilot.budget.BudgetDecision; +import io.tokenpilot.budget.BudgetKey; +import io.tokenpilot.budget.BudgetPolicy; import io.tokenpilot.budget.BudgetState; -import io.tokenpilot.budget.BudgetThreshold; import io.tokenpilot.budget.BudgetStateStore; +import io.tokenpilot.budget.BudgetThreshold; +import io.tokenpilot.budget.BudgetWindow; import io.tokenpilot.budget.exception.BudgetExceededException; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; import java.math.BigDecimal; +import java.time.Clock; +import java.time.Instant; +import java.time.ZoneId; +import java.time.ZoneOffset; +import java.util.ArrayList; +import java.util.Currency; +import java.util.List; import java.util.Map; - -import static org.assertj.core.api.Assertions.*; -import static org.mockito.Mockito.*; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; class DefaultBudgetEvaluatorTest { - private BudgetStateStore store; - private DefaultBudgetEvaluator evaluator; + private static final Currency USD = Currency.getInstance("USD"); + private static final Map TAGS = Map.of("tenant_id", "tenant-a"); - private static final BigDecimal MONTHLY_LIMIT = new BigDecimal("100.00"); - private static final Map TAGS = Map.of("service", "test"); + private BudgetStateStore store; @BeforeEach void setUp() { store = mock(BudgetStateStore.class); - evaluator = new DefaultBudgetEvaluator(store, MONTHLY_LIMIT); - } - - // ─── evaluate(tags, costAmount) ─────────────────────────────────────────── - - @Test - void 정상_범위_사용() { - when(store.getAccumulatedCost(TAGS)).thenReturn(new BigDecimal("10.00")); - - BudgetDecision result = evaluator.evaluate(TAGS, new BigDecimal("10.00")); - - assertThat(result.state()).isEqualTo(BudgetState.ALLOW); - assertThat(result.threshold()).isEqualTo(BudgetThreshold.NONE); } @Test - void 누적_50퍼센트_이상() { - when(store.getAccumulatedCost(TAGS)).thenReturn(new BigDecimal("40.00")); - - BudgetDecision result = evaluator.evaluate(TAGS, new BigDecimal("10.00")); - - assertThat(result.state()).isEqualTo(BudgetState.ALLOW); - assertThat(result.threshold()).isEqualTo(BudgetThreshold.HALF); - } - - @Test - void 누적_80퍼센트_이상() { - when(store.getAccumulatedCost(TAGS)).thenReturn(new BigDecimal("70.00")); + void 예상_비용을_포함해_임계치를_판정한다() { + DefaultBudgetEvaluator evaluator = evaluator(policy(null, ZoneOffset.UTC), "2026-07-22T00:00:00Z"); + when(store.getAccumulatedCost(any(), any(), any())).thenReturn(new BigDecimal("70.00")); BudgetDecision result = evaluator.evaluate(TAGS, new BigDecimal("10.00")); assertThat(result.state()).isEqualTo(BudgetState.WARN); assertThat(result.threshold()).isEqualTo(BudgetThreshold.WARNING); + assertThat(result.key()).isEqualTo(key("policy-a", "tenant-a", "2026-07")); } @Test - void 예산_초과시_예외_발생() { - when(store.getAccumulatedCost(TAGS)).thenReturn(new BigDecimal("95.00")); + void 예산_초과시_동일_key를_포함한_예외가_발생한다() { + DefaultBudgetEvaluator evaluator = evaluator(policy(null, ZoneOffset.UTC), "2026-07-22T00:00:00Z"); + when(store.getAccumulatedCost(any(), any(), any())).thenReturn(new BigDecimal("95.00")); assertThatThrownBy(() -> evaluator.evaluate(TAGS, new BigDecimal("10.00"))) .isInstanceOf(BudgetExceededException.class) - .extracting(e -> ((BudgetExceededException) e).getDecision()) + .extracting(exception -> ((BudgetExceededException) exception).getDecision()) .satisfies(decision -> { assertThat(decision.state()).isEqualTo(BudgetState.BLOCK); - assertThat(decision.threshold()).isEqualTo(BudgetThreshold.EXCEEDED); + assertThat(decision.key()).isEqualTo(key("policy-a", "tenant-a", "2026-07")); }); } - // ─── evaluate(tags) ─────────────────────────────────────────────────────── - @Test - void 현재_사용량_정상() { - when(store.getAccumulatedCost(TAGS)).thenReturn(new BigDecimal("30.00")); + void target_identity_누락은_fail_closed다() { + DefaultBudgetEvaluator evaluator = evaluator(policy(null, ZoneOffset.UTC), "2026-07-22T00:00:00Z"); - BudgetDecision result = evaluator.evaluate(TAGS); + assertThatThrownBy(() -> evaluator.evaluate(Map.of("service", "chat"))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("tenant_id"); - assertThat(result.state()).isEqualTo(BudgetState.ALLOW); - assertThat(result.threshold()).isEqualTo(BudgetThreshold.NONE); + verify(store, never()).getAccumulatedCost(any(), any(), any()); } @Test - void 현재_사용량_50퍼센트_이상() { - when(store.getAccumulatedCost(TAGS)).thenReturn(new BigDecimal("50.00")); + void 명시적_fallback이_설정된_경우에만_fallback_bucket을_사용한다() { + DefaultBudgetEvaluator evaluator = evaluator( + policy("shared-fallback", ZoneOffset.UTC), + "2026-07-22T00:00:00Z" + ); + when(store.getAccumulatedCost(any(), any(), any())).thenReturn(BigDecimal.ZERO); + + BudgetDecision decision = evaluator.evaluate(Map.of("service", "chat")); - BudgetDecision result = evaluator.evaluate(TAGS); + assertThat(decision.key().targetId()).isEqualTo("shared-fallback"); + } - assertThat(result.state()).isEqualTo(BudgetState.ALLOW); - assertThat(result.threshold()).isEqualTo(BudgetThreshold.HALF); + @ParameterizedTest + @CsvSource({ + "2026-07-31T23:59:59Z, 2026-07", + "2026-08-01T00:00:00Z, 2026-08", + "2028-02-29T23:59:59Z, 2028-02", + "2028-03-01T00:00:00Z, 2028-03", + "2026-12-31T23:59:59Z, 2026-12", + "2027-01-01T00:00:00Z, 2027-01" + }) + void Clock_fixed로_월_경계를_결정한다(String instant, String expectedWindow) { + DefaultBudgetEvaluator evaluator = evaluator(policy(null, ZoneOffset.UTC), instant); + when(store.getAccumulatedCost(any(), any(), any())).thenReturn(BigDecimal.ZERO); + + BudgetDecision decision = evaluator.evaluate(TAGS); + + assertThat(decision.key().window()).isEqualTo(BudgetWindow.parse(expectedWindow)); } @Test - void 현재_사용량_80퍼센트_이상() { - when(store.getAccumulatedCost(TAGS)).thenReturn(new BigDecimal("80.00")); + void UTC와_설정_ZoneId의_월이_다르면_설정_ZoneId를_따른다() { + String instant = "2026-07-31T15:30:00Z"; + when(store.getAccumulatedCost(any(), any(), any())).thenReturn(BigDecimal.ZERO); - BudgetDecision result = evaluator.evaluate(TAGS); + BudgetDecision utc = evaluator(policy(null, ZoneOffset.UTC), instant).evaluate(TAGS); + BudgetDecision seoul = evaluator(policy(null, ZoneId.of("Asia/Seoul")), instant).evaluate(TAGS); - assertThat(result.state()).isEqualTo(BudgetState.WARN); - assertThat(result.threshold()).isEqualTo(BudgetThreshold.WARNING); + assertThat(utc.key().window()).isEqualTo(BudgetWindow.parse("2026-07")); + assertThat(seoul.key().window()).isEqualTo(BudgetWindow.parse("2026-08")); } @Test - void 현재_사용량_100퍼센트_이상() { - when(store.getAccumulatedCost(TAGS)).thenReturn(new BigDecimal("100.00")); + void 동시_요청에서도_key_생성_결과가_결정적이다() throws Exception { + InMemoryBudgetStateStore stateStore = new InMemoryBudgetStateStore(); + DefaultBudgetEvaluator evaluator = new DefaultBudgetEvaluator( + stateStore, + policy(null, ZoneOffset.UTC), + Clock.fixed(Instant.parse("2026-07-22T00:00:00Z"), ZoneOffset.UTC) + ); + Set keys = ConcurrentHashMap.newKeySet(); + List> futures = new ArrayList<>(); + + try (var executor = Executors.newFixedThreadPool(8)) { + for (int index = 0; index < 200; index++) { + futures.add(executor.submit(() -> evaluator.evaluate(TAGS).key())); + } + executor.shutdown(); + for (Future future : futures) { + keys.add(future.get(5, TimeUnit.SECONDS)); + } + assertThat(executor.awaitTermination(5, TimeUnit.SECONDS)).isTrue(); + } + + assertThat(keys).containsExactly(key("policy-a", "tenant-a", "2026-07")); + } - BudgetDecision result = evaluator.evaluate(TAGS); + private DefaultBudgetEvaluator evaluator(BudgetPolicy policy, String instant) { + return new DefaultBudgetEvaluator( + store, + policy, + Clock.fixed(Instant.parse(instant), ZoneOffset.UTC) + ); + } + + private static BudgetPolicy policy(String fallbackTargetId, ZoneId zoneId) { + return new BudgetPolicy( + "policy-a", + "tenant", + "tenant_id", + fallbackTargetId, + new BigDecimal("100.00"), + USD, + zoneId + ); + } - assertThat(result.state()).isEqualTo(BudgetState.BLOCK); - assertThat(result.threshold()).isEqualTo(BudgetThreshold.EXCEEDED); + private static BudgetKey key(String policyId, String targetId, String window) { + return new BudgetKey(policyId, "tenant", targetId, BudgetWindow.parse(window)); } } diff --git a/token-pilot-budget/src/test/java/io/tokenpilot/budget/internal/InMemoryBudgetStateStoreTest.java b/token-pilot-budget/src/test/java/io/tokenpilot/budget/internal/InMemoryBudgetStateStoreTest.java new file mode 100644 index 0000000..9a816bc --- /dev/null +++ b/token-pilot-budget/src/test/java/io/tokenpilot/budget/internal/InMemoryBudgetStateStoreTest.java @@ -0,0 +1,101 @@ +package io.tokenpilot.budget.internal; + +import io.tokenpilot.budget.BudgetKey; +import io.tokenpilot.budget.BudgetWindow; +import io.tokenpilot.core.domain.Cost; +import org.junit.jupiter.api.Test; + +import java.math.BigDecimal; +import java.util.Currency; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +class InMemoryBudgetStateStoreTest { + + private static final Currency USD = Currency.getInstance("USD"); + private static final BigDecimal LIMIT = new BigDecimal("100.00"); + + @Test + void 같은_policy_target_window는_계속_누적된다() { + InMemoryBudgetStateStore store = new InMemoryBudgetStateStore(); + BudgetKey key = key("policy-a", "tenant-a", "2026-07"); + + add(store, key, "10.00", USD); + add(store, key, "5.00", USD); + + assertThat(store.getAccumulatedCost(key, LIMIT, USD)).isEqualByComparingTo("15.00"); + } + + @Test + void 다음_월과_다른_target_policy는_각각_격리된다() { + InMemoryBudgetStateStore store = new InMemoryBudgetStateStore(); + BudgetKey july = key("policy-a", "tenant-a", "2026-07"); + BudgetKey august = key("policy-a", "tenant-a", "2026-08"); + BudgetKey otherTarget = key("policy-a", "tenant-b", "2026-07"); + BudgetKey otherPolicy = key("policy-b", "tenant-a", "2026-07"); + + add(store, july, "10.00", USD); + add(store, otherTarget, "20.00", USD); + add(store, otherPolicy, "30.00", USD); + + assertThat(store.getAccumulatedCost(july, LIMIT, USD)).isEqualByComparingTo("10.00"); + assertThat(store.getAccumulatedCost(august, LIMIT, USD)).isEqualByComparingTo("0"); + assertThat(store.getAccumulatedCost(otherTarget, LIMIT, USD)).isEqualByComparingTo("20.00"); + assertThat(store.getAccumulatedCost(otherPolicy, LIMIT, USD)).isEqualByComparingTo("30.00"); + } + + @Test + void currency_mismatch는_상태를_변경하지_않고_실패한다() { + InMemoryBudgetStateStore store = new InMemoryBudgetStateStore(); + BudgetKey key = key("policy-a", "tenant-a", "2026-07"); + add(store, key, "10.00", USD); + + assertThatThrownBy(() -> store.addCost( + key, + LIMIT, + USD, + new Cost(new BigDecimal("1000"), Currency.getInstance("KRW")) + )).isInstanceOf(IllegalArgumentException.class); + + assertThat(store.getAccumulatedCost(key, LIMIT, USD)).isEqualByComparingTo("10.00"); + } + + @Test + void 최초_조회에서_limit과_currency_snapshot을_고정한다() { + InMemoryBudgetStateStore store = new InMemoryBudgetStateStore(); + BudgetKey key = key("policy-a", "tenant-a", "2026-07"); + + assertThat(store.getAccumulatedCost(key, LIMIT, USD)).isEqualByComparingTo("0"); + + assertThatThrownBy(() -> store.addCost( + key, + new BigDecimal("200.00"), + USD, + new Cost(BigDecimal.ONE, USD) + )).isInstanceOf(IllegalArgumentException.class); + + Currency krw = Currency.getInstance("KRW"); + assertThatThrownBy(() -> store.addCost( + key, + LIMIT, + krw, + new Cost(BigDecimal.ONE, krw) + )).isInstanceOf(IllegalArgumentException.class); + + assertThat(store.getAccumulatedCost(key, LIMIT, USD)).isEqualByComparingTo("0"); + } + + private static void add( + InMemoryBudgetStateStore store, + BudgetKey key, + String amount, + Currency currency + ) { + store.addCost(key, LIMIT, currency, new Cost(new BigDecimal(amount), currency)); + } + + private static BudgetKey key(String policyId, String targetId, String window) { + return new BudgetKey(policyId, "tenant", targetId, BudgetWindow.parse(window)); + } +} diff --git a/token-pilot-notification/src/main/java/io/tokenpilot/notification/BudgetNotificationEvent.java b/token-pilot-notification/src/main/java/io/tokenpilot/notification/BudgetNotificationEvent.java index 5fa760a..859864a 100644 --- a/token-pilot-notification/src/main/java/io/tokenpilot/notification/BudgetNotificationEvent.java +++ b/token-pilot-notification/src/main/java/io/tokenpilot/notification/BudgetNotificationEvent.java @@ -1,5 +1,6 @@ package io.tokenpilot.notification; +import io.tokenpilot.budget.BudgetKey; import io.tokenpilot.budget.BudgetState; import io.tokenpilot.budget.BudgetThreshold; @@ -10,12 +11,11 @@ * 예산 임계치 도달 시 발생하는 알림 이벤트 */ public record BudgetNotificationEvent( - String targetId, - String budgetWindow, + BudgetKey key, BudgetThreshold threshold, BudgetState state, String reason, BigDecimal currentUsage, BigDecimal limit, Map tags -) {} \ No newline at end of file +) {} diff --git a/token-pilot-notification/src/main/java/io/tokenpilot/notification/BudgetNotificationService.java b/token-pilot-notification/src/main/java/io/tokenpilot/notification/BudgetNotificationService.java index 8a7912c..0e1bc7b 100644 --- a/token-pilot-notification/src/main/java/io/tokenpilot/notification/BudgetNotificationService.java +++ b/token-pilot-notification/src/main/java/io/tokenpilot/notification/BudgetNotificationService.java @@ -26,8 +26,6 @@ public BudgetNotificationService( */ public void notifyIfNeeded( BudgetDecision decision, - String targetId, - String budgetWindow, Map tags ) { BudgetThreshold current = decision.threshold(); @@ -36,8 +34,7 @@ public void notifyIfNeeded( return; } - BudgetThreshold last = - store.getLastNotifiedThreshold(targetId, budgetWindow); + BudgetThreshold last = store.getLastNotifiedThreshold(decision.key()); // 같은 window에서 중복 방지 if (current.compareTo(last) <= 0) { @@ -46,8 +43,7 @@ public void notifyIfNeeded( BudgetNotificationEvent event = new BudgetNotificationEvent( - targetId, - budgetWindow, + decision.key(), current, decision.state(), decision.reason(), @@ -59,9 +55,8 @@ public void notifyIfNeeded( handler.handle(event); store.updateLastNotifiedThreshold( - targetId, - budgetWindow, + decision.key(), current ); } -} \ No newline at end of file +} diff --git a/token-pilot-notification/src/main/java/io/tokenpilot/notification/InMemoryNotificationStateStore.java b/token-pilot-notification/src/main/java/io/tokenpilot/notification/InMemoryNotificationStateStore.java index 97a3972..3616663 100644 --- a/token-pilot-notification/src/main/java/io/tokenpilot/notification/InMemoryNotificationStateStore.java +++ b/token-pilot-notification/src/main/java/io/tokenpilot/notification/InMemoryNotificationStateStore.java @@ -1,5 +1,6 @@ package io.tokenpilot.notification; +import io.tokenpilot.budget.BudgetKey; import io.tokenpilot.budget.BudgetThreshold; import java.util.Map; @@ -10,29 +11,18 @@ */ public class InMemoryNotificationStateStore implements NotificationStateStore { - private final Map store = new ConcurrentHashMap<>(); - - private String key(String targetId, String window) { - return targetId + ":" + window; - } + private final Map store = new ConcurrentHashMap<>(); @Override - public BudgetThreshold getLastNotifiedThreshold( - String targetId, - String budgetWindow - ) { - return store.getOrDefault( - key(targetId, budgetWindow), - BudgetThreshold.NONE - ); + public BudgetThreshold getLastNotifiedThreshold(BudgetKey key) { + return store.getOrDefault(key, BudgetThreshold.NONE); } @Override public void updateLastNotifiedThreshold( - String targetId, - String budgetWindow, + BudgetKey key, BudgetThreshold threshold ) { - store.put(key(targetId, budgetWindow), threshold); + store.put(key, threshold); } -} \ No newline at end of file +} diff --git a/token-pilot-notification/src/main/java/io/tokenpilot/notification/NotificationStateStore.java b/token-pilot-notification/src/main/java/io/tokenpilot/notification/NotificationStateStore.java index f0fb8f4..891cef3 100644 --- a/token-pilot-notification/src/main/java/io/tokenpilot/notification/NotificationStateStore.java +++ b/token-pilot-notification/src/main/java/io/tokenpilot/notification/NotificationStateStore.java @@ -1,5 +1,6 @@ package io.tokenpilot.notification; +import io.tokenpilot.budget.BudgetKey; import io.tokenpilot.budget.BudgetThreshold; /** @@ -7,14 +8,10 @@ */ public interface NotificationStateStore { - BudgetThreshold getLastNotifiedThreshold( - String targetId, - String budgetWindow - ); + BudgetThreshold getLastNotifiedThreshold(BudgetKey key); void updateLastNotifiedThreshold( - String targetId, - String budgetWindow, + BudgetKey key, BudgetThreshold threshold ); -} \ No newline at end of file +} diff --git a/token-pilot-notification/src/test/java/io/tokenpilot/notification/BudgetNotificationServiceTest.java b/token-pilot-notification/src/test/java/io/tokenpilot/notification/BudgetNotificationServiceTest.java index 688ff94..85c4553 100644 --- a/token-pilot-notification/src/test/java/io/tokenpilot/notification/BudgetNotificationServiceTest.java +++ b/token-pilot-notification/src/test/java/io/tokenpilot/notification/BudgetNotificationServiceTest.java @@ -1,112 +1,91 @@ package io.tokenpilot.notification; import io.tokenpilot.budget.BudgetDecision; +import io.tokenpilot.budget.BudgetKey; import io.tokenpilot.budget.BudgetState; import io.tokenpilot.budget.BudgetThreshold; +import io.tokenpilot.budget.BudgetWindow; import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; import java.math.BigDecimal; +import java.util.Currency; import java.util.Map; -import static org.mockito.Mockito.*; +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Mockito.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.same; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; -/** - * BudgetNotificationService 동작 검증 테스트 - */ class BudgetNotificationServiceTest { @Test - void threshold_increase_triggers_event_only_once_per_window() { - + void 같은_key에서는_증가한_threshold만_알린다() { BudgetNotificationHandler handler = mock(BudgetNotificationHandler.class); - NotificationStateStore store = new InMemoryNotificationStateStore(); - - BudgetNotificationService service = - new BudgetNotificationService(handler, store); - - String targetId = "user1"; - String window = "2026-06"; - - BudgetDecision decision50 = new BudgetDecision( - BudgetState.WARN, - BudgetThreshold.HALF, - "50%", - BigDecimal.valueOf(50), - BigDecimal.valueOf(100) - ); - - BudgetDecision decision80 = new BudgetDecision( - BudgetState.WARN, - BudgetThreshold.WARNING, - "80%", - BigDecimal.valueOf(80), - BigDecimal.valueOf(100) + BudgetNotificationService service = new BudgetNotificationService( + handler, + new InMemoryNotificationStateStore() ); + BudgetKey key = key("2026-06"); - // 50% 최초 → 발생 - service.notifyIfNeeded(decision50, targetId, window, Map.of()); - - // 50% 반복 → 발생 안됨 - service.notifyIfNeeded(decision50, targetId, window, Map.of()); + service.notifyIfNeeded(decision(key, BudgetThreshold.HALF, "50"), Map.of()); + service.notifyIfNeeded(decision(key, BudgetThreshold.HALF, "50"), Map.of()); + service.notifyIfNeeded(decision(key, BudgetThreshold.WARNING, "80"), Map.of()); - // 80% → 발생 - service.notifyIfNeeded(decision80, targetId, window, Map.of()); - - // 총 2번 발생해야 함 (50, 80) verify(handler, times(2)).handle(any()); } @Test - void same_threshold_should_not_trigger_duplicate_event() { - + void 새_window에서는_같은_threshold를_다시_알린다() { BudgetNotificationHandler handler = mock(BudgetNotificationHandler.class); - NotificationStateStore store = new InMemoryNotificationStateStore(); - - BudgetNotificationService service = - new BudgetNotificationService(handler, store); - - String targetId = "user1"; - String window = "2026-06"; - - BudgetDecision decision50 = new BudgetDecision( - BudgetState.WARN, - BudgetThreshold.HALF, - "50%", - BigDecimal.valueOf(50), - BigDecimal.valueOf(100) + BudgetNotificationService service = new BudgetNotificationService( + handler, + new InMemoryNotificationStateStore() ); - service.notifyIfNeeded(decision50, targetId, window, Map.of()); - service.notifyIfNeeded(decision50, targetId, window, Map.of()); + service.notifyIfNeeded(decision(key("2026-06"), BudgetThreshold.HALF, "50"), Map.of()); + service.notifyIfNeeded(decision(key("2026-07"), BudgetThreshold.HALF, "50"), Map.of()); - verify(handler, times(1)).handle(any()); + verify(handler, times(2)).handle(any()); } @Test - void new_window_should_allow_notification_again() { - + void event와_notification_store가_decision의_동일한_key를_사용한다() { BudgetNotificationHandler handler = mock(BudgetNotificationHandler.class); - NotificationStateStore store = new InMemoryNotificationStateStore(); - - BudgetNotificationService service = - new BudgetNotificationService(handler, store); + NotificationStateStore store = mock(NotificationStateStore.class); + BudgetNotificationService service = new BudgetNotificationService(handler, store); + BudgetKey key = key("2026-07"); + when(store.getLastNotifiedThreshold(key)).thenReturn(BudgetThreshold.NONE); + + service.notifyIfNeeded(decision(key, BudgetThreshold.HALF, "50"), Map.of()); + + ArgumentCaptor event = ArgumentCaptor.forClass(BudgetNotificationEvent.class); + verify(handler).handle(event.capture()); + verify(store).getLastNotifiedThreshold(same(key)); + verify(store).updateLastNotifiedThreshold(same(key), same(BudgetThreshold.HALF)); + assertThat(event.getValue().key()).isSameAs(key); + } - String targetId = "user1"; + private static BudgetKey key(String window) { + return new BudgetKey("policy-a", "tenant", "tenant-a", BudgetWindow.parse(window)); + } - BudgetDecision decision50 = new BudgetDecision( + private static BudgetDecision decision( + BudgetKey key, + BudgetThreshold threshold, + String usage + ) { + return new BudgetDecision( + key, BudgetState.WARN, - BudgetThreshold.HALF, - "50%", - BigDecimal.valueOf(50), - BigDecimal.valueOf(100) + threshold, + threshold.name(), + new BigDecimal(usage), + new BigDecimal("100"), + Currency.getInstance("USD") ); - - // 6월 - service.notifyIfNeeded(decision50, targetId, "2026-06", Map.of()); - - // 7월 → 다시 발생해야 함 - service.notifyIfNeeded(decision50, targetId, "2026-07", Map.of()); - - verify(handler, times(2)).handle(any()); } -} \ No newline at end of file +} diff --git a/token-pilot-sample-app/src/main/java/io/tokenpilot/sample/SampleController.java b/token-pilot-sample-app/src/main/java/io/tokenpilot/sample/SampleController.java index d901d32..5af1d91 100644 --- a/token-pilot-sample-app/src/main/java/io/tokenpilot/sample/SampleController.java +++ b/token-pilot-sample-app/src/main/java/io/tokenpilot/sample/SampleController.java @@ -85,7 +85,12 @@ public Map budget() { Map tags = Map.of("tenant_id", "budget-sample-tenant"); BudgetDecision initialDecision = evaluator.evaluate(tags, new BigDecimal("0.001")); - stateStore.addCost(tags, new BigDecimal("0.0045")); + stateStore.addCost( + initialDecision.key(), + initialDecision.limit(), + initialDecision.currency(), + new Cost(new BigDecimal("0.0045"), initialDecision.currency()) + ); try { evaluator.evaluate(tags, new BigDecimal("0.001")); @@ -100,7 +105,7 @@ public Map budget() { "enabled", "true", "initialState", initialDecision.state().name(), "blockedState", blockedDecision.state().name(), - "currentUsage", blockedDecision.currentUsage().toPlainString(), + "currentUsage", blockedDecision.currentUsage().stripTrailingZeros().toPlainString(), "limit", blockedDecision.limit().toPlainString() ); } diff --git a/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/DefaultLedgerAdvisor.java b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/DefaultLedgerAdvisor.java index 760f150..fd23606 100644 --- a/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/DefaultLedgerAdvisor.java +++ b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/DefaultLedgerAdvisor.java @@ -1,5 +1,6 @@ package io.tokenpilot.springai.internal; +import io.tokenpilot.budget.BudgetDecision; import io.tokenpilot.budget.BudgetEvaluator; import io.tokenpilot.budget.BudgetStateStore; import io.tokenpilot.core.*; @@ -23,6 +24,8 @@ */ public class DefaultLedgerAdvisor implements LedgerAdvisor { + static final String BUDGET_DECISION_CONTEXT = "tokenpilot.budget.decision"; + private final LedgerManager ledgerManager; private final UsageExtractor usageExtractor; private final BudgetEvaluator budgetEvaluator; @@ -49,7 +52,10 @@ public DefaultLedgerAdvisor(LedgerManager ledgerManager, UsageExtractor usageExt public ChatClientRequest before(ChatClientRequest request, AdvisorChain chain) { if (budgetEvaluator != null) { Map tags = extractTagsFromRequest(request); - budgetEvaluator.evaluate(tags); + BudgetDecision decision = budgetEvaluator.evaluate(tags); + return request.mutate() + .context(BUDGET_DECISION_CONTEXT, decision) + .build(); } return request; } @@ -68,7 +74,13 @@ public ChatClientResponse after(ChatClientResponse response, AdvisorChain chain) Optional plan = pricingRegistry.getPlan(modelId); if (plan.isPresent()) { Cost cost = costCalculator.calculate(usage, plan.get()); - budgetStateStore.addCost(tags, cost.value()); + BudgetDecision decision = extractBudgetDecision(response); + budgetStateStore.addCost( + decision.key(), + decision.limit(), + decision.currency(), + cost + ); } } @@ -100,6 +112,15 @@ private Map extractTags(ChatClientResponse response) { return tags; } + private BudgetDecision extractBudgetDecision(ChatClientResponse response) { + Map context = response.context(); + Object value = context == null ? null : context.get(BUDGET_DECISION_CONTEXT); + if (value instanceof BudgetDecision decision) { + return decision; + } + throw new IllegalStateException("Resolved budget decision is missing from response context"); + } + private Map extractTagsFromRequest(ChatClientRequest request) { Map tags = new HashMap<>(); Map context = request.context(); diff --git a/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/DefaultLedgerAdvisorTest.java b/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/DefaultLedgerAdvisorTest.java index 41028c1..62fbc25 100644 --- a/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/DefaultLedgerAdvisorTest.java +++ b/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/DefaultLedgerAdvisorTest.java @@ -2,8 +2,11 @@ import io.tokenpilot.budget.BudgetDecision; import io.tokenpilot.budget.BudgetEvaluator; +import io.tokenpilot.budget.BudgetKey; import io.tokenpilot.budget.BudgetState; import io.tokenpilot.budget.BudgetStateStore; +import io.tokenpilot.budget.BudgetThreshold; +import io.tokenpilot.budget.BudgetWindow; import io.tokenpilot.core.*; import io.tokenpilot.core.domain.*; import io.tokenpilot.springai.LedgerAdvisor; @@ -16,6 +19,7 @@ import org.springframework.ai.chat.metadata.ChatResponseMetadata; import org.springframework.ai.chat.model.ChatResponse; import org.springframework.ai.chat.model.Generation; +import org.springframework.ai.chat.prompt.Prompt; import java.math.BigDecimal; import java.util.Currency; @@ -62,6 +66,7 @@ void recordAfterAIResponse() { void recordBudgetAfterAIResponse() { LedgerManager ledgerManager = mock(LedgerManager.class); UsageExtractor extractor = mock(UsageExtractor.class); + BudgetEvaluator budgetEvaluator = mock(BudgetEvaluator.class); BudgetStateStore budgetStateStore = mock(BudgetStateStore.class); CostCalculator costCalculator = mock(CostCalculator.class); PricingRegistry pricingRegistry = mock(PricingRegistry.class); @@ -69,24 +74,33 @@ void recordBudgetAfterAIResponse() { TokenUsage mockUsage = TokenUsage.from(100, 200); PricingPlan mockPlan = new PricingPlan("gpt-4o", new BigDecimal("0.01"), new BigDecimal("0.03"), Currency.getInstance("USD")); Cost mockCost = new Cost(new BigDecimal("0.5"), Currency.getInstance("USD")); + BudgetDecision budgetDecision = decision(); when(extractor.extract(any())).thenReturn(mockUsage); when(pricingRegistry.getPlan("gpt-4o")).thenReturn(Optional.of(mockPlan)); when(costCalculator.calculate(mockUsage, mockPlan)).thenReturn(mockCost); + when(budgetEvaluator.evaluate(anyMap())).thenReturn(budgetDecision); DefaultLedgerAdvisor advisor = new DefaultLedgerAdvisor(ledgerManager, extractor, - null, budgetStateStore, costCalculator, pricingRegistry); + budgetEvaluator, budgetStateStore, costCalculator, pricingRegistry); + + ChatClientRequest request = new ChatClientRequest( + new Prompt("test"), + Map.of("tenant_id", "tenant-abc") + ); + ChatClientRequest resolvedRequest = advisor.before(request, mock(AdvisorChain.class)); ChatResponseMetadata metadata = ChatResponseMetadata.builder().model("gpt-4o").build(); ChatResponse chatResponse = new ChatResponse(List.of(new Generation(new org.springframework.ai.chat.messages.AssistantMessage("test"))), metadata); - Map context = Map.of("tenant_id", "tenant-abc"); - ChatClientResponse response = new ChatClientResponse(chatResponse, context); + ChatClientResponse response = new ChatClientResponse(chatResponse, resolvedRequest.context()); advisor.after(response, mock(AdvisorChain.class)); verify(budgetStateStore, times(1)).addCost( - argThat(tags -> tags.get("tenant_id").equals("tenant-abc")), - argThat(amount -> amount.compareTo(new BigDecimal("0.5")) == 0) + same(budgetDecision.key()), + same(budgetDecision.limit()), + same(budgetDecision.currency()), + same(mockCost) ); } @@ -94,18 +108,23 @@ void recordBudgetAfterAIResponse() { @DisplayName("AI 호출 전 BudgetEvaluator를 통해 예산을 체크해야 한다") void checkBudgetBeforeAIRequest() { BudgetEvaluator budgetEvaluator = mock(BudgetEvaluator.class); + BudgetDecision decision = decision(); + when(budgetEvaluator.evaluate(anyMap())).thenReturn(decision); DefaultLedgerAdvisor advisor = new DefaultLedgerAdvisor(mock(LedgerManager.class), mock(UsageExtractor.class), budgetEvaluator, null, null, null); - ChatClientRequest request = mock(ChatClientRequest.class); - Map context = Map.of("tenant_id", "tenant-abc"); - when(request.context()).thenReturn(context); + ChatClientRequest request = new ChatClientRequest( + new Prompt("test"), + Map.of("tenant_id", "tenant-abc") + ); - advisor.before(request, mock(AdvisorChain.class)); + ChatClientRequest resolvedRequest = advisor.before(request, mock(AdvisorChain.class)); verify(budgetEvaluator, times(1)).evaluate( argThat(tags -> tags.get("tenant_id").equals("tenant-abc")) ); + assertThat(resolvedRequest.context().get(DefaultLedgerAdvisor.BUDGET_DECISION_CONTEXT)) + .isSameAs(decision); } @Test @@ -116,4 +135,21 @@ void checkAdvisorMetadata() { assertThat(advisor.getName()).isEqualTo("LedgerAdvisor"); assertThat(advisor.getOrder()).isEqualTo(0); } + + private static BudgetDecision decision() { + return new BudgetDecision( + new BudgetKey( + "policy-a", + "tenant", + "tenant-abc", + BudgetWindow.parse("2026-07") + ), + BudgetState.ALLOW, + BudgetThreshold.NONE, + "allowed", + BigDecimal.ZERO, + new BigDecimal("100.00"), + Currency.getInstance("USD") + ); + } }