diff --git a/AGENTS.md b/AGENTS.md index 72e3a05..69b7c35 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, 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. The Spring AI 2.0 adapter now owns the supported non-streaming request lifecycle from text-scope adaptation and conservative preflight through reservation, provider invocation, and actual reconciliation. Real-provider compatibility 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. @@ -72,13 +72,13 @@ 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-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, around-call reconciliation, provider-specific reserved-output resolver extension, and enforcement streaming rejection implemented | | `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-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; 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-autoconfigure` | Basic implementation complete | Bean registration, property binding, pricing/budget/accounting/notification wiring, provider-specific reserved-output resolver injection, budget-disabled ledger-only compatibility, and `ChatClientBuilderCustomizer` implemented | | `token-pilot-starter` | Basic implementation complete | Thin final user entrypoint that brings runtime modules together | -| `token-pilot-sample-app` | Basic E2E complete | Direct ledger metrics, budget, and fake Spring AI advisor E2E implemented | +| `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 @@ -328,11 +328,15 @@ 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. +- 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, 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. +- 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. - 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. @@ -402,6 +406,13 @@ 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. + ### 2026-08-22 - Restricted usage-based reservation reconciliation to provider-reported or provider-derived usage so local and heuristic estimates cannot be committed as actual spend. diff --git a/token-pilot-autoconfigure/src/main/java/io/tokenpilot/autoconfigure/TokenPilotAutoConfiguration.java b/token-pilot-autoconfigure/src/main/java/io/tokenpilot/autoconfigure/TokenPilotAutoConfiguration.java index e3d2735..8d16fc5 100644 --- a/token-pilot-autoconfigure/src/main/java/io/tokenpilot/autoconfigure/TokenPilotAutoConfiguration.java +++ b/token-pilot-autoconfigure/src/main/java/io/tokenpilot/autoconfigure/TokenPilotAutoConfiguration.java @@ -3,17 +3,22 @@ import io.micrometer.core.instrument.MeterRegistry; import io.tokenpilot.budget.BudgetEvaluator; import io.tokenpilot.budget.BudgetStateStore; +import io.tokenpilot.budget.ReservationAccounting; import io.tokenpilot.budget.internal.LedgerBudgetComponents; 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.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.ReservedOutputTokensResolver; import io.tokenpilot.springai.UsageExtractor; import io.tokenpilot.springai.internal.LedgerSpringAiComponents; import org.springframework.ai.chat.client.ChatClient; @@ -81,6 +86,30 @@ public PricingEvaluator pricingEvaluator() { return LedgerComponents.defaultPricingEvaluator(); } + @Bean + @ConditionalOnMissingBean + public ModelRegistry modelRegistry() { + return LedgerComponents.defaultModelRegistry(); + } + + @Bean + @ConditionalOnMissingBean + public TokenEstimator tokenEstimator() { + return LedgerComponents.utf8ByteHeuristicTokenEstimator(); + } + + @Bean + @ConditionalOnMissingBean + public TokenBudget tokenBudget(ModelRegistry modelRegistry) { + return LedgerComponents.tokenBudget(modelRegistry); + } + + @Bean + @ConditionalOnMissingBean + public PreflightCostEstimator preflightCostEstimator() { + return LedgerComponents.defaultPreflightCostEstimator(); + } + /** * 비용 기록 및 리스너 관리를 담당하는 LedgerManager를 등록합니다. */ @@ -118,23 +147,36 @@ public LedgerAdvisor ledgerAdvisor( UsageExtractor usageExtractor, ObjectProvider budgetEvaluator, ObjectProvider budgetStateStore, + ObjectProvider reservationAccounting, CostCalculator costCalculator, PricingRegistry pricingRegistry, - PricingEvaluator pricingEvaluator + PricingEvaluator pricingEvaluator, + ModelRegistry modelRegistry, + TokenEstimator tokenEstimator, + TokenBudget tokenBudget, + PreflightCostEstimator preflightCostEstimator, + ObjectProvider reservedOutputTokensResolver, + TokenPilotProperties properties ) { BudgetEvaluator evaluator = budgetEvaluator.getIfAvailable(); BudgetStateStore stateStore = budgetStateStore.getIfAvailable(); + ReservationAccounting accounting = reservationAccounting.getIfAvailable(); - if (evaluator != null && stateStore != null) { - return LedgerSpringAiComponents.defaultLedgerAdvisor( - ledgerManager, + if (evaluator != null && stateStore != null && accounting != null) { + return LedgerSpringAiComponents.accountingLedgerAdvisor( usageExtractor, evaluator, stateStore, - costCalculator, + accounting, pricingRegistry, - pricingEvaluator, - MissingPricingPolicy.FAIL_CLOSED + modelRegistry, + tokenEstimator, + tokenBudget, + preflightCostEstimator, + properties.getSpringAi().getDefaultModelId(), + properties.getSpringAi().getDefaultReservedOutputTokens(), + reservedOutputTokensResolver.getIfAvailable(), + properties.getSpringAi().getFramingHeadroomTokens() ); } 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..92aff01 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 @@ -33,6 +33,9 @@ public class TokenPilotProperties { @NestedConfigurationProperty private NotificationProperties notification = new NotificationProperties(); + @NestedConfigurationProperty + private SpringAiProperties springAi = new SpringAiProperties(); + public boolean isEnabled() { return enabled; } @@ -73,6 +76,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(); @@ -223,4 +234,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/test/java/io/tokenpilot/autoconfigure/TokenPilotAutoConfigurationTest.java b/token-pilot-autoconfigure/src/test/java/io/tokenpilot/autoconfigure/TokenPilotAutoConfigurationTest.java index e1e89c0..7fa2261 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 @@ -11,9 +11,14 @@ import io.tokenpilot.budget.BudgetWindow; 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.budget.ReservationAccounting; import io.tokenpilot.core.domain.Cost; import io.tokenpilot.core.domain.PricingPlan; import io.tokenpilot.core.domain.PricingReconciliationResult; @@ -21,7 +26,6 @@ import io.tokenpilot.core.domain.PricingSnapshot; import io.tokenpilot.core.domain.TokenType; import io.tokenpilot.core.domain.TokenUsage; -import io.tokenpilot.core.exception.MissingPricingException; import io.tokenpilot.notification.BudgetNotificationHandler; import io.tokenpilot.notification.BudgetNotificationService; import io.tokenpilot.notification.NotificationStateStore; @@ -57,7 +61,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 { @@ -79,6 +82,10 @@ void shouldRegisterDefaultBeans() { assertThat(context).hasSingleBean(CostCalculator.class); assertThat(context).hasSingleBean(PricingEvaluator.class); assertThat(context).hasSingleBean(LedgerManager.class); + assertThat(context).hasSingleBean(ModelRegistry.class); + assertThat(context).hasSingleBean(TokenEstimator.class); + assertThat(context).hasSingleBean(TokenBudget.class); + assertThat(context).hasSingleBean(PreflightCostEstimator.class); assertThat(context).hasSingleBean(UsageExtractor.class); assertThat(context).hasSingleBean(LedgerAdvisor.class); @@ -86,6 +93,7 @@ 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); }); @@ -248,6 +256,7 @@ void shouldRegisterBudgetBeansWhenEnabled() { .run(context -> { assertThat(context).hasSingleBean(BudgetStateStore.class); assertThat(context).hasSingleBean(BudgetEvaluator.class); + assertThat(context).hasSingleBean(ReservationAccounting.class); }); } @@ -327,30 +336,6 @@ void shouldWireBudgetEvaluatorIntoLedgerAdvisorWhenBudgetEnabled() { }); } - @Test - @DisplayName("Budget가 활성화되면 missing pricing policy 기본값은 FAIL_CLOSED여야 한다") - void shouldUseFailClosedMissingPricingPolicyWhenBudgetEnabled() { - this.contextRunner - .withUserConfiguration(RecordingBudgetEvaluatorConfiguration.class) - .withPropertyValues("token-pilot.budget.enabled=true") - .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" - ) - ); - - assertThatThrownBy(() -> advisor.before(request, mock(AdvisorChain.class))) - .isInstanceOf(MissingPricingException.class) - .hasMessage("MISSING_PLAN") - .extracting(exception -> ((MissingPricingException) exception).getResolution()) - .isEqualTo(PricingResolution.MISSING_PLAN); - }); - } - @Test @DisplayName("token-pilot.budget.enabled=false 일 때 Budget 관련 빈이 등록되지 않아야 한다") void shouldNotRegisterBudgetBeansWhenDisabled() { 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..98691c1 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,791 @@ package io.tokenpilot.sample; +import io.tokenpilot.budget.ActualUsageCommand; import io.tokenpilot.budget.BudgetDecision; +import io.tokenpilot.budget.BudgetEvaluator; +import io.tokenpilot.budget.BudgetReservationRequest; +import io.tokenpilot.budget.BudgetSnapshot; import io.tokenpilot.budget.BudgetStateStore; +import io.tokenpilot.budget.ReservationAccounting; +import io.tokenpilot.budget.ReservationAccountingEvent; +import io.tokenpilot.budget.ReservationAccountingReason; +import io.tokenpilot.budget.ReservationId; +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.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.ChatClientResponse; +import org.springframework.ai.chat.client.ChatClientRequest; +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" +}) 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; + + @Autowired + private AccountingCostProbe accountingCostProbe; - 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\""); + @BeforeEach + void resetProbes() { + providerProbe.reset(); + usageExtractorProbe.reset(); + accountingListenerProbe.reset(); + accountingCostProbe.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") ); - BudgetDecision budgetDecision = contextValue(response, BudgetDecision.class); - Cost accumulatedCost = budgetStateStore.getAccumulatedCost( - budgetDecision.key(), - budgetDecision.limit() + BudgetSnapshot snapshot = stateStoreProbe.snapshot( + decision.key(), + decision.limit() ); + ReservationAccounting accountingProbe = (ReservationAccounting) stateStoreProbe; - 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")); + 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) + ); + 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(); } - private T contextValue(ChatClientResponse response, Class type) { - return response.context().values().stream() - .filter(type::isInstance) - .map(type::cast) - .findFirst() - .orElseThrow(); + @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 HttpResponse get(String path) throws IOException, InterruptedException { - HttpRequest request = HttpRequest.newBuilder() - .uri(URI.create("http://localhost:" + port + path)) - .GET() + @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); + } + + @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 요청은 provider를 호출하거나 새 예약을 만들지 않는다") + void budgetBlockStopsBeforeProviderAndReservation() { + 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()); + 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, never()).checkAndReserve(any(BudgetReservationRequest.class)); + 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(IllegalStateException.class) + .hasMessageContaining("PRICING_NOT_FOUND"); + + 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("request model과 response model은 같은 reservation 정산 명령에 보존된다") + void requestAndResponseModelsRemainInTheSameAccountingLifecycle() { + providerProbe.respondAs("provider-routed-model-v2"); + + 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).commit(argThat( + command -> command.responseModelId().equals("provider-routed-model-v2") + && command.requestId().equals("request-model-correlation") + && command.attemptId().equals("attempt-model-correlation") + )); + } + + @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 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(); - } - }; + BudgetStateStore accountingProbe( + AccountingListenerProbe listenerProbe, + AccountingCostProbe costProbe + ) { + BudgetStateStore delegate = LedgerBudgetComponents.inMemoryBudgetStateStore( + Clock.systemUTC(), + () -> new ReservationId(UUID.randomUUID().toString()), + costProbe::calculate, + List.of(listenerProbe::onCommitted) + ); + return mock( + BudgetStateStore.class, + withSettings() + .extraInterfaces(ReservationAccounting.class) + .defaultAnswer(delegatesTo(delegate)) + ); + } + + @Bean + UsageExtractorProbe usageExtractor() { + return new UsageExtractorProbe(); + } + + @Bean + AccountingListenerProbe accountingListenerProbe() { + return new AccountingListenerProbe(); + } + + @Bean + AccountingCostProbe accountingCostProbe() { + return new AccountingCostProbe(); + } + + @Bean + ProviderProbe fakeChatModel() { + return new ProviderProbe(); } @Bean @@ -176,32 +798,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/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/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..63e37a4 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,7 +24,10 @@ 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.Map; @@ -49,6 +63,11 @@ 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(); public DefaultLedgerAdvisor(LedgerManager ledgerManager, UsageExtractor usageExtractor) { this(ledgerManager, usageExtractor, null, null, null, null); @@ -89,8 +108,41 @@ 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 + ); + } + + 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 + ) { 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 +155,282 @@ public DefaultLedgerAdvisor(LedgerManager ledgerManager, UsageExtractor usageExt missingPricingPolicy, "missingPricingPolicy 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 + ); + } + + @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 = requestPreflight.resolve( + correlatedRequest + ); + BudgetDecision decision = budgetEvaluator.evaluate( + extractTags(correlatedRequest.context()), + costBound.safeUpperBoundCost() + ); + enforceAdmission(decision); + + ReservationId reservationId = reserve( + requestId, + idempotencyKey, + costBound, + decision + ); + ChatClientRequest providerRequest = prepareDispatch( + correlatedRequest, + reservationId + ); + + ChatClientResponse response; + try { + response = chain.nextCall(providerRequest); + } catch (RuntimeException downstreamFailure) { + preservePendingLiability(reservationId, downstreamFailure); + throw downstreamFailure; + } + settle(response, requestId, attemptId, reservationId); + return response; + } + + private void enforceAdmission(BudgetDecision decision) { + if (!decision.isAdmissionDecision()) { + throw new IllegalStateException("budget admission decision is required"); + } + switch (decision.state()) { + case ALLOW, WARN -> { + return; + } + case BLOCK -> throw new BudgetExceededException(decision); + case CURRENCY_MISMATCH -> throw new IllegalStateException( + "Budget decision currency mismatch: " + decision.reason() + ); + } + } + + private ReservationId reserve( + 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() + ) + ); + BudgetReservationResult result = budgetStateStore.checkAndReserve( + reservation + ); + if (!result.isAccepted()) { + throw new IllegalStateException( + "budget reservation rejected: " + result.reason() + ); + } + return result.reservation().id(); + } + + private ChatClientRequest prepareDispatch( + ChatClientRequest request, + ReservationId reservationId + ) { + try { + ChatClientRequest providerRequest = contextAccessor.withReservationId( + request, + reservationId + ); + ReservationTransition transition = reservationAccounting.markInFlight( + reservationId + ); + if (transition.resultingState() != ReservationState.IN_FLIGHT) { + throw new IllegalStateException( + "reservation did not enter IN_FLIGHT: " + transition.status() + ); + } + return providerRequest; + } catch (RuntimeException failure) { + releaseBeforeDispatch(reservationId, failure); + throw failure; + } + } + + 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 + ) { + TokenUsage usage; + try { + usage = usageExtractor.extract(response); + } catch (RuntimeException ignored) { + markReconciliationRequiredAfterResponse(reservationId); + return; + } + if (usage.source() == UsageSource.UNAVAILABLE) { + markReconciliationRequiredAfterResponse(reservationId); + return; + } + try { + reservationAccounting.commit(new ActualUsageCommand( + requestId, + attemptId, + reservationId, + usage, + requireResponseModelId(response) + )); + } catch (RuntimeException ignored) { + markReconciliationRequiredAfterResponse(reservationId); + } + } + + private void markReconciliationRequired(ReservationId reservationId) { + reservationAccounting.markReconciliationRequired( + reservationId, + ReservationAccountingReason.ACTUAL_USAGE_UNAVAILABLE + ); + } + + private void markReconciliationRequiredAfterResponse( + ReservationId reservationId + ) { + try { + markReconciliationRequired(reservationId); + } catch (RuntimeException ignored) { + // Provider response is preserved. Failure observation belongs to #40. + } + } + + 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 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..698dc74 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,24 @@ 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.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; /** * Spring AI 어댑터 컴포넌트 생성을 위한 팩토리 클래스입니다. @@ -23,6 +33,80 @@ 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 + ) { + ReservedOutputResolver outputResolver = new ReservedOutputResolver( + providerOutputResolver, + defaultReservedOutputTokens + ); + RequestContextAccessor contextAccessor = new RequestContextAccessor(); + RequestPreflight preflight = new RequestPreflight( + new ModelResolver(modelRegistry, defaultModelId), + outputResolver, + tokenEstimator, + tokenBudget, + pricingRegistry, + costEstimator, + framingHeadroomTokens + ); + return new DefaultLedgerAdvisor( + usageExtractor, + budgetEvaluator, + budgetStateStore, + reservationAccounting, + preflight, + contextAccessor, + new IdempotencyKeyResolver( + contextAccessor, + () -> new IdempotencyKey(UUID.randomUUID().toString()) + ) + ); + } + public static LedgerAdvisor defaultLedgerAdvisor( LedgerManager ledgerManager, UsageExtractor usageExtractor 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..709b905 --- /dev/null +++ b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/RequestPreflight.java @@ -0,0 +1,127 @@ +package io.tokenpilot.springai.internal; + +import io.tokenpilot.core.PreflightCostEstimator; +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.TokenCountAccuracy; +import io.tokenpilot.core.domain.TokenCountResult; +import io.tokenpilot.core.domain.TokenCountScope; +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 PreflightCostEstimator costEstimator; + private final long framingHeadroomTokens; + + RequestPreflight( + ModelResolver modelResolver, + ReservedOutputResolver outputResolver, + TokenEstimator tokenEstimator, + TokenBudget tokenBudget, + PricingRegistry pricingRegistry, + 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.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 + ); + + PreflightPricingContext pricing = new PreflightPricingContext( + model.canonicalModelId(), + model.pricingPolicyId(), + model.catalogVersion(), + model.acceptedCompatibilityBasis(), + model.pricingCurrency(), + PreflightPricingContext.UpperBoundCapability.FINITE, + pricingRegistry.resolveSnapshot(model) + ); + 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..4cbd083 --- /dev/null +++ b/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/DefaultLedgerAdvisorLifecycleTest.java @@ -0,0 +1,322 @@ +package io.tokenpilot.springai.internal; + +import io.tokenpilot.budget.AccountingTransitionStatus; +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.ReservationTransition; +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.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 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.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("did not enter IN_FLIGHT"); + + verify(accounting).releaseBeforeDispatch(RESERVATION_ID); + verifyNoInteractions(provider, usageExtractor); + } + + @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 + ); + } + + 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.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(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() { + ChatResponse chatResponse = mock(ChatResponse.class); + ChatResponseMetadata metadata = mock(ChatResponseMetadata.class); + when(chatResponse.getMetadata()).thenReturn(metadata); + when(metadata.getModel()).thenReturn("model-v1"); + 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..e318862 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 @@ -27,6 +27,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; @@ -1010,7 +1011,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/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/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()); + } +}