nikcli 0.0.6

package/src/session/message.ts

@@ -0,0 +1,189 @@
+import z from "zod"
+import { NamedError } from "@nikcli-ai/util/error"
+
+export namespace Message {
+  export const OutputLengthError = NamedError.create("MessageOutputLengthError", z.object({}))
+  export const AuthError = NamedError.create(
+    "ProviderAuthError",
+    z.object({
+      providerID: z.string(),
+      message: z.string(),
+    }),
+  )
+
+  export const ToolCall = z
+    .object({
+      state: z.literal("call"),
+      step: z.number().optional(),
+      toolCallId: z.string(),
+      toolName: z.string(),
+      args: z.custom<Required<unknown>>(),
+    })
+    .meta({
+      ref: "ToolCall",
+    })
+  export type ToolCall = z.infer<typeof ToolCall>
+
+  export const ToolPartialCall = z
+    .object({
+      state: z.literal("partial-call"),
+      step: z.number().optional(),
+      toolCallId: z.string(),
+      toolName: z.string(),
+      args: z.custom<Required<unknown>>(),
+    })
+    .meta({
+      ref: "ToolPartialCall",
+    })
+  export type ToolPartialCall = z.infer<typeof ToolPartialCall>
+
+  export const ToolResult = z
+    .object({
+      state: z.literal("result"),
+      step: z.number().optional(),
+      toolCallId: z.string(),
+      toolName: z.string(),
+      args: z.custom<Required<unknown>>(),
+      result: z.string(),
+    })
+    .meta({
+      ref: "ToolResult",
+    })
+  export type ToolResult = z.infer<typeof ToolResult>
+
+  export const ToolInvocation = z.discriminatedUnion("state", [ToolCall, ToolPartialCall, ToolResult]).meta({
+    ref: "ToolInvocation",
+  })
+  export type ToolInvocation = z.infer<typeof ToolInvocation>
+
+  export const TextPart = z
+    .object({
+      type: z.literal("text"),
+      text: z.string(),
+    })
+    .meta({
+      ref: "TextPart",
+    })
+  export type TextPart = z.infer<typeof TextPart>
+
+  export const ReasoningPart = z
+    .object({
+      type: z.literal("reasoning"),
+      text: z.string(),
+      providerMetadata: z.record(z.string(), z.any()).optional(),
+    })
+    .meta({
+      ref: "ReasoningPart",
+    })
+  export type ReasoningPart = z.infer<typeof ReasoningPart>
+
+  export const ToolInvocationPart = z
+    .object({
+      type: z.literal("tool-invocation"),
+      toolInvocation: ToolInvocation,
+    })
+    .meta({
+      ref: "ToolInvocationPart",
+    })
+  export type ToolInvocationPart = z.infer<typeof ToolInvocationPart>
+
+  export const SourceUrlPart = z
+    .object({
+      type: z.literal("source-url"),
+      sourceId: z.string(),
+      url: z.string(),
+      title: z.string().optional(),
+      providerMetadata: z.record(z.string(), z.any()).optional(),
+    })
+    .meta({
+      ref: "SourceUrlPart",
+    })
+  export type SourceUrlPart = z.infer<typeof SourceUrlPart>
+
+  export const FilePart = z
+    .object({
+      type: z.literal("file"),
+      mediaType: z.string(),
+      filename: z.string().optional(),
+      url: z.string(),
+    })
+    .meta({
+      ref: "FilePart",
+    })
+  export type FilePart = z.infer<typeof FilePart>
+
+  export const StepStartPart = z
+    .object({
+      type: z.literal("step-start"),
+    })
+    .meta({
+      ref: "StepStartPart",
+    })
+  export type StepStartPart = z.infer<typeof StepStartPart>
+
+  export const MessagePart = z
+    .discriminatedUnion("type", [TextPart, ReasoningPart, ToolInvocationPart, SourceUrlPart, FilePart, StepStartPart])
+    .meta({
+      ref: "MessagePart",
+    })
+  export type MessagePart = z.infer<typeof MessagePart>
+
+  export const Info = z
+    .object({
+      id: z.string(),
+      role: z.enum(["user", "assistant"]),
+      parts: z.array(MessagePart),
+      metadata: z
+        .object({
+          time: z.object({
+            created: z.number(),
+            completed: z.number().optional(),
+          }),
+          error: z
+            .discriminatedUnion("name", [AuthError.Schema, NamedError.Unknown.Schema, OutputLengthError.Schema])
+            .optional(),
+          sessionID: z.string(),
+          tool: z.record(
+            z.string(),
+            z
+              .object({
+                title: z.string(),
+                snapshot: z.string().optional(),
+                time: z.object({
+                  start: z.number(),
+                  end: z.number(),
+                }),
+              })
+              .catchall(z.any()),
+          ),
+          assistant: z
+            .object({
+              system: z.string().array(),
+              modelID: z.string(),
+              providerID: z.string(),
+              path: z.object({
+                cwd: z.string(),
+                root: z.string(),
+              }),
+              cost: z.number(),
+              summary: z.boolean().optional(),
+              tokens: z.object({
+                input: z.number(),
+                output: z.number(),
+                reasoning: z.number(),
+                cache: z.object({
+                  read: z.number(),
+                  write: z.number(),
+                }),
+              }),
+            })
+            .optional(),
+          snapshot: z.string().optional(),
+        })
+        .meta({ ref: "MessageMetadata" }),
+    })
+    .meta({
+      ref: "Message",
+    })
+  export type Info = z.infer<typeof Info>
+}

package/src/session/processor.ts

@@ -0,0 +1,408 @@
+import { MessageV2 } from "./message-v2"
+import { Log } from "@/util/log"
+import { Identifier } from "@/id/id"
+import { Session } from "."
+import { Agent } from "@/agent/agent"
+import { Snapshot } from "@/snapshot"
+import { SessionSummary } from "./summary"
+import { Bus } from "@/bus"
+import { SessionRetry } from "./retry"
+import { SessionStatus } from "./status"
+import { Plugin } from "@/plugin"
+import type { Provider } from "@/provider/provider"
+import { LLM } from "./llm"
+import { Config } from "@/config/config"
+import { SessionCompaction } from "./compaction"
+import { PermissionNext } from "@/permission/next"
+import { Question } from "@/question"
+
+export namespace SessionProcessor {
+  const DOOM_LOOP_THRESHOLD = 3
+  const log = Log.create({ service: "session.processor" })
+
+  export type Info = Awaited<ReturnType<typeof create>>
+  export type Result = Awaited<ReturnType<Info["process"]>>
+
+  export function create(input: {
+    assistantMessage: MessageV2.Assistant
+    sessionID: string
+    model: Provider.Model
+    abort: AbortSignal
+  }) {
+    const toolcalls: Record<string, MessageV2.ToolPart> = {}
+    let snapshot: string | undefined
+    let blocked = false
+    let attempt = 0
+    let needsCompaction = false
+
+    const result = {
+      get message() {
+        return input.assistantMessage
+      },
+      partFromToolCall(toolCallID: string) {
+        return toolcalls[toolCallID]
+      },
+      async process(streamInput: LLM.StreamInput) {
+        log.info("process")
+        needsCompaction = false
+        const shouldBreak = (await Config.get()).experimental?.continue_loop_on_deny !== true
+        while (true) {
+          try {
+            let currentText: MessageV2.TextPart | undefined
+            let reasoningMap: Record<string, MessageV2.ReasoningPart> = {}
+            const stream = await LLM.stream(streamInput)
+
+            for await (const value of stream.fullStream) {
+              input.abort.throwIfAborted()
+              switch (value.type) {
+                case "start":
+                  SessionStatus.set(input.sessionID, { type: "busy" })
+                  break
+
+                case "reasoning-start":
+                  if (value.id in reasoningMap) {
+                    continue
+                  }
+                  reasoningMap[value.id] = {
+                    id: Identifier.ascending("part"),
+                    messageID: input.assistantMessage.id,
+                    sessionID: input.assistantMessage.sessionID,
+                    type: "reasoning",
+                    text: "",
+                    time: {
+                      start: Date.now(),
+                    },
+                    metadata: value.providerMetadata,
+                  }
+                  break
+
+                case "reasoning-delta":
+                  if (value.id in reasoningMap) {
+                    const part = reasoningMap[value.id]
+                    part.text += value.text
+                    if (value.providerMetadata) part.metadata = value.providerMetadata
+                    if (part.text) await Session.updatePart({ part, delta: value.text })
+                  }
+                  break
+
+                case "reasoning-end":
+                  if (value.id in reasoningMap) {
+                    const part = reasoningMap[value.id]
+                    part.text = part.text.trimEnd()
+
+                    part.time = {
+                      ...part.time,
+                      end: Date.now(),
+                    }
+                    if (value.providerMetadata) part.metadata = value.providerMetadata
+                    await Session.updatePart(part)
+                    delete reasoningMap[value.id]
+                  }
+                  break
+
+                case "tool-input-start":
+                  const part = await Session.updatePart({
+                    id: toolcalls[value.id]?.id ?? Identifier.ascending("part"),
+                    messageID: input.assistantMessage.id,
+                    sessionID: input.assistantMessage.sessionID,
+                    type: "tool",
+                    tool: value.toolName,
+                    callID: value.id,
+                    state: {
+                      status: "pending",
+                      input: {},
+                      raw: "",
+                    },
+                  })
+                  toolcalls[value.id] = part as MessageV2.ToolPart
+                  break
+
+                case "tool-input-delta":
+                  break
+
+                case "tool-input-end":
+                  break
+
+                case "tool-call": {
+                  const match = toolcalls[value.toolCallId]
+                  if (match) {
+                    const part = await Session.updatePart({
+                      ...match,
+                      tool: value.toolName,
+                      state: {
+                        status: "running",
+                        input: value.input,
+                        time: {
+                          start: Date.now(),
+                        },
+                      },
+                      metadata: value.providerMetadata,
+                    })
+                    toolcalls[value.toolCallId] = part as MessageV2.ToolPart
+
+                    const parts = await MessageV2.parts(input.assistantMessage.id)
+                    const lastThree = parts.slice(-DOOM_LOOP_THRESHOLD)
+
+                    if (
+                      lastThree.length === DOOM_LOOP_THRESHOLD &&
+                      lastThree.every(
+                        (p) =>
+                          p.type === "tool" &&
+                          p.tool === value.toolName &&
+                          p.state.status !== "pending" &&
+                          // Avoid JSON.stringify in a hot path (and avoid crashes on non-JSON values).
+                          Bun.deepEquals(p.state.input, value.input),
+                      )
+                    ) {
+                      const agent = await Agent.get(input.assistantMessage.agent)
+                      await PermissionNext.ask({
+                        permission: "doom_loop",
+                        patterns: [value.toolName],
+                        sessionID: input.assistantMessage.sessionID,
+                        metadata: {
+                          tool: value.toolName,
+                          input: value.input,
+                        },
+                        always: [value.toolName],
+                        ruleset: agent.permission,
+                      })
+                    }
+                  }
+                  break
+                }
+                case "tool-result": {
+                  const match = toolcalls[value.toolCallId]
+                  if (match && match.state.status === "running") {
+                    await Session.updatePart({
+                      ...match,
+                      state: {
+                        status: "completed",
+                        input: value.input ?? match.state.input,
+                        output: value.output.output,
+                        metadata: value.output.metadata,
+                        title: value.output.title,
+                        time: {
+                          start: match.state.time.start,
+                          end: Date.now(),
+                        },
+                        attachments: value.output.attachments,
+                      },
+                    })
+
+                    delete toolcalls[value.toolCallId]
+                  }
+                  break
+                }
+
+                case "tool-error": {
+                  const match = toolcalls[value.toolCallId]
+                  if (match && match.state.status === "running") {
+                    await Session.updatePart({
+                      ...match,
+                      state: {
+                        status: "error",
+                        input: value.input ?? match.state.input,
+                        error: (value.error as any).toString(),
+                        time: {
+                          start: match.state.time.start,
+                          end: Date.now(),
+                        },
+                      },
+                    })
+
+                    if (
+                      value.error instanceof PermissionNext.RejectedError ||
+                      value.error instanceof Question.RejectedError
+                    ) {
+                      blocked = shouldBreak
+                    }
+                    delete toolcalls[value.toolCallId]
+                  }
+                  break
+                }
+                case "error":
+                  throw value.error
+
+                case "start-step":
+                  snapshot = await Snapshot.track()
+                  await Session.updatePart({
+                    id: Identifier.ascending("part"),
+                    messageID: input.assistantMessage.id,
+                    sessionID: input.sessionID,
+                    snapshot,
+                    type: "step-start",
+                  })
+                  break
+
+                case "finish-step":
+                  const usage = Session.getUsage({
+                    model: input.model,
+                    usage: value.usage,
+                    metadata: value.providerMetadata,
+                  })
+                  input.assistantMessage.finish = value.finishReason
+                  input.assistantMessage.cost += usage.cost
+                  input.assistantMessage.tokens = usage.tokens
+                  await Session.updatePart({
+                    id: Identifier.ascending("part"),
+                    reason: value.finishReason,
+                    snapshot: await Snapshot.track(),
+                    messageID: input.assistantMessage.id,
+                    sessionID: input.assistantMessage.sessionID,
+                    type: "step-finish",
+                    tokens: usage.tokens,
+                    cost: usage.cost,
+                  })
+                  await Session.updateMessage(input.assistantMessage)
+                  if (snapshot) {
+                    const patch = await Snapshot.patch(snapshot)
+                    if (patch.files.length) {
+                      await Session.updatePart({
+                        id: Identifier.ascending("part"),
+                        messageID: input.assistantMessage.id,
+                        sessionID: input.sessionID,
+                        type: "patch",
+                        hash: patch.hash,
+                        files: patch.files,
+                      })
+                    }
+                    snapshot = undefined
+                  }
+                  SessionSummary.summarize({
+                    sessionID: input.sessionID,
+                    messageID: input.assistantMessage.parentID,
+                  })
+                  if (await SessionCompaction.isOverflow({ tokens: usage.tokens, model: input.model })) {
+                    needsCompaction = true
+                  }
+                  break
+
+                case "text-start":
+                  currentText = {
+                    id: Identifier.ascending("part"),
+                    messageID: input.assistantMessage.id,
+                    sessionID: input.assistantMessage.sessionID,
+                    type: "text",
+                    text: "",
+                    time: {
+                      start: Date.now(),
+                    },
+                    metadata: value.providerMetadata,
+                  }
+                  break
+
+                case "text-delta":
+                  if (currentText) {
+                    currentText.text += value.text
+                    if (value.providerMetadata) currentText.metadata = value.providerMetadata
+                    if (currentText.text)
+                      await Session.updatePart({
+                        part: currentText,
+                        delta: value.text,
+                      })
+                  }
+                  break
+
+                case "text-end":
+                  if (currentText) {
+                    currentText.text = currentText.text.trimEnd()
+                    const textOutput = await Plugin.trigger(
+                      "experimental.text.complete",
+                      {
+                        sessionID: input.sessionID,
+                        messageID: input.assistantMessage.id,
+                        partID: currentText.id,
+                      },
+                      { text: currentText.text },
+                    )
+                    currentText.text = textOutput.text
+                    currentText.time = {
+                      start: Date.now(),
+                      end: Date.now(),
+                    }
+                    if (value.providerMetadata) currentText.metadata = value.providerMetadata
+                    await Session.updatePart(currentText)
+                  }
+                  currentText = undefined
+                  break
+
+                case "finish":
+                  break
+
+                default:
+                  log.info("unhandled", {
+                    ...value,
+                  })
+                  continue
+              }
+              if (needsCompaction) break
+            }
+          } catch (e: any) {
+            log.error("process", {
+              error: e,
+              stack: JSON.stringify(e.stack),
+            })
+            const error = MessageV2.fromError(e, { providerID: input.model.providerID })
+            // TODO: Handle context overflow error
+            const retry = SessionRetry.retryable(error)
+            if (retry !== undefined) {
+              attempt++
+              const delay = SessionRetry.delay(attempt, error.name === "APIError" ? error : undefined)
+              SessionStatus.set(input.sessionID, {
+                type: "retry",
+                attempt,
+                message: retry,
+                next: Date.now() + delay,
+              })
+              await SessionRetry.sleep(delay, input.abort).catch(() => {})
+              continue
+            }
+            input.assistantMessage.error = error
+            Bus.publish(Session.Event.Error, {
+              sessionID: input.assistantMessage.sessionID,
+              error: input.assistantMessage.error,
+            })
+          }
+          if (snapshot) {
+            const patch = await Snapshot.patch(snapshot)
+            if (patch.files.length) {
+              await Session.updatePart({
+                id: Identifier.ascending("part"),
+                messageID: input.assistantMessage.id,
+                sessionID: input.sessionID,
+                type: "patch",
+                hash: patch.hash,
+                files: patch.files,
+              })
+            }
+            snapshot = undefined
+          }
+          const p = await MessageV2.parts(input.assistantMessage.id)
+          for (const part of p) {
+            if (part.type === "tool" && part.state.status !== "completed" && part.state.status !== "error") {
+              await Session.updatePart({
+                ...part,
+                state: {
+                  ...part.state,
+                  status: "error",
+                  error: "Tool execution aborted",
+                  time: {
+                    start: Date.now(),
+                    end: Date.now(),
+                  },
+                },
+              })
+            }
+          }
+          input.assistantMessage.time.completed = Date.now()
+          await Session.updateMessage(input.assistantMessage)
+          if (needsCompaction) return "compact"
+          if (blocked) return "stop"
+          if (input.assistantMessage.error) return "stop"
+          return "continue"
+        }
+      },
+    }
+    return result
+  }
+}