diff --git a/AGENTS.md b/AGENTS.md index 72e3a05..d7f32b1 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -50,6 +50,7 @@ Never describe a roadmap item as an implemented or published capability. - `token-pilot-notification`은 알림 이벤트 발행과 중복 방지 로직만 담당한다. - 실제 메일/Slack/Webhook 전송은 사용자 애플리케이션의 `BudgetNotificationHandler` 구현체가 담당한다. - 라이브러리 내부에서 SMTP 설정이나 외부 메일 서비스를 기본 흐름으로 포함하지 않는다. +- 알림 전달은 process-local best-effort이며 durable outbox나 재시작 후 replay를 제공하지 않는다. ## Architecture Decision: Framework Independence and Observability @@ -75,7 +76,7 @@ Token Pilot의 제품 포지션은 framework-independent Java LLM control and ac | `token-pilot-spring-ai` | Basic implementation complete | Spring AI 2.0.0 `UsageExtractor`, `LedgerAdvisor`, pricing snapshot resolution, response usage recording, reconciliation decisions, and legacy provider-boundary BLOCK enforcement | | `token-pilot-micrometer` | Basic implementation complete | `MetricsOptions`, tag whitelist, and metric metadata exist; metric ownership must be narrowed | | `token-pilot-budget` | Atomic reservation and reconciliation implemented | Typed monthly keys, Clock/ZoneId windows, safe-upper-bound reservations, commit/release/write-off lifecycle, pending reconciliation liability, estimate/actual token and cost deltas, duplicate callback protection, and framework-independent best-effort accounting events implemented; candidate production and durable stores remain | -| `token-pilot-notification` | Basic implementation complete | Event API and deduplication exist; not yet connected to the full advisor/budget lifecycle | +| `token-pilot-notification` | Atomic accounting integration implemented | Commit, reconciliation-required, late reconciliation, and reservation BLOCK results produce process-local deduplicated threshold events with isolated handlers and a sanitized error hook; durable delivery remains | | `token-pilot-autoconfigure` | Basic implementation complete | Bean registration, property binding, pricing/budget/notification wiring, and `ChatClientBuilderCustomizer` implemented | | `token-pilot-starter` | Basic implementation complete | Thin final user entrypoint that brings runtime modules together | | `token-pilot-sample-app` | Basic E2E complete | Direct ledger metrics, budget, and fake Spring AI advisor E2E implemented | @@ -216,6 +217,7 @@ Autoconfigure tests should use `ApplicationContextRunner` and verify: - Spring AI classpath registers `UsageExtractor`, `LedgerAdvisor`, and `ChatClientBuilderCustomizer` - notification beans do not register by default - notification beans register when notification is enabled and `BudgetNotificationHandler` bean exists +- the default budget store receives the notification service as an accounting listener ## Notification Contract @@ -233,9 +235,13 @@ class MailBudgetNotificationHandler implements BudgetNotificationHandler { - `BudgetNotificationHandler` 빈이 없으면 `BudgetNotificationService`는 등록되지 않는다 (no-op). - `token-pilot.notification.enabled=true` 설정 시에만 notification 빈이 등록된다. -- 알림 중복 방지는 evaluator가 확정한 `BudgetKey(policyId, targetType, targetId, window)`로 처리된다. -- 같은 window 안에서는 낮거나 같은 threshold 재발송이 방지된다. +- 알림은 legacy evaluator 호출이 아니라 적용된 commit/reconciliation-required/late reconciliation과 원자적 reservation BLOCK 결과를 소비한다. +- 알림 중복 방지는 `BudgetKey(policyId, targetType, targetId, window) + BudgetThreshold`로 처리된다. +- 같은 window 안에서는 낮거나 같은 threshold 재발송이 방지되고 duplicate accounting callback은 같은 reservation ID로 다시 누적되지 않는다. - 새 window에서는 50/80/100% 알림이 다시 가능하다. +- dedup 상태는 in-memory store 인스턴스 생명주기 동안 보존되며 TTL, 재시작 후 replay, exactly-once delivery를 제공하지 않는다. +- handler 실패는 다음 handler, 회계 결과, provider 응답을 바꾸지 않으며 bounded/sanitized `BudgetNotificationErrorHook`으로만 관찰한다. +- custom notification store는 atomic lifecycle을 위해 `AtomicNotificationStateStore`를 구현해야 하며, legacy `NotificationStateStore`만 등록하면 자동 설정이 명확히 실패한다. ## Recommended Configuration Shape @@ -402,6 +408,11 @@ Stage and deploy a Central release: ## Update History +### 2026-08-24 + +- Connected budget notifications to applied commit, reconciliation-required, late reconciliation, and atomic reservation BLOCK results instead of legacy evaluator calls. +- Added process-local atomic threshold deduplication, multi-handler failure isolation, sanitized error observation, and Spring Boot listener wiring. + ### 2026-08-22 - Restricted usage-based reservation reconciliation to provider-reported or provider-derived usage so local and heuristic estimates cannot be committed as actual spend. 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 e3d2735..a1c0910 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 @@ -3,6 +3,7 @@ import io.micrometer.core.instrument.MeterRegistry; import io.tokenpilot.budget.BudgetEvaluator; import io.tokenpilot.budget.BudgetStateStore; +import io.tokenpilot.budget.ReservationAccountingListener; import io.tokenpilot.budget.internal.LedgerBudgetComponents; import io.tokenpilot.core.CostCalculator; import io.tokenpilot.core.LedgerListener; @@ -13,6 +14,12 @@ import io.tokenpilot.core.domain.MissingPricingPolicy; import io.tokenpilot.core.internal.LedgerComponents; import io.tokenpilot.micrometer.internal.LedgerMicrometerComponents; +import io.tokenpilot.notification.AtomicNotificationStateStore; +import io.tokenpilot.notification.BudgetNotificationErrorHook; +import io.tokenpilot.notification.BudgetNotificationHandler; +import io.tokenpilot.notification.BudgetNotificationService; +import io.tokenpilot.notification.InMemoryNotificationStateStore; +import io.tokenpilot.notification.NotificationStateStore; import io.tokenpilot.springai.LedgerAdvisor; import io.tokenpilot.springai.UsageExtractor; import io.tokenpilot.springai.internal.LedgerSpringAiComponents; @@ -25,10 +32,6 @@ import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; import org.springframework.boot.context.properties.EnableConfigurationProperties; import org.springframework.context.annotation.Bean; -import io.tokenpilot.notification.BudgetNotificationHandler; -import io.tokenpilot.notification.BudgetNotificationService; -import io.tokenpilot.notification.InMemoryNotificationStateStore; -import io.tokenpilot.notification.NotificationStateStore; import java.time.Clock; @@ -176,8 +179,12 @@ public LedgerListener microCostMetricsPublisher(MeterRegistry meterRegistry, Tok @ConditionalOnMissingBean @ConditionalOnClass(LedgerBudgetComponents.class) @ConditionalOnProperty(prefix = "token-pilot.budget", name = "enabled", havingValue = "true") - public BudgetStateStore budgetStateStore() { - return LedgerBudgetComponents.inMemoryBudgetStateStore(); + public BudgetStateStore budgetStateStore( + ObjectProvider accountingListeners + ) { + return LedgerBudgetComponents.inMemoryBudgetStateStore( + accountingListeners.orderedStream().toList() + ); } /** @@ -205,9 +212,9 @@ public BudgetEvaluator budgetEvaluator( * - token-pilot.notification.enabled=true 일 때만 등록 */ @Bean - @ConditionalOnMissingBean + @ConditionalOnMissingBean(NotificationStateStore.class) @ConditionalOnProperty(prefix = "token-pilot.notification", name = "enabled", havingValue = "true") - public NotificationStateStore notificationStateStore() { + public AtomicNotificationStateStore notificationStateStore() { return new InMemoryNotificationStateStore(); } @@ -222,9 +229,21 @@ public NotificationStateStore notificationStateStore() { @ConditionalOnBean(BudgetNotificationHandler.class) @ConditionalOnProperty(prefix = "token-pilot.notification", name = "enabled", havingValue = "true") public BudgetNotificationService budgetNotificationService( - BudgetNotificationHandler handler, - NotificationStateStore notificationStateStore + ObjectProvider handlers, + AtomicNotificationStateStore notificationStateStore, + ObjectProvider budgetStateStore, + TokenPilotProperties properties, + ObjectProvider errorHook ) { - return new BudgetNotificationService(handler, notificationStateStore); + var policy = properties.toBudgetPolicy(); + return new BudgetNotificationService( + handlers.orderedStream().toList(), + notificationStateStore, + key -> budgetStateStore.getObject().snapshot( + key, + policy.monthlyLimit() + ), + errorHook.getIfAvailable(BudgetNotificationErrorHook::noOp) + ); } } 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 e1e89c0..13d3b55 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 @@ -5,10 +5,13 @@ import io.tokenpilot.budget.BudgetDecision; import io.tokenpilot.budget.BudgetEvaluator; import io.tokenpilot.budget.BudgetKey; +import io.tokenpilot.budget.BudgetReservationResult; import io.tokenpilot.budget.BudgetState; import io.tokenpilot.budget.BudgetStateStore; import io.tokenpilot.budget.BudgetThreshold; import io.tokenpilot.budget.BudgetWindow; +import io.tokenpilot.budget.ReservationAccountingListener; +import io.tokenpilot.budget.ReservationStatus; import io.tokenpilot.core.CostCalculator; import io.tokenpilot.core.LedgerManager; import io.tokenpilot.core.PricingEvaluator; @@ -22,8 +25,11 @@ import io.tokenpilot.core.domain.TokenType; import io.tokenpilot.core.domain.TokenUsage; import io.tokenpilot.core.exception.MissingPricingException; +import io.tokenpilot.notification.AtomicNotificationStateStore; import io.tokenpilot.notification.BudgetNotificationHandler; +import io.tokenpilot.notification.BudgetNotificationEvent; import io.tokenpilot.notification.BudgetNotificationService; +import io.tokenpilot.notification.BudgetNotificationSource; import io.tokenpilot.notification.NotificationStateStore; import io.tokenpilot.springai.LedgerAdvisor; import io.tokenpilot.springai.UsageExtractor; @@ -47,8 +53,10 @@ import java.time.Instant; import java.time.ZoneOffset; import java.util.Currency; +import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.concurrent.CopyOnWriteArrayList; import java.util.stream.Stream; import static io.tokenpilot.core.domain.TokenType.COMPLETION; @@ -428,12 +436,73 @@ void shouldRegisterNotificationServiceWhenEnabledAndHandlerExists() { .withPropertyValues("token-pilot.notification.enabled=true") .run(context -> { assertThat(context).hasSingleBean(NotificationStateStore.class); + assertThat(context).hasSingleBean(AtomicNotificationStateStore.class); assertThat(context).hasSingleBean(BudgetNotificationService.class); + assertThat(context).hasSingleBean(ReservationAccountingListener.class); + assertThat(context.getBean(ReservationAccountingListener.class)) + .isSameAs(context.getBean(BudgetNotificationService.class)); assertThat(context.getBean(TokenPilotProperties.class).getNotification().isEnabled()) .isTrue(); }); } + @Test + @DisplayName("legacy custom notification store는 atomic lifecycle을 조용히 무시하지 않아야 한다") + void shouldFailFastForLegacyCustomNotificationStore() { + this.contextRunner + .withUserConfiguration( + FakeBudgetNotificationHandlerConfiguration.class, + LegacyNotificationStateStoreConfiguration.class + ) + .withPropertyValues("token-pilot.notification.enabled=true") + .run(context -> { + assertThat(context).hasFailed(); + assertThat(context.getStartupFailure()) + .hasMessageContaining("AtomicNotificationStateStore"); + }); + } + + @Test + @DisplayName("원자적 budget BLOCK 결과가 자동 설정된 notification handler에 전달되어야 한다") + void shouldConnectAtomicBudgetBlockToNotificationHandler() { + this.contextRunner + .withUserConfiguration(FakeBudgetNotificationHandlerConfiguration.class) + .withPropertyValues( + "token-pilot.budget.enabled=true", + "token-pilot.notification.enabled=true" + ) + .run(context -> { + BudgetStateStore stateStore = context.getBean(BudgetStateStore.class); + TokenPilotProperties properties = context.getBean(TokenPilotProperties.class); + Cost limit = properties.toBudgetPolicy().monthlyLimit(); + BudgetKey key = new BudgetKey( + "budget-policy", + "tenant", + "tenant-a", + BudgetWindow.parse("2026-08") + ); + RecordingBudgetNotificationHandler handler = + (RecordingBudgetNotificationHandler) context.getBean( + BudgetNotificationHandler.class + ); + + BudgetReservationResult result = stateStore.checkAndReserve( + key, + limit, + limit, + "request-1" + ); + + assertThat(result.status()).isEqualTo(ReservationStatus.BLOCKED); + assertThat(handler.events()) + .extracting(BudgetNotificationEvent::threshold) + .containsExactly(BudgetThreshold.EXCEEDED); + assertThat(handler.events()) + .extracting(BudgetNotificationEvent::source) + .containsExactly(BudgetNotificationSource.RESERVATION_BLOCK); + }); + } + private static Stream providePricingConfigs() { return Stream.of( argumentSet( @@ -565,7 +634,46 @@ public MeterRegistry meterRegistry() { static class FakeBudgetNotificationHandlerConfiguration { @Bean public BudgetNotificationHandler budgetNotificationHandler() { - return event -> {}; + return new RecordingBudgetNotificationHandler(); + } + } + + @Configuration(proxyBeanMethods = false) + static class LegacyNotificationStateStoreConfiguration { + @Bean + public NotificationStateStore legacyNotificationStateStore() { + return new NotificationStateStore() { + private BudgetThreshold threshold = BudgetThreshold.NONE; + + @Override + public BudgetThreshold getLastNotifiedThreshold(BudgetKey key) { + return threshold; + } + + @Override + public void updateLastNotifiedThreshold( + BudgetKey key, + BudgetThreshold threshold + ) { + this.threshold = threshold; + } + }; + } + } + + static class RecordingBudgetNotificationHandler + implements BudgetNotificationHandler { + + private final List events = + new CopyOnWriteArrayList<>(); + + @Override + public void handle(BudgetNotificationEvent event) { + events.add(event); + } + + List events() { + return List.copyOf(events); } } } diff --git a/token-pilot-budget/src/main/java/io/tokenpilot/budget/ReservationAccountingListener.java b/token-pilot-budget/src/main/java/io/tokenpilot/budget/ReservationAccountingListener.java index 57d206b..d795e5b 100644 --- a/token-pilot-budget/src/main/java/io/tokenpilot/budget/ReservationAccountingListener.java +++ b/token-pilot-budget/src/main/java/io/tokenpilot/budget/ReservationAccountingListener.java @@ -1,8 +1,38 @@ package io.tokenpilot.budget; -/** 예약 정산 이벤트를 수신하는 framework-independent 계약입니다. */ +/** + * 예약 정산과 admission 차단 결과를 수신하는 framework-independent 계약입니다. + * + *

회계 변경은 새로 적용된 commit/reconciliation에만 전달됩니다. 예약 결과 callback은 + * {@link ReservationStatus#BLOCKED} 결과에만 사용되며 기존 단일 추상 메서드를 유지하므로 + * 기존 lambda listener와 source 호환됩니다.

+ */ @FunctionalInterface public interface ReservationAccountingListener { void onCommitted(ReservationAccountingEvent event); + + /** + * 적용된 accounting transition과 같은 linearization point의 bucket snapshot을 전달합니다. + * 기존 listener는 {@link #onCommitted(ReservationAccountingEvent)}로 위임됩니다. + */ + default void onAccountingApplied( + ReservationAccountingEvent event, + BudgetSnapshot snapshot + ) { + onCommitted(event); + } + + /** actual 미확정 estimate가 pending liability로 이동한 결과를 전달합니다. */ + default void onReconciliationRequired( + ReservationReconciliationRequiredEvent event + ) { + } + + /** 원자적 admission에서 상태 변경 없이 차단된 결과를 전달합니다. */ + default void onReservationBlocked( + BudgetReservationRequest request, + BudgetReservationResult result + ) { + } } diff --git a/token-pilot-budget/src/main/java/io/tokenpilot/budget/ReservationReconciliationRequiredEvent.java b/token-pilot-budget/src/main/java/io/tokenpilot/budget/ReservationReconciliationRequiredEvent.java new file mode 100644 index 0000000..fc22ee3 --- /dev/null +++ b/token-pilot-budget/src/main/java/io/tokenpilot/budget/ReservationReconciliationRequiredEvent.java @@ -0,0 +1,42 @@ +package io.tokenpilot.budget; + +import java.util.Objects; + +/** actual usage를 확보하지 못해 pending liability로 이동한 원자적 회계 결과입니다. */ +public record ReservationReconciliationRequiredEvent( + ReservationId reservationId, + BudgetKey budgetKey, + ReservationTransition transition, + ReservationAccountingReason reason, + BudgetSnapshot snapshot +) { + + public ReservationReconciliationRequiredEvent { + Objects.requireNonNull( + reservationId, + "reservationId must not be null" + ); + Objects.requireNonNull(budgetKey, "budgetKey must not be null"); + Objects.requireNonNull(transition, "transition must not be null"); + Objects.requireNonNull(reason, "reason must not be null"); + Objects.requireNonNull(snapshot, "snapshot must not be null"); + if (!transition.status().isApplied() + || transition.resultingState() + != ReservationState.RECONCILIATION_REQUIRED) { + throw new IllegalArgumentException( + "event requires an applied RECONCILIATION_REQUIRED transition" + ); + } + if (reason != ReservationAccountingReason.ACTUAL_USAGE_UNAVAILABLE + && reason != ReservationAccountingReason.CALLBACK_TIMED_OUT) { + throw new IllegalArgumentException( + "reason must require later actual reconciliation" + ); + } + if (!budgetKey.equals(snapshot.key())) { + throw new IllegalArgumentException( + "event and snapshot must use the same budget key" + ); + } + } +} 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 7d5e1a1..6b99e3c 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 @@ -16,8 +16,10 @@ import io.tokenpilot.budget.ReservationActualTokens; import io.tokenpilot.budget.ReservationId; import io.tokenpilot.budget.ReservationReconciliation; +import io.tokenpilot.budget.ReservationReconciliationRequiredEvent; import io.tokenpilot.budget.ReservationState; import io.tokenpilot.budget.ReservationStateMachine; +import io.tokenpilot.budget.ReservationStatus; import io.tokenpilot.budget.ReservationTransition; import io.tokenpilot.budget.ReservationTokenEstimate; import io.tokenpilot.core.CostCalculator; @@ -145,7 +147,12 @@ public BudgetReservationResult checkAndReserve(BudgetReservationRequest request) ) ); - return Objects.requireNonNull(result.get(), "reservation result must be set"); + BudgetReservationResult reservationResult = Objects.requireNonNull( + result.get(), + "reservation result must be set" + ); + publishBlockedReservation(request, reservationResult); + return reservationResult; } @Override @@ -260,13 +267,14 @@ public ReservationTransition markReconciliationRequired( ) { requireReconciliationRequiredReason(reason); Bucket bucket = bucketFor(reservationId); + ReservationTransition transition; + ReservationReconciliationRequiredEvent event; synchronized (bucket) { ReservationAccountingState accountingState = accountingState( bucket, reservationId ); - ReservationTransition transition = - accountingState.evaluateReconciliationRequired(); + transition = accountingState.evaluateReconciliationRequired(); if (!transition.status().isApplied()) { return transition; } @@ -276,8 +284,16 @@ public ReservationTransition markReconciliationRequired( accountingState, transition.resultingState() ); - return transition; + event = new ReservationReconciliationRequiredEvent( + reservationId, + accountingState.reservation().key(), + transition, + reason, + bucket.snapshot(accountingState.reservation().key()) + ); } + publishReconciliationRequired(event); + return transition; } @Override @@ -425,6 +441,7 @@ private ReservationReconciliation reconcileUsage( Objects.requireNonNull(reason, "reason must not be null"); Bucket bucket = bucketFor(command.reservationId()); ReservationReconciliation reconciliation; + BudgetSnapshot accountingSnapshot; synchronized (bucket) { reconciliation = reconcileUsageInBucket( bucket, @@ -432,13 +449,15 @@ private ReservationReconciliation reconcileUsage( type, reason ); + accountingSnapshot = bucket.snapshot(reconciliation.budgetKey()); } - publishAccountingEvent(reconciliation); + publishAccountingEvent(reconciliation, accountingSnapshot); return reconciliation; } private void publishAccountingEvent( - ReservationReconciliation reconciliation + ReservationReconciliation reconciliation, + BudgetSnapshot accountingSnapshot ) { if (!reconciliation.transition().status().isApplied() || accountingListeners.isEmpty()) { @@ -448,21 +467,51 @@ private void publishAccountingEvent( reconciliation ); for (ReservationAccountingListener listener : accountingListeners) { - notifyBestEffort(listener, event); + notifyBestEffort(listener, event, accountingSnapshot); } } private static void notifyBestEffort( ReservationAccountingListener listener, - ReservationAccountingEvent event + ReservationAccountingEvent event, + BudgetSnapshot accountingSnapshot ) { try { - listener.onCommitted(event); + listener.onAccountingApplied(event, accountingSnapshot); } catch (RuntimeException ignored) { // Listener 실패는 이미 적용된 회계 상태를 되돌리지 않습니다. } } + private void publishBlockedReservation( + BudgetReservationRequest request, + BudgetReservationResult result + ) { + if (result.status() != ReservationStatus.BLOCKED + || accountingListeners.isEmpty()) { + return; + } + for (ReservationAccountingListener listener : accountingListeners) { + try { + listener.onReservationBlocked(request, result); + } catch (RuntimeException ignored) { + // Listener 실패는 원자적으로 확정된 admission 결과를 바꾸지 않습니다. + } + } + } + + private void publishReconciliationRequired( + ReservationReconciliationRequiredEvent event + ) { + for (ReservationAccountingListener listener : accountingListeners) { + try { + listener.onReconciliationRequired(event); + } catch (RuntimeException ignored) { + // Listener 실패는 이미 적용된 pending liability를 되돌리지 않습니다. + } + } + } + private ReservationReconciliation reconcileUsageInBucket( Bucket bucket, ActualUsageCommand command, 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 d0380c1..ac64027 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 @@ -7,6 +7,7 @@ import io.tokenpilot.budget.ReservationAccountingListener; import io.tokenpilot.budget.ReservationId; import io.tokenpilot.core.CostCalculator; +import io.tokenpilot.core.internal.LedgerComponents; import java.time.Clock; import java.util.List; @@ -25,6 +26,17 @@ public static BudgetStateStore inMemoryBudgetStateStore() { return new InMemoryBudgetStateStore(); } + public static BudgetStateStore inMemoryBudgetStateStore( + List accountingListeners + ) { + return new InMemoryBudgetStateStore( + Clock.systemUTC(), + ReservationId::random, + LedgerComponents.defaultCostCalculator(), + accountingListeners + ); + } + public static BudgetStateStore inMemoryBudgetStateStore( Clock clock, Supplier reservationIdGenerator diff --git a/token-pilot-budget/src/test/java/io/tokenpilot/budget/internal/BudgetReservationStoreTest.java b/token-pilot-budget/src/test/java/io/tokenpilot/budget/internal/BudgetReservationStoreTest.java index 2e3cdb4..9592134 100644 --- a/token-pilot-budget/src/test/java/io/tokenpilot/budget/internal/BudgetReservationStoreTest.java +++ b/token-pilot-budget/src/test/java/io/tokenpilot/budget/internal/BudgetReservationStoreTest.java @@ -1,12 +1,15 @@ package io.tokenpilot.budget.internal; +import io.tokenpilot.budget.AccountingTransitionStatus; import io.tokenpilot.budget.BudgetKey; import io.tokenpilot.budget.BudgetReservationRequest; import io.tokenpilot.budget.BudgetReservationResult; import io.tokenpilot.budget.BudgetSnapshot; import io.tokenpilot.budget.BudgetWindow; import io.tokenpilot.budget.IdempotencyKey; +import io.tokenpilot.budget.ReservationAccountingListener; import io.tokenpilot.budget.ReservationId; +import io.tokenpilot.budget.ReservationReconciliationRequiredEvent; import io.tokenpilot.budget.ReservationStatus; import io.tokenpilot.budget.ReservationState; import io.tokenpilot.core.domain.Cost; @@ -19,6 +22,7 @@ import java.time.Instant; import java.time.ZoneOffset; import java.util.Currency; +import java.util.List; import java.util.Map; import java.util.Optional; import java.util.Set; @@ -27,6 +31,7 @@ import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; import static org.assertj.core.api.Assertions.assertThat; @@ -107,6 +112,135 @@ class BudgetReservationStoreTest { assertThat(store.snapshot(KEY, LIMIT).effectiveUsage()).isEqualTo(Cost.zero(USD)); } + @Test + void BLOCK_listener_실패는_결과를_바꾸지_않고_다음_listener를_막지_않는다() { + AtomicInteger failedDeliveries = new AtomicInteger(); + AtomicInteger successfulDeliveries = new AtomicInteger(); + AtomicInteger sequence = new AtomicInteger(); + ReservationAccountingListener failingListener = new ReservationAccountingListener() { + @Override + public void onCommitted(io.tokenpilot.budget.ReservationAccountingEvent event) { + } + + @Override + public void onReservationBlocked( + BudgetReservationRequest request, + BudgetReservationResult result + ) { + failedDeliveries.incrementAndGet(); + throw new IllegalStateException("listener failed"); + } + }; + ReservationAccountingListener successfulListener = + new ReservationAccountingListener() { + @Override + public void onCommitted( + io.tokenpilot.budget.ReservationAccountingEvent event + ) { + } + + @Override + public void onReservationBlocked( + BudgetReservationRequest request, + BudgetReservationResult result + ) { + successfulDeliveries.incrementAndGet(); + } + }; + InMemoryBudgetStateStore store = new InMemoryBudgetStateStore( + CLOCK, + () -> new ReservationId( + "reservation-" + sequence.incrementAndGet() + ), + (usage, pricing) -> Cost.zero(USD), + List.of(failingListener, successfulListener) + ); + + BudgetReservationResult result = store.checkAndReserve( + KEY, + LIMIT, + LIMIT, + "request-1" + ); + + assertThat(result.status()).isEqualTo(ReservationStatus.BLOCKED); + assertThat(failedDeliveries).hasValue(1); + assertThat(successfulDeliveries).hasValue(1); + assertThat(store.snapshot(KEY, LIMIT).effectiveUsage()) + .isEqualTo(Cost.zero(USD)); + } + + @Test + void 정산대기_listener_실패는_pending_상태를_바꾸지_않고_중복_발행하지_않는다() { + AtomicInteger failedDeliveries = new AtomicInteger(); + AtomicInteger successfulDeliveries = new AtomicInteger(); + AtomicInteger sequence = new AtomicInteger(); + AtomicReference delivered = + new AtomicReference<>(); + ReservationAccountingListener failingListener = + new ReservationAccountingListener() { + @Override + public void onCommitted( + io.tokenpilot.budget.ReservationAccountingEvent event + ) { + } + + @Override + public void onReconciliationRequired( + ReservationReconciliationRequiredEvent event + ) { + failedDeliveries.incrementAndGet(); + throw new IllegalStateException("listener failed"); + } + }; + ReservationAccountingListener successfulListener = + new ReservationAccountingListener() { + @Override + public void onCommitted( + io.tokenpilot.budget.ReservationAccountingEvent event + ) { + } + + @Override + public void onReconciliationRequired( + ReservationReconciliationRequiredEvent event + ) { + successfulDeliveries.incrementAndGet(); + delivered.set(event); + } + }; + InMemoryBudgetStateStore store = new InMemoryBudgetStateStore( + CLOCK, + () -> new ReservationId( + "reservation-" + sequence.incrementAndGet() + ), + (usage, pricing) -> Cost.zero(USD), + List.of(failingListener, successfulListener) + ); + BudgetReservationResult reserved = store.checkAndReserve( + KEY, + LIMIT, + usd("60.00"), + "request-1" + ); + store.markInFlight(reserved.reservationId()); + + var applied = store.markReconciliationRequired(reserved.reservationId()); + var duplicate = store.markReconciliationRequired(reserved.reservationId()); + + assertThat(applied.status()).isEqualTo(AccountingTransitionStatus.APPLIED); + assertThat(duplicate.status()).isEqualTo(AccountingTransitionStatus.REUSED); + assertThat(failedDeliveries).hasValue(1); + assertThat(successfulDeliveries).hasValue(1); + assertThat(delivered.get().snapshot().pendingReconciliationLiability()) + .isEqualTo(usd("60.00")); + BudgetSnapshot snapshot = store.snapshot(KEY, LIMIT); + assertThat(snapshot.activeReservedCost()).isEqualTo(Cost.zero(USD)); + assertThat(snapshot.pendingReconciliationLiability()) + .isEqualTo(usd("60.00")); + assertThat(snapshot.effectiveUsage()).isEqualTo(usd("60.00")); + } + @Test void 이미_예약된_금액까지_포함해_다음_예약을_BLOCKED한다() { InMemoryBudgetStateStore store = store(); diff --git a/token-pilot-notification/src/main/java/io/tokenpilot/notification/AtomicNotificationStateStore.java b/token-pilot-notification/src/main/java/io/tokenpilot/notification/AtomicNotificationStateStore.java new file mode 100644 index 0000000..35b791e --- /dev/null +++ b/token-pilot-notification/src/main/java/io/tokenpilot/notification/AtomicNotificationStateStore.java @@ -0,0 +1,26 @@ +package io.tokenpilot.notification; + +import io.tokenpilot.budget.BudgetReservationRequest; +import io.tokenpilot.budget.BudgetReservationResult; +import io.tokenpilot.budget.BudgetSnapshot; +import io.tokenpilot.budget.ReservationAccountingEvent; +import io.tokenpilot.budget.ReservationReconciliationRequiredEvent; + +/** atomic budget/accounting 결과의 threshold를 원자적으로 claim하는 저장소 계약입니다. */ +public interface AtomicNotificationStateStore extends NotificationStateStore { + + NotificationClaim recordAppliedTransition( + ReservationAccountingEvent event, + BudgetSnapshot snapshot, + BudgetNotificationSource source + ); + + NotificationClaim recordReconciliationRequired( + ReservationReconciliationRequiredEvent event + ); + + NotificationClaim recordBlockedReservation( + BudgetReservationRequest request, + BudgetReservationResult result + ); +} diff --git a/token-pilot-notification/src/main/java/io/tokenpilot/notification/BudgetNotificationError.java b/token-pilot-notification/src/main/java/io/tokenpilot/notification/BudgetNotificationError.java new file mode 100644 index 0000000..ddc6dd2 --- /dev/null +++ b/token-pilot-notification/src/main/java/io/tokenpilot/notification/BudgetNotificationError.java @@ -0,0 +1,66 @@ +package io.tokenpilot.notification; + +import java.util.Objects; + +/** + * 알림 실패 hook에 전달하는 제한된 오류 정보입니다. + * + *

원본 예외 메시지와 notification event를 보존하지 않아 prompt, provider response, + * API key가 오류 경로로 전달되지 않습니다. 각 문자열은 최대 128자로 제한됩니다.

+ */ +public record BudgetNotificationError( + String stage, + String handlerType, + String exceptionType, + String message +) { + + private static final int MAX_LENGTH = 128; + + public BudgetNotificationError { + stage = bounded(stage, "stage"); + handlerType = bounded(handlerType, "handlerType"); + exceptionType = bounded(exceptionType, "exceptionType"); + message = bounded(message, "message"); + } + + static BudgetNotificationError handlerFailure( + BudgetNotificationHandler handler, + RuntimeException failure + ) { + return new BudgetNotificationError( + "HANDLER", + typeName(handler), + typeName(failure), + "budget notification handler failed" + ); + } + + static BudgetNotificationError stateFailure(RuntimeException failure) { + return new BudgetNotificationError( + "STATE", + "notification-state-store", + typeName(failure), + "budget notification state update failed" + ); + } + + private static String typeName(Object value) { + Objects.requireNonNull(value, "value must not be null"); + String simpleName = value.getClass().getSimpleName(); + return simpleName.isBlank() ? "anonymous" : simpleName; + } + + private static String bounded(String value, String name) { + if (value == null || value.isBlank()) { + throw new IllegalArgumentException(name + " must not be blank"); + } + String sanitized = value.replaceAll("[\\p{Cntrl}]", " ").trim(); + if (sanitized.isBlank()) { + throw new IllegalArgumentException(name + " must contain visible text"); + } + return sanitized.length() <= MAX_LENGTH + ? sanitized + : sanitized.substring(0, MAX_LENGTH); + } +} diff --git a/token-pilot-notification/src/main/java/io/tokenpilot/notification/BudgetNotificationErrorHook.java b/token-pilot-notification/src/main/java/io/tokenpilot/notification/BudgetNotificationErrorHook.java new file mode 100644 index 0000000..6a17978 --- /dev/null +++ b/token-pilot-notification/src/main/java/io/tokenpilot/notification/BudgetNotificationErrorHook.java @@ -0,0 +1,12 @@ +package io.tokenpilot.notification; + +/** bounded/sanitized notification 오류를 관찰하는 best-effort hook입니다. */ +@FunctionalInterface +public interface BudgetNotificationErrorHook { + + void onError(BudgetNotificationError error); + + static BudgetNotificationErrorHook noOp() { + return ignored -> { }; + } +} 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 7a791aa..2e927f2 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 @@ -6,6 +6,7 @@ import io.tokenpilot.core.domain.Cost; import java.util.Map; +import java.util.Objects; /** * 예산 임계치 도달 시 발생하는 알림 이벤트입니다. @@ -14,31 +15,88 @@ * @param threshold 도달한 임계치 * @param state 예산 상태 * @param reason 상태 설명 - * @param projectedUsage 후보 비용을 포함한 예상 사용량 + * @param usage 알림을 만든 accounting 또는 admission 결과의 사용량 * @param limit 예산 한도 - * @param tags 알림에 전달할 태그 + * @param source 알림을 만든 원자적 결과 종류 * - *

Migration note: 후보 비용 포함 사용량의 의미를 명확히 하기 위해 - * record component 이름을 {@code currentUsage}에서 {@code projectedUsage}로 변경했습니다. - * 기존 handler를 위한 {@link #currentUsage()} 호환 accessor는 0.1.x 동안 유지하며 - * 0.2.0에서 제거할 예정입니다. + *

이 이벤트는 prompt, raw provider response, API key, 임의 tag map을 포함하지 않습니다. + * 기존 handler를 위한 {@link #projectedUsage()}, {@link #currentUsage()}, {@link #tags()} + * 호환 accessor는 0.1.x 동안 유지하며 0.2.0에서 제거할 예정입니다.

*/ public record BudgetNotificationEvent( BudgetKey key, BudgetThreshold threshold, BudgetState state, String reason, - Cost projectedUsage, + Cost usage, Cost limit, - Map tags + BudgetNotificationSource source ) { + public BudgetNotificationEvent { + Objects.requireNonNull(key, "key must not be null"); + Objects.requireNonNull(threshold, "threshold must not be null"); + Objects.requireNonNull(state, "state must not be null"); + if (reason == null || reason.isBlank()) { + throw new IllegalArgumentException("reason must not be blank"); + } + Objects.requireNonNull(usage, "usage must not be null"); + Objects.requireNonNull(limit, "limit must not be null"); + Objects.requireNonNull(source, "source must not be null"); + if (!usage.currency().equals(limit.currency())) { + throw new IllegalArgumentException("usage and limit must use the same currency"); + } + } + + /** + * @deprecated legacy decision 기반 이벤트 생성 호환용입니다. tags는 보존하지 않습니다. + */ + @Deprecated(since = "0.1.0", forRemoval = false) + public BudgetNotificationEvent( + BudgetKey key, + BudgetThreshold threshold, + BudgetState state, + String reason, + Cost projectedUsage, + Cost limit, + Map tags + ) { + this( + key, + threshold, + state, + reason, + projectedUsage, + limit, + BudgetNotificationSource.LEGACY_DECISION + ); + Objects.requireNonNull(tags, "tags must not be null"); + } + + /** + * @return {@link #usage()}와 동일한 accounting/admission 사용량 + * @deprecated source별 의미가 명확한 {@link #usage()}를 사용하세요. + */ + @Deprecated(since = "0.1.0", forRemoval = true) + public Cost projectedUsage() { + return usage; + } + /** - * @return {@link #projectedUsage()}와 동일한 후보 비용 포함 예상 사용량 - * @deprecated 후보 비용 포함 사용량은 {@link #projectedUsage()}를 사용하세요. + * @return {@link #usage()}와 동일한 accounting/admission 사용량 + * @deprecated source별 의미가 명확한 {@link #usage()}를 사용하세요. */ @Deprecated(since = "0.1.0", forRemoval = true) public Cost currentUsage() { - return projectedUsage; + return usage; + } + + /** + * @return 민감 정보가 포함되지 않는 빈 map + * @deprecated notification event는 임의 tags를 전달하지 않습니다. + */ + @Deprecated(since = "0.1.0", forRemoval = true) + public Map tags() { + return Map.of(); } } 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 43c3ffd..6086513 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 @@ -1,62 +1,281 @@ package io.tokenpilot.notification; import io.tokenpilot.budget.BudgetDecision; +import io.tokenpilot.budget.BudgetKey; +import io.tokenpilot.budget.BudgetReservationRequest; +import io.tokenpilot.budget.BudgetReservationResult; +import io.tokenpilot.budget.BudgetSnapshot; +import io.tokenpilot.budget.BudgetState; import io.tokenpilot.budget.BudgetThreshold; +import io.tokenpilot.budget.ReservationAccountingEvent; +import io.tokenpilot.budget.ReservationAccountingListener; +import io.tokenpilot.budget.ReservationAccountingReason; +import io.tokenpilot.budget.ReservationReconciliationRequiredEvent; +import io.tokenpilot.core.domain.Cost; +import java.util.List; import java.util.Map; +import java.util.Objects; +import java.util.function.Function; /** - * 예산 판단 결과를 기반으로 알림 이벤트를 발행하는 서비스 + * 원자적 budget/accounting 결과를 threshold 알림으로 변환하는 best-effort listener입니다. + * + *

threshold claim은 handler 호출 전에 완료됩니다. handler 실패는 다음 handler 전달을 + * 막지 않고 이미 적용된 회계 또는 admission 결과를 변경하지 않으며, 실패한 전달은 durable + * outbox가 없는 MVP에서 재시도되지 않습니다.

*/ -public class BudgetNotificationService { +public class BudgetNotificationService implements ReservationAccountingListener { - private final BudgetNotificationHandler handler; + private final List handlers; private final NotificationStateStore store; + private final AtomicNotificationStateStore atomicStore; + private final BudgetNotificationErrorHook errorHook; + private final Function snapshotResolver; + /** + * @deprecated legacy {@link BudgetDecision} 알림 호환용 생성자입니다. 신규 연결은 복수 + * handler와 error hook을 받는 생성자를 사용하세요. + */ + @Deprecated(since = "0.1.0", forRemoval = false) public BudgetNotificationService( BudgetNotificationHandler handler, NotificationStateStore store ) { - this.handler = handler; - this.store = store; + this.handlers = List.of( + Objects.requireNonNull(handler, "handler must not be null") + ); + this.store = Objects.requireNonNull(store, "store must not be null"); + this.atomicStore = null; + this.errorHook = BudgetNotificationErrorHook.noOp(); + this.snapshotResolver = null; + } + + public BudgetNotificationService( + List handlers, + AtomicNotificationStateStore store, + Function snapshotResolver + ) { + this( + handlers, + store, + snapshotResolver, + BudgetNotificationErrorHook.noOp() + ); + } + + public BudgetNotificationService( + List handlers, + AtomicNotificationStateStore store, + Function snapshotResolver, + BudgetNotificationErrorHook errorHook + ) { + this.handlers = List.copyOf( + Objects.requireNonNull(handlers, "handlers must not be null") + ); + if (this.handlers.isEmpty()) { + throw new IllegalArgumentException("handlers must not be empty"); + } + this.store = Objects.requireNonNull(store, "store must not be null"); + this.atomicStore = store; + this.snapshotResolver = Objects.requireNonNull( + snapshotResolver, + "snapshotResolver must not be null" + ); + this.errorHook = Objects.requireNonNull(errorHook, "errorHook must not be null"); } /** - * 임계치가 증가한 경우에만 이벤트를 발생시킨다 + * 기존 #37 callback은 주입된 resolver로 현재 bucket snapshot을 조회합니다. 현재 budget store는 + * 정확한 전이 시점 snapshot을 포함하는 + * {@link #onAccountingApplied(ReservationAccountingEvent, BudgetSnapshot)}를 호출합니다. */ + @Override + public void onCommitted(ReservationAccountingEvent event) { + Objects.requireNonNull(event, "event must not be null"); + requireAtomicLifecycle(); + BudgetSnapshot snapshot; + try { + snapshot = Objects.requireNonNull( + snapshotResolver.apply(event.reconciliation().budgetKey()), + "snapshotResolver returned null" + ); + } catch (RuntimeException failure) { + report(BudgetNotificationError.stateFailure(failure)); + return; + } + processAppliedAccounting(event, snapshot); + } + + @Override + public void onAccountingApplied( + ReservationAccountingEvent event, + BudgetSnapshot snapshot + ) { + Objects.requireNonNull(event, "event must not be null"); + Objects.requireNonNull(snapshot, "snapshot must not be null"); + requireAtomicLifecycle(); + processAppliedAccounting(event, snapshot); + } + + private void processAppliedAccounting( + ReservationAccountingEvent event, + BudgetSnapshot snapshot + ) { + ReservationAccountingReason reason = event.reconciliation().reason(); + BudgetNotificationSource source = reason + == ReservationAccountingReason.LATE_ACTUAL_USAGE_REPORTED + ? BudgetNotificationSource.LATE_RECONCILIATION + : BudgetNotificationSource.ACCOUNTING_COMMIT; + + NotificationStateStore.NotificationClaim claim; + try { + claim = atomicStore.recordAppliedTransition(event, snapshot, source); + } catch (RuntimeException failure) { + report(BudgetNotificationError.stateFailure(failure)); + return; + } + + publish( + event.reconciliation().budgetKey(), + claim, + reason.name(), + source, + snapshot.limit() + ); + } + + @Override + public void onReconciliationRequired( + ReservationReconciliationRequiredEvent event + ) { + Objects.requireNonNull(event, "event must not be null"); + requireAtomicLifecycle(); + + NotificationStateStore.NotificationClaim claim; + try { + claim = atomicStore.recordReconciliationRequired(event); + } catch (RuntimeException failure) { + report(BudgetNotificationError.stateFailure(failure)); + return; + } + publish( + event.budgetKey(), + claim, + event.reason().name(), + BudgetNotificationSource.RECONCILIATION_REQUIRED, + event.snapshot().limit() + ); + } + + @Override + public void onReservationBlocked( + BudgetReservationRequest request, + BudgetReservationResult result + ) { + Objects.requireNonNull(request, "request must not be null"); + Objects.requireNonNull(result, "result must not be null"); + requireAtomicLifecycle(); + + NotificationStateStore.NotificationClaim claim; + try { + claim = atomicStore.recordBlockedReservation(request, result); + } catch (RuntimeException failure) { + report(BudgetNotificationError.stateFailure(failure)); + return; + } + publish( + request.key(), + claim, + result.reason(), + BudgetNotificationSource.RESERVATION_BLOCK, + request.limit() + ); + } + + /** + * @deprecated legacy evaluator 결과는 admission/accounting 알림 근거가 아닙니다. 신규 코드는 + * {@link ReservationAccountingListener} 연결을 사용하세요. tags는 전달하지 않습니다. + */ + @Deprecated(since = "0.1.0", forRemoval = false) public void notifyIfNeeded( BudgetDecision decision, Map tags ) { + Objects.requireNonNull(decision, "decision must not be null"); + Objects.requireNonNull(tags, "tags must not be null"); BudgetThreshold current = decision.threshold(); - if (current == BudgetThreshold.NONE) { return; } BudgetThreshold last = store.getLastNotifiedThreshold(decision.key()); - - // 같은 window에서 중복 방지 if (current.compareTo(last) <= 0) { return; } + store.updateLastNotifiedThreshold(decision.key(), current); + dispatch(new BudgetNotificationEvent( + decision.key(), + current, + decision.state(), + decision.reason(), + decision.projectedUsage(), + decision.limit(), + BudgetNotificationSource.LEGACY_DECISION + )); + } - BudgetNotificationEvent event = - new BudgetNotificationEvent( - decision.key(), - current, - decision.state(), - decision.reason(), - decision.projectedUsage(), - decision.limit(), - tags - ); + private void publish( + BudgetKey key, + NotificationStateStore.NotificationClaim claim, + String reason, + BudgetNotificationSource source, + Cost limit + ) { + for (BudgetThreshold threshold : claim.thresholds()) { + dispatch(new BudgetNotificationEvent( + key, + threshold, + state(threshold), + reason, + claim.usage(), + limit, + source + )); + } + } - handler.handle(event); + private void dispatch(BudgetNotificationEvent event) { + for (BudgetNotificationHandler handler : handlers) { + try { + handler.handle(event); + } catch (RuntimeException failure) { + report(BudgetNotificationError.handlerFailure(handler, failure)); + } + } + } - store.updateLastNotifiedThreshold( - decision.key(), - current - ); + private void report(BudgetNotificationError error) { + try { + errorHook.onError(error); + } catch (RuntimeException ignored) { + // Error hook도 best-effort이며 accounting/provider 결과에 영향을 주지 않습니다. + } + } + + private void requireAtomicLifecycle() { + if (atomicStore == null || snapshotResolver == null) { + throw new IllegalStateException( + "atomic accounting callbacks require AtomicNotificationStateStore" + ); + } + } + + private static BudgetState state(BudgetThreshold threshold) { + return switch (threshold) { + case HALF, WARNING -> BudgetState.WARN; + case EXCEEDED -> BudgetState.BLOCK; + case NONE -> throw new IllegalArgumentException("NONE must not be published"); + }; } } diff --git a/token-pilot-notification/src/main/java/io/tokenpilot/notification/BudgetNotificationSource.java b/token-pilot-notification/src/main/java/io/tokenpilot/notification/BudgetNotificationSource.java new file mode 100644 index 0000000..16f70a0 --- /dev/null +++ b/token-pilot-notification/src/main/java/io/tokenpilot/notification/BudgetNotificationSource.java @@ -0,0 +1,10 @@ +package io.tokenpilot.notification; + +/** 알림을 만든 원자적 budget/accounting 결과의 종류입니다. */ +public enum BudgetNotificationSource { + ACCOUNTING_COMMIT, + RECONCILIATION_REQUIRED, + LATE_RECONCILIATION, + RESERVATION_BLOCK, + LEGACY_DECISION +} 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 3616663..4f6410d 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,21 +1,161 @@ package io.tokenpilot.notification; import io.tokenpilot.budget.BudgetKey; +import io.tokenpilot.budget.BudgetReservationRequest; +import io.tokenpilot.budget.BudgetReservationResult; +import io.tokenpilot.budget.BudgetSnapshot; import io.tokenpilot.budget.BudgetThreshold; +import io.tokenpilot.budget.ReservationAccountingEvent; +import io.tokenpilot.budget.ReservationId; +import io.tokenpilot.budget.ReservationReconciliationRequiredEvent; +import io.tokenpilot.budget.ReservationStatus; +import io.tokenpilot.core.domain.Cost; +import java.math.BigDecimal; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; import java.util.Map; +import java.util.Objects; +import java.util.Set; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicReference; /** - * 메모리 기반 상태 저장소 + * 메모리 기반 accounting threshold dedup 저장소입니다. + * + *

dedup key는 {@code BudgetKey(policy/target/window) + BudgetThreshold}입니다. + * 처리한 accounting transition은 같은 bucket의 {@link ReservationId}로 중복 claim을 막습니다. + * 상태는 이 store 인스턴스의 생명주기 동안 보존되며 재시작 후 replay나 TTL 정리는 제공하지 + * 않습니다. 따라서 MVP 전달 보장은 process-local best-effort이며 durable exactly-once가 아닙니다.

*/ -public class InMemoryNotificationStateStore implements NotificationStateStore { +@SuppressWarnings("deprecation") +public class InMemoryNotificationStateStore + implements AtomicNotificationStateStore { - private final Map store = new ConcurrentHashMap<>(); + private static final BigDecimal HALF_RATIO = new BigDecimal("0.5"); + private static final BigDecimal WARNING_RATIO = new BigDecimal("0.8"); + + private final Map store = new ConcurrentHashMap<>(); + + @Override + public NotificationClaim recordAppliedTransition( + ReservationAccountingEvent event, + BudgetSnapshot snapshot, + BudgetNotificationSource source + ) { + Objects.requireNonNull(event, "event must not be null"); + Objects.requireNonNull(snapshot, "snapshot must not be null"); + var reconciliation = event.reconciliation(); + if (!reconciliation.transition().status().isApplied()) { + return NotificationClaim.none(snapshot.effectiveUsage()); + } + if (!reconciliation.budgetKey().equals(snapshot.key())) { + throw new IllegalArgumentException( + "accounting event and snapshot must use the same budget key" + ); + } + requireCurrency(snapshot.limit(), reconciliation.actual()); + requireAccountingSource(source); + + return recordTransition( + reconciliation.reservationId(), + reconciliation.budgetKey(), + snapshot, + source + ); + } + + @Override + public NotificationClaim recordReconciliationRequired( + ReservationReconciliationRequiredEvent event + ) { + Objects.requireNonNull(event, "event must not be null"); + return recordTransition( + event.reservationId(), + event.budgetKey(), + event.snapshot(), + BudgetNotificationSource.RECONCILIATION_REQUIRED + ); + } + + private NotificationClaim recordTransition( + ReservationId reservationId, + BudgetKey key, + BudgetSnapshot snapshot, + BudgetNotificationSource source + ) { + TransitionDedupKey dedupKey = new TransitionDedupKey( + reservationId, + source + ); + + AtomicReference claim = new AtomicReference<>(); + store.compute(key, (ignored, existing) -> { + BucketState state = state(existing, snapshot.limit()); + if (!state.processedTransitions.add(dedupKey)) { + claim.set(NotificationClaim.none(snapshot.effectiveUsage())); + return state; + } + + List thresholds = newlyReached( + state.lastNotifiedThreshold, + snapshot.effectiveUsage(), + state.limit + ); + if (!thresholds.isEmpty()) { + state.lastNotifiedThreshold = thresholds.get(thresholds.size() - 1); + } + claim.set(new NotificationClaim(snapshot.effectiveUsage(), thresholds)); + return state; + }); + return Objects.requireNonNull(claim.get(), "notification claim must be set"); + } + + @Override + public NotificationClaim recordBlockedReservation( + BudgetReservationRequest request, + BudgetReservationResult result + ) { + Objects.requireNonNull(request, "request must not be null"); + Objects.requireNonNull(result, "result must not be null"); + if (result.status() != ReservationStatus.BLOCKED) { + return NotificationClaim.none(result.snapshot().effectiveUsage()); + } + if (!request.key().equals(result.snapshot().key()) + || !request.limit().equals(result.snapshot().limit())) { + throw new IllegalArgumentException("blocked result must match its reservation request"); + } + + Cost projectedUsage = result.snapshot() + .effectiveUsage() + .add(request.safeUpperBoundCost()); + if (projectedUsage.compareTo(request.limit()) < 0) { + throw new IllegalArgumentException("blocked usage must reach the budget limit"); + } + + AtomicReference claim = new AtomicReference<>(); + store.compute(request.key(), (ignored, existing) -> { + BucketState state = state(existing, request.limit()); + if (state.lastNotifiedThreshold.compareTo(BudgetThreshold.EXCEEDED) >= 0) { + claim.set(NotificationClaim.none(projectedUsage)); + return state; + } + state.lastNotifiedThreshold = BudgetThreshold.EXCEEDED; + claim.set(new NotificationClaim( + projectedUsage, + List.of(BudgetThreshold.EXCEEDED) + )); + return state; + }); + return Objects.requireNonNull(claim.get(), "notification claim must be set"); + } @Override public BudgetThreshold getLastNotifiedThreshold(BudgetKey key) { - return store.getOrDefault(key, BudgetThreshold.NONE); + Objects.requireNonNull(key, "key must not be null"); + BucketState state = store.get(key); + return state == null ? BudgetThreshold.NONE : state.lastNotifiedThreshold; } @Override @@ -23,6 +163,131 @@ public void updateLastNotifiedThreshold( BudgetKey key, BudgetThreshold threshold ) { - store.put(key, threshold); + Objects.requireNonNull(key, "key must not be null"); + Objects.requireNonNull(threshold, "threshold must not be null"); + store.compute(key, (ignored, existing) -> { + BucketState state = existing == null ? BucketState.legacy() : existing; + if (threshold.compareTo(state.lastNotifiedThreshold) > 0) { + state.lastNotifiedThreshold = threshold; + } + return state; + }); + } + + private static BucketState state(BucketState existing, Cost limit) { + if (existing == null) { + return BucketState.accounting(limit); + } + if (existing.limit == null) { + existing.limit = limit; + return existing; + } + if (!existing.limit.equals(limit)) { + throw new IllegalArgumentException( + "notification budget limit snapshot changed for an existing key" + ); + } + return existing; + } + + private static List newlyReached( + BudgetThreshold last, + Cost usage, + Cost limit + ) { + List thresholds = new ArrayList<>(3); + addIfReached(thresholds, last, BudgetThreshold.HALF, usage, limit, HALF_RATIO); + addIfReached( + thresholds, + last, + BudgetThreshold.WARNING, + usage, + limit, + WARNING_RATIO + ); + addIfReached( + thresholds, + last, + BudgetThreshold.EXCEEDED, + usage, + limit, + BigDecimal.ONE + ); + return List.copyOf(thresholds); + } + + private static void addIfReached( + List thresholds, + BudgetThreshold last, + BudgetThreshold candidate, + Cost usage, + Cost limit, + BigDecimal ratio + ) { + if (candidate.compareTo(last) <= 0) { + return; + } + Cost boundary = Cost.of( + limit.value().multiply(ratio), + limit.currency() + ); + if (usage.compareTo(boundary) >= 0) { + thresholds.add(candidate); + } + } + + private static void requireCurrency(Cost limit, Cost amount) { + if (!limit.currency().equals(amount.currency())) { + throw new IllegalArgumentException( + "notification usage and limit must use the same currency" + ); + } + } + + private static void requireAccountingSource(BudgetNotificationSource source) { + Objects.requireNonNull(source, "source must not be null"); + if (source != BudgetNotificationSource.ACCOUNTING_COMMIT + && source != BudgetNotificationSource.LATE_RECONCILIATION) { + throw new IllegalArgumentException( + "source must identify an applied accounting result" + ); + } + } + + private static final class BucketState { + private volatile Cost limit; + private volatile BudgetThreshold lastNotifiedThreshold; + private final Set processedTransitions; + + private BucketState( + Cost limit, + BudgetThreshold lastNotifiedThreshold + ) { + this.limit = limit; + this.lastNotifiedThreshold = lastNotifiedThreshold; + this.processedTransitions = new HashSet<>(); + } + + private static BucketState accounting(Cost limit) { + return new BucketState( + limit, + BudgetThreshold.NONE + ); + } + + private static BucketState legacy() { + return new BucketState(null, BudgetThreshold.NONE); + } + } + + private record TransitionDedupKey( + ReservationId reservationId, + BudgetNotificationSource source + ) { + + private TransitionDedupKey { + Objects.requireNonNull(reservationId, "reservationId must not be null"); + Objects.requireNonNull(source, "source must not be null"); + } } } 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 891cef3..11202a7 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 @@ -2,16 +2,46 @@ import io.tokenpilot.budget.BudgetKey; import io.tokenpilot.budget.BudgetThreshold; +import io.tokenpilot.core.domain.Cost; + +import java.util.List; +import java.util.Objects; /** - * 알림 중복 방지를 위한 상태 저장소 + * legacy decision 기반 알림 dedup 저장소입니다. + * Atomic lifecycle 연결은 {@link AtomicNotificationStateStore}를 구현해야 합니다. */ public interface NotificationStateStore { + /** @deprecated legacy decision 기반 조회 API입니다. */ + @Deprecated(since = "0.1.0", forRemoval = false) BudgetThreshold getLastNotifiedThreshold(BudgetKey key); + /** @deprecated legacy decision 기반 갱신 API입니다. */ + @Deprecated(since = "0.1.0", forRemoval = false) void updateLastNotifiedThreshold( BudgetKey key, BudgetThreshold threshold ); + + /** 원자적으로 claim된 threshold와 해당 판단의 사용량입니다. */ + record NotificationClaim( + Cost usage, + List thresholds + ) { + + public NotificationClaim { + Objects.requireNonNull(usage, "usage must not be null"); + thresholds = List.copyOf( + Objects.requireNonNull(thresholds, "thresholds must not be null") + ); + if (thresholds.contains(BudgetThreshold.NONE)) { + throw new IllegalArgumentException("NONE must not be claimed"); + } + } + + public static NotificationClaim none(Cost usage) { + return new NotificationClaim(usage, List.of()); + } + } } diff --git a/token-pilot-notification/src/test/java/io/tokenpilot/notification/AccountingBudgetNotificationTest.java b/token-pilot-notification/src/test/java/io/tokenpilot/notification/AccountingBudgetNotificationTest.java new file mode 100644 index 0000000..0d3a8d1 --- /dev/null +++ b/token-pilot-notification/src/test/java/io/tokenpilot/notification/AccountingBudgetNotificationTest.java @@ -0,0 +1,500 @@ +package io.tokenpilot.notification; + +import io.tokenpilot.budget.AccountingTransitionStatus; +import io.tokenpilot.budget.ActualUsageCommand; +import io.tokenpilot.budget.BudgetKey; +import io.tokenpilot.budget.BudgetReservationRequest; +import io.tokenpilot.budget.BudgetReservationResult; +import io.tokenpilot.budget.BudgetSnapshot; +import io.tokenpilot.budget.BudgetStateStore; +import io.tokenpilot.budget.BudgetThreshold; +import io.tokenpilot.budget.BudgetWindow; +import io.tokenpilot.budget.IdempotencyKey; +import io.tokenpilot.budget.ReservationAccounting; +import io.tokenpilot.budget.ReservationAccountingEvent; +import io.tokenpilot.budget.ReservationId; +import io.tokenpilot.budget.ReservationReconciliation; +import io.tokenpilot.budget.ReservationStatus; +import io.tokenpilot.budget.ReservationTokenEstimate; +import io.tokenpilot.budget.internal.LedgerBudgetComponents; +import io.tokenpilot.core.domain.Cost; +import io.tokenpilot.core.domain.PricingSnapshot; +import io.tokenpilot.core.domain.TokenType; +import io.tokenpilot.core.domain.TokenUsage; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; + +import java.math.BigDecimal; +import java.time.Clock; +import java.time.Instant; +import java.time.ZoneOffset; +import java.util.ArrayList; +import java.util.Currency; +import java.util.List; +import java.util.Map; +import java.util.concurrent.Callable; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; + +import static org.assertj.core.api.Assertions.assertThat; + +class AccountingBudgetNotificationTest { + + private static final Currency USD = Currency.getInstance("USD"); + private static final Cost LIMIT = usd("100.00"); + private static final Clock CLOCK = Clock.fixed( + Instant.parse("2026-08-24T00:00:00Z"), + ZoneOffset.UTC + ); + private static final ReservationTokenEstimate TOKEN_ESTIMATE = + new ReservationTokenEstimate(1, 1, 1); + + @Test + void NONE은_알리지_않는다() { + List events = new ArrayList<>(); + Fixture fixture = fixture(List.of(events::add), BudgetNotificationErrorHook.noOp()); + + commit(fixture, key("tenant-a", "2026-08"), "request-1", 40); + + assertThat(events).isEmpty(); + } + + @Test + void 같은_key와_threshold는_한_번만_전달한다() { + List events = new ArrayList<>(); + Fixture fixture = fixture(List.of(events::add), BudgetNotificationErrorHook.noOp()); + BudgetKey key = key("tenant-a", "2026-08"); + + commit(fixture, key, "request-1", 50); + commit(fixture, key, "request-2", 10); + commit(fixture, key, "request-3", 10); + + assertThat(events) + .extracting(BudgetNotificationEvent::threshold) + .containsExactly(BudgetThreshold.HALF); + } + + @Test + void HALF_WARNING_EXCEEDED는_상승할_때_각각_한_번_전달한다() { + List events = new ArrayList<>(); + Fixture fixture = fixture(List.of(events::add), BudgetNotificationErrorHook.noOp()); + BudgetKey key = key("tenant-a", "2026-08"); + + commit(fixture, key, "request-1", 50); + commit(fixture, key, "request-2", 30); + commit(fixture, key, "request-3", 20); + + assertThat(events) + .extracting(BudgetNotificationEvent::threshold) + .containsExactly( + BudgetThreshold.HALF, + BudgetThreshold.WARNING, + BudgetThreshold.EXCEEDED + ); + assertThat(events) + .extracting(BudgetNotificationEvent::usage) + .containsExactly(usd("50"), usd("80"), usd("100")); + } + + @Test + void commit은_notification_자체_누적이_아닌_atomic_snapshot으로_판정한다() { + List events = new ArrayList<>(); + Fixture fixture = fixture(List.of(events::add), BudgetNotificationErrorHook.noOp()); + BudgetKey key = key("tenant-a", "2026-08"); + BudgetReservationResult first = fixture.stateStore().checkAndReserve( + request(key, "request-1", usd("40")) + ); + BudgetReservationResult second = fixture.stateStore().checkAndReserve( + request(key, "request-2", usd("40")) + ); + fixture.accounting().markInFlight(first.reservationId()); + + fixture.accounting().commit(command("request-1", first.reservationId(), 20)); + + assertThat(first.status()).isEqualTo(ReservationStatus.CREATED); + assertThat(second.status()).isEqualTo(ReservationStatus.CREATED); + assertThat(events) + .extracting(BudgetNotificationEvent::threshold) + .containsExactly(BudgetThreshold.HALF); + assertThat(events) + .extracting(BudgetNotificationEvent::usage) + .containsExactly(usd("60")); + } + + @Test + void actual_unavailable은_pending_snapshot에서_threshold를_판정한다() { + List events = new ArrayList<>(); + Fixture fixture = fixture(List.of(events::add), BudgetNotificationErrorHook.noOp()); + BudgetKey key = key("tenant-a", "2026-08"); + BudgetReservationResult reserved = fixture.stateStore().checkAndReserve( + request(key, "request-1", usd("60")) + ); + fixture.accounting().markInFlight(reserved.reservationId()); + + var applied = fixture.accounting().markReconciliationRequired( + reserved.reservationId() + ); + var duplicate = fixture.accounting().markReconciliationRequired( + reserved.reservationId() + ); + + assertThat(applied.status()).isEqualTo(AccountingTransitionStatus.APPLIED); + assertThat(duplicate.status()).isEqualTo(AccountingTransitionStatus.REUSED); + assertThat(fixture.stateStore().snapshot(key, LIMIT) + .pendingReconciliationLiability()).isEqualTo(usd("60")); + assertThat(events) + .extracting(BudgetNotificationEvent::threshold) + .containsExactly(BudgetThreshold.HALF); + assertThat(events) + .extracting(BudgetNotificationEvent::source) + .containsExactly(BudgetNotificationSource.RECONCILIATION_REQUIRED); + } + + @Test + void 새_window에서는_같은_threshold를_다시_전달한다() { + List events = new ArrayList<>(); + Fixture fixture = fixture(List.of(events::add), BudgetNotificationErrorHook.noOp()); + + commit(fixture, key("tenant-a", "2026-08"), "request-1", 50); + commit(fixture, key("tenant-a", "2026-09"), "request-2", 50); + + assertThat(events) + .extracting(event -> event.key().window()) + .containsExactly(BudgetWindow.parse("2026-08"), BudgetWindow.parse("2026-09")); + assertThat(events) + .extracting(BudgetNotificationEvent::threshold) + .containsExactly(BudgetThreshold.HALF, BudgetThreshold.HALF); + } + + @Test + void duplicate_commit과_late_reconcile은_알림을_다시_만들지_않는다() { + List events = new ArrayList<>(); + Fixture fixture = fixture(List.of(events::add), BudgetNotificationErrorHook.noOp()); + + BudgetKey directKey = key("tenant-a", "2026-08"); + ReservationId directId = reserve(fixture, directKey, "request-1"); + fixture.accounting().markInFlight(directId); + ActualUsageCommand directCommand = command("request-1", directId, 50); + ReservationReconciliation direct = fixture.accounting().commit(directCommand); + ReservationReconciliation directDuplicate = fixture.accounting().commit(directCommand); + + BudgetKey lateKey = key("tenant-b", "2026-08"); + ReservationId lateId = reserve(fixture, lateKey, "request-2"); + fixture.accounting().markInFlight(lateId); + fixture.accounting().markReconciliationRequired(lateId); + ActualUsageCommand lateCommand = command("request-2", lateId, 50); + ReservationReconciliation late = fixture.accounting().reconcileLateActual(lateCommand); + ReservationReconciliation lateDuplicate = fixture.accounting() + .reconcileLateActual(lateCommand); + + assertThat(direct.transition().status()).isEqualTo(AccountingTransitionStatus.APPLIED); + assertThat(directDuplicate.transition().status()) + .isEqualTo(AccountingTransitionStatus.REUSED); + assertThat(late.transition().status()).isEqualTo(AccountingTransitionStatus.APPLIED); + assertThat(lateDuplicate.transition().status()) + .isEqualTo(AccountingTransitionStatus.REUSED); + assertThat(events) + .extracting(BudgetNotificationEvent::source) + .containsExactly( + BudgetNotificationSource.ACCOUNTING_COMMIT, + BudgetNotificationSource.LATE_RECONCILIATION + ); + } + + @Test + void handler_실패는_다음_handler와_commit_BLOCK_결과에_영향을_주지_않는다() { + List delivered = new ArrayList<>(); + List errors = new ArrayList<>(); + BudgetNotificationHandler failing = event -> { + throw new IllegalStateException("api-key=secret\nraw-response"); + }; + Fixture fixture = fixture(List.of(failing, delivered::add), errors::add); + BudgetKey committedKey = key("tenant-a", "2026-08"); + + ReservationReconciliation committed = commit( + fixture, + committedKey, + "request-1", + 50 + ); + BudgetKey blockedKey = key("tenant-b", "2026-08"); + BudgetReservationResult blocked = fixture.stateStore().checkAndReserve( + request(blockedKey, "blocked-request", LIMIT) + ); + + assertThat(committed.transition().status()) + .isEqualTo(AccountingTransitionStatus.APPLIED); + assertThat(fixture.stateStore().snapshot(committedKey, LIMIT).committedCost()) + .isEqualTo(usd("50")); + assertThat(blocked.status()).isEqualTo(ReservationStatus.BLOCKED); + assertThat(fixture.stateStore().snapshot(blockedKey, LIMIT).effectiveUsage()) + .isEqualTo(Cost.zero(USD)); + assertThat(delivered) + .extracting(BudgetNotificationEvent::threshold) + .containsExactly(BudgetThreshold.HALF, BudgetThreshold.EXCEEDED); + assertThat(errors).hasSize(2).allSatisfy(error -> { + assertThat(error.message()).hasSizeLessThanOrEqualTo(128); + assertThat(error.toString()).doesNotContain("secret", "raw-response"); + }); + } + + @Test + @Timeout(value = 10, unit = TimeUnit.SECONDS) + void 병렬_event에서도_threshold_dedup은_원자적으로_동작한다() throws Exception { + List accountingEvents = new ArrayList<>(); + Fixture source = fixture( + List.of(event -> { }), + BudgetNotificationErrorHook.noOp(), + accountingEvents + ); + ReservationReconciliation applied = commit( + source, + key("tenant-a", "2026-08"), + "request-1", + 50 + ); + assertThat(applied.transition().status()).isEqualTo(AccountingTransitionStatus.APPLIED); + assertThat(accountingEvents).hasSize(1); + + AtomicInteger deliveries = new AtomicInteger(); + BudgetNotificationService service = new BudgetNotificationService( + List.of(event -> deliveries.incrementAndGet()), + new InMemoryNotificationStateStore(), + ignored -> deliverySnapshot(accountingEvents) + ); + AccountingDelivery delivery = accountingEvents.get(0); + List> commands = java.util.stream.IntStream.range(0, 200) + .mapToObj(ignored -> (Callable) () -> { + service.onAccountingApplied(delivery.event(), delivery.snapshot()); + return null; + }) + .toList(); + + runConcurrently(commands); + + assertThat(deliveries).hasValue(1); + } + + @Test + void 기존_onCommitted_callback도_resolver_snapshot으로_알림을_만든다() { + List accountingEvents = new ArrayList<>(); + Fixture source = fixture( + List.of(event -> { }), + BudgetNotificationErrorHook.noOp(), + accountingEvents + ); + commit(source, key("tenant-a", "2026-08"), "request-1", 50); + AccountingDelivery delivery = accountingEvents.get(0); + List events = new ArrayList<>(); + BudgetNotificationService service = new BudgetNotificationService( + List.of(events::add), + new InMemoryNotificationStateStore(), + ignored -> delivery.snapshot() + ); + + service.onCommitted(delivery.event()); + + assertThat(events) + .extracting(BudgetNotificationEvent::threshold) + .containsExactly(BudgetThreshold.HALF); + } + + @Test + @SuppressWarnings("removal") + void event는_민감한_payload나_임의_tags를_보존하지_않는다() { + BudgetNotificationEvent event = new BudgetNotificationEvent( + key("tenant-a", "2026-08"), + BudgetThreshold.HALF, + io.tokenpilot.budget.BudgetState.WARN, + "threshold reached", + usd("50"), + LIMIT, + Map.of("api_key", "secret", "prompt", "private") + ); + + assertThat(event.tags()).isEmpty(); + assertThat(event.getClass().getRecordComponents()) + .extracting(java.lang.reflect.RecordComponent::getName) + .doesNotContain("tags", "prompt", "rawProviderResponse", "apiKey"); + } + + private static ReservationReconciliation commit( + Fixture fixture, + BudgetKey key, + String requestId, + long actualCost + ) { + ReservationId reservationId = reserve(fixture, key, requestId); + fixture.accounting().markInFlight(reservationId); + return fixture.accounting().commit(command(requestId, reservationId, actualCost)); + } + + private static ReservationId reserve( + Fixture fixture, + BudgetKey key, + String requestId + ) { + BudgetReservationResult result = fixture.stateStore().checkAndReserve( + request(key, requestId, usd("1.00")) + ); + assertThat(result.status()).isEqualTo(ReservationStatus.CREATED); + return result.reservationId(); + } + + private static BudgetReservationRequest request( + BudgetKey key, + String requestId, + Cost estimate + ) { + PricingSnapshot snapshot = pricingSnapshot(); + return new BudgetReservationRequest( + key, + LIMIT, + estimate, + requestId, + new IdempotencyKey("idempotency-" + requestId), + snapshot, + TOKEN_ESTIMATE + ); + } + + private static ActualUsageCommand command( + String requestId, + ReservationId reservationId, + long actualCost + ) { + return new ActualUsageCommand( + requestId, + "attempt-" + requestId, + reservationId, + TokenUsage.from(actualCost, 0), + "gpt-4o-mini-response" + ); + } + + private static Fixture fixture( + List handlers, + BudgetNotificationErrorHook errorHook + ) { + return fixture(handlers, errorHook, null); + } + + private static Fixture fixture( + List handlers, + BudgetNotificationErrorHook errorHook, + List accountingEvents + ) { + AtomicReference stateStoreReference = new AtomicReference<>(); + BudgetNotificationService service = new BudgetNotificationService( + handlers, + new InMemoryNotificationStateStore(), + key -> stateStoreReference.get().snapshot(key, LIMIT), + errorHook + ); + AtomicInteger sequence = new AtomicInteger(); + List listeners = + accountingEvents == null + ? List.of(service) + : List.of( + service, + new io.tokenpilot.budget.ReservationAccountingListener() { + @Override + public void onCommitted(ReservationAccountingEvent event) { + } + + @Override + public void onAccountingApplied( + ReservationAccountingEvent event, + BudgetSnapshot snapshot + ) { + accountingEvents.add(new AccountingDelivery(event, snapshot)); + } + } + ); + BudgetStateStore stateStore = LedgerBudgetComponents.inMemoryBudgetStateStore( + CLOCK, + () -> new ReservationId("reservation-" + sequence.incrementAndGet()), + (usage, plan) -> usd(Long.toString(usage.inputTokens())), + listeners + ); + stateStoreReference.set(stateStore); + return new Fixture( + stateStore, + LedgerBudgetComponents.reservationAccounting(stateStore) + ); + } + + private static PricingSnapshot pricingSnapshot() { + return new PricingSnapshot( + "gpt-4o-mini-request", + "pricing-v1", + "catalog-v1", + CLOCK.instant(), + Map.of( + TokenType.PROMPT, BigDecimal.ONE, + TokenType.COMPLETION, BigDecimal.ONE + ), + USD + ); + } + + private static BudgetKey key(String tenantId, String window) { + return new BudgetKey( + "budget-policy", + "tenant", + tenantId, + BudgetWindow.parse(window) + ); + } + + private static Cost usd(String amount) { + return Cost.of(new BigDecimal(amount), USD); + } + + private static BudgetSnapshot deliverySnapshot( + List deliveries + ) { + return deliveries.get(0).snapshot(); + } + + private static void runConcurrently(List> commands) + throws Exception { + CountDownLatch ready = new CountDownLatch(commands.size()); + CountDownLatch start = new CountDownLatch(1); + try (var executor = Executors.newVirtualThreadPerTaskExecutor()) { + List> futures = new ArrayList<>(commands.size()); + for (Callable command : commands) { + futures.add(executor.submit(() -> { + ready.countDown(); + if (!start.await(5, TimeUnit.SECONDS)) { + throw new IllegalStateException("start barrier timed out"); + } + return command.call(); + })); + } + assertThat(ready.await(5, TimeUnit.SECONDS)).isTrue(); + start.countDown(); + for (Future future : futures) { + future.get(5, TimeUnit.SECONDS); + } + } finally { + start.countDown(); + } + } + + private record Fixture( + BudgetStateStore stateStore, + ReservationAccounting accounting + ) { + } + + private record AccountingDelivery( + ReservationAccountingEvent event, + BudgetSnapshot snapshot + ) { + } +}