@farhanic017/octocode 3.5.0

package/src/acp/tool.ts

@@ -0,0 +1,321 @@
+import type { ToolCall, ToolCallContent, ToolCallLocation, ToolCallUpdate, ToolKind } from "@agentclientprotocol/sdk"
+
+export type ToolInput = Record<string, unknown>
+
+export type ToolAttachment = {
+  readonly mime?: string
+  readonly url?: string
+  readonly [key: string]: unknown
+}
+
+export type CompletedToolState = {
+  readonly status: "completed"
+  readonly input: ToolInput
+  readonly output: string
+  readonly metadata?: unknown
+  readonly attachments?: ReadonlyArray<ToolAttachment>
+}
+
+export type RunningToolState = {
+  readonly status: "running"
+  readonly input: ToolInput
+  readonly title?: string
+}
+
+export type ErrorToolState = {
+  readonly status: "error"
+  readonly input: ToolInput
+  readonly error: string
+  readonly metadata?: unknown
+}
+
+export type ImageAttachment = {
+  readonly mimeType: string
+  readonly data: string
+}
+
+export function toToolKind(toolName: string): ToolKind {
+  const tool = toolName.toLocaleLowerCase()
+
+  switch (tool) {
+    case "bash":
+    case "shell":
+      return "execute"
+
+    case "webfetch":
+      return "fetch"
+
+    case "edit":
+    case "apply_patch":
+    case "patch":
+    case "write":
+      return "edit"
+
+    case "grep":
+    case "glob":
+    case "context":
+    case "context7_resolve_library_id":
+    case "context7_get_library_docs":
+      return "search"
+
+    case "read":
+      return "read"
+
+    case "task":
+      return "think"
+
+    default:
+      return "other"
+  }
+}
+
+export function toLocations(toolName: string, input: ToolInput): ToolCallLocation[] {
+  const tool = toolName.toLocaleLowerCase()
+
+  switch (tool) {
+    case "read":
+    case "edit":
+    case "write":
+      return locationFrom(input.filePath ?? input.filepath)
+
+    case "external_directory":
+      return locationFrom(input.filePath ?? input.filepath, input.parentDir, input.directories)
+
+    case "grep":
+    case "glob":
+    case "context":
+    case "context7_resolve_library_id":
+    case "context7_get_library_docs":
+      return locationFrom(input.path)
+
+    case "bash":
+    case "shell":
+      return []
+
+    default:
+      return []
+  }
+}
+
+export function completedToolContent(toolName: string, state: CompletedToolState): ToolCallContent[] {
+  const text =
+    toolName.toLocaleLowerCase() === "read" ? (readDisplayText(state.metadata) ?? state.output) : state.output
+  const content: ToolCallContent[] = [
+    {
+      type: "content",
+      content: {
+        type: "text",
+        text,
+      },
+    },
+  ]
+
+  if (toToolKind(toolName) === "edit") {
+    content.push(...diffContent(state.input))
+  }
+
+  content.push(...imageContents(state.attachments ?? []))
+  return content
+}
+
+export function pendingToolCall(input: {
+  readonly toolCallId: string
+  readonly toolName: string
+  readonly state: { readonly input: ToolInput; readonly title?: string }
+}): ToolCall {
+  return {
+    toolCallId: input.toolCallId,
+    title: input.state.title || input.toolName,
+    kind: toToolKind(input.toolName),
+    status: "pending",
+    locations: toLocations(input.toolName, input.state.input),
+    rawInput: input.state.input,
+  }
+}
+
+export function runningToolUpdate(input: {
+  readonly toolCallId: string
+  readonly toolName: string
+  readonly state: RunningToolState
+  readonly output?: string
+}): ToolCallUpdate {
+  const content = input.output
+    ? [
+        {
+          type: "content" as const,
+          content: {
+            type: "text" as const,
+            text: input.output,
+          },
+        },
+      ]
+    : undefined
+
+  return {
+    toolCallId: input.toolCallId,
+    status: "in_progress",
+    kind: toToolKind(input.toolName),
+    title: input.state.title ?? input.toolName,
+    locations: toLocations(input.toolName, input.state.input),
+    rawInput: input.state.input,
+    ...(content ? { content } : {}),
+  }
+}
+
+export function duplicateRunningToolUpdate(input: {
+  readonly toolCallId: string
+  readonly toolName: string
+  readonly state: RunningToolState
+}): ToolCallUpdate {
+  return {
+    toolCallId: input.toolCallId,
+    status: "in_progress",
+    kind: toToolKind(input.toolName),
+    title: input.state.title ?? input.toolName,
+    locations: toLocations(input.toolName, input.state.input),
+    rawInput: input.state.input,
+  }
+}
+
+export function completedToolUpdate(input: {
+  readonly toolCallId: string
+  readonly toolName: string
+  readonly state: CompletedToolState & { readonly title: string }
+}): ToolCallUpdate {
+  return {
+    toolCallId: input.toolCallId,
+    status: "completed",
+    kind: toToolKind(input.toolName),
+    title: input.state.title,
+    content: completedToolContent(input.toolName, input.state),
+    rawInput: input.state.input,
+    rawOutput: completedToolRawOutput(input.state),
+  }
+}
+
+export function errorToolUpdate(input: {
+  readonly toolCallId: string
+  readonly toolName: string
+  readonly state: ErrorToolState
+}): ToolCallUpdate {
+  return {
+    toolCallId: input.toolCallId,
+    status: "failed",
+    kind: toToolKind(input.toolName),
+    title: input.toolName,
+    rawInput: input.state.input,
+    content: [
+      {
+        type: "content",
+        content: {
+          type: "text",
+          text: input.state.error,
+        },
+      },
+    ],
+    rawOutput: {
+      error: input.state.error,
+      metadata: input.state.metadata,
+    },
+  }
+}
+
+export function completedToolRawOutput(state: CompletedToolState) {
+  return {
+    output: state.output,
+    ...(state.metadata !== undefined ? { metadata: state.metadata } : {}),
+    ...(state.attachments?.length ? { attachments: state.attachments } : {}),
+  }
+}
+
+export function imageContents(attachments: ReadonlyArray<ToolAttachment>): ToolCallContent[] {
+  return extractImageAttachments(attachments).map((attachment): ToolCallContent => {
+    return {
+      type: "content",
+      content: {
+        type: "image",
+        mimeType: attachment.mimeType,
+        data: attachment.data,
+      },
+    }
+  })
+}
+
+export function extractImageAttachments(attachments: ReadonlyArray<ToolAttachment>): ImageAttachment[] {
+  return attachments.flatMap((attachment): ImageAttachment[] => {
+    const data = dataUrlImage(attachment)
+    return data ? [data] : []
+  })
+}
+
+export function shellOutputSnapshot(state: { readonly metadata?: unknown }) {
+  if (!state.metadata || typeof state.metadata !== "object") return undefined
+  return stringValue((state.metadata as Record<string, unknown>).output)
+}
+
+export const mapToolKind = toToolKind
+export const extractLocations = toLocations
+export const buildCompletedToolContent = completedToolContent
+export const buildCompletedRawOutput = completedToolRawOutput
+export const extractShellOutputSnapshot = shellOutputSnapshot
+export const buildPendingToolCall = pendingToolCall
+export const buildRunningToolUpdate = runningToolUpdate
+export const buildDuplicateRunningToolUpdate = duplicateRunningToolUpdate
+export const buildCompletedToolUpdate = completedToolUpdate
+export const buildErrorToolUpdate = errorToolUpdate
+
+function locationFrom(...values: unknown[]): ToolCallLocation[] {
+  return Array.from(
+    new Set(
+      values.flatMap((value): string[] => {
+        if (Array.isArray(value)) {
+          return value.filter((item): item is string => typeof item === "string" && item.length > 0)
+        }
+        const path = stringValue(value)
+        return path ? [path] : []
+      }),
+    ),
+    (path) => ({ path }),
+  )
+}
+
+function diffContent(input: ToolInput): ToolCallContent[] {
+  const oldText = stringValue(input.oldString)
+  const newText = stringValue(input.newString) ?? stringValue(input.content)
+  if (oldText === undefined || newText === undefined) return []
+
+  return [
+    {
+      type: "diff",
+      path: stringValue(input.filePath) ?? "",
+      oldText,
+      newText,
+    },
+  ]
+}
+
+function readDisplayText(metadata: unknown) {
+  if (!metadata || typeof metadata !== "object") return undefined
+  const display = (metadata as Record<string, unknown>).display
+  if (!display || typeof display !== "object") return undefined
+  const info = display as Record<string, unknown>
+  if (info.type === "file") return stringValue(info.text)
+  if (info.type === "directory" && Array.isArray(info.entries)) {
+    return info.entries.filter((item): item is string => typeof item === "string").join("\n")
+  }
+  return undefined
+}
+
+function dataUrlImage(attachment: ToolAttachment) {
+  const match = stringValue(attachment.url)?.match(/^data:([^;,]+)(?:;[^,]*)*;base64,(.*)$/)
+  const mime = match?.[1] ?? stringValue(attachment.mime)
+  if (!mime?.startsWith("image/")) return undefined
+
+  const data = match?.[2]
+  if (data === undefined) return undefined
+  return { mimeType: mime, data }
+}
+
+function stringValue(value: unknown) {
+  return typeof value === "string" ? value : undefined
+}

package/src/acp/usage.ts

@@ -0,0 +1,239 @@
+import type { AgentSideConnection, Usage } from "@agentclientprotocol/sdk"
+import * as Log from "@octocode-ai/core/util/log"
+import type { AssistantMessage as OctoAssistantMessage, Message } from "@octocode-ai/sdk/v2"
+import { InstanceRef } from "@/effect/instance-ref"
+import { InstanceStore } from "@/project/instance-store"
+import { ProviderV2 } from "@octocode-ai/core/provider"
+import { ModelV2 } from "@octocode-ai/core/model"
+import { Provider } from "@/provider/provider"
+import { Context, Effect, Layer, SynchronizedRef } from "effect"
+
+const log = Log.create({ service: "acp-usage" })
+
+export type AssistantTokenCost = Pick<OctoAssistantMessage, "cost" | "tokens">
+
+export type AssistantMessage = AssistantTokenCost &
+  Pick<OctoAssistantMessage, "role"> &
+  Partial<Pick<OctoAssistantMessage, "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>()(
+  "@octocode/ACPUsageMessageLoader",
+) {}
+
+export class ContextLimitLoader extends Context.Service<ContextLimitLoader, ContextLimitLoaderInterface>()(
+  "@octocode/ACPUsageContextLimitLoader",
+) {}
+
+export class Service extends Context.Service<Service, Interface>()("@octocode/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.sync(() => {
+                  log.error("failed to get providers for usage context limit", { error })
+                  return 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.sync(() => {
+            log.error("failed to fetch messages for usage update", { error })
+            return 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((error) => {
+            log.error("failed to send usage update", { error })
+          }),
+      )
+    })
+
+    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"