@ixo/oracle-runtime-workers 0.12.0 → 0.15.0

package/src/channels/turns.ts

@@ -0,0 +1,270 @@
+import type { DoSqliteDatabase } from '../sqlite/database';
+import type { RunRecord } from '../do/run-store';
+import type { TurnIdentity, TurnRequest } from '../do/contracts';
+import { parseReplyPlan, planText } from '../delivery/schema';
+import type { ReplyPlan } from '../delivery/types';
+import {
+  ChannelError,
+  channelRequestHash,
+  type ChannelTurnInput,
+  type ChannelTurnResponse,
+} from './contract';
+
+interface ReceiptRow extends Record<string, string | number | null> {
+  binding_id: string;
+  request_id: string;
+  request_hash: string;
+  run_id: string;
+  session_id: string | null;
+  /** 1 once the assistant reply was mirrored into the Companion room. */
+  reply_mirrored: number;
+}
+
+export interface ChannelTurnsHost {
+  createSession(identity: TurnIdentity, markerTxnId: string): Promise<string>;
+  assertSession(identity: TurnIdentity, sessionId: string): Promise<void>;
+  /** Throws a `ChannelError` when the user has no stored delegation for this oracle. */
+  requireDelegation(identity: TurnIdentity): Promise<void>;
+  getRun(runId: string): Promise<RunRecord | undefined>;
+  /** The Reply Plan (JSON) a finished run was built into, if any. */
+  getPlan(runId: string): Promise<string | null>;
+  wasPruned(runId: string): Promise<boolean>;
+  begin(runId: string, request: TurnRequest): Promise<RunRecord>;
+  mirror(
+    request: TurnRequest,
+    text: string,
+    author: 'user' | 'oracle',
+  ): Promise<void>;
+  changed(): Promise<void>;
+}
+
+export class ChannelTurns {
+  private admission: Promise<unknown> = Promise.resolve();
+  private setupPromise: Promise<void> | undefined;
+
+  constructor(
+    private readonly db: DoSqliteDatabase,
+    private readonly host: ChannelTurnsHost,
+  ) {}
+
+  private setup(): Promise<void> {
+    this.setupPromise ??= (async () => {
+      await this.db.run(`CREATE TABLE IF NOT EXISTS channel_requests (
+        binding_id TEXT NOT NULL, request_id TEXT NOT NULL, request_hash TEXT NOT NULL,
+        run_id TEXT NOT NULL UNIQUE, session_id TEXT,
+        reply_mirrored INTEGER NOT NULL DEFAULT 0,
+        PRIMARY KEY (binding_id, request_id)
+      )`);
+      // Tables created before the column existed (CREATE TABLE IF NOT
+      // EXISTS leaves them as they were).
+      const columns = await this.db.exec<{ name: string }>(
+        'PRAGMA table_info(channel_requests)',
+      );
+      if (!columns.some((c) => c.name === 'reply_mirrored'))
+        await this.db.run(
+          'ALTER TABLE channel_requests ADD COLUMN reply_mirrored INTEGER NOT NULL DEFAULT 0',
+        );
+      await this.db.run(`CREATE TABLE IF NOT EXISTS channel_sessions (
+        binding_id TEXT PRIMARY KEY, session_id TEXT NOT NULL
+      )`);
+    })().catch((error: unknown) => {
+      this.setupPromise = undefined;
+      throw error;
+    });
+    return this.setupPromise;
+  }
+
+  /** Admissions, reply deliveries and session releases run one at a time. */
+  private serialize<T>(work: () => Promise<T>): Promise<T> {
+    const next = this.admission.then(work);
+    this.admission = next.catch(() => undefined);
+    return next;
+  }
+
+  submit(
+    identity: TurnIdentity,
+    input: ChannelTurnInput,
+    requestHash: string,
+  ): Promise<ChannelTurnResponse> {
+    return this.serialize(() => this.admit(identity, input, requestHash));
+  }
+
+  /**
+   * Mirror a finished run's reply into the Companion room once: the receipt
+   * records the delivery, so later polls of the same request do not send it
+   * again. Called when the run ends; a poll retries a delivery that failed.
+   * `text` is the model's reply; a run with a Reply Plan mirrors the plan
+   * instead, as the channel user received it.
+   */
+  deliverReply(
+    request: TurnRequest,
+    runId: string,
+    text: string,
+  ): Promise<void> {
+    return this.serialize(async () => {
+      await this.deliver(request, runId, await this.reply(runId, text));
+    });
+  }
+
+  /** A finished run's Reply Plan, and the reply as the channel user sees it. */
+  private async reply(
+    runId: string,
+    text: string,
+  ): Promise<{ plan: ReplyPlan | null; rendered: string }> {
+    const plan = parseReplyPlan(await this.host.getPlan(runId));
+    return { plan, rendered: plan ? planText(plan) : text };
+  }
+
+  /**
+   * The session was deleted: release every binding that pointed at it, so
+   * the binding's next message opens a new Companion session. Returns
+   * whether a binding was released.
+   */
+  forgetSession(sessionId: string): Promise<boolean> {
+    return this.serialize(async () => {
+      await this.setup();
+      const released = await this.db.run(
+        'DELETE FROM channel_sessions WHERE session_id = ?',
+        [sessionId],
+      );
+      if (released.changes === 0) return false;
+      await this.host.changed();
+      return true;
+    });
+  }
+
+  private async deliver(
+    request: TurnRequest,
+    runId: string,
+    reply: { rendered: string },
+  ): Promise<void> {
+    await this.setup();
+    const receipt = await this.db.get<{ reply_mirrored: number }>(
+      'SELECT reply_mirrored FROM channel_requests WHERE run_id = ?',
+      [runId],
+    );
+    if (Number(receipt?.reply_mirrored ?? 0) === 1) return;
+    await this.host.mirror(request, reply.rendered, 'oracle');
+    await this.db.run(
+      'UPDATE channel_requests SET reply_mirrored = 1 WHERE run_id = ?',
+      [runId],
+    );
+    await this.host.changed();
+  }
+
+  private async admit(
+    identity: TurnIdentity,
+    input: ChannelTurnInput,
+    requestHash: string,
+  ): Promise<ChannelTurnResponse> {
+    await this.setup();
+    if (
+      !identity.channel ||
+      identity.channel.bindingId !== input.bindingId ||
+      identity.channel.bindingRevision !== input.bindingRevision ||
+      identity.channel.provider !== input.provider
+    )
+      throw new ChannelError(
+        403,
+        'Channel identity does not match this request',
+      );
+    const runId = `channel_${await channelRequestHash(JSON.stringify([input.bindingId, input.requestId]))}`;
+    // Checked before a receipt is written: the receipt of a pruned run is
+    // gone (run-store.ts), and a new one must not take its place.
+    if (await this.host.wasPruned(runId))
+      throw new ChannelError(
+        410,
+        'Channel response has expired; this request cannot execute again',
+      );
+    const admitted = await this.db.run(
+      `INSERT OR IGNORE INTO channel_requests
+      (binding_id, request_id, request_hash, run_id) VALUES (?, ?, ?, ?)`,
+      [input.bindingId, input.requestId, requestHash, runId],
+    );
+    // A poll of an admitted request changes nothing.
+    if (admitted.changes === 1) await this.host.changed();
+    const receipt = await this.db.get<ReceiptRow>(
+      'SELECT * FROM channel_requests WHERE binding_id = ? AND request_id = ?',
+      [input.bindingId, input.requestId],
+    );
+    if (!receipt || receipt.request_hash !== requestHash)
+      throw new ChannelError(
+        409,
+        'This request ID already belongs to another message',
+      );
+    let sessionId = receipt.session_id;
+    // The session is checked when this request is bound to it, and again
+    // before the request starts a run; a poll of a run that exists already
+    // only reads its receipt and row.
+    let sessionChecked = false;
+    if (!sessionId) {
+      const binding = await this.db.get<{ session_id: string }>(
+        'SELECT session_id FROM channel_sessions WHERE binding_id = ?',
+        [input.bindingId],
+      );
+      if (binding && input.sessionId && binding.session_id !== input.sessionId)
+        throw new ChannelError(409, 'This channel already has another session');
+      sessionId =
+        binding?.session_id ??
+        input.sessionId ??
+        // The marker transaction id is fixed per binding AND the request that
+        // opens the session: a retry of that request reuses the marker event,
+        // while a binding whose session was deleted gets a new one.
+        (await this.host.createSession(
+          identity,
+          `channel-session-${await channelRequestHash(JSON.stringify([input.bindingId, input.requestId]))}`,
+        ));
+      await this.host.assertSession(identity, sessionId);
+      sessionChecked = true;
+      await this.db.transaction(async () => {
+        await this.db.run(
+          'INSERT OR IGNORE INTO channel_sessions (binding_id, session_id) VALUES (?, ?)',
+          [input.bindingId, sessionId],
+        );
+        await this.db.run(
+          'UPDATE channel_requests SET session_id = ? WHERE run_id = ?',
+          [sessionId, runId],
+        );
+      });
+      await this.host.changed();
+    }
+    const request: TurnRequest = {
+      identity,
+      sessionId,
+      message: input.message,
+      client: 'channel',
+      requestId: input.requestId,
+      multitask: 'enqueue',
+      channel: {
+        provider: input.provider,
+        bindingId: input.bindingId,
+        remoteMessageRef: input.remoteMessageRef,
+      },
+    };
+    let record = await this.host.getRun(runId);
+    if (!record) {
+      if (!sessionChecked) await this.host.assertSession(identity, sessionId);
+      await this.host.requireDelegation(identity);
+      await this.host.mirror(request, request.message, 'user');
+      record = await this.host.begin(runId, request);
+    }
+    const reply =
+      record.status === 'finished'
+        ? await this.reply(runId, record.partialText ?? '')
+        : null;
+    // The Companion room records the reply the channel user received, once.
+    if (reply) await this.deliver(request, runId, reply);
+    return {
+      requestId: input.requestId,
+      runId,
+      sessionId,
+      status: record.status,
+      ...(record.messageId ? { messageId: record.messageId } : {}),
+      // `text` is the whole reply as one message for gateways that predate
+      // plans: the rendered plan when there is one (it carries the artefact
+      // links), else the model's text.
+      ...(reply ? { text: reply.rendered } : {}),
+      ...(reply?.plan ? { plan: reply.plan } : {}),
+    };
+  }
+}

package/src/core/budgeted-llm.ts

@@ -0,0 +1,223 @@
+/**
+ * The turn's LLM adapter, metered: every model it hands out carries a
+ * callback that reserves the call against the `TurnBudget` before the
+ * provider is contacted and settles the reservation to the provider's
+ * reported usage when the call ends. Helper models (summarizer, attachment
+ * extraction, a plugin's own `ctx.llm.get`) are covered the same way as the
+ * main model, since they all come from this adapter.
+ *
+ * The same handler is also put in the graph's run config (`callbacks`) so a
+ * model the host constructed elsewhere and passed in directly is still
+ * charged; a call reported through both paths is counted once.
+ */
+import { BaseCallbackHandler } from '@langchain/core/callbacks/base';
+import type { BaseChatModel } from '@langchain/core/language_models/chat_models';
+import type { LLMResult } from '@langchain/core/outputs';
+import type { ChatOpenAIFields, ModelRole } from '../plugin-api/types';
+import {
+  estimateContentTokens,
+  messageContentTokens,
+  messageToolCallTokens,
+} from './context-budget';
+import type { LlmAdapter } from './runtime-context';
+import type { TurnBudget } from './turn-budget';
+
+export interface BudgetedLlmOptions {
+  budget: TurnBudget;
+  /** Tokens reserved for the reply of each call (the context budget's output reserve). */
+  outputReserveTokens: number;
+  /** The turn's abort signal: a call after the abort is refused before it starts. */
+  signal?: AbortSignal;
+}
+
+export interface BudgetedLlm extends LlmAdapter {
+  /** Put this in the run config's `callbacks` as well. */
+  readonly callback: BaseCallbackHandler;
+}
+
+/**
+ * Key of a request's tool list for the per-turn schema estimate: the tool
+ * names in order. Within one turn a name always carries the same schema (the
+ * tools are bound when the turn's graph is built), so the names identify
+ * the list; anything without a recognisable name is not cached.
+ */
+function toolListKey(tools: unknown): string | undefined {
+  if (!Array.isArray(tools)) return undefined;
+  const names: string[] = [];
+  for (const entry of tools) {
+    const fn: unknown =
+      entry && typeof entry === 'object' && 'function' in entry
+        ? entry.function
+        : entry;
+    const name: unknown =
+      fn && typeof fn === 'object' && 'name' in fn ? fn.name : undefined;
+    if (typeof name !== 'string') return undefined;
+    names.push(name);
+  }
+  return names.join('\n');
+}
+
+/**
+ * What one model call sends, reduced to what costs tokens. Message estimates
+ * are memoised per message (`messageContentTokens`), the tool schemas per
+ * tool list in `toolTokens`, so a long history is not re-serialised on every
+ * step of the turn.
+ */
+function estimateRequestTokens(
+  messages: LLMResultMessages,
+  invocationParams: unknown,
+  toolTokens: Map<string, number>,
+): number {
+  let tokens = 0;
+  for (const batch of messages)
+    for (const message of batch) {
+      tokens += messageContentTokens(message) + 4;
+      if ('tool_calls' in message && Array.isArray(message.tool_calls))
+        tokens += messageToolCallTokens(message, message.tool_calls);
+    }
+  if (
+    invocationParams &&
+    typeof invocationParams === 'object' &&
+    'tools' in invocationParams
+  ) {
+    const { tools } = invocationParams;
+    const key = toolListKey(tools);
+    let schema = key === undefined ? undefined : toolTokens.get(key);
+    if (schema === undefined) {
+      schema = estimateContentTokens(tools);
+      if (key !== undefined) toolTokens.set(key, schema);
+    }
+    tokens += schema;
+  }
+  return tokens;
+}
+
+/** The summarizer tags its model call with this `lc_source` (LangChain's summarization middleware). */
+const SUMMARIZATION_SOURCE = 'summarization';
+
+type LLMResultMessages = Parameters<
+  NonNullable<BaseCallbackHandler['handleChatModelStart']>
+>[1];
+
+function numberField(value: unknown, key: string): number | undefined {
+  if (!value || typeof value !== 'object' || !(key in value)) return undefined;
+  const n = Reflect.get(value, key);
+  return typeof n === 'number' && Number.isFinite(n) ? n : undefined;
+}
+
+/** Provider usage of one call: `llmOutput.tokenUsage` (OpenAI-compatible) or the message's `usage_metadata`. */
+function reportedTotalTokens(output: LLMResult): number | undefined {
+  const usage: unknown = output.llmOutput?.tokenUsage;
+  const total = numberField(usage, 'totalTokens');
+  if (total !== undefined) return total;
+  const prompt = numberField(usage, 'promptTokens');
+  const completion = numberField(usage, 'completionTokens');
+  if (prompt !== undefined && completion !== undefined)
+    return prompt + completion;
+  const generation: unknown = output.generations[0]?.[0];
+  const message =
+    generation && typeof generation === 'object' && 'message' in generation
+      ? generation.message
+      : undefined;
+  const meta =
+    message && typeof message === 'object' && 'usage_metadata' in message
+      ? message.usage_metadata
+      : undefined;
+  return numberField(meta, 'total_tokens');
+}
+
+class TurnBudgetHandler extends BaseCallbackHandler {
+  name = 'TurnBudgetHandler';
+  /** A refused reservation must fail the model call, not be logged and ignored. */
+  raiseError = true;
+  awaitHandlers = true;
+  private readonly reservations = new Map<
+    string,
+    { tokens: number; summary: boolean }
+  >();
+  /** Tool-schema estimate per tool list, for this turn (see `toolListKey`). */
+  private readonly toolTokens = new Map<string, number>();
+
+  constructor(private readonly options: BudgetedLlmOptions) {
+    super();
+  }
+
+  handleChatModelStart(
+    _llm: unknown,
+    messages: LLMResultMessages,
+    runId: string,
+    _parentRunId?: string,
+    extraParams?: Record<string, unknown>,
+    _tags?: string[],
+    metadata?: Record<string, unknown>,
+  ): void {
+    // The handler can reach the same call twice (model-level and run-level
+    // registration); the first registration charges it.
+    if (this.reservations.has(runId)) return;
+    const tokens = this.options.budget.reserveModel(
+      estimateRequestTokens(
+        messages,
+        extraParams?.invocation_params,
+        this.toolTokens,
+      ),
+      this.options.outputReserveTokens,
+      this.options.signal,
+    );
+    this.reservations.set(runId, {
+      tokens,
+      summary: metadata?.lc_source === SUMMARIZATION_SOURCE,
+    });
+  }
+
+  handleLLMEnd(output: LLMResult, runId: string): void {
+    const reservation = this.reservations.get(runId);
+    if (reservation === undefined) return;
+    this.reservations.delete(runId);
+    const reported = reportedTotalTokens(output);
+    if (reported !== undefined)
+      this.options.budget.settleModel(reservation.tokens, reported);
+  }
+
+  handleLLMError(_error: unknown, runId: string): void {
+    const reservation = this.reservations.get(runId);
+    this.reservations.delete(runId);
+    // A failed call keeps its reservation: the provider may have consumed
+    // the input before failing, and a retry reserves again. A failed summary
+    // is the exception: the summarizer swallows the failure, the turn goes on
+    // with its history unchanged and does not try again (summarization.ts),
+    // so its reservation — the whole history it was handed — is given back
+    // instead of starving the rest of the turn.
+    if (reservation?.summary)
+      this.options.budget.releaseModel(reservation.tokens);
+  }
+}
+
+function attach(model: BaseChatModel, handler: BaseCallbackHandler): void {
+  const existing = model.callbacks;
+  if (existing === undefined) {
+    model.callbacks = [handler];
+    return;
+  }
+  if (Array.isArray(existing)) {
+    if (!existing.includes(handler)) model.callbacks = [...existing, handler];
+    return;
+  }
+  const manager = existing.copy();
+  manager.addHandler(handler);
+  model.callbacks = manager;
+}
+
+export function budgetedLlm(
+  adapter: LlmAdapter,
+  options: BudgetedLlmOptions,
+): BudgetedLlm {
+  const callback = new TurnBudgetHandler(options);
+  return {
+    callback,
+    get(role: ModelRole, params?: ChatOpenAIFields): BaseChatModel {
+      const model = adapter.get(role, params);
+      attach(model, callback);
+      return model;
+    },
+  };
+}