diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..5b0a6c1 --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,128 @@ +# Contributing to budget-breaker + +Thanks for considering a contribution! This guide covers everything you need to get started. + +## Prerequisites + +- JDK 21+ (we recommend [Eclipse Temurin](https://adoptium.net/)) +- Kotlin 2.0+ (managed by Gradle -- no separate install needed) + +## Dev Setup + +```bash +git clone git@github.com:UnityInFlow/budget-breaker.git +cd budget-breaker + +# Build everything +./gradlew build + +# Run tests +./gradlew test + +# Format code (must pass before committing) +./gradlew ktlintFormat +``` + +## Project Structure + +``` +budget-breaker/ # Core library (pure Kotlin + coroutines) + src/main/kotlin/dev/unityinflow/budget/ + AgentBudget.kt # Budget configuration data class + TokenTracker.kt # Thread-safe AtomicLong token counter + BudgetCircuitBreaker.kt # Coroutine supervisor with withBudget() API + BudgetScope.kt # DSL scope for tracking calls + BudgetException.kt # Sealed exception hierarchy + BudgetEvent.kt # SharedFlow event types + ModelPricing.kt # LLM cost estimation + BudgetReport.kt # Post-run summary + +budget-breaker-spring-boot-starter/ # Spring Boot auto-config (Phase 2) +``` + +## How to Add Model Pricing + +Default pricing lives in `ModelPricing.kt` in the companion object's `DEFAULTS` map. + +To add a new model: + +1. Add an entry to the `DEFAULTS` map in `ModelPricing.kt`: + ```kotlin + "model-name" to PriceConfig(inputPerMillion = 2.0, outputPerMillion = 8.0), + ``` +2. Add a test in `ModelPricingTest.kt` verifying the cost estimate. +3. Update the pricing table in `README.md`. + +Users can also provide custom pricing at runtime without modifying the library: + +```kotlin +val pricing = ModelPricing( + overrides = mapOf( + "my-custom-model" to ModelPricing.PriceConfig(1.0, 5.0) + ) +) +val breaker = BudgetCircuitBreaker(pricing = pricing) +``` + +## How to Extend + +### Adding a New BudgetEvent Type + +1. Add a new subclass to the `BudgetEvent` sealed class in `BudgetEvent.kt`. +2. Emit it from `BudgetScope` at the appropriate point. +3. Add tests in `BudgetCircuitBreakerTest.kt` verifying the event is emitted. + +### Adding a New Exception Type + +1. Add a new subclass to the `BudgetException` sealed class in `BudgetException.kt`. +2. Handle it in `BudgetScope.trackCall()` or `BudgetCircuitBreaker.withBudget()`. +3. Add tests covering the throw + catch path. + +## Commit Conventions + +We follow [Conventional Commits](https://www.conventionalcommits.org/): + +``` +feat: add Mistral model pricing defaults +fix: TokenTracker overflow on Long.MAX_VALUE +test: add edge cases for zero-token budgets +docs: update README with new model pricing +chore: bump coroutines to 1.10.2 +refactor: extract limit checking into separate function +``` + +## Branch Naming + +``` +feat/add-mistral-pricing +fix/tracker-overflow +docs/contributing-guide +``` + +## Code Style + +- Kotlin idioms first (`let`, `also`, `apply`, `run`) +- `val` only -- no `var` +- No `!!` without a comment explaining safety +- Coroutines for all async work -- never `Thread.sleep()` +- KDoc on all public APIs +- `ktlint` must pass before every commit + +## Testing + +- JUnit 5 + Kotest matchers +- Every public API: at least 3 passing and 3 failing cases +- Edge cases: zero tokens, Long.MAX_VALUE, concurrent access +- Run with `./gradlew test` + +## Pull Requests + +1. Fork and create a feature branch from `main`. +2. Write tests first (TDD preferred). +3. Ensure `./gradlew build` and `./gradlew test` pass. +4. Open a PR against `main` with a clear description. +5. One approval required to merge. + +## License + +By contributing, you agree that your contributions will be licensed under the [MIT License](LICENSE). diff --git a/README.md b/README.md index 2b6cc6a..e41cb46 100644 --- a/README.md +++ b/README.md @@ -1,54 +1,266 @@ # budget-breaker -Kotlin coroutine budget circuit breaker for AI agents — soft/hard token limits with cost estimation. +Kotlin coroutine budget circuit breaker for AI agents -- soft/hard token limits with cost estimation. **Tool 05** in the [UnityInFlow](https://github.com/UnityInFlow) ecosystem. ## Problem -AI agents calling LLMs can silently burn through token budgets. Without a circuit breaker, -a runaway agent loop can consume thousands of dollars in minutes. `budget-breaker` provides -coroutine-aware soft and hard token limits with clean cancellation semantics. +Overnight agent runs routinely cost $50--$200 when a loop goes unchecked. Existing solutions are either polling-based (too slow), framework-coupled (Spring/Ktor only), or lack coroutine awareness (can't cancel a structured scope cleanly). -## Features - -- **Soft limits** — callback + SharedFlow event when threshold is reached; execution continues -- **Hard limits** — coroutine scope cancelled via `BudgetHardLimitException` -- **Thread-safe** — AtomicLong counters for lock-free concurrent tracking -- **Cost estimation** — built-in pricing for Claude, GPT-4o, Gemini models -- **Zero Spring dependency** — core library is pure Kotlin + coroutines -- **Spring Boot starter** — optional auto-configuration (Phase 2) +No reactive budget enforcement exists for Kotlin coroutines. `budget-breaker` fixes this: it wraps your agent code in a budget-tracked scope, fires a callback at your soft limit, and cancels the coroutine at the hard limit. Zero Spring dependency in the core module. ## Installation +### Gradle (Kotlin DSL) + ```kotlin -// build.gradle.kts dependencies { - implementation("dev.unityinflow:budget-breaker:0.0.1-SNAPSHOT") + implementation("dev.unityinflow:budget-breaker:0.0.1") } ``` -## Quick Start +### Gradle (Groovy) + +```groovy +dependencies { + implementation 'dev.unityinflow:budget-breaker:0.0.1' +} +``` + +### Maven + +```xml + + dev.unityinflow + budget-breaker + 0.0.1 + +``` + +## Usage + +### Basic `withBudget` ```kotlin +import dev.unityinflow.budget.* + val budget = AgentBudget( model = "claude-sonnet-4-6", hardLimitTokens = 100_000, softLimitTokens = 80_000, ) -val breaker = BudgetCircuitBreaker(agentId = "my-agent", budget = budget) +val breaker = BudgetCircuitBreaker() + +suspend fun main() { + try { + breaker.withBudget(agentId = "research-agent", budget = budget) { + // Simulate LLM calls inside the budget scope + trackCall(promptTokens = 500, completionTokens = 200) + trackCall(promptTokens = 1_000, completionTokens = 400) + // ...your agent logic here + } + } catch (e: BudgetHardLimitException) { + println("Budget exceeded: ${e.message}") + } + + // Get the post-run report + val report = breaker.getReport("research-agent") + println("Total tokens: ${report?.totalTokens}") + println("Estimated cost: \$${report?.estimatedCostUsd}") +} +``` + +### With Soft Limit Callback + +When the soft limit is reached, your callback fires but execution continues. Use this for alerts, logging, or graceful wind-down. + +```kotlin +val breaker = BudgetCircuitBreaker( + onSoftLimit = { report -> + println("WARNING: Agent '${report.agentId}' hit soft limit") + println(" Tokens used: ${report.totalTokens}") + println(" Estimated cost: \$${report.estimatedCostUsd}") + // Send a Slack alert, write to a log, etc. + } +) + +breaker.withBudget(agentId = "summarizer", budget = budget) { + repeat(100) { + trackCall(promptTokens = 1_000, completionTokens = 500) + // Callback fires once when soft limit (80k) is crossed + // Execution continues until hard limit (100k) + } +} +``` + +### With SharedFlow Events + +For reactive consumers, `BudgetCircuitBreaker` exposes a `SharedFlow` that emits every call tracked, soft limit reached, and hard limit exceeded. -breaker.withBudget { - // Your agent code here - trackCall(promptTokens = 500, completionTokens = 200) +```kotlin +import kotlinx.coroutines.* +import kotlinx.coroutines.flow.collect + +val breaker = BudgetCircuitBreaker() + +fun main() = runBlocking { + // Collect events in a background coroutine + val collector = launch { + breaker.events.collect { event -> + when (event) { + is BudgetEvent.CallTracked -> println("Call: +${event.promptTokens}p/${event.completionTokens}c") + is BudgetEvent.SoftLimitReached -> println("SOFT LIMIT: ${event.percentUsed}% used") + is BudgetEvent.HardLimitExceeded -> println("HARD LIMIT: \$${event.estimatedCostUsd}") + } + } + } + + try { + breaker.withBudget(agentId = "streaming-agent", budget = budget) { + repeat(200) { + trackCall(promptTokens = 500, completionTokens = 200) + } + } + } catch (_: BudgetHardLimitException) { + // Expected + } + + collector.cancel() } ``` -## Status +## Model Pricing + +Built-in pricing defaults (per million tokens, USD): + +| Model | Input | Output | +|---|---:|---:| +| `claude-opus-4-6` | $15.00 | $75.00 | +| `claude-sonnet-4-6` | $3.00 | $15.00 | +| `claude-haiku-4-5` | $0.80 | $4.00 | +| `gpt-4o` | $2.50 | $10.00 | +| `gpt-4o-mini` | $0.15 | $0.60 | +| `gemini-2.5-pro` | $1.25 | $10.00 | +| `gemini-2.5-flash` | $0.15 | $0.60 | +| `ollama` | $0.00 | $0.00 | + +### Custom Pricing + +Override or add models at construction time: + +```kotlin +val pricing = ModelPricing( + overrides = mapOf( + "mistral-large" to ModelPricing.PriceConfig( + inputPerMillion = 2.0, + outputPerMillion = 6.0, + ) + ) +) + +val breaker = BudgetCircuitBreaker(pricing = pricing) +``` + +## API Reference + +### `AgentBudget` + +Configuration data class for an agent's token budget. + +| Property | Type | Default | Description | +|---|---|---|---| +| `model` | `String` | `"claude-sonnet-4-6"` | LLM model identifier (for cost estimation) | +| `hardLimitTokens` | `Long` | `100_000` | Exceeding this cancels the coroutine scope | +| `softLimitTokens` | `Long` | `80_000` | Triggers callback and Flow event when reached | + +### `BudgetCircuitBreaker` + +Main entry point. Wraps agent code in a budget-tracked scope. + +| Method | Description | +|---|---| +| `withBudget(agentId, budget, block)` | Execute `block` within a budget-tracked `BudgetScope` | +| `getReport(agentId)` | Get the `BudgetReport` for a completed agent run | +| `events` | `SharedFlow` for reactive consumers | + +### `BudgetScope` + +DSL scope available inside `withBudget`. Call `trackCall()` after each LLM invocation. + +| Method | Description | +|---|---| +| `trackCall(promptTokens, completionTokens)` | Record token usage; checks limits after recording | + +### `TokenTracker` + +Thread-safe token counter using `AtomicLong`. + +| Property / Method | Description | +|---|---| +| `promptTokens` | Total prompt tokens recorded | +| `completionTokens` | Total completion tokens recorded | +| `totalTokens` | Sum of prompt + completion | +| `percentUsed()` | Usage as percentage of hard limit | + +### `BudgetException` (sealed) + +| Subclass | When | Fatal? | +|---|---|---| +| `BudgetSoftLimitException` | Soft limit reached | No | +| `BudgetHardLimitException` | Hard limit exceeded | Yes -- cancels scope | + +### `BudgetEvent` (sealed) + +| Subclass | When | +|---|---| +| `CallTracked` | Every `trackCall()` invocation | +| `SoftLimitReached` | First time soft limit is crossed | +| `HardLimitExceeded` | Hard limit exceeded (before exception) | + +### `BudgetReport` + +Post-run summary returned by `getReport()`. + +| Field | Type | Description | +|---|---|---| +| `agentId` | `String` | Agent identifier | +| `model` | `String` | LLM model used | +| `promptTokens` | `Long` | Total prompt tokens | +| `completionTokens` | `Long` | Total completion tokens | +| `totalTokens` | `Long` | Sum of prompt + completion | +| `estimatedCostUsd` | `Double` | Estimated cost in USD | +| `softLimitBreachCount` | `Int` | Number of soft limit breaches | +| `hardLimitBreached` | `Boolean` | Whether hard limit was hit | +| `durationMs` | `Long` | Execution duration in milliseconds | +| `percentUsed` | `Double` | Usage as percentage of hard limit | + +### `ModelPricing` + +Cost estimation with hardcoded defaults and runtime overrides. + +| Method | Description | +|---|---| +| `estimateCost(model, promptTokens, completionTokens)` | Estimate cost in USD | +| `ModelPricing.estimateCost(...)` (companion) | Static convenience using default pricing | + +## Spring Boot Integration + +Coming in v0.1.0: + +- `@ConfigurationProperties`-based auto-configuration +- `/actuator/budget` endpoint +- Micrometer metrics integration (counters, gauges, timers) + +The core library has zero Spring dependency -- it works in any Kotlin project. + +## Requirements -v0.0.1-SNAPSHOT — Phase 1 (core library) in progress. +- JDK 21+ +- Kotlin 2.0+ +- `kotlinx-coroutines-core` 1.10+ ## License -MIT +[MIT](LICENSE) diff --git a/budget-breaker/src/main/kotlin/dev/unityinflow/budget/BudgetCircuitBreaker.kt b/budget-breaker/src/main/kotlin/dev/unityinflow/budget/BudgetCircuitBreaker.kt new file mode 100644 index 0000000..433b453 --- /dev/null +++ b/budget-breaker/src/main/kotlin/dev/unityinflow/budget/BudgetCircuitBreaker.kt @@ -0,0 +1,55 @@ +package dev.unityinflow.budget + +import kotlinx.coroutines.coroutineScope +import kotlinx.coroutines.flow.MutableSharedFlow +import kotlinx.coroutines.flow.SharedFlow +import kotlinx.coroutines.flow.asSharedFlow +import java.util.concurrent.ConcurrentHashMap + +/** + * Coroutine-aware circuit breaker for AI agent token budgets. + * + * Wraps agent code in a budget-tracked scope. Soft limit triggers callback + Flow event. + * Hard limit cancels the coroutine scope. + */ +class BudgetCircuitBreaker( + private val defaultBudget: AgentBudget = AgentBudget(), + private val pricing: ModelPricing = ModelPricing(), + private val onSoftLimit: ((BudgetReport) -> Unit)? = null, +) { + private val reports = ConcurrentHashMap() + private val _events = MutableSharedFlow(extraBufferCapacity = 64) + + /** SharedFlow of budget events for reactive consumers. */ + val events: SharedFlow = _events.asSharedFlow() + + /** + * Execute [block] within a budget-tracked scope. + * Use [BudgetScope.trackCall] inside the block to record token usage. + * + * @throws BudgetHardLimitException if hard limit is exceeded + */ + suspend fun withBudget( + agentId: String, + budget: AgentBudget = defaultBudget, + block: suspend BudgetScope.() -> T, + ): T { + val tracker = TokenTracker(agentId, budget) + val scope = BudgetScope(tracker, pricing, onSoftLimit, _events) + val startTime = System.currentTimeMillis() + + return try { + coroutineScope { + scope.block() + } + } catch (e: BudgetHardLimitException) { + throw e + } finally { + val durationMs = System.currentTimeMillis() - startTime + reports[agentId] = scope.buildReport(durationMs) + } + } + + /** Get the budget report for a completed agent run. */ + fun getReport(agentId: String): BudgetReport? = reports[agentId] +} diff --git a/budget-breaker/src/main/kotlin/dev/unityinflow/budget/BudgetReport.kt b/budget-breaker/src/main/kotlin/dev/unityinflow/budget/BudgetReport.kt new file mode 100644 index 0000000..9bb091a --- /dev/null +++ b/budget-breaker/src/main/kotlin/dev/unityinflow/budget/BudgetReport.kt @@ -0,0 +1,17 @@ +package dev.unityinflow.budget + +/** + * Summary of an agent's budget usage after a run completes. + */ +data class BudgetReport( + val agentId: String, + val model: String, + val promptTokens: Long, + val completionTokens: Long, + val totalTokens: Long, + val estimatedCostUsd: Double, + val softLimitBreachCount: Int, + val hardLimitBreached: Boolean, + val durationMs: Long, + val percentUsed: Double, +) diff --git a/budget-breaker/src/main/kotlin/dev/unityinflow/budget/BudgetScope.kt b/budget-breaker/src/main/kotlin/dev/unityinflow/budget/BudgetScope.kt new file mode 100644 index 0000000..6195fcb --- /dev/null +++ b/budget-breaker/src/main/kotlin/dev/unityinflow/budget/BudgetScope.kt @@ -0,0 +1,87 @@ +package dev.unityinflow.budget + +import kotlinx.coroutines.flow.MutableSharedFlow + +/** + * Scope DSL for tracking LLM calls within a budget. + * Call [trackCall] after each LLM invocation to record token usage. + */ +class BudgetScope internal constructor( + private val tracker: TokenTracker, + private val pricing: ModelPricing, + private val onSoftLimit: ((BudgetReport) -> Unit)?, + private val eventFlow: MutableSharedFlow, +) { + private val softLimitBreachCount = java.util.concurrent.atomic.AtomicInteger(0) + private val softLimitFired = java.util.concurrent.atomic.AtomicBoolean(false) + + /** + * Record token usage from an LLM call. + * Checks soft/hard limits after recording. + * + * @throws BudgetHardLimitException if hard limit is exceeded + */ + suspend fun trackCall(promptTokens: Long, completionTokens: Long) { + tracker.add(promptTokens, completionTokens) + + eventFlow.emit( + BudgetEvent.CallTracked( + agentId = tracker.agentId, + tokensUsed = tracker.totalTokens, + promptTokens = promptTokens, + completionTokens = completionTokens, + ) + ) + + if (tracker.isAboveHardLimit()) { + val cost = pricing.estimateCost( + tracker.model, tracker.promptTokens, tracker.completionTokens + ) + val event = BudgetEvent.HardLimitExceeded( + agentId = tracker.agentId, + tokensUsed = tracker.totalTokens, + budgetTokens = tracker.hardLimitTokens, + estimatedCostUsd = cost, + ) + eventFlow.emit(event) + throw BudgetHardLimitException( + agentId = tracker.agentId, + tokensUsed = tracker.totalTokens, + budgetTokens = tracker.hardLimitTokens, + estimatedCostUsd = cost, + ) + } + + if (tracker.isAboveSoftLimit() && !softLimitFired.getAndSet(true)) { + softLimitBreachCount.incrementAndGet() + val report = buildReport(0) + onSoftLimit?.invoke(report) + eventFlow.emit( + BudgetEvent.SoftLimitReached( + agentId = tracker.agentId, + tokensUsed = tracker.totalTokens, + budgetTokens = tracker.softLimitTokens, + percentUsed = tracker.percentUsed(), + ) + ) + } + } + + internal fun buildReport(durationMs: Long): BudgetReport { + val cost = pricing.estimateCost( + tracker.model, tracker.promptTokens, tracker.completionTokens + ) + return BudgetReport( + agentId = tracker.agentId, + model = tracker.model, + promptTokens = tracker.promptTokens, + completionTokens = tracker.completionTokens, + totalTokens = tracker.totalTokens, + estimatedCostUsd = cost, + softLimitBreachCount = softLimitBreachCount.get(), + hardLimitBreached = tracker.isAboveHardLimit(), + durationMs = durationMs, + percentUsed = tracker.percentUsed(), + ) + } +} diff --git a/budget-breaker/src/main/kotlin/dev/unityinflow/budget/ModelPricing.kt b/budget-breaker/src/main/kotlin/dev/unityinflow/budget/ModelPricing.kt new file mode 100644 index 0000000..5abc3e6 --- /dev/null +++ b/budget-breaker/src/main/kotlin/dev/unityinflow/budget/ModelPricing.kt @@ -0,0 +1,43 @@ +package dev.unityinflow.budget + +/** + * Estimates LLM API costs based on token usage. + * Has hardcoded defaults for popular models, supports custom overrides. + */ +class ModelPricing( + private val overrides: Map = emptyMap(), +) { + data class PriceConfig( + val inputPerMillion: Double, + val outputPerMillion: Double, + ) + + /** Estimate cost in USD for the given token usage. */ + fun estimateCost(model: String, promptTokens: Long, completionTokens: Long): Double { + val config = overrides[model] ?: DEFAULTS[model] ?: return 0.0 + val inputCost = (promptTokens.toDouble() / 1_000_000) * config.inputPerMillion + val outputCost = (completionTokens.toDouble() / 1_000_000) * config.outputPerMillion + return inputCost + outputCost + } + + companion object { + private val DEFAULTS = mapOf( + // Claude (April 2026) + "claude-opus-4-6" to PriceConfig(15.0, 75.0), + "claude-sonnet-4-6" to PriceConfig(3.0, 15.0), + "claude-haiku-4-5" to PriceConfig(0.80, 4.0), + // OpenAI + "gpt-4o" to PriceConfig(2.50, 10.0), + "gpt-4o-mini" to PriceConfig(0.15, 0.60), + // Google + "gemini-2.5-pro" to PriceConfig(1.25, 10.0), + "gemini-2.5-flash" to PriceConfig(0.15, 0.60), + // Local (free) + "ollama" to PriceConfig(0.0, 0.0), + ) + + /** Convenience: estimate cost using default pricing (no overrides). */ + fun estimateCost(model: String, promptTokens: Long, completionTokens: Long): Double = + ModelPricing().estimateCost(model, promptTokens, completionTokens) + } +} diff --git a/budget-breaker/src/main/kotlin/dev/unityinflow/budget/TokenTracker.kt b/budget-breaker/src/main/kotlin/dev/unityinflow/budget/TokenTracker.kt new file mode 100644 index 0000000..1dcd244 --- /dev/null +++ b/budget-breaker/src/main/kotlin/dev/unityinflow/budget/TokenTracker.kt @@ -0,0 +1,43 @@ +package dev.unityinflow.budget + +import java.util.concurrent.atomic.AtomicLong + +/** + * Thread-safe token counter with soft/hard limit checks. + * Uses AtomicLong for lock-free concurrent access. + */ +class TokenTracker( + val agentId: String, + private val budget: AgentBudget, +) { + /** The model identifier from the budget configuration. */ + val model: String get() = budget.model + + /** The configured hard limit in tokens. */ + val hardLimitTokens: Long get() = budget.hardLimitTokens + + /** The configured soft limit in tokens. */ + val softLimitTokens: Long get() = budget.softLimitTokens + + private val _promptTokens = AtomicLong(0) + private val _completionTokens = AtomicLong(0) + + val promptTokens: Long get() = _promptTokens.get() + val completionTokens: Long get() = _completionTokens.get() + val totalTokens: Long get() = _promptTokens.get() + _completionTokens.get() + + /** Record token usage from an LLM call. */ + fun add(promptTokens: Long, completionTokens: Long) { + _promptTokens.addAndGet(promptTokens) + _completionTokens.addAndGet(completionTokens) + } + + /** Check if total tokens exceed the soft limit. */ + fun isAboveSoftLimit(): Boolean = totalTokens >= budget.softLimitTokens + + /** Check if total tokens exceed the hard limit. */ + fun isAboveHardLimit(): Boolean = totalTokens >= budget.hardLimitTokens + + /** Current usage as a percentage of the hard limit. */ + fun percentUsed(): Double = (totalTokens.toDouble() / budget.hardLimitTokens) * 100.0 +} diff --git a/budget-breaker/src/test/kotlin/dev/unityinflow/budget/BudgetCircuitBreakerTest.kt b/budget-breaker/src/test/kotlin/dev/unityinflow/budget/BudgetCircuitBreakerTest.kt new file mode 100644 index 0000000..d993014 --- /dev/null +++ b/budget-breaker/src/test/kotlin/dev/unityinflow/budget/BudgetCircuitBreakerTest.kt @@ -0,0 +1,148 @@ +package dev.unityinflow.budget + +import io.kotest.assertions.throwables.shouldThrow +import io.kotest.matchers.doubles.shouldBeGreaterThan +import io.kotest.matchers.shouldBe +import io.kotest.matchers.shouldNotBe +import kotlinx.coroutines.launch +import kotlinx.coroutines.test.runTest +import org.junit.jupiter.api.Test +import java.util.concurrent.atomic.AtomicBoolean + +class BudgetCircuitBreakerTest { + + private val budget = AgentBudget( + model = "claude-sonnet-4-6", + hardLimitTokens = 1000, + softLimitTokens = 800, + ) + + @Test + fun `completes normally within budget`() = runTest { + val breaker = BudgetCircuitBreaker(budget) + + val result = breaker.withBudget("agent-1") { + trackCall(promptTokens = 100, completionTokens = 50) + "done" + } + + result shouldBe "done" + } + + @Test + fun `throws on hard limit breach`() = runTest { + val breaker = BudgetCircuitBreaker(budget) + + shouldThrow { + breaker.withBudget("agent-1") { + trackCall(promptTokens = 600, completionTokens = 500) + } + } + } + + @Test + fun `fires callback on soft limit`() = runTest { + val callbackFired = AtomicBoolean(false) + val breaker = BudgetCircuitBreaker( + budget, + onSoftLimit = { callbackFired.set(true) }, + ) + + breaker.withBudget("agent-1") { + trackCall(promptTokens = 500, completionTokens = 350) + } + + callbackFired.get() shouldBe true + } + + @Test + fun `emits event on soft limit`() = runTest { + val breaker = BudgetCircuitBreaker(budget) + val collected = mutableListOf() + + val eventJob = launch { + breaker.events.collect { collected.add(it) } + } + + // Yield to let the collector subscribe before emitting + kotlinx.coroutines.yield() + + breaker.withBudget("agent-1") { + trackCall(promptTokens = 500, completionTokens = 350) + } + + // Yield to let collector process buffered events + kotlinx.coroutines.yield() + eventJob.cancel() + + val softEvent = collected.filterIsInstance() + softEvent.size shouldBe 1 + softEvent.first().agentId shouldBe "agent-1" + } + + @Test + fun `generates report after completion`() = runTest { + val breaker = BudgetCircuitBreaker(budget) + + breaker.withBudget("agent-1") { + trackCall(promptTokens = 200, completionTokens = 100) + } + + val report = breaker.getReport("agent-1") + report shouldNotBe null + report?.totalTokens shouldBe 300 + report?.hardLimitBreached shouldBe false + } + + @Test + fun `generates report even after hard limit exception`() = runTest { + val breaker = BudgetCircuitBreaker(budget) + + try { + breaker.withBudget("agent-1") { + trackCall(promptTokens = 600, completionTokens = 500) + } + } catch (_: BudgetHardLimitException) { + // expected + } + + val report = breaker.getReport("agent-1") + report shouldNotBe null + report?.hardLimitBreached shouldBe true + } + + @Test + fun `report includes cost estimation`() = runTest { + val breaker = BudgetCircuitBreaker(budget) + + breaker.withBudget("agent-1") { + trackCall(promptTokens = 200, completionTokens = 100) + } + + val report = breaker.getReport("agent-1") + report shouldNotBe null + report!!.estimatedCostUsd shouldBeGreaterThan 0.0 // safe: asserted non-null above + } + + @Test + fun `supports multiple agents concurrently`() = runTest { + val breaker = BudgetCircuitBreaker(budget) + + val job1 = launch { + breaker.withBudget("agent-1") { + trackCall(promptTokens = 100, completionTokens = 50) + } + } + val job2 = launch { + breaker.withBudget("agent-2") { + trackCall(promptTokens = 200, completionTokens = 100) + } + } + + job1.join() + job2.join() + + breaker.getReport("agent-1")?.totalTokens shouldBe 150 + breaker.getReport("agent-2")?.totalTokens shouldBe 300 + } +} diff --git a/budget-breaker/src/test/kotlin/dev/unityinflow/budget/BudgetReportTest.kt b/budget-breaker/src/test/kotlin/dev/unityinflow/budget/BudgetReportTest.kt new file mode 100644 index 0000000..f97955a --- /dev/null +++ b/budget-breaker/src/test/kotlin/dev/unityinflow/budget/BudgetReportTest.kt @@ -0,0 +1,47 @@ +package dev.unityinflow.budget + +import io.kotest.matchers.doubles.shouldBeGreaterThan +import io.kotest.matchers.shouldBe +import org.junit.jupiter.api.Test + +class BudgetReportTest { + + @Test + fun `creates report with all fields`() { + val report = BudgetReport( + agentId = "test-agent", + model = "claude-sonnet-4-6", + promptTokens = 500, + completionTokens = 200, + totalTokens = 700, + estimatedCostUsd = 0.0045, + softLimitBreachCount = 0, + hardLimitBreached = false, + durationMs = 1500, + percentUsed = 70.0, + ) + + report.agentId shouldBe "test-agent" + report.totalTokens shouldBe 700 + report.estimatedCostUsd shouldBeGreaterThan 0.0 + report.hardLimitBreached shouldBe false + } + + @Test + fun `totalTokens equals prompt + completion`() { + val report = BudgetReport( + agentId = "test", + model = "claude-sonnet-4-6", + promptTokens = 300, + completionTokens = 200, + totalTokens = 500, + estimatedCostUsd = 0.0, + softLimitBreachCount = 0, + hardLimitBreached = false, + durationMs = 0, + percentUsed = 50.0, + ) + + report.totalTokens shouldBe (report.promptTokens + report.completionTokens) + } +} diff --git a/budget-breaker/src/test/kotlin/dev/unityinflow/budget/ModelPricingTest.kt b/budget-breaker/src/test/kotlin/dev/unityinflow/budget/ModelPricingTest.kt new file mode 100644 index 0000000..63b0cfd --- /dev/null +++ b/budget-breaker/src/test/kotlin/dev/unityinflow/budget/ModelPricingTest.kt @@ -0,0 +1,54 @@ +package dev.unityinflow.budget + +import io.kotest.matchers.shouldBe +import org.junit.jupiter.api.Test + +class ModelPricingTest { + + @Test + fun `estimates cost for Claude Sonnet`() { + val cost = ModelPricing.estimateCost("claude-sonnet-4-6", promptTokens = 1_000_000, completionTokens = 0) + cost shouldBe 3.0 // $3 per million input tokens + } + + @Test + fun `estimates cost for Claude Opus output`() { + val cost = ModelPricing.estimateCost("claude-opus-4-6", promptTokens = 0, completionTokens = 1_000_000) + cost shouldBe 75.0 // $75 per million output tokens + } + + @Test + fun `estimates combined input + output cost`() { + val cost = ModelPricing.estimateCost("claude-sonnet-4-6", promptTokens = 500_000, completionTokens = 100_000) + // 0.5M * $3/M + 0.1M * $15/M = $1.5 + $1.5 = $3.0 + cost shouldBe 3.0 + } + + @Test + fun `returns zero for unknown model`() { + val cost = ModelPricing.estimateCost("unknown-model", promptTokens = 1000, completionTokens = 1000) + cost shouldBe 0.0 + } + + @Test + fun `supports custom price overrides`() { + val pricing = ModelPricing( + overrides = mapOf( + "my-model" to ModelPricing.PriceConfig(inputPerMillion = 10.0, outputPerMillion = 50.0) + ) + ) + val cost = pricing.estimateCost("my-model", promptTokens = 1_000_000, completionTokens = 0) + cost shouldBe 10.0 + } + + @Test + fun `override takes precedence over default`() { + val pricing = ModelPricing( + overrides = mapOf( + "claude-sonnet-4-6" to ModelPricing.PriceConfig(inputPerMillion = 99.0, outputPerMillion = 99.0) + ) + ) + val cost = pricing.estimateCost("claude-sonnet-4-6", promptTokens = 1_000_000, completionTokens = 0) + cost shouldBe 99.0 + } +} diff --git a/budget-breaker/src/test/kotlin/dev/unityinflow/budget/TokenTrackerTest.kt b/budget-breaker/src/test/kotlin/dev/unityinflow/budget/TokenTrackerTest.kt new file mode 100644 index 0000000..816066a --- /dev/null +++ b/budget-breaker/src/test/kotlin/dev/unityinflow/budget/TokenTrackerTest.kt @@ -0,0 +1,87 @@ +package dev.unityinflow.budget + +import io.kotest.matchers.longs.shouldBeGreaterThan +import io.kotest.matchers.shouldBe +import kotlinx.coroutines.launch +import kotlinx.coroutines.test.runTest +import org.junit.jupiter.api.Test + +class TokenTrackerTest { + + private val budget = AgentBudget( + model = "claude-sonnet-4-6", + hardLimitTokens = 1000, + softLimitTokens = 800, + ) + + @Test + fun `tracks prompt and completion tokens separately`() { + val tracker = TokenTracker("test-agent", budget) + tracker.add(promptTokens = 100, completionTokens = 50) + + tracker.promptTokens shouldBe 100 + tracker.completionTokens shouldBe 50 + tracker.totalTokens shouldBe 150 + } + + @Test + fun `accumulates tokens across multiple calls`() { + val tracker = TokenTracker("test-agent", budget) + tracker.add(promptTokens = 100, completionTokens = 50) + tracker.add(promptTokens = 200, completionTokens = 100) + + tracker.totalTokens shouldBe 450 + } + + @Test + fun `detects soft limit breach`() { + val tracker = TokenTracker("test-agent", budget) + tracker.add(promptTokens = 500, completionTokens = 350) + + tracker.isAboveSoftLimit() shouldBe true + tracker.isAboveHardLimit() shouldBe false + } + + @Test + fun `detects hard limit breach`() { + val tracker = TokenTracker("test-agent", budget) + tracker.add(promptTokens = 600, completionTokens = 500) + + tracker.isAboveHardLimit() shouldBe true + } + + @Test + fun `reports percent used`() { + val tracker = TokenTracker("test-agent", budget) + tracker.add(promptTokens = 250, completionTokens = 250) + + tracker.percentUsed() shouldBe 50.0 + } + + @Test + fun `is thread-safe under concurrent access`() = runTest { + val tracker = TokenTracker("test-agent", AgentBudget(hardLimitTokens = 1_000_000, softLimitTokens = 800_000)) + + val jobs = (1..100).map { + launch { + repeat(100) { + tracker.add(promptTokens = 1, completionTokens = 1) + } + } + } + jobs.forEach { it.join() } + + tracker.totalTokens shouldBe 20_000 + } + + @Test + fun `starts at zero`() { + val tracker = TokenTracker("test-agent", budget) + + tracker.promptTokens shouldBe 0 + tracker.completionTokens shouldBe 0 + tracker.totalTokens shouldBe 0 + tracker.isAboveSoftLimit() shouldBe false + tracker.isAboveHardLimit() shouldBe false + } +} diff --git a/build.gradle.kts b/build.gradle.kts index 280341b..a7459bf 100644 --- a/build.gradle.kts +++ b/build.gradle.kts @@ -4,7 +4,7 @@ plugins { allprojects { group = "dev.unityinflow" - version = "0.0.1-SNAPSHOT" + version = "0.0.1" repositories { mavenCentral() diff --git a/gradle.properties b/gradle.properties index 7879770..30c4c7e 100644 --- a/gradle.properties +++ b/gradle.properties @@ -1,4 +1,4 @@ kotlin.code.style=official org.gradle.jvmargs=-Xmx1g group=dev.unityinflow -version=0.0.1-SNAPSHOT +version=0.0.1