@onmars/lunar-agent-claude 0.7.0 → 0.8.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@onmars/lunar-agent-claude",
-  "version": "0.7.0",
+  "version": "0.8.0",
   "type": "module",
   "main": "src/index.ts",
   "types": "src/index.ts",
@@ -12,7 +12,7 @@
     "LICENSE"
   ],
   "dependencies": {
-    "@onmars/lunar-core": "^0.7.0"
+    "@onmars/lunar-core": "^0.8.0"
   },
   "description": "Claude CLI agent adapter for Lunar",
   "author": "onMars Tech",

package/src/__tests__/token-counter-claude.test.ts

@@ -0,0 +1,211 @@
+import { describe, expect, test } from 'bun:test'
+import { ClaudeTokenCounter } from '../lib/token-counter-claude'
+
+// ─── mock fetch helper ─────────────────────────────────────
+
+interface FakeCall {
+  url: string
+  init: RequestInit
+}
+
+function makeFetch(responses: Array<{ status?: number; body: unknown }>) {
+  const calls: FakeCall[] = []
+  let idx = 0
+  const fetchImpl = (async (input: unknown, init?: RequestInit) => {
+    calls.push({ url: String(input), init: init ?? {} })
+    const resp = responses[Math.min(idx, responses.length - 1)]
+    idx++
+    return new Response(JSON.stringify(resp.body), {
+      status: resp.status ?? 200,
+      headers: { 'content-type': 'application/json' },
+    })
+  }) as unknown as typeof fetch
+  return { fetchImpl, calls }
+}
+
+// ─── tests ─────────────────────────────────────────────────
+
+describe('ClaudeTokenCounter', () => {
+  test('count returns input_tokens from API response', async () => {
+    const { fetchImpl, calls } = makeFetch([{ body: { input_tokens: 12345 } }])
+
+    const counter = new ClaudeTokenCounter({
+      apiKey: 'sk-test',
+      model: 'claude-opus-4-7',
+      fetchImpl,
+    })
+
+    const n = await counter.count([{ role: 'user', content: 'hello world' }])
+    expect(n).toBe(12345)
+    expect(calls).toHaveLength(1)
+
+    // Request shape
+    const init = calls[0].init
+    expect(init.method).toBe('POST')
+    const headers = init.headers as Record<string, string>
+    expect(headers['x-api-key']).toBe('sk-test')
+    expect(headers['anthropic-version']).toBe('2023-06-01')
+    expect(headers['content-type']).toBe('application/json')
+
+    const body = JSON.parse(init.body as string)
+    expect(body.model).toBe('claude-opus-4-7')
+    expect(body.messages).toHaveLength(1)
+  })
+
+  test('cache hit avoids second fetch for identical messages', async () => {
+    const { fetchImpl, calls } = makeFetch([{ body: { input_tokens: 100 } }])
+
+    const counter = new ClaudeTokenCounter({
+      apiKey: 'sk-test',
+      model: 'claude-sonnet-4-6',
+      fetchImpl,
+    })
+
+    const msgs = [{ role: 'user', content: 'hi' }]
+    const a = await counter.count(msgs)
+    const b = await counter.count(msgs)
+
+    expect(a).toBe(100)
+    expect(b).toBe(100)
+    expect(calls).toHaveLength(1)
+    expect(counter.cacheSize()).toBe(1)
+  })
+
+  test('different messages trigger new API call', async () => {
+    const { fetchImpl, calls } = makeFetch([
+      { body: { input_tokens: 10 } },
+      { body: { input_tokens: 20 } },
+    ])
+
+    const counter = new ClaudeTokenCounter({
+      apiKey: 'sk-test',
+      model: 'claude-haiku-4-5',
+      fetchImpl,
+    })
+
+    expect(await counter.count([{ role: 'user', content: 'first' }])).toBe(10)
+    expect(await counter.count([{ role: 'user', content: 'second' }])).toBe(20)
+    expect(calls).toHaveLength(2)
+  })
+
+  test('LRU evicts oldest after 50 unique entries', async () => {
+    // Each call returns a fresh token count. We just care about cache bookkeeping.
+    const { fetchImpl } = makeFetch([{ body: { input_tokens: 1 } }])
+
+    const counter = new ClaudeTokenCounter({
+      apiKey: 'sk-test',
+      model: 'claude-opus-4-7',
+      fetchImpl,
+    })
+
+    for (let i = 0; i < 60; i++) {
+      await counter.count([{ role: 'user', content: `msg-${i}` }])
+    }
+    expect(counter.cacheSize()).toBe(50)
+  })
+
+  test('windowSize returns 200K for standard models', () => {
+    const a = new ClaudeTokenCounter({ apiKey: 'x', model: 'claude-opus-4-7' })
+    expect(a.windowSize()).toBe(200_000)
+
+    const b = new ClaudeTokenCounter({ apiKey: 'x', model: 'claude-sonnet-4-6' })
+    expect(b.windowSize()).toBe(200_000)
+
+    const c = new ClaudeTokenCounter({ apiKey: 'x', model: 'claude-haiku-4-5' })
+    expect(c.windowSize()).toBe(200_000)
+  })
+
+  test('windowSize returns 1M when model ends with [1m]', () => {
+    const counter = new ClaudeTokenCounter({
+      apiKey: 'x',
+      model: 'claude-opus-4-7[1m]',
+    })
+    expect(counter.windowSize()).toBe(1_000_000)
+  })
+
+  test('aliases resolve: opus → claude-opus-4-7', () => {
+    const counter = new ClaudeTokenCounter({ apiKey: 'x', model: 'opus' })
+    expect(counter.windowSize()).toBe(200_000)
+  })
+
+  test('aliases + [1m]: sonnet[1m] → 1M window', () => {
+    const counter = new ClaudeTokenCounter({ apiKey: 'x', model: 'sonnet[1m]' })
+    expect(counter.windowSize()).toBe(1_000_000)
+  })
+
+  test('disabled mode (no apiKey) returns 0 without calling fetch', async () => {
+    let callCount = 0
+    const fetchImpl = (async () => {
+      callCount++
+      return new Response('{}', { status: 200 })
+    }) as unknown as typeof fetch
+
+    const counter = new ClaudeTokenCounter({
+      model: 'claude-opus-4-7',
+      fetchImpl,
+    })
+    expect(counter.isDisabled()).toBe(true)
+
+    const n = await counter.count([{ role: 'user', content: 'hi' }])
+    expect(n).toBe(0)
+    expect(callCount).toBe(0)
+    expect(counter.windowSize()).toBe(200_000)
+  })
+
+  test('empty messages returns 0 without fetching', async () => {
+    let callCount = 0
+    const fetchImpl = (async () => {
+      callCount++
+      return new Response(JSON.stringify({ input_tokens: 999 }), { status: 200 })
+    }) as unknown as typeof fetch
+
+    const counter = new ClaudeTokenCounter({
+      apiKey: 'sk-test',
+      model: 'claude-opus-4-7',
+      fetchImpl,
+    })
+
+    expect(await counter.count([])).toBe(0)
+    expect(callCount).toBe(0)
+  })
+
+  test('non-OK response returns 0 without crashing', async () => {
+    const { fetchImpl } = makeFetch([{ status: 500, body: { error: 'server down' } }])
+    const counter = new ClaudeTokenCounter({
+      apiKey: 'sk-test',
+      model: 'claude-opus-4-7',
+      fetchImpl,
+    })
+    const n = await counter.count([{ role: 'user', content: 'hi' }])
+    expect(n).toBe(0)
+  })
+
+  test('fetch throw is swallowed into 0', async () => {
+    const fetchImpl = (async () => {
+      throw new Error('network failure')
+    }) as unknown as typeof fetch
+
+    const counter = new ClaudeTokenCounter({
+      apiKey: 'sk-test',
+      model: 'claude-opus-4-7',
+      fetchImpl,
+    })
+    const n = await counter.count([{ role: 'user', content: 'hi' }])
+    expect(n).toBe(0)
+  })
+
+  test('context1m flag adds anthropic-beta header', async () => {
+    const { fetchImpl, calls } = makeFetch([{ body: { input_tokens: 7 } }])
+    const counter = new ClaudeTokenCounter({
+      apiKey: 'sk-test',
+      model: 'claude-opus-4-7[1m]',
+      fetchImpl,
+    })
+    await counter.count([{ role: 'user', content: 'hi' }])
+    const headers = calls[0].init.headers as Record<string, string>
+    expect(headers['anthropic-beta']).toBe('context-1m-2025-08-07')
+    // Model sent to API is the alias-resolved, suffix-stripped version
+    const body = JSON.parse(calls[0].init.body as string)
+    expect(body.model).toBe('claude-opus-4-7')
+  })
+})

package/src/index.ts

@@ -1,2 +1,4 @@
 export type { ClaudeAgentOptions } from './adapter'
 export { CLAUDE_MODEL_ALIASES, ClaudeAgent, resolveModelAlias } from './adapter'
+export type { ClaudeTokenCounterOptions } from './lib/token-counter-claude'
+export { ClaudeTokenCounter } from './lib/token-counter-claude'

package/src/lib/token-counter-claude.ts

@@ -0,0 +1,212 @@
+/**
+ * ClaudeTokenCounter — concrete TokenCounter implementation backed by
+ * Anthropic's Token Count API (https://docs.anthropic.com/en/api/messages-count-tokens).
+ *
+ * This lives in agent-claude to keep the core package provider-agnostic.
+ * A tiny in-memory LRU cache (50 entries) keeps the same message bundle
+ * from re-hitting the API when the Router checks the ratio repeatedly.
+ *
+ * If the constructor doesn't receive an API key, the counter operates in
+ * disabled mode: `count()` returns 0 and the FlushHook will simply never
+ * see a token-threshold signal above zero.
+ */
+
+import { createHash } from 'node:crypto'
+import { log } from '@onmars/lunar-core'
+import type { CountableMessage, TokenCounter } from '@onmars/lunar-core/lib/token-counter'
+import { resolveModelAlias } from '../adapter'
+
+// ════════════════════════════════════════════════════════════
+// Model → context window mapping (static; single source of truth here)
+// ════════════════════════════════════════════════════════════
+
+const MODEL_WINDOW: Readonly<Record<string, number>> = Object.freeze({
+  'claude-opus-4-7': 200_000,
+  'claude-opus-4-6': 200_000,
+  'claude-sonnet-4-6': 200_000,
+  'claude-sonnet-4-5': 200_000,
+  'claude-haiku-4-5': 200_000,
+})
+
+const ONE_MILLION = 1_000_000
+const DEFAULT_WINDOW = 200_000
+
+// ════════════════════════════════════════════════════════════
+// Tiny LRU (insertion-ordered Map trick — no dep)
+// ════════════════════════════════════════════════════════════
+
+class TinyLRU<V> {
+  private readonly max: number
+  private readonly store = new Map<string, V>()
+
+  constructor(max: number) {
+    this.max = max
+  }
+
+  get(key: string): V | undefined {
+    const v = this.store.get(key)
+    if (v !== undefined) {
+      // refresh recency
+      this.store.delete(key)
+      this.store.set(key, v)
+    }
+    return v
+  }
+
+  set(key: string, value: V): void {
+    if (this.store.has(key)) {
+      this.store.delete(key)
+    } else if (this.store.size >= this.max) {
+      // evict oldest
+      const oldest = this.store.keys().next().value as string | undefined
+      if (oldest !== undefined) this.store.delete(oldest)
+    }
+    this.store.set(key, value)
+  }
+
+  get size(): number {
+    return this.store.size
+  }
+}
+
+// ════════════════════════════════════════════════════════════
+// Public options + types
+// ════════════════════════════════════════════════════════════
+
+export interface ClaudeTokenCounterOptions {
+  /** Anthropic API key. When absent, the counter runs in disabled mode. */
+  apiKey?: string
+  /** Model ID (aliases resolved; accepts "[1m]" suffix for 1M context). */
+  model: string
+  /** Override the endpoint (tests). */
+  endpoint?: string
+  /** Override fetch (tests). */
+  fetchImpl?: typeof fetch
+  /** Anthropic API version header. Defaults to '2023-06-01'. */
+  apiVersion?: string
+}
+
+interface CountTokensResponse {
+  input_tokens: number
+}
+
+// ════════════════════════════════════════════════════════════
+// Implementation
+// ════════════════════════════════════════════════════════════
+
+export class ClaudeTokenCounter implements TokenCounter {
+  private readonly apiKey?: string
+  private readonly model: string
+  private readonly effectiveModel: string
+  private readonly endpoint: string
+  private readonly fetchImpl: typeof fetch
+  private readonly apiVersion: string
+  private readonly context1m: boolean
+  private readonly cache = new TinyLRU<number>(50)
+  private readonly disabled: boolean
+
+  constructor(options: ClaudeTokenCounterOptions) {
+    this.apiKey = options.apiKey
+    this.model = options.model
+    this.context1m = options.model.endsWith('[1m]')
+    const baseModel = this.context1m ? options.model.slice(0, -'[1m]'.length) : options.model
+    this.effectiveModel = resolveModelAlias(baseModel)
+    this.endpoint = options.endpoint ?? 'https://api.anthropic.com/v1/messages/count_tokens'
+    this.fetchImpl = options.fetchImpl ?? fetch
+    this.apiVersion = options.apiVersion ?? '2023-06-01'
+    this.disabled = !this.apiKey
+
+    if (this.disabled) {
+      log.warn(
+        { model: this.effectiveModel },
+        'ClaudeTokenCounter: no ANTHROPIC_API_KEY — token-threshold signals disabled',
+      )
+    }
+  }
+
+  /** Whether this counter is disabled (no API key). */
+  isDisabled(): boolean {
+    return this.disabled
+  }
+
+  async count(messages: CountableMessage[]): Promise<number> {
+    if (this.disabled || messages.length === 0) return 0
+
+    const key = this.cacheKey(messages)
+    const cached = this.cache.get(key)
+    if (cached !== undefined) return cached
+
+    try {
+      const headers: Record<string, string> = {
+        'content-type': 'application/json',
+        'x-api-key': this.apiKey!,
+        'anthropic-version': this.apiVersion,
+      }
+      if (this.context1m) {
+        headers['anthropic-beta'] = 'context-1m-2025-08-07'
+      }
+
+      const body = JSON.stringify({
+        model: this.effectiveModel,
+        messages: normaliseForApi(messages),
+      })
+
+      const res = await this.fetchImpl(this.endpoint, {
+        method: 'POST',
+        headers,
+        body,
+      })
+
+      if (!res.ok) {
+        log.warn(
+          { status: res.status, model: this.effectiveModel },
+          'ClaudeTokenCounter: count_tokens API returned non-OK — treating as 0',
+        )
+        return 0
+      }
+
+      const data = (await res.json()) as CountTokensResponse
+      const tokens = typeof data?.input_tokens === 'number' ? data.input_tokens : 0
+      this.cache.set(key, tokens)
+      return tokens
+    } catch (err) {
+      log.warn(
+        { err: err instanceof Error ? err.message : String(err) },
+        'ClaudeTokenCounter: count_tokens call failed — treating as 0',
+      )
+      return 0
+    }
+  }
+
+  windowSize(): number {
+    const base = MODEL_WINDOW[this.effectiveModel] ?? DEFAULT_WINDOW
+    return this.context1m ? ONE_MILLION : base
+  }
+
+  /** Exposed for tests. */
+  cacheSize(): number {
+    return this.cache.size
+  }
+
+  private cacheKey(messages: CountableMessage[]): string {
+    const serialised = JSON.stringify(messages)
+    return createHash('sha1')
+      .update(
+        this.effectiveModel + '\u0000' + (this.context1m ? '1m' : 'std') + '\u0000' + serialised,
+      )
+      .digest('hex')
+  }
+}
+
+// ════════════════════════════════════════════════════════════
+// Helpers
+// ════════════════════════════════════════════════════════════
+
+function normaliseForApi(messages: CountableMessage[]): Array<{ role: string; content: unknown }> {
+  // Anthropic accepts string or block-array content; we just pass through.
+  // If content is null/undefined, stringify to empty.
+  return messages.map((m) => ({
+    role: m.role,
+    content: m.content == null ? '' : m.content,
+  }))
+}