airaknit 1.1.2-rc.9

package/src/server/auto-continue/limit-detector.ts

@@ -0,0 +1,194 @@
+export interface LimitDetection {
+  chatId: string
+  resetAt: number
+  tz: string
+  raw: unknown
+}
+
+export interface LimitDetector {
+  detect(chatId: string, error: unknown): LimitDetection | null
+  detectFromResultText?(chatId: string, text: string, nowMs?: number): LimitDetection | null
+  detectFromSdkRateLimitInfo?(chatId: string, info: unknown): LimitDetection | null
+}
+
+interface ErrorLike {
+  message?: string
+  status?: number
+  headers?: Record<string, string>
+}
+
+function extractHeaders(error: unknown): Record<string, string> {
+  if (error && typeof error === "object" && "headers" in error) {
+    const headers = (error as ErrorLike).headers
+    if (headers && typeof headers === "object") return headers
+  }
+  return {}
+}
+
+function parseBody(error: unknown): Record<string, unknown> | null {
+  if (!error || typeof error !== "object") return null
+  const message = (error as ErrorLike).message
+  if (!message) return null
+  try {
+    const parsed = JSON.parse(message)
+    return parsed && typeof parsed === "object" ? (parsed as Record<string, unknown>) : null
+  } catch {
+    return null
+  }
+}
+
+function parseIsoMillis(value: unknown): number | null {
+  if (typeof value !== "string" || !value) return null
+  const millis = new Date(value).getTime()
+  return Number.isFinite(millis) ? millis : null
+}
+
+function zonedWallClockToUtcMs(
+  year: number, month: number, day: number, hour: number, minute: number, tz: string,
+): number {
+  const utcGuess = Date.UTC(year, month - 1, day, hour, minute)
+  const dtf = new Intl.DateTimeFormat("en-US", {
+    timeZone: tz, year: "numeric", month: "2-digit", day: "2-digit",
+    hour: "2-digit", minute: "2-digit", second: "2-digit", hour12: false,
+  })
+  const parts = Object.fromEntries(
+    dtf.formatToParts(new Date(utcGuess))
+      .filter((part) => part.type !== "literal")
+      .map((part) => [part.type, part.value]),
+  )
+  const asLocal = Date.UTC(
+    Number(parts.year), Number(parts.month) - 1, Number(parts.day),
+    parts.hour === "24" ? 0 : Number(parts.hour), Number(parts.minute),
+  )
+  return utcGuess - (asLocal - utcGuess)
+}
+
+export function parseResetFromText(text: string, nowMs: number = Date.now()): { resetAt: number; tz: string } | null {
+  if (typeof text !== "string") return null
+  const match = text.match(/resets\s+(\d{1,2})(?::(\d{2}))?(am|pm)\s*\(([^)]+)\)/i)
+  if (!match) return null
+  const hour12 = Number(match[1])
+  const minute = match[2] ? Number(match[2]) : 0
+  const meridiem = match[3].toLowerCase()
+  const tz = match[4].trim()
+  if (!Number.isFinite(hour12) || hour12 < 1 || hour12 > 12) return null
+  if (!Number.isFinite(minute) || minute < 0 || minute > 59) return null
+  const hour24 = meridiem === "pm"
+    ? (hour12 === 12 ? 12 : hour12 + 12)
+    : (hour12 === 12 ? 0 : hour12)
+  let tzYear: number, tzMonth: number, tzDay: number
+  try {
+    const dtf = new Intl.DateTimeFormat("en-US", {
+      timeZone: tz, year: "numeric", month: "2-digit", day: "2-digit",
+    })
+    const parts = Object.fromEntries(
+      dtf.formatToParts(new Date(nowMs))
+        .filter((part) => part.type !== "literal")
+        .map((part) => [part.type, part.value]),
+    )
+    tzYear = Number(parts.year)
+    tzMonth = Number(parts.month)
+    tzDay = Number(parts.day)
+  } catch {
+    return null
+  }
+  let resetAt = zonedWallClockToUtcMs(tzYear, tzMonth, tzDay, hour24, minute, tz)
+  if (resetAt <= nowMs) {
+    const next = new Date(Date.UTC(tzYear, tzMonth - 1, tzDay) + 24 * 3600_000)
+    resetAt = zonedWallClockToUtcMs(
+      next.getUTCFullYear(), next.getUTCMonth() + 1, next.getUTCDate(), hour24, minute, tz,
+    )
+  }
+  return { resetAt, tz }
+}
+
+export class ClaudeLimitDetector implements LimitDetector {
+  detect(chatId: string, error: unknown): LimitDetection | null {
+    const body = parseBody(error)
+    const inner = body && typeof body.error === "object" && body.error !== null
+      ? (body.error as Record<string, unknown>)
+      : null
+    const isRateLimit = inner?.type === "rate_limit_error"
+      || (error as ErrorLike | null)?.status === 429 && inner?.type === "rate_limit_error"
+
+    if (isRateLimit) {
+      const headers = extractHeaders(error)
+      const resetAt = parseIsoMillis(headers["anthropic-ratelimit-unified-reset"])
+        ?? parseIsoMillis(inner?.resets_at)
+        ?? parseIsoMillis(inner?.reset_at)
+      if (resetAt !== null) {
+        const tz = headers["x-anthropic-timezone"]
+          ?? (typeof inner?.timezone === "string" ? (inner.timezone as string) : null)
+          ?? "system"
+        return { chatId, resetAt, tz, raw: error }
+      }
+    }
+
+    // Fallback: the Claude Code SDK rethrows CLI result errors as
+    // `Error("Claude Code returned an error result: <text>")`. Parse the
+    // text directly for "You've hit your limit · resets ..." / "usage limit
+    // reached|<unix>" forms.
+    const message = (error as ErrorLike | null)?.message
+    if (typeof message === "string") {
+      return this.detectFromResultText(chatId, message)
+    }
+    return null
+  }
+
+  detectFromResultText(chatId: string, text: string, nowMs: number = Date.now()): LimitDetection | null {
+    const parsed = parseResetFromText(text, nowMs)
+    if (parsed) return { chatId, resetAt: parsed.resetAt, tz: parsed.tz, raw: text }
+    const pipe = parseClaudeUsageLimitPipe(text)
+    if (pipe !== null) return { chatId, resetAt: pipe, tz: "system", raw: text }
+    return null
+  }
+
+  detectFromSdkRateLimitInfo(chatId: string, info: unknown): LimitDetection | null {
+    if (!info || typeof info !== "object") return null
+    const rec = info as Record<string, unknown>
+    if (rec.status !== "rejected") return null
+    const raw = rec.resetsAt
+    if (typeof raw !== "number" || !Number.isFinite(raw) || raw <= 0) return null
+    // SDK emits `resetsAt` as epoch seconds for claude.ai subscription limits;
+    // coerce to ms defensively (anything below year 5138 in ms is below 1e14).
+    const resetAt = raw < 1e12 ? Math.round(raw * 1000) : raw
+    return { chatId, resetAt, tz: "system", raw: info }
+  }
+}
+
+export function parseClaudeUsageLimitPipe(text: string): number | null {
+  // Claude CLI sometimes returns "Claude AI usage limit reached|<unix-seconds>".
+  if (typeof text !== "string") return null
+  const match = text.match(/usage limit reached\|(\d{9,13})/i)
+  if (!match) return null
+  const value = Number(match[1])
+  if (!Number.isFinite(value) || value <= 0) return null
+  return value < 1e12 ? value * 1000 : value
+}
+
+interface JsonRpcErrorLike {
+  code?: number
+  message?: string
+  data?: Record<string, unknown>
+}
+
+export class CodexLimitDetector implements LimitDetector {
+  detect(chatId: string, error: unknown): LimitDetection | null {
+    if (!error || typeof error !== "object") return null
+    const rpc = error as JsonRpcErrorLike
+    const data = rpc.data && typeof rpc.data === "object" ? rpc.data : null
+    const isRateLimit = data?.code === "rate_limit" || rpc.code === -32001
+    if (!isRateLimit) return null
+
+    let resetAt: number | null
+    if (typeof data?.resets_at_ms === "number" && Number.isFinite(data.resets_at_ms)) {
+      resetAt = data.resets_at_ms
+    } else {
+      resetAt = parseIsoMillis(data?.resets_at)
+    }
+    if (resetAt === null) return null
+
+    const tz = typeof data?.timezone === "string" ? (data.timezone as string) : "system"
+    return { chatId, resetAt, tz, raw: error }
+  }
+}

package/src/server/bm25.test.ts

@@ -0,0 +1,92 @@
+// src/server/bm25.test.ts
+import { describe, expect, test } from "bun:test"
+import { BM25Index } from "./bm25"
+
+describe("BM25Index", () => {
+  describe("tokenize", () => {
+    test("lowercases and splits on whitespace/punctuation", () => {
+      const index = new BM25Index<string>()
+      index.add("d1", "Hello, World! This is a test.")
+      const results = index.search("hello")
+      expect(results.length).toBe(1)
+      expect(results[0].id).toBe("d1")
+    })
+
+    test("filters stopwords", () => {
+      const index = new BM25Index<string>()
+      index.add("d1", "the quick brown fox")
+      index.add("d2", "a lazy dog")
+      // "the" and "a" are stopwords, should not dominate results
+      const results = index.search("quick")
+      expect(results.length).toBe(1)
+      expect(results[0].id).toBe("d1")
+    })
+  })
+
+  describe("search", () => {
+    test("ranks relevant documents higher", () => {
+      const index = new BM25Index<string>()
+      index.add("d1", "postgres database migration schema users table")
+      index.add("d2", "react component button styling CSS")
+      index.add("d3", "database connection pool postgres config")
+
+      const results = index.search("postgres database")
+      expect(results.length).toBeGreaterThanOrEqual(2)
+      // d1 and d3 both mention postgres+database terms
+      const ids = results.map((r) => r.id)
+      expect(ids).toContain("d1")
+      expect(ids).toContain("d3")
+      // d2 should not appear or rank very low
+      expect(ids.indexOf("d1")).toBeLessThan(ids.indexOf("d2") === -1 ? Infinity : ids.indexOf("d2"))
+    })
+
+    test("returns empty for no matches", () => {
+      const index = new BM25Index<string>()
+      index.add("d1", "hello world")
+      const results = index.search("nonexistent")
+      expect(results).toEqual([])
+    })
+
+    test("respects limit parameter", () => {
+      const index = new BM25Index<string>()
+      for (let i = 0; i < 20; i++) {
+        index.add(`d${i}`, `document ${i} about testing`)
+      }
+      const results = index.search("testing", 5)
+      expect(results.length).toBe(5)
+    })
+
+    test("handles multi-field documents via concatenation", () => {
+      const index = new BM25Index<string>()
+      index.add("d1", "auth middleware implementation error handling retry logic")
+      index.add("d2", "auth login form validation")
+      const results = index.search("auth error handling")
+      expect(results[0].id).toBe("d1")
+    })
+  })
+
+  describe("remove", () => {
+    test("removes document from index", () => {
+      const index = new BM25Index<string>()
+      index.add("d1", "hello world")
+      index.add("d2", "hello there")
+      index.remove("d1")
+      const results = index.search("hello")
+      expect(results.length).toBe(1)
+      expect(results[0].id).toBe("d2")
+    })
+  })
+
+  describe("size", () => {
+    test("tracks document count", () => {
+      const index = new BM25Index<string>()
+      expect(index.size).toBe(0)
+      index.add("d1", "hello")
+      expect(index.size).toBe(1)
+      index.add("d2", "world")
+      expect(index.size).toBe(2)
+      index.remove("d1")
+      expect(index.size).toBe(1)
+    })
+  })
+})

package/src/server/bm25.ts

@@ -0,0 +1,101 @@
+// src/server/bm25.ts
+
+const STOPWORDS = new Set([
+  "a", "an", "the", "is", "it", "of", "in", "to", "and", "or", "for",
+  "on", "at", "by", "with", "from", "as", "be", "was", "were", "been",
+  "are", "has", "had", "have", "do", "does", "did", "but", "not", "this",
+  "that", "these", "those", "i", "we", "you", "he", "she", "they",
+])
+
+function tokenize(text: string): string[] {
+  return text
+    .toLowerCase()
+    .split(/[^a-z0-9_]+/)
+    .filter((t) => t.length > 1 && !STOPWORDS.has(t))
+}
+
+interface DocEntry {
+  tokens: string[]
+  length: number
+}
+
+export interface BM25Result<ID> {
+  id: ID
+  score: number
+}
+
+export class BM25Index<ID extends string = string> {
+  private readonly k1 = 1.2
+  private readonly b = 0.75
+  private readonly docs = new Map<ID, DocEntry>()
+  private readonly invertedIndex = new Map<string, Set<ID>>()
+  private totalLength = 0
+
+  get size(): number {
+    return this.docs.size
+  }
+
+  private get avgLength(): number {
+    return this.docs.size === 0 ? 0 : this.totalLength / this.docs.size
+  }
+
+  add(id: ID, text: string): void {
+    this.remove(id)
+    const tokens = tokenize(text)
+    this.docs.set(id, { tokens, length: tokens.length })
+    this.totalLength += tokens.length
+
+    for (const token of tokens) {
+      let set = this.invertedIndex.get(token)
+      if (!set) {
+        set = new Set()
+        this.invertedIndex.set(token, set)
+      }
+      set.add(id)
+    }
+  }
+
+  remove(id: ID): void {
+    const doc = this.docs.get(id)
+    if (!doc) return
+    this.totalLength -= doc.length
+    this.docs.delete(id)
+    for (const token of doc.tokens) {
+      const set = this.invertedIndex.get(token)
+      if (set) {
+        set.delete(id)
+        if (set.size === 0) this.invertedIndex.delete(token)
+      }
+    }
+  }
+
+  search(query: string, limit = 10): BM25Result<ID>[] {
+    const queryTokens = tokenize(query)
+    if (queryTokens.length === 0) return []
+
+    const scores = new Map<ID, number>()
+    const N = this.docs.size
+    const avgDl = this.avgLength
+
+    for (const qt of queryTokens) {
+      const postings = this.invertedIndex.get(qt)
+      if (!postings) continue
+
+      const df = postings.size
+      const idf = Math.log(1 + (N - df + 0.5) / (df + 0.5))
+
+      for (const docId of postings) {
+        const doc = this.docs.get(docId)!
+        const tf = doc.tokens.filter((t) => t === qt).length
+        const tfNorm = (tf * (this.k1 + 1)) / (tf + this.k1 * (1 - this.b + this.b * (doc.length / avgDl)))
+        const score = idf * tfNorm
+        scores.set(docId, (scores.get(docId) ?? 0) + score)
+      }
+    }
+
+    return [...scores.entries()]
+      .sort((a, b) => b[1] - a[1])
+      .slice(0, limit)
+      .map(([id, score]) => ({ id, score }))
+  }
+}

package/src/server/chat-events-jetstream.test.ts

@@ -0,0 +1,135 @@
+import { describe, test, expect, afterEach } from "bun:test"
+import { connect, type NatsConnection } from "@nats-io/transport-node"
+import { jetstream, DeliverPolicy } from "@nats-io/jetstream"
+import { NatsBridge } from "./nats-bridge"
+import { ensureChatMessageStream, CHAT_MESSAGE_EVENTS_STREAM } from "./nats-streams"
+import { chatMessageSubject } from "../shared/nats-subjects"
+import { compressPayload, decompressPayload } from "../shared/compression"
+
+const encoder = new TextEncoder()
+const decoder = new TextDecoder()
+
+let bridge: NatsBridge | null = null
+let nc: NatsConnection | null = null
+
+afterEach(async () => {
+  await nc?.drain()
+  nc = null
+  if (bridge) {
+    await bridge.dispose()
+    bridge = null
+  }
+})
+
+async function setup() {
+  bridge = await NatsBridge.create()
+  nc = await connect({ servers: bridge.natsUrl })
+  await ensureChatMessageStream(nc)
+  return nc
+}
+
+describe("chat message events via JetStream ordered consumer", () => {
+  test("consumer receives events published to JetStream", async () => {
+    const conn = await setup()
+    const js = jetstream(conn)
+    const chatId = "test-chat-js-1"
+    const subject = chatMessageSubject(chatId)
+
+    // Create ordered consumer with DeliverPolicy.New
+    const consumer = await js.consumers.get(CHAT_MESSAGE_EVENTS_STREAM, {
+      filter_subjects: subject,
+      deliver_policy: DeliverPolicy.New,
+    })
+
+    // Publish events (mimicking server-side publishChatMessage)
+    const events = [
+      { chatId, entry: { id: "1", kind: "message", content: "hello" } },
+      { chatId, entry: { id: "2", kind: "message", content: "world" } },
+    ]
+
+    for (const event of events) {
+      const payload = compressPayload(encoder.encode(JSON.stringify(event)))
+      await js.publish(subject, payload)
+    }
+
+    // Consume and verify
+    const received: unknown[] = []
+    const messages = await consumer.consume()
+    for await (const msg of messages) {
+      const decoded = await decompressPayload(msg.data)
+      const data = JSON.parse(decoder.decode(decoded))
+      received.push(data)
+      if (received.length >= 2) break
+    }
+    await messages.close()
+
+    expect(received).toHaveLength(2)
+    expect(received[0]).toEqual(events[0])
+    expect(received[1]).toEqual(events[1])
+  })
+
+  test("consumer filters by chat subject", async () => {
+    const conn = await setup()
+    const js = jetstream(conn)
+
+    const chatId1 = "chat-js-filter-1"
+    const chatId2 = "chat-js-filter-2"
+
+    const consumer = await js.consumers.get(CHAT_MESSAGE_EVENTS_STREAM, {
+      filter_subjects: chatMessageSubject(chatId1),
+      deliver_policy: DeliverPolicy.New,
+    })
+
+    // Publish to chat 2 (should be filtered out)
+    await js.publish(
+      chatMessageSubject(chatId2),
+      encoder.encode(JSON.stringify({ chatId: chatId2, entry: { id: "x" } }))
+    )
+
+    // Publish to chat 1 (should be received)
+    await js.publish(
+      chatMessageSubject(chatId1),
+      encoder.encode(JSON.stringify({ chatId: chatId1, entry: { id: "1" } }))
+    )
+
+    const messages = await consumer.consume()
+    const received: unknown[] = []
+    for await (const msg of messages) {
+      const data = JSON.parse(decoder.decode(msg.data))
+      received.push(data)
+      if (received.length >= 1) break
+    }
+    await messages.close()
+
+    expect(received).toHaveLength(1)
+    expect(received[0]).toEqual({ chatId: chatId1, entry: { id: "1" } })
+  })
+
+  test("compressed events are decompressed correctly by consumer", async () => {
+    const conn = await setup()
+    const js = jetstream(conn)
+    const chatId = "test-chat-compress"
+    const subject = chatMessageSubject(chatId)
+
+    const consumer = await js.consumers.get(CHAT_MESSAGE_EVENTS_STREAM, {
+      filter_subjects: subject,
+      deliver_policy: DeliverPolicy.New,
+    })
+
+    // Create a large payload that will trigger compression (> 64KB threshold)
+    const largeContent = "x".repeat(100_000)
+    const event = { chatId, entry: { id: "big", content: largeContent } }
+    const payload = compressPayload(encoder.encode(JSON.stringify(event)))
+
+    await js.publish(subject, payload)
+
+    const messages = await consumer.consume()
+    for await (const msg of messages) {
+      const decoded = await decompressPayload(msg.data)
+      const data = JSON.parse(decoder.decode(decoded)) as typeof event
+      expect(data.entry.content).toBe(largeContent)
+      break
+    }
+    await messages.close()
+  })
+})