rird 1.0.200

package/src/session/index.ts

@@ -0,0 +1,461 @@
+import { BusEvent } from "@/bus/bus-event"
+import { Bus } from "@/bus"
+import { Decimal } from "decimal.js"
+import z from "zod"
+import { type LanguageModelUsage, type ProviderMetadata } from "ai"
+import { Config } from "../config/config"
+import { Flag } from "../flag/flag"
+import { Identifier } from "../id/id"
+import { Installation } from "../installation"
+
+import { Storage } from "../storage/storage"
+import { Log } from "../util/log"
+import { MessageV2 } from "./message-v2"
+import { Instance } from "../project/instance"
+import { SessionPrompt } from "./prompt"
+import { fn } from "@/util/fn"
+import { Command } from "../command"
+import { Snapshot } from "@/snapshot"
+import { SensitiveFilter } from "./sensitive-filter"
+
+import type { Provider } from "@/provider/provider"
+
+export namespace Session {
+  const log = Log.create({ service: "session" })
+
+  const parentTitlePrefix = "New session - "
+  const childTitlePrefix = "Child session - "
+
+  function createDefaultTitle(isChild = false) {
+    return (isChild ? childTitlePrefix : parentTitlePrefix) + new Date().toISOString()
+  }
+
+  export function isDefaultTitle(title: string) {
+    return new RegExp(
+      `^(${parentTitlePrefix}|${childTitlePrefix})\\d{4}-\\d{2}-\\d{2}T\\d{2}:\\d{2}:\\d{2}\\.\\d{3}Z$`,
+    ).test(title)
+  }
+
+  export const Info = z
+    .object({
+      id: Identifier.schema("session"),
+      projectID: z.string(),
+      directory: z.string(),
+      parentID: Identifier.schema("session").optional(),
+      summary: z
+        .object({
+          additions: z.number(),
+          deletions: z.number(),
+          files: z.number(),
+          diffs: Snapshot.FileDiff.array().optional(),
+        })
+        .optional(),
+      share: z
+        .object({
+          url: z.string(),
+        })
+        .optional(),
+      title: z.string(),
+      version: z.string(),
+      time: z.object({
+        created: z.number(),
+        updated: z.number(),
+        compacting: z.number().optional(),
+        archived: z.number().optional(),
+      }),
+      revert: z
+        .object({
+          messageID: z.string(),
+          partID: z.string().optional(),
+          snapshot: z.string().optional(),
+          diff: z.string().optional(),
+        })
+        .optional(),
+    })
+    .meta({
+      ref: "Session",
+    })
+  export type Info = z.output<typeof Info>
+
+  export const ShareInfo = z
+    .object({
+      secret: z.string(),
+      url: z.string(),
+    })
+    .meta({
+      ref: "SessionShare",
+    })
+  export type ShareInfo = z.output<typeof ShareInfo>
+
+  export const Event = {
+    Created: BusEvent.define(
+      "session.created",
+      z.object({
+        info: Info,
+      }),
+    ),
+    Updated: BusEvent.define(
+      "session.updated",
+      z.object({
+        info: Info,
+      }),
+    ),
+    Deleted: BusEvent.define(
+      "session.deleted",
+      z.object({
+        info: Info,
+      }),
+    ),
+    Diff: BusEvent.define(
+      "session.diff",
+      z.object({
+        sessionID: z.string(),
+        diff: Snapshot.FileDiff.array(),
+      }),
+    ),
+    Error: BusEvent.define(
+      "session.error",
+      z.object({
+        sessionID: z.string().optional(),
+        error: MessageV2.Assistant.shape.error,
+      }),
+    ),
+  }
+
+  export const create = fn(
+    z
+      .object({
+        parentID: Identifier.schema("session").optional(),
+        title: z.string().optional(),
+      })
+      .optional(),
+    async (input) => {
+      return createNext({
+        parentID: input?.parentID,
+        directory: Instance.directory,
+        title: input?.title,
+      })
+    },
+  )
+
+  export const fork = fn(
+    z.object({
+      sessionID: Identifier.schema("session"),
+      messageID: Identifier.schema("message").optional(),
+    }),
+    async (input) => {
+      const session = await createNext({
+        directory: Instance.directory,
+      })
+      const msgs = await messages({ sessionID: input.sessionID })
+      for (const msg of msgs) {
+        if (input.messageID && msg.info.id >= input.messageID) break
+        const cloned = await updateMessage({
+          ...msg.info,
+          sessionID: session.id,
+          id: Identifier.ascending("message"),
+        })
+
+        for (const part of msg.parts) {
+          await updatePart({
+            ...part,
+            id: Identifier.ascending("part"),
+            messageID: cloned.id,
+            sessionID: session.id,
+          })
+        }
+      }
+      return session
+    },
+  )
+
+  export const touch = fn(Identifier.schema("session"), async (sessionID) => {
+    await update(sessionID, (draft) => {
+      draft.time.updated = Date.now()
+    })
+  })
+
+  export async function createNext(input: { id?: string; title?: string; parentID?: string; directory: string }) {
+    const result: Info = {
+      id: Identifier.descending("session", input.id),
+      version: Installation.VERSION,
+      projectID: Instance.project.id,
+      directory: input.directory,
+      parentID: input.parentID,
+      title: input.title ?? createDefaultTitle(!!input.parentID),
+      time: {
+        created: Date.now(),
+        updated: Date.now(),
+      },
+    }
+    log.info("created", result)
+    await Storage.write(["session", Instance.project.id, result.id], result)
+    Bus.publish(Event.Created, {
+      info: result,
+    })
+    const cfg = await Config.get()
+    if (!result.parentID && (Flag.RIRD_AUTO_SHARE || cfg.share === "auto"))
+      share(result.id)
+        .then((share) => {
+          update(result.id, (draft) => {
+            draft.share = share
+          })
+        })
+        .catch(() => {
+          // Silently ignore sharing errors during session creation
+        })
+    Bus.publish(Event.Updated, {
+      info: result,
+    })
+    return result
+  }
+
+  export const get = fn(Identifier.schema("session"), async (id) => {
+    const read = await Storage.read<Info>(["session", Instance.project.id, id])
+    return read as Info
+  })
+
+  export const getShare = fn(Identifier.schema("session"), async (id) => {
+    return Storage.read<ShareInfo>(["share", id])
+  })
+
+  export const share = fn(Identifier.schema("session"), async (id) => {
+    const cfg = await Config.get()
+    if (cfg.share === "disabled") {
+      throw new Error("Sharing is disabled in configuration")
+    }
+    const { ShareNext } = await import("@/share/share-next")
+    const share = await ShareNext.create(id)
+    await update(id, (draft) => {
+      draft.share = {
+        url: share.url,
+      }
+    })
+    return share
+  })
+
+  export const unshare = fn(Identifier.schema("session"), async (id) => {
+    // Use ShareNext to remove the share (same as share function uses ShareNext to create)
+    const { ShareNext } = await import("@/share/share-next")
+    await ShareNext.remove(id)
+    await update(id, (draft) => {
+      draft.share = undefined
+    })
+  })
+
+  export async function update(id: string, editor: (session: Info) => void) {
+    const project = Instance.project
+    const result = await Storage.update<Info>(["session", project.id, id], (draft) => {
+      editor(draft)
+      draft.time.updated = Date.now()
+    })
+    Bus.publish(Event.Updated, {
+      info: result,
+    })
+    return result
+  }
+
+  export const diff = fn(Identifier.schema("session"), async (sessionID) => {
+    const diffs = await Storage.read<Snapshot.FileDiff[]>(["session_diff", sessionID])
+    return diffs ?? []
+  })
+
+  export const messages = fn(
+    z.object({
+      sessionID: Identifier.schema("session"),
+      limit: z.number().optional(),
+    }),
+    async (input) => {
+      const result = [] as MessageV2.WithParts[]
+      for await (const msg of MessageV2.stream(input.sessionID)) {
+        if (input.limit && result.length >= input.limit) break
+        result.push(msg)
+      }
+      result.reverse()
+      return result
+    },
+  )
+
+  export async function* list() {
+    const project = Instance.project
+    for (const item of await Storage.list(["session", project.id])) {
+      yield Storage.read<Info>(item)
+    }
+  }
+
+  export const children = fn(Identifier.schema("session"), async (parentID) => {
+    const project = Instance.project
+    const result = [] as Session.Info[]
+    for (const item of await Storage.list(["session", project.id])) {
+      const session = await Storage.read<Info>(item)
+      if (session.parentID !== parentID) continue
+      result.push(session)
+    }
+    return result
+  })
+
+  export const remove = fn(Identifier.schema("session"), async (sessionID) => {
+    const project = Instance.project
+    try {
+      const session = await get(sessionID)
+      for (const child of await children(sessionID)) {
+        await remove(child.id)
+      }
+      await unshare(sessionID).catch(() => {})
+      for (const msg of await Storage.list(["message", sessionID])) {
+        for (const part of await Storage.list(["part", msg.at(-1)!])) {
+          await Storage.remove(part)
+        }
+        await Storage.remove(msg)
+      }
+      await Storage.remove(["session", project.id, sessionID])
+      // Clean up sensitive data mappings for this session
+      SensitiveFilter.clearSession(sessionID)
+      Bus.publish(Event.Deleted, {
+        info: session,
+      })
+    } catch (e) {
+      log.error(e)
+    }
+  })
+
+  export const updateMessage = fn(MessageV2.Info, async (msg) => {
+    await Storage.write(["message", msg.sessionID, msg.id], msg)
+    Bus.publish(MessageV2.Event.Updated, {
+      info: msg,
+    })
+    return msg
+  })
+
+  export const removeMessage = fn(
+    z.object({
+      sessionID: Identifier.schema("session"),
+      messageID: Identifier.schema("message"),
+    }),
+    async (input) => {
+      await Storage.remove(["message", input.sessionID, input.messageID])
+      Bus.publish(MessageV2.Event.Removed, {
+        sessionID: input.sessionID,
+        messageID: input.messageID,
+      })
+      return input.messageID
+    },
+  )
+
+  export const removePart = fn(
+    z.object({
+      sessionID: Identifier.schema("session"),
+      messageID: Identifier.schema("message"),
+      partID: Identifier.schema("part"),
+    }),
+    async (input) => {
+      await Storage.remove(["part", input.messageID, input.partID])
+      Bus.publish(MessageV2.Event.PartRemoved, {
+        sessionID: input.sessionID,
+        messageID: input.messageID,
+        partID: input.partID,
+      })
+      return input.partID
+    },
+  )
+
+  const UpdatePartInput = z.union([
+    MessageV2.Part,
+    z.object({
+      part: MessageV2.TextPart,
+      delta: z.string(),
+    }),
+    z.object({
+      part: MessageV2.ReasoningPart,
+      delta: z.string(),
+    }),
+  ])
+
+  export const updatePart = fn(UpdatePartInput, async (input) => {
+    const part = "delta" in input ? input.part : input
+    const delta = "delta" in input ? input.delta : undefined
+    await Storage.write(["part", part.messageID, part.id], part)
+    Bus.publish(MessageV2.Event.PartUpdated, {
+      part,
+      delta,
+    })
+    return part
+  })
+
+  export const getUsage = fn(
+    z.object({
+      model: z.custom<Provider.Model>(),
+      usage: z.custom<LanguageModelUsage>(),
+      metadata: z.custom<ProviderMetadata>().optional(),
+    }),
+    (input) => {
+      const cachedInputTokens = input.usage.cachedInputTokens ?? 0
+      const excludesCachedTokens = !!(input.metadata?.["anthropic"] || input.metadata?.["bedrock"])
+      const adjustedInputTokens = excludesCachedTokens
+        ? (input.usage.inputTokens ?? 0)
+        : (input.usage.inputTokens ?? 0) - cachedInputTokens
+      const safe = (value: number) => {
+        if (!Number.isFinite(value)) return 0
+        return value
+      }
+
+      const tokens = {
+        input: safe(adjustedInputTokens),
+        output: safe(input.usage.outputTokens ?? 0),
+        reasoning: safe(input.usage?.reasoningTokens ?? 0),
+        cache: {
+          write: safe(
+            (input.metadata?.["anthropic"]?.["cacheCreationInputTokens"] ??
+              // @ts-expect-error
+              input.metadata?.["bedrock"]?.["usage"]?.["cacheWriteInputTokens"] ??
+              0) as number,
+          ),
+          read: safe(cachedInputTokens),
+        },
+      }
+
+      const costInfo =
+        input.model.cost?.experimentalOver200K && tokens.input + tokens.cache.read > 200_000
+          ? input.model.cost.experimentalOver200K
+          : input.model.cost
+      return {
+        cost: safe(
+          new Decimal(0)
+            .add(new Decimal(tokens.input).mul(costInfo?.input ?? 0).div(1_000_000))
+            .add(new Decimal(tokens.output).mul(costInfo?.output ?? 0).div(1_000_000))
+            .add(new Decimal(tokens.cache.read).mul(costInfo?.cache?.read ?? 0).div(1_000_000))
+            .add(new Decimal(tokens.cache.write).mul(costInfo?.cache?.write ?? 0).div(1_000_000))
+            // TODO: update models.dev to have better pricing model, for now:
+            // charge reasoning tokens at the same rate as output tokens
+            .add(new Decimal(tokens.reasoning).mul(costInfo?.output ?? 0).div(1_000_000))
+            .toNumber(),
+        ),
+        tokens,
+      }
+    },
+  )
+
+  export class BusyError extends Error {
+    constructor(public readonly sessionID: string) {
+      super(`Session ${sessionID} is busy`)
+    }
+  }
+
+  export const initialize = fn(
+    z.object({
+      sessionID: Identifier.schema("session"),
+      modelID: z.string(),
+      providerID: z.string(),
+      messageID: Identifier.schema("message"),
+    }),
+    async (input) => {
+      await SessionPrompt.command({
+        sessionID: input.sessionID,
+        messageID: input.messageID,
+        model: input.providerID + "/" + input.modelID,
+        command: Command.Default.INIT,
+        arguments: "",
+      })
+    },
+  )
+}

package/src/session/llm.ts

@@ -0,0 +1,201 @@
+import { Provider } from "@/provider/provider"
+import { Log } from "@/util/log"
+import { streamText, wrapLanguageModel, type ModelMessage, type StreamTextResult, type Tool, type ToolSet } from "ai"
+import { clone, mergeDeep, pipe } from "remeda"
+import { ProviderTransform } from "@/provider/transform"
+import { Config } from "@/config/config"
+import { Instance } from "@/project/instance"
+import type { Agent } from "@/agent/agent"
+import type { MessageV2 } from "./message-v2"
+import { Plugin } from "@/plugin"
+import { SystemPrompt } from "./system"
+import { ToolRegistry } from "@/tool/registry"
+import { Flag } from "@/flag/flag"
+
+export namespace LLM {
+  const log = Log.create({ service: "llm" })
+
+  export const OUTPUT_TOKEN_MAX = Flag.RIRD_EXPERIMENTAL_OUTPUT_TOKEN_MAX || 32_000
+
+  export type StreamInput = {
+    user: MessageV2.User
+    sessionID: string
+    model: Provider.Model
+    agent: Agent.Info
+    system: string[]
+    abort: AbortSignal
+    messages: ModelMessage[]
+    small?: boolean
+    tools: Record<string, Tool>
+    retries?: number
+  }
+
+  export type StreamOutput = StreamTextResult<ToolSet, unknown>
+
+  export async function stream(input: StreamInput) {
+    const l = log
+      .clone()
+      .tag("providerID", input.model.providerID)
+      .tag("modelID", input.model.id)
+      .tag("sessionID", input.sessionID)
+      .tag("small", (input.small ?? false).toString())
+      .tag("agent", input.agent.name)
+    l.info("stream", {
+      modelID: input.model.id,
+      providerID: input.model.providerID,
+    })
+    const [language, cfg] = await Promise.all([Provider.getLanguage(input.model), Config.get()])
+
+    const system = SystemPrompt.header(input.model.providerID)
+    system.push(
+      [
+        // use agent prompt otherwise provider prompt
+        ...(input.agent.prompt ? [input.agent.prompt] : SystemPrompt.provider(input.model)),
+        // any custom prompt passed into this call
+        ...input.system,
+        // any custom prompt from last user message
+        ...(input.user.system ? [input.user.system] : []),
+      ]
+        .filter((x) => x)
+        .join("\n"),
+    )
+
+    const header = system[0]
+    const original = clone(system)
+    await Plugin.trigger("experimental.chat.system.transform", {}, { system })
+    if (system.length === 0) {
+      system.push(...original)
+    }
+    // rejoin to maintain 2-part structure for caching if header unchanged
+    if (system.length > 2 && system[0] === header) {
+      const rest = system.slice(1)
+      system.length = 0
+      system.push(header, rest.join("\n"))
+    }
+
+    const provider = await Provider.getProvider(input.model.providerID)
+
+    const params = await Plugin.trigger(
+      "chat.params",
+      {
+        sessionID: input.sessionID,
+        agent: input.agent,
+        model: input.model,
+        provider: Provider.getProvider(input.model.providerID),
+        message: input.user,
+      },
+      {
+        temperature: input.model.capabilities.temperature
+          ? (input.agent.temperature ?? ProviderTransform.temperature(input.model))
+          : undefined,
+        topP: input.agent.topP ?? ProviderTransform.topP(input.model),
+        topK: ProviderTransform.topK(input.model),
+        options: pipe(
+          {},
+          mergeDeep(ProviderTransform.options(input.model, input.sessionID, provider.options)),
+          input.small ? mergeDeep(ProviderTransform.smallOptions(input.model)) : mergeDeep({}),
+          mergeDeep(input.model.options),
+          mergeDeep(input.agent.options),
+        ),
+      },
+    )
+
+    l.info("params", {
+      params,
+    })
+
+    const maxOutputTokens = ProviderTransform.maxOutputTokens(
+      input.model.api.npm,
+      params.options,
+      input.model.limit.output,
+      OUTPUT_TOKEN_MAX,
+    )
+
+    const tools = await resolveTools(input)
+
+    return streamText({
+      onError(error) {
+        l.error("stream error", {
+          error,
+        })
+      },
+      async experimental_repairToolCall(failed) {
+        const lower = failed.toolCall.toolName.toLowerCase()
+        if (lower !== failed.toolCall.toolName && tools[lower]) {
+          l.info("repairing tool call", {
+            tool: failed.toolCall.toolName,
+            repaired: lower,
+          })
+          return {
+            ...failed.toolCall,
+            toolName: lower,
+          }
+        }
+        return {
+          ...failed.toolCall,
+          input: JSON.stringify({
+            tool: failed.toolCall.toolName,
+            error: failed.error.message,
+          }),
+          toolName: "invalid",
+        }
+      },
+      temperature: params.temperature,
+      topP: params.topP,
+      topK: params.topK,
+      providerOptions: ProviderTransform.providerOptions(input.model, params.options),
+      activeTools: Object.keys(tools).filter((x) => x !== "invalid"),
+      tools,
+      maxOutputTokens,
+      abortSignal: input.abort,
+      headers: {
+        ...(input.model.providerID.startsWith("rird")
+          ? {
+              "x-rird-project": Instance.project.id,
+              "x-rird-session": input.sessionID,
+              "x-rird-request": input.user.id,
+              "x-rird-client": Flag.RIRD_CLIENT,
+            }
+          : undefined),
+        ...input.model.headers,
+      },
+      maxRetries: input.retries ?? 0,
+      messages: [
+        ...system.map(
+          (x): ModelMessage => ({
+            role: "system",
+            content: x,
+          }),
+        ),
+        ...input.messages,
+      ],
+      model: wrapLanguageModel({
+        model: language,
+        middleware: [
+          {
+            async transformParams(args) {
+              if (args.type === "stream") {
+                // @ts-expect-error
+                args.params.prompt = ProviderTransform.message(args.params.prompt, input.model)
+              }
+              return args.params
+            },
+          },
+        ],
+      }),
+      experimental_telemetry: { isEnabled: cfg.experimental?.openTelemetry },
+    })
+  }
+
+  async function resolveTools(input: Pick<StreamInput, "tools" | "agent" | "user">) {
+    const enabled = pipe(
+      input.agent.tools,
+      mergeDeep(await ToolRegistry.enabled(input.agent)),
+      mergeDeep(input.user.tools ?? {}),
+    )
+    for (const [key, value] of Object.entries(enabled)) {
+      if (value === false) delete input.tools[key]
+    }
+    return input.tools
+  }
+}