diff --git a/AGENTS.md b/AGENTS.md index 72e3a05..a2c6b41 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, Clock-based monthly budget windows, pure budget decisions, typed missing-pricing policies, pricing snapshots, framework-independent token count results, a UTF-8 byte heuristic estimator, a preflight cost-bound projection, versioned model metadata, conservative context admission, a plain-Java core consumer verification path, framework-independent in-memory atomic reservations, and estimate/actual reconciliation with best-effort accounting events are implemented. Candidate-aware request production and Spring AI lifecycle integration remain 30-day MVP targets, not current capabilities. +Current truth: post-call usage normalization, cost calculation, ledger events, Token Pilot-owned Micrometer control/accounting metrics, Clock-based monthly budget windows, pure budget decisions, typed missing-pricing policies, pricing snapshots, framework-independent token count results, a UTF-8 byte heuristic estimator, a preflight cost-bound projection, versioned model metadata, conservative context admission, a plain-Java core consumer verification path, framework-independent in-memory atomic reservations, estimate/actual reconciliation with observable best-effort accounting events, and atomic threshold notifications are implemented. The Spring AI 2.0 adapter owns the supported non-streaming request lifecycle from text-scope adaptation and conservative preflight through reservation, provider invocation, model-aware actual reconciliation, and metrics. Real-provider compatibility, durable delivery, and advanced streaming reconciliation remain outside the verified capability. 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. @@ -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 @@ -71,14 +72,14 @@ Token Pilot의 제품 포지션은 framework-independent Java LLM control and ac | Module | Status | Notes | | --- | --- | --- | -| `token-pilot-core` | Basic implementation complete | Domain records, pricing, calculator, registry, ledger manager, pricing snapshots, versioned model catalog, token count results, UTF-8 byte heuristic estimation, preflight cost-bound projection, conservative context admission, and public plain-Java consumer verification | -| `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-autoconfigure` | Basic implementation complete | Bean registration, property binding, pricing/budget/notification wiring, and `ChatClientBuilderCustomizer` implemented | +| `token-pilot-core` | Basic implementation complete | Domain records, pricing, calculator, registry, ledger manager, pricing snapshots, versioned model catalog, token count results, UTF-8 byte heuristic estimation, preflight cost-bound projection, conservative context admission, decision/pricing-miss listener contracts, and public plain-Java consumer verification | +| `token-pilot-spring-ai` | Non-streaming accounting integration implemented | Spring AI 2.0.0 text-scope adaptation, model/output resolution, framing headroom, candidate-aware preflight and reservation, pricing-miss observation, model-aware around-call reconciliation, provider-specific reserved-output resolver extension, and enforcement streaming rejection implemented | +| `token-pilot-micrometer` | Token Pilot-owned metrics implemented | Optional owner-specific publishers cover cost, preflight, reservation, pricing-miss, reconciliation, listener-failure, and notification outcomes; legacy `ai.token.*` publishing is opt-in | +| `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, bounded listener-failure observation, and framework-independent best-effort accounting events implemented; durable stores remain | +| `token-pilot-notification` | Atomic accounting integration implemented | Commit, reconciliation-required, late reconciliation, and reservation BLOCK results produce process-local atomically deduplicated threshold events with isolated handlers, lifecycle observation, and a sanitized error hook; durable delivery remains | +| `token-pilot-autoconfigure` | Basic implementation complete | Owner-specific optional bean graphs wire pricing, atomic budget/accounting, Spring AI, notification, and metrics while preserving budget-disabled ledger-only compatibility and `ChatClientBuilderCustomizer` | | `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 | +| `token-pilot-sample-app` | Basic E2E complete | Direct ledger metrics and fake Spring AI lifecycle E2E cover preflight BLOCK, reservation, reconciliation, request-scope rejection, Advisor ordering, listener isolation, and enforcement streaming rejection without an API key | ## Current Work Focus @@ -99,13 +100,14 @@ Gradle dependency cleanup has landed. Library modules should not regain app-only Current Micrometer status: -- `MetricsOptions` exists as the small Micrometer options object. -- Default allowed tag keys remain `tenant_id`. -- The existing `MicroCostMetricsPublisher(MeterRegistry)` constructor is preserved. -- Tests cover null/empty tags and multiple allowed tags. -- Metric descriptions and base units should remain stable. -- Current `ai.token.*` metrics may overlap Spring AI Observability; treat renaming/default suppression as an explicit compatibility decision. -- New default metrics should describe Token Pilot-owned cost, preflight, reservation, pricing-miss, and reconciliation outcomes. +- Default Token Pilot metrics describe cost, preflight, reservation, pricing-miss, reconciliation, isolated listener-failure, and notification lifecycle outcomes. +- Default user tag keys are empty, and Token Pilot-owned metrics use only bounded enum/registered identifier values. +- `tokenpilot.cost.total` consumes newly applied actual reservation commits, never legacy ledger estimates or missing-pricing zero values. +- Duplicate/reused reconciliation callbacks do not republish accounting metrics; actual-unavailable transitions publish `reconciliation_required` without zero cost or error samples. +- The existing `MicroCostMetricsPublisher` constructors remain available for direct compatibility. +- Legacy `ai.token.*` metrics are off in starter/autoconfigure unless `token-pilot.metrics.legacy-ai-token-metrics-enabled=true`. +- The legacy metric opt-in, including its raw `model` compatibility tag, is a 0.1.x bridge planned for removal in 0.2.0. +- Micrometer `double` counters are operational telemetry; exact `BigDecimal` ledger values remain the accounting source of truth. ## Distribution Contract @@ -149,6 +151,7 @@ The autoconfigure module provides: - `META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports` - `TokenPilotAutoConfiguration` +- Owner-specific core metrics, budget, Spring AI, notification, and metrics auto-configurations - `TokenPilotProperties` - Pricing property binding - Budget property binding @@ -192,15 +195,19 @@ Default bean graph: | --- | --- | --- | | `CostCalculator` | missing bean | Core cost calculation | | `PricingRegistry` | missing bean | Pricing plan lookup | +| `ModelRegistry` | missing bean | Versioned default model/context metadata | +| `TokenEstimator` | missing bean | UTF-8 byte heuristic token estimate | +| `TokenBudget` | missing bean | Context admission and preflight decision events | | `LedgerManager` | missing bean | Cost and usage recording | | `UsageExtractor` | Spring AI classpath + missing bean | Spring AI response usage extraction | | `LedgerAdvisor` | Spring AI classpath + missing bean | ChatClient advisor | -| `MicroCostMetricsPublisher` | Micrometer classpath + `token-pilot.metrics.enabled` | Cost/token metrics listener | -| `BudgetStateStore` | `token-pilot.budget.enabled` + missing bean | Default in-memory budget state | +| owner-specific Token Pilot metric publishers | Micrometer + owner contract classpath + `token-pilot.metrics.enabled` | Low-cardinality control/accounting metrics | +| `MicroCostMetricsPublisher` | Micrometer + explicit legacy flag | Compatibility `ai.token.*` metrics listener | +| `AtomicBudgetStateStore` | `token-pilot.budget.enabled` + missing `BudgetStateStore` | One default in-memory bean for reservation admission and accounting transitions | | `BudgetEvaluator` | `token-pilot.budget.enabled` + missing bean | Default budget evaluator | | `ChatClientBuilderCustomizer` | Spring AI classpath + `LedgerAdvisor` bean | Adds advisor to ChatClient builders | | `NotificationStateStore` | `token-pilot.notification.enabled` + missing bean | Window-based notification deduplication state | -| `BudgetNotificationService` | `token-pilot.notification.enabled` + `BudgetNotificationHandler` bean | Publishes budget notification events to user-defined handler | +| `BudgetNotificationService` | notification enabled + `BudgetNotificationHandler` + `BudgetStateStore` | Publishes atomic budget/accounting notifications to user-defined handlers | `core.internal` implementation classes should remain package-private. Cross-module construction should go through `LedgerComponents` or another deliberate public factory/API. Do not make internal implementation classes public just to satisfy autoconfigure access. @@ -216,6 +223,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 @@ -232,10 +240,15 @@ class MailBudgetNotificationHandler implements BudgetNotificationHandler { ``` - `BudgetNotificationHandler` 빈이 없으면 `BudgetNotificationService`는 등록되지 않는다 (no-op). -- `token-pilot.notification.enabled=true` 설정 시에만 notification 빈이 등록된다. -- 알림 중복 방지는 evaluator가 확정한 `BudgetKey(policyId, targetType, targetId, window)`로 처리된다. -- 같은 window 안에서는 낮거나 같은 threshold 재발송이 방지된다. +- `token-pilot.notification.enabled=true` 설정 시에만 notification 상태 빈이 등록된다. +- `BudgetNotificationService` 자동 설정은 atomic accounting 근거가 필요하므로 budget 활성화와 `BudgetStateStore`도 요구한다. +- 알림은 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 @@ -251,9 +264,8 @@ token-pilot: COMPLETION: 0.00060 metrics: enabled: true - tag-whitelist: - - tenant_id - - model + legacy-ai-token-metrics-enabled: false + tag-whitelist: [] budget: enabled: false policy-id: default-monthly @@ -285,7 +297,7 @@ Test-only E2E endpoint: - `GET /test/token-pilot/chat`: exercises the Spring AI `ChatClient` advisor path with a fake/mock provider or documented real provider setup. -The direct ledger E2E test verifies that `/actuator/prometheus` contains token-pilot metrics after a ledger event is recorded. The fake ChatClient E2E test verifies that Spring AI `ChatClient` calls flow through `LedgerAdvisor` into token-pilot metrics without requiring a real provider API key. +The default starter E2E verifies a Token Budget preflight metric while legacy `ai.token.*` meters and sensitive tags remain absent. The fake ChatClient E2E verifies the Spring AI advisor path, explicit legacy compatibility mode, and a bounded pricing-miss metric without requiring a real provider API key. ## Maven Publishing Direction @@ -328,11 +340,21 @@ The active checklist is in `docs/30_DAY_MVP_REPORT.md`; detailed long-term works - Budget money interfaces now use `Cost` while preserving `BudgetKey`, `BudgetPolicy`, Clock/ZoneId monthly windows, and per-key policy snapshots. - The legacy `DefaultLedgerManager.record(String, ...)` path preserves an explicit zero USD fail-open result for a missing plan; the pricing-snapshot path applies `MissingPricingPolicy` and records `UNPRICED` or rejects before provider invocation, so neither behavior is a priced zero-rate plan. - Spring AI usage extraction converts map/JSON-compatible native usage objects into the normalized core model. Real-provider compatibility fixtures remain required because provider and Spring AI usage shapes can change independently. -- The legacy provider boundary blocks an already-exhausted budget decision before provider invocation. Its candidate-free `STATUS` input is a regression guard, not admission evidence; the flow remains check-then-add and is not connected to the new atomic reservation lifecycle until #39. -- In-memory reservation reconciliation uses the reservation-time pricing snapshot, accepts only provider-reported or provider-derived actual usage, moves estimate liability atomically between active, pending, and committed totals, and skips cost calculation for exact duplicate callbacks. Legacy reservations without pricing/token metadata have an explicit cost-only settlement path; new reservations should use the usage-based API. Spring AI callback integration remains #39. -- Accounting listeners run synchronously after the bucket lock is released. Runtime listener failures do not roll back a committed transition, stop later listeners, or trigger redelivery on duplicate callbacks, but delivery remains best-effort at-most-once without a durable outbox; failure observation remains #40. -- Current Micrometer `ai.token.*` metrics may duplicate Spring AI Observability; preserve compatibility while deciding default suppression or replacement. -- The verified Spring AI 2.0.0 path is synchronous `ChatClient` usage recording with a fake provider. Streaming cancellation and reconciliation remain outside the current compatibility guarantee. +- Budget-enabled Spring AI calls use the candidate-aware atomic reservation lifecycle. Budget-disabled calls preserve the legacy ledger-only path; its candidate-free status check is a compatibility guard and must not be described as atomic admission. +- In-memory reservation reconciliation uses the reservation-time pricing snapshot, accepts only provider-reported or provider-derived actual usage, moves estimate liability atomically between active, pending, and committed totals, and skips cost calculation for exact duplicate callbacks. Legacy reservations without pricing/token metadata have an explicit cost-only settlement path; new reservations and the Spring AI lifecycle use the usage-based API. +- Accounting listeners run synchronously after the bucket lock is released. Runtime listener failures do not roll back a committed transition, stop later listeners, or trigger redelivery on duplicate callbacks; bounded failure events feed `tokenpilot.listener.failures`, but delivery remains best-effort at-most-once without a durable outbox. +- Legacy `LedgerListener` runtime failures are likewise isolated from ledger/provider results and later listeners; JVM `Error`s still propagate. +- Compatibility cost-only reservation commits do not carry token/model correlation and therefore do not publish accounting metrics; new reservations should use usage-based reconciliation. +- Legacy Micrometer `ai.token.*` metrics may duplicate Spring AI Observability and are therefore disabled by default; their explicit migration opt-in still permits application-provided tag values with caller-owned cardinality. +- The verified Spring AI 2.0.0 path is the non-streaming `ChatClient` lifecycle with a fake provider. Enforcement-enabled streaming fails closed before reservation or provider invocation; budget-disabled ledger-only streaming is a regression-tested compatibility path, not a 0.1.0 streaming accounting guarantee. Chunk accounting, cancellation, and partial-usage reconciliation remain post-MVP. +- Spring request admission is limited to the supported `TEXT_ONLY` scope. Role/length framing and configured headroom preserve message boundaries and cover known Advisor additions, but they are not an exact provider payload count. Media, tool messages/calls/schema, and structured-output augmentation fail closed before reservation and provider invocation. +- Spring AI cannot expose the exact provider-internal dispatch boundary. `LedgerAdvisor` runs at `LOWEST_PRECEDENCE - 1` so ordinary user Advisors finish before Token Pilot and the terminal model Advisor runs after it, but callers can still alter ordering explicitly. Downstream failures after `markInFlight()` therefore preserve estimate liability as `RECONCILIATION_REQUIRED`. +- Request `ChatOptions.maxTokens` takes precedence over an optional user-provided `ReservedOutputTokensResolver`, which in turn precedes the configured default. Resolver values must be positive; provider-specific option capability implementations remain user-supplied. +- Generated idempotency keys deduplicate repeated commands only when the same key is reused within the accounting contract. A key generated for one Advisor invocation does not guarantee deduplication of an upstream retry that creates a new request context. +- Reservation idempotency fingerprints ignore `PricingSnapshot.checkedAt` because it is observation metadata. Model, pricing policy, catalog version, currency, semantic rate values, token estimates, limits, request identity, and other liability-bearing inputs still produce `CONFLICT` when they change under the same key. +- Budget auto-configuration exposes one `AtomicBudgetStateStore` bean as both `BudgetStateStore` and `ReservationAccounting`; a custom `BudgetStateStore` used with enforcement must also expose the accounting capability or startup fails instead of silently selecting the ledger-only path. +- A response model that differs from the reservation pricing snapshot is not committed at the request model price; it remains pending with `PRICING_RECONCILIATION_REQUIRED` for later settlement. +- Pricing-mismatch pending events preserve provider actual usage and response model metadata. Consumers may call the explicit `reconcileLateActual(command, responsePricingSnapshot)` API with an immutable response-model pricing snapshot; model, currency, rate, state, and duplicate-callback validation remain fail-closed. - The repository, README, JReleaser configuration, and every published module POM use the MIT License. `verifyPublicationMetadata` guards this release contract and ensures the sample app is not published. - Sample app E2E uses a fake Spring AI `ChatModel`; real provider API behavior is not yet verified. - `token-pilot-spring-ai-starter` does not exist in the current build; never use it as an install instruction until implemented and published. @@ -356,6 +378,18 @@ Verify the fixed Java/Spring compatibility matrix and generated core-only consum ./gradlew verifyCompatibilityMatrix verifyCoreConsumer ``` +Verify the published Micrometer artifact with optional owners absent/composed and core autoconfiguration without optional adapters: + +```bash +./gradlew verifyMicrometerConsumer +``` + +Verify the published Spring AI adapter API and the starter with one explicitly selected provider: + +```bash +./gradlew verifyPublishedIntegrationConsumer +``` + Run sample app after implementation work: ```bash @@ -402,6 +436,23 @@ Stage and deploy a Central release: ## Update History +### 2026-08-25 + +- Connected supported Spring AI 2.0 non-streaming `ChatClient` requests to text-scope adaptation, model/output resolution, conservative preflight, atomic reservation, in-flight tracking, provider usage extraction, and `ReservationAccounting` reconciliation in one around-call lifecycle. +- Added namespaced request correlation without `ThreadLocal`, generated-or-reused idempotency keys, provider-specific reserved-output resolution, configurable framing headroom, and fail-closed media/tool/structured-output scope handling. +- Placed Token Pilot immediately before the terminal model Advisor, rejected enforcement-enabled streaming before reservation/provider invocation, and preserved the budget-disabled ledger-only compatibility path. +- Added fake-provider E2E coverage for admission failures, fallback resolution, dispatch boundaries, pending liability, model correlation, usage normalization, unsupported request scope, Advisor ordering, listener isolation, duplicate accounting commands, and streaming rejection. +- 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. +- Added owner-specific Micrometer publishers for Token Pilot cost, preflight, atomic reservation, pricing-miss, reconciliation, listener-failure, and notification lifecycle outcomes with bounded default tags. +- Connected framework-independent decision/lifecycle listener contracts while preserving duplicate-safe applied accounting semantics and observable best-effort listener failure isolation. +- Disabled legacy `ai.token.*` metrics and the default user tag whitelist in starter/autoconfigure; retained explicit compatibility opt-in and existing direct publisher constructors. +- Added published-artifact consumer checks for Micrometer with optional owners both absent and explicitly composed. +- Published Spring AI, Reactor, and budget types exposed by the adapter at API scope, carried the supported Boot/Spring AI platforms in both Maven and Gradle metadata, and added external adapter/starter compile-runtime consumers with explicit provider selection. +- Rejected duplicate in-flight provider dispatch, preserved response-model pricing mismatches as pending liability, and made incomplete budget/accounting graphs fail at startup instead of silently downgrading enforcement. +- Added the `AtomicBudgetStateStore` composite contract so Spring exposes reservation admission and reconciliation through one unambiguous state bean. +- Preserved response-model pricing mismatches as auditable pending actual metadata and added explicit late repricing reconciliation with request/response pricing provenance and duplicate-safe accounting events. + ### 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/README.md b/README.md index 28de07d..a0cf04f 100644 --- a/README.md +++ b/README.md @@ -51,11 +51,87 @@ library behind it. ## Project Status TokenPilot 0.1.0 is under active development. The current foundation includes -Spring AI usage integration, model cost calculation, metrics, basic budget -policy, Spring Boot autoconfiguration, and a framework-independent in-memory -atomic budget reservation/idempotency foundation. The MVP is extending this -foundation with provider-connected preflight control and actual usage -reconciliation. +model cost calculation, owner-specific metrics, atomic budget reservation and +notification, and Spring Boot autoconfiguration. For supported non-streaming +Spring AI `ChatClient` calls, TokenPilot now performs conservative preflight, +claims one provider dispatch, and reconciles provider-reported actual usage +against the reservation-time pricing snapshot. + +## Micrometer metrics + +When a `MeterRegistry` is available, the Spring Boot starter publishes +TokenPilot-owned, low-cardinality metrics for control and accounting outcomes: + +| Metric | Tags | Meaning | +| --- | --- | --- | +| `tokenpilot.cost.total` | `currency` | Newly committed actual cost from usage-based reconciliation | +| `tokenpilot.preflight.requests` | `decision`, `reason` | Context admission decisions | +| `tokenpilot.budget.reservations` | `state` | Atomic reservation results | +| `tokenpilot.reconciliation.error.tokens` | `direction` | Absolute estimate/actual token error | +| `tokenpilot.reconciliation.outcomes` | `outcome`, `reason` | Applied reconciliation outcomes | +| `tokenpilot.pricing.missing` | `policy` | Missing pricing observed at the provider boundary | +| `tokenpilot.listener.failures` | `listener`, `phase` | Isolated accounting-listener failures | +| `tokenpilot.notification.events` | `outcome`, `threshold` | Notification delivery and deduplication outcomes | + +The default user-tag whitelist is empty. The metrics above never include raw +model, tenant, user, request, reservation, or idempotency identifiers. Their +tag values come from bounded domain enums or registered currency codes. + +```yaml +token-pilot: + metrics: + enabled: true + tag-whitelist: [] + legacy-ai-token-metrics-enabled: false +``` + +The former `ai.token.*` meters are disabled by default because Spring AI +Observability may already publish standard token telemetry. Set +`token-pilot.metrics.legacy-ai-token-metrics-enabled=true` to opt in during +migration. This is a 0.1.x compatibility bridge, including the legacy raw +`model` tag, and is planned for removal in 0.2.0. Migrate dashboards to Spring +AI token telemetry and the `tokenpilot.*` control/accounting meters before +then. `tag-whitelist` applies only to that legacy path and limits keys, not the +cardinality of application-provided values. Existing direct +`MicroCostMetricsPublisher` constructors retain their legacy `tenant_id` +allowlist behavior; the starter default remains empty. + +Accounting metrics consume newly applied reservation transitions, so reused +callbacks do not add cost twice and unavailable actual usage is recorded as +`reconciliation_required`, not as zero cost or zero error. Listener delivery +is synchronous, best-effort, and at-most-once without a durable outbox. +Micrometer counters use `double` internally and are operational telemetry, not +the monetary source of truth; the ledger's `BigDecimal` values remain +authoritative. The legacy cost-only commit methods cannot carry token/model +correlation and do not emit these accounting metrics; new reservations should +use the usage-based reconciliation API. + +`LedgerListener` and other optional observer `RuntimeException`s are isolated: +they do not change ledger/provider results and later listeners still run. +JVM `Error`s are not swallowed. + +For direct autoconfiguration composition, use +`TokenPilotBudgetPolicyFactory.from(properties)` instead of the former +`TokenPilotProperties.toBudgetPolicy()`. Keeping the budget return type out of +the shared properties class allows autoconfiguration to start when the optional +budget module is absent. + +## Spring AI starter + +Spring Boot applications use one Token Pilot convenience starter and select +their Spring AI provider separately. For example: + +```gradle +dependencies { + implementation 'cloud.token-pilot:token-pilot-starter:' + implementation 'org.springframework.ai:spring-ai-starter-model-openai:2.0.0' +} +``` + +Token Pilot does not choose or bundle a provider. The published adapter and +starter are compile/runtime verified from their generated Maven and Gradle +metadata against the supported Java 25, Spring Boot 4.1.0, and Spring AI 2.0.0 +baseline. ## Framework-independent core @@ -108,10 +184,19 @@ Boot, Spring AI, Micrometer, or Reactor dependencies. Spring Boot 3, Spring AI are not part of the 0.1.0 support guarantee. The verified Spring AI path is the synchronous `ChatClient` call lifecycle with -a fake provider. Streaming cancellation, provider-connected preflight -blocking, reservation lifecycle reconciliation, and estimate/actual -reconciliation are still MVP work and are not included in this compatibility -claim. +a fake provider, including preflight blocking, atomic reservation, dispatch, +and estimate/actual reconciliation. Real-provider compatibility, chunk +accounting, streaming cancellation, and partial-usage reconciliation are not +included in this compatibility claim. + +If a provider returns a model different from the request pricing snapshot, +TokenPilot keeps the estimate as `PRICING_RECONCILIATION_REQUIRED` instead of +charging the request model's price. The pending event preserves the provider +usage and response model; an application that has an immutable response-model +pricing snapshot can finish the lifecycle with +`ReservationAccounting.reconcileLateActual(command, responsePricingSnapshot)`. +Mismatched model, currency, pricing terms, state, or duplicate callbacks remain +fail-closed. ## License diff --git a/build.gradle b/build.gradle index b0a3f97..0595052 100644 --- a/build.gradle +++ b/build.gradle @@ -13,7 +13,7 @@ plugins { def publishedModuleDescriptions = [ 'token-pilot-core': 'Framework-independent LLM usage normalization, pricing, cost calculation, and ledger contracts.', 'token-pilot-spring-ai': 'Optional Spring AI adapter for recording provider-reported LLM usage through Token Pilot.', - 'token-pilot-micrometer': 'Optional Micrometer publisher for Token Pilot cost and usage accounting events.', + 'token-pilot-micrometer': 'Optional Micrometer publishers for Token Pilot control, cost, budget, reconciliation, and notification events.', 'token-pilot-budget': 'Budget policy evaluation and in-memory budget state components for Token Pilot.', 'token-pilot-notification': 'Budget notification events, handlers, and deduplication components for Token Pilot.', 'token-pilot-autoconfigure': 'Spring Boot auto-configuration for Token Pilot core and optional adapters.', @@ -30,7 +30,6 @@ def compatibilityMatrix = [ ] def springManagedModules = [ - 'token-pilot-spring-ai', 'token-pilot-autoconfigure', 'token-pilot-starter', 'token-pilot-sample-app' @@ -312,6 +311,7 @@ tasks.register('verifyPublicationMetadata') { group = 'verification' description = 'Verifies generated Maven POM metadata for every published Token Pilot module.' dependsOn publicationPomTasks + dependsOn ':token-pilot-spring-ai:generateMetadataFileForMavenJavaPublication' inputs.file(layout.projectDirectory.file('LICENSE')) doLast { @@ -371,6 +371,150 @@ tasks.register('verifyPublicationMetadata') { ) requirePomValue(moduleName, 'SCM URL', pom.scm.url.text(), 'https://github.com/tokenpliot/tokenpilot') } + + def micrometerPomFile = project(':token-pilot-micrometer') + .layout.buildDirectory + .file('publications/mavenJava/pom-default.xml') + .get() + .asFile + def micrometerPom = new XmlSlurper(false, false).parse(micrometerPomFile) + def micrometerDependencies = micrometerPom.dependencies.dependency.collect { + "${it.groupId.text()}:${it.artifactId.text()}:${it.scope.text()}".toString() + } as Set + if (!micrometerDependencies.contains( + 'cloud.token-pilot:token-pilot-core:compile' + ) || !micrometerDependencies.contains( + 'io.micrometer:micrometer-core:compile' + )) { + throw new GradleException( + 'token-pilot-micrometer POM must expose core and micrometer-core as compile dependencies.' + ) + } + def forbiddenMicrometerOwners = micrometerPom.dependencies.dependency.findAll { + it.groupId.text() == 'cloud.token-pilot' && + it.artifactId.text() in [ + 'token-pilot-budget', + 'token-pilot-notification' + ] + } + if (!forbiddenMicrometerOwners.isEmpty()) { + throw new GradleException( + 'token-pilot-micrometer POM must keep budget and notification owners optional.' + ) + } + + def springAiPomFile = project(':token-pilot-spring-ai') + .layout.buildDirectory + .file('publications/mavenJava/pom-default.xml') + .get() + .asFile + def springAiPom = new XmlSlurper(false, false).parse(springAiPomFile) + def springAiDependencies = springAiPom.dependencies.dependency.collect { + "${it.groupId.text()}:${it.artifactId.text()}:${it.scope.text()}".toString() + } as Set + def expectedSpringAiDependencies = [ + 'cloud.token-pilot:token-pilot-core:compile', + 'cloud.token-pilot:token-pilot-budget:compile', + 'org.springframework.ai:spring-ai-model:compile', + 'org.springframework.ai:spring-ai-client-chat:compile', + 'io.projectreactor:reactor-core:compile' + ] as Set + if (springAiDependencies != expectedSpringAiDependencies) { + throw new GradleException( + 'token-pilot-spring-ai POM must expose every public ABI dependency at compile scope. ' + + "Expected ${expectedSpringAiDependencies}, actual ${springAiDependencies}." + ) + } + def springAiManagedDependencies = springAiPom.dependencyManagement.dependencies.dependency.collect { + "${it.groupId.text()}:${it.artifactId.text()}:${it.version.text()}:${it.type.text()}:${it.scope.text()}".toString() + } + def expectedSpringAiManagedDependencies = [ + 'org.springframework.boot:spring-boot-dependencies:4.1.0:pom:import', + 'org.springframework.ai:spring-ai-bom:2.0.0:pom:import' + ] as Set + if (springAiManagedDependencies.size() != 2 || + (springAiManagedDependencies as Set) != expectedSpringAiManagedDependencies) { + throw new GradleException( + 'token-pilot-spring-ai POM must import each supported BOM exactly once. ' + + "Expected ${expectedSpringAiManagedDependencies}, actual ${springAiManagedDependencies}." + ) + } + + def springAiModuleFile = project(':token-pilot-spring-ai') + .layout.buildDirectory + .file('publications/mavenJava/module.json') + .get() + .asFile + def springAiModule = new JsonSlurper().parse(springAiModuleFile) + def expectedSpringAiVariantDependencies = [ + 'org.springframework.boot:spring-boot-dependencies', + 'org.springframework.ai:spring-ai-bom', + 'cloud.token-pilot:token-pilot-core', + 'cloud.token-pilot:token-pilot-budget', + 'org.springframework.ai:spring-ai-model', + 'org.springframework.ai:spring-ai-client-chat', + 'io.projectreactor:reactor-core' + ] as Set + [ + apiElements: 'java-api', + runtimeElements: 'java-runtime' + ].each { variantName, usage -> + def variant = springAiModule.variants.find { it.name == variantName } + if (variant == null || variant.attributes['org.gradle.usage'] != usage) { + throw new GradleException( + "token-pilot-spring-ai module metadata is missing ${variantName}." + ) + } + def variantDependencies = variant.dependencies.collect { + "${it.group}:${it.module}".toString() + } as Set + if (variantDependencies != expectedSpringAiVariantDependencies) { + throw new GradleException( + "token-pilot-spring-ai ${variantName} dependencies differ. " + + "Expected ${expectedSpringAiVariantDependencies}, actual ${variantDependencies}." + ) + } + def bootPlatform = variant.dependencies.find { + it.group == 'org.springframework.boot' && + it.module == 'spring-boot-dependencies' + } + def springAiPlatform = variant.dependencies.find { + it.group == 'org.springframework.ai' && it.module == 'spring-ai-bom' + } + if (bootPlatform?.version?.requires != compatibilityMatrix.springBoot || + bootPlatform?.attributes?.get('org.gradle.category') != 'platform' || + springAiPlatform?.version?.requires != compatibilityMatrix.springAi || + springAiPlatform?.attributes?.get('org.gradle.category') != 'platform') { + throw new GradleException( + "token-pilot-spring-ai ${variantName} must carry the supported Boot and Spring AI platforms." + ) + } + } + + def starterPomFile = project(':token-pilot-starter') + .layout.buildDirectory + .file('publications/mavenJava/pom-default.xml') + .get() + .asFile + def starterPom = new XmlSlurper(false, false).parse(starterPomFile) + def starterDependencies = starterPom.dependencies.dependency.collect { + "${it.groupId.text()}:${it.artifactId.text()}:${it.scope.text()}".toString() + } as Set + def expectedStarterDependencies = [ + 'cloud.token-pilot:token-pilot-core:compile', + 'cloud.token-pilot:token-pilot-spring-ai:compile', + 'cloud.token-pilot:token-pilot-micrometer:compile', + 'cloud.token-pilot:token-pilot-budget:compile', + 'cloud.token-pilot:token-pilot-notification:compile', + 'cloud.token-pilot:token-pilot-autoconfigure:compile' + ] as Set + if (starterDependencies != expectedStarterDependencies) { + throw new GradleException( + 'token-pilot-starter POM must expose only the Token Pilot runtime graph at compile scope; ' + + 'applications select their Spring AI provider separately. ' + + "Expected ${expectedStarterDependencies}, actual ${starterDependencies}." + ) + } } } @@ -737,7 +881,588 @@ tasks.register('verifyCoreConsumer', Exec) { commandLine rootProject.file('gradlew').absolutePath, '--no-daemon', 'run' } +def micrometerConsumerDirectory = layout.buildDirectory.dir( + 'compatibility/micrometer-consumer' +) + +tasks.register('prepareMicrometerConsumer') { + group = 'verification' + description = 'Generates external consumers for optional Micrometer owners and core-only autoconfiguration.' + dependsOn ':token-pilot-core:publishMavenJavaPublicationToStagingRepository' + dependsOn ':token-pilot-autoconfigure:publishMavenJavaPublicationToStagingRepository' + dependsOn ':token-pilot-micrometer:publishMavenJavaPublicationToStagingRepository' + dependsOn ':token-pilot-budget:publishMavenJavaPublicationToStagingRepository' + dependsOn ':token-pilot-notification:publishMavenJavaPublicationToStagingRepository' + outputs.dir(micrometerConsumerDirectory) + outputs.upToDateWhen { false } + + doLast { + def consumerDirectory = micrometerConsumerDirectory.get().asFile + delete(consumerDirectory) + + def coreOnlySource = new File( + consumerDirectory, + 'core-only/src/main/java/compatibility' + ) + def allOwnersSource = new File( + consumerDirectory, + 'all-owners/src/main/java/compatibility' + ) + def autoConfigureSource = new File( + consumerDirectory, + 'autoconfigure-core-only/src/main/java/compatibility' + ) + coreOnlySource.mkdirs() + allOwnersSource.mkdirs() + autoConfigureSource.mkdirs() + + new File(consumerDirectory, 'settings.gradle').text = """ +rootProject.name = 'token-pilot-micrometer-consumer' +include 'core-only', 'all-owners', 'autoconfigure-core-only' +""".stripIndent().trim() + System.lineSeparator() + + def stagingRepositories = [ + project(':token-pilot-core'), + project(':token-pilot-autoconfigure'), + project(':token-pilot-micrometer'), + project(':token-pilot-budget'), + project(':token-pilot-notification') + ].collect { moduleProject -> + "maven { url = uri('${moduleProject.layout.buildDirectory.dir('staging-deploy').get().asFile.toURI()}') }" + }.join(System.lineSeparator() + ' ') + + new File(consumerDirectory, 'build.gradle').text = """ +subprojects { + apply plugin: 'application' + + configurations.configureEach { + resolutionStrategy.cacheChangingModulesFor 0, 'seconds' + } + + repositories { + ${stagingRepositories} + mavenCentral() + } + + java { + toolchain { + languageVersion = JavaLanguageVersion.of(${compatibilityMatrix.java}) + } + } + + tasks.withType(JavaCompile).configureEach { + options.release = ${compatibilityMatrix.java} + } +} +""".stripIndent().trim() + System.lineSeparator() + + new File(consumerDirectory, 'core-only/build.gradle').text = """ +dependencies { + implementation 'cloud.token-pilot:token-pilot-micrometer:${project.version}' +} + +application { + mainClass = 'compatibility.MicrometerCoreConsumer' +} + +tasks.register('verifyOptionalOwnersAbsent') { + doLast { + def forbidden = configurations.runtimeClasspath.incoming.resolutionResult + .allComponents + .findAll { component -> + component.moduleVersion?.group == 'cloud.token-pilot' && + component.moduleVersion?.name in [ + 'token-pilot-budget', + 'token-pilot-notification' + ] + } + if (!forbidden.isEmpty()) { + throw new GradleException('Optional owners leaked into micrometer runtime: ' + forbidden) + } + } +} + +tasks.named('run') { + dependsOn tasks.named('verifyOptionalOwnersAbsent') +} +""".stripIndent().trim() + System.lineSeparator() + + new File(coreOnlySource, 'MicrometerCoreConsumer.java').text = """ +package compatibility; + +import io.micrometer.core.instrument.simple.SimpleMeterRegistry; +import io.tokenpilot.core.domain.MissingPricingPolicy; +import io.tokenpilot.core.domain.PricingMissingEvent; +import io.tokenpilot.core.domain.PricingResolution; +import io.tokenpilot.micrometer.internal.CoreMetricsPublisher; +import io.tokenpilot.micrometer.internal.LedgerMicrometerComponents; +import java.lang.reflect.Method; +import java.util.Arrays; + +public final class MicrometerCoreConsumer { + private MicrometerCoreConsumer() { + } + + public static void main(String[] args) throws Exception { + for (Method method : LedgerMicrometerComponents.class.getDeclaredMethods()) { + String signature = method.toGenericString(); + if (signature.contains("io.tokenpilot.budget") + || signature.contains("io.tokenpilot.notification")) { + throw new IllegalStateException("Optional owner leaked into core factory method table"); + } + } + for (String optionalType : Arrays.asList( + "io.tokenpilot.budget.ReservationAccountingListener", + "io.tokenpilot.notification.BudgetNotificationLifecycleListener")) { + try { + Class.forName(optionalType); + throw new IllegalStateException("Optional owner is unexpectedly present: " + optionalType); + } catch (ClassNotFoundException expected) { + // Expected for the core + micrometer-only consumer. + } + } + + var registry = new SimpleMeterRegistry(); + CoreMetricsPublisher publisher = + LedgerMicrometerComponents.coreMetricsPublisher(registry); + publisher.onPricingMissing(new PricingMissingEvent( + MissingPricingPolicy.FAIL_CLOSED, + PricingResolution.MISSING_PLAN + )); + double count = registry.get("tokenpilot.pricing.missing") + .tag("policy", "fail_closed") + .counter() + .count(); + if (count != 1.0d) { + throw new IllegalStateException("Core Micrometer publisher verification failed"); + } + System.out.println("token-pilot-micrometer core-only consumer OK"); + } +} +""".stripIndent().trim() + System.lineSeparator() + + new File(consumerDirectory, 'all-owners/build.gradle').text = """ +dependencies { + implementation 'cloud.token-pilot:token-pilot-micrometer:${project.version}' + implementation 'cloud.token-pilot:token-pilot-budget:${project.version}' + implementation 'cloud.token-pilot:token-pilot-notification:${project.version}' +} + +application { + mainClass = 'compatibility.MicrometerOwnerConsumer' +} +""".stripIndent().trim() + System.lineSeparator() + + new File(allOwnersSource, 'MicrometerOwnerConsumer.java').text = """ +package compatibility; + +import io.micrometer.core.instrument.simple.SimpleMeterRegistry; +import io.tokenpilot.budget.ReservationAccountingListener; +import io.tokenpilot.micrometer.internal.BudgetMicrometerComponents; +import io.tokenpilot.micrometer.internal.NotificationMicrometerComponents; +import io.tokenpilot.notification.BudgetNotificationLifecycleListener; + +public final class MicrometerOwnerConsumer { + private MicrometerOwnerConsumer() { + } + + public static void main(String[] args) { + var registry = new SimpleMeterRegistry(); + Object budget = BudgetMicrometerComponents.budgetMetricsPublisher(registry); + Object notification = NotificationMicrometerComponents + .notificationMetricsPublisher(registry); + if (!(budget instanceof ReservationAccountingListener) + || !(notification instanceof BudgetNotificationLifecycleListener)) { + throw new IllegalStateException("Owner Micrometer composition verification failed"); + } + System.out.println("token-pilot-micrometer all-owners consumer OK"); + } +} +""".stripIndent().trim() + System.lineSeparator() + + new File( + consumerDirectory, + 'autoconfigure-core-only/build.gradle' + ).text = """ +dependencies { + implementation 'cloud.token-pilot:token-pilot-autoconfigure:${project.version}' + implementation 'org.springframework.boot:spring-boot:${compatibilityMatrix.springBoot}' + implementation 'org.springframework.boot:spring-boot-autoconfigure:${compatibilityMatrix.springBoot}' +} + +application { + mainClass = 'compatibility.AutoConfigureCoreConsumer' +} + +tasks.register('verifyOptionalAdaptersAbsent') { + doLast { + def forbidden = configurations.runtimeClasspath.incoming.resolutionResult + .allComponents + .findAll { component -> + component.moduleVersion?.group == 'cloud.token-pilot' && + component.moduleVersion?.name in [ + 'token-pilot-budget', + 'token-pilot-notification', + 'token-pilot-spring-ai', + 'token-pilot-micrometer', + 'token-pilot-starter' + ] + } + if (!forbidden.isEmpty()) { + throw new GradleException('Optional adapters leaked into core autoconfigure runtime: ' + forbidden) + } + } +} + +tasks.named('run') { + dependsOn tasks.named('verifyOptionalAdaptersAbsent') +} +""".stripIndent().trim() + System.lineSeparator() + + new File(autoConfigureSource, 'AutoConfigureCoreConsumer.java').text = """ +package compatibility; + +import io.tokenpilot.core.PricingRegistry; +import io.tokenpilot.core.TokenBudget; +import io.tokenpilot.core.TokenEstimator; +import org.springframework.boot.SpringApplication; +import org.springframework.boot.WebApplicationType; +import org.springframework.boot.autoconfigure.EnableAutoConfiguration; +import org.springframework.boot.SpringBootConfiguration; +import org.springframework.context.ConfigurableApplicationContext; + +@SpringBootConfiguration +@EnableAutoConfiguration +public class AutoConfigureCoreConsumer { + public static void main(String[] args) { + SpringApplication application = new SpringApplication( + AutoConfigureCoreConsumer.class + ); + application.setWebApplicationType(WebApplicationType.NONE); + application.setLogStartupInfo(false); + try (ConfigurableApplicationContext context = application.run(args)) { + context.getBean(PricingRegistry.class); + context.getBean(TokenBudget.class); + context.getBean(TokenEstimator.class); + } + System.out.println("token-pilot-autoconfigure core-only consumer OK"); + } +} +""".stripIndent().trim() + System.lineSeparator() + } +} + +tasks.register('verifyMicrometerConsumer', Exec) { + group = 'verification' + description = 'Runs Micrometer owner consumers and core-only autoconfiguration without optional adapters.' + dependsOn tasks.named('prepareMicrometerConsumer') + workingDir micrometerConsumerDirectory + commandLine rootProject.file('gradlew').absolutePath, + '--no-daemon', + ':core-only:run', + ':all-owners:run', + ':autoconfigure-core-only:run' +} + +def publishedIntegrationConsumerDirectory = layout.buildDirectory.dir( + 'compatibility/published-integration-consumer' +) + +tasks.register('preparePublishedIntegrationConsumer') { + group = 'verification' + description = 'Generates external compile/runtime consumers for the published Spring AI adapter and starter.' + dependsOn publishedModuleDescriptions.keySet().collect { + ":${it}:publishMavenJavaPublicationToStagingRepository" + } + outputs.dir(publishedIntegrationConsumerDirectory) + outputs.upToDateWhen { false } + + doLast { + def consumerDirectory = publishedIntegrationConsumerDirectory.get().asFile + delete(consumerDirectory) + + def adapterSource = new File( + consumerDirectory, + 'adapter/src/main/java/compatibility' + ) + def starterSource = new File( + consumerDirectory, + 'starter/src/main/java/compatibility' + ) + def adapterPomSource = new File( + consumerDirectory, + 'adapter-pom/src/main/java/compatibility' + ) + def starterPomSource = new File( + consumerDirectory, + 'starter-pom/src/main/java/compatibility' + ) + adapterSource.mkdirs() + starterSource.mkdirs() + adapterPomSource.mkdirs() + starterPomSource.mkdirs() + + new File(consumerDirectory, 'settings.gradle').text = """ +rootProject.name = 'token-pilot-published-integration-consumer' +include 'adapter', 'starter', 'adapter-pom', 'starter-pom' +""".stripIndent().trim() + System.lineSeparator() + + def stagingRepositories = publishedModuleDescriptions.keySet().collect { moduleName -> + def moduleProject = project(":${moduleName}") + "maven { url = uri('${moduleProject.layout.buildDirectory.dir('staging-deploy').get().asFile.toURI()}') }" + }.join(System.lineSeparator() + ' ') + def pomOnlyStagingRepositories = publishedModuleDescriptions.keySet().collect { moduleName -> + def moduleProject = project(":${moduleName}") + """maven { + url = uri('${moduleProject.layout.buildDirectory.dir('staging-deploy').get().asFile.toURI()}') + metadataSources { + mavenPom() + artifact() + } + }""".stripIndent().trim() + }.join(System.lineSeparator() + ' ') + + new File(consumerDirectory, 'build.gradle').text = """ +subprojects { + apply plugin: 'application' + + configurations.configureEach { + resolutionStrategy.cacheChangingModulesFor 0, 'seconds' + } + + java { + toolchain { + languageVersion = JavaLanguageVersion.of(${compatibilityMatrix.java}) + } + } + + tasks.withType(JavaCompile).configureEach { + options.release = ${compatibilityMatrix.java} + } +} +""".stripIndent().trim() + System.lineSeparator() + + new File(consumerDirectory, 'adapter/build.gradle').text = """ +repositories { + ${stagingRepositories} + mavenCentral() +} + +dependencies { + implementation 'cloud.token-pilot:token-pilot-spring-ai:${project.version}' +} + +application { + mainClass = 'compatibility.SpringAiAdapterConsumer' +} +""".stripIndent().trim() + System.lineSeparator() + + new File(consumerDirectory, 'adapter-pom/build.gradle').text = """ +repositories { + ${pomOnlyStagingRepositories} + mavenCentral { + metadataSources { + mavenPom() + artifact() + } + } +} + +dependencies { + implementation 'cloud.token-pilot:token-pilot-spring-ai:${project.version}' +} + +application { + mainClass = 'compatibility.SpringAiAdapterConsumer' +} +""".stripIndent().trim() + System.lineSeparator() + + def adapterConsumerSource = """ +package compatibility; + +import reactor.core.publisher.Flux; +import io.tokenpilot.budget.AtomicBudgetStateStore; +import io.tokenpilot.budget.BudgetStateStore; +import io.tokenpilot.budget.ReservationAccounting; +import io.tokenpilot.core.domain.TokenUsage; +import io.tokenpilot.springai.LedgerAdvisor; +import io.tokenpilot.springai.ReservedOutputTokensResolver; +import io.tokenpilot.springai.UsageExtractor; +import io.tokenpilot.springai.internal.LedgerSpringAiComponents; +import java.util.OptionalLong; +import org.springframework.ai.chat.client.ChatClientRequest; +import org.springframework.ai.chat.client.ChatClientResponse; +import org.springframework.ai.chat.client.advisor.api.BaseAdvisor; + +public final class SpringAiAdapterConsumer { + private SpringAiAdapterConsumer() { + } + + public static void main(String[] args) { + UsageExtractor extractor = ignored -> TokenUsage.from(1, 1); + ReservedOutputTokensResolver resolver = ignored -> OptionalLong.of(128L); + Object defaultExtractor = LedgerSpringAiComponents.defaultUsageExtractor(); + + Class[] publicApiTypes = { + LedgerAdvisor.class, + ChatClientRequest.class, + ChatClientResponse.class, + BaseAdvisor.class, + Flux.class, + BudgetStateStore.class, + ReservationAccounting.class, + AtomicBudgetStateStore.class + }; + if (!BaseAdvisor.class.isAssignableFrom(LedgerAdvisor.class) + || extractor == null + || resolver == null + || defaultExtractor == null + || publicApiTypes.length != 8) { + throw new IllegalStateException("Published Spring AI adapter API verification failed"); + } + System.out.println("token-pilot-spring-ai published consumer OK"); + } +} +""".stripIndent().trim() + System.lineSeparator() + new File(adapterSource, 'SpringAiAdapterConsumer.java').text = adapterConsumerSource + new File(adapterPomSource, 'SpringAiAdapterConsumer.java').text = adapterConsumerSource + + def starterBuildBody = """ +configurations { + starterOnly { + canBeConsumed = false + canBeResolved = true + } +} + +dependencies { + // Token Pilot remains one user entrypoint; the application selects one provider. + implementation 'cloud.token-pilot:token-pilot-starter:${project.version}' + implementation 'org.springframework.ai:spring-ai-starter-model-openai:${compatibilityMatrix.springAi}' + + starterOnly 'cloud.token-pilot:token-pilot-starter:${project.version}' +} + +application { + mainClass = 'compatibility.StarterConsumer' +} + +tasks.register('verifyStarterDoesNotBundleProvider') { + doLast { + def providerArtifacts = configurations.starterOnly.incoming.resolutionResult + .allComponents + .findAll { component -> + component.moduleVersion?.group == 'org.springframework.ai' && + (component.moduleVersion.name == 'spring-ai-openai' || + component.moduleVersion.name.startsWith('spring-ai-starter-model-')) + } + if (!providerArtifacts.isEmpty()) { + throw new GradleException( + 'token-pilot-starter must not choose a Spring AI provider: ' + providerArtifacts + ) + } + } +} + +tasks.named('run') { + dependsOn tasks.named('verifyStarterDoesNotBundleProvider') +} +""".stripIndent().trim() + System.lineSeparator() + + new File(consumerDirectory, 'starter/build.gradle').text = """ +repositories { + ${stagingRepositories} + mavenCentral() +} + +""".stripIndent() + starterBuildBody + new File(consumerDirectory, 'starter-pom/build.gradle').text = """ +repositories { + ${pomOnlyStagingRepositories} + mavenCentral { + metadataSources { + mavenPom() + artifact() + } + } +} + +""".stripIndent() + starterBuildBody + + def starterConsumerSource = """ +package compatibility; + +import io.tokenpilot.autoconfigure.LedgerChatClientBuilderCustomizer; +import io.tokenpilot.budget.AtomicBudgetStateStore; +import io.tokenpilot.budget.BudgetStateStore; +import io.tokenpilot.budget.ReservationAccounting; +import io.tokenpilot.notification.BudgetNotificationHandler; +import io.tokenpilot.notification.BudgetNotificationService; +import io.tokenpilot.springai.LedgerAdvisor; +import java.util.Map; +import org.springframework.ai.openai.OpenAiChatModel; +import org.springframework.boot.SpringApplication; +import org.springframework.boot.WebApplicationType; +import org.springframework.boot.autoconfigure.EnableAutoConfiguration; +import org.springframework.boot.SpringBootConfiguration; +import org.springframework.context.ConfigurableApplicationContext; +import org.springframework.context.annotation.Bean; + +@SpringBootConfiguration +@EnableAutoConfiguration +public class StarterConsumer { + @Bean + BudgetNotificationHandler budgetNotificationHandler() { + return event -> { }; + } + + public static void main(String[] args) { + SpringApplication application = new SpringApplication(StarterConsumer.class); + application.setWebApplicationType(WebApplicationType.NONE); + application.setLogStartupInfo(false); + application.setDefaultProperties(Map.of( + "spring.ai.openai.api-key", "published-consumer-test-key", + "token-pilot.budget.enabled", "true", + "token-pilot.notification.enabled", "true", + "token-pilot.metrics.enabled", "false" + )); + + try (ConfigurableApplicationContext context = application.run(args)) { + AtomicBudgetStateStore atomicStore = context.getBean(AtomicBudgetStateStore.class); + if (atomicStore != context.getBean(BudgetStateStore.class) + || atomicStore != context.getBean(ReservationAccounting.class)) { + throw new IllegalStateException("Budget admission/accounting bean identity was split"); + } + context.getBean(LedgerAdvisor.class); + context.getBean(LedgerChatClientBuilderCustomizer.class); + context.getBean(BudgetNotificationService.class); + context.getBean(OpenAiChatModel.class); + } + System.out.println("token-pilot-starter published consumer OK"); + } +} +""".stripIndent().trim() + System.lineSeparator() + new File(starterSource, 'StarterConsumer.java').text = starterConsumerSource + new File(starterPomSource, 'StarterConsumer.java').text = starterConsumerSource + } +} + +tasks.register('verifyPublishedIntegrationConsumer', Exec) { + group = 'verification' + description = 'Compiles and runs published Spring AI adapter and starter consumers.' + dependsOn tasks.named('preparePublishedIntegrationConsumer') + workingDir publishedIntegrationConsumerDirectory + commandLine rootProject.file('gradlew').absolutePath, + '--no-daemon', + ':adapter:run', + ':starter:run', + ':adapter-pom:run', + ':starter-pom:run' +} + tasks.named('check') { dependsOn tasks.named('verifyCompatibilityMatrix') dependsOn tasks.named('verifyCoreConsumer') + dependsOn tasks.named('verifyMicrometerConsumer') + dependsOn tasks.named('verifyPublishedIntegrationConsumer') } 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..6225c79 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 @@ -1,52 +1,36 @@ package io.tokenpilot.autoconfigure; -import io.micrometer.core.instrument.MeterRegistry; -import io.tokenpilot.budget.BudgetEvaluator; -import io.tokenpilot.budget.BudgetStateStore; -import io.tokenpilot.budget.internal.LedgerBudgetComponents; +import io.tokenpilot.core.CoreComponents; import io.tokenpilot.core.CostCalculator; import io.tokenpilot.core.LedgerListener; import io.tokenpilot.core.LedgerManager; +import io.tokenpilot.core.ModelRegistry; +import io.tokenpilot.core.PreflightCostEstimator; +import io.tokenpilot.core.PreflightDecisionListener; import io.tokenpilot.core.PricingEvaluator; import io.tokenpilot.core.PricingProvider; import io.tokenpilot.core.PricingRegistry; -import io.tokenpilot.core.domain.MissingPricingPolicy; +import io.tokenpilot.core.TokenBudget; +import io.tokenpilot.core.TokenEstimator; import io.tokenpilot.core.internal.LedgerComponents; -import io.tokenpilot.micrometer.internal.LedgerMicrometerComponents; -import io.tokenpilot.springai.LedgerAdvisor; -import io.tokenpilot.springai.UsageExtractor; -import io.tokenpilot.springai.internal.LedgerSpringAiComponents; -import org.springframework.ai.chat.client.ChatClient; import org.springframework.beans.factory.ObjectProvider; import org.springframework.boot.autoconfigure.AutoConfiguration; -import org.springframework.boot.autoconfigure.condition.ConditionalOnBean; -import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean; 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; - -/** - * Token Pilot 라이브러리의 자동 설정을 담당하는 클래스. - */ -@AutoConfiguration(afterName = { - "org.springframework.boot.micrometer.metrics.autoconfigure.MetricsAutoConfiguration", - "org.springframework.boot.micrometer.metrics.autoconfigure.CompositeMeterRegistryAutoConfiguration", - "org.springframework.boot.micrometer.metrics.autoconfigure.export.prometheus.PrometheusMetricsExportAutoConfiguration" -}) -@ConditionalOnProperty(prefix = "token-pilot", name = "enabled", havingValue = "true", matchIfMissing = true) +/** Token Pilot의 framework-independent core 자동 설정입니다. */ +@AutoConfiguration +@ConditionalOnProperty( + prefix = "token-pilot", + name = "enabled", + havingValue = "true", + matchIfMissing = true +) @EnableConfigurationProperties(TokenPilotProperties.class) public class TokenPilotAutoConfiguration { - /** - * 외부 설정으로부터 가격 정책을 읽어오는 PricingProvider를 등록합니다. - */ @Bean @ConditionalOnMissingBean public PricingProvider pricingProvider(TokenPilotProperties properties) { @@ -54,177 +38,69 @@ public PricingProvider pricingProvider(TokenPilotProperties properties) { return () -> plans; } - /** - * 메모리 기반의 가격 정책 저장소(PricingRegistry)를 등록합니다. - */ @Bean @ConditionalOnMissingBean - public PricingRegistry pricingRegistry(ObjectProvider pricingProviders) { - return LedgerComponents.inMemoryPricingRegistry(pricingProviders.orderedStream().toList()); + public PricingRegistry pricingRegistry( + ObjectProvider pricingProviders + ) { + return LedgerComponents.inMemoryPricingRegistry( + pricingProviders.orderedStream().toList() + ); } - /** - * 토큰 사용량을 바탕으로 비용을 계산하는 CostCalculator를 등록합니다. - */ @Bean @ConditionalOnMissingBean public CostCalculator costCalculator() { return LedgerComponents.defaultCostCalculator(); } - /** - * Pricing snapshot rate와 actual model 정합성을 평가하는 정책을 등록합니다. - */ @Bean @ConditionalOnMissingBean public PricingEvaluator pricingEvaluator() { return LedgerComponents.defaultPricingEvaluator(); } - /** - * 비용 기록 및 리스너 관리를 담당하는 LedgerManager를 등록합니다. - */ @Bean @ConditionalOnMissingBean - public LedgerManager ledgerManager(PricingRegistry pricingRegistry, - CostCalculator costCalculator, - ObjectProvider ledgerListeners) { - return LedgerComponents.defaultLedgerManager( - pricingRegistry, - costCalculator, - ledgerListeners.orderedStream() - .toList() - ); + public ModelRegistry modelRegistry() { + return CoreComponents.defaultModelRegistry(); } - /** - * ChatClientResponse에서 토큰 사용량을 추출하는 UsageExtractor를 등록합니다. - */ @Bean @ConditionalOnMissingBean - @ConditionalOnClass({ChatClient.class, LedgerSpringAiComponents.class}) - public UsageExtractor usageExtractor() { - return LedgerSpringAiComponents.defaultUsageExtractor(); + public TokenEstimator tokenEstimator() { + return CoreComponents.utf8ByteHeuristicTokenEstimator(); } - /** - * Spring AI 호출 전후로 비용을 측정하고 기록하는 LedgerAdvisor를 등록합니다. - */ @Bean @ConditionalOnMissingBean - @ConditionalOnClass({ChatClient.class, LedgerSpringAiComponents.class}) - public LedgerAdvisor ledgerAdvisor( - LedgerManager ledgerManager, - UsageExtractor usageExtractor, - ObjectProvider budgetEvaluator, - ObjectProvider budgetStateStore, - CostCalculator costCalculator, - PricingRegistry pricingRegistry, - PricingEvaluator pricingEvaluator + public TokenBudget tokenBudget( + ModelRegistry modelRegistry, + ObjectProvider decisionListeners ) { - BudgetEvaluator evaluator = budgetEvaluator.getIfAvailable(); - BudgetStateStore stateStore = budgetStateStore.getIfAvailable(); - - if (evaluator != null && stateStore != null) { - return LedgerSpringAiComponents.defaultLedgerAdvisor( - ledgerManager, - usageExtractor, - evaluator, - stateStore, - costCalculator, - pricingRegistry, - pricingEvaluator, - MissingPricingPolicy.FAIL_CLOSED - ); - } - - return LedgerSpringAiComponents.defaultLedgerAdvisor( - ledgerManager, - usageExtractor, - costCalculator, - pricingRegistry, - pricingEvaluator + return CoreComponents.tokenBudget( + modelRegistry, + decisionListeners.orderedStream().toList() ); } - /** - * LedgerAdvisor가 빈으로 등록되어 있을 경우, ChatClient.Builder를 위한 커스터마이저를 생성합니다. - */ - @Bean - @ConditionalOnBean(LedgerAdvisor.class) - @ConditionalOnClass(ChatClient.class) - public LedgerChatClientBuilderCustomizer ledgerChatClientBuilderCustomizer(LedgerAdvisor ledgerAdvisor) { - return new LedgerChatClientBuilderCustomizer(ledgerAdvisor); - } - - /** - * Micrometer 메트릭 발행을 위한 리스너를 등록합니다. - */ - @Bean - @ConditionalOnClass({MeterRegistry.class, LedgerMicrometerComponents.class}) - @ConditionalOnBean(MeterRegistry.class) - @ConditionalOnProperty(prefix = "token-pilot.metrics", name = "enabled", havingValue = "true", matchIfMissing = true) - public LedgerListener microCostMetricsPublisher(MeterRegistry meterRegistry, TokenPilotProperties properties) { - return LedgerMicrometerComponents.microCostMetricsPublisher(meterRegistry, properties.getMetrics() - .getTagWhitelist()); - } - - /** - * 예산 관리 모듈이 있을 경우 기본 예산 저장소를 등록합니다. - */ @Bean @ConditionalOnMissingBean - @ConditionalOnClass(LedgerBudgetComponents.class) - @ConditionalOnProperty(prefix = "token-pilot.budget", name = "enabled", havingValue = "true") - public BudgetStateStore budgetStateStore() { - return LedgerBudgetComponents.inMemoryBudgetStateStore(); + public PreflightCostEstimator preflightCostEstimator() { + return CoreComponents.defaultPreflightCostEstimator(); } - /** - * 예산 관리 모듈이 있을 경우 예산 평가기를 등록합니다. - */ @Bean @ConditionalOnMissingBean - @ConditionalOnClass(LedgerBudgetComponents.class) - @ConditionalOnProperty(prefix = "token-pilot.budget", name = "enabled", havingValue = "true") - public BudgetEvaluator budgetEvaluator( - BudgetStateStore budgetStateStore, - TokenPilotProperties properties, - ObjectProvider clock + public LedgerManager ledgerManager( + PricingRegistry pricingRegistry, + CostCalculator costCalculator, + ObjectProvider ledgerListeners ) { - return LedgerBudgetComponents.defaultBudgetEvaluator( - budgetStateStore, - properties.toBudgetPolicy(), - clock.getIfAvailable(Clock::systemUTC) + return LedgerComponents.defaultLedgerManager( + pricingRegistry, + costCalculator, + ledgerListeners.orderedStream().toList() ); } - - /** - * 알림 상태 저장소를 등록합니다. - * - 중복 알림 방지를 위해 window 단위로 상태를 저장 - * - token-pilot.notification.enabled=true 일 때만 등록 - */ - @Bean - @ConditionalOnMissingBean - @ConditionalOnProperty(prefix = "token-pilot.notification", name = "enabled", havingValue = "true") - public NotificationStateStore notificationStateStore() { - return new InMemoryNotificationStateStore(); - } - - /** - * 예산 알림 서비스를 등록합니다. - * - BudgetNotificationHandler 빈이 있을 때만 등록 - * - 없으면 no-op으로 동작 (알림 서비스 자체가 등록되지 않음) - * - token-pilot.notification.enabled=true 일 때만 등록 - */ - @Bean - @ConditionalOnMissingBean - @ConditionalOnBean(BudgetNotificationHandler.class) - @ConditionalOnProperty(prefix = "token-pilot.notification", name = "enabled", havingValue = "true") - public BudgetNotificationService budgetNotificationService( - BudgetNotificationHandler handler, - NotificationStateStore notificationStateStore - ) { - return new BudgetNotificationService(handler, notificationStateStore); - } } diff --git a/token-pilot-autoconfigure/src/main/java/io/tokenpilot/autoconfigure/TokenPilotBudgetAutoConfiguration.java b/token-pilot-autoconfigure/src/main/java/io/tokenpilot/autoconfigure/TokenPilotBudgetAutoConfiguration.java new file mode 100644 index 0000000..3b914d7 --- /dev/null +++ b/token-pilot-autoconfigure/src/main/java/io/tokenpilot/autoconfigure/TokenPilotBudgetAutoConfiguration.java @@ -0,0 +1,62 @@ +package io.tokenpilot.autoconfigure; + +import io.tokenpilot.budget.BudgetEvaluator; +import io.tokenpilot.budget.AtomicBudgetStateStore; +import io.tokenpilot.budget.BudgetStateStore; +import io.tokenpilot.budget.ReservationAccountingListener; +import io.tokenpilot.budget.ReservationAccountingListenerErrorHandler; +import io.tokenpilot.budget.internal.LedgerBudgetComponents; +import org.springframework.beans.factory.ObjectProvider; +import org.springframework.boot.autoconfigure.AutoConfiguration; +import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; +import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean; +import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; +import org.springframework.context.annotation.Bean; + +import java.time.Clock; + +/** Optional budget module의 bean graph를 격리하는 자동 설정입니다. */ +@AutoConfiguration(after = { + TokenPilotAutoConfiguration.class, + TokenPilotBudgetMetricsAutoConfiguration.class +}) +@ConditionalOnClass(LedgerBudgetComponents.class) +@ConditionalOnProperty( + prefix = "token-pilot", + name = "enabled", + havingValue = "true", + matchIfMissing = true +) +@ConditionalOnProperty( + prefix = "token-pilot.budget", + name = "enabled", + havingValue = "true" +) +public class TokenPilotBudgetAutoConfiguration { + + @Bean + @ConditionalOnMissingBean(BudgetStateStore.class) + public AtomicBudgetStateStore budgetStateStore( + ObjectProvider accountingListeners, + ObjectProvider listenerErrorHandlers + ) { + return LedgerBudgetComponents.inMemoryAtomicBudgetStateStore( + () -> accountingListeners.orderedStream().toList(), + () -> listenerErrorHandlers.orderedStream().toList() + ); + } + + @Bean + @ConditionalOnMissingBean + public BudgetEvaluator budgetEvaluator( + BudgetStateStore budgetStateStore, + TokenPilotProperties properties, + ObjectProvider clock + ) { + return LedgerBudgetComponents.defaultBudgetEvaluator( + budgetStateStore, + TokenPilotBudgetPolicyFactory.from(properties), + clock.getIfAvailable(Clock::systemUTC) + ); + } +} diff --git a/token-pilot-autoconfigure/src/main/java/io/tokenpilot/autoconfigure/TokenPilotBudgetMetricsAutoConfiguration.java b/token-pilot-autoconfigure/src/main/java/io/tokenpilot/autoconfigure/TokenPilotBudgetMetricsAutoConfiguration.java new file mode 100644 index 0000000..28a5726 --- /dev/null +++ b/token-pilot-autoconfigure/src/main/java/io/tokenpilot/autoconfigure/TokenPilotBudgetMetricsAutoConfiguration.java @@ -0,0 +1,43 @@ +package io.tokenpilot.autoconfigure; + +import io.micrometer.core.instrument.MeterRegistry; +import io.tokenpilot.budget.ReservationAccountingListener; +import io.tokenpilot.micrometer.internal.BudgetMicrometerComponents; +import io.tokenpilot.micrometer.internal.BudgetMetricsPublisher; +import org.springframework.boot.autoconfigure.AutoConfiguration; +import org.springframework.boot.autoconfigure.condition.ConditionalOnBean; +import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; +import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean; +import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; +import org.springframework.context.annotation.Bean; + +/** Atomic reservation/accounting metrics의 조건부 자동 설정입니다. */ +@AutoConfiguration(after = TokenPilotCoreMetricsAutoConfiguration.class) +@ConditionalOnClass({ + MeterRegistry.class, + ReservationAccountingListener.class, + BudgetMetricsPublisher.class +}) +@ConditionalOnBean(MeterRegistry.class) +@ConditionalOnProperty( + prefix = "token-pilot", + name = "enabled", + havingValue = "true", + matchIfMissing = true +) +@ConditionalOnProperty( + prefix = "token-pilot.metrics", + name = "enabled", + havingValue = "true", + matchIfMissing = true +) +public class TokenPilotBudgetMetricsAutoConfiguration { + + @Bean + @ConditionalOnMissingBean + public BudgetMetricsPublisher tokenPilotBudgetMetricsPublisher( + MeterRegistry meterRegistry + ) { + return BudgetMicrometerComponents.budgetMetricsPublisher(meterRegistry); + } +} diff --git a/token-pilot-autoconfigure/src/main/java/io/tokenpilot/autoconfigure/TokenPilotBudgetPolicyFactory.java b/token-pilot-autoconfigure/src/main/java/io/tokenpilot/autoconfigure/TokenPilotBudgetPolicyFactory.java new file mode 100644 index 0000000..0e63186 --- /dev/null +++ b/token-pilot-autoconfigure/src/main/java/io/tokenpilot/autoconfigure/TokenPilotBudgetPolicyFactory.java @@ -0,0 +1,40 @@ +package io.tokenpilot.autoconfigure; + +import io.tokenpilot.budget.BudgetPolicy; +import io.tokenpilot.core.domain.Cost; + +import java.time.ZoneId; +import java.util.Currency; +import java.util.Objects; + +/** + * Token Pilot 설정을 optional budget 정책으로 변환하는 공개 진입점입니다. + * + *

budget 타입을 {@link TokenPilotProperties}에서 분리해 budget 모듈이 없는 + * core-only 자동 설정도 class loading에 실패하지 않도록 합니다.

+ */ +public final class TokenPilotBudgetPolicyFactory { + + private TokenPilotBudgetPolicyFactory() { + } + + public static BudgetPolicy from(TokenPilotProperties properties) { + Objects.requireNonNull(properties, "properties must not be null"); + return from(properties.getBudget()); + } + + public static BudgetPolicy from(TokenPilotProperties.BudgetProperties properties) { + Objects.requireNonNull(properties, "properties must not be null"); + return new BudgetPolicy( + properties.getPolicyId(), + properties.getTargetType(), + properties.getTargetTagKey(), + properties.getFallbackTargetId(), + Cost.of( + properties.getMonthlyLimit(), + Currency.getInstance(properties.getCurrency()) + ), + ZoneId.of(properties.getZoneId()) + ); + } +} diff --git a/token-pilot-autoconfigure/src/main/java/io/tokenpilot/autoconfigure/TokenPilotCoreMetricsAutoConfiguration.java b/token-pilot-autoconfigure/src/main/java/io/tokenpilot/autoconfigure/TokenPilotCoreMetricsAutoConfiguration.java new file mode 100644 index 0000000..5d09d94 --- /dev/null +++ b/token-pilot-autoconfigure/src/main/java/io/tokenpilot/autoconfigure/TokenPilotCoreMetricsAutoConfiguration.java @@ -0,0 +1,65 @@ +package io.tokenpilot.autoconfigure; + +import io.micrometer.core.instrument.MeterRegistry; +import io.tokenpilot.micrometer.MetricsOptions; +import io.tokenpilot.micrometer.internal.CoreMetricsPublisher; +import io.tokenpilot.micrometer.internal.LedgerMicrometerComponents; +import io.tokenpilot.micrometer.internal.MicroCostMetricsPublisher; +import org.springframework.boot.autoconfigure.AutoConfiguration; +import org.springframework.boot.autoconfigure.condition.ConditionalOnBean; +import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; +import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean; +import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; +import org.springframework.boot.context.properties.EnableConfigurationProperties; +import org.springframework.context.annotation.Bean; + +/** Core control events와 legacy ledger metrics의 조건부 자동 설정입니다. */ +@AutoConfiguration(after = TokenPilotAutoConfiguration.class, afterName = { + "org.springframework.boot.micrometer.metrics.autoconfigure.MetricsAutoConfiguration", + "org.springframework.boot.micrometer.metrics.autoconfigure.CompositeMeterRegistryAutoConfiguration" +}) +@ConditionalOnClass({MeterRegistry.class, CoreMetricsPublisher.class}) +@ConditionalOnBean(MeterRegistry.class) +@ConditionalOnProperty( + prefix = "token-pilot", + name = "enabled", + havingValue = "true", + matchIfMissing = true +) +@ConditionalOnProperty( + prefix = "token-pilot.metrics", + name = "enabled", + havingValue = "true", + matchIfMissing = true +) +@EnableConfigurationProperties(TokenPilotProperties.class) +public class TokenPilotCoreMetricsAutoConfiguration { + + @Bean + @ConditionalOnMissingBean + public CoreMetricsPublisher tokenPilotCoreMetricsPublisher( + MeterRegistry meterRegistry + ) { + return LedgerMicrometerComponents.coreMetricsPublisher(meterRegistry); + } + + @Bean(name = "microCostMetricsPublisher") + @ConditionalOnMissingBean(MicroCostMetricsPublisher.class) + @ConditionalOnProperty( + prefix = "token-pilot.metrics", + name = "legacy-ai-token-metrics-enabled", + havingValue = "true" + ) + public MicroCostMetricsPublisher microCostMetricsPublisher( + MeterRegistry meterRegistry, + TokenPilotProperties properties + ) { + return LedgerMicrometerComponents.legacyMicroCostMetricsPublisher( + meterRegistry, + new MetricsOptions( + properties.getMetrics().getTagWhitelist(), + true + ) + ); + } +} diff --git a/token-pilot-autoconfigure/src/main/java/io/tokenpilot/autoconfigure/TokenPilotNotificationAutoConfiguration.java b/token-pilot-autoconfigure/src/main/java/io/tokenpilot/autoconfigure/TokenPilotNotificationAutoConfiguration.java new file mode 100644 index 0000000..826ceaf --- /dev/null +++ b/token-pilot-autoconfigure/src/main/java/io/tokenpilot/autoconfigure/TokenPilotNotificationAutoConfiguration.java @@ -0,0 +1,69 @@ +package io.tokenpilot.autoconfigure; + +import io.tokenpilot.budget.BudgetStateStore; +import io.tokenpilot.notification.AtomicNotificationStateStore; +import io.tokenpilot.notification.BudgetNotificationErrorHook; +import io.tokenpilot.notification.BudgetNotificationHandler; +import io.tokenpilot.notification.BudgetNotificationLifecycleListener; +import io.tokenpilot.notification.BudgetNotificationService; +import io.tokenpilot.notification.InMemoryNotificationStateStore; +import io.tokenpilot.notification.NotificationStateStore; +import org.springframework.beans.factory.ObjectProvider; +import org.springframework.boot.autoconfigure.AutoConfiguration; +import org.springframework.boot.autoconfigure.condition.ConditionalOnBean; +import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; +import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean; +import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; +import org.springframework.context.annotation.Bean; + +/** Optional notification module의 bean graph를 격리하는 자동 설정입니다. */ +@AutoConfiguration(after = { + TokenPilotAutoConfiguration.class, + TokenPilotNotificationMetricsAutoConfiguration.class +}) +@ConditionalOnClass(NotificationStateStore.class) +@ConditionalOnProperty( + prefix = "token-pilot", + name = "enabled", + havingValue = "true", + matchIfMissing = true +) +@ConditionalOnProperty( + prefix = "token-pilot.notification", + name = "enabled", + havingValue = "true" +) +public class TokenPilotNotificationAutoConfiguration { + + @Bean + @ConditionalOnMissingBean(NotificationStateStore.class) + public NotificationStateStore notificationStateStore() { + return new InMemoryNotificationStateStore(); + } + + @Bean + @ConditionalOnMissingBean + @ConditionalOnBean({BudgetNotificationHandler.class, BudgetStateStore.class}) + public BudgetNotificationService budgetNotificationService( + ObjectProvider handlers, + NotificationStateStore notificationStateStore, + BudgetStateStore budgetStateStore, + TokenPilotProperties properties, + ObjectProvider errorHook, + ObjectProvider lifecycleListeners + ) { + if (!(notificationStateStore instanceof AtomicNotificationStateStore atomicStore)) { + throw new IllegalStateException( + "notification state store must implement AtomicNotificationStateStore" + ); + } + var policy = TokenPilotBudgetPolicyFactory.from(properties); + return new BudgetNotificationService( + handlers.orderedStream().toList(), + atomicStore, + key -> budgetStateStore.snapshot(key, policy.monthlyLimit()), + errorHook.getIfAvailable(BudgetNotificationErrorHook::noOp), + lifecycleListeners.orderedStream().toList() + ); + } +} diff --git a/token-pilot-autoconfigure/src/main/java/io/tokenpilot/autoconfigure/TokenPilotNotificationMetricsAutoConfiguration.java b/token-pilot-autoconfigure/src/main/java/io/tokenpilot/autoconfigure/TokenPilotNotificationMetricsAutoConfiguration.java new file mode 100644 index 0000000..0b39766 --- /dev/null +++ b/token-pilot-autoconfigure/src/main/java/io/tokenpilot/autoconfigure/TokenPilotNotificationMetricsAutoConfiguration.java @@ -0,0 +1,45 @@ +package io.tokenpilot.autoconfigure; + +import io.micrometer.core.instrument.MeterRegistry; +import io.tokenpilot.micrometer.internal.NotificationMetricsPublisher; +import io.tokenpilot.micrometer.internal.NotificationMicrometerComponents; +import io.tokenpilot.notification.BudgetNotificationLifecycleListener; +import org.springframework.boot.autoconfigure.AutoConfiguration; +import org.springframework.boot.autoconfigure.condition.ConditionalOnBean; +import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; +import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean; +import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; +import org.springframework.context.annotation.Bean; + +/** Budget notification lifecycle metrics의 조건부 자동 설정입니다. */ +@AutoConfiguration(after = TokenPilotCoreMetricsAutoConfiguration.class) +@ConditionalOnClass({ + MeterRegistry.class, + BudgetNotificationLifecycleListener.class, + NotificationMetricsPublisher.class +}) +@ConditionalOnBean(MeterRegistry.class) +@ConditionalOnProperty( + prefix = "token-pilot", + name = "enabled", + havingValue = "true", + matchIfMissing = true +) +@ConditionalOnProperty( + prefix = "token-pilot.metrics", + name = "enabled", + havingValue = "true", + matchIfMissing = true +) +public class TokenPilotNotificationMetricsAutoConfiguration { + + @Bean + @ConditionalOnMissingBean + public NotificationMetricsPublisher tokenPilotNotificationMetricsPublisher( + MeterRegistry meterRegistry + ) { + return NotificationMicrometerComponents.notificationMetricsPublisher( + meterRegistry + ); + } +} 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 44a7209..0d512f9 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,14 +1,10 @@ package io.tokenpilot.autoconfigure; -import io.tokenpilot.budget.BudgetPolicy; -import io.tokenpilot.core.domain.Cost; 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; @@ -33,6 +29,9 @@ public class TokenPilotProperties { @NestedConfigurationProperty private NotificationProperties notification = new NotificationProperties(); + @NestedConfigurationProperty + private SpringAiProperties springAi = new SpringAiProperties(); + public boolean isEnabled() { return enabled; } @@ -73,6 +72,14 @@ public void setNotification(NotificationProperties notification) { this.notification = notification; } + public SpringAiProperties getSpringAi() { + return springAi; + } + + public void setSpringAi(SpringAiProperties springAi) { + this.springAi = springAi; + } + public List toPricingPlans() { if (pricing == null || pricing.getPlans() == null) { return List.of(); @@ -84,20 +91,6 @@ public List toPricingPlans() { .toList(); } - public BudgetPolicy toBudgetPolicy() { - return new BudgetPolicy( - budget.getPolicyId(), - budget.getTargetType(), - budget.getTargetTagKey(), - budget.getFallbackTargetId(), - Cost.of( - budget.getMonthlyLimit(), - Currency.getInstance(budget.getCurrency()) - ), - ZoneId.of(budget.getZoneId()) - ); - } - public static class PricingProperties { private List plans = new ArrayList<>(); @@ -112,7 +105,8 @@ public void setPlans(List plans) { public static class MetricsProperties { private boolean enabled = true; - private Set tagWhitelist = new HashSet<>(List.of("tenant_id")); + private Set tagWhitelist = new HashSet<>(); + private boolean legacyAiTokenMetricsEnabled = false; public boolean isEnabled() { return enabled; @@ -129,6 +123,16 @@ public Set getTagWhitelist() { public void setTagWhitelist(Set tagWhitelist) { this.tagWhitelist = tagWhitelist; } + + public boolean isLegacyAiTokenMetricsEnabled() { + return legacyAiTokenMetricsEnabled; + } + + public void setLegacyAiTokenMetricsEnabled( + boolean legacyAiTokenMetricsEnabled + ) { + this.legacyAiTokenMetricsEnabled = legacyAiTokenMetricsEnabled; + } } public static class BudgetProperties { @@ -223,4 +227,34 @@ public void setEnabled(boolean enabled) { this.enabled = enabled; } } + + public static class SpringAiProperties { + private String defaultModelId; + private Long defaultReservedOutputTokens; + private long framingHeadroomTokens; + + public String getDefaultModelId() { + return defaultModelId; + } + + public void setDefaultModelId(String defaultModelId) { + this.defaultModelId = defaultModelId; + } + + public Long getDefaultReservedOutputTokens() { + return defaultReservedOutputTokens; + } + + public void setDefaultReservedOutputTokens(Long defaultReservedOutputTokens) { + this.defaultReservedOutputTokens = defaultReservedOutputTokens; + } + + public long getFramingHeadroomTokens() { + return framingHeadroomTokens; + } + + public void setFramingHeadroomTokens(long framingHeadroomTokens) { + this.framingHeadroomTokens = framingHeadroomTokens; + } + } } diff --git a/token-pilot-autoconfigure/src/main/java/io/tokenpilot/autoconfigure/TokenPilotSpringAiAutoConfiguration.java b/token-pilot-autoconfigure/src/main/java/io/tokenpilot/autoconfigure/TokenPilotSpringAiAutoConfiguration.java new file mode 100644 index 0000000..2f49df1 --- /dev/null +++ b/token-pilot-autoconfigure/src/main/java/io/tokenpilot/autoconfigure/TokenPilotSpringAiAutoConfiguration.java @@ -0,0 +1,125 @@ +package io.tokenpilot.autoconfigure; + +import io.tokenpilot.budget.BudgetEvaluator; +import io.tokenpilot.budget.AtomicBudgetStateStore; +import io.tokenpilot.core.CostCalculator; +import io.tokenpilot.core.LedgerManager; +import io.tokenpilot.core.ModelRegistry; +import io.tokenpilot.core.PreflightCostEstimator; +import io.tokenpilot.core.PricingEvaluator; +import io.tokenpilot.core.PricingMissingListener; +import io.tokenpilot.core.PricingRegistry; +import io.tokenpilot.core.TokenBudget; +import io.tokenpilot.core.TokenEstimator; +import io.tokenpilot.core.domain.MissingPricingPolicy; +import io.tokenpilot.springai.LedgerAdvisor; +import io.tokenpilot.springai.ReservedOutputTokensResolver; +import io.tokenpilot.springai.UsageExtractor; +import io.tokenpilot.springai.internal.LedgerSpringAiComponents; +import org.springframework.ai.chat.client.ChatClient; +import org.springframework.beans.factory.ObjectProvider; +import org.springframework.boot.autoconfigure.AutoConfiguration; +import org.springframework.boot.autoconfigure.condition.ConditionalOnBean; +import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; +import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean; +import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; +import org.springframework.context.annotation.Bean; + +/** Optional Spring AI adapter의 bean graph를 격리하는 자동 설정입니다. */ +@AutoConfiguration(after = { + TokenPilotAutoConfiguration.class, + TokenPilotBudgetAutoConfiguration.class, + TokenPilotCoreMetricsAutoConfiguration.class +}) +@ConditionalOnClass({ChatClient.class, LedgerSpringAiComponents.class}) +@ConditionalOnProperty( + prefix = "token-pilot", + name = "enabled", + havingValue = "true", + matchIfMissing = true +) +public class TokenPilotSpringAiAutoConfiguration { + + @Bean + @ConditionalOnMissingBean + public UsageExtractor usageExtractor() { + return LedgerSpringAiComponents.defaultUsageExtractor(); + } + + @Bean(name = "ledgerAdvisor") + @ConditionalOnMissingBean + @ConditionalOnProperty( + prefix = "token-pilot.budget", + name = "enabled", + havingValue = "false", + matchIfMissing = true + ) + public LedgerAdvisor ledgerOnlyAdvisor( + LedgerManager ledgerManager, + UsageExtractor usageExtractor, + CostCalculator costCalculator, + PricingRegistry pricingRegistry, + PricingEvaluator pricingEvaluator, + ObjectProvider pricingMissingListeners + ) { + return LedgerSpringAiComponents.defaultLedgerAdvisor( + ledgerManager, + usageExtractor, + null, + null, + costCalculator, + pricingRegistry, + pricingEvaluator, + MissingPricingPolicy.FAIL_OPEN, + pricingMissingListeners.orderedStream().toList() + ); + } + + @Bean(name = "ledgerAdvisor") + @ConditionalOnMissingBean(LedgerAdvisor.class) + @ConditionalOnProperty( + prefix = "token-pilot.budget", + name = "enabled", + havingValue = "true" + ) + public LedgerAdvisor accountingLedgerAdvisor( + UsageExtractor usageExtractor, + BudgetEvaluator budgetEvaluator, + AtomicBudgetStateStore budgetStateStore, + PricingRegistry pricingRegistry, + PricingEvaluator pricingEvaluator, + ModelRegistry modelRegistry, + TokenEstimator tokenEstimator, + TokenBudget tokenBudget, + PreflightCostEstimator preflightCostEstimator, + ObjectProvider reservedOutputTokensResolver, + ObjectProvider pricingMissingListeners, + TokenPilotProperties properties + ) { + return LedgerSpringAiComponents.accountingLedgerAdvisor( + usageExtractor, + budgetEvaluator, + budgetStateStore, + budgetStateStore, + pricingRegistry, + modelRegistry, + tokenEstimator, + tokenBudget, + preflightCostEstimator, + pricingEvaluator, + pricingMissingListeners.orderedStream().toList(), + properties.getSpringAi().getDefaultModelId(), + properties.getSpringAi().getDefaultReservedOutputTokens(), + reservedOutputTokensResolver.getIfAvailable(), + properties.getSpringAi().getFramingHeadroomTokens() + ); + } + + @Bean + @ConditionalOnBean(LedgerAdvisor.class) + public LedgerChatClientBuilderCustomizer ledgerChatClientBuilderCustomizer( + LedgerAdvisor ledgerAdvisor + ) { + return new LedgerChatClientBuilderCustomizer(ledgerAdvisor); + } +} diff --git a/token-pilot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports b/token-pilot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports index f531281..3f7af7a 100644 --- a/token-pilot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports +++ b/token-pilot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports @@ -1 +1,7 @@ io.tokenpilot.autoconfigure.TokenPilotAutoConfiguration +io.tokenpilot.autoconfigure.TokenPilotCoreMetricsAutoConfiguration +io.tokenpilot.autoconfigure.TokenPilotBudgetMetricsAutoConfiguration +io.tokenpilot.autoconfigure.TokenPilotBudgetAutoConfiguration +io.tokenpilot.autoconfigure.TokenPilotSpringAiAutoConfiguration +io.tokenpilot.autoconfigure.TokenPilotNotificationMetricsAutoConfiguration +io.tokenpilot.autoconfigure.TokenPilotNotificationAutoConfiguration diff --git a/token-pilot-autoconfigure/src/test/java/io/tokenpilot/autoconfigure/ChatClientBuilderCustomizerTest.java b/token-pilot-autoconfigure/src/test/java/io/tokenpilot/autoconfigure/ChatClientBuilderCustomizerTest.java index 8d39441..4d02eec 100644 --- a/token-pilot-autoconfigure/src/test/java/io/tokenpilot/autoconfigure/ChatClientBuilderCustomizerTest.java +++ b/token-pilot-autoconfigure/src/test/java/io/tokenpilot/autoconfigure/ChatClientBuilderCustomizerTest.java @@ -14,7 +14,10 @@ class ChatClientBuilderCustomizerTest { private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() - .withConfiguration(AutoConfigurations.of(TokenPilotAutoConfiguration.class)); + .withConfiguration(AutoConfigurations.of( + TokenPilotAutoConfiguration.class, + TokenPilotSpringAiAutoConfiguration.class + )); @Test @DisplayName("Advisor가 존재할 때 ChatClementCustomizer를 등록해야 합니다.") 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..c3eb8f6 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 @@ -2,18 +2,32 @@ import io.micrometer.core.instrument.MeterRegistry; import io.micrometer.core.instrument.simple.SimpleMeterRegistry; +import io.tokenpilot.budget.ActualUsageCommand; import io.tokenpilot.budget.BudgetDecision; import io.tokenpilot.budget.BudgetEvaluator; import io.tokenpilot.budget.BudgetKey; +import io.tokenpilot.budget.BudgetReservationRequest; +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.IdempotencyKey; +import io.tokenpilot.budget.ReservationAccounting; +import io.tokenpilot.budget.ReservationAccountingEvent; +import io.tokenpilot.budget.ReservationAccountingListener; +import io.tokenpilot.budget.ReservationId; +import io.tokenpilot.budget.ReservationStatus; +import io.tokenpilot.budget.ReservationTokenEstimate; import io.tokenpilot.core.CostCalculator; import io.tokenpilot.core.LedgerManager; +import io.tokenpilot.core.ModelRegistry; +import io.tokenpilot.core.PreflightCostEstimator; import io.tokenpilot.core.PricingEvaluator; import io.tokenpilot.core.PricingProvider; import io.tokenpilot.core.PricingRegistry; +import io.tokenpilot.core.TokenBudget; +import io.tokenpilot.core.TokenEstimator; import io.tokenpilot.core.domain.Cost; import io.tokenpilot.core.domain.PricingPlan; import io.tokenpilot.core.domain.PricingReconciliationResult; @@ -22,8 +36,15 @@ 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.micrometer.internal.BudgetMetricsPublisher; +import io.tokenpilot.micrometer.internal.CoreMetricsPublisher; +import io.tokenpilot.micrometer.internal.MicroCostMetricsPublisher; +import io.tokenpilot.micrometer.internal.NotificationMetricsPublisher; 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; @@ -34,10 +55,12 @@ import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; import org.springframework.ai.chat.client.ChatClientRequest; +import org.springframework.ai.chat.client.ChatClientResponse; import org.springframework.ai.chat.client.advisor.api.AdvisorChain; import org.springframework.ai.chat.prompt.ChatOptions; import org.springframework.ai.chat.prompt.Prompt; import org.springframework.boot.autoconfigure.AutoConfigurations; +import org.springframework.boot.test.context.FilteredClassLoader; import org.springframework.boot.test.context.runner.ApplicationContextRunner; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; @@ -47,8 +70,11 @@ 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.concurrent.atomic.AtomicInteger; import java.util.stream.Stream; import static io.tokenpilot.core.domain.TokenType.COMPLETION; @@ -57,7 +83,6 @@ import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.junit.jupiter.params.provider.Arguments.argumentSet; import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; class TokenPilotAutoConfigurationTest { @@ -68,7 +93,15 @@ class TokenPilotAutoConfigurationTest { private static final String PROP_CURRENCY = PREFIX + ".currency"; private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() - .withConfiguration(AutoConfigurations.of(TokenPilotAutoConfiguration.class)); + .withConfiguration(AutoConfigurations.of( + TokenPilotAutoConfiguration.class, + TokenPilotCoreMetricsAutoConfiguration.class, + TokenPilotBudgetMetricsAutoConfiguration.class, + TokenPilotNotificationMetricsAutoConfiguration.class, + TokenPilotBudgetAutoConfiguration.class, + TokenPilotSpringAiAutoConfiguration.class, + TokenPilotNotificationAutoConfiguration.class + )); @Test @DisplayName("기본 설정에서 Core 및 Spring AI 빈은 등록되고, Budget 빈은 등록되지 않아야 한다") @@ -78,7 +111,11 @@ void shouldRegisterDefaultBeans() { assertThat(context).hasSingleBean(PricingRegistry.class); assertThat(context).hasSingleBean(CostCalculator.class); assertThat(context).hasSingleBean(PricingEvaluator.class); + assertThat(context).hasSingleBean(ModelRegistry.class); + assertThat(context).hasSingleBean(TokenEstimator.class); + assertThat(context).hasSingleBean(TokenBudget.class); assertThat(context).hasSingleBean(LedgerManager.class); + assertThat(context).hasSingleBean(PreflightCostEstimator.class); assertThat(context).hasSingleBean(UsageExtractor.class); assertThat(context).hasSingleBean(LedgerAdvisor.class); @@ -86,11 +123,43 @@ void shouldRegisterDefaultBeans() { assertThat(context).doesNotHaveBean(BudgetStateStore.class); assertThat(context).doesNotHaveBean(BudgetEvaluator.class); + assertThat(context).doesNotHaveBean(ReservationAccounting.class); assertThat(context).doesNotHaveBean(NotificationStateStore.class); assertThat(context).doesNotHaveBean(BudgetNotificationService.class); + assertThat(context).doesNotHaveBean(CoreMetricsPublisher.class); + assertThat(context).doesNotHaveBean(BudgetMetricsPublisher.class); + assertThat(context).doesNotHaveBean(NotificationMetricsPublisher.class); }); } + @Test + @DisplayName("선택 adapter 모듈이 없어도 core 자동 설정은 시작되어야 한다") + void shouldStartCoreOnlyWithoutOptionalAdapterModules() { + this.contextRunner + .withClassLoader(new FilteredClassLoader( + "io.tokenpilot.budget", + "io.tokenpilot.notification", + "io.tokenpilot.springai", + "io.tokenpilot.micrometer", + "org.springframework.ai", + "io.micrometer" + )) + .run(context -> { + assertThat(context).hasNotFailed(); + assertThat(context).hasSingleBean(PricingProvider.class); + assertThat(context).hasSingleBean(PricingRegistry.class); + assertThat(context).hasSingleBean(CostCalculator.class); + assertThat(context).hasSingleBean(ModelRegistry.class); + assertThat(context).hasSingleBean(TokenEstimator.class); + assertThat(context).hasSingleBean(TokenBudget.class); + assertThat(context).hasSingleBean(LedgerManager.class); + assertThat(context).doesNotHaveBean("ledgerAdvisor"); + assertThat(context).doesNotHaveBean("budgetStateStore"); + assertThat(context).doesNotHaveBean("notificationStateStore"); + assertThat(context).doesNotHaveBean("tokenPilotCoreMetricsPublisher"); + }); + } + @Test @DisplayName("Ledger-only advisor는 Prompt model과 기본 policy로 pricing snapshot을 resolve해야 한다") void shouldResolvePricingSnapshotInLedgerOnlyAdvisor() { @@ -139,6 +208,33 @@ void shouldResolvePricingSnapshotInLedgerOnlyAdvisor() { }); } + @Test + @DisplayName("pricing miss는 advisor에서 Token Pilot 고유 meter로 연결되어야 한다") + void shouldPublishPricingMissingMetricFromAdvisor() { + this.contextRunner + .withUserConfiguration(MeterRegistryConfiguration.class) + .run(context -> { + LedgerAdvisor advisor = context.getBean(LedgerAdvisor.class); + ChatClientRequest request = new ChatClientRequest( + new Prompt( + "test", + ChatOptions.builder().model("missing-model").build() + ), + Map.of() + ); + + advisor.before(request, mock(AdvisorChain.class)); + + MeterRegistry registry = context.getBean(MeterRegistry.class); + assertThat(registry.get("tokenpilot.pricing.missing") + .tag("policy", "fail_open") + .counter() + .count()).isEqualTo(1.0); + assertThat(registry.get("tokenpilot.pricing.missing") + .counter().getId().getTag("model")).isNull(); + }); + } + @Test @DisplayName("설정 값이 없을 경우 빈 목록을 가진 PricingProvider가 생성되어야 한다") void shouldRegisterDefaultPricingProviderWhenNoProperties() { @@ -248,6 +344,147 @@ void shouldRegisterBudgetBeansWhenEnabled() { .run(context -> { assertThat(context).hasSingleBean(BudgetStateStore.class); assertThat(context).hasSingleBean(BudgetEvaluator.class); + assertThat(context).hasSingleBean(ReservationAccounting.class); + }); + } + + @Test + @DisplayName("Budget가 활성화된 custom store는 원자적 accounting 계약이 없으면 시작에 실패해야 한다") + void shouldFailFastWhenBudgetStoreDoesNotSupportReservationAccounting() { + this.contextRunner + .withUserConfiguration(StoreOnlyBudgetConfiguration.class) + .withPropertyValues("token-pilot.budget.enabled=true") + .run(context -> { + assertThat(context).hasFailed(); + assertThat(context.getStartupFailure()) + .hasStackTraceContaining("ReservationAccounting"); + }); + } + + @Test + @DisplayName("서로 다른 store와 accounting 빈은 원자적 budget graph로 조합하지 않아야 한다") + void shouldRejectMismatchedBudgetStoreAndAccountingBeans() { + this.contextRunner + .withUserConfiguration(MismatchedBudgetGraphConfiguration.class) + .withPropertyValues("token-pilot.budget.enabled=true") + .run(context -> { + assertThat(context).hasFailed(); + assertThat(context.getStartupFailure()) + .hasStackTraceContaining("AtomicBudgetStateStore"); + }); + } + + @Test + @DisplayName("BudgetStateStore를 의존하는 listener도 순환 참조 없이 지연 연결되어야 한다") + void shouldLazilyResolveStoreDependentAccountingListener() { + this.contextRunner + .withUserConfiguration(StoreDependentAccountingListenerConfiguration.class) + .withPropertyValues("token-pilot.budget.enabled=true") + .run(context -> { + assertThat(context).hasNotFailed(); + BudgetStateStore stateStore = context.getBean(BudgetStateStore.class); + StoreDependentAccountingListener listener = context.getBean( + StoreDependentAccountingListener.class + ); + Currency usd = Currency.getInstance("USD"); + Cost limit = Cost.of(new BigDecimal("10.00"), usd); + BudgetKey key = new BudgetKey( + "policy-a", + "tenant", + "tenant-a", + BudgetWindow.parse("2026-08") + ); + + stateStore.checkAndReserve(new BudgetReservationRequest( + key, + limit, + Cost.of(BigDecimal.ONE, usd), + "request-lazy", + new IdempotencyKey("dedupe-lazy"), + null, + null, + null, + Optional.empty(), + Optional.empty() + )); + + assertThat(listener.stateStore()).isSameAs(stateStore); + assertThat(listener.evaluations()).isEqualTo(1); + }); + } + + @Test + @DisplayName("자동 설정 budget store가 reservation, commit, listener failure metric을 연결해야 한다") + void shouldWireBudgetAccountingMetricsIntoAutoConfiguredStore() { + this.contextRunner + .withUserConfiguration( + MeterRegistryConfiguration.class, + FailingAccountingListenerConfiguration.class + ) + .withPropertyValues( + "token-pilot.budget.enabled=true", + "token-pilot.budget.monthly-limit=100.00" + ) + .run(context -> { + Currency usd = Currency.getInstance("USD"); + BudgetKey key = new BudgetKey( + "policy-a", + "tenant", + "tenant-a", + BudgetWindow.parse("2026-08") + ); + Cost limit = Cost.of(new BigDecimal("100.00"), usd); + PricingSnapshot snapshot = new PricingSnapshot( + "gpt-4o-mini", + "pricing-v1", + "catalog-v1", + Instant.parse("2026-08-25T00:00:00Z"), + Map.of( + PROMPT, new BigDecimal("0.001"), + COMPLETION, new BigDecimal("0.002") + ), + usd + ); + BudgetReservationRequest request = new BudgetReservationRequest( + key, + limit, + Cost.of(new BigDecimal("1.00"), usd), + "request-a", + new IdempotencyKey("dedupe-a"), + snapshot, + new ReservationTokenEstimate(100, 120, 50) + ); + BudgetStateStore stateStore = context.getBean(BudgetStateStore.class); + ReservationAccounting accounting = (ReservationAccounting) stateStore; + + BudgetReservationResult reservation = stateStore.checkAndReserve(request); + ReservationId reservationId = reservation.reservationId(); + accounting.markInFlight(reservationId); + ActualUsageCommand command = new ActualUsageCommand( + "request-a", + "attempt-a", + reservationId, + TokenUsage.from(110, 60), + "gpt-4o-mini" + ); + var applied = accounting.commit(command); + var reused = accounting.commit(command); + + MeterRegistry registry = context.getBean(MeterRegistry.class); + assertThat(applied.transition().status().isApplied()).isTrue(); + assertThat(reused.transition().status().isApplied()).isFalse(); + assertThat(registry.get("tokenpilot.budget.reservations") + .tag("state", "created").counter().count()).isEqualTo(1.0); + assertThat(registry.get("tokenpilot.cost.total") + .tag("currency", "USD").counter().count()).isEqualTo(0.00023); + assertThat(registry.get("tokenpilot.reconciliation.outcomes") + .tag("outcome", "committed").counter().count()).isEqualTo(1.0); + assertThat(registry.get("tokenpilot.listener.failures") + .tags("listener", "custom", "phase", "reservation_evaluated") + .counter().count()).isEqualTo(1.0); + assertThat(registry.get("tokenpilot.listener.failures") + .tags("listener", "custom", "phase", "accounting_applied") + .counter().count()).isEqualTo(1.0); }); } @@ -328,47 +565,132 @@ void shouldWireBudgetEvaluatorIntoLedgerAdvisorWhenBudgetEnabled() { } @Test - @DisplayName("Budget가 활성화되면 missing pricing policy 기본값은 FAIL_CLOSED여야 한다") - void shouldUseFailClosedMissingPricingPolicyWhenBudgetEnabled() { + @DisplayName("token-pilot.budget.enabled=false 일 때 Budget 관련 빈이 등록되지 않아야 한다") + void shouldNotRegisterBudgetBeansWhenDisabled() { this.contextRunner - .withUserConfiguration(RecordingBudgetEvaluatorConfiguration.class) - .withPropertyValues("token-pilot.budget.enabled=true") + .withPropertyValues("token-pilot.budget.enabled=false") .run(context -> { - LedgerAdvisor advisor = context.getBean(LedgerAdvisor.class); - ChatClientRequest request = new ChatClientRequest( - new Prompt("test"), - Map.of( - "tenant_id", "tenant-abc", - "tokenpilot.model.id", "missing-model" - ) - ); + assertThat(context).doesNotHaveBean(BudgetStateStore.class); + assertThat(context).doesNotHaveBean(BudgetEvaluator.class); + }); + } - assertThatThrownBy(() -> advisor.before(request, mock(AdvisorChain.class))) - .isInstanceOf(MissingPricingException.class) - .hasMessage("MISSING_PLAN") - .extracting(exception -> ((MissingPricingException) exception).getResolution()) - .isEqualTo(PricingResolution.MISSING_PLAN); + @Test + @DisplayName("MeterRegistry가 존재하면 Token Pilot 고유 publisher만 기본 등록되어야 한다") + void shouldRegisterMicrometerBeanWhenMeterRegistryExists() { + this.contextRunner + .withUserConfiguration(MeterRegistryConfiguration.class) + .run(context -> { + assertThat(context).hasSingleBean(CoreMetricsPublisher.class); + assertThat(context).hasSingleBean(BudgetMetricsPublisher.class); + assertThat(context).hasSingleBean(NotificationMetricsPublisher.class); + assertThat(context).doesNotHaveBean(MicroCostMetricsPublisher.class); + assertThat(context.getBean(TokenPilotProperties.class) + .getMetrics().getTagWhitelist()).isEmpty(); + assertThat(context.getBean(TokenPilotProperties.class) + .getMetrics().isLegacyAiTokenMetricsEnabled()).isFalse(); }); } @Test - @DisplayName("token-pilot.budget.enabled=false 일 때 Budget 관련 빈이 등록되지 않아야 한다") - void shouldNotRegisterBudgetBeansWhenDisabled() { + @DisplayName("자동 설정 TokenBudget의 admission 결과가 preflight meter로 연결되어야 한다") + void shouldPublishPreflightMetricFromAutoConfiguredTokenBudget() { this.contextRunner - .withPropertyValues("token-pilot.budget.enabled=false") + .withUserConfiguration(MeterRegistryConfiguration.class) .run(context -> { - assertThat(context).doesNotHaveBean(BudgetStateStore.class); - assertThat(context).doesNotHaveBean(BudgetEvaluator.class); + TokenEstimator estimator = context.getBean(TokenEstimator.class); + TokenBudget tokenBudget = context.getBean(TokenBudget.class); + + tokenBudget.check( + "gpt-4o-mini", + estimator.estimate("preflight metric"), + 256 + ); + + MeterRegistry registry = context.getBean(MeterRegistry.class); + assertThat(registry.get("tokenpilot.preflight.requests") + .tags("decision", "indeterminate", "reason", "incomplete_scope") + .counter().count()).isEqualTo(1.0); }); } @Test - @DisplayName("MeterRegistry가 존재할 때 Micrometer 관련 빈이 등록되어야 한다") - void shouldRegisterMicrometerBeanWhenMeterRegistryExists() { + @DisplayName("legacy ai.token metrics는 명시적으로 활성화할 때만 등록되어야 한다") + void shouldRegisterLegacyMicrometerBeanOnlyWhenExplicitlyEnabled() { this.contextRunner .withUserConfiguration(MeterRegistryConfiguration.class) + .withPropertyValues( + "token-pilot.metrics.legacy-ai-token-metrics-enabled=true", + "token-pilot.metrics.tag-whitelist[0]=tenant_id" + ) .run(context -> { + assertThat(context).hasSingleBean(MicroCostMetricsPublisher.class); assertThat(context).hasBean("microCostMetricsPublisher"); + assertThat(context.getBean(TokenPilotProperties.class) + .getMetrics().getTagWhitelist()).containsExactly("tenant_id"); + assertThat(context.getBean(TokenPilotProperties.class) + .getMetrics().isLegacyAiTokenMetricsEnabled()).isTrue(); + }); + } + + @Test + @DisplayName("legacy flag 기본값에서는 Ledger 기록이 ai.token meter를 만들지 않아야 한다") + void shouldNotPublishLegacyMetersByDefault() { + this.contextRunner + .withUserConfiguration(MeterRegistryConfiguration.class) + .withPropertyValues(buildProperties( + "gpt-4o", + "0.005", + "0.015", + "USD" + )) + .run(context -> { + context.getBean(LedgerManager.class).record( + "gpt-4o", + TokenUsage.from(10, 20), + Map.of("tenant_id", "raw-tenant") + ); + + MeterRegistry registry = context.getBean(MeterRegistry.class); + assertThat(registry.find("ai.token.cost.total").counter()).isNull(); + assertThat(registry.find("ai.token.usage.total").counter()).isNull(); + }); + } + + @Test + @DisplayName("사용자 정의 metrics publisher가 있으면 기본 publisher가 물러나야 한다") + void shouldBackOffForUserDefinedMetricsPublisher() { + this.contextRunner + .withUserConfiguration(UserMetricsPublisherConfiguration.class) + .run(context -> { + assertThat(context).hasSingleBean(CoreMetricsPublisher.class); + assertThat(context).hasBean("customCoreMetricsPublisher"); + assertThat(context).doesNotHaveBean("tokenPilotCoreMetricsPublisher"); + }); + } + + @Test + @DisplayName("legacy flag를 활성화하면 기존 ai.token meter를 발행해야 한다") + void shouldPublishLegacyMetersWhenEnabled() { + this.contextRunner + .withUserConfiguration(MeterRegistryConfiguration.class) + .withPropertyValues( + "token-pilot.metrics.legacy-ai-token-metrics-enabled=true", + PROP_MODEL_ID + "=gpt-4o", + PROP_PROMPT + "=0.005", + PROP_COMPLETION + "=0.015", + PROP_CURRENCY + "=USD" + ) + .run(context -> { + context.getBean(LedgerManager.class).record( + "gpt-4o", + TokenUsage.from(10, 20), + Map.of() + ); + + MeterRegistry registry = context.getBean(MeterRegistry.class); + assertThat(registry.find("ai.token.cost.total").counter()).isNotNull(); + assertThat(registry.find("ai.token.usage.total").counters()).isNotEmpty(); }); } @@ -377,8 +699,14 @@ void shouldRegisterMicrometerBeanWhenMeterRegistryExists() { void shouldNotRegisterMicrometerBeanWhenMetricsDisabled() { this.contextRunner .withUserConfiguration(MeterRegistryConfiguration.class) - .withPropertyValues("token-pilot.metrics.enabled=false") + .withPropertyValues( + "token-pilot.metrics.enabled=false", + "token-pilot.metrics.legacy-ai-token-metrics-enabled=true" + ) .run(context -> { + assertThat(context).doesNotHaveBean(CoreMetricsPublisher.class); + assertThat(context).doesNotHaveBean(BudgetMetricsPublisher.class); + assertThat(context).doesNotHaveBean(NotificationMetricsPublisher.class); assertThat(context).doesNotHaveBean("microCostMetricsPublisher"); }); } @@ -425,15 +753,128 @@ void shouldNotRegisterNotificationServiceWhenHandlerMissing() { void shouldRegisterNotificationServiceWhenEnabledAndHandlerExists() { this.contextRunner .withUserConfiguration(FakeBudgetNotificationHandlerConfiguration.class) - .withPropertyValues("token-pilot.notification.enabled=true") + .withPropertyValues( + "token-pilot.budget.enabled=true", + "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.budget.enabled=true", + "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 = TokenPilotBudgetPolicyFactory.from(properties).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); + }); + } + + @Test + @DisplayName("notification success와 dedup 결과가 lifecycle meter로 연결되어야 한다") + void shouldPublishNotificationLifecycleMetrics() { + this.contextRunner + .withUserConfiguration( + MeterRegistryConfiguration.class, + FakeBudgetNotificationHandlerConfiguration.class + ) + .withPropertyValues( + "token-pilot.budget.enabled=true", + "token-pilot.notification.enabled=true" + ) + .run(context -> { + Currency usd = Currency.getInstance("USD"); + BudgetDecision decision = new BudgetDecision( + new BudgetKey( + "policy-a", + "tenant", + "tenant-a", + BudgetWindow.parse("2026-08") + ), + BudgetDecision.EvaluationType.STATUS, + BudgetState.WARN, + BudgetThreshold.HALF, + "half reached", + Cost.of(new BigDecimal("5.00"), usd), + Cost.of(new BigDecimal("5.00"), usd), + Cost.of(new BigDecimal("10.00"), usd) + ); + BudgetNotificationService service = context.getBean( + BudgetNotificationService.class + ); + + service.notifyIfNeeded(decision, Map.of()); + service.notifyIfNeeded(decision, Map.of()); + + MeterRegistry registry = context.getBean(MeterRegistry.class); + assertThat(registry.get("tokenpilot.notification.events") + .tags("outcome", "success", "threshold", "half") + .counter().count()).isEqualTo(1.0); + assertThat(registry.get("tokenpilot.notification.events") + .tags("outcome", "deduplicated", "threshold", "half") + .counter().count()).isEqualTo(1.0); + }); + } + private static Stream providePricingConfigs() { return Stream.of( argumentSet( @@ -505,6 +946,32 @@ public RecordingBudgetEvaluator budgetEvaluator() { } } + @Configuration(proxyBeanMethods = false) + static class StoreOnlyBudgetConfiguration { + @Bean + public BudgetStateStore budgetStateStore() { + return mock(BudgetStateStore.class); + } + + @Bean + public UsageExtractor usageExtractor() { + return response -> TokenUsage.from(10, 20); + } + } + + @Configuration(proxyBeanMethods = false) + static class MismatchedBudgetGraphConfiguration { + @Bean + public BudgetStateStore budgetStateStore() { + return mock(BudgetStateStore.class); + } + + @Bean + public ReservationAccounting reservationAccounting() { + return mock(ReservationAccounting.class); + } + } + @Configuration(proxyBeanMethods = false) static class FixedClockConfiguration { @Bean @@ -560,12 +1027,128 @@ public MeterRegistry meterRegistry() { } } + @Configuration(proxyBeanMethods = false) + static class UserMetricsPublisherConfiguration { + @Bean + public MeterRegistry meterRegistry() { + return new SimpleMeterRegistry(); + } + + @Bean + public CoreMetricsPublisher customCoreMetricsPublisher( + MeterRegistry meterRegistry + ) { + return new CoreMetricsPublisher(meterRegistry); + } + } + + @Configuration(proxyBeanMethods = false) + static class FailingAccountingListenerConfiguration { + @Bean + public ReservationAccountingListener failingAccountingListener() { + return new ReservationAccountingListener() { + @Override + public void onCommitted(ReservationAccountingEvent event) { + throw new IllegalStateException("accounting listener failed"); + } + + @Override + public void onReservationEvaluated( + BudgetReservationRequest request, + BudgetReservationResult result + ) { + throw new IllegalStateException("reservation listener failed"); + } + }; + } + } + + @Configuration(proxyBeanMethods = false) + static class StoreDependentAccountingListenerConfiguration { + @Bean + StoreDependentAccountingListener storeDependentAccountingListener( + BudgetStateStore stateStore + ) { + return new StoreDependentAccountingListener(stateStore); + } + } + + static final class StoreDependentAccountingListener + implements ReservationAccountingListener { + + private final BudgetStateStore stateStore; + private final AtomicInteger evaluations = new AtomicInteger(); + + private StoreDependentAccountingListener(BudgetStateStore stateStore) { + this.stateStore = stateStore; + } + + @Override + public void onCommitted(ReservationAccountingEvent event) { + } + + @Override + public void onReservationEvaluated( + BudgetReservationRequest request, + BudgetReservationResult result + ) { + evaluations.incrementAndGet(); + } + + BudgetStateStore stateStore() { + return stateStore; + } + + int evaluations() { + return evaluations.get(); + } + } + // 테스트용 no-op handler - 실제 알림 전송 없이 빈 등록 여부만 검증 @Configuration(proxyBeanMethods = false) 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/AtomicBudgetStateStore.java b/token-pilot-budget/src/main/java/io/tokenpilot/budget/AtomicBudgetStateStore.java new file mode 100644 index 0000000..f8edc5a --- /dev/null +++ b/token-pilot-budget/src/main/java/io/tokenpilot/budget/AtomicBudgetStateStore.java @@ -0,0 +1,12 @@ +package io.tokenpilot.budget; + +/** + * 예약 승인과 후속 회계 전이를 같은 원자적 상태에서 처리하는 budget store입니다. + * + *

Spring 같은 DI container는 이 합성 계약을 단일 bean으로 노출해 + * {@link BudgetStateStore}와 {@link ReservationAccounting}이 서로 다른 저장소로 + * 연결되는 것을 방지할 수 있습니다.

+ */ +public interface AtomicBudgetStateStore + extends BudgetStateStore, ReservationAccounting { +} diff --git a/token-pilot-budget/src/main/java/io/tokenpilot/budget/BudgetReservation.java b/token-pilot-budget/src/main/java/io/tokenpilot/budget/BudgetReservation.java index 443660e..1657acc 100644 --- a/token-pilot-budget/src/main/java/io/tokenpilot/budget/BudgetReservation.java +++ b/token-pilot-budget/src/main/java/io/tokenpilot/budget/BudgetReservation.java @@ -130,6 +130,13 @@ public static BudgetReservation reserved( ); } + /** + * 같은 idempotency key에 전달된 요청이 기존 예약과 같은 비용 책임을 뜻하는지 비교합니다. + * + *

{@link PricingSnapshot#checkedAt()}은 가격 조회 시점의 관측 metadata이므로 비교에서 + * 제외합니다. 반면 model, pricing policy, catalog version, currency, rate와 token estimate를 + * 포함해 실제 가격 또는 예약 책임을 바꾸는 값은 모두 일치해야 합니다.

+ */ public boolean matches(BudgetReservationRequest request) { return key.equals(request.key()) && limit.equals(request.limit()) @@ -139,10 +146,29 @@ public boolean matches(BudgetReservationRequest request) { && Objects.equals(modelId, request.modelId()) && Objects.equals(pricingPolicyId, request.pricingPolicyId()) && Objects.equals(catalogVersion, request.catalogVersion()) - && pricingSnapshot.equals(request.pricingSnapshot()) + && hasSamePricingTerms( + pricingSnapshot, + request.pricingSnapshot() + ) && tokenEstimate.equals(request.tokenEstimate()); } + private static boolean hasSamePricingTerms( + Optional existing, + Optional candidate + ) { + if (existing.isEmpty() || candidate.isEmpty()) { + return existing.isEmpty() && candidate.isEmpty(); + } + + PricingSnapshot existingSnapshot = existing.orElseThrow(); + PricingSnapshot candidateSnapshot = candidate.orElseThrow(); + return PricingSnapshot.haveSameBillingTerms( + existingSnapshot, + candidateSnapshot + ); + } + /** 이 예약이 지정한 provider 요청에 속하는지 확인합니다. */ public boolean belongsTo(String candidateRequestId) { return requestId.equals(candidateRequestId); diff --git a/token-pilot-budget/src/main/java/io/tokenpilot/budget/ReservationAccounting.java b/token-pilot-budget/src/main/java/io/tokenpilot/budget/ReservationAccounting.java index e5f8494..9840e52 100644 --- a/token-pilot-budget/src/main/java/io/tokenpilot/budget/ReservationAccounting.java +++ b/token-pilot-budget/src/main/java/io/tokenpilot/budget/ReservationAccounting.java @@ -1,6 +1,9 @@ package io.tokenpilot.budget; import io.tokenpilot.core.domain.Cost; +import io.tokenpilot.core.domain.PricingSnapshot; + +import java.util.Objects; /** * 예약의 회계 상태와 금액을 변경하는 단일 진입점입니다. @@ -106,7 +109,11 @@ ReservationTransition release( */ ReservationTransition commitCost(ReservationId reservationId, Cost actualCost); - /** provider actual usage를 예약 시점 가격으로 계산하여 확정합니다. */ + /** + * provider actual usage를 예약 시점 가격으로 계산하여 확정합니다. + * 응답 모델이 예약 pricing snapshot과 다르면 다른 모델 가격을 request 가격으로 + * 확정하지 않도록 거부합니다. + */ ReservationReconciliation commit(ActualUsageCommand command); /** actual을 확보하지 못한 예약을 정산 대기로 전환합니다. */ @@ -124,9 +131,42 @@ ReservationTransition markReconciliationRequired( ReservationAccountingReason reason ); - /** 늦게 도착한 provider actual usage를 예약 시점 가격으로 계산하여 확정합니다. */ + /** + * response model과 actual usage를 보존하면서 pricing reconciliation 대기로 이동합니다. + * 기존 구현은 command metadata를 보존하지 않는 호환 동작으로 위임할 수 있습니다. + */ + default ReservationTransition markReconciliationRequired( + ActualUsageCommand command, + ReservationAccountingReason reason + ) { + Objects.requireNonNull(command, "command must not be null"); + return markReconciliationRequired(command.reservationId(), reason); + } + + /** + * 늦게 도착한 provider actual usage를 예약 시점 가격으로 계산하여 확정합니다. + * 응답 모델은 예약 pricing snapshot의 모델과 같아야 합니다. + */ ReservationReconciliation reconcileLateActual(ActualUsageCommand command); + /** + * pending actual에 response model의 명시적 immutable pricing snapshot을 적용합니다. + * 기본 구현은 기존 예약 snapshot만 지원하는 구현과의 호환을 위해 fail-closed합니다. + */ + default ReservationReconciliation reconcileLateActual( + ActualUsageCommand command, + PricingSnapshot actualPricingSnapshot + ) { + Objects.requireNonNull(command, "command must not be null"); + Objects.requireNonNull( + actualPricingSnapshot, + "actualPricingSnapshot must not be null" + ); + throw new UnsupportedOperationException( + "late reconciliation with an alternate pricing snapshot is unsupported" + ); + } + /** * pricing snapshot과 token estimate가 없던 pending 호환 예약을 caller가 계산한 actual 비용으로 확정합니다. * 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..e0d509e 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,50 @@ package io.tokenpilot.budget; -/** 예약 정산 이벤트를 수신하는 framework-independent 계약입니다. */ +/** + * 예약 admission과 정산 결과를 수신하는 framework-independent 계약입니다. + * + *

기존 lambda listener와의 source 호환성을 위해 {@link #onCommitted}만 추상 + * 메서드로 유지합니다. 나머지 callback은 예약 저장소가 확정한 결과를 관찰하기 위한 + * 선택적 확장점입니다. 회계 변경 callback은 새로 적용된 전이에만 전달되고, 예약 + * callback은 모든 평가 결과와 BLOCK 결과를 각각 관찰할 수 있습니다.

+ */ @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 + ) { + } + + /** 원자적 예약 시도의 모든 결과를 전달합니다. */ + default void onReservationEvaluated( + BudgetReservationRequest request, + BudgetReservationResult result + ) { + } + /** 원자적 admission에서 상태 변경 없이 차단된 결과를 전달합니다. */ + default void onReservationBlocked( + BudgetReservationRequest request, + BudgetReservationResult result + ) { + } + + /** metric tag에 안전한 제한된 listener 종류를 반환합니다. */ + default ReservationAccountingListenerType listenerType() { + return ReservationAccountingListenerType.CUSTOM; + } } diff --git a/token-pilot-budget/src/main/java/io/tokenpilot/budget/ReservationAccountingListenerErrorHandler.java b/token-pilot-budget/src/main/java/io/tokenpilot/budget/ReservationAccountingListenerErrorHandler.java new file mode 100644 index 0000000..a97c4a4 --- /dev/null +++ b/token-pilot-budget/src/main/java/io/tokenpilot/budget/ReservationAccountingListenerErrorHandler.java @@ -0,0 +1,12 @@ +package io.tokenpilot.budget; + +/** bounded 회계 listener 실패를 관찰하는 best-effort hook입니다. */ +@FunctionalInterface +public interface ReservationAccountingListenerErrorHandler { + + void onFailure(ReservationAccountingListenerFailureEvent event); + + static ReservationAccountingListenerErrorHandler noOp() { + return ignored -> { }; + } +} diff --git a/token-pilot-budget/src/main/java/io/tokenpilot/budget/ReservationAccountingListenerFailureEvent.java b/token-pilot-budget/src/main/java/io/tokenpilot/budget/ReservationAccountingListenerFailureEvent.java new file mode 100644 index 0000000..89684b0 --- /dev/null +++ b/token-pilot-budget/src/main/java/io/tokenpilot/budget/ReservationAccountingListenerFailureEvent.java @@ -0,0 +1,20 @@ +package io.tokenpilot.budget; + +import java.util.Objects; + +/** + * 회계 listener callback 실패를 bounded 정보만으로 전달하는 이벤트입니다. + * + *

원본 예외와 메시지를 포함하지 않아 prompt, provider 응답 또는 자격 증명이 + * 관측 경로로 전달되지 않습니다.

+ */ +public record ReservationAccountingListenerFailureEvent( + ReservationAccountingListenerType listenerType, + ReservationAccountingListenerPhase phase +) { + + public ReservationAccountingListenerFailureEvent { + Objects.requireNonNull(listenerType, "listenerType must not be null"); + Objects.requireNonNull(phase, "phase must not be null"); + } +} diff --git a/token-pilot-budget/src/main/java/io/tokenpilot/budget/ReservationAccountingListenerPhase.java b/token-pilot-budget/src/main/java/io/tokenpilot/budget/ReservationAccountingListenerPhase.java new file mode 100644 index 0000000..93642ce --- /dev/null +++ b/token-pilot-budget/src/main/java/io/tokenpilot/budget/ReservationAccountingListenerPhase.java @@ -0,0 +1,9 @@ +package io.tokenpilot.budget; + +/** 회계 listener 실패 관측에 사용하는 제한된 callback 단계입니다. */ +public enum ReservationAccountingListenerPhase { + ACCOUNTING_APPLIED, + RECONCILIATION_REQUIRED, + RESERVATION_EVALUATED, + RESERVATION_BLOCKED +} diff --git a/token-pilot-budget/src/main/java/io/tokenpilot/budget/ReservationAccountingListenerType.java b/token-pilot-budget/src/main/java/io/tokenpilot/budget/ReservationAccountingListenerType.java new file mode 100644 index 0000000..8fac5d8 --- /dev/null +++ b/token-pilot-budget/src/main/java/io/tokenpilot/budget/ReservationAccountingListenerType.java @@ -0,0 +1,8 @@ +package io.tokenpilot.budget; + +/** 회계 listener 실패 관측에 사용하는 제한된 listener 종류입니다. */ +public enum ReservationAccountingListenerType { + METRICS, + NOTIFICATION, + CUSTOM +} diff --git a/token-pilot-budget/src/main/java/io/tokenpilot/budget/ReservationAccountingReason.java b/token-pilot-budget/src/main/java/io/tokenpilot/budget/ReservationAccountingReason.java index af50e02..6e85725 100644 --- a/token-pilot-budget/src/main/java/io/tokenpilot/budget/ReservationAccountingReason.java +++ b/token-pilot-budget/src/main/java/io/tokenpilot/budget/ReservationAccountingReason.java @@ -5,6 +5,7 @@ public enum ReservationAccountingReason { ACTUAL_USAGE_REPORTED, LATE_ACTUAL_USAGE_REPORTED, ACTUAL_USAGE_UNAVAILABLE, + PRICING_RECONCILIATION_REQUIRED, CALLBACK_TIMED_OUT, CANCELLED_BEFORE_DISPATCH, PROVIDER_CONFIRMED_UNBILLED, diff --git a/token-pilot-budget/src/main/java/io/tokenpilot/budget/ReservationReconciliation.java b/token-pilot-budget/src/main/java/io/tokenpilot/budget/ReservationReconciliation.java index be34bea..e8cb910 100644 --- a/token-pilot-budget/src/main/java/io/tokenpilot/budget/ReservationReconciliation.java +++ b/token-pilot-budget/src/main/java/io/tokenpilot/budget/ReservationReconciliation.java @@ -16,7 +16,10 @@ public record ReservationReconciliation( ReservationId reservationId, BudgetKey budgetKey, String responseModelId, + /** 예약 시점에 estimate를 계산한 request pricing snapshot입니다. */ PricingSnapshot pricingSnapshot, + /** provider response model의 actual 비용 계산에 사용한 pricing snapshot입니다. */ + PricingSnapshot actualPricingSnapshot, ReservationTokenEstimate tokenEstimate, ReservationActualTokens actualTokens, Cost estimate, @@ -33,6 +36,10 @@ public record ReservationReconciliation( Objects.requireNonNull(budgetKey, "budgetKey must not be null"); responseModelId = requireText(responseModelId, "responseModelId"); Objects.requireNonNull(pricingSnapshot, "pricingSnapshot must not be null"); + Objects.requireNonNull( + actualPricingSnapshot, + "actualPricingSnapshot must not be null" + ); Objects.requireNonNull(tokenEstimate, "tokenEstimate must not be null"); Objects.requireNonNull(actualTokens, "actualTokens must not be null"); Objects.requireNonNull(estimate, "estimate must not be null"); @@ -44,11 +51,71 @@ public record ReservationReconciliation( "estimate and actual must use the same currency" ); } - if (!estimate.currency().equals(pricingSnapshot.currency())) { + if (!estimate.currency().equals(pricingSnapshot.currency()) + || !actual.currency().equals(actualPricingSnapshot.currency())) { throw new IllegalArgumentException( "reconciliation costs must use the pricing snapshot currency" ); } + if (!responseModelId.equals(actualPricingSnapshot.modelId())) { + throw new IllegalArgumentException( + "responseModelId must match the actual pricing snapshot model" + ); + } + } + + /** + * 기존 request pricing snapshot만 사용한 정산 결과를 만드는 호환 생성자입니다. + * request와 response model이 같은 기존 경로에서는 두 snapshot이 동일합니다. + */ + public ReservationReconciliation( + String requestId, + String attemptId, + ReservationId reservationId, + BudgetKey budgetKey, + String responseModelId, + PricingSnapshot pricingSnapshot, + ReservationTokenEstimate tokenEstimate, + ReservationActualTokens actualTokens, + Cost estimate, + Cost actual, + boolean overLimit, + ReservationTransition transition, + ReservationAccountingReason reason + ) { + this( + requestId, + attemptId, + reservationId, + budgetKey, + responseModelId, + pricingSnapshot, + actualSnapshotForResponseModel(pricingSnapshot, responseModelId), + tokenEstimate, + actualTokens, + estimate, + actual, + overLimit, + transition, + reason + ); + } + + private static PricingSnapshot actualSnapshotForResponseModel( + PricingSnapshot requestSnapshot, + String responseModelId + ) { + if (requestSnapshot.modelId().equals(responseModelId)) { + return requestSnapshot; + } + return new PricingSnapshot( + responseModelId, + requestSnapshot.pricingPolicyId(), + requestSnapshot.catalogVersion(), + requestSnapshot.checkedAt(), + requestSnapshot.rates(), + requestSnapshot.currency() + ); } /** 예약 시점 pricing snapshot의 request model입니다. */ @@ -56,6 +123,16 @@ public String requestModelId() { return pricingSnapshot.modelId(); } + /** actual provider usage에 적용한 response model pricing policy입니다. */ + public String actualPricingPolicyId() { + return actualPricingSnapshot.pricingPolicyId(); + } + + /** actual provider usage에 적용한 response model catalog version입니다. */ + public String actualCatalogVersion() { + return actualPricingSnapshot.catalogVersion(); + } + /** 예약 시점 pricing policy 식별자입니다. */ public String pricingPolicyId() { return pricingSnapshot.pricingPolicyId(); 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..120e4ee --- /dev/null +++ b/token-pilot-budget/src/main/java/io/tokenpilot/budget/ReservationReconciliationRequiredEvent.java @@ -0,0 +1,77 @@ +package io.tokenpilot.budget; + +import java.util.Objects; +import java.util.Optional; + +/** actual usage를 확보하지 못해 pending liability로 이동한 원자적 회계 결과입니다. */ +public record ReservationReconciliationRequiredEvent( + ReservationId reservationId, + BudgetKey budgetKey, + ReservationTransition transition, + ReservationAccountingReason reason, + BudgetSnapshot snapshot, + Optional pendingActualUsage +) { + + 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"); + pendingActualUsage = Objects.requireNonNull( + pendingActualUsage, + "pendingActualUsage 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.PRICING_RECONCILIATION_REQUIRED + && 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" + ); + } + if (pendingActualUsage.isPresent()) { + ActualUsageCommand command = pendingActualUsage.orElseThrow(); + if (!reservationId.equals(command.reservationId())) { + throw new IllegalArgumentException( + "pending actual and event must use the same reservation" + ); + } + if (reason != ReservationAccountingReason.PRICING_RECONCILIATION_REQUIRED) { + throw new IllegalArgumentException( + "pending actual requires pricing reconciliation reason" + ); + } + } + } + + /** 기존 actual-unavailable 이벤트를 만드는 호환 생성자입니다. */ + public ReservationReconciliationRequiredEvent( + ReservationId reservationId, + BudgetKey budgetKey, + ReservationTransition transition, + ReservationAccountingReason reason, + BudgetSnapshot snapshot + ) { + this( + reservationId, + budgetKey, + transition, + reason, + snapshot, + Optional.empty() + ); + } +} diff --git a/token-pilot-budget/src/main/java/io/tokenpilot/budget/internal/ActualUsageFingerprint.java b/token-pilot-budget/src/main/java/io/tokenpilot/budget/internal/ActualUsageFingerprint.java index 2f08d41..ac017cf 100644 --- a/token-pilot-budget/src/main/java/io/tokenpilot/budget/internal/ActualUsageFingerprint.java +++ b/token-pilot-budget/src/main/java/io/tokenpilot/budget/internal/ActualUsageFingerprint.java @@ -3,6 +3,9 @@ import io.tokenpilot.budget.ActualUsageCommand; import io.tokenpilot.budget.ReservationActualTokens; import io.tokenpilot.budget.ReservationId; +import io.tokenpilot.core.domain.PricingSnapshot; + +import java.util.Objects; /** 중복 actual callback을 민감하거나 무제한인 metadata 없이 식별합니다. */ record ActualUsageFingerprint( @@ -10,19 +13,71 @@ record ActualUsageFingerprint( String attemptId, ReservationId reservationId, ReservationActualTokens actualTokens, - String responseModelId + String responseModelId, + PricingSnapshot pricingSnapshot ) { + ActualUsageFingerprint { + Objects.requireNonNull(requestId, "requestId must not be null"); + Objects.requireNonNull(attemptId, "attemptId must not be null"); + Objects.requireNonNull(reservationId, "reservationId must not be null"); + Objects.requireNonNull(actualTokens, "actualTokens must not be null"); + Objects.requireNonNull(responseModelId, "responseModelId must not be null"); + Objects.requireNonNull(pricingSnapshot, "pricingSnapshot must not be null"); + } + static ActualUsageFingerprint from( ActualUsageCommand command, - ReservationActualTokens actualTokens + ReservationActualTokens actualTokens, + PricingSnapshot pricingSnapshot ) { return new ActualUsageFingerprint( command.requestId(), command.attemptId(), command.reservationId(), actualTokens, - command.responseModelId() + command.responseModelId(), + pricingSnapshot + ); + } + + boolean sameCallback(ActualUsageFingerprint other) { + Objects.requireNonNull(other, "other must not be null"); + return requestId.equals(other.requestId) + && attemptId.equals(other.attemptId) + && reservationId.equals(other.reservationId) + && actualTokens.equals(other.actualTokens) + && responseModelId.equals(other.responseModelId); + } + + @Override + public boolean equals(Object other) { + if (this == other) { + return true; + } + if (!(other instanceof ActualUsageFingerprint that)) { + return false; + } + return requestId.equals(that.requestId) + && attemptId.equals(that.attemptId) + && reservationId.equals(that.reservationId) + && actualTokens.equals(that.actualTokens) + && responseModelId.equals(that.responseModelId) + && PricingSnapshot.haveSameBillingTerms( + pricingSnapshot, + that.pricingSnapshot + ); + } + + @Override + public int hashCode() { + return Objects.hash( + requestId, + attemptId, + reservationId, + actualTokens, + responseModelId, + pricingSnapshot.billingTermsHashCode() ); } } 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 370f49a..19c6497 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 @@ -17,9 +17,10 @@ import java.util.Objects; /** - * 현재 구현이 지원하는 snapshot은 확정 비용과 후보 비용입니다. + * 이 evaluator는 부수 효과 없는 policy projection만 담당합니다. * {@code projectedUsage = committedUsage + candidateCost}이며 active reservation과 - * reconciliation liability는 #36, #37에서 추가됩니다. + * reconciliation liability는 {@link BudgetStateStore#snapshot}과 원자적 reservation + * store가 별도로 반영합니다. */ public class DefaultBudgetEvaluator implements BudgetEvaluator { 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..4a7fdb4 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 @@ -2,6 +2,7 @@ import io.tokenpilot.budget.ActualUsageCommand; import io.tokenpilot.budget.AccountingTransitionStatus; +import io.tokenpilot.budget.AtomicBudgetStateStore; import io.tokenpilot.budget.BudgetKey; import io.tokenpilot.budget.BudgetReservation; import io.tokenpilot.budget.BudgetReservationRequest; @@ -9,14 +10,19 @@ import io.tokenpilot.budget.BudgetSnapshot; import io.tokenpilot.budget.BudgetStateStore; import io.tokenpilot.budget.IdempotencyKey; -import io.tokenpilot.budget.ReservationAccounting; import io.tokenpilot.budget.ReservationAccountingEvent; import io.tokenpilot.budget.ReservationAccountingListener; +import io.tokenpilot.budget.ReservationAccountingListenerErrorHandler; +import io.tokenpilot.budget.ReservationAccountingListenerFailureEvent; +import io.tokenpilot.budget.ReservationAccountingListenerPhase; +import io.tokenpilot.budget.ReservationAccountingListenerType; import io.tokenpilot.budget.ReservationAccountingReason; 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.ReservationStatus; import io.tokenpilot.budget.ReservationStateMachine; import io.tokenpilot.budget.ReservationTransition; import io.tokenpilot.budget.ReservationTokenEstimate; @@ -43,17 +49,21 @@ *

bucket별 monitor가 조회·통화 검증·한도 검증·예약 갱신을 함께 보호하고, * 별도의 idempotency index가 같은 요청의 중복 예약을 차단합니다.

*/ -public class InMemoryBudgetStateStore implements BudgetStateStore, ReservationAccounting { +public class InMemoryBudgetStateStore implements AtomicBudgetStateStore { private final ConcurrentMap store = new ConcurrentHashMap<>(); private final ConcurrentMap idempotencyIndex = new ConcurrentHashMap<>(); private final ConcurrentMap reservationIndex = new ConcurrentHashMap<>(); + private final ConcurrentMap + pendingActualFingerprints = new ConcurrentHashMap<>(); private final Clock clock; private final Supplier reservationIdGenerator; private final CostCalculator costCalculator; - private final List accountingListeners; + private final Supplier> accountingListenerSupplier; + private final Supplier> + listenerErrorHandlerSupplier; public InMemoryBudgetStateStore() { this( @@ -79,7 +89,7 @@ public InMemoryBudgetStateStore( Supplier reservationIdGenerator, CostCalculator costCalculator ) { - this(clock, reservationIdGenerator, costCalculator, List.of()); + this(clock, reservationIdGenerator, costCalculator, List.of(), List.of()); } public InMemoryBudgetStateStore( @@ -87,6 +97,45 @@ public InMemoryBudgetStateStore( Supplier reservationIdGenerator, CostCalculator costCalculator, List accountingListeners + ) { + this( + clock, + reservationIdGenerator, + costCalculator, + accountingListeners, + List.of() + ); + } + + public InMemoryBudgetStateStore( + Clock clock, + Supplier reservationIdGenerator, + CostCalculator costCalculator, + List accountingListeners, + List listenerErrorHandlers + ) { + this( + clock, + reservationIdGenerator, + costCalculator, + fixedListSupplier( + accountingListeners, + "accountingListeners must not be null" + ), + fixedListSupplier( + listenerErrorHandlers, + "listenerErrorHandlers must not be null" + ) + ); + } + + public InMemoryBudgetStateStore( + Clock clock, + Supplier reservationIdGenerator, + CostCalculator costCalculator, + Supplier> accountingListenerSupplier, + Supplier> + listenerErrorHandlerSupplier ) { this.clock = Objects.requireNonNull(clock, "clock must not be null"); this.reservationIdGenerator = Objects.requireNonNull( @@ -97,11 +146,13 @@ public InMemoryBudgetStateStore( costCalculator, "costCalculator must not be null" ); - this.accountingListeners = List.copyOf( - Objects.requireNonNull( - accountingListeners, - "accountingListeners must not be null" - ) + this.accountingListenerSupplier = Objects.requireNonNull( + accountingListenerSupplier, + "accountingListenerSupplier must not be null" + ); + this.listenerErrorHandlerSupplier = Objects.requireNonNull( + listenerErrorHandlerSupplier, + "listenerErrorHandlerSupplier must not be null" ); } @@ -145,7 +196,13 @@ 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" + ); + publishReservationEvaluated(request, reservationResult); + publishBlockedReservation(request, reservationResult); + return reservationResult; } @Override @@ -249,7 +306,8 @@ public ReservationReconciliation commit(ActualUsageCommand command) { return reconcileUsage( command, CommitType.DIRECT, - ReservationAccountingReason.ACTUAL_USAGE_REPORTED + ReservationAccountingReason.ACTUAL_USAGE_REPORTED, + Optional.empty() ); } @@ -257,26 +315,135 @@ public ReservationReconciliation commit(ActualUsageCommand command) { public ReservationTransition markReconciliationRequired( ReservationId reservationId, ReservationAccountingReason reason + ) { + return markReconciliationRequiredInternal( + reservationId, + reason, + Optional.empty() + ); + } + + @Override + public ReservationTransition markReconciliationRequired( + ActualUsageCommand command, + ReservationAccountingReason reason + ) { + Objects.requireNonNull(command, "command must not be null"); + if (reason != ReservationAccountingReason.PRICING_RECONCILIATION_REQUIRED) { + throw new IllegalArgumentException( + "pending actual metadata requires pricing reconciliation reason" + ); + } + return markReconciliationRequiredInternal( + command.reservationId(), + reason, + Optional.of(command) + ); + } + + private ReservationTransition markReconciliationRequiredInternal( + ReservationId reservationId, + ReservationAccountingReason reason, + Optional pendingActualUsage ) { requireReconciliationRequiredReason(reason); + Objects.requireNonNull(pendingActualUsage, "pendingActualUsage must not be null"); Bucket bucket = bucketFor(reservationId); + ReservationTransition transition; + ReservationReconciliationRequiredEvent event; synchronized (bucket) { ReservationAccountingState accountingState = accountingState( bucket, reservationId ); - ReservationTransition transition = - accountingState.evaluateReconciliationRequired(); + pendingActualUsage.ifPresent(command -> + validatePendingActual(command, accountingState.reservation()) + ); + Optional pendingFingerprint = + pendingActualUsage.map(command -> pendingFingerprint( + command, + accountingState.reservation() + )); + transition = accountingState.evaluateReconciliationRequired(); if (!transition.status().isApplied()) { + if (transition.status() == AccountingTransitionStatus.REUSED) { + pendingFingerprint.ifPresent(fingerprint -> + mergePendingActualFingerprint(reservationId, fingerprint) + ); + } return transition; } + pendingFingerprint.ifPresent(fingerprint -> + mergePendingActualFingerprint(reservationId, fingerprint) + ); moveActiveReservationToPending( bucket, accountingState, transition.resultingState() ); - return transition; + event = new ReservationReconciliationRequiredEvent( + reservationId, + accountingState.reservation().key(), + transition, + reason, + bucket.snapshot(accountingState.reservation().key()), + pendingActualUsage + ); + } + publishReconciliationRequired(event); + return transition; + } + + private ActualUsageFingerprint pendingFingerprint( + ActualUsageCommand command, + BudgetReservation reservation + ) { + PricingSnapshot requestSnapshot = reservation.pricingSnapshot().orElseThrow( + () -> new IllegalStateException( + "reservation does not contain a pricing snapshot" + ) + ); + return ActualUsageFingerprint.from( + command, + ReservationActualTokens.from(command.usage()), + requestSnapshot + ); + } + + private void mergePendingActualFingerprint( + ReservationId reservationId, + ActualUsageFingerprint candidate + ) { + ActualUsageFingerprint existing = pendingActualFingerprints.putIfAbsent( + reservationId, + candidate + ); + if (existing != null && !existing.sameCallback(candidate)) { + throw new IllegalArgumentException( + "pending actual callback conflicts with the existing callback" + ); + } + } + + private static void validatePendingActual( + ActualUsageCommand command, + BudgetReservation reservation + ) { + if (!reservation.belongsTo(command.requestId())) { + throw new IllegalArgumentException( + "requestId must match the reservation request" + ); + } + PricingSnapshot requestSnapshot = reservation.pricingSnapshot().orElseThrow( + () -> new IllegalStateException( + "reservation does not contain a pricing snapshot" + ) + ); + if (requestSnapshot.modelId().equals(command.responseModelId())) { + throw new IllegalArgumentException( + "pricing reconciliation requires a different response model" + ); } } @@ -326,7 +493,7 @@ private AccountingTransitionOutcome applyCostInBucket( ReservationAccountingState accountingState, Cost actualCost, CommitType type, - Optional fingerprint + Optional fingerprint ) { ReservationTransition transition = type == CommitType.DIRECT ? accountingState.evaluateCommit(actualCost, fingerprint) @@ -358,7 +525,24 @@ public ReservationReconciliation reconcileLateActual(ActualUsageCommand command) return reconcileUsage( command, CommitType.LATE_ACTUAL, - ReservationAccountingReason.LATE_ACTUAL_USAGE_REPORTED + ReservationAccountingReason.LATE_ACTUAL_USAGE_REPORTED, + Optional.empty() + ); + } + + @Override + public ReservationReconciliation reconcileLateActual( + ActualUsageCommand command, + PricingSnapshot actualPricingSnapshot + ) { + return reconcileUsage( + command, + CommitType.LATE_ACTUAL, + ReservationAccountingReason.LATE_ACTUAL_USAGE_REPORTED, + Optional.of(Objects.requireNonNull( + actualPricingSnapshot, + "actualPricingSnapshot must not be null" + )) ); } @@ -384,6 +568,7 @@ public ReservationTransition writeOff( accountingState, transition.resultingState() ); + pendingActualFingerprints.remove(reservationId); return transition; } } @@ -418,28 +603,39 @@ private void validateCurrency(Cost limit, Cost amount) { private ReservationReconciliation reconcileUsage( ActualUsageCommand command, CommitType type, - ReservationAccountingReason reason + ReservationAccountingReason reason, + Optional actualPricingSnapshot ) { Objects.requireNonNull(command, "command must not be null"); Objects.requireNonNull(type, "type must not be null"); Objects.requireNonNull(reason, "reason must not be null"); + Objects.requireNonNull( + actualPricingSnapshot, + "actualPricingSnapshot must not be null" + ); Bucket bucket = bucketFor(command.reservationId()); ReservationReconciliation reconciliation; + BudgetSnapshot accountingSnapshot; synchronized (bucket) { reconciliation = reconcileUsageInBucket( bucket, command, type, - reason + reason, + actualPricingSnapshot ); + accountingSnapshot = bucket.snapshot(reconciliation.budgetKey()); } - publishAccountingEvent(reconciliation); + publishAccountingEvent(reconciliation, accountingSnapshot); return reconciliation; } private void publishAccountingEvent( - ReservationReconciliation reconciliation + ReservationReconciliation reconciliation, + BudgetSnapshot accountingSnapshot ) { + List accountingListeners = + accountingListeners(); if (!reconciliation.transition().status().isApplied() || accountingListeners.isEmpty()) { return; @@ -448,18 +644,124 @@ private void publishAccountingEvent( reconciliation ); for (ReservationAccountingListener listener : accountingListeners) { - notifyBestEffort(listener, event); + notifyBestEffort( + listener, + ReservationAccountingListenerPhase.ACCOUNTING_APPLIED, + () -> listener.onAccountingApplied(event, accountingSnapshot) + ); } } - private static void notifyBestEffort( + private void publishReservationEvaluated( + BudgetReservationRequest request, + BudgetReservationResult result + ) { + for (ReservationAccountingListener listener : accountingListeners()) { + notifyBestEffort( + listener, + ReservationAccountingListenerPhase.RESERVATION_EVALUATED, + () -> listener.onReservationEvaluated(request, result) + ); + } + } + + private void publishBlockedReservation( + BudgetReservationRequest request, + BudgetReservationResult result + ) { + if (result.status() != ReservationStatus.BLOCKED) { + return; + } + for (ReservationAccountingListener listener : accountingListeners()) { + notifyBestEffort( + listener, + ReservationAccountingListenerPhase.RESERVATION_BLOCKED, + () -> listener.onReservationBlocked(request, result) + ); + } + } + + private void publishReconciliationRequired( + ReservationReconciliationRequiredEvent event + ) { + for (ReservationAccountingListener listener : accountingListeners()) { + notifyBestEffort( + listener, + ReservationAccountingListenerPhase.RECONCILIATION_REQUIRED, + () -> listener.onReconciliationRequired(event) + ); + } + } + + private void notifyBestEffort( ReservationAccountingListener listener, - ReservationAccountingEvent event + ReservationAccountingListenerPhase phase, + Runnable callback ) { try { - listener.onCommitted(event); + callback.run(); } catch (RuntimeException ignored) { - // Listener 실패는 이미 적용된 회계 상태를 되돌리지 않습니다. + publishListenerFailure(listenerType(listener), phase); + } + } + + private ReservationAccountingListenerType listenerType( + ReservationAccountingListener listener + ) { + try { + ReservationAccountingListenerType listenerType = listener.listenerType(); + return listenerType == null + ? ReservationAccountingListenerType.CUSTOM + : listenerType; + } catch (RuntimeException ignored) { + return ReservationAccountingListenerType.CUSTOM; + } + } + + private void publishListenerFailure( + ReservationAccountingListenerType listenerType, + ReservationAccountingListenerPhase phase + ) { + List listenerErrorHandlers = + listenerErrorHandlers(); + if (listenerErrorHandlers.isEmpty()) { + return; + } + ReservationAccountingListenerFailureEvent event = + new ReservationAccountingListenerFailureEvent(listenerType, phase); + for (ReservationAccountingListenerErrorHandler errorHandler + : listenerErrorHandlers) { + try { + errorHandler.onFailure(event); + } catch (RuntimeException ignored) { + // Error handler도 best-effort이며 회계/admission 결과에 영향을 주지 않습니다. + } + } + } + + private List accountingListeners() { + return suppliedList(accountingListenerSupplier); + } + + private List listenerErrorHandlers() { + return suppliedList(listenerErrorHandlerSupplier); + } + + private static Supplier> fixedListSupplier( + List values, + String message + ) { + List copy = List.copyOf(Objects.requireNonNull(values, message)); + return () -> copy; + } + + private static List suppliedList(Supplier> supplier) { + try { + List values = supplier.get(); + return values == null ? List.of() : List.copyOf(values); + } catch (RuntimeException ignored) { + // Optional observer resolution cannot change reservation/accounting results. + return List.of(); } } @@ -467,7 +769,8 @@ private ReservationReconciliation reconcileUsageInBucket( Bucket bucket, ActualUsageCommand command, CommitType type, - ReservationAccountingReason reason + ReservationAccountingReason reason, + Optional actualPricingSnapshot ) { ReservationAccountingState accountingState = accountingState( bucket, @@ -485,6 +788,25 @@ private ReservationReconciliation reconcileUsageInBucket( "reservation does not contain a pricing snapshot" ) ); + PricingSnapshot resolvedActualSnapshot = actualPricingSnapshot.orElse(snapshot); + if (!resolvedActualSnapshot.modelId().equals(command.responseModelId())) { + throw new IllegalArgumentException( + actualPricingSnapshot.isEmpty() + ? "responseModelId must match the reservation pricing snapshot" + : "responseModelId must match the actual pricing snapshot" + ); + } + if (!reservation.amount().currency().equals(resolvedActualSnapshot.currency())) { + throw new IllegalArgumentException( + "actual pricing snapshot must use the reservation currency" + ); + } + if (actualPricingSnapshot.isPresent() + && type != CommitType.LATE_ACTUAL) { + throw new IllegalArgumentException( + "alternate pricing snapshot is only valid for late actual reconciliation" + ); + } ReservationTokenEstimate tokenEstimate = reservation.tokenEstimate().orElseThrow( () -> new IllegalStateException( "reservation does not contain a token estimate" @@ -495,8 +817,18 @@ private ReservationReconciliation reconcileUsageInBucket( ); ActualUsageFingerprint fingerprint = ActualUsageFingerprint.from( command, - actualTokens + actualTokens, + resolvedActualSnapshot + ); + ActualUsageFingerprint pendingFingerprint = pendingActualFingerprints.get( + command.reservationId() ); + if (pendingFingerprint != null + && !pendingFingerprint.sameCallback(fingerprint)) { + throw new IllegalArgumentException( + "actual usage callback conflicts with the pending pricing callback" + ); + } Optional reusedCommit = accountingState.reusedCommit(type, fingerprint); if (reusedCommit.isPresent()) { @@ -512,6 +844,7 @@ private ReservationReconciliation reconcileUsageInBucket( command, reservation, snapshot, + resolvedActualSnapshot, tokenEstimate, actualTokens, appliedCommit.actualCost(), @@ -520,7 +853,7 @@ private ReservationReconciliation reconcileUsageInBucket( ); } - Cost actualCost = calculateActualCost(command, snapshot); + Cost actualCost = calculateActualCost(command, resolvedActualSnapshot); AccountingTransitionOutcome outcome = applyCostInBucket( bucket, accountingState, @@ -528,10 +861,14 @@ private ReservationReconciliation reconcileUsageInBucket( type, Optional.of(fingerprint) ); + if (outcome.transition().status().isApplied()) { + pendingActualFingerprints.remove(command.reservationId()); + } return reconciliation( command, reservation, snapshot, + resolvedActualSnapshot, tokenEstimate, actualTokens, actualCost, @@ -556,7 +893,8 @@ private Cost calculateActualCost( private static ReservationReconciliation reconciliation( ActualUsageCommand command, BudgetReservation reservation, - PricingSnapshot snapshot, + PricingSnapshot requestPricingSnapshot, + PricingSnapshot actualPricingSnapshot, ReservationTokenEstimate tokenEstimate, ReservationActualTokens actualTokens, Cost actualCost, @@ -569,7 +907,8 @@ private static ReservationReconciliation reconciliation( reservation.id(), reservation.key(), command.responseModelId(), - snapshot, + requestPricingSnapshot, + actualPricingSnapshot, tokenEstimate, actualTokens, reservation.amount(), @@ -888,6 +1227,7 @@ private static void requireReconciliationRequiredReason( ) { Objects.requireNonNull(reason, "reason must not be null"); if (reason != ReservationAccountingReason.ACTUAL_USAGE_UNAVAILABLE + && reason != ReservationAccountingReason.PRICING_RECONCILIATION_REQUIRED && reason != ReservationAccountingReason.CALLBACK_TIMED_OUT) { throw new IllegalArgumentException( "reason is not valid for markReconciliationRequired" 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..71db80e 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 @@ -3,10 +3,13 @@ import io.tokenpilot.budget.BudgetEvaluator; import io.tokenpilot.budget.BudgetPolicy; import io.tokenpilot.budget.BudgetStateStore; +import io.tokenpilot.budget.AtomicBudgetStateStore; import io.tokenpilot.budget.ReservationAccounting; import io.tokenpilot.budget.ReservationAccountingListener; +import io.tokenpilot.budget.ReservationAccountingListenerErrorHandler; 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 +28,67 @@ public static BudgetStateStore inMemoryBudgetStateStore() { return new InMemoryBudgetStateStore(); } + /** 예약 admission과 accounting을 한 객체로 노출하는 기본 합성 store입니다. */ + public static AtomicBudgetStateStore inMemoryAtomicBudgetStateStore() { + return new InMemoryBudgetStateStore(); + } + + /** 예약과 회계를 같은 원자적 상태에 연결하는 합성 store factory입니다. */ + public static AtomicBudgetStateStore inMemoryAtomicBudgetStateStore( + Supplier> accountingListenerSupplier, + Supplier> + listenerErrorHandlerSupplier + ) { + return new InMemoryBudgetStateStore( + Clock.systemUTC(), + ReservationId::random, + LedgerComponents.defaultCostCalculator(), + accountingListenerSupplier, + listenerErrorHandlerSupplier + ); + } + + /** PR #68 accounting/notification listener 연결과 호환되는 convenience factory입니다. */ + public static BudgetStateStore inMemoryBudgetStateStore( + List accountingListeners + ) { + return new InMemoryBudgetStateStore( + Clock.systemUTC(), + ReservationId::random, + LedgerComponents.defaultCostCalculator(), + accountingListeners + ); + } + + /** listener 실패 관측 hook까지 연결하는 convenience factory입니다. */ + public static BudgetStateStore inMemoryBudgetStateStore( + List accountingListeners, + List listenerErrorHandlers + ) { + return new InMemoryBudgetStateStore( + Clock.systemUTC(), + ReservationId::random, + LedgerComponents.defaultCostCalculator(), + accountingListeners, + listenerErrorHandlers + ); + } + + /** Spring 등 DI container가 listener를 store 생성 뒤 지연 해석할 수 있는 factory입니다. */ + public static BudgetStateStore inMemoryBudgetStateStore( + Supplier> accountingListenerSupplier, + Supplier> + listenerErrorHandlerSupplier + ) { + return new InMemoryBudgetStateStore( + Clock.systemUTC(), + ReservationId::random, + LedgerComponents.defaultCostCalculator(), + accountingListenerSupplier, + listenerErrorHandlerSupplier + ); + } + public static BudgetStateStore inMemoryBudgetStateStore( Clock clock, Supplier reservationIdGenerator @@ -58,6 +122,22 @@ public static BudgetStateStore inMemoryBudgetStateStore( ); } + public static BudgetStateStore inMemoryBudgetStateStore( + Clock clock, + Supplier reservationIdGenerator, + CostCalculator costCalculator, + List accountingListeners, + List listenerErrorHandlers + ) { + return new InMemoryBudgetStateStore( + clock, + reservationIdGenerator, + costCalculator, + accountingListeners, + listenerErrorHandlers + ); + } + /** * 예약을 생성한 store와 동일한 객체의 회계 명령 진입점을 반환합니다. */ 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..6688fec 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,14 +1,18 @@ 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.budget.ReservationTokenEstimate; import io.tokenpilot.core.domain.Cost; import io.tokenpilot.core.domain.PricingSnapshot; import io.tokenpilot.core.domain.TokenType; @@ -19,6 +23,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 +32,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 +113,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(); @@ -147,6 +282,221 @@ class BudgetReservationStoreTest { assertThat(store.snapshot(KEY, LIMIT).activeReservationIds()).hasSize(1); } + @Test + void 같은_가격_조건의_다른_조회_시각은_기존_예약을_REUSED한다() { + InMemoryBudgetStateStore store = store(); + ReservationTokenEstimate estimate = new ReservationTokenEstimate(10, 12, 5); + PricingSnapshot firstSnapshot = pricingSnapshot( + CLOCK.instant(), + "gpt-4o-mini", + "pricing-v1", + "catalog-v1", + "0.10", + "0.20", + USD + ); + PricingSnapshot laterSnapshot = pricingSnapshot( + CLOCK.instant().plusSeconds(1), + "gpt-4o-mini", + "pricing-v1", + "catalog-v1", + "0.1000", + "0.2000", + USD + ); + + BudgetReservationResult created = store.checkAndReserve( + accountingRequest( + KEY, + LIMIT, + usd("10.00"), + "request-1", + "idempotency-1", + firstSnapshot, + estimate + ) + ); + BudgetReservationResult reused = store.checkAndReserve( + accountingRequest( + KEY, + LIMIT, + usd("10.00"), + "request-1", + "idempotency-1", + laterSnapshot, + estimate + ) + ); + + assertThat(created.status()).isEqualTo(ReservationStatus.CREATED); + assertThat(reused.status()).isEqualTo(ReservationStatus.REUSED); + assertThat(reused.reservation().id()).isEqualTo(created.reservation().id()); + assertThat(reused.reservation().pricingSnapshot()) + .contains(firstSnapshot); + assertThat(store.snapshot(KEY, LIMIT).activeReservedCost()) + .isEqualTo(usd("10.00")); + } + + @Test + void 같은_idempotency_key의_비용_책임_의미_변화는_CONFLICT다() { + InMemoryBudgetStateStore store = store(); + ReservationTokenEstimate estimate = new ReservationTokenEstimate(10, 12, 5); + PricingSnapshot snapshot = pricingSnapshot( + CLOCK.instant(), + "gpt-4o-mini", + "pricing-v1", + "catalog-v1", + "0.10", + "0.20", + USD + ); + BudgetReservationRequest original = accountingRequest( + KEY, + LIMIT, + usd("10.00"), + "request-1", + "idempotency-1", + snapshot, + estimate + ); + Currency eur = Currency.getInstance("EUR"); + + BudgetReservationResult created = store.checkAndReserve(original); + List changedRequests = List.of( + accountingRequest( + KEY, + LIMIT, + usd("10.00"), + "request-2", + "idempotency-1", + snapshot, + estimate + ), + accountingRequest( + KEY, + LIMIT, + usd("11.00"), + "request-1", + "idempotency-1", + snapshot, + estimate + ), + accountingRequest( + KEY, + usd("200.00"), + usd("10.00"), + "request-1", + "idempotency-1", + snapshot, + estimate + ), + accountingRequest( + KEY, + LIMIT, + usd("10.00"), + "request-1", + "idempotency-1", + pricingSnapshot( + CLOCK.instant(), + "gpt-4o", + "pricing-v1", + "catalog-v1", + "0.10", + "0.20", + USD + ), + estimate + ), + accountingRequest( + KEY, + LIMIT, + usd("10.00"), + "request-1", + "idempotency-1", + pricingSnapshot( + CLOCK.instant(), + "gpt-4o-mini", + "pricing-v2", + "catalog-v1", + "0.10", + "0.20", + USD + ), + estimate + ), + accountingRequest( + KEY, + LIMIT, + usd("10.00"), + "request-1", + "idempotency-1", + pricingSnapshot( + CLOCK.instant(), + "gpt-4o-mini", + "pricing-v1", + "catalog-v2", + "0.10", + "0.20", + USD + ), + estimate + ), + accountingRequest( + KEY, + Cost.of(new BigDecimal("100.00"), eur), + Cost.of(new BigDecimal("10.00"), eur), + "request-1", + "idempotency-1", + pricingSnapshot( + CLOCK.instant(), + "gpt-4o-mini", + "pricing-v1", + "catalog-v1", + "0.10", + "0.20", + eur + ), + estimate + ), + accountingRequest( + KEY, + LIMIT, + usd("10.00"), + "request-1", + "idempotency-1", + pricingSnapshot( + CLOCK.instant(), + "gpt-4o-mini", + "pricing-v1", + "catalog-v1", + "0.11", + "0.20", + USD + ), + estimate + ), + accountingRequest( + KEY, + LIMIT, + usd("10.00"), + "request-1", + "idempotency-1", + snapshot, + new ReservationTokenEstimate(11, 13, 6) + ) + ); + + assertThat(created.status()).isEqualTo(ReservationStatus.CREATED); + assertThat(changedRequests) + .allSatisfy(changed -> assertThat( + store.checkAndReserve(changed).status() + ).isEqualTo(ReservationStatus.CONFLICT)); + assertThat(store.snapshot(KEY, LIMIT).activeReservedCost()) + .isEqualTo(usd("10.00")); + assertThat(store.snapshot(KEY, LIMIT).activeReservationIds()) + .containsExactly(created.reservation().id()); + } + @Test void 같은_idempotency_key의_다른_요청은_CONFLICT이고_기존_예약을_보존한다() { InMemoryBudgetStateStore store = store(); @@ -341,6 +691,26 @@ private static BudgetReservationRequest request( ); } + private static BudgetReservationRequest accountingRequest( + BudgetKey key, + Cost limit, + Cost safeUpperBoundCost, + String requestId, + String idempotencyKey, + PricingSnapshot pricingSnapshot, + ReservationTokenEstimate tokenEstimate + ) { + return new BudgetReservationRequest( + key, + limit, + safeUpperBoundCost, + requestId, + new IdempotencyKey(idempotencyKey), + pricingSnapshot, + tokenEstimate + ); + } + private static Cost usd(String amount) { return Cost.of(new BigDecimal(amount), USD); } @@ -349,18 +719,38 @@ private static PricingSnapshot pricingSnapshot( String promptRate, String completionRate ) { - return new PricingSnapshot( + return pricingSnapshot( + CLOCK.instant(), "gpt-4o-mini", "pricing-v1", "catalog-v1", - CLOCK.instant(), + promptRate, + completionRate, + USD + ); + } + + private static PricingSnapshot pricingSnapshot( + Instant checkedAt, + String modelId, + String pricingPolicyId, + String catalogVersion, + String promptRate, + String completionRate, + Currency currency + ) { + return new PricingSnapshot( + modelId, + pricingPolicyId, + catalogVersion, + checkedAt, Map.of( TokenType.PROMPT, new BigDecimal(promptRate), TokenType.COMPLETION, new BigDecimal(completionRate) ), - USD + currency ); } } diff --git a/token-pilot-budget/src/test/java/io/tokenpilot/budget/internal/ReservationAccountingListenerLifecycleTest.java b/token-pilot-budget/src/test/java/io/tokenpilot/budget/internal/ReservationAccountingListenerLifecycleTest.java new file mode 100644 index 0000000..9e358fd --- /dev/null +++ b/token-pilot-budget/src/test/java/io/tokenpilot/budget/internal/ReservationAccountingListenerLifecycleTest.java @@ -0,0 +1,375 @@ +package io.tokenpilot.budget.internal; + +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.BudgetWindow; +import io.tokenpilot.budget.IdempotencyKey; +import io.tokenpilot.budget.ReservationAccountingEvent; +import io.tokenpilot.budget.ReservationAccountingListener; +import io.tokenpilot.budget.ReservationAccountingListenerFailureEvent; +import io.tokenpilot.budget.ReservationAccountingListenerPhase; +import io.tokenpilot.budget.ReservationAccountingListenerType; +import io.tokenpilot.budget.ReservationId; +import io.tokenpilot.budget.ReservationReconciliation; +import io.tokenpilot.budget.ReservationReconciliationRequiredEvent; +import io.tokenpilot.budget.ReservationStatus; +import io.tokenpilot.budget.ReservationTokenEstimate; +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 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.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; + +import static io.tokenpilot.budget.AccountingTransitionStatus.APPLIED; +import static io.tokenpilot.budget.AccountingTransitionStatus.REUSED; +import static io.tokenpilot.budget.ReservationState.COMMITTED; +import static io.tokenpilot.budget.ReservationState.RECONCILIATION_REQUIRED; +import static org.assertj.core.api.Assertions.assertThat; + +class ReservationAccountingListenerLifecycleTest { + + private static final Currency USD = Currency.getInstance("USD"); + private static final Currency KRW = Currency.getInstance("KRW"); + private static final Cost LIMIT = usd("100.00"); + private static final BudgetKey KEY = new BudgetKey( + "budget-policy", + "tenant", + "tenant-a", + BudgetWindow.parse("2026-08") + ); + private static final Clock CLOCK = Clock.fixed( + Instant.parse("2026-08-15T12:34:56Z"), + ZoneOffset.UTC + ); + private static final ReservationTokenEstimate TOKEN_ESTIMATE = + new ReservationTokenEstimate(90, 100, 50); + + @Test + void 모든_예약_결과를_평가_callback으로_전달하고_BLOCKED만_별도_전달한다() { + List evaluated = new ArrayList<>(); + List blocked = new ArrayList<>(); + ReservationAccountingListener listener = new ReservationAccountingListener() { + @Override + public void onCommitted(ReservationAccountingEvent event) { + } + + @Override + public void onReservationEvaluated( + BudgetReservationRequest request, + BudgetReservationResult result + ) { + evaluated.add(result.status()); + } + + @Override + public void onReservationBlocked( + BudgetReservationRequest request, + BudgetReservationResult result + ) { + blocked.add(result.status()); + } + }; + InMemoryBudgetStateStore store = store(List.of(listener), List.of()); + + store.checkAndReserve(legacyRequest("request-1", usd("20.00"))); + store.checkAndReserve(legacyRequest("request-1", usd("20.00"))); + store.checkAndReserve(legacyRequest("request-1", usd("30.00"))); + store.checkAndReserve(legacyRequest("request-2", usd("80.00"))); + store.checkAndReserve(legacyRequest( + "request-3", + Cost.of(new BigDecimal("10.00"), KRW) + )); + + assertThat(evaluated).containsExactly( + ReservationStatus.CREATED, + ReservationStatus.REUSED, + ReservationStatus.CONFLICT, + ReservationStatus.BLOCKED, + ReservationStatus.CURRENCY_MISMATCH + ); + assertThat(blocked).containsExactly(ReservationStatus.BLOCKED); + } + + @Test + void RECONCILIATION_REQUIRED는_APPLIED_전이만_snapshot과_함께_한번_전달한다() { + List events = new ArrayList<>(); + ReservationAccountingListener listener = new ReservationAccountingListener() { + @Override + public void onCommitted(ReservationAccountingEvent event) { + } + + @Override + public void onReconciliationRequired( + ReservationReconciliationRequiredEvent event + ) { + events.add(event); + } + }; + InMemoryBudgetStateStore store = store(List.of(listener), List.of()); + ReservationId reservationId = reserveUsage(store, "request-1", "dedupe-1", "60.00"); + store.markInFlight(reservationId); + + var applied = store.markReconciliationRequired(reservationId); + var duplicate = store.markReconciliationRequired(reservationId); + + assertThat(applied.status()).isEqualTo(APPLIED); + assertThat(duplicate.status()).isEqualTo(REUSED); + assertThat(events).hasSize(1); + ReservationReconciliationRequiredEvent event = events.getFirst(); + assertThat(event.transition()).isEqualTo(applied); + assertThat(event.snapshot().activeReservedCost()).isEqualTo(usd("0.00")); + assertThat(event.snapshot().pendingReconciliationLiability()) + .isEqualTo(usd("60.00")); + } + + @Test + void APPLIED_commit은_같은_linearization_point의_snapshot과_함께_한번_전달한다() { + List events = new ArrayList<>(); + AtomicReference snapshot = new AtomicReference<>(); + ReservationAccountingListener listener = new ReservationAccountingListener() { + @Override + public void onCommitted(ReservationAccountingEvent event) { + } + + @Override + public void onAccountingApplied( + ReservationAccountingEvent event, + BudgetSnapshot accountingSnapshot + ) { + events.add(event); + snapshot.set(accountingSnapshot); + } + }; + InMemoryBudgetStateStore store = store(List.of(listener), List.of()); + ReservationId reservationId = reserveUsage(store, "request-1", "dedupe-1", "60.00"); + store.markInFlight(reservationId); + + ReservationReconciliation applied = store.commit(command("request-1", reservationId)); + ReservationReconciliation duplicate = store.commit(command("request-1", reservationId)); + + assertThat(applied.transition().status()).isEqualTo(APPLIED); + assertThat(duplicate.transition().status()).isEqualTo(REUSED); + assertThat(events).containsExactly(new ReservationAccountingEvent(applied)); + assertThat(snapshot.get().committedCost()).isEqualTo(usd("40.00")); + assertThat(snapshot.get().activeReservedCost()).isEqualTo(usd("0.00")); + assertThat(snapshot.get().pendingReconciliationLiability()) + .isEqualTo(usd("0.00")); + } + + @Test + void callback과_error_handler_실패를_격리하고_bounded_failure를_계속_전달한다() { + List failures = new ArrayList<>(); + AtomicInteger observedEvaluations = new AtomicInteger(); + AtomicInteger observedBlocks = new AtomicInteger(); + AtomicInteger observedPending = new AtomicInteger(); + AtomicInteger observedCommits = new AtomicInteger(); + + ReservationAccountingListener failing = failingMetricsListener(); + ReservationAccountingListener following = new ReservationAccountingListener() { + @Override + public void onCommitted(ReservationAccountingEvent event) { + observedCommits.incrementAndGet(); + } + + @Override + public void onReconciliationRequired( + ReservationReconciliationRequiredEvent event + ) { + observedPending.incrementAndGet(); + } + + @Override + public void onReservationEvaluated( + BudgetReservationRequest request, + BudgetReservationResult result + ) { + observedEvaluations.incrementAndGet(); + } + + @Override + public void onReservationBlocked( + BudgetReservationRequest request, + BudgetReservationResult result + ) { + observedBlocks.incrementAndGet(); + } + }; + InMemoryBudgetStateStore store = store( + List.of(failing, following), + List.of( + event -> { + throw new IllegalStateException("error handler failed"); + }, + failures::add + ) + ); + + BudgetReservationResult blocked = store.checkAndReserve( + legacyRequest("blocked", usd("100.00")) + ); + ReservationId pendingId = reserveUsage( + store, + "pending-request", + "pending-dedupe", + "60.00" + ); + store.markInFlight(pendingId); + store.markReconciliationRequired(pendingId); + + ReservationId committedId = reserveUsage( + store, + "committed-request", + "committed-dedupe", + "20.00" + ); + store.markInFlight(committedId); + ReservationReconciliation committed = store.commit( + command("committed-request", committedId) + ); + store.commit(command("committed-request", committedId)); + + assertThat(blocked.status()).isEqualTo(ReservationStatus.BLOCKED); + assertThat(committed.transition().resultingState()).isEqualTo(COMMITTED); + assertThat(store.snapshot(KEY, LIMIT).pendingReconciliationLiability()) + .isEqualTo(usd("60.00")); + assertThat(store.snapshot(KEY, LIMIT).committedCost()).isEqualTo(usd("40.00")); + assertThat(observedEvaluations).hasValue(3); + assertThat(observedBlocks).hasValue(1); + assertThat(observedPending).hasValue(1); + assertThat(observedCommits).hasValue(1); + assertThat(failures) + .extracting(ReservationAccountingListenerFailureEvent::listenerType) + .containsOnly(ReservationAccountingListenerType.METRICS); + assertThat(failures) + .extracting(ReservationAccountingListenerFailureEvent::phase) + .contains( + ReservationAccountingListenerPhase.RESERVATION_EVALUATED, + ReservationAccountingListenerPhase.RESERVATION_BLOCKED, + ReservationAccountingListenerPhase.RECONCILIATION_REQUIRED, + ReservationAccountingListenerPhase.ACCOUNTING_APPLIED + ); + assertThat(failures.stream() + .filter(event -> event.phase() + == ReservationAccountingListenerPhase.ACCOUNTING_APPLIED)) + .hasSize(1); + } + + private static ReservationAccountingListener failingMetricsListener() { + return new ReservationAccountingListener() { + @Override + public void onCommitted(ReservationAccountingEvent event) { + throw new IllegalStateException("commit listener failed"); + } + + @Override + public void onReconciliationRequired( + ReservationReconciliationRequiredEvent event + ) { + throw new IllegalStateException("pending listener failed"); + } + + @Override + public void onReservationEvaluated( + BudgetReservationRequest request, + BudgetReservationResult result + ) { + throw new IllegalStateException("reservation listener failed"); + } + + @Override + public void onReservationBlocked( + BudgetReservationRequest request, + BudgetReservationResult result + ) { + throw new IllegalStateException("blocked listener failed"); + } + + @Override + public ReservationAccountingListenerType listenerType() { + return ReservationAccountingListenerType.METRICS; + } + }; + } + + private static InMemoryBudgetStateStore store( + List listeners, + List errorHandlers + ) { + AtomicInteger sequence = new AtomicInteger(); + return new InMemoryBudgetStateStore( + CLOCK, + () -> new ReservationId("reservation-" + sequence.incrementAndGet()), + (usage, plan) -> usd("40.00"), + listeners, + errorHandlers + ); + } + + private static BudgetReservationRequest legacyRequest(String id, Cost amount) { + return new BudgetReservationRequest(KEY, LIMIT, amount, id); + } + + private static ReservationId reserveUsage( + InMemoryBudgetStateStore store, + String requestId, + String deduplicationId, + String amount + ) { + return store.checkAndReserve(new BudgetReservationRequest( + KEY, + LIMIT, + usd(amount), + requestId, + new IdempotencyKey(deduplicationId), + pricingSnapshot(), + TOKEN_ESTIMATE + )).reservationId(); + } + + private static ActualUsageCommand command( + String requestId, + ReservationId reservationId + ) { + return new ActualUsageCommand( + requestId, + "attempt-1", + reservationId, + TokenUsage.from(100, 50), + "gpt-4o-mini-request" + ); + } + + private static PricingSnapshot pricingSnapshot() { + return new PricingSnapshot( + "gpt-4o-mini-request", + "pricing-v1", + "catalog-v1", + CLOCK.instant(), + Map.of( + TokenType.PROMPT, + new BigDecimal("0.10"), + TokenType.COMPLETION, + new BigDecimal("0.20") + ), + USD + ); + } + + private static Cost usd(String amount) { + return Cost.of(new BigDecimal(amount), USD); + } +} diff --git a/token-pilot-budget/src/test/java/io/tokenpilot/budget/internal/ReservationReconciliationTest.java b/token-pilot-budget/src/test/java/io/tokenpilot/budget/internal/ReservationReconciliationTest.java index 992aa08..025c154 100644 --- a/token-pilot-budget/src/test/java/io/tokenpilot/budget/internal/ReservationReconciliationTest.java +++ b/token-pilot-budget/src/test/java/io/tokenpilot/budget/internal/ReservationReconciliationTest.java @@ -12,6 +12,7 @@ import io.tokenpilot.budget.ReservationAccountingReason; import io.tokenpilot.budget.ReservationActualTokens; import io.tokenpilot.budget.ReservationReconciliation; +import io.tokenpilot.budget.ReservationReconciliationRequiredEvent; import io.tokenpilot.budget.ReservationStatus; import io.tokenpilot.budget.ReservationTokenEstimate; import io.tokenpilot.core.CostCalculator; @@ -22,6 +23,7 @@ import io.tokenpilot.core.domain.TokenUsage; import io.tokenpilot.core.domain.TokenUsageDetails; import io.tokenpilot.core.domain.UsageSource; +import io.tokenpilot.core.internal.LedgerComponents; import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; @@ -85,7 +87,7 @@ void calculatesActualCostOnceWithReservedPricingSnapshot() { "attempt-1", reservationId, TokenUsage.from(100, 50), - "gpt-4o-mini-response" + "gpt-4o-mini-request" ) ); @@ -220,7 +222,7 @@ void correlatesRequestReservationModelsAndPricing() { assertThat(reconciliation.budgetKey()).isEqualTo(KEY); assertThat(reconciliation.requestModelId()).isEqualTo(snapshot.modelId()); assertThat(reconciliation.responseModelId()) - .isEqualTo("gpt-4o-mini-response"); + .isEqualTo("gpt-4o-mini-request"); assertThat(reconciliation.pricingPolicyId()) .isEqualTo(snapshot.pricingPolicyId()); assertThat(reconciliation.catalogVersion()) @@ -370,7 +372,7 @@ void rejectsMismatchedRequestIdBeforeCostCalculation() { "attempt-1", reservationId, TokenUsage.from(100, 50), - "gpt-4o-mini-response" + "gpt-4o-mini-request" ) ) ).isInstanceOf(IllegalArgumentException.class) @@ -380,6 +382,223 @@ void rejectsMismatchedRequestIdBeforeCostCalculation() { assertThat(store.snapshot(KEY, LIMIT)).isEqualTo(before); } + @Test + @DisplayName("다른 response model은 예약 가격으로 계산하기 전에 거부한다") + void rejectsMismatchedResponseModelBeforeCostCalculation() { + AtomicInteger calculationCount = new AtomicInteger(); + InMemoryBudgetStateStore store = store((usage, plan) -> { + calculationCount.incrementAndGet(); + return usd("40.00"); + }); + ReservationId reservationId = reserve( + store, + pricingSnapshot(), + usd("60.00") + ); + store.markInFlight(reservationId); + var before = store.snapshot(KEY, LIMIT); + + assertThatThrownBy(() -> store.commit(new ActualUsageCommand( + "request-1", + "attempt-1", + reservationId, + TokenUsage.from(100, 50), + "provider-routed-model" + ))).isInstanceOf(IllegalArgumentException.class) + .hasMessage( + "responseModelId must match the reservation pricing snapshot" + ); + + assertThat(calculationCount).hasValue(0); + assertThat(store.snapshot(KEY, LIMIT)).isEqualTo(before); + } + + @Test + @DisplayName("pricing mismatch pending은 같은 잘못된 모델의 late actual로 확정하지 않는다") + void rejectsLateActualThatStillMismatchesReservedPricing() { + AtomicInteger calculationCount = new AtomicInteger(); + InMemoryBudgetStateStore store = store((usage, plan) -> { + calculationCount.incrementAndGet(); + return usd("40.00"); + }); + ReservationId reservationId = reserve( + store, + pricingSnapshot(), + usd("60.00") + ); + store.markInFlight(reservationId); + store.markReconciliationRequired( + reservationId, + ReservationAccountingReason.PRICING_RECONCILIATION_REQUIRED + ); + var before = store.snapshot(KEY, LIMIT); + + assertThatThrownBy(() -> store.reconcileLateActual( + new ActualUsageCommand( + "request-1", + "attempt-1", + reservationId, + TokenUsage.from(100, 50), + "provider-routed-model" + ) + )).isInstanceOf(IllegalArgumentException.class) + .hasMessage( + "responseModelId must match the reservation pricing snapshot" + ); + + assertThat(calculationCount).hasValue(0); + assertThat(store.snapshot(KEY, LIMIT)).isEqualTo(before); + assertThat(before.pendingReconciliationLiability()) + .isEqualTo(usd("60.00")); + } + + @Test + @DisplayName("response model pricing snapshot으로 pending actual을 한 번만 재정산한다") + void reconcilesPendingActualWithExplicitResponsePricingSnapshot() { + AtomicInteger calculationCount = new AtomicInteger(); + List events = new ArrayList<>(); + InMemoryBudgetStateStore store = new InMemoryBudgetStateStore( + CLOCK, + () -> new ReservationId("reservation-repriced"), + (usage, plan) -> { + calculationCount.incrementAndGet(); + assertThat(plan.modelId()).isEqualTo("provider-routed-model"); + return usd("25.00"); + }, + List.of(events::add) + ); + PricingSnapshot requestSnapshot = pricingSnapshot(); + PricingSnapshot responseSnapshot = new PricingSnapshot( + "provider-routed-model", + "provider-policy", + "catalog-v2", + CLOCK.instant().plusSeconds(1), + Map.of( + TokenType.PROMPT, new BigDecimal("0.20"), + TokenType.COMPLETION, new BigDecimal("0.40") + ), + USD + ); + ReservationId reservationId = reserve(store, requestSnapshot, usd("60.00")); + store.markInFlight(reservationId); + ActualUsageCommand command = new ActualUsageCommand( + "request-1", + "attempt-1", + reservationId, + TokenUsage.from(100, 50), + "provider-routed-model" + ); + store.markReconciliationRequired( + command, + ReservationAccountingReason.PRICING_RECONCILIATION_REQUIRED + ); + var pendingBeforeInvalidCallback = store.snapshot(KEY, LIMIT); + ActualUsageCommand conflictingCommand = new ActualUsageCommand( + "request-1", + "attempt-1", + reservationId, + TokenUsage.from(101, 50), + "provider-routed-model" + ); + assertThatThrownBy(() -> store.reconcileLateActual( + conflictingCommand, + responseSnapshot + )).isInstanceOf(IllegalArgumentException.class) + .hasMessage( + "actual usage callback conflicts with the pending pricing callback" + ); + assertThat(store.snapshot(KEY, LIMIT)).isEqualTo(pendingBeforeInvalidCallback); + + assertThatThrownBy(() -> store.reconcileLateActual( + command, + new PricingSnapshot( + "wrong-model", + "provider-policy", + "catalog-v2", + CLOCK.instant(), + responseSnapshot.rates(), + USD + ) + )).isInstanceOf(IllegalArgumentException.class) + .hasMessage("responseModelId must match the actual pricing snapshot"); + assertThat(store.snapshot(KEY, LIMIT)).isEqualTo(pendingBeforeInvalidCallback); + + ReservationReconciliation applied = store.reconcileLateActual( + command, + responseSnapshot + ); + ReservationReconciliation duplicate = store.reconcileLateActual( + command, + new PricingSnapshot( + "provider-routed-model", + "provider-policy", + "catalog-v2", + CLOCK.instant().plusSeconds(2), + responseSnapshot.rates(), + USD + ) + ); + + assertThat(calculationCount).hasValue(1); + assertThat(applied.transition().status()).isEqualTo( + io.tokenpilot.budget.AccountingTransitionStatus.APPLIED + ); + assertThat(duplicate.transition().status()).isEqualTo(REUSED); + assertThat(applied.pricingSnapshot()).isEqualTo(requestSnapshot); + assertThat(applied.actualPricingSnapshot()).isEqualTo(responseSnapshot); + assertThat(applied.requestModelId()).isEqualTo(requestSnapshot.modelId()); + assertThat(applied.responseModelId()).isEqualTo(responseSnapshot.modelId()); + assertThat(applied.actual()).isEqualTo(usd("25.00")); + assertThat(events).containsExactly(new ReservationAccountingEvent(applied)); + assertThat(store.snapshot(KEY, LIMIT).pendingReconciliationLiability()) + .isEqualTo(usd("0.00")); + assertThat(store.snapshot(KEY, LIMIT).committedCost()) + .isEqualTo(usd("25.00")); + } + + @Test + @DisplayName("pricing mismatch event는 provider actual과 response model을 보존한다") + void preservesPendingActualInPricingReconciliationEvent() { + List events = new ArrayList<>(); + InMemoryBudgetStateStore store = new InMemoryBudgetStateStore( + CLOCK, + () -> new ReservationId("reservation-pending-pricing"), + LedgerComponents.defaultCostCalculator(), + List.of(new io.tokenpilot.budget.ReservationAccountingListener() { + @Override + public void onCommitted(ReservationAccountingEvent event) { + } + + @Override + public void onReconciliationRequired( + ReservationReconciliationRequiredEvent event + ) { + events.add(event); + } + }) + ); + ReservationId reservationId = reserve(store, pricingSnapshot(), usd("60.00")); + store.markInFlight(reservationId); + ActualUsageCommand command = new ActualUsageCommand( + "request-1", + "attempt-1", + reservationId, + TokenUsage.from(100, 50), + "provider-routed-model" + ); + + store.markReconciliationRequired( + command, + ReservationAccountingReason.PRICING_RECONCILIATION_REQUIRED + ); + + assertThat(events).hasSize(1); + assertThat(events.getFirst().pendingActualUsage()).contains(command); + assertThat(events.getFirst().reason()).isEqualTo( + ReservationAccountingReason.PRICING_RECONCILIATION_REQUIRED + ); + } + @Test @DisplayName("계산된 actual 통화가 예약 가격 통화와 다르면 상태를 변경하지 않는다") void rejectsCalculatedCostWithUnexpectedCurrency() { @@ -411,7 +630,7 @@ void rejectsUnavailableActualUsage() { "attempt-1", new ReservationId("reservation-1"), TokenUsage.unavailable(Map.of()), - "gpt-4o-mini-response" + "gpt-4o-mini-request" ) ).isInstanceOf(IllegalArgumentException.class) .hasMessage("usage must be available for actual reconciliation"); @@ -439,7 +658,7 @@ void rejectsNonProviderUsageSources(UsageSource source) { "attempt-1", new ReservationId("reservation-1"), usage, - "gpt-4o-mini-response" + "gpt-4o-mini-request" ) ).isInstanceOf(IllegalArgumentException.class); } @@ -593,7 +812,7 @@ void calculatesChangedCallbackBeforeReportingConflict() { "attempt-1", reservationId, TokenUsage.from(101, 50), - "gpt-4o-mini-response" + "gpt-4o-mini-request" ) ); @@ -625,7 +844,7 @@ void reportsConflictWhenChangedCallbackHasSameCost() { "attempt-1", reservationId, TokenUsage.from(101, 50), - "gpt-4o-mini-response" + "gpt-4o-mini-request" ) ); @@ -665,7 +884,7 @@ void publishesAccountingEventOnlyForNewlyAppliedCommit() { "attempt-1", reservationId, TokenUsage.from(101, 50), - "gpt-4o-mini-response" + "gpt-4o-mini-request" ) ); @@ -822,7 +1041,7 @@ private static ActualUsageCommand command(ReservationId reservationId) { "attempt-1", reservationId, TokenUsage.from(100, 50), - "gpt-4o-mini-response" + "gpt-4o-mini-request" ); } diff --git a/token-pilot-core/src/main/java/io/tokenpilot/core/CoreComponents.java b/token-pilot-core/src/main/java/io/tokenpilot/core/CoreComponents.java index 63a8cc7..b3750c3 100644 --- a/token-pilot-core/src/main/java/io/tokenpilot/core/CoreComponents.java +++ b/token-pilot-core/src/main/java/io/tokenpilot/core/CoreComponents.java @@ -4,6 +4,7 @@ import io.tokenpilot.core.internal.LedgerComponents; import java.util.Collection; +import java.util.List; /** * framework-independent core component entrypoint입니다. @@ -25,7 +26,18 @@ public static ModelRegistry defaultModelRegistry() { return LedgerComponents.defaultModelRegistry(); } + public static PreflightCostEstimator defaultPreflightCostEstimator() { + return LedgerComponents.defaultPreflightCostEstimator(); + } + public static TokenBudget tokenBudget(ModelRegistry modelRegistry) { return LedgerComponents.tokenBudget(modelRegistry); } + + public static TokenBudget tokenBudget( + ModelRegistry modelRegistry, + List listeners + ) { + return LedgerComponents.tokenBudget(modelRegistry, listeners); + } } diff --git a/token-pilot-core/src/main/java/io/tokenpilot/core/PreflightDecisionListener.java b/token-pilot-core/src/main/java/io/tokenpilot/core/PreflightDecisionListener.java new file mode 100644 index 0000000..4621305 --- /dev/null +++ b/token-pilot-core/src/main/java/io/tokenpilot/core/PreflightDecisionListener.java @@ -0,0 +1,10 @@ +package io.tokenpilot.core; + +import io.tokenpilot.core.domain.PreflightDecisionEvent; + +/** context admission 판정을 best-effort로 관찰하는 framework-independent 계약입니다. */ +@FunctionalInterface +public interface PreflightDecisionListener { + + void onDecision(PreflightDecisionEvent event); +} diff --git a/token-pilot-core/src/main/java/io/tokenpilot/core/PricingMissingListener.java b/token-pilot-core/src/main/java/io/tokenpilot/core/PricingMissingListener.java new file mode 100644 index 0000000..6d23f32 --- /dev/null +++ b/token-pilot-core/src/main/java/io/tokenpilot/core/PricingMissingListener.java @@ -0,0 +1,10 @@ +package io.tokenpilot.core; + +import io.tokenpilot.core.domain.PricingMissingEvent; + +/** missing-pricing 판정을 best-effort로 관찰하는 framework-independent 계약입니다. */ +@FunctionalInterface +public interface PricingMissingListener { + + void onPricingMissing(PricingMissingEvent event); +} diff --git a/token-pilot-core/src/main/java/io/tokenpilot/core/domain/PreflightDecisionEvent.java b/token-pilot-core/src/main/java/io/tokenpilot/core/domain/PreflightDecisionEvent.java new file mode 100644 index 0000000..efd3eea --- /dev/null +++ b/token-pilot-core/src/main/java/io/tokenpilot/core/domain/PreflightDecisionEvent.java @@ -0,0 +1,15 @@ +package io.tokenpilot.core.domain; + +import java.util.Objects; + +/** + * provider 호출 전 context admission 판정을 전달하는 framework-independent 이벤트입니다. + * + * @param result bounded status/reason과 판정 metadata를 포함한 immutable 결과 + */ +public record PreflightDecisionEvent(BudgetResult result) { + + public PreflightDecisionEvent { + Objects.requireNonNull(result, "result must not be null"); + } +} diff --git a/token-pilot-core/src/main/java/io/tokenpilot/core/domain/PricingMissingEvent.java b/token-pilot-core/src/main/java/io/tokenpilot/core/domain/PricingMissingEvent.java new file mode 100644 index 0000000..94580f1 --- /dev/null +++ b/token-pilot-core/src/main/java/io/tokenpilot/core/domain/PricingMissingEvent.java @@ -0,0 +1,23 @@ +package io.tokenpilot.core.domain; + +import java.util.Objects; + +/** + * 명시적인 missing-pricing 판정을 전달하는 framework-independent 이벤트입니다. + * + * @param policy missing pricing 처리 정책 + * @param resolution missing plan/rate 판정 + */ +public record PricingMissingEvent( + MissingPricingPolicy policy, + PricingResolution resolution +) { + + public PricingMissingEvent { + Objects.requireNonNull(policy, "policy must not be null"); + Objects.requireNonNull(resolution, "resolution must not be null"); + if (!resolution.isMissing()) { + throw new IllegalArgumentException("resolution must identify missing pricing"); + } + } +} diff --git a/token-pilot-core/src/main/java/io/tokenpilot/core/domain/PricingResolution.java b/token-pilot-core/src/main/java/io/tokenpilot/core/domain/PricingResolution.java index 2a6b6fe..48153fd 100644 --- a/token-pilot-core/src/main/java/io/tokenpilot/core/domain/PricingResolution.java +++ b/token-pilot-core/src/main/java/io/tokenpilot/core/domain/PricingResolution.java @@ -9,4 +9,8 @@ public enum PricingResolution { public boolean isResolved() { return this == RESOLVED; } + + public boolean isMissing() { + return this == MISSING_PLAN || this == MISSING_RATE; + } } diff --git a/token-pilot-core/src/main/java/io/tokenpilot/core/domain/PricingSnapshot.java b/token-pilot-core/src/main/java/io/tokenpilot/core/domain/PricingSnapshot.java index 3d71c4c..a945d1b 100644 --- a/token-pilot-core/src/main/java/io/tokenpilot/core/domain/PricingSnapshot.java +++ b/token-pilot-core/src/main/java/io/tokenpilot/core/domain/PricingSnapshot.java @@ -55,4 +55,62 @@ public static PricingSnapshot from(PricingPlan plan, String catalogVersion, Inst plan.currency() ); } + + /** + * 가격 조회 시각을 제외하고 실제 비용 책임을 결정하는 terms가 같은지 반환합니다. + * {@link #checkedAt()}은 같은 요청의 재시도에서 달라질 수 있는 관측 metadata이므로 + * 멱등성 fingerprint에는 포함하지 않습니다. + */ + public boolean hasSameBillingTerms(PricingSnapshot other) { + return haveSameBillingTerms(this, other); + } + + /** + * 두 snapshot의 비용 책임 terms를 비교합니다. 테스트 double처럼 rate map이 없는 + * 호환 snapshot은 양쪽 모두 없을 때 같은 것으로 취급합니다. + */ + public static boolean haveSameBillingTerms( + PricingSnapshot first, + PricingSnapshot second + ) { + Objects.requireNonNull(first, "first must not be null"); + Objects.requireNonNull(second, "second must not be null"); + if (!Objects.equals(first.modelId(), second.modelId()) + || !Objects.equals(first.pricingPolicyId(), second.pricingPolicyId()) + || !Objects.equals(first.catalogVersion(), second.catalogVersion()) + || !Objects.equals(first.currency(), second.currency())) { + return false; + } + Map firstRates = first.rates(); + Map secondRates = second.rates(); + if (firstRates == null || secondRates == null) { + return firstRates == secondRates; + } + if (!firstRates.keySet().equals(secondRates.keySet())) { + return false; + } + return firstRates.entrySet().stream().allMatch(entry -> + entry.getValue().compareTo(secondRates.get(entry.getKey())) == 0 + ); + } + + /** + * {@link #hasSameBillingTerms(PricingSnapshot)}와 일치하는 bounded hash code입니다. + */ + public int billingTermsHashCode() { + int result = Objects.hash(modelId, pricingPolicyId, catalogVersion, currency); + if (rates == null) { + return result; + } + for (TokenType tokenType : TokenType.values()) { + BigDecimal rate = rates.get(tokenType); + if (rate != null) { + result = 31 * result + Objects.hash( + tokenType, + rate.stripTrailingZeros() + ); + } + } + return result; + } } diff --git a/token-pilot-core/src/main/java/io/tokenpilot/core/internal/DefaultLedgerManager.java b/token-pilot-core/src/main/java/io/tokenpilot/core/internal/DefaultLedgerManager.java index e9e66f4..3cba97f 100644 --- a/token-pilot-core/src/main/java/io/tokenpilot/core/internal/DefaultLedgerManager.java +++ b/token-pilot-core/src/main/java/io/tokenpilot/core/internal/DefaultLedgerManager.java @@ -80,7 +80,13 @@ private void publish(String modelId, TokenUsage usage, Cost cost, Map listener.onRecord(event)); + for (LedgerListener listener : listeners) { + try { + listener.onRecord(event); + } catch (RuntimeException ignored) { + // Optional observers do not change an already calculated ledger result. + } + } } } } diff --git a/token-pilot-core/src/main/java/io/tokenpilot/core/internal/DefaultTokenBudget.java b/token-pilot-core/src/main/java/io/tokenpilot/core/internal/DefaultTokenBudget.java index 01f9165..da34bed 100644 --- a/token-pilot-core/src/main/java/io/tokenpilot/core/internal/DefaultTokenBudget.java +++ b/token-pilot-core/src/main/java/io/tokenpilot/core/internal/DefaultTokenBudget.java @@ -1,14 +1,17 @@ package io.tokenpilot.core.internal; import io.tokenpilot.core.ModelRegistry; +import io.tokenpilot.core.PreflightDecisionListener; import io.tokenpilot.core.TokenBudget; import io.tokenpilot.core.domain.AdmissionReason; import io.tokenpilot.core.domain.AdmissionStatus; import io.tokenpilot.core.domain.BudgetResult; import io.tokenpilot.core.domain.ModelDefinition; +import io.tokenpilot.core.domain.PreflightDecisionEvent; import io.tokenpilot.core.domain.TokenCountResult; import io.tokenpilot.core.domain.TokenCountScope; +import java.util.List; import java.util.Objects; import java.util.Optional; import java.util.OptionalLong; @@ -19,9 +22,20 @@ final class DefaultTokenBudget implements TokenBudget { private final ModelRegistry modelRegistry; + private final List listeners; DefaultTokenBudget(ModelRegistry modelRegistry) { + this(modelRegistry, List.of()); + } + + DefaultTokenBudget( + ModelRegistry modelRegistry, + List listeners + ) { this.modelRegistry = Objects.requireNonNull(modelRegistry, "modelRegistry must not be null"); + this.listeners = List.copyOf( + Objects.requireNonNull(listeners, "listeners must not be null") + ); } @Override @@ -146,7 +160,7 @@ private BudgetResult buildResult( OptionalLong maxContextTokens, OptionalLong remainingTokens ) { - return new BudgetResult( + BudgetResult result = new BudgetResult( status, reason, canonicalModelId, @@ -158,5 +172,21 @@ private BudgetResult buildResult( input.estimatorDescriptor(), input.tokenizationBasis() ); + publishBestEffort(result); + return result; + } + + private void publishBestEffort(BudgetResult result) { + if (listeners.isEmpty()) { + return; + } + PreflightDecisionEvent event = new PreflightDecisionEvent(result); + for (PreflightDecisionListener listener : listeners) { + try { + listener.onDecision(event); + } catch (RuntimeException ignored) { + // Optional observers do not change a completed admission decision. + } + } } } diff --git a/token-pilot-core/src/main/java/io/tokenpilot/core/internal/LedgerComponents.java b/token-pilot-core/src/main/java/io/tokenpilot/core/internal/LedgerComponents.java index 461f71c..6267212 100644 --- a/token-pilot-core/src/main/java/io/tokenpilot/core/internal/LedgerComponents.java +++ b/token-pilot-core/src/main/java/io/tokenpilot/core/internal/LedgerComponents.java @@ -8,6 +8,7 @@ import io.tokenpilot.core.PricingProvider; import io.tokenpilot.core.PricingRegistry; import io.tokenpilot.core.PreflightCostEstimator; +import io.tokenpilot.core.PreflightDecisionListener; import io.tokenpilot.core.TokenBudget; import io.tokenpilot.core.TokenEstimator; import io.tokenpilot.core.domain.ModelDefinition; @@ -63,6 +64,16 @@ public static TokenBudget tokenBudget(ModelRegistry modelRegistry) { return new DefaultTokenBudget(modelRegistry); } + /** + * model registry와 best-effort decision listener를 사용하는 context admission evaluator를 생성합니다. + */ + public static TokenBudget tokenBudget( + ModelRegistry modelRegistry, + List listeners + ) { + return new DefaultTokenBudget(modelRegistry, listeners); + } + /** * 등록할 model definition으로 in-memory model registry를 생성합니다. * diff --git a/token-pilot-core/src/test/java/io/tokenpilot/core/domain/PricingMissingEventTest.java b/token-pilot-core/src/test/java/io/tokenpilot/core/domain/PricingMissingEventTest.java new file mode 100644 index 0000000..9f9c13e --- /dev/null +++ b/token-pilot-core/src/test/java/io/tokenpilot/core/domain/PricingMissingEventTest.java @@ -0,0 +1,38 @@ +package io.tokenpilot.core.domain; + +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +class PricingMissingEventTest { + + @Test + @DisplayName("missing pricing policy와 bounded resolution을 보존한다") + void preservesPolicyAndResolution() { + PricingMissingEvent event = new PricingMissingEvent( + MissingPricingPolicy.FAIL_CLOSED, + PricingResolution.MISSING_RATE + ); + + assertThat(event.policy()).isEqualTo(MissingPricingPolicy.FAIL_CLOSED); + assertThat(event.resolution()).isEqualTo(PricingResolution.MISSING_RATE); + } + + @Test + @DisplayName("resolved와 currency mismatch는 missing pricing event가 될 수 없다") + void rejectsNonMissingPricing() { + for (PricingResolution resolution : new PricingResolution[] { + PricingResolution.RESOLVED, + PricingResolution.CURRENCY_MISMATCH + }) { + assertThatThrownBy(() -> new PricingMissingEvent( + MissingPricingPolicy.FAIL_OPEN, + resolution + )) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("resolution must identify missing pricing"); + } + } +} diff --git a/token-pilot-core/src/test/java/io/tokenpilot/core/domain/PricingSnapshotTest.java b/token-pilot-core/src/test/java/io/tokenpilot/core/domain/PricingSnapshotTest.java index 0245eb2..f7e19b0 100644 --- a/token-pilot-core/src/test/java/io/tokenpilot/core/domain/PricingSnapshotTest.java +++ b/token-pilot-core/src/test/java/io/tokenpilot/core/domain/PricingSnapshotTest.java @@ -76,4 +76,35 @@ void preserveEmptyRates() { assertThat(snapshot.rates()).isEmpty(); } + + @Test + @DisplayName("멱등성 비교는 조회 시각과 decimal scale을 가격 책임에서 제외해야 한다") + void comparesBillingTermsWithoutObservationMetadata() { + PricingSnapshot first = new PricingSnapshot( + "gpt-4o", + "standard", + "catalog-v1", + Instant.parse("2026-07-30T00:00:00Z"), + Map.of( + TokenType.PROMPT, new BigDecimal("0.010"), + TokenType.COMPLETION, new BigDecimal("0.030") + ), + Currency.getInstance("USD") + ); + PricingSnapshot second = new PricingSnapshot( + "gpt-4o", + "standard", + "catalog-v1", + Instant.parse("2026-07-30T00:01:00Z"), + Map.of( + TokenType.PROMPT, new BigDecimal("0.01"), + TokenType.COMPLETION, new BigDecimal("0.03") + ), + Currency.getInstance("USD") + ); + + assertThat(first.hasSameBillingTerms(second)).isTrue(); + assertThat(first.billingTermsHashCode()) + .isEqualTo(second.billingTermsHashCode()); + } } diff --git a/token-pilot-core/src/test/java/io/tokenpilot/core/internal/DefaultLedgerManagerTest.java b/token-pilot-core/src/test/java/io/tokenpilot/core/internal/DefaultLedgerManagerTest.java index ab7d258..56617ca 100644 --- a/token-pilot-core/src/test/java/io/tokenpilot/core/internal/DefaultLedgerManagerTest.java +++ b/token-pilot-core/src/test/java/io/tokenpilot/core/internal/DefaultLedgerManagerTest.java @@ -132,4 +132,40 @@ void shouldRecordZeroCostWithExplicitZeroRateSnapshot() { event.cost().equals(cost) )); } + + @Test + @DisplayName("한 리스너의 RuntimeException은 비용 결과와 다음 리스너 호출에 영향을 주지 않아야 한다") + void shouldIsolateListenerRuntimeExceptionAndContinuePublishing() { + PricingPlan plan = new PricingPlan( + "gpt-4o", + new BigDecimal("5.0"), + new BigDecimal("15.0") + ); + registry.registerPlan(plan); + LedgerListener failingListener = mock(LedgerListener.class); + LedgerListener succeedingListener = mock(LedgerListener.class); + DefaultLedgerManager isolatedManager = new DefaultLedgerManager( + registry, + calculator, + List.of(failingListener, succeedingListener) + ); + TokenUsage usage = TokenUsage.from(1_000, 1_000); + doThrow(new IllegalStateException("listener failed")) + .when(failingListener) + .onRecord(any(CostRecordedEvent.class)); + + Cost cost = isolatedManager.record("gpt-4o", usage, Map.of()); + + assertThat(cost).isEqualTo(Cost.of( + new BigDecimal("20.000000"), + Currency.getInstance("USD") + )); + var inOrder = inOrder(failingListener, succeedingListener); + inOrder.verify(failingListener).onRecord(any(CostRecordedEvent.class)); + inOrder.verify(succeedingListener).onRecord(argThat(event -> + event.modelId().equals("gpt-4o") + && event.usage().equals(usage) + && event.cost().equals(cost) + )); + } } diff --git a/token-pilot-core/src/test/java/io/tokenpilot/core/internal/DefaultTokenBudgetTest.java b/token-pilot-core/src/test/java/io/tokenpilot/core/internal/DefaultTokenBudgetTest.java index 4771dcb..e826809 100644 --- a/token-pilot-core/src/test/java/io/tokenpilot/core/internal/DefaultTokenBudgetTest.java +++ b/token-pilot-core/src/test/java/io/tokenpilot/core/internal/DefaultTokenBudgetTest.java @@ -1,11 +1,13 @@ package io.tokenpilot.core.internal; +import io.tokenpilot.core.CoreComponents; import io.tokenpilot.core.ModelRegistry; import io.tokenpilot.core.TokenBudget; import io.tokenpilot.core.domain.AdmissionReason; import io.tokenpilot.core.domain.AdmissionStatus; import io.tokenpilot.core.domain.BudgetResult; import io.tokenpilot.core.domain.ModelDefinition; +import io.tokenpilot.core.domain.PreflightDecisionEvent; import io.tokenpilot.core.domain.TokenCountAccuracy; import io.tokenpilot.core.domain.TokenCountResult; import io.tokenpilot.core.domain.TokenCountScope; @@ -18,6 +20,8 @@ import java.time.Instant; import java.util.List; import java.util.Set; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.atomic.AtomicInteger; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; @@ -160,6 +164,49 @@ void rejectsNegativeReservedOutputAndRequiresFits() { .hasMessageContaining("INCOMPLETE_SCOPE"); } + @Test + @DisplayName("모든 admission 결과를 immutable event로 한 번씩 발행한다") + void publishesEveryAdmissionResultOnce() { + List events = new CopyOnWriteArrayList<>(); + TokenBudget observed = CoreComponents.tokenBudget( + registry, + List.of(events::add) + ); + + BudgetResult fits = observed.check("model-v1", request(6, 6), 4); + BudgetResult exceeds = observed.check("model-v1", request(6, 6), 5); + BudgetResult indeterminate = observed.check("unknown", request(1, 1), 0); + + assertThat(events).containsExactly( + new PreflightDecisionEvent(fits), + new PreflightDecisionEvent(exceeds), + new PreflightDecisionEvent(indeterminate) + ); + } + + @Test + @DisplayName("한 preflight listener 실패가 결과나 다음 listener를 바꾸지 않는다") + void isolatesPreflightListenerFailures() { + AtomicInteger failedDeliveries = new AtomicInteger(); + List received = new CopyOnWriteArrayList<>(); + TokenBudget observed = LedgerComponents.tokenBudget( + registry, + List.of( + event -> { + failedDeliveries.incrementAndGet(); + throw new IllegalStateException("listener failed"); + }, + received::add + ) + ); + + BudgetResult result = observed.check("model-v1", request(6, 6), 4); + + assertThat(result.status()).isEqualTo(AdmissionStatus.FITS); + assertThat(failedDeliveries).hasValue(1); + assertThat(received).containsExactly(new PreflightDecisionEvent(result)); + } + private static TokenCountResult request(long tokens, long safeUpperBound) { return counted(tokens, safeUpperBound, TokenCountScope.REQUEST); } diff --git a/token-pilot-micrometer/build.gradle b/token-pilot-micrometer/build.gradle index 4f836b2..9a31486 100644 --- a/token-pilot-micrometer/build.gradle +++ b/token-pilot-micrometer/build.gradle @@ -1,5 +1,9 @@ dependencies { api project(':token-pilot-core') - implementation 'io.micrometer:micrometer-observation:1.17.0' - implementation 'io.micrometer:micrometer-core:1.17.0' + api 'io.micrometer:micrometer-core:1.17.0' + compileOnly project(':token-pilot-budget') + compileOnly project(':token-pilot-notification') + + testImplementation project(':token-pilot-budget') + testImplementation project(':token-pilot-notification') } diff --git a/token-pilot-micrometer/src/main/java/io/tokenpilot/micrometer/MetricsOptions.java b/token-pilot-micrometer/src/main/java/io/tokenpilot/micrometer/MetricsOptions.java index 01cf3ea..343e470 100644 --- a/token-pilot-micrometer/src/main/java/io/tokenpilot/micrometer/MetricsOptions.java +++ b/token-pilot-micrometer/src/main/java/io/tokenpilot/micrometer/MetricsOptions.java @@ -6,23 +6,49 @@ * Micrometer metric publishing options. * * @param allowedTagKeys metric tag keys allowed beyond the built-in low-cardinality tags + * @param legacyAiTokenMetricsEnabled whether the legacy {@code ai.token.*} meters are enabled */ -public record MetricsOptions(Set allowedTagKeys) { +public record MetricsOptions( + Set allowedTagKeys, + boolean legacyAiTokenMetricsEnabled +) { - public static final Set DEFAULT_ALLOWED_TAG_KEYS = Set.of("tenant_id"); + public static final Set DEFAULT_ALLOWED_TAG_KEYS = Set.of(); + private static final Set LEGACY_DEFAULT_ALLOWED_TAG_KEYS = + Set.of("tenant_id"); public MetricsOptions { allowedTagKeys = normalize(allowedTagKeys); } + /** + * 기존 직접 생성 코드의 source/binary 호환성을 유지하는 생성자입니다. + * + *

이 생성자를 직접 사용하는 것은 legacy {@code ai.token.*} meter에 대한 명시적 + * opt-in으로 취급됩니다. Starter 자동 설정 기본값은 {@link #defaults()}를 사용하여 + * legacy meter를 비활성화합니다.

+ */ + public MetricsOptions(Set allowedTagKeys) { + this( + allowedTagKeys == null + ? LEGACY_DEFAULT_ALLOWED_TAG_KEYS + : allowedTagKeys, + true + ); + } + public static MetricsOptions defaults() { - return new MetricsOptions(DEFAULT_ALLOWED_TAG_KEYS); + return new MetricsOptions(DEFAULT_ALLOWED_TAG_KEYS, false); } public static MetricsOptions withAllowedTagKeys(Set allowedTagKeys) { return new MetricsOptions(allowedTagKeys); } + public static MetricsOptions legacyDefaults() { + return new MetricsOptions(LEGACY_DEFAULT_ALLOWED_TAG_KEYS, true); + } + private static Set normalize(Set allowedTagKeys) { if (allowedTagKeys == null) { return DEFAULT_ALLOWED_TAG_KEYS; diff --git a/token-pilot-micrometer/src/main/java/io/tokenpilot/micrometer/internal/BudgetMetricsPublisher.java b/token-pilot-micrometer/src/main/java/io/tokenpilot/micrometer/internal/BudgetMetricsPublisher.java new file mode 100644 index 0000000..8b46648 --- /dev/null +++ b/token-pilot-micrometer/src/main/java/io/tokenpilot/micrometer/internal/BudgetMetricsPublisher.java @@ -0,0 +1,161 @@ +package io.tokenpilot.micrometer.internal; + +import io.micrometer.core.instrument.Counter; +import io.micrometer.core.instrument.DistributionSummary; +import io.micrometer.core.instrument.MeterRegistry; +import io.tokenpilot.budget.AccountingTransitionStatus; +import io.tokenpilot.budget.BudgetReservationRequest; +import io.tokenpilot.budget.BudgetReservationResult; +import io.tokenpilot.budget.BudgetSnapshot; +import io.tokenpilot.budget.ReservationAccountingEvent; +import io.tokenpilot.budget.ReservationAccountingListener; +import io.tokenpilot.budget.ReservationAccountingListenerErrorHandler; +import io.tokenpilot.budget.ReservationAccountingListenerFailureEvent; +import io.tokenpilot.budget.ReservationAccountingListenerType; +import io.tokenpilot.budget.ReservationAccountingReason; +import io.tokenpilot.budget.ReservationReconciliation; +import io.tokenpilot.budget.ReservationReconciliationRequiredEvent; +import io.tokenpilot.budget.ReservationState; + +import java.math.BigInteger; +import java.util.Objects; + +/** Atomic budget/accounting 결과를 Token Pilot 고유 Micrometer meter로 투영합니다. */ +public final class BudgetMetricsPublisher + implements ReservationAccountingListener, ReservationAccountingListenerErrorHandler { + + private final MeterRegistry meterRegistry; + + public BudgetMetricsPublisher(MeterRegistry meterRegistry) { + this.meterRegistry = Objects.requireNonNull( + meterRegistry, + "meterRegistry must not be null" + ); + } + + @Override + public void onCommitted(ReservationAccountingEvent event) { + recordCommitted(event); + } + + @Override + public void onAccountingApplied( + ReservationAccountingEvent event, + BudgetSnapshot snapshot + ) { + Objects.requireNonNull(snapshot, "snapshot must not be null"); + recordCommitted(event); + } + + @Override + public void onReconciliationRequired( + ReservationReconciliationRequiredEvent event + ) { + Objects.requireNonNull(event, "event must not be null"); + if (event.transition().status() != AccountingTransitionStatus.APPLIED + || event.transition().resultingState() + != ReservationState.RECONCILIATION_REQUIRED) { + return; + } + recordReconciliationOutcome( + event.transition().resultingState(), + event.reason() + ); + } + + @Override + public void onReservationEvaluated( + BudgetReservationRequest request, + BudgetReservationResult result + ) { + Objects.requireNonNull(request, "request must not be null"); + Objects.requireNonNull(result, "result must not be null"); + Counter.builder(TokenPilotMeterMetadata.BUDGET_RESERVATIONS) + .description(TokenPilotMeterMetadata.BUDGET_RESERVATIONS_DESCRIPTION) + .baseUnit(TokenPilotMeterMetadata.RESERVATIONS_BASE_UNIT) + .tag("state", TokenPilotMeterMetadata.tagValue(result.status())) + .register(meterRegistry) + .increment(); + } + + @Override + public ReservationAccountingListenerType listenerType() { + return ReservationAccountingListenerType.METRICS; + } + + @Override + public void onFailure(ReservationAccountingListenerFailureEvent event) { + Objects.requireNonNull(event, "event must not be null"); + Counter.builder(TokenPilotMeterMetadata.LISTENER_FAILURES) + .description(TokenPilotMeterMetadata.LISTENER_FAILURES_DESCRIPTION) + .baseUnit(TokenPilotMeterMetadata.FAILURES_BASE_UNIT) + .tag("listener", TokenPilotMeterMetadata.tagValue(event.listenerType())) + .tag("phase", TokenPilotMeterMetadata.tagValue(event.phase())) + .register(meterRegistry) + .increment(); + } + + private void recordCommitted(ReservationAccountingEvent event) { + Objects.requireNonNull(event, "event must not be null"); + ReservationReconciliation reconciliation = event.reconciliation(); + if (reconciliation.transition().status() != AccountingTransitionStatus.APPLIED + || reconciliation.transition().resultingState() + != ReservationState.COMMITTED) { + return; + } + + Counter.builder(TokenPilotMeterMetadata.COST_TOTAL) + .description(TokenPilotMeterMetadata.COST_TOTAL_DESCRIPTION) + .baseUnit(TokenPilotMeterMetadata.COST_BASE_UNIT) + .tag("currency", reconciliation.currency().getCurrencyCode()) + .register(meterRegistry) + .increment(reconciliation.actual().value().doubleValue()); + + BigInteger delta = totalTokenDelta(reconciliation); + String direction = delta.signum() > 0 + ? "underestimate" + : delta.signum() < 0 ? "overestimate" : "exact"; + DistributionSummary.builder(TokenPilotMeterMetadata.RECONCILIATION_ERROR_TOKENS) + .description(TokenPilotMeterMetadata.RECONCILIATION_ERROR_TOKENS_DESCRIPTION) + .baseUnit(TokenPilotMeterMetadata.TOKENS_BASE_UNIT) + .tag("direction", direction) + .register(meterRegistry) + .record(delta.abs().doubleValue()); + + recordReconciliationOutcome( + reconciliation.transition().resultingState(), + reconciliation.reason() + ); + } + + private void recordReconciliationOutcome( + ReservationState outcome, + ReservationAccountingReason reason + ) { + Counter.builder(TokenPilotMeterMetadata.RECONCILIATION_OUTCOMES) + .description(TokenPilotMeterMetadata.RECONCILIATION_OUTCOMES_DESCRIPTION) + .baseUnit(TokenPilotMeterMetadata.RECONCILIATIONS_BASE_UNIT) + .tag("outcome", TokenPilotMeterMetadata.tagValue(outcome)) + .tag("reason", TokenPilotMeterMetadata.tagValue(reason)) + .register(meterRegistry) + .increment(); + } + + private static BigInteger totalTokenDelta( + ReservationReconciliation reconciliation + ) { + BigInteger actual = BigInteger.valueOf( + reconciliation.actualTokens().inputTokens() + ) + .add(BigInteger.valueOf( + reconciliation.actualTokens().outputTokens() + )); + BigInteger estimated = BigInteger.valueOf( + reconciliation.tokenEstimate().inputEstimatedTokens() + ) + .add(BigInteger.valueOf( + reconciliation.tokenEstimate().reservedOutputTokens() + )); + return actual.subtract(estimated); + } +} diff --git a/token-pilot-micrometer/src/main/java/io/tokenpilot/micrometer/internal/BudgetMicrometerComponents.java b/token-pilot-micrometer/src/main/java/io/tokenpilot/micrometer/internal/BudgetMicrometerComponents.java new file mode 100644 index 0000000..754185a --- /dev/null +++ b/token-pilot-micrometer/src/main/java/io/tokenpilot/micrometer/internal/BudgetMicrometerComponents.java @@ -0,0 +1,16 @@ +package io.tokenpilot.micrometer.internal; + +import io.micrometer.core.instrument.MeterRegistry; + +/** Atomic budget owner의 선택 Micrometer publisher 팩토리입니다. */ +public final class BudgetMicrometerComponents { + + private BudgetMicrometerComponents() { + } + + public static BudgetMetricsPublisher budgetMetricsPublisher( + MeterRegistry meterRegistry + ) { + return new BudgetMetricsPublisher(meterRegistry); + } +} diff --git a/token-pilot-micrometer/src/main/java/io/tokenpilot/micrometer/internal/CoreMetricsPublisher.java b/token-pilot-micrometer/src/main/java/io/tokenpilot/micrometer/internal/CoreMetricsPublisher.java new file mode 100644 index 0000000..7672c76 --- /dev/null +++ b/token-pilot-micrometer/src/main/java/io/tokenpilot/micrometer/internal/CoreMetricsPublisher.java @@ -0,0 +1,48 @@ +package io.tokenpilot.micrometer.internal; + +import io.micrometer.core.instrument.Counter; +import io.micrometer.core.instrument.MeterRegistry; +import io.tokenpilot.core.PreflightDecisionListener; +import io.tokenpilot.core.PricingMissingListener; +import io.tokenpilot.core.domain.PreflightDecisionEvent; +import io.tokenpilot.core.domain.PricingMissingEvent; + +import java.util.Objects; + +/** Token Pilot core의 preflight와 pricing-miss 이벤트를 Micrometer meter로 투영합니다. */ +public final class CoreMetricsPublisher + implements PreflightDecisionListener, PricingMissingListener { + + private final MeterRegistry meterRegistry; + + public CoreMetricsPublisher(MeterRegistry meterRegistry) { + this.meterRegistry = Objects.requireNonNull( + meterRegistry, + "meterRegistry must not be null" + ); + } + + @Override + public void onDecision(PreflightDecisionEvent event) { + Objects.requireNonNull(event, "event must not be null"); + var result = event.result(); + Counter.builder(TokenPilotMeterMetadata.PREFLIGHT_REQUESTS) + .description(TokenPilotMeterMetadata.PREFLIGHT_REQUESTS_DESCRIPTION) + .baseUnit(TokenPilotMeterMetadata.REQUESTS_BASE_UNIT) + .tag("decision", TokenPilotMeterMetadata.tagValue(result.status())) + .tag("reason", TokenPilotMeterMetadata.tagValue(result.reason())) + .register(meterRegistry) + .increment(); + } + + @Override + public void onPricingMissing(PricingMissingEvent event) { + Objects.requireNonNull(event, "event must not be null"); + Counter.builder(TokenPilotMeterMetadata.PRICING_MISSING) + .description(TokenPilotMeterMetadata.PRICING_MISSING_DESCRIPTION) + .baseUnit(TokenPilotMeterMetadata.EVENTS_BASE_UNIT) + .tag("policy", TokenPilotMeterMetadata.tagValue(event.policy())) + .register(meterRegistry) + .increment(); + } +} diff --git a/token-pilot-micrometer/src/main/java/io/tokenpilot/micrometer/internal/LedgerMicrometerComponents.java b/token-pilot-micrometer/src/main/java/io/tokenpilot/micrometer/internal/LedgerMicrometerComponents.java index 8fde10a..404d760 100644 --- a/token-pilot-micrometer/src/main/java/io/tokenpilot/micrometer/internal/LedgerMicrometerComponents.java +++ b/token-pilot-micrometer/src/main/java/io/tokenpilot/micrometer/internal/LedgerMicrometerComponents.java @@ -21,4 +21,16 @@ public static LedgerListener microCostMetricsPublisher(MeterRegistry meterRegist public static LedgerListener microCostMetricsPublisher(MeterRegistry meterRegistry, MetricsOptions options) { return new MicroCostMetricsPublisher(meterRegistry, options); } + + public static MicroCostMetricsPublisher legacyMicroCostMetricsPublisher( + MeterRegistry meterRegistry, + MetricsOptions options + ) { + return new MicroCostMetricsPublisher(meterRegistry, options); + } + + public static CoreMetricsPublisher coreMetricsPublisher(MeterRegistry meterRegistry) { + return new CoreMetricsPublisher(meterRegistry); + } + } diff --git a/token-pilot-micrometer/src/main/java/io/tokenpilot/micrometer/internal/MicroCostMetricsPublisher.java b/token-pilot-micrometer/src/main/java/io/tokenpilot/micrometer/internal/MicroCostMetricsPublisher.java index 1744581..8c533e6 100644 --- a/token-pilot-micrometer/src/main/java/io/tokenpilot/micrometer/internal/MicroCostMetricsPublisher.java +++ b/token-pilot-micrometer/src/main/java/io/tokenpilot/micrometer/internal/MicroCostMetricsPublisher.java @@ -10,7 +10,9 @@ import io.tokenpilot.core.domain.TokenType; import io.tokenpilot.micrometer.MetricsOptions; +import java.util.Locale; import java.util.Map; +import java.util.Objects; import java.util.Set; /** @@ -22,9 +24,10 @@ public class MicroCostMetricsPublisher implements LedgerListener { private final MeterRegistry meterRegistry; private final Set allowedTagKeys; + private final boolean enabled; public MicroCostMetricsPublisher(MeterRegistry meterRegistry) { - this(meterRegistry, MetricsOptions.defaults()); + this(meterRegistry, MetricsOptions.legacyDefaults()); } public MicroCostMetricsPublisher(MeterRegistry meterRegistry, Set allowedTagKeys) { @@ -32,13 +35,23 @@ public MicroCostMetricsPublisher(MeterRegistry meterRegistry, Set allowe } public MicroCostMetricsPublisher(MeterRegistry meterRegistry, MetricsOptions options) { - this.meterRegistry = meterRegistry; - MetricsOptions resolvedOptions = (options != null) ? options : MetricsOptions.defaults(); + this.meterRegistry = Objects.requireNonNull( + meterRegistry, + "meterRegistry must not be null" + ); + MetricsOptions resolvedOptions = (options != null) + ? options + : MetricsOptions.legacyDefaults(); this.allowedTagKeys = resolvedOptions.allowedTagKeys(); + this.enabled = resolvedOptions.legacyAiTokenMetricsEnabled(); } @Override public void onRecord(CostRecordedEvent event) { + if (!enabled) { + return; + } + Objects.requireNonNull(event, "event must not be null"); Tags commonTags = Tags.of("model", event.modelId()) .and(allowedTags(event.tags())); @@ -47,7 +60,7 @@ public void onRecord(CostRecordedEvent event) { for (TokenType tokenType : TokenType.values()) { long count = event.usage().getCount(tokenType); if (count > 0) { - String typeName = tokenType.name().toLowerCase(); + String typeName = tokenType.name().toLowerCase(Locale.ROOT); DistributionSummary.builder("ai.token.usage.distribution") .description("Distribution of AI token usage per recorded model call") .baseUnit("tokens") diff --git a/token-pilot-micrometer/src/main/java/io/tokenpilot/micrometer/internal/NotificationMetricsPublisher.java b/token-pilot-micrometer/src/main/java/io/tokenpilot/micrometer/internal/NotificationMetricsPublisher.java new file mode 100644 index 0000000..04e35b3 --- /dev/null +++ b/token-pilot-micrometer/src/main/java/io/tokenpilot/micrometer/internal/NotificationMetricsPublisher.java @@ -0,0 +1,34 @@ +package io.tokenpilot.micrometer.internal; + +import io.micrometer.core.instrument.Counter; +import io.micrometer.core.instrument.MeterRegistry; +import io.tokenpilot.notification.BudgetNotificationLifecycleEvent; +import io.tokenpilot.notification.BudgetNotificationLifecycleListener; + +import java.util.Objects; + +/** Budget notification lifecycle 이벤트를 Micrometer meter로 투영합니다. */ +public final class NotificationMetricsPublisher + implements BudgetNotificationLifecycleListener { + + private final MeterRegistry meterRegistry; + + public NotificationMetricsPublisher(MeterRegistry meterRegistry) { + this.meterRegistry = Objects.requireNonNull( + meterRegistry, + "meterRegistry must not be null" + ); + } + + @Override + public void onNotificationLifecycle(BudgetNotificationLifecycleEvent event) { + Objects.requireNonNull(event, "event must not be null"); + Counter.builder(TokenPilotMeterMetadata.NOTIFICATION_EVENTS) + .description(TokenPilotMeterMetadata.NOTIFICATION_EVENTS_DESCRIPTION) + .baseUnit(TokenPilotMeterMetadata.EVENTS_BASE_UNIT) + .tag("outcome", TokenPilotMeterMetadata.tagValue(event.outcome())) + .tag("threshold", TokenPilotMeterMetadata.tagValue(event.threshold())) + .register(meterRegistry) + .increment(); + } +} diff --git a/token-pilot-micrometer/src/main/java/io/tokenpilot/micrometer/internal/NotificationMicrometerComponents.java b/token-pilot-micrometer/src/main/java/io/tokenpilot/micrometer/internal/NotificationMicrometerComponents.java new file mode 100644 index 0000000..023e2ee --- /dev/null +++ b/token-pilot-micrometer/src/main/java/io/tokenpilot/micrometer/internal/NotificationMicrometerComponents.java @@ -0,0 +1,16 @@ +package io.tokenpilot.micrometer.internal; + +import io.micrometer.core.instrument.MeterRegistry; + +/** Budget notification owner의 선택 Micrometer publisher 팩토리입니다. */ +public final class NotificationMicrometerComponents { + + private NotificationMicrometerComponents() { + } + + public static NotificationMetricsPublisher notificationMetricsPublisher( + MeterRegistry meterRegistry + ) { + return new NotificationMetricsPublisher(meterRegistry); + } +} diff --git a/token-pilot-micrometer/src/main/java/io/tokenpilot/micrometer/internal/TokenPilotMeterMetadata.java b/token-pilot-micrometer/src/main/java/io/tokenpilot/micrometer/internal/TokenPilotMeterMetadata.java new file mode 100644 index 0000000..e12ee29 --- /dev/null +++ b/token-pilot-micrometer/src/main/java/io/tokenpilot/micrometer/internal/TokenPilotMeterMetadata.java @@ -0,0 +1,54 @@ +package io.tokenpilot.micrometer.internal; + +import java.util.Locale; + +final class TokenPilotMeterMetadata { + + static final String COST_TOTAL = "tokenpilot.cost.total"; + static final String COST_TOTAL_DESCRIPTION = + "Total actual LLM cost committed by Token Pilot"; + static final String COST_BASE_UNIT = "currency"; + + static final String PREFLIGHT_REQUESTS = "tokenpilot.preflight.requests"; + static final String PREFLIGHT_REQUESTS_DESCRIPTION = + "Total preflight admission decisions made by Token Pilot"; + static final String REQUESTS_BASE_UNIT = "requests"; + + static final String BUDGET_RESERVATIONS = "tokenpilot.budget.reservations"; + static final String BUDGET_RESERVATIONS_DESCRIPTION = + "Total budget reservation results produced by Token Pilot"; + static final String RESERVATIONS_BASE_UNIT = "reservations"; + + static final String RECONCILIATION_ERROR_TOKENS = + "tokenpilot.reconciliation.error.tokens"; + static final String RECONCILIATION_ERROR_TOKENS_DESCRIPTION = + "Absolute difference between estimated and actual total tokens per reconciliation"; + static final String TOKENS_BASE_UNIT = "tokens"; + + static final String RECONCILIATION_OUTCOMES = + "tokenpilot.reconciliation.outcomes"; + static final String RECONCILIATION_OUTCOMES_DESCRIPTION = + "Total reservation reconciliation outcomes applied by Token Pilot"; + static final String RECONCILIATIONS_BASE_UNIT = "reconciliations"; + + static final String PRICING_MISSING = "tokenpilot.pricing.missing"; + static final String PRICING_MISSING_DESCRIPTION = + "Total pricing misses observed by Token Pilot"; + + static final String LISTENER_FAILURES = "tokenpilot.listener.failures"; + static final String LISTENER_FAILURES_DESCRIPTION = + "Total listener failures isolated by Token Pilot"; + static final String FAILURES_BASE_UNIT = "failures"; + + static final String NOTIFICATION_EVENTS = "tokenpilot.notification.events"; + static final String NOTIFICATION_EVENTS_DESCRIPTION = + "Total budget notification lifecycle events observed by Token Pilot"; + static final String EVENTS_BASE_UNIT = "events"; + + private TokenPilotMeterMetadata() { + } + + static String tagValue(Enum value) { + return value.name().toLowerCase(Locale.ROOT); + } +} diff --git a/token-pilot-micrometer/src/test/java/io/tokenpilot/micrometer/internal/BudgetMetricsPublisherTest.java b/token-pilot-micrometer/src/test/java/io/tokenpilot/micrometer/internal/BudgetMetricsPublisherTest.java new file mode 100644 index 0000000..8ed149c --- /dev/null +++ b/token-pilot-micrometer/src/test/java/io/tokenpilot/micrometer/internal/BudgetMetricsPublisherTest.java @@ -0,0 +1,405 @@ +package io.tokenpilot.micrometer.internal; + +import io.micrometer.core.instrument.Meter; +import io.micrometer.core.instrument.simple.SimpleMeterRegistry; +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.ReservationAccountingEvent; +import io.tokenpilot.budget.ReservationAccountingListenerFailureEvent; +import io.tokenpilot.budget.ReservationAccountingListenerPhase; +import io.tokenpilot.budget.ReservationAccountingListenerType; +import io.tokenpilot.budget.ReservationAccountingReason; +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.ReservationTokenEstimate; +import io.tokenpilot.budget.ReservationTransition; +import io.tokenpilot.core.domain.Cost; +import io.tokenpilot.core.domain.PricingSnapshot; +import io.tokenpilot.core.domain.TokenType; +import io.tokenpilot.core.domain.TokenUsageDetails; +import io.tokenpilot.core.domain.UsageSource; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.math.BigDecimal; +import java.time.Instant; +import java.util.Currency; +import java.util.Map; +import java.util.Optional; +import java.util.Set; + +import static org.assertj.core.api.Assertions.assertThat; + +class BudgetMetricsPublisherTest { + + private static final Currency USD = Currency.getInstance("USD"); + private static final ReservationId RESERVATION_ID = + new ReservationId("sensitive-reservation-id"); + private static final BudgetKey BUDGET_KEY = new BudgetKey( + "monthly", + "tenant", + "sensitive-tenant-id", + BudgetWindow.parse("2026-08") + ); + private static final Cost LIMIT = usd("10.00"); + + private SimpleMeterRegistry meterRegistry; + private BudgetMetricsPublisher publisher; + + @BeforeEach + void setUp() { + meterRegistry = new SimpleMeterRegistry(); + publisher = new BudgetMetricsPublisher(meterRegistry); + } + + @Test + void publishesActualCostErrorAndOutcomeOnlyForAppliedCommit() { + ReservationAccountingEvent applied = accountingEvent( + ReservationTransition.applied( + ReservationState.IN_FLIGHT, + ReservationState.COMMITTED + ), + 110, + 60, + "0.80" + ); + + publisher.onAccountingApplied(applied, snapshot()); + + var cost = meterRegistry.find("tokenpilot.cost.total") + .tag("currency", "USD") + .counter(); + assertThat(cost).isNotNull(); + assertThat(cost.count()).isEqualTo(0.8); + assertThat(cost.getId().getType()).isEqualTo(Meter.Type.COUNTER); + assertThat(cost.getId().getBaseUnit()).isEqualTo("currency"); + assertThat(cost.getId().getDescription()) + .isEqualTo("Total actual LLM cost committed by Token Pilot"); + assertThat(cost.getId().getTags()).extracting(tag -> tag.getKey()) + .containsExactly("currency"); + + var error = meterRegistry.find("tokenpilot.reconciliation.error.tokens") + .tag("direction", "underestimate") + .summary(); + assertThat(error).isNotNull(); + assertThat(error.count()).isEqualTo(1L); + assertThat(error.totalAmount()).isEqualTo(20.0); + assertThat(error.getId().getType()).isEqualTo(Meter.Type.DISTRIBUTION_SUMMARY); + assertThat(error.getId().getBaseUnit()).isEqualTo("tokens"); + assertThat(error.getId().getDescription()).isEqualTo( + "Absolute difference between estimated and actual total tokens per reconciliation" + ); + + var outcome = meterRegistry.find("tokenpilot.reconciliation.outcomes") + .tag("outcome", "committed") + .tag("reason", "actual_usage_reported") + .counter(); + assertThat(outcome).isNotNull(); + assertThat(outcome.count()).isEqualTo(1.0); + assertThat(outcome.getId().getBaseUnit()).isEqualTo("reconciliations"); + assertThat(outcome.getId().getDescription()) + .isEqualTo("Total reservation reconciliation outcomes applied by Token Pilot"); + } + + @Test + void ignoresReusedCommitSoCostAndOutcomeAreNotDuplicated() { + ReservationAccountingEvent applied = accountingEvent( + ReservationTransition.applied( + ReservationState.IN_FLIGHT, + ReservationState.COMMITTED + ), + 100, + 50, + "0.80" + ); + ReservationAccountingEvent reused = accountingEvent( + ReservationTransition.unchanged( + ReservationState.COMMITTED, + AccountingTransitionStatus.REUSED + ), + 100, + 50, + "0.80" + ); + + publisher.onCommitted(applied); + publisher.onCommitted(reused); + + assertThat(meterRegistry.find("tokenpilot.cost.total").counter().count()) + .isEqualTo(0.8); + assertThat(meterRegistry.find("tokenpilot.reconciliation.outcomes") + .counter().count()).isEqualTo(1.0); + assertThat(meterRegistry.find("tokenpilot.reconciliation.error.tokens") + .summary().count()).isEqualTo(1L); + } + + @Test + void explicitZeroCommitRemainsDistinctFromMissingPricing() { + publisher.onCommitted(accountingEvent( + ReservationTransition.applied( + ReservationState.IN_FLIGHT, + ReservationState.COMMITTED + ), + 100, + 50, + "0.00" + )); + + var cost = meterRegistry.find("tokenpilot.cost.total") + .tag("currency", "USD") + .counter(); + assertThat(cost).isNotNull(); + assertThat(cost.count()).isZero(); + assertThat(meterRegistry.find("tokenpilot.reconciliation.outcomes") + .tag("outcome", "committed") + .counter().count()).isEqualTo(1.0); + assertThat(meterRegistry.find("tokenpilot.pricing.missing").meter()) + .isNull(); + } + + @Test + void distinguishesOverestimateAndExactTokenErrors() { + publisher.onCommitted(accountingEvent( + ReservationTransition.applied( + ReservationState.IN_FLIGHT, + ReservationState.COMMITTED + ), + 90, + 50, + "0.40" + )); + publisher.onCommitted(accountingEvent( + ReservationTransition.applied( + ReservationState.IN_FLIGHT, + ReservationState.COMMITTED + ), + 100, + 50, + "0.40" + )); + + var overestimate = meterRegistry.find("tokenpilot.reconciliation.error.tokens") + .tag("direction", "overestimate") + .summary(); + var exact = meterRegistry.find("tokenpilot.reconciliation.error.tokens") + .tag("direction", "exact") + .summary(); + + assertThat(overestimate.count()).isEqualTo(1L); + assertThat(overestimate.totalAmount()).isEqualTo(10.0); + assertThat(exact.count()).isEqualTo(1L); + assertThat(exact.totalAmount()).isZero(); + } + + @Test + void reconciliationRequiredPublishesOutcomeWithoutZeroCostOrError() { + ReservationTransition transition = ReservationTransition.applied( + ReservationState.IN_FLIGHT, + ReservationState.RECONCILIATION_REQUIRED + ); + publisher.onReconciliationRequired(new ReservationReconciliationRequiredEvent( + RESERVATION_ID, + BUDGET_KEY, + transition, + ReservationAccountingReason.ACTUAL_USAGE_UNAVAILABLE, + snapshot() + )); + + var outcome = meterRegistry.find("tokenpilot.reconciliation.outcomes") + .tag("outcome", "reconciliation_required") + .tag("reason", "actual_usage_unavailable") + .counter(); + assertThat(outcome).isNotNull(); + assertThat(outcome.count()).isEqualTo(1.0); + assertThat(meterRegistry.find("tokenpilot.cost.total").meter()).isNull(); + assertThat(meterRegistry.find("tokenpilot.reconciliation.error.tokens") + .meter()).isNull(); + } + + @Test + void pricingMismatchPublishesBoundedPendingOutcomeWithoutCostOrError() { + publisher.onReconciliationRequired( + new ReservationReconciliationRequiredEvent( + RESERVATION_ID, + BUDGET_KEY, + ReservationTransition.applied( + ReservationState.IN_FLIGHT, + ReservationState.RECONCILIATION_REQUIRED + ), + ReservationAccountingReason.PRICING_RECONCILIATION_REQUIRED, + snapshot() + ) + ); + + var outcome = meterRegistry.find("tokenpilot.reconciliation.outcomes") + .tag("outcome", "reconciliation_required") + .tag("reason", "pricing_reconciliation_required") + .counter(); + assertThat(outcome).isNotNull(); + assertThat(outcome.count()).isEqualTo(1.0); + assertThat(meterRegistry.find("tokenpilot.cost.total").meter()).isNull(); + assertThat(meterRegistry.find("tokenpilot.reconciliation.error.tokens") + .meter()).isNull(); + } + + @Test + void publishesEveryReservationEvaluationWithBoundedStateOnly() { + BudgetReservationRequest request = new BudgetReservationRequest( + BUDGET_KEY, + LIMIT, + usd("1.00"), + "sensitive-request-id", + new IdempotencyKey("sensitive-idempotency-id"), + null, + null, + null, + Optional.empty() + ); + BudgetReservationResult result = BudgetReservationResult.blocked( + snapshot(), + "raw rejection message that must not become a tag" + ); + + publisher.onReservationEvaluated(request, result); + + var counter = meterRegistry.find("tokenpilot.budget.reservations") + .tag("state", "blocked") + .counter(); + assertThat(counter).isNotNull(); + assertThat(counter.count()).isEqualTo(1.0); + assertThat(counter.getId().getBaseUnit()).isEqualTo("reservations"); + assertThat(counter.getId().getDescription()) + .isEqualTo("Total budget reservation results produced by Token Pilot"); + assertThat(counter.getId().getTags()).extracting(tag -> tag.getKey()) + .containsExactly("state"); + } + + @Test + void publishesBoundedListenerFailureAndIdentifiesAsMetricsListener() { + publisher.onFailure(new ReservationAccountingListenerFailureEvent( + ReservationAccountingListenerType.NOTIFICATION, + ReservationAccountingListenerPhase.RECONCILIATION_REQUIRED + )); + + var counter = meterRegistry.find("tokenpilot.listener.failures") + .tag("listener", "notification") + .tag("phase", "reconciliation_required") + .counter(); + assertThat(counter).isNotNull(); + assertThat(counter.count()).isEqualTo(1.0); + assertThat(counter.getId().getBaseUnit()).isEqualTo("failures"); + assertThat(counter.getId().getDescription()) + .isEqualTo("Total listener failures isolated by Token Pilot"); + assertThat(publisher.listenerType()) + .isEqualTo(ReservationAccountingListenerType.METRICS); + } + + @Test + void neverCopiesRawIdentifiersModelsOrMessagesIntoMeterMetadata() { + publisher.onCommitted(accountingEvent( + ReservationTransition.applied( + ReservationState.IN_FLIGHT, + ReservationState.COMMITTED + ), + 110, + 60, + "0.80" + )); + BudgetReservationRequest request = new BudgetReservationRequest( + BUDGET_KEY, + LIMIT, + usd("1.00"), + "sensitive-request-id", + new IdempotencyKey("sensitive-idempotency-id"), + null, + null, + null, + Optional.empty() + ); + publisher.onReservationEvaluated( + request, + BudgetReservationResult.blocked( + snapshot(), + "raw rejection message that must not become metadata" + ) + ); + + assertThat(meterRegistry.getMeters()) + .extracting(meter -> meter.getId().toString()) + .allSatisfy(metadata -> assertThat(metadata) + .doesNotContain( + "sensitive-reservation-id", + "sensitive-tenant-id", + "sensitive-request-id", + "sensitive-idempotency-id", + "raw-request-model", + "raw-response-model", + "raw rejection message" + )); + } + + private static ReservationAccountingEvent accountingEvent( + ReservationTransition transition, + long actualInput, + long actualOutput, + String actualCost + ) { + return new ReservationAccountingEvent(new ReservationReconciliation( + "sensitive-request-id", + "sensitive-attempt-id", + RESERVATION_ID, + BUDGET_KEY, + "raw-response-model", + pricingSnapshot(), + new ReservationTokenEstimate(100, 120, 50), + new ReservationActualTokens( + actualInput, + actualOutput, + TokenUsageDetails.unreported(), + UsageSource.PROVIDER_REPORTED + ), + usd("1.00"), + usd(actualCost), + false, + transition, + ReservationAccountingReason.ACTUAL_USAGE_REPORTED + )); + } + + private static PricingSnapshot pricingSnapshot() { + return new PricingSnapshot( + "raw-request-model", + "pricing-policy", + "catalog-version", + Instant.parse("2026-08-25T00:00:00Z"), + Map.of( + TokenType.PROMPT, new BigDecimal("0.001"), + TokenType.COMPLETION, new BigDecimal("0.002") + ), + USD + ); + } + + private static BudgetSnapshot snapshot() { + return new BudgetSnapshot( + BUDGET_KEY, + LIMIT, + Cost.zero(USD), + Cost.zero(USD), + Cost.zero(USD), + Set.of() + ); + } + + private static Cost usd(String amount) { + return Cost.of(new BigDecimal(amount), USD); + } +} diff --git a/token-pilot-micrometer/src/test/java/io/tokenpilot/micrometer/internal/CoreMetricsPublisherTest.java b/token-pilot-micrometer/src/test/java/io/tokenpilot/micrometer/internal/CoreMetricsPublisherTest.java new file mode 100644 index 0000000..7c047d8 --- /dev/null +++ b/token-pilot-micrometer/src/test/java/io/tokenpilot/micrometer/internal/CoreMetricsPublisherTest.java @@ -0,0 +1,87 @@ +package io.tokenpilot.micrometer.internal; + +import io.micrometer.core.instrument.Meter; +import io.micrometer.core.instrument.simple.SimpleMeterRegistry; +import io.tokenpilot.core.domain.AdmissionReason; +import io.tokenpilot.core.domain.AdmissionStatus; +import io.tokenpilot.core.domain.BudgetResult; +import io.tokenpilot.core.domain.MissingPricingPolicy; +import io.tokenpilot.core.domain.PreflightDecisionEvent; +import io.tokenpilot.core.domain.PricingMissingEvent; +import io.tokenpilot.core.domain.PricingResolution; +import io.tokenpilot.core.domain.TokenEstimatorDescriptor; +import io.tokenpilot.core.domain.TokenizationBasis; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.Optional; +import java.util.OptionalLong; + +import static org.assertj.core.api.Assertions.assertThat; + +class CoreMetricsPublisherTest { + + private SimpleMeterRegistry meterRegistry; + private CoreMetricsPublisher publisher; + + @BeforeEach + void setUp() { + meterRegistry = new SimpleMeterRegistry(); + publisher = new CoreMetricsPublisher(meterRegistry); + } + + @Test + void publishesBoundedPreflightDecision() { + publisher.onDecision(new PreflightDecisionEvent(indeterminateResult())); + + var counter = meterRegistry.find("tokenpilot.preflight.requests") + .tag("decision", "indeterminate") + .tag("reason", "unknown_model") + .counter(); + + assertThat(counter).isNotNull(); + assertThat(counter.count()).isEqualTo(1.0); + assertThat(counter.getId().getType()).isEqualTo(Meter.Type.COUNTER); + assertThat(counter.getId().getBaseUnit()).isEqualTo("requests"); + assertThat(counter.getId().getDescription()) + .isEqualTo("Total preflight admission decisions made by Token Pilot"); + assertThat(counter.getId().getTags()).extracting(tag -> tag.getKey()) + .containsExactly("decision", "reason"); + } + + @Test + void publishesPricingMissWithoutResolutionOrRequestTags() { + publisher.onPricingMissing(new PricingMissingEvent( + MissingPricingPolicy.FAIL_CLOSED, + PricingResolution.MISSING_PLAN + )); + + var counter = meterRegistry.find("tokenpilot.pricing.missing") + .tag("policy", "fail_closed") + .counter(); + + assertThat(counter).isNotNull(); + assertThat(counter.count()).isEqualTo(1.0); + assertThat(counter.getId().getType()).isEqualTo(Meter.Type.COUNTER); + assertThat(counter.getId().getBaseUnit()).isEqualTo("events"); + assertThat(counter.getId().getDescription()) + .isEqualTo("Total pricing misses observed by Token Pilot"); + assertThat(counter.getId().getTags()).extracting(tag -> tag.getKey()) + .containsExactly("policy"); + } + + private static BudgetResult indeterminateResult() { + return new BudgetResult( + AdmissionStatus.INDETERMINATE, + AdmissionReason.UNKNOWN_MODEL, + Optional.empty(), + OptionalLong.empty(), + OptionalLong.empty(), + 0, + OptionalLong.empty(), + OptionalLong.empty(), + new TokenEstimatorDescriptor("test-estimator", "1"), + new TokenizationBasis("TEST") + ); + } +} diff --git a/token-pilot-micrometer/src/test/java/io/tokenpilot/micrometer/internal/LedgerMicrometerComponentsTest.java b/token-pilot-micrometer/src/test/java/io/tokenpilot/micrometer/internal/LedgerMicrometerComponentsTest.java new file mode 100644 index 0000000..293a0c4 --- /dev/null +++ b/token-pilot-micrometer/src/test/java/io/tokenpilot/micrometer/internal/LedgerMicrometerComponentsTest.java @@ -0,0 +1,40 @@ +package io.tokenpilot.micrometer.internal; + +import io.micrometer.core.instrument.simple.SimpleMeterRegistry; +import io.tokenpilot.micrometer.MetricsOptions; +import org.junit.jupiter.api.Test; + +import static org.assertj.core.api.Assertions.assertThat; + +class LedgerMicrometerComponentsTest { + + @Test + void createsOwnerSpecificPublishers() { + var meterRegistry = new SimpleMeterRegistry(); + + assertThat(LedgerMicrometerComponents.coreMetricsPublisher(meterRegistry)) + .isInstanceOf(CoreMetricsPublisher.class); + assertThat(BudgetMicrometerComponents.budgetMetricsPublisher(meterRegistry)) + .isInstanceOf(BudgetMetricsPublisher.class); + assertThat(NotificationMicrometerComponents.notificationMetricsPublisher(meterRegistry)) + .isInstanceOf(NotificationMetricsPublisher.class); + assertThat(LedgerMicrometerComponents.legacyMicroCostMetricsPublisher( + meterRegistry, + MetricsOptions.legacyDefaults() + )).isInstanceOf(MicroCostMetricsPublisher.class); + } + + @Test + void coreFactoryMethodTableDoesNotReferenceOptionalOwners() { + assertThat(LedgerMicrometerComponents.class.getDeclaredMethods()) + .allSatisfy(method -> { + assertThat(method.getReturnType().getName()) + .doesNotContain("BudgetMetricsPublisher") + .doesNotContain("NotificationMetricsPublisher"); + assertThat(method.getParameterTypes()) + .extracting(Class::getName) + .noneMatch(name -> name.startsWith("io.tokenpilot.budget")) + .noneMatch(name -> name.startsWith("io.tokenpilot.notification")); + }); + } +} diff --git a/token-pilot-micrometer/src/test/java/io/tokenpilot/micrometer/internal/MicroCostMetricsPublisherTest.java b/token-pilot-micrometer/src/test/java/io/tokenpilot/micrometer/internal/MicroCostMetricsPublisherTest.java index 319d63b..314bd9a 100644 --- a/token-pilot-micrometer/src/test/java/io/tokenpilot/micrometer/internal/MicroCostMetricsPublisherTest.java +++ b/token-pilot-micrometer/src/test/java/io/tokenpilot/micrometer/internal/MicroCostMetricsPublisherTest.java @@ -43,13 +43,11 @@ void shouldPublishMetricsWhenEventRecorded() { assertThat(meterRegistry.find("ai.token.usage.total") .tag("model", "gpt-4o") .tag("token_type", "prompt") - .tag("tenant_id", "tenant-1") .counter().count()).isEqualTo(100.0); assertThat(meterRegistry.find("ai.token.usage.total") .tag("model", "gpt-4o") .tag("token_type", "completion") - .tag("tenant_id", "tenant-1") .counter().count()).isEqualTo(200.0); // Then: 토큰 사용량 분포(Summary) 확인 @@ -61,7 +59,6 @@ void shouldPublishMetricsWhenEventRecorded() { // Then: 비용 카운터 확인 var costCounter = meterRegistry.find("ai.token.cost.total") .tag("model", "gpt-4o") - .tag("tenant_id", "tenant-1") .tag("currency", "USD") .counter(); @@ -71,13 +68,14 @@ void shouldPublishMetricsWhenEventRecorded() { var promptSummary = meterRegistry.find("ai.token.usage.distribution") .tag("model", "gpt-4o") - .tag("tenant_id", "tenant-1") .tag("token_type", "prompt") .summary(); assertThat(promptSummary.getId().getDescription()) .isEqualTo("Distribution of AI token usage per recorded model call"); assertThat(promptSummary.getId().getBaseUnit()).isEqualTo("tokens"); + assertThat(costCounter.getId().getTag("tenant_id")) + .isEqualTo("tenant-1"); } @Test @@ -172,4 +170,58 @@ void shouldPublishMetricsWhenTagsAreEmpty() { .tag("currency", "USD") .counter()).isNotNull(); } + + @Test + @DisplayName("기본 MetricsOptions는 사용자 태그와 legacy meter를 비활성화해야 한다") + void shouldUseSafeDefaults() { + MetricsOptions options = MetricsOptions.defaults(); + + assertThat(options.allowedTagKeys()).isEmpty(); + assertThat(options.legacyAiTokenMetricsEnabled()).isFalse(); + } + + @Test + @DisplayName("legacy 비활성 options를 사용하면 ai.token meter를 발행하지 않아야 한다") + void shouldNotPublishLegacyMetricsWhenDisabled() { + publisher = new MicroCostMetricsPublisher( + meterRegistry, + MetricsOptions.defaults() + ); + + publisher.onRecord(new CostRecordedEvent( + "gpt-4o", + TokenUsage.from(100, 200), + new Cost(new BigDecimal("0.5"), Currency.getInstance("USD")), + Map.of("tenant_id", "tenant-1") + )); + + assertThat(meterRegistry.find("ai.token.usage.total").meter()).isNull(); + assertThat(meterRegistry.find("ai.token.usage.distribution").meter()).isNull(); + assertThat(meterRegistry.find("ai.token.cost.total").meter()).isNull(); + } + + @Test + @DisplayName("1-arg MetricsOptions 생성자는 기존 직접 생성의 legacy opt-in을 유지해야 한다") + void shouldPreserveOneArgumentOptionsCompatibility() { + MetricsOptions options = new MetricsOptions(Set.of("tenant_id")); + + assertThat(options.allowedTagKeys()).containsExactly("tenant_id"); + assertThat(options.legacyAiTokenMetricsEnabled()).isTrue(); + } + + @Test + @DisplayName("기존 direct 기본 생성자는 tenant_id 허용 동작을 유지해야 한다") + void shouldPreserveDirectPublisherDefaultTagCompatibility() { + publisher.onRecord(new CostRecordedEvent( + "gpt-4o", + TokenUsage.from(1, 1), + Cost.of(new BigDecimal("0.1"), Currency.getInstance("USD")), + Map.of("tenant_id", "tenant-legacy", "user_id", "user-sensitive") + )); + + var counter = meterRegistry.get("ai.token.cost.total") + .tag("tenant_id", "tenant-legacy") + .counter(); + assertThat(counter.getId().getTag("user_id")).isNull(); + } } diff --git a/token-pilot-micrometer/src/test/java/io/tokenpilot/micrometer/internal/NotificationMetricsPublisherTest.java b/token-pilot-micrometer/src/test/java/io/tokenpilot/micrometer/internal/NotificationMetricsPublisherTest.java new file mode 100644 index 0000000..c48b398 --- /dev/null +++ b/token-pilot-micrometer/src/test/java/io/tokenpilot/micrometer/internal/NotificationMetricsPublisherTest.java @@ -0,0 +1,68 @@ +package io.tokenpilot.micrometer.internal; + +import io.micrometer.core.instrument.Meter; +import io.micrometer.core.instrument.simple.SimpleMeterRegistry; +import io.tokenpilot.budget.BudgetThreshold; +import io.tokenpilot.notification.BudgetNotificationLifecycleEvent; +import io.tokenpilot.notification.BudgetNotificationOutcome; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import static org.assertj.core.api.Assertions.assertThat; + +class NotificationMetricsPublisherTest { + + private SimpleMeterRegistry meterRegistry; + private NotificationMetricsPublisher publisher; + + @BeforeEach + void setUp() { + meterRegistry = new SimpleMeterRegistry(); + publisher = new NotificationMetricsPublisher(meterRegistry); + } + + @Test + void publishesBoundedNotificationLifecycleOutcome() { + publisher.onNotificationLifecycle(new BudgetNotificationLifecycleEvent( + BudgetNotificationOutcome.FAILURE, + BudgetThreshold.WARNING + )); + + var counter = meterRegistry.find("tokenpilot.notification.events") + .tag("outcome", "failure") + .tag("threshold", "warning") + .counter(); + + assertThat(counter).isNotNull(); + assertThat(counter.count()).isEqualTo(1.0); + assertThat(counter.getId().getType()).isEqualTo(Meter.Type.COUNTER); + assertThat(counter.getId().getBaseUnit()).isEqualTo("events"); + assertThat(counter.getId().getDescription()) + .isEqualTo("Total budget notification lifecycle events observed by Token Pilot"); + assertThat(counter.getId().getTags()).extracting(tag -> tag.getKey()) + .containsExactly("outcome", "threshold"); + } + + @Test + void distinguishesSuccessFailureAndDeduplicatedOutcomes() { + publisher.onNotificationLifecycle(new BudgetNotificationLifecycleEvent( + BudgetNotificationOutcome.SUCCESS, + BudgetThreshold.HALF + )); + publisher.onNotificationLifecycle(new BudgetNotificationLifecycleEvent( + BudgetNotificationOutcome.FAILURE, + BudgetThreshold.HALF + )); + publisher.onNotificationLifecycle(new BudgetNotificationLifecycleEvent( + BudgetNotificationOutcome.DEDUPLICATED, + BudgetThreshold.HALF + )); + + assertThat(meterRegistry.find("tokenpilot.notification.events") + .tag("outcome", "success").counter().count()).isEqualTo(1.0); + assertThat(meterRegistry.find("tokenpilot.notification.events") + .tag("outcome", "failure").counter().count()).isEqualTo(1.0); + assertThat(meterRegistry.find("tokenpilot.notification.events") + .tag("outcome", "deduplicated").counter().count()).isEqualTo(1.0); + } +} 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..2625449 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,101 @@ * @param threshold 도달한 임계치 * @param state 예산 상태 * @param reason 상태 설명 - * @param projectedUsage 후보 비용을 포함한 예상 사용량 + * @param usage 알림을 만든 accounting 또는 admission 결과의 사용량 * @param limit 예산 한도 - * @param tags 알림에 전달할 태그 + * @param source 알림을 만든 원자적 결과 종류 + * @param tags legacy decision 알림에서만 보존하는 호환 tag * - *

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

원자적 accounting 이벤트는 prompt, raw provider response, API key, 임의 tag map을 + * 포함하지 않습니다. legacy decision 생성자만 기존 handler 호환을 위해 불변 tag 복사본을 + * 보존합니다. {@link #projectedUsage()}와 {@link #currentUsage()} 호환 accessor는 0.1.x 동안 + * 유지하며 0.2.0에서 제거할 예정입니다.

*/ public record BudgetNotificationEvent( BudgetKey key, BudgetThreshold threshold, BudgetState state, String reason, - Cost projectedUsage, + Cost usage, Cost limit, + BudgetNotificationSource source, Map tags ) { + 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"); + tags = Map.copyOf(Objects.requireNonNull(tags, "tags must not be null")); + if (!usage.currency().equals(limit.currency())) { + throw new IllegalArgumentException("usage and limit must use the same currency"); + } + if (source != BudgetNotificationSource.LEGACY_DECISION && !tags.isEmpty()) { + throw new IllegalArgumentException( + "atomic notification events must not contain arbitrary tags" + ); + } + } + + public BudgetNotificationEvent( + BudgetKey key, + BudgetThreshold threshold, + BudgetState state, + String reason, + Cost usage, + Cost limit, + BudgetNotificationSource source + ) { + this(key, threshold, state, reason, usage, limit, source, Map.of()); + } + + /** + * @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, + tags + ); + } + /** - * @return {@link #projectedUsage()}와 동일한 후보 비용 포함 예상 사용량 - * @deprecated 후보 비용 포함 사용량은 {@link #projectedUsage()}를 사용하세요. + * @return {@link #usage()}와 동일한 accounting/admission 사용량 + * @deprecated source별 의미가 명확한 {@link #usage()}를 사용하세요. + */ + @Deprecated(since = "0.1.0", forRemoval = true) + public Cost projectedUsage() { + return usage; + } + + /** + * @return {@link #usage()}와 동일한 accounting/admission 사용량 + * @deprecated source별 의미가 명확한 {@link #usage()}를 사용하세요. */ @Deprecated(since = "0.1.0", forRemoval = true) public Cost currentUsage() { - return projectedUsage; + return usage; } + } diff --git a/token-pilot-notification/src/main/java/io/tokenpilot/notification/BudgetNotificationLifecycleEvent.java b/token-pilot-notification/src/main/java/io/tokenpilot/notification/BudgetNotificationLifecycleEvent.java new file mode 100644 index 0000000..3ebabea --- /dev/null +++ b/token-pilot-notification/src/main/java/io/tokenpilot/notification/BudgetNotificationLifecycleEvent.java @@ -0,0 +1,25 @@ +package io.tokenpilot.notification; + +import io.tokenpilot.budget.BudgetThreshold; + +import java.util.Objects; + +/** + * notification handler 처리 또는 dedup 결과를 전달하는 framework-independent 이벤트입니다. + * + * @param outcome success, failure 또는 process-local dedup 결과 + * @param threshold 처리 대상인 bounded budget threshold + */ +public record BudgetNotificationLifecycleEvent( + BudgetNotificationOutcome outcome, + BudgetThreshold threshold +) { + + public BudgetNotificationLifecycleEvent { + Objects.requireNonNull(outcome, "outcome must not be null"); + Objects.requireNonNull(threshold, "threshold must not be null"); + if (threshold == BudgetThreshold.NONE) { + throw new IllegalArgumentException("threshold must identify a notification boundary"); + } + } +} diff --git a/token-pilot-notification/src/main/java/io/tokenpilot/notification/BudgetNotificationLifecycleListener.java b/token-pilot-notification/src/main/java/io/tokenpilot/notification/BudgetNotificationLifecycleListener.java new file mode 100644 index 0000000..d1d6e0b --- /dev/null +++ b/token-pilot-notification/src/main/java/io/tokenpilot/notification/BudgetNotificationLifecycleListener.java @@ -0,0 +1,8 @@ +package io.tokenpilot.notification; + +/** Budget notification lifecycle을 best-effort로 관찰하는 framework-independent 계약입니다. */ +@FunctionalInterface +public interface BudgetNotificationLifecycleListener { + + void onNotificationLifecycle(BudgetNotificationLifecycleEvent event); +} diff --git a/token-pilot-notification/src/main/java/io/tokenpilot/notification/BudgetNotificationOutcome.java b/token-pilot-notification/src/main/java/io/tokenpilot/notification/BudgetNotificationOutcome.java new file mode 100644 index 0000000..0056630 --- /dev/null +++ b/token-pilot-notification/src/main/java/io/tokenpilot/notification/BudgetNotificationOutcome.java @@ -0,0 +1,8 @@ +package io.tokenpilot.notification; + +/** Budget notification 처리의 bounded 운영 결과입니다. */ +public enum BudgetNotificationOutcome { + SUCCESS, + FAILURE, + DEDUPLICATED +} 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..562df54 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,48 +1,273 @@ 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.ReservationAccountingListenerType; +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; +import java.util.function.Supplier; /** - * 예산 판단 결과를 기반으로 알림 이벤트를 발행하는 서비스 + * 원자적 budget/accounting 결과를 threshold 알림으로 변환하는 best-effort listener입니다. + * + *

원자적 callback의 threshold claim은 handler 호출 전에 완료됩니다. handler 실패는 다음 + * handler 전달을 막지 않고 이미 적용된 회계 또는 admission 결과를 변경하지 않으며, 실패한 + * 전달은 durable outbox가 없는 MVP에서 재시도되지 않습니다. Legacy decision 경로는 호환성을 + * 위해 전달 성공 후에만 상태를 갱신하므로 실패한 호출을 다음 호출에서 다시 시도합니다. + * 두 경로의 claim은 같은 store monitor에서 직렬화되어 혼용 중에도 같은 threshold를 중복 + * 전달하지 않습니다.

*/ -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; + private final List lifecycleListeners; + /** + * @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(handler, store, List.of()); + } + + /** legacy decision 알림에 lifecycle 관측을 추가하는 호환 생성자입니다. */ + public BudgetNotificationService( + BudgetNotificationHandler handler, + NotificationStateStore store, + List lifecycleListeners + ) { + this( + List.of(Objects.requireNonNull(handler, "handler must not be null")), + Objects.requireNonNull(store, "store must not be null"), + store instanceof AtomicNotificationStateStore atomic ? atomic : null, + null, + BudgetNotificationErrorHook.noOp(), + lifecycleListeners + ); + } + + public BudgetNotificationService( + List handlers, + AtomicNotificationStateStore store, + Function snapshotResolver + ) { + this( + handlers, + store, + snapshotResolver, + BudgetNotificationErrorHook.noOp(), + List.of() + ); + } + + public BudgetNotificationService( + List handlers, + AtomicNotificationStateStore store, + Function snapshotResolver, + BudgetNotificationErrorHook errorHook + ) { + this(handlers, store, snapshotResolver, errorHook, List.of()); + } + + public BudgetNotificationService( + List handlers, + AtomicNotificationStateStore store, + Function snapshotResolver, + BudgetNotificationErrorHook errorHook, + List lifecycleListeners + ) { + this( + handlers, + store, + store, + Objects.requireNonNull( + snapshotResolver, + "snapshotResolver must not be null" + ), + errorHook, + lifecycleListeners + ); + } + + private BudgetNotificationService( + List handlers, + NotificationStateStore store, + AtomicNotificationStateStore atomicStore, + Function snapshotResolver, + BudgetNotificationErrorHook errorHook, + List lifecycleListeners + ) { + 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 = atomicStore; + this.snapshotResolver = snapshotResolver; + this.errorHook = Objects.requireNonNull(errorHook, "errorHook must not be null"); + this.lifecycleListeners = List.copyOf( + Objects.requireNonNull(lifecycleListeners, "lifecycleListeners must not be null") + ); } /** - * 임계치가 증가한 경우에만 이벤트를 발생시킨다 + * 기존 #37 callback은 주입된 resolver로 현재 bucket snapshot을 조회합니다. 현재 budget store는 + * 정확한 전이 시점 snapshot을 포함하는 + * {@link #onAccountingApplied(ReservationAccountingEvent, BudgetSnapshot)}를 호출합니다. */ - public void notifyIfNeeded( - BudgetDecision decision, - Map tags + @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 ) { - BudgetThreshold current = decision.threshold(); + Objects.requireNonNull(event, "event must not be null"); + Objects.requireNonNull(snapshot, "snapshot must not be null"); + requireAtomicLifecycle(); + processAppliedAccounting(event, snapshot); + } - if (current == BudgetThreshold.NONE) { + 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 = claimAtomic( + () -> atomicStore.recordAppliedTransition(event, snapshot, source) + ); + } catch (RuntimeException failure) { + report(BudgetNotificationError.stateFailure(failure)); return; } - BudgetThreshold last = store.getLastNotifiedThreshold(decision.key()); + 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(); - // 같은 window에서 중복 방지 - if (current.compareTo(last) <= 0) { + NotificationStateStore.NotificationClaim claim; + try { + claim = claimAtomic(() -> atomicStore.recordReconciliationRequired(event)); + } catch (RuntimeException failure) { + report(BudgetNotificationError.stateFailure(failure)); return; } + publish( + event.budgetKey(), + claim, + event.reason().name(), + BudgetNotificationSource.RECONCILIATION_REQUIRED, + event.snapshot().limit() + ); + } - BudgetNotificationEvent event = - new BudgetNotificationEvent( + @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 = claimAtomic(() -> 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는 legacy 이벤트에서만 + * 불변 복사본으로 보존됩니다. + */ + @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; + } + + BudgetNotificationOutcome outcome; + RuntimeException handlerFailure = null; + synchronized (store) { + BudgetThreshold last = store.getLastNotifiedThreshold(decision.key()); + if (current.compareTo(last) <= 0) { + outcome = BudgetNotificationOutcome.DEDUPLICATED; + } else { + BudgetNotificationEvent event = new BudgetNotificationEvent( decision.key(), current, decision.state(), @@ -51,12 +276,127 @@ public void notifyIfNeeded( decision.limit(), tags ); + try { + for (BudgetNotificationHandler handler : handlers) { + handler.handle(event); + } + store.updateLastNotifiedThreshold(decision.key(), current); + outcome = BudgetNotificationOutcome.SUCCESS; + } catch (RuntimeException failure) { + outcome = BudgetNotificationOutcome.FAILURE; + handlerFailure = failure; + } + } + } + + publishLifecycleBestEffort(outcome, current); + if (handlerFailure != null) { + throw handlerFailure; + } + } + + private void publish( + BudgetKey key, + NotificationStateStore.NotificationClaim claim, + String reason, + BudgetNotificationSource source, + Cost limit + ) { + if (claim.thresholds().isEmpty()) { + if (claim.observedThreshold() != BudgetThreshold.NONE) { + publishLifecycleBestEffort( + BudgetNotificationOutcome.DEDUPLICATED, + claim.observedThreshold() + ); + } + return; + } + for (BudgetThreshold threshold : claim.thresholds()) { + dispatchAtomic(new BudgetNotificationEvent( + key, + threshold, + state(threshold), + reason, + claim.usage(), + limit, + source + )); + } + } + + private void dispatchAtomic(BudgetNotificationEvent event) { + for (BudgetNotificationHandler handler : handlers) { + try { + handler.handle(event); + publishLifecycleBestEffort( + BudgetNotificationOutcome.SUCCESS, + event.threshold() + ); + } catch (RuntimeException failure) { + publishLifecycleBestEffort( + BudgetNotificationOutcome.FAILURE, + event.threshold() + ); + report(BudgetNotificationError.handlerFailure(handler, failure)); + } + } + } + + @Override + public ReservationAccountingListenerType listenerType() { + return ReservationAccountingListenerType.NOTIFICATION; + } + + private void report(BudgetNotificationError error) { + try { + errorHook.onError(error); + } catch (RuntimeException ignored) { + // Error hook도 best-effort이며 accounting/provider 결과에 영향을 주지 않습니다. + } + } - handler.handle(event); + private NotificationStateStore.NotificationClaim claimAtomic( + Supplier claim + ) { + synchronized (store) { + return claim.get(); + } + } + + private void requireAtomicLifecycle() { + if (atomicStore == null || snapshotResolver == null) { + throw new IllegalStateException( + "atomic accounting callbacks require AtomicNotificationStateStore" + ); + } + } - store.updateLastNotifiedThreshold( - decision.key(), - current + private static BudgetState state(BudgetThreshold threshold) { + return switch (threshold) { + case HALF -> BudgetState.ALLOW; + case WARNING -> BudgetState.WARN; + case EXCEEDED -> BudgetState.BLOCK; + case NONE -> throw new IllegalArgumentException("NONE must not be published"); + }; + } + + private void publishLifecycleBestEffort( + BudgetNotificationOutcome outcome, + BudgetThreshold threshold + ) { + if (lifecycleListeners.isEmpty()) { + return; + } + BudgetNotificationLifecycleEvent event = new BudgetNotificationLifecycleEvent( + outcome, + threshold ); + for (BudgetNotificationLifecycleListener listener : lifecycleListeners) { + try { + listener.onNotificationLifecycle(event); + } catch (RuntimeException ignored) { + // Optional observers do not alter notification delivery or failure propagation. + } + } } } 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..da7873a 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,176 @@ 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()); + BudgetThreshold observedThreshold = reachedThreshold( + snapshot.effectiveUsage(), + state.limit + ); + if (!state.processedTransitions.add(dedupKey)) { + claim.set(NotificationClaim.none( + snapshot.effectiveUsage(), + observedThreshold + )); + 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, + observedThreshold + )); + 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, + BudgetThreshold.EXCEEDED + )); + return state; + } + state.lastNotifiedThreshold = BudgetThreshold.EXCEEDED; + claim.set(new NotificationClaim( + projectedUsage, + List.of(BudgetThreshold.EXCEEDED), + 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 +178,151 @@ 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 BudgetThreshold reachedThreshold(Cost usage, Cost limit) { + if (usage.compareTo(limit) >= 0) { + return BudgetThreshold.EXCEEDED; + } + Cost warningBoundary = Cost.of( + limit.value().multiply(WARNING_RATIO), + limit.currency() + ); + if (usage.compareTo(warningBoundary) >= 0) { + return BudgetThreshold.WARNING; + } + Cost halfBoundary = Cost.of( + limit.value().multiply(HALF_RATIO), + limit.currency() + ); + return usage.compareTo(halfBoundary) >= 0 + ? BudgetThreshold.HALF + : BudgetThreshold.NONE; + } + + 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..da73370 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,74 @@ 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, + BudgetThreshold observedThreshold + ) { + + 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"); + } + Objects.requireNonNull( + observedThreshold, + "observedThreshold must not be null" + ); + if (!thresholds.isEmpty() + && observedThreshold.compareTo(thresholds.get(thresholds.size() - 1)) < 0) { + throw new IllegalArgumentException( + "observedThreshold must include every claimed threshold" + ); + } + } + + public NotificationClaim(Cost usage, List thresholds) { + this( + usage, + thresholds, + thresholds.isEmpty() + ? BudgetThreshold.NONE + : thresholds.get(thresholds.size() - 1) + ); + } + + public static NotificationClaim none(Cost usage) { + return none(usage, BudgetThreshold.NONE); + } + + public static NotificationClaim none( + Cost usage, + BudgetThreshold observedThreshold + ) { + return new NotificationClaim(usage, List.of(), observedThreshold); + } + } } 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..d9131f9 --- /dev/null +++ b/token-pilot-notification/src/test/java/io/tokenpilot/notification/AccountingBudgetNotificationTest.java @@ -0,0 +1,524 @@ +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.HashMap; +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; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.assertj.core.api.Assertions.tuple; + +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")); + assertThat(events) + .extracting(BudgetNotificationEvent::threshold, BudgetNotificationEvent::state) + .containsExactly( + tuple(BudgetThreshold.HALF, io.tokenpilot.budget.BudgetState.ALLOW), + tuple(BudgetThreshold.WARNING, io.tokenpilot.budget.BudgetState.WARN), + tuple(BudgetThreshold.EXCEEDED, io.tokenpilot.budget.BudgetState.BLOCK) + ); + } + + @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 legacy_event는_tags를_불변_복사하고_atomic_event는_tags를_보존하지_않는다() { + Map legacyTags = new HashMap<>(); + legacyTags.put("tenant_id", "tenant-a"); + BudgetNotificationEvent event = new BudgetNotificationEvent( + key("tenant-a", "2026-08"), + BudgetThreshold.HALF, + io.tokenpilot.budget.BudgetState.WARN, + "threshold reached", + usd("50"), + LIMIT, + legacyTags + ); + legacyTags.put("tenant_id", "changed"); + + assertThat(event.tags()).containsExactly(Map.entry("tenant_id", "tenant-a")); + assertThatThrownBy(() -> event.tags().put("another", "value")) + .isInstanceOf(UnsupportedOperationException.class); + + List atomicEvents = new ArrayList<>(); + Fixture fixture = fixture( + List.of(atomicEvents::add), + BudgetNotificationErrorHook.noOp() + ); + commit(fixture, key("tenant-b", "2026-08"), "request-atomic", 50); + + assertThat(atomicEvents).singleElement().satisfies(atomic -> { + assertThat(atomic.tags()).isEmpty(); + assertThat(atomic.toString()).doesNotContain("tenant-a", "changed"); + }); + } + + 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-request" + ); + } + + 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 + ) { + } +} diff --git a/token-pilot-notification/src/test/java/io/tokenpilot/notification/BudgetNotificationLifecycleEventTest.java b/token-pilot-notification/src/test/java/io/tokenpilot/notification/BudgetNotificationLifecycleEventTest.java new file mode 100644 index 0000000..e0dea8f --- /dev/null +++ b/token-pilot-notification/src/test/java/io/tokenpilot/notification/BudgetNotificationLifecycleEventTest.java @@ -0,0 +1,31 @@ +package io.tokenpilot.notification; + +import io.tokenpilot.budget.BudgetThreshold; +import org.junit.jupiter.api.Test; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +class BudgetNotificationLifecycleEventTest { + + @Test + void bounded_outcome과_threshold를_보존한다() { + BudgetNotificationLifecycleEvent event = new BudgetNotificationLifecycleEvent( + BudgetNotificationOutcome.SUCCESS, + BudgetThreshold.WARNING + ); + + assertThat(event.outcome()).isEqualTo(BudgetNotificationOutcome.SUCCESS); + assertThat(event.threshold()).isEqualTo(BudgetThreshold.WARNING); + } + + @Test + void NONE_threshold를_거부한다() { + assertThatThrownBy(() -> new BudgetNotificationLifecycleEvent( + BudgetNotificationOutcome.DEDUPLICATED, + BudgetThreshold.NONE + )) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("threshold must identify a notification boundary"); + } +} 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 6bc580e..d76327c 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 @@ -3,20 +3,32 @@ import io.tokenpilot.budget.BudgetDecision; import io.tokenpilot.budget.BudgetDecision.EvaluationType; import io.tokenpilot.budget.BudgetKey; +import io.tokenpilot.budget.BudgetReservationRequest; +import io.tokenpilot.budget.BudgetReservationResult; import io.tokenpilot.budget.BudgetState; import io.tokenpilot.budget.BudgetThreshold; import io.tokenpilot.budget.BudgetWindow; import io.tokenpilot.core.domain.Cost; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; import org.mockito.ArgumentCaptor; import java.math.BigDecimal; import java.util.Currency; +import java.util.List; import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicBoolean; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.mockito.Mockito.any; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.same; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; @@ -33,8 +45,14 @@ class BudgetNotificationServiceTest { ); BudgetKey key = key("2026-06"); - service.notifyIfNeeded(decision(key, BudgetThreshold.HALF, "50"), Map.of()); - service.notifyIfNeeded(decision(key, BudgetThreshold.HALF, "50"), Map.of()); + service.notifyIfNeeded( + decision(key, BudgetThreshold.HALF, "50"), + Map.of("tenant_id", "acme") + ); + service.notifyIfNeeded( + decision(key, BudgetThreshold.HALF, "50"), + Map.of("tenant_id", "acme") + ); service.notifyIfNeeded(decision(key, BudgetThreshold.WARNING, "80"), Map.of()); verify(handler, times(2)).handle(any()); @@ -63,7 +81,10 @@ class BudgetNotificationServiceTest { BudgetKey key = key("2026-07"); when(store.getLastNotifiedThreshold(key)).thenReturn(BudgetThreshold.NONE); - service.notifyIfNeeded(decision(key, BudgetThreshold.HALF, "50"), Map.of()); + service.notifyIfNeeded( + decision(key, BudgetThreshold.HALF, "50"), + Map.of("tenant_id", "acme") + ); ArgumentCaptor event = ArgumentCaptor.forClass(BudgetNotificationEvent.class); verify(handler).handle(event.capture()); @@ -74,6 +95,226 @@ class BudgetNotificationServiceTest { .isEqualTo(Cost.of(new BigDecimal("50"), Currency.getInstance("USD"))); assertThat(event.getValue().currentUsage()) .isEqualTo(event.getValue().projectedUsage()); + assertThat(event.getValue().tags()) + .containsExactly(Map.entry("tenant_id", "acme")); + } + + @Test + @Timeout(value = 10, unit = TimeUnit.SECONDS) + void legacy_threshold_claim은_병렬_호출에서도_한_번만_전달한다() throws Exception { + AtomicInteger deliveries = new AtomicInteger(); + CountDownLatch handlerEntered = new CountDownLatch(1); + CountDownLatch releaseHandler = new CountDownLatch(1); + List lifecycleEvents = + new CopyOnWriteArrayList<>(); + BudgetNotificationService service = new BudgetNotificationService( + event -> { + deliveries.incrementAndGet(); + handlerEntered.countDown(); + try { + if (!releaseHandler.await(5, TimeUnit.SECONDS)) { + throw new IllegalStateException("handler release timed out"); + } + } catch (InterruptedException exception) { + Thread.currentThread().interrupt(); + throw new IllegalStateException("handler interrupted", exception); + } + }, + new InMemoryNotificationStateStore(), + List.of(lifecycleEvents::add) + ); + BudgetDecision half = decision(key("2026-08"), BudgetThreshold.HALF, "50"); + + try (var executor = Executors.newVirtualThreadPerTaskExecutor()) { + var first = executor.submit(() -> service.notifyIfNeeded(half, Map.of())); + assertThat(handlerEntered.await(5, TimeUnit.SECONDS)).isTrue(); + var second = executor.submit(() -> service.notifyIfNeeded(half, Map.of())); + releaseHandler.countDown(); + first.get(5, TimeUnit.SECONDS); + second.get(5, TimeUnit.SECONDS); + } finally { + releaseHandler.countDown(); + } + + assertThat(deliveries).hasValue(1); + assertThat(lifecycleEvents).containsExactlyInAnyOrder( + new BudgetNotificationLifecycleEvent( + BudgetNotificationOutcome.SUCCESS, + BudgetThreshold.HALF + ), + new BudgetNotificationLifecycleEvent( + BudgetNotificationOutcome.DEDUPLICATED, + BudgetThreshold.HALF + ) + ); + } + + @Test + void atomic_claim도_legacy와_같은_store_monitor에서_직렬화한다() { + AtomicBoolean monitorHeld = new AtomicBoolean(); + InMemoryNotificationStateStore store = new InMemoryNotificationStateStore() { + @Override + public NotificationClaim recordBlockedReservation( + BudgetReservationRequest request, + BudgetReservationResult result + ) { + monitorHeld.set(Thread.holdsLock(this)); + return NotificationClaim.none( + Cost.zero(Currency.getInstance("USD")) + ); + } + }; + BudgetNotificationService service = new BudgetNotificationService( + List.of(event -> { }), + store, + ignored -> null + ); + BudgetReservationRequest request = mock(BudgetReservationRequest.class); + BudgetReservationResult result = mock(BudgetReservationResult.class); + when(request.key()).thenReturn(key("2026-08")); + when(request.limit()).thenReturn(Cost.of( + new BigDecimal("100"), + Currency.getInstance("USD") + )); + when(result.reason()).thenReturn("budget exceeded"); + + service.onReservationBlocked(request, result); + + assertThat(monitorHeld).isTrue(); + } + + @Test + void handler_성공과_중복을_lifecycle_event로_발행한다() { + BudgetNotificationHandler handler = mock(BudgetNotificationHandler.class); + List lifecycleEvents = new CopyOnWriteArrayList<>(); + BudgetNotificationService service = new BudgetNotificationService( + handler, + new InMemoryNotificationStateStore(), + List.of(lifecycleEvents::add) + ); + BudgetDecision half = decision(key("2026-08"), BudgetThreshold.HALF, "50"); + + service.notifyIfNeeded(half, Map.of()); + service.notifyIfNeeded(half, Map.of()); + + assertThat(lifecycleEvents).containsExactly( + new BudgetNotificationLifecycleEvent( + BudgetNotificationOutcome.SUCCESS, + BudgetThreshold.HALF + ), + new BudgetNotificationLifecycleEvent( + BudgetNotificationOutcome.DEDUPLICATED, + BudgetThreshold.HALF + ) + ); + verify(handler).handle(any()); + } + + @Test + void handler_실패를_발행하고_기존처럼_호출자에게_전파한다() { + RuntimeException failure = new IllegalStateException("handler failed"); + BudgetNotificationHandler handler = event -> { + throw failure; + }; + NotificationStateStore store = mock(NotificationStateStore.class); + List lifecycleEvents = new CopyOnWriteArrayList<>(); + BudgetNotificationService service = new BudgetNotificationService( + handler, + store, + List.of(lifecycleEvents::add) + ); + BudgetKey key = key("2026-08"); + when(store.getLastNotifiedThreshold(key)).thenReturn(BudgetThreshold.NONE); + + assertThatThrownBy(() -> service.notifyIfNeeded( + decision(key, BudgetThreshold.WARNING, "80"), + Map.of() + )).isSameAs(failure); + + assertThat(lifecycleEvents).containsExactly( + new BudgetNotificationLifecycleEvent( + BudgetNotificationOutcome.FAILURE, + BudgetThreshold.WARNING + ) + ); + verify(store, never()).updateLastNotifiedThreshold(any(), any()); + } + + @Test + void legacy_handler_실패는_claim하지_않아_다음_호출이_재시도한다() { + AtomicInteger attempts = new AtomicInteger(); + BudgetNotificationService service = new BudgetNotificationService( + event -> { + if (attempts.incrementAndGet() == 1) { + throw new IllegalStateException("first delivery failed"); + } + }, + new InMemoryNotificationStateStore() + ); + BudgetDecision warning = decision( + key("2026-08"), + BudgetThreshold.WARNING, + "80" + ); + + assertThatThrownBy(() -> service.notifyIfNeeded(warning, Map.of())) + .isInstanceOf(IllegalStateException.class) + .hasMessage("first delivery failed"); + + service.notifyIfNeeded(warning, Map.of()); + + assertThat(attempts).hasValue(2); + } + + @Test + void lifecycle_listener_실패는_handler와_다음_listener를_바꾸지_않는다() { + BudgetNotificationHandler handler = mock(BudgetNotificationHandler.class); + AtomicInteger failedDeliveries = new AtomicInteger(); + List received = new CopyOnWriteArrayList<>(); + BudgetNotificationService service = new BudgetNotificationService( + handler, + new InMemoryNotificationStateStore(), + List.of( + event -> { + failedDeliveries.incrementAndGet(); + throw new IllegalStateException("listener failed"); + }, + received::add + ) + ); + + service.notifyIfNeeded( + decision(key("2026-08"), BudgetThreshold.HALF, "50"), + Map.of() + ); + + assertThat(failedDeliveries).hasValue(1); + assertThat(received).containsExactly( + new BudgetNotificationLifecycleEvent( + BudgetNotificationOutcome.SUCCESS, + BudgetThreshold.HALF + ) + ); + verify(handler).handle(any()); + } + + @Test + void NONE은_lifecycle_event를_발행하지_않는다() { + List lifecycleEvents = new CopyOnWriteArrayList<>(); + BudgetNotificationHandler handler = mock(BudgetNotificationHandler.class); + BudgetNotificationService service = new BudgetNotificationService( + handler, + new InMemoryNotificationStateStore(), + List.of(lifecycleEvents::add) + ); + + service.notifyIfNeeded( + decision(key("2026-08"), BudgetThreshold.NONE, "0"), + Map.of() + ); + + assertThat(lifecycleEvents).isEmpty(); + verify(handler, never()).handle(any()); } private static BudgetKey key(String window) { 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 1065052..8f9f89e 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 @@ -53,6 +53,9 @@ public Map tokenLedgerBeans() { "ledgerManager", applicationContext.containsBean("ledgerManager"), "ledgerAdvisor", applicationContext.containsBean("ledgerAdvisor"), "pricingRegistry", applicationContext.containsBean("pricingRegistry"), + "tokenPilotCoreMetricsPublisher", applicationContext.containsBean("tokenPilotCoreMetricsPublisher"), + "tokenPilotBudgetMetricsPublisher", applicationContext.containsBean("tokenPilotBudgetMetricsPublisher"), + "tokenPilotNotificationMetricsPublisher", applicationContext.containsBean("tokenPilotNotificationMetricsPublisher"), "microCostMetricsPublisher", applicationContext.containsBean("microCostMetricsPublisher"), "budgetEvaluator", applicationContext.containsBean("budgetEvaluator"), "budgetStateStore", applicationContext.containsBean("budgetStateStore") diff --git a/token-pilot-sample-app/src/main/resources/application.yml b/token-pilot-sample-app/src/main/resources/application.yml index ae3ce58..53589fd 100644 --- a/token-pilot-sample-app/src/main/resources/application.yml +++ b/token-pilot-sample-app/src/main/resources/application.yml @@ -21,5 +21,5 @@ token-pilot: COMPLETION: 0.00060 metrics: enabled: true - tag-whitelist: - - tenant_id + legacy-ai-token-metrics-enabled: false + tag-whitelist: [] diff --git a/token-pilot-sample-app/src/test/java/io/tokenpilot/sample/SampleApplicationChatClientE2ETest.java b/token-pilot-sample-app/src/test/java/io/tokenpilot/sample/SampleApplicationChatClientE2ETest.java index 2298e39..3eefbcb 100644 --- a/token-pilot-sample-app/src/test/java/io/tokenpilot/sample/SampleApplicationChatClientE2ETest.java +++ b/token-pilot-sample-app/src/test/java/io/tokenpilot/sample/SampleApplicationChatClientE2ETest.java @@ -1,169 +1,965 @@ package io.tokenpilot.sample; +import io.micrometer.core.instrument.MeterRegistry; +import io.tokenpilot.budget.ActualUsageCommand; +import io.tokenpilot.budget.AtomicBudgetStateStore; import io.tokenpilot.budget.BudgetDecision; +import io.tokenpilot.budget.BudgetEvaluator; +import io.tokenpilot.budget.BudgetReservationRequest; +import io.tokenpilot.budget.BudgetReservationResult; +import io.tokenpilot.budget.BudgetSnapshot; import io.tokenpilot.budget.BudgetStateStore; +import io.tokenpilot.budget.ReservationAccounting; +import io.tokenpilot.budget.ReservationAccountingEvent; +import io.tokenpilot.budget.ReservationAccountingListener; +import io.tokenpilot.budget.ReservationAccountingReason; +import io.tokenpilot.budget.ReservationId; +import io.tokenpilot.budget.ReservationReconciliationRequiredEvent; +import io.tokenpilot.budget.ReservationStatus; +import io.tokenpilot.budget.exception.BudgetExceededException; +import io.tokenpilot.budget.internal.LedgerBudgetComponents; +import io.tokenpilot.core.CostCalculator; +import io.tokenpilot.core.domain.TokenUsage; import io.tokenpilot.core.domain.Cost; import io.tokenpilot.core.domain.PricingPlan; -import io.tokenpilot.core.domain.PricingReconciliationResult; -import io.tokenpilot.core.domain.PricingResolution; -import io.tokenpilot.core.domain.PricingSnapshot; +import io.tokenpilot.core.exception.MissingPricingException; +import io.tokenpilot.notification.BudgetNotificationEvent; +import io.tokenpilot.notification.BudgetNotificationHandler; +import io.tokenpilot.notification.BudgetNotificationSource; +import io.tokenpilot.springai.UsageExtractor; +import io.tokenpilot.springai.internal.LedgerSpringAiComponents; +import io.tokenpilot.core.internal.LedgerComponents; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Test; import org.springframework.ai.chat.client.ChatClient; import org.springframework.ai.chat.client.ChatClientBuilderCustomizer; +import org.springframework.ai.chat.client.ChatClientRequest; import org.springframework.ai.chat.client.ChatClientResponse; +import org.springframework.ai.chat.client.advisor.api.CallAdvisor; +import org.springframework.ai.chat.client.advisor.api.CallAdvisorChain; import org.springframework.ai.chat.messages.AssistantMessage; +import org.springframework.ai.chat.messages.Message; +import org.springframework.ai.chat.messages.UserMessage; import org.springframework.ai.chat.metadata.ChatResponseMetadata; import org.springframework.ai.chat.metadata.DefaultUsage; +import org.springframework.ai.chat.metadata.Usage; import org.springframework.ai.chat.model.ChatModel; import org.springframework.ai.chat.model.ChatResponse; import org.springframework.ai.chat.model.Generation; import org.springframework.ai.chat.prompt.ChatOptions; import org.springframework.ai.chat.prompt.Prompt; +import org.springframework.ai.content.Media; +import org.springframework.ai.model.tool.ToolCallingChatOptions; +import org.springframework.ai.tool.ToolCallback; import org.springframework.beans.factory.ObjectProvider; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.boot.test.context.SpringBootTest; import org.springframework.boot.test.context.TestConfiguration; -import org.springframework.boot.test.web.server.LocalServerPort; import org.springframework.context.annotation.Bean; -import org.springframework.web.bind.annotation.GetMapping; -import org.springframework.web.bind.annotation.RestController; - -import java.io.IOException; -import java.net.URI; -import java.net.http.HttpClient; -import java.net.http.HttpRequest; -import java.net.http.HttpResponse; +import org.springframework.util.MimeTypeUtils; + +import java.time.Clock; import java.util.Currency; +import java.util.ArrayList; import java.util.List; import java.util.Map; +import java.util.UUID; +import java.util.concurrent.atomic.AtomicInteger; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.AdditionalAnswers.delegatesTo; +import static org.mockito.ArgumentCaptor.forClass; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.argThat; +import static org.mockito.Mockito.clearInvocations; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.reset; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; +import static org.mockito.Mockito.when; +import static org.mockito.Mockito.withSettings; -@SpringBootTest( - webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT, - properties = { - "token-pilot.enabled=true", - "token-pilot.pricing.plans[0].model-id=fake-chat-model", - "token-pilot.pricing.plans[0].currency=USD", - "token-pilot.pricing.plans[0].rates.PROMPT=0.00015", - "token-pilot.pricing.plans[0].rates.COMPLETION=0.00060", - "token-pilot.metrics.enabled=true", - "token-pilot.metrics.tag-whitelist[0]=tenant_id", - "token-pilot.budget.enabled=true", - "token-pilot.budget.monthly-limit=10.00", - "token-pilot.budget.currency=USD", - "token-pilot.budget.target-tag-key=tenant_id", - "management.endpoints.web.exposure.include=prometheus,health" - } -) +@SpringBootTest(properties = { + "token-pilot.enabled=true", + "token-pilot.pricing.plans[0].model-id=gpt-4o-2024-08-06", + "token-pilot.pricing.plans[0].currency=USD", + "token-pilot.pricing.plans[0].rates.PROMPT=0.00015", + "token-pilot.pricing.plans[0].rates.COMPLETION=0.00060", + "token-pilot.budget.enabled=true", + "token-pilot.budget.monthly-limit=10.00", + "token-pilot.budget.currency=USD", + "token-pilot.budget.target-tag-key=tenant_id", + "token-pilot.notification.enabled=true" +}) class SampleApplicationChatClientE2ETest { - private final HttpClient httpClient = HttpClient.newHttpClient(); + @Autowired + private ChatClient.Builder chatClientBuilder; - @LocalServerPort - private int port; + @Autowired + private ProviderProbe providerProbe; @Autowired - private ChatClient.Builder chatClientBuilder; + private BudgetStateStore stateStoreProbe; @Autowired - private BudgetStateStore budgetStateStore; + private BudgetEvaluator budgetEvaluator; - @Test - void chatClientAdvisorRecordsTokenPilotMetricsEndToEnd() throws Exception { - HttpResponse beans = get("/test/token-pilot/beans"); - assertThat(beans.statusCode()).isEqualTo(200); - assertThat(beans.body()) - .contains("\"ledgerAdvisor\":true") - .contains("\"microCostMetricsPublisher\":true"); + @Autowired + private UsageExtractorProbe usageExtractorProbe; - HttpResponse chat = get("/test/token-pilot/chat"); - assertThat(chat.statusCode()).isEqualTo(200); - assertThat(chat.body()) - .contains("\"available\":\"true\"") - .contains("\"content\":\"fake chat response\""); + @Autowired + private AccountingListenerProbe accountingListenerProbe; - HttpResponse prometheus = get("/actuator/prometheus"); - assertThat(prometheus.statusCode()).isEqualTo(200); - assertThat(prometheus.body()) - .contains("ai_token_usage_total") - .contains("ai_token_cost_total") - .contains("model=\"fake-chat-model\"") - .contains("tenant_id=\"chat-sample-tenant\"") - .doesNotContain("user_id=\"chat-sample-user\""); + @Autowired + private AccountingCostProbe accountingCostProbe; + + @Autowired + private ReservationLifecycleProbe reservationLifecycleProbe; + + @Autowired + private MeterRegistry meterRegistry; + + @Autowired + private NotificationProbe notificationProbe; + + @BeforeEach + void resetProbes() { + providerProbe.reset(); + usageExtractorProbe.reset(); + accountingListenerProbe.reset(); + accountingCostProbe.reset(); + reservationLifecycleProbe.reset(); + notificationProbe.reset(); + reset(stateStoreProbe); } @Test - void budgetAdvisorResolvesModelAndPolicyFromRegularChatClientCall() { + @DisplayName("정상 호출은 예약 후 IN_FLIGHT를 거쳐 actual 비용을 한 번 commit한다") + void providerAndAccountingLifecycleRunOnce() { ChatClientResponse response = chatClientBuilder.clone() .build() .prompt() .user("Record this fake budget-aware Spring AI call.") - .advisors(advisors -> advisors.param("tenant_id", "budget-chat-tenant")) + .options(ChatOptions.builder() + .model("gpt-4o-2024-08-06") + .maxTokens(100)) + .advisors(advisors -> advisors + .param("tenant_id", "budget-chat-tenant") + .param("tokenpilot.request.id", "request-1") + .param("tokenpilot.attempt.id", "attempt-1")) .call() .chatClientResponse(); - PricingSnapshot snapshot = contextValue(response, PricingSnapshot.class); - PricingResolution resolution = contextValue(response, PricingResolution.class); - PricingReconciliationResult reconciliationResult = contextValue( - response, - PricingReconciliationResult.class + BudgetDecision decision = budgetEvaluator.evaluate( + Map.of("tenant_id", "budget-chat-tenant") + ); + BudgetSnapshot snapshot = stateStoreProbe.snapshot( + decision.key(), + decision.limit() + ); + ReservationAccounting accountingProbe = (ReservationAccounting) stateStoreProbe; + + assertThat(response.chatResponse().getResult().getOutput().getText()) + .isEqualTo("fake chat response"); + assertThat(providerProbe.invocationCount()).isEqualTo(1); + verify(stateStoreProbe).checkAndReserve(any(BudgetReservationRequest.class)); + verify(accountingProbe).markInFlight(any(ReservationId.class)); + verify(accountingProbe).commit(any(ActualUsageCommand.class)); + verify(accountingProbe, never()).release( + any(ReservationId.class), + any(ReservationAccountingReason.class) ); - BudgetDecision budgetDecision = contextValue(response, BudgetDecision.class); - Cost accumulatedCost = budgetStateStore.getAccumulatedCost( - budgetDecision.key(), - budgetDecision.limit() + verify(accountingProbe, never()).markReconciliationRequired( + any(ReservationId.class), + any(ReservationAccountingReason.class) ); + assertThat(snapshot.committedCost().value()).isEqualByComparingTo("0.00135"); + assertThat(snapshot.activeReservedCost().value()).isZero(); + assertThat(snapshot.pendingReconciliationLiability().value()).isZero(); + } - assertThat(snapshot.modelId()).isEqualTo("fake-chat-model"); - assertThat(snapshot.pricingPolicyId()).isEqualTo(PricingPlan.DEFAULT_PRICING_POLICY_ID); - assertThat(snapshot.currency()).isEqualTo(Currency.getInstance("USD")); - assertThat(resolution).isEqualTo(PricingResolution.RESOLVED); - assertThat(reconciliationResult).isEqualTo(PricingReconciliationResult.RECONCILED); - assertThat(accumulatedCost.value()).isEqualByComparingTo("0.00135"); - assertThat(accumulatedCost.currency()).isEqualTo(Currency.getInstance("USD")); + @Test + @DisplayName("request와 기본 model이 없으면 provider 호출과 예약 전에 차단한다") + void missingRequestAndDefaultModelStopsBeforeProviderAndReservation() { + assertThatThrownBy(() -> chatClientBuilder.clone() + .build() + .prompt() + .user("This request has no model.") + .options(ChatOptions.builder().maxTokens(100)) + .advisors(advisors -> advisors + .param("tenant_id", "missing-model-tenant") + .param("tokenpilot.request.id", "request-missing-model") + .param("tokenpilot.attempt.id", "attempt-missing-model")) + .call() + .chatClientResponse()) + .isInstanceOf(IllegalStateException.class) + .hasMessage("MODEL_UNRESOLVED"); + + assertThat(providerProbe.invocationCount()).isZero(); + verifyNoInteractions(stateStoreProbe); } - private T contextValue(ChatClientResponse response, Class type) { - return response.context().values().stream() - .filter(type::isInstance) - .map(type::cast) - .findFirst() - .orElseThrow(); + @Test + @DisplayName("request와 기본 reserved output이 없으면 provider 호출과 예약 전에 차단한다") + void missingRequestAndDefaultOutputStopsBeforeProviderAndReservation() { + assertThatThrownBy(() -> chatClientBuilder.clone() + .build() + .prompt() + .user("This request has no reserved output.") + .options(ChatOptions.builder().model("gpt-4o-2024-08-06")) + .advisors(advisors -> advisors + .param("tenant_id", "missing-output-tenant") + .param("tokenpilot.request.id", "request-missing-output") + .param("tokenpilot.attempt.id", "attempt-missing-output")) + .call() + .chatClientResponse()) + .isInstanceOf(IllegalStateException.class) + .hasMessage("OUTPUT_RESERVATION_UNRESOLVED"); + + assertThat(providerProbe.invocationCount()).isZero(); + verifyNoInteractions(stateStoreProbe); } - private HttpResponse get(String path) throws IOException, InterruptedException { - HttpRequest request = HttpRequest.newBuilder() - .uri(URI.create("http://localhost:" + port + path)) - .GET() + @Test + @DisplayName("context 한도를 넘는 요청은 provider 호출과 예약 전에 차단한다") + void contextAdmissionStopsBeforeProviderAndReservation() { + assertThatThrownBy(() -> chatClientBuilder.clone() + .build() + .prompt() + .user("x".repeat(150_000)) + .options(ChatOptions.builder() + .model("gpt-4o-2024-08-06") + .maxTokens(100)) + .advisors(advisors -> advisors + .param("tenant_id", "context-blocked-tenant") + .param("tokenpilot.request.id", "request-context-blocked") + .param("tokenpilot.attempt.id", "attempt-context-blocked")) + .call() + .chatClientResponse()) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("CONTEXT_EXCEEDED"); + + assertThat(providerProbe.invocationCount()).isZero(); + verifyNoInteractions(stateStoreProbe); + } + + @Test + @DisplayName("예산 BLOCK 요청은 원자적 BLOCKED 전이를 발행하고 provider나 예약을 만들지 않는다") + void budgetBlockPublishesAtomicTransitionWithoutProviderOrReservation() { + Map tags = Map.of("tenant_id", "budget-blocked-tenant"); + BudgetDecision initial = budgetEvaluator.evaluate(tags); + stateStoreProbe.addCost(initial.key(), initial.limit(), initial.limit()); + BudgetSnapshot before = stateStoreProbe.snapshot(initial.key(), initial.limit()); + var blockedCounterBefore = meterRegistry + .find("tokenpilot.budget.reservations") + .tag("state", "blocked") + .counter(); + double blockedMeterBefore = blockedCounterBefore == null + ? 0.0 + : blockedCounterBefore.count(); + clearInvocations(stateStoreProbe); + + assertThatThrownBy(() -> chatClientBuilder.clone() + .build() + .prompt() + .user("This request is over budget.") + .options(ChatOptions.builder() + .model("gpt-4o-2024-08-06") + .maxTokens(100)) + .advisors(advisors -> advisors + .param("tenant_id", "budget-blocked-tenant") + .param("tokenpilot.request.id", "request-budget-blocked") + .param("tokenpilot.attempt.id", "attempt-budget-blocked")) + .call() + .chatClientResponse()) + .isInstanceOf(BudgetExceededException.class); + + BudgetSnapshot after = stateStoreProbe.snapshot(initial.key(), initial.limit()); + assertThat(providerProbe.invocationCount()).isZero(); + verify(stateStoreProbe).checkAndReserve(argThat(reservation -> + reservation.key().equals(initial.key()) + && reservation.limit().equals(initial.limit()) + )); + assertThat(reservationLifecycleProbe.evaluatedStatuses()) + .containsExactly(ReservationStatus.BLOCKED); + assertThat(reservationLifecycleProbe.blockedStatuses()) + .containsExactly(ReservationStatus.BLOCKED); + assertThat(meterRegistry.get("tokenpilot.budget.reservations") + .tag("state", "blocked") + .counter() + .count()).isEqualTo(blockedMeterBefore + 1.0); + assertThat(notificationProbe.events()) + .extracting(BudgetNotificationEvent::source) + .containsExactly(BudgetNotificationSource.RESERVATION_BLOCK); + assertThat(after).isEqualTo(before); + } + + @Test + @DisplayName("pricing이 없는 요청은 provider 호출과 예약 전에 차단한다") + void missingPricingStopsBeforeProviderAndReservation() { + assertThatThrownBy(() -> chatClientBuilder.clone() + .build() + .prompt() + .user("This model has no configured pricing.") + .options(ChatOptions.builder() + .model("gpt-4o-mini-2024-07-18") + .maxTokens(100)) + .advisors(advisors -> advisors + .param("tenant_id", "missing-pricing-tenant") + .param("tokenpilot.request.id", "request-missing-pricing") + .param("tokenpilot.attempt.id", "attempt-missing-pricing")) + .call() + .chatClientResponse()) + .isInstanceOf(MissingPricingException.class) + .hasMessageContaining("MISSING_PLAN"); + + assertThat(providerProbe.invocationCount()).isZero(); + verifyNoInteractions(stateStoreProbe); + } + + @Test + @DisplayName("예약 후 dispatch 이전 오류는 예약을 한 번 release한다") + void errorBeforeDispatchReleasesReservationOnce() { + ReservationAccounting accountingProbe = (ReservationAccounting) stateStoreProbe; + doThrow(new IllegalStateException("dispatch preparation failed")) + .when(accountingProbe) + .markInFlight(any(ReservationId.class)); + + assertThatThrownBy(() -> call("pre-dispatch-tenant", "pre-dispatch")) + .isInstanceOf(IllegalStateException.class) + .hasMessage("dispatch preparation failed"); + + BudgetSnapshot snapshot = snapshot("pre-dispatch-tenant"); + assertThat(providerProbe.invocationCount()).isZero(); + verify(accountingProbe).releaseBeforeDispatch(any(ReservationId.class)); + assertThat(snapshot.committedCost().value()).isZero(); + assertThat(snapshot.activeReservedCost().value()).isZero(); + assertThat(snapshot.pendingReconciliationLiability().value()).isZero(); + } + + @Test + @DisplayName("IN_FLIGHT 이후 downstream 오류는 estimate를 pending liability로 보존한다") + void downstreamErrorPreservesPendingLiability() { + providerProbe.failWith(new IllegalStateException("provider failed")); + + assertThatThrownBy(() -> call("downstream-error-tenant", "downstream-error")) + .isInstanceOf(IllegalStateException.class) + .hasMessage("provider failed"); + + ReservationAccounting accountingProbe = (ReservationAccounting) stateStoreProbe; + BudgetSnapshot snapshot = snapshot("downstream-error-tenant"); + assertThat(providerProbe.invocationCount()).isEqualTo(1); + verify(accountingProbe).markReconciliationRequired( + any(ReservationId.class), + any(ReservationAccountingReason.class) + ); + verify(accountingProbe, never()).commit(any(ActualUsageCommand.class)); + verify(accountingProbe, never()).release( + any(ReservationId.class), + any(ReservationAccountingReason.class) + ); + assertThat(snapshot.committedCost().value()).isZero(); + assertThat(snapshot.activeReservedCost().value()).isZero(); + assertThat(snapshot.pendingReconciliationLiability().value()).isPositive(); + } + + @Test + @DisplayName("성공 응답의 usage가 없으면 응답을 보존하고 pending liability로 남긴다") + void unavailableUsagePreservesResponseAndPendingLiability() { + providerProbe.withoutUsage(); + + ChatClientResponse response = call("usage-unavailable-tenant", "usage-unavailable"); + + ReservationAccounting accountingProbe = (ReservationAccounting) stateStoreProbe; + BudgetSnapshot snapshot = snapshot("usage-unavailable-tenant"); + assertThat(response.chatResponse().getResult().getOutput().getText()) + .isEqualTo("fake chat response"); + assertThat(providerProbe.invocationCount()).isEqualTo(1); + verify(accountingProbe).markReconciliationRequired( + any(ReservationId.class), + any(ReservationAccountingReason.class) + ); + verify(accountingProbe, never()).commit(any(ActualUsageCommand.class)); + assertThat(snapshot.committedCost().value()).isZero(); + assertThat(snapshot.activeReservedCost().value()).isZero(); + assertThat(snapshot.pendingReconciliationLiability().value()).isPositive(); + } + + @Test + @DisplayName("usage 추출 오류는 provider 응답을 뒤집지 않고 pending liability로 남긴다") + void extractionErrorPreservesResponseAndPendingLiability() { + usageExtractorProbe.failWith(new IllegalStateException("extraction failed")); + + ChatClientResponse response = call("extraction-error-tenant", "extraction-error"); + + ReservationAccounting accountingProbe = (ReservationAccounting) stateStoreProbe; + BudgetSnapshot snapshot = snapshot("extraction-error-tenant"); + assertThat(response.chatResponse().getResult().getOutput().getText()) + .isEqualTo("fake chat response"); + verify(accountingProbe).markReconciliationRequired( + any(ReservationId.class), + any(ReservationAccountingReason.class) + ); + verify(accountingProbe, never()).commit(any(ActualUsageCommand.class)); + assertThat(snapshot.pendingReconciliationLiability().value()).isPositive(); + } + + @Test + @DisplayName("actual 비용 통화가 reservation과 다르면 응답을 보존하고 pending liability로 남긴다") + void actualCurrencyMismatchPreservesResponseAndPendingLiability() { + ReservationAccounting accountingProbe = (ReservationAccounting) stateStoreProbe; + accountingCostProbe.returnCurrency(Currency.getInstance("EUR")); + + ChatClientResponse response = call("commit-error-tenant", "commit-error"); + + BudgetSnapshot snapshot = snapshot("commit-error-tenant"); + assertThat(response.chatResponse().getResult().getOutput().getText()) + .isEqualTo("fake chat response"); + verify(accountingProbe).markReconciliationRequired( + any(ReservationId.class), + any(ReservationAccountingReason.class) + ); + assertThat(snapshot.committedCost().value()).isZero(); + assertThat(snapshot.activeReservedCost().value()).isZero(); + assertThat(snapshot.pendingReconciliationLiability().value()).isPositive(); + } + + @Test + @DisplayName("response model이 다르면 request pricing으로 commit하지 않고 pending으로 남긴다") + void responseModelMismatchRequiresPricingReconciliation() { + providerProbe.respondAs("provider-routed-model-v2"); + + ChatClientResponse response = call( + "model-correlation-tenant", + "model-correlation" + ); + + ReservationAccounting accountingProbe = (ReservationAccounting) stateStoreProbe; + verify(stateStoreProbe).checkAndReserve(argThat( + reservation -> reservation.modelId().equals("gpt-4o-2024-08-06") + )); + verify(accountingProbe).markReconciliationRequired( + argThat((ActualUsageCommand command) -> command.requestId() + .equals("request-model-correlation")), + org.mockito.ArgumentMatchers.eq( + ReservationAccountingReason.PRICING_RECONCILIATION_REQUIRED + ) + ); + verify(accountingProbe, never()).commit(any(ActualUsageCommand.class)); + assertThat(response.chatResponse().getResult().getOutput().getText()) + .isEqualTo("fake chat response"); + assertThat(snapshot("model-correlation-tenant") + .pendingReconciliationLiability().value()).isPositive(); + } + + @Test + @DisplayName("cache read·create와 reasoning usage를 정규화해 actual commit에 전달한다") + void normalizedCacheAndReasoningUsageReachesAccountingCommit() { + Usage usage = mock(Usage.class); + when(usage.getPromptTokens()).thenReturn(50); + when(usage.getCompletionTokens()).thenReturn(60); + when(usage.getNativeUsage()).thenReturn(Map.of( + "input_tokens", 50, + "cache_read_input_tokens", 100, + "cache_creation_input_tokens", 25, + "candidatesTokenCount", 60, + "thoughtsTokenCount", 20 + )); + providerProbe.withUsage(usage); + + call("normalized-usage-tenant", "normalized-usage"); + + ReservationAccounting accountingProbe = (ReservationAccounting) stateStoreProbe; + verify(accountingProbe).commit(argThat(command -> + command.usage().inputTokens() == 175 + && command.usage().outputTokens() == 80 + && command.usage().details().cacheReadInputTokens() == 100 + && command.usage().details().cacheCreationInputTokens() == 25 + && command.usage().details().reasoningOutputTokens() == 20 + )); + } + + @Test + @DisplayName("tool schema 요청은 provider 호출과 예약 전에 지원하지 않는 scope로 차단한다") + void unsupportedToolSchemaStopsBeforeProviderAndReservation() { + ToolCallback toolCallback = mock(ToolCallback.class); + + assertThatThrownBy(() -> chatClientBuilder.clone() + .build() + .prompt() + .user("Do not dispatch this tool request.") + .options(ToolCallingChatOptions.builder() + .model("gpt-4o-2024-08-06") + .maxTokens(100) + .toolCallbacks(toolCallback)) + .advisors(advisors -> advisors + .param("tenant_id", "unsupported-scope-tenant") + .param("tokenpilot.request.id", "request-unsupported-scope") + .param("tokenpilot.attempt.id", "attempt-unsupported-scope")) + .call() + .chatClientResponse()) + .isInstanceOf(IllegalStateException.class) + .hasMessage("UNSUPPORTED_REQUEST_SCOPE: TOOL_SCHEMA"); + + assertThat(providerProbe.invocationCount()).isZero(); + verifyNoInteractions(stateStoreProbe); + } + + @Test + @DisplayName("media 요청은 provider 호출과 예약 전에 지원하지 않는 scope로 차단한다") + void unsupportedMediaStopsBeforeProviderAndReservation() { + Media media = Media.builder() + .mimeType(MimeTypeUtils.IMAGE_PNG) + .data(new byte[]{1}) .build(); - return httpClient.send(request, HttpResponse.BodyHandlers.ofString()); + + assertThatThrownBy(() -> chatClientBuilder.clone() + .build() + .prompt() + .user(user -> user + .text("Do not dispatch this media request.") + .media(media)) + .options(ChatOptions.builder() + .model("gpt-4o-2024-08-06") + .maxTokens(100)) + .advisors(advisors -> advisors + .param("tenant_id", "unsupported-media-tenant") + .param("tokenpilot.request.id", "request-unsupported-media") + .param("tokenpilot.attempt.id", "attempt-unsupported-media")) + .call() + .chatClientResponse()) + .isInstanceOf(IllegalStateException.class) + .hasMessage("UNSUPPORTED_REQUEST_SCOPE: MEDIA"); + + assertThat(providerProbe.invocationCount()).isZero(); + verifyNoInteractions(stateStoreProbe); + } + + @Test + @DisplayName("structured output 요청은 provider 호출과 예약 전에 지원하지 않는 scope로 차단한다") + void unsupportedStructuredOutputStopsBeforeProviderAndReservation() { + assertThatThrownBy(() -> chatClientBuilder.clone() + .build() + .prompt() + .user("Do not dispatch this structured output request.") + .options(ChatOptions.builder() + .model("gpt-4o-2024-08-06") + .maxTokens(100)) + .advisors(advisors -> advisors + .param("tenant_id", "unsupported-structured-output-tenant") + .param("tokenpilot.request.id", "request-unsupported-structured-output") + .param("tokenpilot.attempt.id", "attempt-unsupported-structured-output")) + .call() + .entity(StructuredResponse.class, parameters -> + parameters.useProviderStructuredOutput())) + .isInstanceOf(IllegalStateException.class) + .hasMessage("UNSUPPORTED_REQUEST_SCOPE: STRUCTURED_OUTPUT"); + + assertThat(providerProbe.invocationCount()).isZero(); + verifyNoInteractions(stateStoreProbe); + } + + @Test + @DisplayName("뒤쪽 user Advisor가 추가한 framing도 preflight에 반영해 provider 호출 전에 차단한다") + void downstreamUserAdvisorFramingIsIncludedInPreflight() { + assertThatThrownBy(() -> chatClientBuilder.clone() + .build() + .prompt() + .user("This text alone fits the context window.") + .options(ChatOptions.builder() + .model("gpt-4o-2024-08-06") + .maxTokens(100)) + .advisors(advisors -> advisors + .advisors(new ContextExpandingAdvisor()) + .param("tenant_id", "advisor-order-tenant") + .param("tokenpilot.request.id", "request-advisor-order") + .param("tokenpilot.attempt.id", "attempt-advisor-order")) + .call() + .chatClientResponse()) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("CONTEXT_EXCEEDED"); + + assertThat(providerProbe.invocationCount()).isZero(); + verifyNoInteractions(stateStoreProbe); + } + + @Test + @DisplayName("budget enforcement가 활성화된 streaming은 provider 호출과 예약 전에 차단한다") + void streamingEnforcementStopsBeforeProviderAndReservation() { + assertThatThrownBy(() -> chatClientBuilder.clone() + .build() + .prompt() + .user("Do not dispatch this streaming request.") + .options(ChatOptions.builder() + .model("gpt-4o-2024-08-06") + .maxTokens(100)) + .advisors(advisors -> advisors + .param("tenant_id", "streaming-tenant") + .param("tokenpilot.request.id", "request-streaming") + .param("tokenpilot.attempt.id", "attempt-streaming")) + .stream() + .chatClientResponse() + .blockLast()) + .isInstanceOf(IllegalStateException.class) + .hasMessage("STREAMING_UNSUPPORTED_FOR_ENFORCEMENT"); + + assertThat(providerProbe.invocationCount()).isZero(); + verifyNoInteractions(stateStoreProbe); + } + + @Test + @DisplayName("listener 실패와 중복 commit 명령은 완료된 정산과 event를 반복하지 않는다") + void listenerFailureAndDuplicateCommandDoNotReverseOrRepeatCommit() { + ChatClientResponse response = call("listener-failure-tenant", "listener-failure"); + ReservationAccounting accountingProbe = (ReservationAccounting) stateStoreProbe; + var commandCaptor = forClass(ActualUsageCommand.class); + verify(accountingProbe).commit(commandCaptor.capture()); + BudgetSnapshot committed = snapshot("listener-failure-tenant"); + + accountingProbe.commit(commandCaptor.getValue()); + BudgetSnapshot reused = snapshot("listener-failure-tenant"); + + assertThat(response.chatResponse().getResult().getOutput().getText()) + .isEqualTo("fake chat response"); + assertThat(accountingListenerProbe.deliveryCount()).isEqualTo(1); + assertThat(reused).isEqualTo(committed); + } + + private ChatClientResponse call(String tenantId, String correlationId) { + return chatClientBuilder.clone() + .build() + .prompt() + .user("Run a fake provider lifecycle scenario.") + .options(ChatOptions.builder() + .model("gpt-4o-2024-08-06") + .maxTokens(100)) + .advisors(advisors -> advisors + .param("tenant_id", tenantId) + .param("tokenpilot.request.id", "request-" + correlationId) + .param("tokenpilot.attempt.id", "attempt-" + correlationId)) + .call() + .chatClientResponse(); + } + + private BudgetSnapshot snapshot(String tenantId) { + BudgetDecision decision = budgetEvaluator.evaluate(Map.of("tenant_id", tenantId)); + return stateStoreProbe.snapshot(decision.key(), decision.limit()); + } + + private record StructuredResponse(String value) { + } + + private static final class ContextExpandingAdvisor implements CallAdvisor { + + @Override + public ChatClientResponse adviseCall( + ChatClientRequest request, + CallAdvisorChain chain + ) { + List messages = new ArrayList<>( + request.prompt().getInstructions() + ); + messages.add(new UserMessage("x".repeat(150_000))); + ChatClientRequest expandedRequest = request.mutate() + .prompt(new Prompt(messages, request.prompt().getOptions())) + .build(); + return chain.nextCall(expandedRequest); + } + + @Override + public String getName() { + return "Context Expanding User Advisor"; + } + + @Override + public int getOrder() { + return 1; + } + } + + static final class UsageExtractorProbe implements UsageExtractor { + private final UsageExtractor delegate = LedgerSpringAiComponents.defaultUsageExtractor(); + private RuntimeException failure; + + @Override + public TokenUsage extract(ChatClientResponse response) { + if (failure != null) { + throw failure; + } + return delegate.extract(response); + } + + void failWith(RuntimeException failure) { + this.failure = failure; + } + + void reset() { + failure = null; + } + } + + static final class AccountingListenerProbe { + private final AtomicInteger deliveryCount = new AtomicInteger(); + + void onCommitted(ReservationAccountingEvent event) { + deliveryCount.incrementAndGet(); + throw new IllegalStateException("listener failed"); + } + + int deliveryCount() { + return deliveryCount.get(); + } + + void reset() { + deliveryCount.set(0); + } + } + + static final class AccountingCostProbe { + private final CostCalculator delegate = LedgerComponents.defaultCostCalculator(); + private Currency returnedCurrency; + + Cost calculate(TokenUsage usage, PricingPlan plan) { + Cost calculated = delegate.calculate(usage, plan); + if (returnedCurrency == null) { + return calculated; + } + return Cost.of(calculated.value(), returnedCurrency); + } + + void returnCurrency(Currency currency) { + returnedCurrency = currency; + } + + void reset() { + returnedCurrency = null; + } + } + + static final class ReservationLifecycleProbe + implements ReservationAccountingListener { + private final List evaluatedStatuses = new ArrayList<>(); + private final List blockedStatuses = new ArrayList<>(); + + @Override + public void onCommitted(ReservationAccountingEvent event) { + } + + @Override + public void onReservationEvaluated( + BudgetReservationRequest request, + BudgetReservationResult result + ) { + evaluatedStatuses.add(result.status()); + } + + @Override + public void onReservationBlocked( + BudgetReservationRequest request, + BudgetReservationResult result + ) { + blockedStatuses.add(result.status()); + } + + List evaluatedStatuses() { + return List.copyOf(evaluatedStatuses); + } + + List blockedStatuses() { + return List.copyOf(blockedStatuses); + } + + void reset() { + evaluatedStatuses.clear(); + blockedStatuses.clear(); + } + } + + static final class LateBoundAccountingListeners + implements ReservationAccountingListener { + private final ObjectProvider listeners; + + LateBoundAccountingListeners( + ObjectProvider listeners + ) { + this.listeners = listeners; + } + + @Override + public void onCommitted(ReservationAccountingEvent event) { + listeners.orderedStream().forEach(listener -> listener.onCommitted(event)); + } + + @Override + public void onAccountingApplied( + ReservationAccountingEvent event, + BudgetSnapshot snapshot + ) { + listeners.orderedStream() + .forEach(listener -> listener.onAccountingApplied(event, snapshot)); + } + + @Override + public void onReconciliationRequired( + ReservationReconciliationRequiredEvent event + ) { + listeners.orderedStream() + .forEach(listener -> listener.onReconciliationRequired(event)); + } + + @Override + public void onReservationEvaluated( + BudgetReservationRequest request, + BudgetReservationResult result + ) { + listeners.orderedStream() + .forEach(listener -> listener.onReservationEvaluated(request, result)); + } + + @Override + public void onReservationBlocked( + BudgetReservationRequest request, + BudgetReservationResult result + ) { + listeners.orderedStream() + .forEach(listener -> listener.onReservationBlocked(request, result)); + } + } + + static final class NotificationProbe implements BudgetNotificationHandler { + private final List events = new ArrayList<>(); + + @Override + public void handle(BudgetNotificationEvent event) { + events.add(event); + } + + List events() { + return List.copyOf(events); + } + + void reset() { + events.clear(); + } + } + + static final class ProviderProbe implements ChatModel { + private final AtomicInteger invocationCount = new AtomicInteger(); + private RuntimeException failure; + private boolean usageAvailable = true; + private String responseModelId = "gpt-4o-2024-08-06"; + private Usage usage = new DefaultUsage(1_000, 2_000); + + @Override + public ChatResponse call(Prompt prompt) { + invocationCount.incrementAndGet(); + if (failure != null) { + throw failure; + } + ChatResponseMetadata.Builder metadata = ChatResponseMetadata.builder() + .model(responseModelId); + if (usageAvailable) { + metadata.usage(usage); + } + return new ChatResponse( + List.of(new Generation(new AssistantMessage("fake chat response"))), + metadata.build() + ); + } + + @Override + public ChatOptions getOptions() { + return ToolCallingChatOptions.builder().build(); + } + + int invocationCount() { + return invocationCount.get(); + } + + void failWith(RuntimeException failure) { + this.failure = failure; + } + + void withoutUsage() { + usageAvailable = false; + } + + void respondAs(String responseModelId) { + this.responseModelId = responseModelId; + } + + void withUsage(Usage usage) { + this.usage = usage; + } + + void reset() { + invocationCount.set(0); + failure = null; + usageAvailable = true; + responseModelId = "gpt-4o-2024-08-06"; + usage = new DefaultUsage(1_000, 2_000); + } } @TestConfiguration(proxyBeanMethods = false) static class FakeChatClientConfiguration { @Bean - ChatModel fakeChatModel() { - return new ChatModel() { - @Override - public ChatResponse call(Prompt prompt) { - return new ChatResponse( - List.of(new Generation(new AssistantMessage("fake chat response"))), - ChatResponseMetadata.builder() - .model("fake-chat-model") - .usage(new DefaultUsage(1_000, 2_000)) - .build() - ); - } - - @Override - public ChatOptions getOptions() { - return ChatOptions.builder() - .model("fake-chat-model") - .build(); - } - }; + AtomicBudgetStateStore accountingProbe( + AccountingListenerProbe listenerProbe, + AccountingCostProbe costProbe, + ObjectProvider accountingListeners + ) { + AtomicBudgetStateStore delegate = (AtomicBudgetStateStore) + LedgerBudgetComponents.inMemoryBudgetStateStore( + Clock.systemUTC(), + () -> new ReservationId(UUID.randomUUID().toString()), + costProbe::calculate, + List.of( + listenerProbe::onCommitted, + new LateBoundAccountingListeners(accountingListeners) + ) + ); + return mock( + AtomicBudgetStateStore.class, + withSettings() + .defaultAnswer(delegatesTo(delegate)) + ); + } + + @Bean + UsageExtractorProbe usageExtractor() { + return new UsageExtractorProbe(); + } + + @Bean + AccountingListenerProbe accountingListenerProbe() { + return new AccountingListenerProbe(); + } + + @Bean + AccountingCostProbe accountingCostProbe() { + return new AccountingCostProbe(); + } + + @Bean + ReservationLifecycleProbe reservationLifecycleProbe() { + return new ReservationLifecycleProbe(); + } + + @Bean + NotificationProbe notificationProbe() { + return new NotificationProbe(); + } + + @Bean + ProviderProbe fakeChatModel() { + return new ProviderProbe(); } @Bean @@ -176,32 +972,5 @@ ChatClient.Builder chatClientBuilder( .forEach(customizer -> customizer.customize(builder)); return builder; } - - @RestController - static class FakeChatController { - private final ChatClient.Builder chatClientBuilder; - - FakeChatController(ChatClient.Builder chatClientBuilder) { - this.chatClientBuilder = chatClientBuilder; - } - - @GetMapping("/test/token-pilot/chat") - Map chat() { - String content = chatClientBuilder.clone() - .build() - .prompt() - .user("Record this fake Spring AI call.") - .advisors(advisors -> advisors - .param("tenant_id", "chat-sample-tenant") - .param("user_id", "chat-sample-user")) - .call() - .content(); - - return Map.of( - "available", "true", - "content", content - ); - } - } } } diff --git a/token-pilot-sample-app/src/test/java/io/tokenpilot/sample/SampleApplicationE2ETest.java b/token-pilot-sample-app/src/test/java/io/tokenpilot/sample/SampleApplicationE2ETest.java index 6c73807..5d95472 100644 --- a/token-pilot-sample-app/src/test/java/io/tokenpilot/sample/SampleApplicationE2ETest.java +++ b/token-pilot-sample-app/src/test/java/io/tokenpilot/sample/SampleApplicationE2ETest.java @@ -1,6 +1,9 @@ package io.tokenpilot.sample; +import io.tokenpilot.core.TokenBudget; +import io.tokenpilot.core.TokenEstimator; import org.junit.jupiter.api.Test; +import org.springframework.beans.factory.annotation.Autowired; import org.springframework.boot.test.context.SpringBootTest; import org.springframework.boot.test.web.server.LocalServerPort; @@ -21,7 +24,6 @@ "token-pilot.pricing.plans[0].rates.PROMPT=0.00015", "token-pilot.pricing.plans[0].rates.COMPLETION=0.00060", "token-pilot.metrics.enabled=true", - "token-pilot.metrics.tag-whitelist[0]=tenant_id", "management.endpoints.web.exposure.include=prometheus,health" } ) @@ -32,6 +34,12 @@ class SampleApplicationE2ETest { @LocalServerPort private int port; + @Autowired + private TokenEstimator tokenEstimator; + + @Autowired + private TokenBudget tokenBudget; + @Test void starterEndpointsAndPrometheusMetricsWorkEndToEnd() throws Exception { HttpResponse smoke = get("/test/token-pilot/smoke"); @@ -46,7 +54,10 @@ void starterEndpointsAndPrometheusMetricsWorkEndToEnd() throws Exception { .contains("\"ledgerManager\":true") .contains("\"ledgerAdvisor\":true") .contains("\"pricingRegistry\":true") - .contains("\"microCostMetricsPublisher\":true"); + .contains("\"tokenPilotCoreMetricsPublisher\":true") + .contains("\"tokenPilotBudgetMetricsPublisher\":true") + .contains("\"tokenPilotNotificationMetricsPublisher\":true") + .contains("\"microCostMetricsPublisher\":false"); HttpResponse record = get("/test/token-pilot/record"); assertThat(record.statusCode()).isEqualTo(200); @@ -55,13 +66,22 @@ void starterEndpointsAndPrometheusMetricsWorkEndToEnd() throws Exception { .contains("\"cost\":\"0.001350\"") .contains("\"currency\":\"USD\""); + tokenBudget.check( + "gpt-4o-mini", + tokenEstimator.estimate("sample preflight"), + 256 + ); + HttpResponse prometheus = get("/actuator/prometheus"); assertThat(prometheus.statusCode()).isEqualTo(200); assertThat(prometheus.body()) - .contains("ai_token_usage_total") - .contains("ai_token_usage_distribution") - .contains("ai_token_cost_total") - .contains("tenant_id=\"sample-tenant\"") + .contains("tokenpilot_preflight_requests_total") + .contains("decision=\"indeterminate\"") + .contains("reason=\"incomplete_scope\"") + .doesNotContain("ai_token_usage_total") + .doesNotContain("ai_token_usage_distribution") + .doesNotContain("ai_token_cost_total") + .doesNotContain("tenant_id=\"sample-tenant\"") .doesNotContain("user_id=\"sample-user\""); } diff --git a/token-pilot-sample-app/src/test/java/io/tokenpilot/sample/SampleApplicationLedgerOnlyStreamingE2ETest.java b/token-pilot-sample-app/src/test/java/io/tokenpilot/sample/SampleApplicationLedgerOnlyStreamingE2ETest.java new file mode 100644 index 0000000..a4decb7 --- /dev/null +++ b/token-pilot-sample-app/src/test/java/io/tokenpilot/sample/SampleApplicationLedgerOnlyStreamingE2ETest.java @@ -0,0 +1,155 @@ +package io.tokenpilot.sample; + +import io.tokenpilot.budget.BudgetStateStore; +import io.tokenpilot.budget.ReservationAccounting; +import io.tokenpilot.core.LedgerManager; +import io.tokenpilot.core.domain.PricingSnapshot; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.springframework.ai.chat.client.ChatClient; +import org.springframework.ai.chat.client.ChatClientBuilderCustomizer; +import org.springframework.ai.chat.client.ChatClientResponse; +import org.springframework.ai.chat.messages.AssistantMessage; +import org.springframework.ai.chat.metadata.ChatGenerationMetadata; +import org.springframework.ai.chat.metadata.ChatResponseMetadata; +import org.springframework.ai.chat.metadata.DefaultUsage; +import org.springframework.ai.chat.model.ChatModel; +import org.springframework.ai.chat.model.ChatResponse; +import org.springframework.ai.chat.model.Generation; +import org.springframework.ai.chat.prompt.ChatOptions; +import org.springframework.ai.chat.prompt.Prompt; +import org.springframework.beans.factory.ObjectProvider; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.boot.test.context.SpringBootTest; +import org.springframework.boot.test.context.TestConfiguration; +import org.springframework.context.ApplicationContext; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Import; +import reactor.core.publisher.Flux; + +import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyMap; +import static org.mockito.ArgumentMatchers.argThat; +import static org.mockito.Mockito.clearInvocations; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; + +@SpringBootTest(properties = { + "token-pilot.enabled=true", + "token-pilot.budget.enabled=false", + "token-pilot.pricing.plans[0].model-id=gpt-4o-2024-08-06", + "token-pilot.pricing.plans[0].currency=USD", + "token-pilot.pricing.plans[0].rates.PROMPT=0.00015", + "token-pilot.pricing.plans[0].rates.COMPLETION=0.00060" +}) +@Import(SampleApplicationLedgerOnlyStreamingE2ETest.StreamingConfiguration.class) +class SampleApplicationLedgerOnlyStreamingE2ETest { + + @Autowired + private ChatClient.Builder chatClientBuilder; + + @Autowired + private StreamingProviderProbe providerProbe; + + @Autowired + private LedgerManager ledgerManagerProbe; + + @Autowired + private ApplicationContext applicationContext; + + @BeforeEach + void resetProbes() { + providerProbe.reset(); + clearInvocations(ledgerManagerProbe); + } + + @Test + @DisplayName("budget 비활성 ledger-only streaming은 기존 provider와 terminal usage 기록 경로를 유지한다") + void ledgerOnlyStreamingKeepsExistingProviderAndLedgerPath() { + List responses = chatClientBuilder.clone() + .build() + .prompt() + .user("Run the ledger-only streaming path.") + .options(ChatOptions.builder() + .model("gpt-4o-2024-08-06")) + .stream() + .chatClientResponse() + .collectList() + .block(); + + assertThat(responses).hasSize(1); + assertThat(responses.getFirst().chatResponse().getResult().getOutput().getText()) + .isEqualTo("ledger-only stream response"); + assertThat(providerProbe.invocationCount()).isEqualTo(1); + assertThat(applicationContext.getBeansOfType(BudgetStateStore.class)).isEmpty(); + assertThat(applicationContext.getBeansOfType(ReservationAccounting.class)).isEmpty(); + verify(ledgerManagerProbe).record( + any(PricingSnapshot.class), + argThat(usage -> usage.inputTokens() == 10 && usage.outputTokens() == 5), + anyMap() + ); + } + + static final class StreamingProviderProbe implements ChatModel { + private final AtomicInteger invocationCount = new AtomicInteger(); + + @Override + public ChatResponse call(Prompt prompt) { + throw new UnsupportedOperationException("non-streaming call is not expected"); + } + + @Override + public Flux stream(Prompt prompt) { + invocationCount.incrementAndGet(); + Generation generation = new Generation( + new AssistantMessage("ledger-only stream response"), + ChatGenerationMetadata.builder() + .finishReason("STOP") + .build() + ); + ChatResponseMetadata metadata = ChatResponseMetadata.builder() + .model("gpt-4o-2024-08-06") + .usage(new DefaultUsage(10, 5)) + .build(); + return Flux.just(new ChatResponse(List.of(generation), metadata)); + } + + int invocationCount() { + return invocationCount.get(); + } + + void reset() { + invocationCount.set(0); + } + } + + @TestConfiguration(proxyBeanMethods = false) + static class StreamingConfiguration { + + @Bean + LedgerManager ledgerManagerProbe() { + return mock(LedgerManager.class); + } + + @Bean + StreamingProviderProbe streamingProviderProbe() { + return new StreamingProviderProbe(); + } + + @Bean + ChatClient.Builder chatClientBuilder( + ChatModel chatModel, + ObjectProvider customizers + ) { + ChatClient.Builder builder = ChatClient.builder(chatModel); + customizers.orderedStream() + .forEach(customizer -> customizer.customize(builder)); + return builder; + } + } +} diff --git a/token-pilot-sample-app/src/test/java/io/tokenpilot/sample/SampleApplicationSpringAiFallbackE2ETest.java b/token-pilot-sample-app/src/test/java/io/tokenpilot/sample/SampleApplicationSpringAiFallbackE2ETest.java new file mode 100644 index 0000000..1977d4f --- /dev/null +++ b/token-pilot-sample-app/src/test/java/io/tokenpilot/sample/SampleApplicationSpringAiFallbackE2ETest.java @@ -0,0 +1,80 @@ +package io.tokenpilot.sample; + +import io.tokenpilot.budget.BudgetStateStore; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.springframework.ai.chat.client.ChatClient; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.boot.test.context.SpringBootTest; +import org.springframework.context.annotation.Import; + +import java.nio.charset.StandardCharsets; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.argThat; +import static org.mockito.Mockito.clearInvocations; +import static org.mockito.Mockito.verify; + +@SpringBootTest(properties = { + "token-pilot.enabled=true", + "token-pilot.pricing.plans[0].model-id=gpt-4o-2024-08-06", + "token-pilot.pricing.plans[0].currency=USD", + "token-pilot.pricing.plans[0].rates.PROMPT=0.00015", + "token-pilot.pricing.plans[0].rates.COMPLETION=0.00060", + "token-pilot.budget.enabled=true", + "token-pilot.budget.monthly-limit=10.00", + "token-pilot.budget.currency=USD", + "token-pilot.budget.target-tag-key=tenant_id", + "token-pilot.spring-ai.default-model-id=gpt-4o", + "token-pilot.spring-ai.default-reserved-output-tokens=64", + "token-pilot.spring-ai.framing-headroom-tokens=8" +}) +@Import(SampleApplicationChatClientE2ETest.FakeChatClientConfiguration.class) +class SampleApplicationSpringAiFallbackE2ETest { + + @Autowired + private ChatClient.Builder chatClientBuilder; + + @Autowired + private SampleApplicationChatClientE2ETest.ProviderProbe providerProbe; + + @Autowired + private BudgetStateStore stateStoreProbe; + + @BeforeEach + void resetProbes() { + providerProbe.reset(); + clearInvocations(stateStoreProbe); + } + + @Test + @DisplayName("request model과 maxTokens가 없으면 설정 fallback으로 versioned safe bound 예약을 만든다") + void configuredFallbacksCreateVersionedSafeBoundReservation() { + String message = "Use configured model and output fallbacks."; + long expectedInputSafeUpperBound = ( + "USER:" + message.length() + ":" + message + '\n' + ).getBytes(StandardCharsets.UTF_8).length + 8L; + + chatClientBuilder.clone() + .build() + .prompt() + .user(message) + .advisors(advisors -> advisors + .param("tenant_id", "fallback-tenant") + .param("tokenpilot.request.id", "request-fallback") + .param("tokenpilot.attempt.id", "attempt-fallback")) + .call() + .chatClientResponse(); + + verify(stateStoreProbe).checkAndReserve(argThat( + reservation -> reservation.modelId().equals("gpt-4o-2024-08-06") + && reservation.tokenEstimate().orElseThrow() + .reservedOutputTokens() == 64 + && reservation.tokenEstimate().orElseThrow() + .inputSafeUpperBoundTokens() == expectedInputSafeUpperBound + && reservation.safeUpperBoundCost().value().signum() > 0 + )); + assertThat(providerProbe.invocationCount()).isEqualTo(1); + } +} diff --git a/token-pilot-spring-ai/build.gradle b/token-pilot-spring-ai/build.gradle index 90edee4..0e4601d 100644 --- a/token-pilot-spring-ai/build.gradle +++ b/token-pilot-spring-ai/build.gradle @@ -1,12 +1,15 @@ dependencies { + api platform('org.springframework.boot:spring-boot-dependencies:4.1.0') + api platform('org.springframework.ai:spring-ai-bom:2.0.0') + api project(':token-pilot-core') - implementation project(':token-pilot-budget') + api project(':token-pilot-budget') - implementation 'org.springframework.ai:spring-ai-model' - implementation 'org.springframework.ai:spring-ai-client-chat' + api 'org.springframework.ai:spring-ai-model' + api 'org.springframework.ai:spring-ai-client-chat' // Reactor (스트리밍 처리를 위한 Flux 등) - implementation 'io.projectreactor:reactor-core' + api 'io.projectreactor:reactor-core' testImplementation 'org.junit.jupiter:junit-jupiter-api' testRuntimeOnly 'org.junit.jupiter:junit-jupiter-engine' diff --git a/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/LedgerAdvisor.java b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/LedgerAdvisor.java index 30b4a98..27421f6 100644 --- a/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/LedgerAdvisor.java +++ b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/LedgerAdvisor.java @@ -4,6 +4,7 @@ import org.springframework.ai.chat.client.ChatClientResponse; import org.springframework.ai.chat.client.advisor.api.AdvisorChain; import org.springframework.ai.chat.client.advisor.api.BaseAdvisor; +import org.springframework.core.Ordered; /** * ChatClient 호출 시 토큰 사용량을 가로채서 기록하는 어드바이저 인터페이스. @@ -26,6 +27,6 @@ default String getName() { @Override default int getOrder() { - return 0; + return Ordered.LOWEST_PRECEDENCE - 1; } } diff --git a/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/ReservedOutputTokensResolver.java b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/ReservedOutputTokensResolver.java new file mode 100644 index 0000000..5c0910d --- /dev/null +++ b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/ReservedOutputTokensResolver.java @@ -0,0 +1,15 @@ +package io.tokenpilot.springai; + +import org.springframework.ai.chat.client.ChatClientRequest; + +import java.util.OptionalLong; + +/** Provider 전용 request option에서 reserved output token 수를 해석합니다. */ +@FunctionalInterface +public interface ReservedOutputTokensResolver { + + /** + * 지원하는 provider request이면 양수 token 수를, 지원하지 않으면 빈 값을 반환합니다. + */ + OptionalLong resolve(ChatClientRequest request); +} diff --git a/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/AdaptedMessage.java b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/AdaptedMessage.java new file mode 100644 index 0000000..96b816d --- /dev/null +++ b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/AdaptedMessage.java @@ -0,0 +1,11 @@ +package io.tokenpilot.springai.internal; + +import java.util.Objects; + +record AdaptedMessage(RequestMessageRole role, String text) { + + AdaptedMessage { + Objects.requireNonNull(role, "role must not be null"); + Objects.requireNonNull(text, "text must not be null"); + } +} diff --git a/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/AdaptedRequest.java b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/AdaptedRequest.java new file mode 100644 index 0000000..f0ab8bd --- /dev/null +++ b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/AdaptedRequest.java @@ -0,0 +1,19 @@ +package io.tokenpilot.springai.internal; + +import java.util.List; +import java.util.Objects; + +final class AdaptedRequest { + + private final List messages; + + AdaptedRequest(List messages) { + this.messages = List.copyOf( + Objects.requireNonNull(messages, "messages must not be null") + ); + } + + List messages() { + return messages; + } +} 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 441d574..e305d2c 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,8 +1,19 @@ package io.tokenpilot.springai.internal; +import io.tokenpilot.budget.ActualUsageCommand; import io.tokenpilot.budget.BudgetDecision; import io.tokenpilot.budget.BudgetEvaluator; +import io.tokenpilot.budget.BudgetReservationRequest; +import io.tokenpilot.budget.BudgetReservationResult; +import io.tokenpilot.budget.BudgetState; import io.tokenpilot.budget.BudgetStateStore; +import io.tokenpilot.budget.IdempotencyKey; +import io.tokenpilot.budget.ReservationAccounting; +import io.tokenpilot.budget.ReservationAccountingReason; +import io.tokenpilot.budget.ReservationId; +import io.tokenpilot.budget.ReservationState; +import io.tokenpilot.budget.ReservationTokenEstimate; +import io.tokenpilot.budget.ReservationTransition; import io.tokenpilot.budget.exception.BudgetExceededException; import io.tokenpilot.core.*; import io.tokenpilot.core.domain.*; @@ -13,9 +24,13 @@ import org.springframework.ai.chat.client.ChatClientRequest; import org.springframework.ai.chat.client.ChatClientResponse; import org.springframework.ai.chat.client.advisor.api.AdvisorChain; +import org.springframework.ai.chat.client.advisor.api.CallAdvisorChain; +import org.springframework.ai.chat.client.advisor.api.StreamAdvisorChain; import org.springframework.ai.chat.model.ChatResponse; +import reactor.core.publisher.Flux; import java.util.HashMap; +import java.util.List; import java.util.Map; import java.util.Objects; import java.util.Optional; @@ -28,9 +43,10 @@ * 또한 {@link BudgetEvaluator}를 통해 예산 초과 여부를 사전에 차단하고, * 호출 성공 시 {@link BudgetStateStore}에 비용을 누적합니다. *

- * 현재 {@link #before(ChatClientRequest, AdvisorChain)}의 상태 조회는 이미 BLOCK인 예산의 - * provider 호출 회귀를 막는 legacy guard입니다. 후보 비용을 포함한 admission 근거가 아니며, - * 최종 candidate-aware lifecycle은 #39의 {@code adviseCall()} 경계가 담당합니다. + * {@link #before(ChatClientRequest, AdvisorChain)}는 budget-disabled ledger-only + * compatibility 경로의 상태 guard입니다. budget-enabled 호출은 + * {@link #adviseCall(ChatClientRequest, CallAdvisorChain)}에서 후보 비용을 포함한 + * preflight와 원자적 reservation lifecycle을 수행합니다. */ public class DefaultLedgerAdvisor implements LedgerAdvisor { @@ -49,6 +65,12 @@ public class DefaultLedgerAdvisor implements LedgerAdvisor { private final PricingRegistry pricingRegistry; private final PricingEvaluator pricingEvaluator; private final MissingPricingPolicy missingPricingPolicy; + private final RequestPreflight requestPreflight; + private final ReservationAccounting reservationAccounting; + private final RequestContextAccessor contextAccessor; + private final IdempotencyKeyResolver idempotencyKeyResolver; + private final StreamingRequestPolicy streamingRequestPolicy = new StreamingRequestPolicy(); + private final List pricingMissingListeners; public DefaultLedgerAdvisor(LedgerManager ledgerManager, UsageExtractor usageExtractor) { this(ledgerManager, usageExtractor, null, null, null, null); @@ -89,8 +111,71 @@ public DefaultLedgerAdvisor(LedgerManager ledgerManager, UsageExtractor usageExt CostCalculator costCalculator, PricingRegistry pricingRegistry, PricingEvaluator pricingEvaluator, MissingPricingPolicy missingPricingPolicy) { + this( + ledgerManager, + usageExtractor, + budgetEvaluator, + budgetStateStore, + costCalculator, + pricingRegistry, + pricingEvaluator, + missingPricingPolicy, + null, + null, + null, + null, + List.of() + ); + } + + public DefaultLedgerAdvisor( + LedgerManager ledgerManager, + UsageExtractor usageExtractor, + BudgetEvaluator budgetEvaluator, + BudgetStateStore budgetStateStore, + CostCalculator costCalculator, + PricingRegistry pricingRegistry, + PricingEvaluator pricingEvaluator, + MissingPricingPolicy missingPricingPolicy, + List pricingMissingListeners + ) { + this( + ledgerManager, + usageExtractor, + budgetEvaluator, + budgetStateStore, + costCalculator, + pricingRegistry, + pricingEvaluator, + missingPricingPolicy, + null, + null, + null, + null, + pricingMissingListeners + ); + } + + DefaultLedgerAdvisor( + LedgerManager ledgerManager, + UsageExtractor usageExtractor, + BudgetEvaluator budgetEvaluator, + BudgetStateStore budgetStateStore, + CostCalculator costCalculator, + PricingRegistry pricingRegistry, + PricingEvaluator pricingEvaluator, + MissingPricingPolicy missingPricingPolicy, + RequestPreflight requestPreflight, + ReservationAccounting reservationAccounting, + RequestContextAccessor contextAccessor, + IdempotencyKeyResolver idempotencyKeyResolver, + List pricingMissingListeners + ) { this.ledgerManager = ledgerManager; - this.usageExtractor = usageExtractor; + this.usageExtractor = Objects.requireNonNull( + usageExtractor, + "usageExtractor must not be null" + ); this.budgetEvaluator = budgetEvaluator; this.budgetStateStore = budgetStateStore; this.costCalculator = costCalculator; @@ -103,6 +188,415 @@ public DefaultLedgerAdvisor(LedgerManager ledgerManager, UsageExtractor usageExt missingPricingPolicy, "missingPricingPolicy must not be null" ); + this.pricingMissingListeners = List.copyOf( + Objects.requireNonNull( + pricingMissingListeners, + "pricingMissingListeners must not be null" + ) + ); + this.requestPreflight = requestPreflight; + this.contextAccessor = contextAccessor; + this.idempotencyKeyResolver = idempotencyKeyResolver; + if (requestPreflight == null) { + this.reservationAccounting = null; + return; + } + Objects.requireNonNull(budgetEvaluator, "budgetEvaluator must not be null"); + Objects.requireNonNull(contextAccessor, "contextAccessor must not be null"); + Objects.requireNonNull( + idempotencyKeyResolver, + "idempotencyKeyResolver must not be null" + ); + Objects.requireNonNull( + budgetStateStore, + "budgetStateStore must not be null" + ); + this.reservationAccounting = Objects.requireNonNull( + reservationAccounting, + "reservationAccounting must not be null" + ); + } + + DefaultLedgerAdvisor( + UsageExtractor usageExtractor, + BudgetEvaluator budgetEvaluator, + BudgetStateStore budgetStateStore, + ReservationAccounting reservationAccounting, + RequestPreflight requestPreflight, + RequestContextAccessor contextAccessor, + IdempotencyKeyResolver idempotencyKeyResolver + ) { + this( + null, + usageExtractor, + budgetEvaluator, + budgetStateStore, + null, + null, + LedgerComponents.defaultPricingEvaluator(), + MissingPricingPolicy.FAIL_CLOSED, + requestPreflight, + reservationAccounting, + contextAccessor, + idempotencyKeyResolver, + List.of() + ); + } + + DefaultLedgerAdvisor( + UsageExtractor usageExtractor, + BudgetEvaluator budgetEvaluator, + BudgetStateStore budgetStateStore, + ReservationAccounting reservationAccounting, + RequestPreflight requestPreflight, + PricingEvaluator pricingEvaluator, + RequestContextAccessor contextAccessor, + IdempotencyKeyResolver idempotencyKeyResolver, + List pricingMissingListeners + ) { + this( + null, + usageExtractor, + budgetEvaluator, + budgetStateStore, + null, + null, + pricingEvaluator, + MissingPricingPolicy.FAIL_CLOSED, + requestPreflight, + reservationAccounting, + contextAccessor, + idempotencyKeyResolver, + pricingMissingListeners + ); + } + + @Override + public ChatClientResponse adviseCall( + ChatClientRequest request, + CallAdvisorChain chain + ) { + if (requestPreflight == null) { + return adviseLegacyCall(request, chain); + } + return adviseAccountingCall(request, chain); + } + + @Override + public Flux adviseStream( + ChatClientRequest request, + StreamAdvisorChain chain + ) { + streamingRequestPolicy.requireSupported(requestPreflight != null); + return LedgerAdvisor.super.adviseStream(request, chain); + } + + private ChatClientResponse adviseLegacyCall( + ChatClientRequest request, + CallAdvisorChain chain + ) { + ChatClientRequest resolvedRequest = before(request, chain); + ChatClientResponse response = chain.nextCall(resolvedRequest); + return after(response, chain); + } + + private ChatClientResponse adviseAccountingCall( + ChatClientRequest request, + CallAdvisorChain chain + ) { + IdempotencyKey idempotencyKey = idempotencyKeyResolver.resolve(request); + ChatClientRequest correlatedRequest = contextAccessor.withIdempotencyKey( + request, + idempotencyKey + ); + String requestId = requireCorrelation( + contextAccessor.requestId(correlatedRequest), + "request ID" + ); + String attemptId = requireCorrelation( + contextAccessor.attemptId(correlatedRequest), + "attempt ID" + ); + PreflightCostResult.Bounded costBound; + try { + costBound = requestPreflight.resolve(correlatedRequest); + } catch (MissingPricingException exception) { + publishPricingMissingBestEffort(exception.getResolution()); + throw exception; + } + BudgetDecision decision = budgetEvaluator.evaluate( + extractTags(correlatedRequest.context()), + costBound.safeUpperBoundCost() + ); + validateAdmissionDecision(decision); + + BudgetReservationResult reservationResult = checkAndReserve( + requestId, + idempotencyKey, + costBound, + decision + ); + enforceEvaluatorBlock(decision, reservationResult); + requireAcceptedReservation(reservationResult); + ReservationId reservationId = reservationResult.reservation().id(); + ChatClientRequest providerRequest = prepareDispatch( + correlatedRequest, + reservationResult + ); + + ChatClientResponse response; + try { + response = chain.nextCall(providerRequest); + } catch (RuntimeException downstreamFailure) { + preservePendingLiability(reservationId, downstreamFailure); + throw downstreamFailure; + } + settle( + response, + requestId, + attemptId, + reservationId, + costBound.pricingSnapshot() + ); + return response; + } + + private void validateAdmissionDecision(BudgetDecision decision) { + if (!decision.isAdmissionDecision()) { + throw new IllegalStateException("budget admission decision is required"); + } + switch (decision.state()) { + case ALLOW, WARN, BLOCK -> { + return; + } + case CURRENCY_MISMATCH -> throw new IllegalStateException( + "Budget decision currency mismatch: " + decision.reason() + ); + } + } + + private void enforceEvaluatorBlock( + BudgetDecision decision, + BudgetReservationResult reservationResult + ) { + if (decision.state() != BudgetState.BLOCK) { + return; + } + BudgetExceededException failure = new BudgetExceededException(decision); + releaseCreatedBeforeDispatch(reservationResult, failure); + throw failure; + } + + private void requireAcceptedReservation(BudgetReservationResult result) { + if (!result.isAccepted()) { + throw new IllegalStateException( + "budget reservation rejected: " + result.reason() + ); + } + } + + private BudgetReservationResult checkAndReserve( + String requestId, + IdempotencyKey idempotencyKey, + PreflightCostResult.Bounded costBound, + BudgetDecision decision + ) { + BudgetReservationRequest reservation = new BudgetReservationRequest( + decision.key(), + decision.limit(), + costBound.safeUpperBoundCost(), + requestId, + idempotencyKey, + costBound.pricingSnapshot(), + new ReservationTokenEstimate( + costBound.inputEstimatedTokens(), + costBound.inputSafeUpperBoundTokens(), + costBound.reservedOutputTokens() + ) + ); + return budgetStateStore.checkAndReserve(reservation); + } + + private ChatClientRequest prepareDispatch( + ChatClientRequest request, + BudgetReservationResult reservationResult + ) { + ReservationId reservationId = reservationResult.reservation().id(); + ChatClientRequest providerRequest; + try { + providerRequest = contextAccessor.withReservationId( + request, + reservationId + ); + } catch (RuntimeException failure) { + releaseCreatedBeforeDispatch(reservationResult, failure); + throw failure; + } + + ReservationTransition transition; + try { + transition = reservationAccounting.markInFlight(reservationId); + } catch (RuntimeException failure) { + releaseCreatedBeforeDispatch(reservationResult, failure); + throw failure; + } + boolean claimed = transition.status().isApplied() + && transition.previousState() == ReservationState.RESERVED + && transition.resultingState() == ReservationState.IN_FLIGHT; + if (!claimed) { + IllegalStateException failure = new IllegalStateException( + "reservation dispatch was not claimed: " + transition.status() + ); + if (transition.resultingState() == ReservationState.RESERVED) { + releaseCreatedBeforeDispatch(reservationResult, failure); + } + throw failure; + } + return providerRequest; + } + + private void releaseCreatedBeforeDispatch( + BudgetReservationResult reservationResult, + RuntimeException originalFailure + ) { + if (reservationResult.status() != io.tokenpilot.budget.ReservationStatus.CREATED) { + return; + } + releaseBeforeDispatch(reservationResult.reservation().id(), originalFailure); + } + + private void releaseBeforeDispatch( + ReservationId reservationId, + RuntimeException originalFailure + ) { + try { + reservationAccounting.releaseBeforeDispatch(reservationId); + } catch (RuntimeException releaseFailure) { + originalFailure.addSuppressed(releaseFailure); + } + } + + private void preservePendingLiability( + ReservationId reservationId, + RuntimeException originalFailure + ) { + try { + markReconciliationRequired(reservationId); + } catch (RuntimeException reconciliationFailure) { + originalFailure.addSuppressed(reconciliationFailure); + } + } + + private void settle( + ChatClientResponse response, + String requestId, + String attemptId, + ReservationId reservationId, + PricingSnapshot pricingSnapshot + ) { + TokenUsage usage; + try { + usage = usageExtractor.extract(response); + } catch (RuntimeException ignored) { + markReconciliationRequiredAfterResponse(reservationId); + return; + } + if (usage.source() == UsageSource.UNAVAILABLE) { + markReconciliationRequiredAfterResponse(reservationId); + return; + } + try { + String responseModelId = requireResponseModelId(response); + ActualUsageCommand actualCommand = new ActualUsageCommand( + requestId, + attemptId, + reservationId, + usage, + responseModelId + ); + PricingReconciliationResult reconciliation = + pricingEvaluator.determineReconciliation( + Optional.of(pricingSnapshot), + responseModelId + ); + if (reconciliation != PricingReconciliationResult.RECONCILED) { + if (reconciliation + == PricingReconciliationResult.RECONCILIATION_REQUIRED) { + markPricingReconciliationRequired( + actualCommand, + ReservationAccountingReason.PRICING_RECONCILIATION_REQUIRED + ); + } else { + markReconciliationRequiredAfterResponse( + reservationId, + ReservationAccountingReason.PRICING_RECONCILIATION_REQUIRED + ); + } + return; + } + reservationAccounting.commit(actualCommand); + } catch (RuntimeException ignored) { + markReconciliationRequiredAfterResponse(reservationId); + } + } + + private void markPricingReconciliationRequired( + ActualUsageCommand command, + ReservationAccountingReason reason + ) { + try { + reservationAccounting.markReconciliationRequired( + command, + reason + ); + } catch (RuntimeException ignored) { + // Provider response is preserved; accounting listeners observe failures separately. + } + } + + private void markReconciliationRequired(ReservationId reservationId) { + reservationAccounting.markReconciliationRequired( + reservationId, + ReservationAccountingReason.ACTUAL_USAGE_UNAVAILABLE + ); + } + + private void markReconciliationRequiredAfterResponse( + ReservationId reservationId + ) { + markReconciliationRequiredAfterResponse( + reservationId, + ReservationAccountingReason.ACTUAL_USAGE_UNAVAILABLE + ); + } + + private void markReconciliationRequiredAfterResponse( + ReservationId reservationId, + ReservationAccountingReason reason + ) { + try { + reservationAccounting.markReconciliationRequired( + reservationId, + reason + ); + } catch (RuntimeException ignored) { + // Provider response is preserved; accounting listeners observe failures separately. + } + } + + private String requireResponseModelId(ChatClientResponse response) { + String modelId = extractMetadataModelId(response); + if (modelId == null) { + throw new IllegalStateException("response model ID is unavailable"); + } + return modelId; + } + + private String requireCorrelation(String value, String name) { + if (value == null) { + throw new IllegalStateException(name + " is required"); + } + return value; } @Override @@ -151,6 +645,7 @@ public ChatClientResponse after(ChatClientResponse response, AdvisorChain chain) try { cost = ledgerManager.record(resolvedSnapshot, usage, tags); } catch (MissingPricingException exception) { + publishPricingMissingBestEffort(exception.getResolution()); return handleActualPricingFailure(response, exception); } ChatClientResponse reconciledResponse = withReconciliationResult( @@ -167,7 +662,12 @@ public ChatClientResponse after(ChatClientResponse response, AdvisorChain chain) } return reconciledResponse; } else { - ledgerManager.record(modelId, usage, tags); + try { + ledgerManager.record(modelId, usage, tags); + } catch (MissingPricingException exception) { + publishPricingMissingBestEffort(exception.getResolution()); + throw exception; + } recordLegacyBudgetCost(modelId, usage, response); } @@ -223,6 +723,9 @@ private ChatClientRequest resolvePricing(ChatClientRequest request) { ? Optional.empty() : pricingRegistry.resolveSnapshot(modelId, pricingPolicyId); PricingResolution resolution = pricingEvaluator.validateSnapshotRates(snapshot); + if (!resolution.isResolved()) { + publishPricingMissingBestEffort(resolution); + } rejectMissingPricingIfFailClosed(resolution); return withPricingContext(request, pricingPolicyId, resolution, snapshot); @@ -253,6 +756,23 @@ private void rejectMissingPricingIfFailClosed(PricingResolution resolution) { throw new MissingPricingException(resolution); } + private void publishPricingMissingBestEffort(PricingResolution resolution) { + if (!resolution.isMissing() || pricingMissingListeners.isEmpty()) { + return; + } + PricingMissingEvent event = new PricingMissingEvent( + missingPricingPolicy, + resolution + ); + for (PricingMissingListener listener : pricingMissingListeners) { + try { + listener.onPricingMissing(event); + } catch (RuntimeException ignored) { + // Optional observers do not change pricing or provider-call semantics. + } + } + } + private String extractModelId(ChatClientRequest request) { Object contextValue = request.context().get(MODEL_ID_CONTEXT); if (contextValue instanceof String modelId && !modelId.isBlank()) { @@ -275,7 +795,7 @@ private String extractModelId(ChatClientRequest request) { private void enforceExistingBlock(BudgetDecision decision) { switch (decision.state()) { case ALLOW, WARN -> { - // Legacy boundary: candidate-aware admission is implemented by issue #39. + // budget-enabled adviseCall performs the atomic reservation immediately after this check. } case BLOCK -> throw new BudgetExceededException(decision); case CURRENCY_MISMATCH -> throw new IllegalStateException( diff --git a/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/DefaultRequestAdapter.java b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/DefaultRequestAdapter.java new file mode 100644 index 0000000..60fe123 --- /dev/null +++ b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/DefaultRequestAdapter.java @@ -0,0 +1,39 @@ +package io.tokenpilot.springai.internal; + +import org.springframework.ai.chat.client.ChatClientRequest; +import org.springframework.ai.chat.messages.Message; +import org.springframework.ai.chat.messages.MessageType; + +import java.util.ArrayList; +import java.util.List; +import java.util.Objects; + +final class DefaultRequestAdapter { + + AdaptedRequest adapt(ChatClientRequest request) { + Objects.requireNonNull(request, "request must not be null"); + + List messages = new ArrayList<>(); + for (Message message : request.prompt().getInstructions()) { + messages.add(adaptMessage(message)); + } + + return new AdaptedRequest(messages); + } + + private AdaptedMessage adaptMessage(Message message) { + return new AdaptedMessage( + adaptRole(message.getMessageType()), + message.getText() + ); + } + + private RequestMessageRole adaptRole(MessageType messageType) { + return switch (messageType) { + case SYSTEM -> RequestMessageRole.SYSTEM; + case USER -> RequestMessageRole.USER; + case ASSISTANT -> RequestMessageRole.ASSISTANT; + case TOOL -> throw new IllegalArgumentException("TOOL message is unsupported"); + }; + } +} diff --git a/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/IdempotencyKeyResolver.java b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/IdempotencyKeyResolver.java new file mode 100644 index 0000000..06a1951 --- /dev/null +++ b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/IdempotencyKeyResolver.java @@ -0,0 +1,42 @@ +package io.tokenpilot.springai.internal; + +import io.tokenpilot.budget.IdempotencyKey; +import org.springframework.ai.chat.client.ChatClientRequest; + +import java.util.Objects; +import java.util.function.Supplier; + +/** 현재 Advisor 호출에서 사용할 idempotency key를 결정합니다. */ +final class IdempotencyKeyResolver { + + private final RequestContextAccessor contextAccessor; + private final Supplier keyGenerator; + + IdempotencyKeyResolver( + RequestContextAccessor contextAccessor, + Supplier keyGenerator + ) { + this.contextAccessor = Objects.requireNonNull( + contextAccessor, + "contextAccessor must not be null" + ); + this.keyGenerator = Objects.requireNonNull( + keyGenerator, + "keyGenerator must not be null" + ); + } + + IdempotencyKey resolve(ChatClientRequest request) { + Objects.requireNonNull(request, "request must not be null"); + + IdempotencyKey existingKey = contextAccessor.idempotencyKey(request); + if (existingKey != null) { + return existingKey; + } + + return Objects.requireNonNull( + keyGenerator.get(), + "generated idempotencyKey must not be null" + ); + } +} diff --git a/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/LedgerSpringAiComponents.java b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/LedgerSpringAiComponents.java index 2ee86cd..591907c 100644 --- a/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/LedgerSpringAiComponents.java +++ b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/LedgerSpringAiComponents.java @@ -2,14 +2,28 @@ import io.tokenpilot.budget.BudgetEvaluator; import io.tokenpilot.budget.BudgetStateStore; +import io.tokenpilot.budget.IdempotencyKey; +import io.tokenpilot.budget.ReservationAccounting; import io.tokenpilot.core.CostCalculator; import io.tokenpilot.core.LedgerManager; +import io.tokenpilot.core.ModelRegistry; +import io.tokenpilot.core.PreflightCostEstimator; import io.tokenpilot.core.PricingEvaluator; +import io.tokenpilot.core.PricingMissingListener; import io.tokenpilot.core.PricingRegistry; +import io.tokenpilot.core.TokenBudget; +import io.tokenpilot.core.TokenEstimator; import io.tokenpilot.core.domain.MissingPricingPolicy; import io.tokenpilot.core.internal.LedgerComponents; import io.tokenpilot.springai.LedgerAdvisor; +import io.tokenpilot.springai.ReservedOutputTokensResolver; import io.tokenpilot.springai.UsageExtractor; +import org.jspecify.annotations.Nullable; + +import java.util.UUID; + +import java.util.List; +import java.util.Objects; /** * Spring AI 어댑터 컴포넌트 생성을 위한 팩토리 클래스입니다. @@ -23,6 +37,130 @@ public static UsageExtractor defaultUsageExtractor() { return new DefaultUsageExtractor(); } + public static LedgerAdvisor accountingLedgerAdvisor( + UsageExtractor usageExtractor, + BudgetEvaluator budgetEvaluator, + BudgetStateStore budgetStateStore, + ReservationAccounting reservationAccounting, + PricingRegistry pricingRegistry, + ModelRegistry modelRegistry, + TokenEstimator tokenEstimator, + TokenBudget tokenBudget, + PreflightCostEstimator costEstimator, + @Nullable String defaultModelId, + @Nullable Long defaultReservedOutputTokens, + long framingHeadroomTokens + ) { + return accountingLedgerAdvisor( + usageExtractor, + budgetEvaluator, + budgetStateStore, + reservationAccounting, + pricingRegistry, + modelRegistry, + tokenEstimator, + tokenBudget, + costEstimator, + defaultModelId, + defaultReservedOutputTokens, + null, + framingHeadroomTokens + ); + } + + public static LedgerAdvisor accountingLedgerAdvisor( + UsageExtractor usageExtractor, + BudgetEvaluator budgetEvaluator, + BudgetStateStore budgetStateStore, + ReservationAccounting reservationAccounting, + PricingRegistry pricingRegistry, + ModelRegistry modelRegistry, + TokenEstimator tokenEstimator, + TokenBudget tokenBudget, + PreflightCostEstimator costEstimator, + @Nullable String defaultModelId, + @Nullable Long defaultReservedOutputTokens, + @Nullable ReservedOutputTokensResolver providerOutputResolver, + long framingHeadroomTokens + ) { + return accountingLedgerAdvisor( + usageExtractor, + budgetEvaluator, + budgetStateStore, + reservationAccounting, + pricingRegistry, + modelRegistry, + tokenEstimator, + tokenBudget, + costEstimator, + LedgerComponents.defaultPricingEvaluator(), + List.of(), + defaultModelId, + defaultReservedOutputTokens, + providerOutputResolver, + framingHeadroomTokens + ); + } + + public static LedgerAdvisor accountingLedgerAdvisor( + UsageExtractor usageExtractor, + BudgetEvaluator budgetEvaluator, + BudgetStateStore budgetStateStore, + ReservationAccounting reservationAccounting, + PricingRegistry pricingRegistry, + ModelRegistry modelRegistry, + TokenEstimator tokenEstimator, + TokenBudget tokenBudget, + PreflightCostEstimator costEstimator, + PricingEvaluator pricingEvaluator, + List pricingMissingListeners, + @Nullable String defaultModelId, + @Nullable Long defaultReservedOutputTokens, + @Nullable ReservedOutputTokensResolver providerOutputResolver, + long framingHeadroomTokens + ) { + if (Objects.requireNonNull( + budgetStateStore, + "budgetStateStore must not be null" + ) != Objects.requireNonNull( + reservationAccounting, + "reservationAccounting must not be null" + )) { + throw new IllegalArgumentException( + "budgetStateStore and reservationAccounting must be the same instance" + ); + } + ReservedOutputResolver outputResolver = new ReservedOutputResolver( + providerOutputResolver, + defaultReservedOutputTokens + ); + RequestContextAccessor contextAccessor = new RequestContextAccessor(); + RequestPreflight preflight = new RequestPreflight( + new ModelResolver(modelRegistry, defaultModelId), + outputResolver, + tokenEstimator, + tokenBudget, + pricingRegistry, + pricingEvaluator, + costEstimator, + framingHeadroomTokens + ); + return new DefaultLedgerAdvisor( + usageExtractor, + budgetEvaluator, + budgetStateStore, + reservationAccounting, + preflight, + pricingEvaluator, + contextAccessor, + new IdempotencyKeyResolver( + contextAccessor, + () -> new IdempotencyKey(UUID.randomUUID().toString()) + ), + pricingMissingListeners + ); + } + public static LedgerAdvisor defaultLedgerAdvisor( LedgerManager ledgerManager, UsageExtractor usageExtractor @@ -115,6 +253,30 @@ public static LedgerAdvisor defaultLedgerAdvisor( PricingRegistry pricingRegistry, PricingEvaluator pricingEvaluator, MissingPricingPolicy missingPricingPolicy + ) { + return defaultLedgerAdvisor( + ledgerManager, + usageExtractor, + budgetEvaluator, + budgetStateStore, + costCalculator, + pricingRegistry, + pricingEvaluator, + missingPricingPolicy, + List.of() + ); + } + + public static LedgerAdvisor defaultLedgerAdvisor( + LedgerManager ledgerManager, + UsageExtractor usageExtractor, + BudgetEvaluator budgetEvaluator, + BudgetStateStore budgetStateStore, + CostCalculator costCalculator, + PricingRegistry pricingRegistry, + PricingEvaluator pricingEvaluator, + MissingPricingPolicy missingPricingPolicy, + List pricingMissingListeners ) { return new DefaultLedgerAdvisor( ledgerManager, @@ -124,7 +286,8 @@ public static LedgerAdvisor defaultLedgerAdvisor( costCalculator, pricingRegistry, pricingEvaluator, - missingPricingPolicy + missingPricingPolicy, + pricingMissingListeners ); } } diff --git a/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/ModelResolver.java b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/ModelResolver.java new file mode 100644 index 0000000..03ec39a --- /dev/null +++ b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/ModelResolver.java @@ -0,0 +1,55 @@ +package io.tokenpilot.springai.internal; + +import io.tokenpilot.core.ModelRegistry; +import io.tokenpilot.core.domain.ModelDefinition; +import org.jspecify.annotations.Nullable; +import org.springframework.ai.chat.client.ChatClientRequest; +import org.springframework.ai.chat.prompt.ChatOptions; + +import java.util.Objects; +import java.util.Optional; + +/** request model을 configured default 순서로 선택하고 Core registry에서 해석합니다. */ +final class ModelResolver { + + private final ModelRegistry modelRegistry; + private final Optional defaultModelId; + + ModelResolver(ModelRegistry modelRegistry) { + this(modelRegistry, null); + } + + ModelResolver(ModelRegistry modelRegistry, @Nullable String defaultModelId) { + this.modelRegistry = Objects.requireNonNull( + modelRegistry, + "modelRegistry must not be null" + ); + this.defaultModelId = Optional.ofNullable(defaultModelId) + .filter(modelId -> !modelId.isBlank()); + } + + Optional resolve(ChatClientRequest request) { + Objects.requireNonNull(request, "request must not be null"); + + String requestModelId = requestModelId(request); + if (requestModelId != null) { + return resolveModel(requestModelId); + } + return defaultModelId.flatMap(this::resolveModel); + } + + private Optional resolveModel(String modelId) { + if (modelId.isBlank()) { + return Optional.empty(); + } + return modelRegistry.find(modelId); + } + + private @Nullable String requestModelId(ChatClientRequest request) { + ChatOptions options = request.prompt().getOptions(); + if (options == null) { + return null; + } + return options.getModel(); + } +} diff --git a/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/RequestContextAccessor.java b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/RequestContextAccessor.java new file mode 100644 index 0000000..5e4145f --- /dev/null +++ b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/RequestContextAccessor.java @@ -0,0 +1,95 @@ +package io.tokenpilot.springai.internal; + +import io.tokenpilot.budget.IdempotencyKey; +import io.tokenpilot.budget.ReservationId; +import org.jspecify.annotations.Nullable; +import org.springframework.ai.chat.client.ChatClientRequest; + +import java.util.Objects; + +/** Spring AI request context에서 Token Pilot correlation 값을 읽고 추가합니다. */ +final class RequestContextAccessor { + + static final String REQUEST_ID_CONTEXT_KEY = "tokenpilot.request.id"; + static final String ATTEMPT_ID_CONTEXT_KEY = "tokenpilot.attempt.id"; + static final String IDEMPOTENCY_CONTEXT_KEY = "tokenpilot.idempotency.key"; + static final String RESERVATION_ID_CONTEXT_KEY = "tokenpilot.reservation.id"; + + @Nullable String requestId(ChatClientRequest request) { + return readContextValue(request, REQUEST_ID_CONTEXT_KEY, String.class); + } + + @Nullable String attemptId(ChatClientRequest request) { + return readContextValue(request, ATTEMPT_ID_CONTEXT_KEY, String.class); + } + + @Nullable IdempotencyKey idempotencyKey(ChatClientRequest request) { + return readContextValue(request, IDEMPOTENCY_CONTEXT_KEY, IdempotencyKey.class); + } + + @Nullable ReservationId reservationId(ChatClientRequest request) { + return readContextValue(request, RESERVATION_ID_CONTEXT_KEY, ReservationId.class); + } + + ChatClientRequest withIdempotencyKey( + ChatClientRequest request, + IdempotencyKey idempotencyKey + ) { + Objects.requireNonNull(request, "request must not be null"); + Objects.requireNonNull(idempotencyKey, "idempotencyKey must not be null"); + + IdempotencyKey existingIdempotencyKey = idempotencyKey(request); + if (existingIdempotencyKey == null) { + return request.mutate() + .context(IDEMPOTENCY_CONTEXT_KEY, idempotencyKey) + .build(); + } + if (!existingIdempotencyKey.equals(idempotencyKey)) { + throw new IllegalStateException( + IDEMPOTENCY_CONTEXT_KEY + " cannot be changed" + ); + } + return request; + } + + ChatClientRequest withReservationId( + ChatClientRequest request, + ReservationId reservationId + ) { + Objects.requireNonNull(request, "request must not be null"); + Objects.requireNonNull(reservationId, "reservationId must not be null"); + + ReservationId existingReservationId = reservationId(request); + if (existingReservationId == null) { + return request.mutate() + .context(RESERVATION_ID_CONTEXT_KEY, reservationId) + .build(); + } + if (!existingReservationId.equals(reservationId)) { + throw new IllegalStateException( + RESERVATION_ID_CONTEXT_KEY + " cannot be changed" + ); + } + return request; + } + + private @Nullable T readContextValue( + ChatClientRequest request, + String key, + Class valueType + ) { + Objects.requireNonNull(request, "request must not be null"); + + if (!request.context().containsKey(key)) { + return null; + } + + Object value = request.context().get(key); + if (!valueType.isInstance(value)) { + throw new IllegalArgumentException( + key + " must contain " + valueType.getSimpleName() + ); + } + return valueType.cast(value); + } +} diff --git a/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/RequestFramingPolicy.java b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/RequestFramingPolicy.java new file mode 100644 index 0000000..242e452 --- /dev/null +++ b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/RequestFramingPolicy.java @@ -0,0 +1,21 @@ +package io.tokenpilot.springai.internal; + +import java.util.Objects; + +final class RequestFramingPolicy { + + String frame(AdaptedRequest request) { + Objects.requireNonNull(request, "request must not be null"); + + StringBuilder framed = new StringBuilder(); + for (AdaptedMessage message : request.messages()) { + framed.append(message.role()) + .append(':') + .append(message.text().length()) + .append(':') + .append(message.text()) + .append('\n'); + } + return framed.toString(); + } +} diff --git a/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/RequestMessageRole.java b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/RequestMessageRole.java new file mode 100644 index 0000000..8c295fb --- /dev/null +++ b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/RequestMessageRole.java @@ -0,0 +1,8 @@ +package io.tokenpilot.springai.internal; + +/** Token Pilot이 지원하는 text message 타입입니다. */ +enum RequestMessageRole { + SYSTEM, + USER, + ASSISTANT +} diff --git a/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/RequestPreflight.java b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/RequestPreflight.java new file mode 100644 index 0000000..0959291 --- /dev/null +++ b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/RequestPreflight.java @@ -0,0 +1,162 @@ +package io.tokenpilot.springai.internal; + +import io.tokenpilot.core.PreflightCostEstimator; +import io.tokenpilot.core.PricingEvaluator; +import io.tokenpilot.core.PricingRegistry; +import io.tokenpilot.core.TokenBudget; +import io.tokenpilot.core.TokenEstimator; +import io.tokenpilot.core.domain.ModelDefinition; +import io.tokenpilot.core.domain.PreflightCostResult; +import io.tokenpilot.core.domain.PreflightPricingContext; +import io.tokenpilot.core.domain.PricingResolution; +import io.tokenpilot.core.domain.TokenCountAccuracy; +import io.tokenpilot.core.domain.TokenCountResult; +import io.tokenpilot.core.domain.TokenCountScope; +import io.tokenpilot.core.exception.MissingPricingException; +import io.tokenpilot.core.internal.LedgerComponents; +import org.springframework.ai.chat.client.ChatClientRequest; + +import java.util.Objects; + +/** Spring AI 요청 해석 결과를 기존 Core preflight 계약에 연결합니다. */ +final class RequestPreflight { + + private final DefaultRequestAdapter requestAdapter = new DefaultRequestAdapter(); + private final RequestScopeResolver scopeResolver = new RequestScopeResolver(); + private final RequestFramingPolicy framingPolicy = new RequestFramingPolicy(); + private final ModelResolver modelResolver; + private final ReservedOutputResolver outputResolver; + private final TokenEstimator tokenEstimator; + private final TokenBudget tokenBudget; + private final PricingRegistry pricingRegistry; + private final PricingEvaluator pricingEvaluator; + private final PreflightCostEstimator costEstimator; + private final long framingHeadroomTokens; + + RequestPreflight( + ModelResolver modelResolver, + ReservedOutputResolver outputResolver, + TokenEstimator tokenEstimator, + TokenBudget tokenBudget, + PricingRegistry pricingRegistry, + PreflightCostEstimator costEstimator, + long framingHeadroomTokens + ) { + this( + modelResolver, + outputResolver, + tokenEstimator, + tokenBudget, + pricingRegistry, + LedgerComponents.defaultPricingEvaluator(), + costEstimator, + framingHeadroomTokens + ); + } + + RequestPreflight( + ModelResolver modelResolver, + ReservedOutputResolver outputResolver, + TokenEstimator tokenEstimator, + TokenBudget tokenBudget, + PricingRegistry pricingRegistry, + PricingEvaluator pricingEvaluator, + PreflightCostEstimator costEstimator, + long framingHeadroomTokens + ) { + this.modelResolver = Objects.requireNonNull(modelResolver); + this.outputResolver = Objects.requireNonNull(outputResolver); + this.tokenEstimator = Objects.requireNonNull(tokenEstimator); + this.tokenBudget = Objects.requireNonNull(tokenBudget); + this.pricingRegistry = Objects.requireNonNull(pricingRegistry); + this.pricingEvaluator = Objects.requireNonNull(pricingEvaluator); + this.costEstimator = Objects.requireNonNull(costEstimator); + if (framingHeadroomTokens < 0) { + throw new IllegalArgumentException("framingHeadroomTokens must be non-negative"); + } + this.framingHeadroomTokens = framingHeadroomTokens; + } + + PreflightCostResult.Bounded resolve(ChatClientRequest request) { + requireTextOnly(request); + ModelDefinition model = modelResolver.resolve(request) + .orElseThrow(() -> new IllegalStateException("MODEL_UNRESOLVED")); + long reservedOutputTokens = outputResolver.resolve(request) + .orElseThrow(() -> new IllegalStateException( + "OUTPUT_RESERVATION_UNRESOLVED" + )); + TokenCountResult requestTokens = countRequest(request); + tokenBudget.requireFits( + model.canonicalModelId(), + requestTokens, + reservedOutputTokens + ); + + var pricingSnapshot = pricingRegistry.resolveSnapshot(model); + PricingResolution pricingResolution = pricingEvaluator.validateSnapshotRates( + pricingSnapshot + ); + if (!pricingResolution.isResolved()) { + throw new MissingPricingException(pricingResolution); + } + PreflightPricingContext pricing = new PreflightPricingContext( + model.canonicalModelId(), + model.pricingPolicyId(), + model.catalogVersion(), + model.acceptedCompatibilityBasis(), + model.pricingCurrency(), + PreflightPricingContext.UpperBoundCapability.FINITE, + pricingSnapshot + ); + PreflightCostResult result = costEstimator.estimate( + pricing, + requestTokens, + reservedOutputTokens + ); + if (result instanceof PreflightCostResult.Bounded bounded) { + return bounded; + } + PreflightCostResult.Unavailable unavailable = + (PreflightCostResult.Unavailable) result; + throw new IllegalStateException( + "preflight cost bound unavailable: " + unavailable.reason() + ); + } + + private void requireTextOnly(ChatClientRequest request) { + RequestScopeResult result = scopeResolver.resolve(request); + if (result.scope().isPresent()) { + return; + } + throw new IllegalStateException( + "UNSUPPORTED_REQUEST_SCOPE: " + + result.unsupportedReason().orElseThrow() + ); + } + + private TokenCountResult countRequest(ChatClientRequest request) { + TokenCountResult text = tokenEstimator.estimate( + framingPolicy.frame(requestAdapter.adapt(request)) + ); + if (text.isUnavailable()) { + return TokenCountResult.unavailable( + text.unavailableReason().orElseThrow(), + TokenCountScope.REQUEST, + text.estimatorDescriptor(), + text.tokenizationBasis() + ); + } + return TokenCountResult.counted( + text.tokens().orElseThrow(), + Math.addExact( + text.safeUpperBoundTokens().orElseThrow(), + framingHeadroomTokens + ), + TokenCountAccuracy.HEURISTIC, + TokenCountScope.REQUEST, + text.estimatorDescriptor(), + text.tokenizationBasis() + ); + } + +} diff --git a/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/RequestScopeResolver.java b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/RequestScopeResolver.java new file mode 100644 index 0000000..ae283f9 --- /dev/null +++ b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/RequestScopeResolver.java @@ -0,0 +1,101 @@ +package io.tokenpilot.springai.internal; + +import org.springframework.ai.chat.client.ChatClientAttributes; +import org.springframework.ai.chat.client.ChatClientRequest; +import org.springframework.ai.chat.messages.AssistantMessage; +import org.springframework.ai.chat.messages.Message; +import org.springframework.ai.chat.messages.MessageType; +import org.springframework.ai.chat.prompt.ChatOptions; +import org.springframework.ai.content.MediaContent; +import org.springframework.ai.model.tool.StructuredOutputChatOptions; +import org.springframework.ai.model.tool.ToolCallingChatOptions; +import org.springframework.util.CollectionUtils; +import org.springframework.util.StringUtils; + +import java.util.Objects; + +/** Spring AI 요청이 Token Pilot 0.1.0에서 지원하는 범위인지 판정합니다. */ +final class RequestScopeResolver { + + RequestScopeResult resolve(ChatClientRequest request) { + Objects.requireNonNull(request, "request must not be null"); + + for (Message message : request.prompt().getInstructions()) { + if (message.getMessageType() == MessageType.TOOL) { + return RequestScopeResult.unsupported( + UnsupportedRequestScopeReason.TOOL_MESSAGE + ); + } + if (hasToolCalls(message)) { + return RequestScopeResult.unsupported( + UnsupportedRequestScopeReason.TOOL_CALL + ); + } + if (hasMedia(message)) { + return RequestScopeResult.unsupported( + UnsupportedRequestScopeReason.MEDIA + ); + } + if (message.getText() == null) { + return RequestScopeResult.unsupported( + UnsupportedRequestScopeReason.MESSAGE_TEXT_UNAVAILABLE + ); + } + } + if (hasToolSchema(request)) { + return RequestScopeResult.unsupported( + UnsupportedRequestScopeReason.TOOL_SCHEMA + ); + } + if (usesStructuredOutput(request)) { + return RequestScopeResult.unsupported( + UnsupportedRequestScopeReason.STRUCTURED_OUTPUT + ); + } + return RequestScopeResult.textOnly(); + } + + private boolean hasMedia(Message message) { + return message instanceof MediaContent mediaContent + && !mediaContent.getMedia().isEmpty(); + } + + private boolean hasToolCalls(Message message) { + return message instanceof AssistantMessage assistantMessage + && assistantMessage.hasToolCalls(); + } + + private boolean usesStructuredOutput(ChatClientRequest request) { + return hasTextAttribute(request, ChatClientAttributes.OUTPUT_FORMAT) + || hasTextAttribute(request, ChatClientAttributes.STRUCTURED_OUTPUT_SCHEMA) + || usesNativeStructuredOutput(request) + || hasStructuredOutputOption(request); + } + + private boolean hasToolSchema(ChatClientRequest request) { + ChatOptions options = request.prompt().getOptions(); + return options instanceof ToolCallingChatOptions toolOptions + && !CollectionUtils.isEmpty(toolOptions.getToolCallbacks()); + } + + private boolean hasStructuredOutputOption(ChatClientRequest request) { + ChatOptions options = request.prompt().getOptions(); + return options instanceof StructuredOutputChatOptions structuredOptions + && StringUtils.hasText(structuredOptions.getOutputSchema()); + } + + private boolean usesNativeStructuredOutput(ChatClientRequest request) { + Object value = request.context().get( + ChatClientAttributes.STRUCTURED_OUTPUT_NATIVE.getKey() + ); + return Boolean.TRUE.equals(value); + } + + private boolean hasTextAttribute( + ChatClientRequest request, + ChatClientAttributes attribute + ) { + Object value = request.context().get(attribute.getKey()); + return value instanceof String text && StringUtils.hasText(text); + } +} diff --git a/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/RequestScopeResult.java b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/RequestScopeResult.java new file mode 100644 index 0000000..5aa3df2 --- /dev/null +++ b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/RequestScopeResult.java @@ -0,0 +1,43 @@ +package io.tokenpilot.springai.internal; + +import io.tokenpilot.core.domain.TokenCountScope; + +import java.util.Objects; +import java.util.Optional; + +/** 지원하는 request scope 또는 제한된 미지원 사유를 나타냅니다. */ +final class RequestScopeResult { + + private final Optional scope; + private final Optional unsupportedReason; + + private RequestScopeResult( + Optional scope, + Optional unsupportedReason + ) { + this.scope = scope; + this.unsupportedReason = unsupportedReason; + } + + static RequestScopeResult textOnly() { + return new RequestScopeResult( + Optional.of(TokenCountScope.TEXT_ONLY), + Optional.empty() + ); + } + + static RequestScopeResult unsupported(UnsupportedRequestScopeReason reason) { + return new RequestScopeResult( + Optional.empty(), + Optional.of(Objects.requireNonNull(reason, "reason must not be null")) + ); + } + + Optional scope() { + return scope; + } + + Optional unsupportedReason() { + return unsupportedReason; + } +} diff --git a/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/ReservedOutputResolver.java b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/ReservedOutputResolver.java new file mode 100644 index 0000000..b97ccc4 --- /dev/null +++ b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/ReservedOutputResolver.java @@ -0,0 +1,71 @@ +package io.tokenpilot.springai.internal; + +import io.tokenpilot.springai.ReservedOutputTokensResolver; +import org.jspecify.annotations.Nullable; +import org.springframework.ai.chat.client.ChatClientRequest; +import org.springframework.ai.chat.prompt.ChatOptions; + +import java.util.Objects; +import java.util.OptionalLong; + +/** request maxTokens를 configured default 순서로 해석합니다. */ +final class ReservedOutputResolver { + + private final @Nullable ReservedOutputTokensResolver providerResolver; + private final OptionalLong defaultReservedOutputTokens; + + ReservedOutputResolver() { + this(null, null); + } + + ReservedOutputResolver(long defaultReservedOutputTokens) { + this(null, defaultReservedOutputTokens); + } + + ReservedOutputResolver( + @Nullable ReservedOutputTokensResolver providerResolver, + @Nullable Long defaultReservedOutputTokens + ) { + this.providerResolver = providerResolver; + if (defaultReservedOutputTokens == null) { + this.defaultReservedOutputTokens = OptionalLong.empty(); + return; + } + this.defaultReservedOutputTokens = positiveTokens(defaultReservedOutputTokens); + } + + OptionalLong resolve(ChatClientRequest request) { + Objects.requireNonNull(request, "request must not be null"); + + Integer requestMaxTokens = requestMaxTokens(request); + if (requestMaxTokens != null) { + return positiveTokens(requestMaxTokens); + } + if (providerResolver == null) { + return defaultReservedOutputTokens; + } + OptionalLong providerTokens = Objects.requireNonNull( + providerResolver.resolve(request), + "provider resolver result must not be null" + ); + if (providerTokens.isEmpty()) { + return defaultReservedOutputTokens; + } + return positiveTokens(providerTokens.getAsLong()); + } + + private OptionalLong positiveTokens(long tokens) { + if (tokens <= 0) { + return OptionalLong.empty(); + } + return OptionalLong.of(tokens); + } + + private @Nullable Integer requestMaxTokens(ChatClientRequest request) { + ChatOptions options = request.prompt().getOptions(); + if (options == null) { + return null; + } + return options.getMaxTokens(); + } +} diff --git a/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/StreamingRequestPolicy.java b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/StreamingRequestPolicy.java new file mode 100644 index 0000000..815b430 --- /dev/null +++ b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/StreamingRequestPolicy.java @@ -0,0 +1,12 @@ +package io.tokenpilot.springai.internal; + +final class StreamingRequestPolicy { + + void requireSupported(boolean enforcementEnabled) { + if (enforcementEnabled) { + throw new IllegalStateException( + "STREAMING_UNSUPPORTED_FOR_ENFORCEMENT" + ); + } + } +} diff --git a/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/UnsupportedRequestScopeReason.java b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/UnsupportedRequestScopeReason.java new file mode 100644 index 0000000..8592577 --- /dev/null +++ b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/UnsupportedRequestScopeReason.java @@ -0,0 +1,11 @@ +package io.tokenpilot.springai.internal; + +/** Token Pilot 0.1.0이 지원하지 않는 Spring AI request scope의 제한된 사유입니다. */ +enum UnsupportedRequestScopeReason { + TOOL_MESSAGE, + MEDIA, + TOOL_CALL, + TOOL_SCHEMA, + MESSAGE_TEXT_UNAVAILABLE, + STRUCTURED_OUTPUT +} diff --git a/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/DefaultLedgerAdvisorLifecycleTest.java b/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/DefaultLedgerAdvisorLifecycleTest.java new file mode 100644 index 0000000..0444015 --- /dev/null +++ b/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/DefaultLedgerAdvisorLifecycleTest.java @@ -0,0 +1,456 @@ +package io.tokenpilot.springai.internal; + +import io.tokenpilot.budget.AccountingTransitionStatus; +import io.tokenpilot.budget.ActualUsageCommand; +import io.tokenpilot.budget.BudgetDecision; +import io.tokenpilot.budget.BudgetEvaluator; +import io.tokenpilot.budget.BudgetKey; +import io.tokenpilot.budget.BudgetReservation; +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.IdempotencyKey; +import io.tokenpilot.budget.ReservationAccounting; +import io.tokenpilot.budget.ReservationAccountingReason; +import io.tokenpilot.budget.ReservationId; +import io.tokenpilot.budget.ReservationState; +import io.tokenpilot.budget.ReservationStatus; +import io.tokenpilot.budget.ReservationTransition; +import io.tokenpilot.budget.internal.LedgerBudgetComponents; +import io.tokenpilot.core.domain.Cost; +import io.tokenpilot.core.domain.PreflightCostResult; +import io.tokenpilot.core.domain.PricingSnapshot; +import io.tokenpilot.core.domain.TokenUsage; +import io.tokenpilot.springai.UsageExtractor; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +import org.springframework.ai.chat.client.ChatClientRequest; +import org.springframework.ai.chat.client.ChatClientResponse; +import org.springframework.ai.chat.client.advisor.api.CallAdvisorChain; +import org.springframework.ai.chat.metadata.ChatResponseMetadata; +import org.springframework.ai.chat.model.ChatResponse; +import org.springframework.ai.chat.prompt.Prompt; + +import java.math.BigDecimal; +import java.util.Currency; +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executors; +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.ArgumentMatchers.argThat; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.inOrder; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; +import static org.mockito.Mockito.when; +import static org.mockito.Mockito.withSettings; + +class DefaultLedgerAdvisorLifecycleTest { + + private static final Currency USD = Currency.getInstance("USD"); + private static final ReservationId RESERVATION_ID = + new ReservationId("reservation-1"); + + private RequestPreflight preflight; + private BudgetEvaluator budgetEvaluator; + private Object accountingWriter; + private BudgetStateStore stateStore; + private ReservationAccounting accounting; + private UsageExtractor usageExtractor; + private CallAdvisorChain provider; + private DefaultLedgerAdvisor advisor; + + @BeforeEach + void setUp() { + preflight = mock(RequestPreflight.class); + budgetEvaluator = mock(BudgetEvaluator.class); + accountingWriter = mock( + BudgetStateStore.class, + withSettings().extraInterfaces(ReservationAccounting.class) + ); + stateStore = (BudgetStateStore) accountingWriter; + accounting = (ReservationAccounting) accountingWriter; + usageExtractor = mock(UsageExtractor.class); + provider = mock(CallAdvisorChain.class); + RequestContextAccessor contextAccessor = new RequestContextAccessor(); + advisor = new DefaultLedgerAdvisor( + usageExtractor, + budgetEvaluator, + stateStore, + accounting, + preflight, + contextAccessor, + new IdempotencyKeyResolver( + contextAccessor, + () -> new IdempotencyKey("generated-key") + ) + ); + } + + @Test + @DisplayName("정상 call은 예약과 provider를 거쳐 actual을 한 번 commit한다") + void commitsSuccessfulCallOnce() { + ChatClientRequest request = request(); + ChatClientResponse response = response(); + PreflightCostResult.Bounded costBound = stubDispatch(request, response); + when(usageExtractor.extract(response)).thenReturn(TokenUsage.from(10, 5)); + + ChatClientResponse result = advisor.adviseCall(request, provider); + + assertThat(result).isSameAs(response); + var order = inOrder(preflight, stateStore, accounting, provider, usageExtractor); + order.verify(preflight).resolve(request); + order.verify(stateStore).checkAndReserve( + argThat(reservation -> reservation.pricingSnapshot().orElseThrow() + .equals(costBound.pricingSnapshot())) + ); + order.verify(accounting).markInFlight(RESERVATION_ID); + order.verify(provider).nextCall(any()); + order.verify(usageExtractor).extract(response); + order.verify(accounting).commit( + argThat(command -> command.requestId().equals("request-1") + && command.attemptId().equals("attempt-1") + && command.reservationId().equals(RESERVATION_ID) + && command.responseModelId().equals("model-v1")) + ); + } + + @Test + @DisplayName("preflight 실패는 예약과 provider 호출 전에 종료한다") + void stopsBeforeReservationWhenPreflightFails() { + ChatClientRequest request = request(); + when(preflight.resolve(request)) + .thenThrow(new IllegalStateException("MODEL_UNRESOLVED")); + + assertThatThrownBy(() -> advisor.adviseCall(request, provider)) + .isInstanceOf(IllegalStateException.class) + .hasMessage("MODEL_UNRESOLVED"); + + verifyNoInteractions(accountingWriter, provider, usageExtractor); + } + + @Test + @DisplayName("dispatch 준비 실패는 예약을 release하고 provider를 호출하지 않는다") + void releasesBeforeDispatch() { + ChatClientRequest request = request(); + stubReservation(request); + when(accounting.markInFlight(RESERVATION_ID)) + .thenReturn(ReservationTransition.unchanged( + ReservationState.RESERVED, + AccountingTransitionStatus.NOT_ALLOWED + )); + + assertThatThrownBy(() -> advisor.adviseCall(request, provider)) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("dispatch was not claimed"); + + verify(accounting).releaseBeforeDispatch(RESERVATION_ID); + verifyNoInteractions(provider, usageExtractor); + } + + @Test + @DisplayName("이미 IN_FLIGHT인 재사용 예약은 provider를 다시 호출하지 않는다") + void reusedInFlightReservationDoesNotDispatchAgain() { + ChatClientRequest request = request(); + stubReservation(request); + BudgetReservation reservation = mock(BudgetReservation.class); + BudgetReservationResult reused = mock(BudgetReservationResult.class); + when(reservation.id()).thenReturn(RESERVATION_ID); + when(reused.isAccepted()).thenReturn(true); + when(reused.status()).thenReturn(ReservationStatus.REUSED); + when(reused.reservation()).thenReturn(reservation); + when(stateStore.checkAndReserve(any())).thenReturn(reused); + when(accounting.markInFlight(RESERVATION_ID)) + .thenReturn(ReservationTransition.unchanged( + ReservationState.IN_FLIGHT, + AccountingTransitionStatus.NOT_ALLOWED + )); + + assertThatThrownBy(() -> advisor.adviseCall(request, provider)) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("dispatch was not claimed"); + + verify(accounting, never()).releaseBeforeDispatch(RESERVATION_ID); + verifyNoInteractions(provider, usageExtractor); + } + + @Test + @Timeout(value = 10, unit = TimeUnit.SECONDS) + @DisplayName("같은 idempotency key의 동시 호출은 실제 store에서도 provider를 한 번만 호출한다") + void concurrentReusedReservationDispatchesProviderOnce() throws Exception { + RequestPreflight concurrentPreflight = mock(RequestPreflight.class); + BudgetEvaluator concurrentEvaluator = mock(BudgetEvaluator.class); + BudgetStateStore concurrentStore = LedgerBudgetComponents + .inMemoryBudgetStateStore(); + ReservationAccounting concurrentAccounting = + (ReservationAccounting) concurrentStore; + UsageExtractor concurrentUsageExtractor = mock(UsageExtractor.class); + CallAdvisorChain concurrentProvider = mock(CallAdvisorChain.class); + RequestContextAccessor contextAccessor = new RequestContextAccessor(); + DefaultLedgerAdvisor concurrentAdvisor = new DefaultLedgerAdvisor( + concurrentUsageExtractor, + concurrentEvaluator, + concurrentStore, + concurrentAccounting, + concurrentPreflight, + contextAccessor, + new IdempotencyKeyResolver( + contextAccessor, + () -> new IdempotencyKey("unused-generated-key") + ) + ); + ChatClientRequest request = request(); + ChatClientResponse response = response(); + PreflightCostResult.Bounded bound = costBound(); + BudgetDecision decision = allowed(bound.safeUpperBoundCost()); + when(concurrentPreflight.resolve(request)).thenReturn(bound); + when(concurrentEvaluator.evaluate(any(), any())) + .thenReturn(decision); + when(concurrentUsageExtractor.extract(response)) + .thenReturn(TokenUsage.unavailable(Map.of())); + CountDownLatch providerEntered = new CountDownLatch(1); + CountDownLatch releaseProvider = new CountDownLatch(1); + when(concurrentProvider.nextCall(any())).thenAnswer(ignored -> { + providerEntered.countDown(); + if (!releaseProvider.await(5, TimeUnit.SECONDS)) { + throw new IllegalStateException("provider release timed out"); + } + return response; + }); + + try (var executor = Executors.newVirtualThreadPerTaskExecutor()) { + var first = executor.submit( + () -> concurrentAdvisor.adviseCall(request, concurrentProvider) + ); + assertThat(providerEntered.await(5, TimeUnit.SECONDS)).isTrue(); + var second = executor.submit( + () -> concurrentAdvisor.adviseCall(request, concurrentProvider) + ); + + assertThatThrownBy(() -> second.get(5, TimeUnit.SECONDS)) + .isInstanceOf(ExecutionException.class) + .hasRootCauseMessage( + "reservation dispatch was not claimed: NOT_ALLOWED" + ); + releaseProvider.countDown(); + assertThat(first.get(5, TimeUnit.SECONDS)).isSameAs(response); + } finally { + releaseProvider.countDown(); + } + + verify(concurrentProvider, times(1)).nextCall(any()); + assertThat(concurrentStore.snapshot( + decision.key(), + decision.limit() + ).pendingReconciliationLiability().value()).isPositive(); + } + + @Test + @DisplayName("IN_FLIGHT 이후 downstream 실패는 pending liability로 보존한다") + void preservesPendingLiabilityAfterDownstreamFailure() { + ChatClientRequest request = request(); + RuntimeException failure = new IllegalStateException("provider failed"); + stubReservation(request); + when(accounting.markInFlight(RESERVATION_ID)) + .thenReturn(inFlight()); + when(provider.nextCall(any())).thenThrow(failure); + + assertThatThrownBy(() -> advisor.adviseCall(request, provider)) + .isSameAs(failure); + + verify(accounting).markReconciliationRequired( + RESERVATION_ID, + ReservationAccountingReason.ACTUAL_USAGE_UNAVAILABLE + ); + verifyNoInteractions(usageExtractor); + } + + @Test + @DisplayName("usage unavailable은 0원 commit 없이 provider 응답을 보존한다") + void preservesResponseWhenUsageIsUnavailable() { + ChatClientRequest request = request(); + ChatClientResponse response = response(); + stubDispatch(request, response); + when(usageExtractor.extract(response)) + .thenReturn(TokenUsage.unavailable(Map.of())); + when(accounting.markReconciliationRequired( + RESERVATION_ID, + ReservationAccountingReason.ACTUAL_USAGE_UNAVAILABLE + )).thenThrow(new IllegalStateException("recording failed")); + + ChatClientResponse result = advisor.adviseCall(request, provider); + + assertThat(result).isSameAs(response); + verify(accounting, times(1)).markReconciliationRequired( + RESERVATION_ID, + ReservationAccountingReason.ACTUAL_USAGE_UNAVAILABLE + ); + verify(accounting, never()).commit(any()); + } + + @Test + @DisplayName("usage extraction 오류는 provider 응답을 보존하고 정산 대기로 전환한다") + void preservesResponseWhenUsageExtractionFails() { + ChatClientRequest request = request(); + ChatClientResponse response = response(); + stubDispatch(request, response); + when(usageExtractor.extract(response)) + .thenThrow(new IllegalStateException("extraction failed")); + + ChatClientResponse result = advisor.adviseCall(request, provider); + + assertThat(result).isSameAs(response); + verify(accounting).markReconciliationRequired( + RESERVATION_ID, + ReservationAccountingReason.ACTUAL_USAGE_UNAVAILABLE + ); + verify(accounting, never()).commit(any()); + } + + @Test + @DisplayName("actual 비용 또는 통화 오류는 provider 응답을 보존하고 정산 대기로 전환한다") + void preservesResponseWhenActualCommitFails() { + ChatClientRequest request = request(); + ChatClientResponse response = response(); + stubDispatch(request, response); + when(usageExtractor.extract(response)).thenReturn(TokenUsage.from(10, 5)); + when(accounting.commit(any())).thenThrow(new IllegalStateException( + "calculated cost must use the pricing snapshot currency" + )); + + ChatClientResponse result = advisor.adviseCall(request, provider); + + assertThat(result).isSameAs(response); + verify(accounting).markReconciliationRequired( + RESERVATION_ID, + ReservationAccountingReason.ACTUAL_USAGE_UNAVAILABLE + ); + } + + @Test + @DisplayName("응답 모델이 예약 snapshot과 다르면 요청 모델 가격으로 commit하지 않는다") + void preservesPricingMismatchForLaterReconciliation() { + ChatClientRequest request = request(); + ChatClientResponse response = response("model-v2"); + stubDispatch(request, response); + when(usageExtractor.extract(response)).thenReturn(TokenUsage.from(10, 5)); + + ChatClientResponse result = advisor.adviseCall(request, provider); + + assertThat(result).isSameAs(response); + verify(accounting, never()).commit(any()); + verify(accounting).markReconciliationRequired( + argThat((ActualUsageCommand actual) -> + actual.reservationId().equals(RESERVATION_ID) + && actual.requestId().equals("request-1") + && actual.responseModelId().equals("model-v2") + && actual.usage().equals(TokenUsage.from(10, 5))), + eq(ReservationAccountingReason.PRICING_RECONCILIATION_REQUIRED) + ); + } + + private PreflightCostResult.Bounded stubDispatch( + ChatClientRequest request, + ChatClientResponse response + ) { + PreflightCostResult.Bounded costBound = stubReservation(request); + when(accounting.markInFlight(RESERVATION_ID)).thenReturn(inFlight()); + when(provider.nextCall(any())).thenReturn(response); + return costBound; + } + + private PreflightCostResult.Bounded stubReservation(ChatClientRequest request) { + PreflightCostResult.Bounded costBound = costBound(); + Cost safeUpperBoundCost = costBound.safeUpperBoundCost(); + when(preflight.resolve(request)).thenReturn(costBound); + when(budgetEvaluator.evaluate(any(), any())) + .thenReturn(allowed(safeUpperBoundCost)); + BudgetReservation reservation = mock(BudgetReservation.class); + BudgetReservationResult result = mock(BudgetReservationResult.class); + when(reservation.id()).thenReturn(RESERVATION_ID); + when(result.isAccepted()).thenReturn(true); + when(result.status()).thenReturn(ReservationStatus.CREATED); + when(result.reservation()).thenReturn(reservation); + when(stateStore.checkAndReserve(any())).thenReturn(result); + return costBound; + } + + private PreflightCostResult.Bounded costBound() { + PreflightCostResult.Bounded result = mock(PreflightCostResult.Bounded.class); + PricingSnapshot pricingSnapshot = mock(PricingSnapshot.class); + when(pricingSnapshot.currency()).thenReturn(USD); + when(pricingSnapshot.modelId()).thenReturn("model-v1"); + when(result.safeUpperBoundCost()).thenReturn(usd("0.01")); + when(result.pricingSnapshot()).thenReturn(pricingSnapshot); + when(result.inputEstimatedTokens()).thenReturn(10L); + when(result.inputSafeUpperBoundTokens()).thenReturn(12L); + when(result.reservedOutputTokens()).thenReturn(5L); + return result; + } + + private BudgetDecision allowed(Cost candidate) { + return new BudgetDecision( + new BudgetKey( + "monthly", + "tenant", + "tenant-1", + BudgetWindow.parse("2026-08") + ), + BudgetDecision.EvaluationType.ADMISSION, + BudgetState.ALLOW, + BudgetThreshold.NONE, + "allowed", + Cost.zero(USD), + candidate, + usd("10.00") + ); + } + + private ChatClientRequest request() { + Map context = new HashMap<>(); + context.put(RequestContextAccessor.REQUEST_ID_CONTEXT_KEY, "request-1"); + context.put(RequestContextAccessor.ATTEMPT_ID_CONTEXT_KEY, "attempt-1"); + context.put( + RequestContextAccessor.IDEMPOTENCY_CONTEXT_KEY, + new IdempotencyKey("idempotency-1") + ); + context.put("tenant_id", "tenant-1"); + return new ChatClientRequest(new Prompt("question"), context); + } + + private ChatClientResponse response() { + return response("model-v1"); + } + + private ChatClientResponse response(String modelId) { + ChatResponse chatResponse = mock(ChatResponse.class); + ChatResponseMetadata metadata = mock(ChatResponseMetadata.class); + when(chatResponse.getMetadata()).thenReturn(metadata); + when(metadata.getModel()).thenReturn(modelId); + return new ChatClientResponse(chatResponse, Map.of()); + } + + private ReservationTransition inFlight() { + return ReservationTransition.applied( + ReservationState.RESERVED, + ReservationState.IN_FLIGHT + ); + } + + private Cost usd(String amount) { + return Cost.of(new BigDecimal(amount), USD); + } +} 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 1aa03d7..b48e7b2 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 @@ -13,6 +13,7 @@ import io.tokenpilot.core.domain.*; import io.tokenpilot.core.exception.MissingPricingException; import io.tokenpilot.core.internal.LedgerComponents; +import io.tokenpilot.springai.LedgerAdvisor; import io.tokenpilot.springai.UsageExtractor; import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Test; @@ -27,6 +28,7 @@ import org.springframework.ai.chat.model.Generation; import org.springframework.ai.chat.prompt.ChatOptions; import org.springframework.ai.chat.prompt.Prompt; +import org.springframework.core.Ordered; import java.math.BigDecimal; import java.time.Instant; @@ -663,6 +665,179 @@ void failOpenAllowsProviderCallAndMarksMissingPricingAsUnpriced() { verifyNoInteractions(ledgerManager); } + @Test + @DisplayName("pre-call missing pricing은 provider 전 한 번만 발행하고 after에서 중복하지 않는다") + void publishesPreCallPricingMissExactlyOnce() { + LedgerManager ledgerManager = mock(LedgerManager.class); + UsageExtractor extractor = mock(UsageExtractor.class); + PricingRegistry pricingRegistry = mock(PricingRegistry.class); + List events = new java.util.concurrent.CopyOnWriteArrayList<>(); + + when(extractor.extract(any())).thenReturn(TokenUsage.from(100, 200)); + when(pricingRegistry.resolveSnapshot("missing-model", PricingPlan.DEFAULT_PRICING_POLICY_ID)) + .thenReturn(Optional.empty()); + + LedgerAdvisor advisor = LedgerSpringAiComponents.defaultLedgerAdvisor( + ledgerManager, + extractor, + null, + null, + mock(CostCalculator.class), + pricingRegistry, + LedgerComponents.defaultPricingEvaluator(), + MissingPricingPolicy.FAIL_OPEN, + List.of(events::add) + ); + ChatClientRequest request = new ChatClientRequest( + new Prompt("test"), + Map.of(DefaultLedgerAdvisor.MODEL_ID_CONTEXT, "missing-model") + ); + + ChatClientRequest resolved = advisor.before(request, mock(AdvisorChain.class)); + advisor.after( + response("missing-model", resolved.context()), + mock(AdvisorChain.class) + ); + + assertThat(events).containsExactly(new PricingMissingEvent( + MissingPricingPolicy.FAIL_OPEN, + PricingResolution.MISSING_PLAN + )); + } + + @Test + @DisplayName("actual cost 계산의 MissingPricingException도 한 번 발행한다") + void publishesActualPricingMissExactlyOnce() { + LedgerManager ledgerManager = mock(LedgerManager.class); + UsageExtractor extractor = mock(UsageExtractor.class); + List events = new java.util.concurrent.CopyOnWriteArrayList<>(); + TokenUsage usage = TokenUsage.from(100, 200); + PricingSnapshot snapshot = PricingSnapshot.from( + new PricingPlan( + "gpt-4o", + new BigDecimal("0.01"), + new BigDecimal("0.03"), + Currency.getInstance("USD") + ), + PricingSnapshot.DEFAULT_CATALOG_VERSION, + Instant.parse("2026-07-30T00:00:00Z") + ); + + when(extractor.extract(any())).thenReturn(usage); + when(ledgerManager.record(same(snapshot), same(usage), anyMap())) + .thenThrow(new MissingPricingException(PricingResolution.MISSING_RATE)); + + DefaultLedgerAdvisor advisor = new DefaultLedgerAdvisor( + ledgerManager, + extractor, + null, + null, + mock(CostCalculator.class), + null, + LedgerComponents.defaultPricingEvaluator(), + MissingPricingPolicy.FAIL_OPEN, + List.of(events::add) + ); + + ChatClientResponse result = advisor.after( + response( + "gpt-4o", + Map.of( + DefaultLedgerAdvisor.PRICING_SNAPSHOT_CONTEXT, snapshot, + DefaultLedgerAdvisor.PRICING_RESOLUTION_CONTEXT, PricingResolution.RESOLVED + ) + ), + mock(AdvisorChain.class) + ); + + assertThat(result.context().get(DefaultLedgerAdvisor.PRICING_RECONCILIATION_RESULT_CONTEXT)) + .isEqualTo(PricingReconciliationResult.UNPRICED); + assertThat(events).containsExactly(new PricingMissingEvent( + MissingPricingPolicy.FAIL_OPEN, + PricingResolution.MISSING_RATE + )); + } + + @Test + @DisplayName("legacy actual pricing miss는 한 번 발행하고 기존 예외를 그대로 전파한다") + void publishesLegacyActualPricingMissWithoutChangingFailureSemantics() { + LedgerManager ledgerManager = mock(LedgerManager.class); + UsageExtractor extractor = mock(UsageExtractor.class); + List events = new java.util.concurrent.CopyOnWriteArrayList<>(); + TokenUsage usage = TokenUsage.from(100, 200); + MissingPricingException failure = new MissingPricingException(PricingResolution.MISSING_PLAN); + + when(extractor.extract(any())).thenReturn(usage); + when(ledgerManager.record(eq("missing-model"), same(usage), anyMap())) + .thenThrow(failure); + + DefaultLedgerAdvisor advisor = new DefaultLedgerAdvisor( + ledgerManager, + extractor, + null, + null, + null, + null, + LedgerComponents.defaultPricingEvaluator(), + MissingPricingPolicy.FAIL_OPEN, + List.of(events::add) + ); + + assertThatThrownBy(() -> advisor.after( + response("missing-model", Map.of()), + mock(AdvisorChain.class) + )).isSameAs(failure); + assertThat(events).containsExactly(new PricingMissingEvent( + MissingPricingPolicy.FAIL_OPEN, + PricingResolution.MISSING_PLAN + )); + } + + @Test + @DisplayName("pricing miss listener 실패는 fail-closed 판정과 다음 listener를 바꾸지 않는다") + void isolatesPricingMissListenerFailures() { + LedgerManager ledgerManager = mock(LedgerManager.class); + PricingRegistry pricingRegistry = mock(PricingRegistry.class); + AtomicInteger failedDeliveries = new AtomicInteger(); + List received = new java.util.concurrent.CopyOnWriteArrayList<>(); + + when(pricingRegistry.resolveSnapshot("missing-model", PricingPlan.DEFAULT_PRICING_POLICY_ID)) + .thenReturn(Optional.empty()); + + DefaultLedgerAdvisor advisor = new DefaultLedgerAdvisor( + ledgerManager, + mock(UsageExtractor.class), + null, + null, + mock(CostCalculator.class), + pricingRegistry, + LedgerComponents.defaultPricingEvaluator(), + MissingPricingPolicy.FAIL_CLOSED, + List.of( + event -> { + failedDeliveries.incrementAndGet(); + throw new IllegalStateException("listener failed"); + }, + received::add + ) + ); + ChatClientRequest request = new ChatClientRequest( + new Prompt("test"), + Map.of(DefaultLedgerAdvisor.MODEL_ID_CONTEXT, "missing-model") + ); + + assertThatThrownBy(() -> advisor.before(request, mock(AdvisorChain.class))) + .isInstanceOf(MissingPricingException.class) + .extracting(exception -> ((MissingPricingException) exception).getResolution()) + .isEqualTo(PricingResolution.MISSING_PLAN); + assertThat(failedDeliveries).hasValue(1); + assertThat(received).containsExactly(new PricingMissingEvent( + MissingPricingPolicy.FAIL_CLOSED, + PricingResolution.MISSING_PLAN + )); + verifyNoInteractions(ledgerManager); + } + @Test @DisplayName("FAIL_OPEN은 model id가 없어도 MISSING_PLAN을 보존하고 UNPRICED로 남겨야 한다") void failOpenPreservesMissingPlanWhenModelIdIsMissing() { @@ -1010,7 +1185,7 @@ void checkAdvisorMetadata() { DefaultLedgerAdvisor advisor = new DefaultLedgerAdvisor(mock(LedgerManager.class), mock(UsageExtractor.class)); assertThat(advisor.getName()).isEqualTo("LedgerAdvisor"); - assertThat(advisor.getOrder()).isEqualTo(0); + assertThat(advisor.getOrder()).isEqualTo(Ordered.LOWEST_PRECEDENCE - 1); } private static final Currency USD = Currency.getInstance("USD"); diff --git a/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/DefaultRequestAdapterTest.java b/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/DefaultRequestAdapterTest.java new file mode 100644 index 0000000..5f8cb9f --- /dev/null +++ b/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/DefaultRequestAdapterTest.java @@ -0,0 +1,50 @@ +package io.tokenpilot.springai.internal; + +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.springframework.ai.chat.client.ChatClientRequest; +import org.springframework.ai.chat.messages.AssistantMessage; +import org.springframework.ai.chat.messages.SystemMessage; +import org.springframework.ai.chat.messages.UserMessage; +import org.springframework.ai.chat.prompt.Prompt; + +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; + +class DefaultRequestAdapterTest { + + private final DefaultRequestAdapter adapter = new DefaultRequestAdapter(); + + @Test + @DisplayName("지원 message의 role과 순서를 보존한다") + void preservesSupportedMessageRolesAndOrder() { + ChatClientRequest request = new ChatClientRequest( + new Prompt(List.of( + new SystemMessage("system instructions"), + new UserMessage("user question"), + new AssistantMessage("previous answer") + )), + Map.of() + ); + + AdaptedRequest adaptedRequest = adapter.adapt(request); + + assertThat(adaptedRequest.messages()).containsExactly( + new AdaptedMessage( + RequestMessageRole.SYSTEM, + "system instructions" + ), + new AdaptedMessage( + RequestMessageRole.USER, + "user question" + ), + new AdaptedMessage( + RequestMessageRole.ASSISTANT, + "previous answer" + ) + ); + } + +} diff --git a/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/IdempotencyKeyResolverTest.java b/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/IdempotencyKeyResolverTest.java new file mode 100644 index 0000000..3b0d7e5 --- /dev/null +++ b/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/IdempotencyKeyResolverTest.java @@ -0,0 +1,100 @@ +package io.tokenpilot.springai.internal; + +import io.tokenpilot.budget.IdempotencyKey; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.springframework.ai.chat.client.ChatClientRequest; +import org.springframework.ai.chat.prompt.Prompt; + +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +class IdempotencyKeyResolverTest { + + private final RequestContextAccessor contextAccessor = new RequestContextAccessor(); + + @Test + @DisplayName("caller가 제공한 idempotency key를 그대로 사용한다") + void usesCallerProvidedIdempotencyKey() { + AtomicInteger generationCount = new AtomicInteger(); + IdempotencyKey callerKey = new IdempotencyKey("caller-key"); + IdempotencyKeyResolver resolver = resolver(generationCount); + ChatClientRequest request = requestWithContext(Map.of( + RequestContextAccessor.IDEMPOTENCY_CONTEXT_KEY, + callerKey + )); + + IdempotencyKey resolved = resolver.resolve(request); + + assertThat(resolved).isEqualTo(callerKey); + assertThat(generationCount).hasValue(0); + } + + @Test + @DisplayName("idempotency key가 없으면 현재 호출을 위한 key를 한 번 생성한다") + void generatesIdempotencyKeyForCurrentInvocation() { + AtomicInteger generationCount = new AtomicInteger(); + IdempotencyKeyResolver resolver = resolver(generationCount); + ChatClientRequest request = requestWithContext(Map.of()); + + IdempotencyKey resolved = resolver.resolve(request); + + assertThat(resolved).isEqualTo(new IdempotencyKey("generated-key")); + assertThat(generationCount).hasValue(1); + } + + @Test + @DisplayName("생성한 idempotency key는 같은 request context에서 재사용한다") + void reusesGeneratedIdempotencyKeyInSameRequestContext() { + AtomicInteger generationCount = new AtomicInteger(); + IdempotencyKeyResolver resolver = resolver(generationCount); + ChatClientRequest request = requestWithContext(Map.of()); + + IdempotencyKey generatedKey = resolver.resolve(request); + ChatClientRequest requestWithKey = contextAccessor.withIdempotencyKey( + request, + generatedKey + ); + IdempotencyKey reusedKey = resolver.resolve(requestWithKey); + + assertThat(reusedKey).isEqualTo(generatedKey); + assertThat(generationCount).hasValue(1); + } + + @Test + @DisplayName("잘못된 타입의 기존 idempotency key를 새 key로 대체하지 않는다") + void rejectsInvalidCallerIdempotencyKey() { + AtomicInteger generationCount = new AtomicInteger(); + IdempotencyKeyResolver resolver = resolver(generationCount); + ChatClientRequest request = requestWithContext(Map.of( + RequestContextAccessor.IDEMPOTENCY_CONTEXT_KEY, + "caller-key" + )); + + assertThatThrownBy(() -> resolver.resolve(request)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining(RequestContextAccessor.IDEMPOTENCY_CONTEXT_KEY); + assertThat(generationCount).hasValue(0); + } + + private IdempotencyKeyResolver resolver(AtomicInteger generationCount) { + return new IdempotencyKeyResolver( + contextAccessor, + () -> { + generationCount.incrementAndGet(); + return new IdempotencyKey("generated-key"); + } + ); + } + + private ChatClientRequest requestWithContext(Map context) { + return new ChatClientRequest( + new Prompt("user question"), + new HashMap<>(context) + ); + } +} diff --git a/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/LedgerSpringAiComponentsAccountingTest.java b/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/LedgerSpringAiComponentsAccountingTest.java new file mode 100644 index 0000000..ef96704 --- /dev/null +++ b/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/LedgerSpringAiComponentsAccountingTest.java @@ -0,0 +1,39 @@ +package io.tokenpilot.springai.internal; + +import io.tokenpilot.budget.BudgetEvaluator; +import io.tokenpilot.budget.BudgetStateStore; +import io.tokenpilot.budget.ReservationAccounting; +import io.tokenpilot.core.ModelRegistry; +import io.tokenpilot.core.PreflightCostEstimator; +import io.tokenpilot.core.PricingRegistry; +import io.tokenpilot.core.TokenBudget; +import io.tokenpilot.core.TokenEstimator; +import io.tokenpilot.springai.UsageExtractor; +import org.junit.jupiter.api.Test; + +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.Mockito.mock; + +class LedgerSpringAiComponentsAccountingTest { + + @Test + void 서로_다른_예약_store와_accounting_writer를_조합하지_않는다() { + assertThatThrownBy(() -> LedgerSpringAiComponents.accountingLedgerAdvisor( + mock(UsageExtractor.class), + mock(BudgetEvaluator.class), + mock(BudgetStateStore.class), + mock(ReservationAccounting.class), + mock(PricingRegistry.class), + mock(ModelRegistry.class), + mock(TokenEstimator.class), + mock(TokenBudget.class), + mock(PreflightCostEstimator.class), + null, + null, + 0L + )).isInstanceOf(IllegalArgumentException.class) + .hasMessage( + "budgetStateStore and reservationAccounting must be the same instance" + ); + } +} diff --git a/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/ModelResolverTest.java b/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/ModelResolverTest.java new file mode 100644 index 0000000..b9cc67b --- /dev/null +++ b/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/ModelResolverTest.java @@ -0,0 +1,72 @@ +package io.tokenpilot.springai.internal; + +import io.tokenpilot.core.CoreComponents; +import io.tokenpilot.core.ModelRegistry; +import io.tokenpilot.core.domain.ModelDefinition; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.springframework.ai.chat.client.ChatClientRequest; +import org.springframework.ai.chat.prompt.ChatOptions; +import org.springframework.ai.chat.prompt.Prompt; + +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; + +class ModelResolverTest { + + private final ModelRegistry modelRegistry = CoreComponents.defaultModelRegistry(); + + @Test + @DisplayName("request model을 configured default보다 우선한다") + void prioritizesRequestModelOverConfiguredDefault() { + ModelResolver resolver = new ModelResolver(modelRegistry, "gpt-4o-mini"); + + assertThat(resolver.resolve(requestWithModel("gpt-4o"))) + .get() + .extracting(ModelDefinition::canonicalModelId) + .isEqualTo("gpt-4o-2024-08-06"); + } + + @Test + @DisplayName("request model이 없으면 configured default를 사용한다") + void usesConfiguredDefaultWhenRequestModelIsAbsent() { + ModelResolver resolver = new ModelResolver(modelRegistry, "gpt-4o-mini"); + + assertThat(resolver.resolve(requestWithoutModel())) + .get() + .extracting(ModelDefinition::canonicalModelId) + .isEqualTo("gpt-4o-mini-2024-07-18"); + } + + @Test + @DisplayName("request model과 configured default가 모두 없으면 해석하지 못한다") + void doesNotResolveWhenRequestAndDefaultModelsAreAbsent() { + ModelResolver resolver = new ModelResolver(modelRegistry); + + assertThat(resolver.resolve(requestWithoutModel())).isEmpty(); + } + + @Test + @DisplayName("blank 또는 unknown request model은 configured default로 대체하지 않는다") + void doesNotFallbackForInvalidRequestModel() { + ModelResolver resolver = new ModelResolver(modelRegistry, "gpt-4o-mini"); + + assertThat(resolver.resolve(requestWithModel(" "))).isEmpty(); + assertThat(resolver.resolve(requestWithModel("unknown-model"))).isEmpty(); + } + + private ChatClientRequest requestWithModel(String modelId) { + return new ChatClientRequest( + new Prompt( + "user question", + ChatOptions.builder().model(modelId).build() + ), + Map.of() + ); + } + + private ChatClientRequest requestWithoutModel() { + return new ChatClientRequest(new Prompt("user question"), Map.of()); + } +} diff --git a/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/RequestContextAccessorTest.java b/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/RequestContextAccessorTest.java new file mode 100644 index 0000000..f5664ea --- /dev/null +++ b/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/RequestContextAccessorTest.java @@ -0,0 +1,148 @@ +package io.tokenpilot.springai.internal; + +import io.tokenpilot.budget.IdempotencyKey; +import io.tokenpilot.budget.ReservationId; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.springframework.ai.chat.client.ChatClientRequest; +import org.springframework.ai.chat.prompt.Prompt; + +import java.util.HashMap; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +class RequestContextAccessorTest { + + private final RequestContextAccessor accessor = new RequestContextAccessor(); + + @Test + @DisplayName("caller가 제공한 request correlation을 namespaced context에서 읽는다") + void readsCallerProvidedRequestCorrelation() { + IdempotencyKey idempotencyKey = new IdempotencyKey("idempotency-1"); + ChatClientRequest request = requestWithContext(Map.of( + RequestContextAccessor.REQUEST_ID_CONTEXT_KEY, "request-1", + RequestContextAccessor.ATTEMPT_ID_CONTEXT_KEY, "attempt-1", + RequestContextAccessor.IDEMPOTENCY_CONTEXT_KEY, idempotencyKey + )); + + assertThat(accessor.requestId(request)).isEqualTo("request-1"); + assertThat(accessor.attemptId(request)).isEqualTo("attempt-1"); + assertThat(accessor.idempotencyKey(request)).isEqualTo(idempotencyKey); + assertThat(accessor.reservationId(request)).isNull(); + } + + @Test + @DisplayName("context에 없는 correlation 값은 null이다") + void returnsNullForMissingCorrelation() { + ChatClientRequest request = requestWithContext(Map.of()); + + assertThat(accessor.requestId(request)).isNull(); + assertThat(accessor.attemptId(request)).isNull(); + assertThat(accessor.idempotencyKey(request)).isNull(); + assertThat(accessor.reservationId(request)).isNull(); + } + + @Test + @DisplayName("context에 존재하는 correlation 값의 타입이 다르면 거부한다") + void rejectsCorrelationWithWrongType() { + ChatClientRequest request = requestWithContext(Map.of( + RequestContextAccessor.IDEMPOTENCY_CONTEXT_KEY, + "idempotency-1" + )); + + assertThatThrownBy(() -> accessor.idempotencyKey(request)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining(RequestContextAccessor.IDEMPOTENCY_CONTEXT_KEY); + } + + @Test + @DisplayName("request mutation 이후에도 correlation과 application context를 보존한다") + void preservesCorrelationAcrossRequestMutation() { + IdempotencyKey idempotencyKey = new IdempotencyKey("idempotency-1"); + ChatClientRequest request = requestWithContext(Map.of( + RequestContextAccessor.REQUEST_ID_CONTEXT_KEY, "request-1", + RequestContextAccessor.ATTEMPT_ID_CONTEXT_KEY, "attempt-1", + RequestContextAccessor.IDEMPOTENCY_CONTEXT_KEY, idempotencyKey, + "application.context", "application-value" + )); + + ChatClientRequest mutated = request.mutate() + .context("application.mutated", true) + .build(); + + assertThat(accessor.requestId(mutated)).isEqualTo("request-1"); + assertThat(accessor.attemptId(mutated)).isEqualTo("attempt-1"); + assertThat(accessor.idempotencyKey(mutated)).isEqualTo(idempotencyKey); + assertThat(mutated.context()) + .containsEntry("application.context", "application-value") + .containsEntry("application.mutated", true); + } + + @Test + @DisplayName("reservation ID를 추가해도 기존 correlation과 원본 request를 보존한다") + void addsReservationIdWithoutChangingExistingContext() { + IdempotencyKey idempotencyKey = new IdempotencyKey("idempotency-1"); + ReservationId reservationId = new ReservationId("reservation-1"); + ChatClientRequest request = requestWithContext(Map.of( + RequestContextAccessor.REQUEST_ID_CONTEXT_KEY, "request-1", + RequestContextAccessor.ATTEMPT_ID_CONTEXT_KEY, "attempt-1", + RequestContextAccessor.IDEMPOTENCY_CONTEXT_KEY, idempotencyKey, + "application.context", "application-value" + )); + + ChatClientRequest reservedRequest = accessor.withReservationId( + request, + reservationId + ); + + assertThat(accessor.reservationId(request)).isNull(); + assertThat(accessor.reservationId(reservedRequest)).isEqualTo(reservationId); + assertThat(accessor.requestId(reservedRequest)).isEqualTo("request-1"); + assertThat(accessor.attemptId(reservedRequest)).isEqualTo("attempt-1"); + assertThat(accessor.idempotencyKey(reservedRequest)).isEqualTo(idempotencyKey); + assertThat(reservedRequest.context()) + .containsEntry("application.context", "application-value"); + } + + @Test + @DisplayName("같은 reservation ID를 다시 추가하면 기존 correlation을 유지한다") + void preservesMatchingReservationId() { + ReservationId reservationId = new ReservationId("reservation-1"); + ChatClientRequest reservedRequest = requestWithContext(Map.of( + RequestContextAccessor.RESERVATION_ID_CONTEXT_KEY, + reservationId + )); + + ChatClientRequest result = accessor.withReservationId( + reservedRequest, + reservationId + ); + + assertThat(accessor.reservationId(result)).isEqualTo(reservationId); + } + + @Test + @DisplayName("다른 reservation ID로 기존 correlation을 변경할 수 없다") + void rejectsConflictingReservationId() { + ChatClientRequest reservedRequest = requestWithContext(Map.of( + RequestContextAccessor.RESERVATION_ID_CONTEXT_KEY, + new ReservationId("reservation-1") + )); + + assertThatThrownBy(() -> accessor.withReservationId( + reservedRequest, + new ReservationId("reservation-2") + )) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining(RequestContextAccessor.RESERVATION_ID_CONTEXT_KEY); + } + + private ChatClientRequest requestWithContext(Map context) { + return new ChatClientRequest( + new Prompt("user question"), + new HashMap<>(context) + ); + } +} diff --git a/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/RequestFramingPolicyTest.java b/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/RequestFramingPolicyTest.java new file mode 100644 index 0000000..88c8efd --- /dev/null +++ b/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/RequestFramingPolicyTest.java @@ -0,0 +1,26 @@ +package io.tokenpilot.springai.internal; + +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +class RequestFramingPolicyTest { + + private final RequestFramingPolicy policy = new RequestFramingPolicy(); + + @Test + @DisplayName("message role과 순서 및 text 경계를 명시적인 framing으로 보존한다") + void preservesRoleOrderAndTextBoundaries() { + AdaptedRequest request = new AdaptedRequest(List.of( + new AdaptedMessage(RequestMessageRole.SYSTEM, "rules"), + new AdaptedMessage(RequestMessageRole.USER, "a\nb") + )); + + String framed = policy.frame(request); + + assertThat(framed).isEqualTo("SYSTEM:5:rules\nUSER:3:a\nb\n"); + } +} diff --git a/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/RequestPreflightIdempotencyTest.java b/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/RequestPreflightIdempotencyTest.java new file mode 100644 index 0000000..61e77d9 --- /dev/null +++ b/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/RequestPreflightIdempotencyTest.java @@ -0,0 +1,130 @@ +package io.tokenpilot.springai.internal; + +import io.tokenpilot.budget.AtomicBudgetStateStore; +import io.tokenpilot.budget.BudgetKey; +import io.tokenpilot.budget.BudgetReservationRequest; +import io.tokenpilot.budget.BudgetWindow; +import io.tokenpilot.budget.IdempotencyKey; +import io.tokenpilot.budget.ReservationStatus; +import io.tokenpilot.budget.ReservationTokenEstimate; +import io.tokenpilot.budget.internal.LedgerBudgetComponents; +import io.tokenpilot.core.ModelRegistry; +import io.tokenpilot.core.PricingRegistry; +import io.tokenpilot.core.domain.Cost; +import io.tokenpilot.core.domain.ModelDefinition; +import io.tokenpilot.core.domain.PreflightCostResult; +import io.tokenpilot.core.domain.PricingSnapshot; +import io.tokenpilot.core.domain.TokenType; +import io.tokenpilot.core.internal.LedgerComponents; +import org.junit.jupiter.api.Test; +import org.springframework.ai.chat.client.ChatClientRequest; +import org.springframework.ai.chat.prompt.ChatOptions; +import org.springframework.ai.chat.prompt.Prompt; + +import java.math.BigDecimal; +import java.time.Instant; +import java.util.Currency; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +class RequestPreflightIdempotencyTest { + + private static final Currency USD = Currency.getInstance("USD"); + + @Test + void 실제_preflight의_서로_다른_pricing_조회_시각은_동일_예약을_REUSED한다() { + ModelRegistry modelRegistry = LedgerComponents.defaultModelRegistry(); + PricingRegistry pricingRegistry = mock(PricingRegistry.class); + PricingSnapshot firstSnapshot = snapshot( + Instant.parse("2026-08-25T00:00:00Z") + ); + PricingSnapshot secondSnapshot = snapshot( + Instant.parse("2026-08-25T00:00:01Z") + ); + AtomicInteger resolution = new AtomicInteger(); + when(pricingRegistry.resolveSnapshot(any(ModelDefinition.class))) + .thenAnswer(ignored -> Optional.of( + resolution.getAndIncrement() == 0 + ? firstSnapshot + : secondSnapshot + )); + RequestPreflight preflight = new RequestPreflight( + new ModelResolver(modelRegistry), + new ReservedOutputResolver(), + LedgerComponents.utf8ByteHeuristicTokenEstimator(), + LedgerComponents.tokenBudget(modelRegistry), + pricingRegistry, + LedgerComponents.defaultPricingEvaluator(), + LedgerComponents.defaultPreflightCostEstimator(), + 0 + ); + ChatClientRequest request = new ChatClientRequest( + new Prompt( + "same request", + ChatOptions.builder() + .model("gpt-4o-2024-08-06") + .maxTokens(100) + .build() + ), + Map.of() + ); + + PreflightCostResult.Bounded firstBound = preflight.resolve(request); + PreflightCostResult.Bounded secondBound = preflight.resolve(request); + AtomicBudgetStateStore store = LedgerBudgetComponents + .inMemoryAtomicBudgetStateStore(); + var created = store.checkAndReserve(reservation(firstBound)); + var reused = store.checkAndReserve(reservation(secondBound)); + + assertThat(firstBound.pricingSnapshot().checkedAt()) + .isNotEqualTo(secondBound.pricingSnapshot().checkedAt()); + assertThat(created.status()).isEqualTo(ReservationStatus.CREATED); + assertThat(reused.status()).isEqualTo(ReservationStatus.REUSED); + assertThat(reused.reservation().id()).isEqualTo(created.reservation().id()); + assertThat(reused.reservation().pricingSnapshot()) + .contains(firstSnapshot); + } + + private BudgetReservationRequest reservation( + PreflightCostResult.Bounded bound + ) { + return new BudgetReservationRequest( + new BudgetKey( + "monthly", + "tenant", + "tenant-1", + BudgetWindow.parse("2026-08") + ), + Cost.of(new BigDecimal("10.00"), USD), + bound.safeUpperBoundCost(), + "request-1", + new IdempotencyKey("idempotency-1"), + bound.pricingSnapshot(), + new ReservationTokenEstimate( + bound.inputEstimatedTokens(), + bound.inputSafeUpperBoundTokens(), + bound.reservedOutputTokens() + ) + ); + } + + private PricingSnapshot snapshot(Instant checkedAt) { + return new PricingSnapshot( + "gpt-4o-2024-08-06", + "default", + "openai-2026-08-14", + checkedAt, + Map.of( + TokenType.PROMPT, new BigDecimal("0.00015"), + TokenType.COMPLETION, new BigDecimal("0.00060") + ), + USD + ); + } +} diff --git a/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/RequestScopeResolverTest.java b/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/RequestScopeResolverTest.java new file mode 100644 index 0000000..0103dab --- /dev/null +++ b/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/RequestScopeResolverTest.java @@ -0,0 +1,220 @@ +package io.tokenpilot.springai.internal; + +import io.tokenpilot.core.domain.TokenCountScope; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.springframework.ai.chat.client.ChatClientAttributes; +import org.springframework.ai.chat.client.ChatClientRequest; +import org.springframework.ai.chat.messages.AssistantMessage; +import org.springframework.ai.chat.messages.ToolResponseMessage; +import org.springframework.ai.chat.messages.UserMessage; +import org.springframework.ai.chat.prompt.ChatOptions; +import org.springframework.ai.chat.prompt.Prompt; +import org.springframework.ai.content.Media; +import org.springframework.ai.model.tool.StructuredOutputChatOptions; +import org.springframework.ai.model.tool.ToolCallingChatOptions; +import org.springframework.ai.tool.ToolCallback; +import org.springframework.util.MimeTypeUtils; + +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Mockito.mock; + +class RequestScopeResolverTest { + + private final RequestScopeResolver resolver = new RequestScopeResolver(); + + @Test + @DisplayName("text message만 있는 요청은 TEXT_ONLY 범위다") + void resolvesTextOnlyRequestScope() { + ChatClientRequest request = new ChatClientRequest( + new Prompt(List.of(new UserMessage("user question"))), + Map.of() + ); + + RequestScopeResult result = resolver.resolve(request); + + assertThat(result.scope()).contains(TokenCountScope.TEXT_ONLY); + assertThat(result.unsupportedReason()).isEmpty(); + } + + @Test + @DisplayName("tool message는 지원하지 않는 request scope다") + void rejectsToolMessageAsUnsupportedScope() { + ToolResponseMessage toolMessage = ToolResponseMessage.builder() + .responses(List.of(new ToolResponseMessage.ToolResponse( + "call-1", + "weather", + "sunny" + ))) + .build(); + ChatClientRequest request = new ChatClientRequest( + new Prompt(List.of(toolMessage)), + Map.of() + ); + + RequestScopeResult result = resolver.resolve(request); + + assertThat(result.scope()).isEmpty(); + assertThat(result.unsupportedReason()).contains( + UnsupportedRequestScopeReason.TOOL_MESSAGE + ); + } + + @Test + @DisplayName("media가 포함된 user message는 지원하지 않는 request scope다") + void rejectsUserMessageWithMediaAsUnsupportedScope() { + Media media = Media.builder() + .mimeType(MimeTypeUtils.IMAGE_PNG) + .data(new byte[]{1}) + .build(); + UserMessage userMessage = UserMessage.builder() + .text("describe this image") + .media(media) + .build(); + ChatClientRequest request = new ChatClientRequest( + new Prompt(List.of(userMessage)), + Map.of() + ); + + RequestScopeResult result = resolver.resolve(request); + + assertThat(result.scope()).isEmpty(); + assertThat(result.unsupportedReason()).contains( + UnsupportedRequestScopeReason.MEDIA + ); + } + + @Test + @DisplayName("tool call이 포함된 assistant message는 지원하지 않는 request scope다") + void rejectsAssistantMessageWithToolCallAsUnsupportedScope() { + AssistantMessage assistantMessage = AssistantMessage.builder() + .content("") + .toolCalls(List.of(new AssistantMessage.ToolCall( + "call-1", + "function", + "weather", + "{}" + ))) + .build(); + ChatClientRequest request = new ChatClientRequest( + new Prompt(List.of(assistantMessage)), + Map.of() + ); + + RequestScopeResult result = resolver.resolve(request); + + assertThat(result.scope()).isEmpty(); + assertThat(result.unsupportedReason()).contains( + UnsupportedRequestScopeReason.TOOL_CALL + ); + } + + @Test + @DisplayName("tool schema가 포함된 request option은 지원하지 않는 request scope다") + void rejectsToolSchemaAsUnsupportedScope() { + ToolCallback toolCallback = mock(ToolCallback.class); + ChatOptions options = ToolCallingChatOptions.builder() + .toolCallbacks(toolCallback) + .build(); + ChatClientRequest request = new ChatClientRequest( + new Prompt(List.of(new UserMessage("user question")), options), + Map.of() + ); + + RequestScopeResult result = resolver.resolve(request); + + assertThat(result.scope()).isEmpty(); + assertThat(result.unsupportedReason()).contains( + UnsupportedRequestScopeReason.TOOL_SCHEMA + ); + } + + @Test + @DisplayName("output format이 추가될 요청은 지원하지 않는 request scope다") + void rejectsOutputFormatAugmentationAsUnsupportedScope() { + ChatClientRequest request = new ChatClientRequest( + new Prompt(List.of(new UserMessage("user question"))), + Map.of(ChatClientAttributes.OUTPUT_FORMAT.getKey(), "Return JSON") + ); + + RequestScopeResult result = resolver.resolve(request); + + assertThat(result.scope()).isEmpty(); + assertThat(result.unsupportedReason()).contains( + UnsupportedRequestScopeReason.STRUCTURED_OUTPUT + ); + } + + @Test + @DisplayName("structured output schema가 추가될 요청은 지원하지 않는 request scope다") + void rejectsStructuredOutputSchemaAsUnsupportedScope() { + ChatClientRequest request = new ChatClientRequest( + new Prompt(List.of(new UserMessage("user question"))), + Map.of(ChatClientAttributes.STRUCTURED_OUTPUT_SCHEMA.getKey(), "{type: object}") + ); + + RequestScopeResult result = resolver.resolve(request); + + assertThat(result.scope()).isEmpty(); + assertThat(result.unsupportedReason()).contains( + UnsupportedRequestScopeReason.STRUCTURED_OUTPUT + ); + } + + @Test + @DisplayName("structured output option은 지원하지 않는 request scope다") + void rejectsStructuredOutputOptionAsUnsupportedScope() { + ChatOptions options = StructuredOutputChatOptions.builder() + .outputSchema("{type: object}") + .build(); + ChatClientRequest request = new ChatClientRequest( + new Prompt(List.of(new UserMessage("user question")), options), + Map.of() + ); + + RequestScopeResult result = resolver.resolve(request); + + assertThat(result.scope()).isEmpty(); + assertThat(result.unsupportedReason()).contains( + UnsupportedRequestScopeReason.STRUCTURED_OUTPUT + ); + } + + @Test + @DisplayName("provider native structured output은 지원하지 않는 request scope다") + void rejectsNativeStructuredOutputAsUnsupportedScope() { + ChatClientRequest request = new ChatClientRequest( + new Prompt(List.of(new UserMessage("user question"))), + Map.of(ChatClientAttributes.STRUCTURED_OUTPUT_NATIVE.getKey(), true) + ); + + RequestScopeResult result = resolver.resolve(request); + + assertThat(result.scope()).isEmpty(); + assertThat(result.unsupportedReason()).contains( + UnsupportedRequestScopeReason.STRUCTURED_OUTPUT + ); + } + + @Test + @DisplayName("text가 없는 assistant message는 예외 대신 제한된 미지원 결과를 반환한다") + void rejectsAssistantMessageWithoutTextAsUnsupportedScope() { + AssistantMessage assistantMessage = AssistantMessage.builder() + .content(null) + .build(); + ChatClientRequest request = new ChatClientRequest( + new Prompt(List.of(assistantMessage)), + Map.of() + ); + + RequestScopeResult result = resolver.resolve(request); + + assertThat(result.scope()).isEmpty(); + assertThat(result.unsupportedReason()).contains( + UnsupportedRequestScopeReason.MESSAGE_TEXT_UNAVAILABLE + ); + } +} diff --git a/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/ReservedOutputResolverTest.java b/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/ReservedOutputResolverTest.java new file mode 100644 index 0000000..a45a734 --- /dev/null +++ b/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/ReservedOutputResolverTest.java @@ -0,0 +1,121 @@ +package io.tokenpilot.springai.internal; + +import io.tokenpilot.springai.ReservedOutputTokensResolver; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.springframework.ai.chat.client.ChatClientRequest; +import org.springframework.ai.chat.prompt.ChatOptions; +import org.springframework.ai.chat.prompt.Prompt; + +import java.util.Map; +import java.util.OptionalLong; + +import static org.assertj.core.api.Assertions.assertThat; + +class ReservedOutputResolverTest { + + @Test + @DisplayName("request maxTokens를 configured default보다 우선한다") + void prioritizesRequestMaxTokensOverConfiguredDefault() { + ReservedOutputResolver resolver = new ReservedOutputResolver(4_096); + + assertThat(resolver.resolve(requestWithMaxTokens(1_024))) + .hasValue(1_024); + } + + @Test + @DisplayName("request maxTokens가 없으면 configured default를 사용한다") + void usesConfiguredDefaultWhenRequestMaxTokensIsAbsent() { + ReservedOutputResolver resolver = new ReservedOutputResolver(4_096); + + assertThat(resolver.resolve(requestWithoutMaxTokens())) + .hasValue(4_096); + } + + @Test + @DisplayName("request와 유효한 configured default가 모두 없으면 해석하지 못한다") + void doesNotResolveWithoutRequestOrValidDefault() { + assertThat(new ReservedOutputResolver().resolve(requestWithoutMaxTokens())) + .isEmpty(); + assertThat(new ReservedOutputResolver(0).resolve(requestWithoutMaxTokens())) + .isEmpty(); + } + + @Test + @DisplayName("0 이하 request maxTokens는 configured default로 대체하지 않는다") + void doesNotFallbackForInvalidRequestMaxTokens() { + ReservedOutputResolver resolver = new ReservedOutputResolver(4_096); + + assertThat(resolver.resolve(requestWithMaxTokens(0))).isEmpty(); + assertThat(resolver.resolve(requestWithMaxTokens(-1))).isEmpty(); + } + + @Test + @DisplayName("request maxTokens가 없으면 provider resolver의 값을 사용한다") + void usesProviderResolverWhenRequestMaxTokensIsAbsent() { + ReservedOutputTokensResolver providerResolver = request -> + OptionalLong.of(2_048); + ReservedOutputResolver resolver = new ReservedOutputResolver( + providerResolver, + 4_096L + ); + + assertThat(resolver.resolve(requestWithoutMaxTokens())) + .hasValue(2_048); + } + + @Test + @DisplayName("request maxTokens는 provider resolver보다 우선한다") + void prioritizesRequestMaxTokensOverProviderResolver() { + ReservedOutputTokensResolver providerResolver = request -> + OptionalLong.of(2_048); + ReservedOutputResolver resolver = new ReservedOutputResolver( + providerResolver, + 4_096L + ); + + assertThat(resolver.resolve(requestWithMaxTokens(1_024))) + .hasValue(1_024); + } + + @Test + @DisplayName("provider resolver가 해석하지 못하면 configured default를 사용한다") + void usesConfiguredDefaultWhenProviderResolverDoesNotResolve() { + ReservedOutputTokensResolver providerResolver = request -> + OptionalLong.empty(); + ReservedOutputResolver resolver = new ReservedOutputResolver( + providerResolver, + 4_096L + ); + + assertThat(resolver.resolve(requestWithoutMaxTokens())) + .hasValue(4_096); + } + + @Test + @DisplayName("provider resolver의 0 이하 값은 configured default로 대체하지 않는다") + void doesNotFallbackForInvalidProviderValue() { + ReservedOutputTokensResolver providerResolver = request -> + OptionalLong.of(0); + ReservedOutputResolver resolver = new ReservedOutputResolver( + providerResolver, + 4_096L + ); + + assertThat(resolver.resolve(requestWithoutMaxTokens())).isEmpty(); + } + + private ChatClientRequest requestWithMaxTokens(int maxTokens) { + return new ChatClientRequest( + new Prompt( + "user question", + ChatOptions.builder().maxTokens(maxTokens).build() + ), + Map.of() + ); + } + + private ChatClientRequest requestWithoutMaxTokens() { + return new ChatClientRequest(new Prompt("user question"), Map.of()); + } +}