rimuru-ai 1.19.0

package/src/acp/usage.ts

@@ -0,0 +1,232 @@
+import type { AgentSideConnection, Usage } from "@agentclientprotocol/sdk"
+import type { AssistantMessage as RimuruAssistantMessage, Message } from "@rimurucode-ai/sdk/v2"
+import { InstanceRef } from "@/effect/instance-ref"
+import { InstanceStore } from "@/project/instance-store"
+import { ProviderV2 } from "@rimurucode-ai/core/provider"
+import { ModelV2 } from "@rimurucode-ai/core/model"
+import { Provider } from "@/provider/provider"
+import { Context, Effect, Layer, SynchronizedRef } from "effect"
+
+export type AssistantTokenCost = Pick<RimuruAssistantMessage, "cost" | "tokens">
+
+export type AssistantMessage = AssistantTokenCost &
+  Pick<RimuruAssistantMessage, "role"> &
+  Partial<Pick<RimuruAssistantMessage, "providerID" | "modelID">>
+
+export type SessionMessage = {
+  readonly info: { readonly role: Message["role"] } | AssistantMessage
+}
+
+export type MessagesInput = {
+  readonly sessionID: string
+  readonly directory: string
+}
+
+export type SDK = {
+  readonly session: {
+    readonly messages: (
+      parameters: { readonly sessionID: string; readonly directory: string },
+      options: { readonly throwOnError: true },
+    ) => Promise<{ readonly data?: readonly SessionMessage[] | null }>
+  }
+}
+
+export interface MessageLoaderInterface {
+  readonly messages: (input: MessagesInput) => Effect.Effect<readonly SessionMessage[], unknown>
+}
+
+export interface ContextLimitLoaderInterface {
+  readonly providers: (directory: string) => Effect.Effect<Record<ProviderV2.ID, Provider.Info>, unknown>
+}
+
+export type UsageConnection = Pick<AgentSideConnection, "sessionUpdate">
+
+export interface Interface {
+  readonly buildUsage: (message: AssistantTokenCost) => Usage
+  readonly latestAssistantMessage: (messages: readonly SessionMessage[]) => AssistantMessage | undefined
+  readonly totalSessionCost: (messages: readonly SessionMessage[]) => number
+  readonly contextLimit: (input: {
+    readonly directory: string
+    readonly providerID: ProviderV2.ID
+    readonly modelID: ModelV2.ID
+  }) => Effect.Effect<number | undefined>
+  readonly sendUpdate: (input: {
+    readonly connection: UsageConnection
+    readonly sessionID: string
+    readonly directory: string
+  }) => Effect.Effect<void>
+}
+
+export class MessageLoader extends Context.Service<MessageLoader, MessageLoaderInterface>()(
+  "@rimuru/ACPUsageMessageLoader",
+) {}
+
+export class ContextLimitLoader extends Context.Service<ContextLimitLoader, ContextLimitLoaderInterface>()(
+  "@rimuru/ACPUsageContextLimitLoader",
+) {}
+
+export class Service extends Context.Service<Service, Interface>()("@rimuru/ACPUsage") {}
+
+export function messageLoaderFromSDK(sdk: SDK): MessageLoaderInterface {
+  return MessageLoader.of({
+    messages: (input) =>
+      Effect.promise(() =>
+        sdk.session
+          .messages({ sessionID: input.sessionID, directory: input.directory }, { throwOnError: true })
+          .then((response) => response.data ?? []),
+      ),
+  })
+}
+
+export const messageLoaderLayer = (sdk: SDK) => Layer.succeed(MessageLoader, messageLoaderFromSDK(sdk))
+
+export function buildUsage(message: AssistantTokenCost): Usage {
+  const cachedReadTokens = message.tokens.cache.read
+  const cachedWriteTokens = message.tokens.cache.write
+  const thoughtTokens = message.tokens.reasoning
+
+  return {
+    inputTokens: message.tokens.input,
+    outputTokens: message.tokens.output,
+    totalTokens: message.tokens.input + message.tokens.output + thoughtTokens + cachedReadTokens + cachedWriteTokens,
+    ...(thoughtTokens > 0 ? { thoughtTokens } : {}),
+    ...(cachedReadTokens > 0 ? { cachedReadTokens } : {}),
+    ...(cachedWriteTokens > 0 ? { cachedWriteTokens } : {}),
+  }
+}
+
+export function latestAssistantMessage(messages: readonly SessionMessage[]): AssistantMessage | undefined {
+  return messages
+    .filter((message): message is { readonly info: AssistantMessage } => message.info.role === "assistant")
+    .at(-1)?.info
+}
+
+export function totalSessionCost(messages: readonly SessionMessage[]): number {
+  return messages
+    .filter((message): message is { readonly info: AssistantMessage } => message.info.role === "assistant")
+    .reduce((sum, message) => sum + message.info.cost, 0)
+}
+
+export function findContextLimit(
+  providers: Record<ProviderV2.ID, Provider.Info>,
+  providerID: ProviderV2.ID,
+  modelID: ModelV2.ID,
+): number | undefined {
+  return providers[providerID]?.models[modelID]?.limit.context
+}
+
+export const contextLimitLoaderLayer = Layer.effect(
+  ContextLimitLoader,
+  Effect.gen(function* () {
+    const store = yield* InstanceStore.Service
+    const provider = yield* Provider.Service
+
+    return ContextLimitLoader.of({
+      providers: Effect.fn("ACPUsageContextLimitLoader.providers")(function* (directory) {
+        const ctx = yield* store.load({ directory })
+        return yield* Effect.gen(function* () {
+          return yield* provider.list()
+        }).pipe(Effect.provideService(InstanceRef, ctx))
+      }),
+    })
+  }),
+)
+
+export const layer = Layer.effect(
+  Service,
+  Effect.gen(function* () {
+    const messageLoader = yield* MessageLoader
+    const contextLimitLoader = yield* ContextLimitLoader
+    const limits = yield* SynchronizedRef.make(new Map<string, Effect.Effect<number | undefined>>())
+
+    const cachedLimit = Effect.fnUntraced(function* (input: {
+      readonly directory: string
+      readonly providerID: ProviderV2.ID
+      readonly modelID: ModelV2.ID
+    }) {
+      return yield* SynchronizedRef.modifyEffect(
+        limits,
+        Effect.fnUntraced(function* (items) {
+          const key = `${input.directory}\u0000${input.providerID}\u0000${input.modelID}`
+          const current = items.get(key)
+          if (current) return [current, items] as const
+          const next = yield* Effect.cached(
+            contextLimitLoader.providers(input.directory).pipe(
+              Effect.map((providers) => findContextLimit(providers, input.providerID, input.modelID)),
+              Effect.catch((error) =>
+                Effect.logError("failed to get providers for usage context limit", { error: error }).pipe(
+                  Effect.as(undefined),
+                ),
+              ),
+            ),
+          )
+          return [next, new Map(items).set(key, next)] as const
+        }),
+      )
+    })
+
+    const contextLimit = Effect.fn("ACPUsage.contextLimit")(function* (input: {
+      readonly directory: string
+      readonly providerID: ProviderV2.ID
+      readonly modelID: ModelV2.ID
+    }) {
+      return yield* yield* cachedLimit(input)
+    })
+
+    const sendUpdate = Effect.fn("ACPUsage.sendUpdate")(function* (input: {
+      readonly connection: UsageConnection
+      readonly sessionID: string
+      readonly directory: string
+    }) {
+      const messages = yield* messageLoader
+        .messages({ sessionID: input.sessionID, directory: input.directory })
+        .pipe(
+          Effect.catch((error) =>
+            Effect.logError("failed to fetch messages for usage update", { error: error }).pipe(Effect.as(undefined)),
+          ),
+        )
+      if (!messages) return
+
+      const message = latestAssistantMessage(messages)
+      if (!message) return
+      if (!message.providerID || !message.modelID) return
+
+      const size = yield* contextLimit({
+        directory: input.directory,
+        providerID: ProviderV2.ID.make(message.providerID),
+        modelID: ModelV2.ID.make(message.modelID),
+      })
+      if (!size) return
+
+      yield* Effect.promise(() =>
+        input.connection
+          .sessionUpdate({
+            sessionId: input.sessionID,
+            update: {
+              sessionUpdate: "usage_update",
+              used: message.tokens.input + message.tokens.cache.read,
+              size,
+              cost: { amount: totalSessionCost(messages), currency: "USD" },
+            },
+          })
+          .catch(() => {}),
+      )
+    })
+
+    return Service.of({
+      buildUsage,
+      latestAssistantMessage,
+      totalSessionCost,
+      contextLimit,
+      sendUpdate,
+    })
+  }),
+)
+
+export const defaultLayer = layer.pipe(
+  Layer.provide(contextLimitLoaderLayer),
+  Layer.provide(Provider.defaultLayer),
+  Layer.provide(InstanceStore.defaultLayer),
+)
+
+export * as UsageService from "./usage"

package/src/agent/agent.ts

@@ -0,0 +1,459 @@
+import { LayerNode } from "@rimurucode-ai/core/effect/layer-node"
+import { PermissionV1 } from "@rimurucode-ai/core/v1/permission"
+import { Config } from "@/config/config"
+import { serviceUse } from "@rimurucode-ai/core/effect/service-use"
+import { Provider } from "@/provider/provider"
+
+import { generateObject, streamObject, type ModelMessage } from "ai"
+import { Truncate } from "@/tool/truncate"
+import { Auth } from "../auth"
+import { ProviderTransform } from "@/provider/transform"
+
+import PROMPT_GENERATE from "./generate.txt"
+import PROMPT_COMPACTION from "./prompt/compaction.txt"
+import PROMPT_EXPLORE from "./prompt/explore.txt"
+import PROMPT_SUMMARY from "./prompt/summary.txt"
+import PROMPT_TITLE from "./prompt/title.txt"
+import { Permission } from "@/permission"
+import { mergeDeep, pipe, sortBy, values } from "remeda"
+import { Global } from "@rimurucode-ai/core/global"
+import path from "path"
+import { Plugin } from "@/plugin"
+import { Skill } from "../skill"
+import { Effect, Context, Layer, Schema } from "effect"
+import { InstanceState } from "@/effect/instance-state"
+import * as Option from "effect/Option"
+import * as OtelTracer from "@effect/opentelemetry/Tracer"
+import { AbsolutePath, type DeepMutable } from "@rimurucode-ai/core/schema"
+import { ProviderV2 } from "@rimurucode-ai/core/provider"
+import { ModelV2 } from "@rimurucode-ai/core/model"
+import { LocationServiceMap } from "@rimurucode-ai/core/location-layer"
+import { PluginBoot } from "@rimurucode-ai/core/plugin/boot"
+import { Reference } from "@rimurucode-ai/core/reference"
+import { Location } from "@rimurucode-ai/core/location"
+
+export const Info = Schema.Struct({
+  name: Schema.String,
+  description: Schema.optional(Schema.String),
+  mode: Schema.Literals(["subagent", "primary", "all"]),
+  native: Schema.optional(Schema.Boolean),
+  hidden: Schema.optional(Schema.Boolean),
+  topP: Schema.optional(Schema.Finite),
+  temperature: Schema.optional(Schema.Finite),
+  color: Schema.optional(Schema.String),
+  permission: PermissionV1.Ruleset,
+  model: Schema.optional(
+    Schema.Struct({
+      modelID: ModelV2.ID,
+      providerID: ProviderV2.ID,
+    }),
+  ),
+  variant: Schema.optional(Schema.String),
+  prompt: Schema.optional(Schema.String),
+  options: Schema.Record(Schema.String, Schema.Unknown),
+  steps: Schema.optional(Schema.Finite),
+}).annotate({ identifier: "Agent" })
+export type Info = DeepMutable<Schema.Schema.Type<typeof Info>>
+
+const GeneratedAgent = Schema.Struct({
+  identifier: Schema.String,
+  whenToUse: Schema.String,
+  systemPrompt: Schema.String,
+})
+
+export interface Interface {
+  readonly get: (agent: string) => Effect.Effect<Info>
+  readonly list: () => Effect.Effect<Info[]>
+  readonly defaultInfo: () => Effect.Effect<Info>
+  readonly defaultAgent: () => Effect.Effect<string>
+  readonly generate: (input: {
+    description: string
+    model?: { providerID: ProviderV2.ID; modelID: ModelV2.ID }
+  }) => Effect.Effect<
+    {
+      identifier: string
+      whenToUse: string
+      systemPrompt: string
+    },
+    Provider.DefaultModelError
+  >
+}
+
+type State = Omit<Interface, "generate">
+
+export class Service extends Context.Service<Service, Interface>()("@rimuru/Agent") {}
+
+export const use = serviceUse(Service)
+
+export const layer = Layer.effect(
+  Service,
+  Effect.gen(function* () {
+    const config = yield* Config.Service
+    const auth = yield* Auth.Service
+    const plugin = yield* Plugin.Service
+    const skill = yield* Skill.Service
+    const provider = yield* Provider.Service
+    const locations = yield* LocationServiceMap
+
+    const state = yield* InstanceState.make<State>(
+      Effect.fn("Agent.state")(function* (ctx) {
+        const cfg = yield* config.get()
+        const skillDirs = yield* skill.dirs()
+        const referenceDirs = yield* Effect.gen(function* () {
+          yield* (yield* PluginBoot.Service).wait()
+          return (yield* (yield* Reference.Service).list()).map((reference) => reference.path)
+        }).pipe(Effect.provide(locations.get(Location.Ref.make({ directory: AbsolutePath.make(ctx.directory) }))))
+        const whitelistedDirs = [
+          Truncate.GLOB,
+          path.join(Global.Path.tmp, "*"),
+          ...skillDirs.map((dir) => path.join(dir, "*")),
+          ...referenceDirs.map((dir) => path.join(dir, "*")),
+        ]
+        const readonlyExternalDirectory = {
+          "*": "ask",
+          ...Object.fromEntries(whitelistedDirs.map((dir) => [dir, "allow"])),
+        } satisfies Record<string, "allow" | "ask" | "deny">
+
+        const defaults = Permission.fromConfig({
+          "*": "allow",
+          doom_loop: "ask",
+          external_directory: {
+            "*": "ask",
+            ...Object.fromEntries(whitelistedDirs.map((dir) => [dir, "allow"])),
+          },
+          question: "deny",
+          plan_enter: "deny",
+          plan_exit: "deny",
+          // mirrors github.com/github/gitignore Node.gitignore pattern for .env files
+          read: {
+            "*": "allow",
+            "*.env": "ask",
+            "*.env.*": "ask",
+            "*.env.example": "allow",
+          },
+        })
+
+        const user = Permission.fromConfig(cfg.permission ?? {})
+
+        const agents: Record<string, Info> = {
+          build: {
+            name: "build",
+            description: "The default agent. Executes tools based on configured permissions.",
+            options: {},
+            permission: Permission.merge(
+              defaults,
+              Permission.fromConfig({
+                question: "allow",
+                plan_enter: "allow",
+              }),
+              user,
+            ),
+            mode: "primary",
+            native: true,
+          },
+          plan: {
+            name: "plan",
+            description: "Plan mode. Disallows all edit tools.",
+            options: {},
+            permission: Permission.merge(
+              defaults,
+              Permission.fromConfig({
+                question: "allow",
+                plan_exit: "allow",
+                task: {
+                  general: "deny",
+                },
+                external_directory: {
+                  [path.join(Global.Path.data, "plans", "*")]: "allow",
+                },
+                edit: {
+                  "*": "deny",
+                  [path.join(".rimuru", "plans", "*.md")]: "allow",
+                  [path.relative(ctx.worktree, path.join(Global.Path.data, path.join("plans", "*.md")))]: "allow",
+                },
+              }),
+              user,
+            ),
+            mode: "primary",
+            native: true,
+          },
+          general: {
+            name: "general",
+            description: `General-purpose agent for researching complex questions and executing multi-step tasks. Use this agent to execute multiple units of work in parallel.`,
+            permission: Permission.merge(
+              defaults,
+              Permission.fromConfig({
+                todowrite: "deny",
+              }),
+              user,
+            ),
+            options: {},
+            mode: "subagent",
+            native: true,
+          },
+          explore: {
+            name: "explore",
+            permission: Permission.merge(
+              defaults,
+              Permission.fromConfig({
+                "*": "deny",
+                grep: "allow",
+                glob: "allow",
+                list: "allow",
+                bash: "allow",
+                webfetch: "allow",
+                websearch: "allow",
+                read: "allow",
+                external_directory: readonlyExternalDirectory,
+              }),
+              user,
+            ),
+            description: `Fast agent specialized for exploring codebases. Use this when you need to quickly find files by patterns (eg. "src/components/**/*.tsx"), search code for keywords (eg. "API endpoints"), or answer questions about the codebase (eg. "how do API endpoints work?"). When calling this agent, specify the desired thoroughness level: "quick" for basic searches, "medium" for moderate exploration, or "very thorough" for comprehensive analysis across multiple locations and naming conventions.`,
+            prompt: PROMPT_EXPLORE,
+            options: {},
+            mode: "subagent",
+            native: true,
+          },
+          compaction: {
+            name: "compaction",
+            mode: "primary",
+            native: true,
+            hidden: true,
+            prompt: PROMPT_COMPACTION,
+            permission: Permission.merge(
+              defaults,
+              Permission.fromConfig({
+                "*": "deny",
+              }),
+              user,
+            ),
+            options: {},
+          },
+          title: {
+            name: "title",
+            mode: "primary",
+            options: {},
+            native: true,
+            hidden: true,
+            temperature: 0.5,
+            permission: Permission.merge(
+              defaults,
+              Permission.fromConfig({
+                "*": "deny",
+              }),
+              user,
+            ),
+            prompt: PROMPT_TITLE,
+          },
+          summary: {
+            name: "summary",
+            mode: "primary",
+            options: {},
+            native: true,
+            hidden: true,
+            permission: Permission.merge(
+              defaults,
+              Permission.fromConfig({
+                "*": "deny",
+              }),
+              user,
+            ),
+            prompt: PROMPT_SUMMARY,
+          },
+        }
+
+        for (const [key, value] of Object.entries(cfg.agent ?? {})) {
+          if (value.disable) {
+            delete agents[key]
+            continue
+          }
+          let item = agents[key]
+          if (!item)
+            item = agents[key] = {
+              name: key,
+              mode: "all",
+              permission: Permission.merge(defaults, user),
+              options: {},
+              native: false,
+            }
+          if (value.model) item.model = Provider.parseModel(value.model)
+          item.variant = value.variant ?? item.variant
+          item.prompt = value.prompt ?? item.prompt
+          item.description = value.description ?? item.description
+          item.temperature = value.temperature ?? item.temperature
+          item.topP = value.top_p ?? item.topP
+          item.mode = value.mode ?? item.mode
+          item.color = value.color ?? item.color
+          item.hidden = value.hidden ?? item.hidden
+          item.name = value.name ?? item.name
+          item.steps = value.steps ?? item.steps
+          item.options = mergeDeep(item.options, value.options ?? {})
+          item.permission = Permission.merge(item.permission, Permission.fromConfig(value.permission ?? {}))
+        }
+
+        // Ensure Truncate.GLOB is allowed unless explicitly configured
+        for (const name in agents) {
+          const agent = agents[name]
+          const explicit = agent.permission.some((r) => {
+            if (r.permission !== "external_directory") return false
+            if (r.action !== "deny") return false
+            return r.pattern === Truncate.GLOB
+          })
+          if (explicit) continue
+
+          agents[name].permission = Permission.merge(
+            agents[name].permission,
+            Permission.fromConfig({ external_directory: { [Truncate.GLOB]: "allow" } }),
+          )
+        }
+
+        const get = Effect.fnUntraced(function* (agent: string) {
+          return agents[agent]
+        })
+
+        const list = Effect.fnUntraced(function* () {
+          const cfg = yield* config.get()
+          return pipe(
+            agents,
+            values(),
+            sortBy(
+              [(x) => (cfg.default_agent ? x.name === cfg.default_agent : x.name === "build"), "desc"],
+              [(x) => x.name, "asc"],
+            ),
+          )
+        })
+
+        const defaultInfo = Effect.fnUntraced(function* () {
+          const c = yield* config.get()
+          if (c.default_agent) {
+            const agent = agents[c.default_agent]
+            if (!agent) throw new Error(`default agent "${c.default_agent}" not found`)
+            if (agent.mode === "subagent") throw new Error(`default agent "${c.default_agent}" is a subagent`)
+            if (agent.hidden === true) throw new Error(`default agent "${c.default_agent}" is hidden`)
+            return agent
+          }
+          const visible = Object.values(agents).find((a) => a.mode !== "subagent" && a.hidden !== true)
+          if (!visible) throw new Error("no primary visible agent found")
+          return visible
+        })
+
+        const defaultAgent = Effect.fnUntraced(function* () {
+          return (yield* defaultInfo()).name
+        })
+
+        return {
+          get,
+          list,
+          defaultInfo,
+          defaultAgent,
+        } satisfies State
+      }),
+    )
+
+    return Service.of({
+      get: Effect.fn("Agent.get")(function* (agent: string) {
+        return yield* InstanceState.useEffect(state, (s) => s.get(agent))
+      }),
+      list: Effect.fn("Agent.list")(function* () {
+        return yield* InstanceState.useEffect(state, (s) => s.list())
+      }),
+      defaultInfo: Effect.fn("Agent.defaultInfo")(function* () {
+        return yield* InstanceState.useEffect(state, (s) => s.defaultInfo())
+      }),
+      defaultAgent: Effect.fn("Agent.defaultAgent")(function* () {
+        return yield* InstanceState.useEffect(state, (s) => s.defaultAgent())
+      }),
+      generate: Effect.fn("Agent.generate")(function* (input: {
+        description: string
+        model?: { providerID: ProviderV2.ID; modelID: ModelV2.ID }
+      }) {
+        const cfg = yield* config.get()
+        const model = input.model ?? (yield* provider.defaultModel())
+        const resolved = yield* provider.getModel(model.providerID, model.modelID)
+        const language = yield* provider.getLanguage(resolved)
+        const tracer = cfg.experimental?.openTelemetry
+          ? Option.getOrUndefined(yield* Effect.serviceOption(OtelTracer.OtelTracer))
+          : undefined
+
+        const system = [PROMPT_GENERATE]
+        yield* plugin.trigger("experimental.chat.system.transform", { model: resolved }, { system })
+        const existing = yield* InstanceState.useEffect(state, (s) => s.list())
+
+        // TODO: clean this up so provider specific logic doesnt bleed over
+        const authInfo = yield* auth.get(model.providerID).pipe(Effect.orDie)
+        const isOpenaiOauth = model.providerID === "openai" && authInfo?.type === "oauth"
+
+        const params = {
+          experimental_telemetry: {
+            isEnabled: cfg.experimental?.openTelemetry,
+            tracer,
+            metadata: {
+              userId: cfg.username ?? "unknown",
+            },
+          },
+          temperature: 0.3,
+          messages: [
+            ...(isOpenaiOauth
+              ? []
+              : system.map(
+                  (item): ModelMessage => ({
+                    role: "system",
+                    content: item,
+                  }),
+                )),
+            {
+              role: "user",
+              content: `Create an agent configuration based on this request: "${input.description}".\n\nIMPORTANT: The following identifiers already exist and must NOT be used: ${existing.map((i) => i.name).join(", ")}\n  Return ONLY the JSON object, no other text, do not wrap in backticks`,
+            },
+          ],
+          model: language,
+          schema: Object.assign(
+            Schema.toStandardSchemaV1(GeneratedAgent),
+            Schema.toStandardJSONSchemaV1(GeneratedAgent),
+          ),
+        } satisfies Parameters<typeof generateObject>[0]
+
+        if (isOpenaiOauth) {
+          return yield* Effect.promise(async () => {
+            const result = streamObject({
+              ...params,
+              providerOptions: ProviderTransform.providerOptions(resolved, {
+                instructions: system.join("\n"),
+                store: false,
+              }),
+              onError: () => {},
+            })
+            for await (const part of result.fullStream) {
+              if (part.type === "error") throw part.error
+            }
+            return result.object
+          })
+        }
+
+        return yield* Effect.promise(() => generateObject(params).then((r) => r.object))
+      }),
+    })
+  }),
+)
+
+export const defaultLayer = layer.pipe(
+  Layer.provide(Plugin.defaultLayer),
+  Layer.provide(Provider.defaultLayer),
+  Layer.provide(Auth.defaultLayer),
+  Layer.provide(Config.defaultLayer),
+  Layer.provide(Skill.defaultLayer),
+  Layer.provide(LocationServiceMap.layer),
+)
+
+const locationServiceMapNode = LayerNode.make(LocationServiceMap.layer, [])
+
+export const node = LayerNode.make(layer, [
+  Config.node,
+  Auth.node,
+  Plugin.node,
+  Skill.node,
+  Provider.node,
+  locationServiceMapNode,
+])
+
+export * as Agent from "./agent"