@pikku/mongodb 0.12.20 → 0.12.22

package/src/mongodb-agent-storage-service.ts

@@ -1,632 +0,0 @@
-import type {
-  AgentMessage,
-  AgentRunState,
-  AgentThread,
-} from '@pikku/core/agent'
-import type { AgentRunScore } from '@pikku/core/agent-scorer'
-import type {
-  AgentRunStateService,
-  AgentStorageService,
-  CreateRunInput,
-  SaveScoreInput,
-} from '@pikku/core/services'
-import type { Db, Collection } from 'mongodb'
-
-interface AgentThreadDoc {
-  _id: string
-  resourceId: string
-  title: string | null
-  metadata: any | null
-  createdAt: Date
-  updatedAt: Date
-}
-
-interface AgentMessageDoc {
-  _id: string
-  threadId: string
-  role: string
-  content: any | null
-  createdAt: Date
-}
-
-interface AgentToolCallDoc {
-  _id: string
-  threadId: string
-  messageId: string
-  runId: string | null
-  toolName: string
-  args: any
-  result: string | null
-  approvalStatus: string | null
-  approvalType: string | null
-  agentRunId: string | null
-  displayToolName: string | null
-  displayArgs: any | null
-  createdAt: Date
-}
-
-interface AgentWorkingMemoryDoc {
-  id: string
-  scope: string
-  data: any
-  updatedAt: Date
-}
-
-interface AgentRunScoreDoc {
-  _id: string
-  runId: string
-  scorerName: string
-  score: number
-  reason: string | null
-  metadata: any | null
-  createdAt: Date
-}
-
-interface AgentRunDoc {
-  _id: string
-  agentName: string
-  threadId: string
-  resourceId: string
-  status: string
-  errorMessage: string | null
-  suspendReason: string | null
-  missingRpcs: any | null
-  usageInputTokens: number
-  usageOutputTokens: number
-  usageModel: string
-  createdAt: Date
-  updatedAt: Date
-}
-
-export class MongoDBAgentStorageService
-  implements AgentStorageService, AgentRunStateService
-{
-  private initialized = false
-  private threads!: Collection<AgentThreadDoc>
-  private agentMessages!: Collection<AgentMessageDoc>
-  private toolCalls!: Collection<AgentToolCallDoc>
-  private workingMemory!: Collection<AgentWorkingMemoryDoc>
-  private agentRuns!: Collection<AgentRunDoc>
-  private agentRunScores!: Collection<AgentRunScoreDoc>
-
-  constructor(private db: Db) {}
-
-  public async init(): Promise<void> {
-    if (this.initialized) return
-
-    this.threads = this.db.collection<AgentThreadDoc>('agent_threads')
-    this.agentMessages = this.db.collection<AgentMessageDoc>('agent_message')
-    this.toolCalls = this.db.collection<AgentToolCallDoc>('agent_tool_call')
-    this.workingMemory = this.db.collection<AgentWorkingMemoryDoc>(
-      'agent_working_memory'
-    )
-    this.agentRuns = this.db.collection<AgentRunDoc>('agent_run')
-    this.agentRunScores =
-      this.db.collection<AgentRunScoreDoc>('agent_run_score')
-
-    await this.threads.createIndex({ resourceId: 1 })
-
-    await this.agentMessages.createIndex({ threadId: 1, createdAt: 1 })
-
-    await this.toolCalls.createIndex({ threadId: 1 })
-    await this.toolCalls.createIndex({ messageId: 1 })
-
-    await this.workingMemory.createIndex({ id: 1, scope: 1 }, { unique: true })
-
-    await this.agentRuns.createIndex({ threadId: 1, createdAt: 1 })
-
-    await this.agentRunScores.createIndex({ runId: 1, createdAt: 1 })
-
-    this.initialized = true
-  }
-
-  async createThread(
-    resourceId: string,
-    options?: {
-      threadId?: string
-      title?: string
-      metadata?: Record<string, unknown>
-    }
-  ): Promise<AgentThread> {
-    const id = options?.threadId ?? crypto.randomUUID()
-    const now = new Date()
-
-    await this.threads.insertOne({
-      _id: id,
-      resourceId,
-      title: options?.title ?? null,
-      metadata: options?.metadata ?? null,
-      createdAt: now,
-      updatedAt: now,
-    })
-
-    return {
-      id,
-      resourceId,
-      title: options?.title,
-      metadata: options?.metadata,
-      createdAt: now,
-      updatedAt: now,
-    }
-  }
-
-  async getThread(threadId: string): Promise<AgentThread> {
-    const row = await this.threads.findOne({ _id: threadId })
-
-    if (!row) {
-      throw new Error(`Thread not found: ${threadId}`)
-    }
-
-    return {
-      id: row._id,
-      resourceId: row.resourceId,
-      title: row.title ?? undefined,
-      metadata: row.metadata ?? undefined,
-      createdAt: new Date(row.createdAt),
-      updatedAt: new Date(row.updatedAt),
-    }
-  }
-
-  async getThreads(resourceId: string): Promise<AgentThread[]> {
-    const result = await this.threads
-      .find({ resourceId })
-      .sort({ updatedAt: -1 })
-      .toArray()
-
-    return result.map((row) => ({
-      id: row._id,
-      resourceId: row.resourceId,
-      title: row.title ?? undefined,
-      metadata: row.metadata ?? undefined,
-      createdAt: new Date(row.createdAt),
-      updatedAt: new Date(row.updatedAt),
-    }))
-  }
-
-  async deleteThread(threadId: string): Promise<void> {
-    await this.toolCalls.deleteMany({ threadId })
-    await this.agentMessages.deleteMany({ threadId })
-    await this.agentRuns.deleteMany({ threadId })
-    await this.threads.deleteOne({ _id: threadId })
-  }
-
-  async getMessages(
-    threadId: string,
-    options?: { lastN?: number; cursor?: string }
-  ): Promise<AgentMessage[]> {
-    let filter: Record<string, any> = { threadId }
-
-    if (options?.cursor) {
-      const cursorRow = await this.agentMessages.findOne({
-        _id: options.cursor,
-      })
-      if (cursorRow) {
-        filter.createdAt = { $lt: cursorRow.createdAt }
-      }
-    }
-
-    let msgResult
-    if (options?.cursor || options?.lastN) {
-      const innerResult = await this.agentMessages
-        .find(filter)
-        .sort({ createdAt: -1 })
-        .limit(options?.lastN ?? 50)
-        .toArray()
-      innerResult.reverse()
-      msgResult = innerResult
-    } else {
-      msgResult = await this.agentMessages
-        .find(filter)
-        .sort({ createdAt: 1 })
-        .toArray()
-    }
-
-    const tcResult = await this.toolCalls
-      .find({ threadId })
-      .sort({ createdAt: 1 })
-      .toArray()
-
-    const tcByMessage = new Map<string, (typeof tcResult)[number][]>()
-    for (const tc of tcResult) {
-      const msgId = tc.messageId
-      if (!tcByMessage.has(msgId)) tcByMessage.set(msgId, [])
-      tcByMessage.get(msgId)!.push(tc)
-    }
-
-    const messages: AgentMessage[] = []
-    for (const row of msgResult) {
-      let parsedContent: AgentMessage['content'] = row.content ?? undefined
-      if (typeof row.content === 'string') {
-        try {
-          const parsed = JSON.parse(row.content)
-          if (Array.isArray(parsed)) {
-            parsedContent = parsed
-          } else if (typeof parsed === 'string') {
-            parsedContent = parsed
-          }
-        } catch {
-          // Not JSON, use raw string
-        }
-      }
-
-      const msg: AgentMessage = {
-        id: row._id,
-        role: row.role as AgentMessage['role'],
-        content: parsedContent,
-        createdAt: new Date(row.createdAt),
-      }
-
-      const tcs = tcByMessage.get(msg.id)
-      if (tcs?.length) {
-        msg.toolCalls = tcs.map((tc) => ({
-          id: tc._id,
-          name: tc.toolName,
-          args: tc.args as Record<string, unknown>,
-        }))
-
-        const completed = tcs.filter((tc) => tc.result != null)
-        if (completed.length) {
-          messages.push(msg)
-          messages.push({
-            id: `tool-results-${msg.id}`,
-            role: 'tool',
-            toolResults: completed.map((tc) => ({
-              id: tc._id,
-              name: tc.toolName,
-              result: tc.result!,
-            })),
-            createdAt: msg.createdAt,
-          })
-          continue
-        }
-      }
-
-      messages.push(msg)
-    }
-
-    return messages
-  }
-
-  async saveMessages(
-    threadId: string,
-    messages: AgentMessage[]
-  ): Promise<void> {
-    if (messages.length === 0) return
-
-    const toolMessages = messages.filter((m) => m.role === 'tool')
-    const nonToolMessages = messages.filter((m) => m.role !== 'tool')
-
-    if (nonToolMessages.length > 0) {
-      await this.agentMessages.insertMany(
-        nonToolMessages.map((msg) => ({
-          _id: msg.id,
-          threadId,
-          role: msg.role,
-          content: msg.content != null ? JSON.stringify(msg.content) : null,
-          createdAt: msg.createdAt ?? new Date(),
-        }))
-      )
-    }
-
-    const toolCallDocs = nonToolMessages.flatMap(
-      (msg) => msg.toolCalls?.map((tc) => ({ ...tc, messageId: msg.id })) ?? []
-    )
-    if (toolCallDocs.length > 0) {
-      await this.toolCalls.insertMany(
-        toolCallDocs.map((tc) => ({
-          _id: tc.id,
-          threadId,
-          messageId: tc.messageId,
-          runId: null,
-          toolName: tc.name,
-          args: tc.args,
-          result: null,
-          approvalStatus: null,
-          approvalType: null,
-          agentRunId: null,
-          displayToolName: null,
-          displayArgs: null,
-          createdAt: new Date(),
-        }))
-      )
-    }
-
-    for (const toolMsg of toolMessages) {
-      if (!toolMsg.toolResults) continue
-      for (const tr of toolMsg.toolResults) {
-        await this.toolCalls.updateOne(
-          { _id: tr.id },
-          { $set: { result: tr.result } }
-        )
-      }
-    }
-
-    await this.threads.updateOne(
-      { _id: threadId },
-      { $set: { updatedAt: new Date() } }
-    )
-  }
-
-  async getWorkingMemory(
-    id: string,
-    scope: 'resource' | 'thread'
-  ): Promise<Record<string, unknown> | null> {
-    const row = await this.workingMemory.findOne({ id, scope })
-    if (!row) return null
-    return row.data
-  }
-
-  async saveWorkingMemory(
-    id: string,
-    scope: 'resource' | 'thread',
-    data: Record<string, unknown>
-  ): Promise<void> {
-    await this.workingMemory.updateOne(
-      { id, scope },
-      {
-        $set: {
-          data,
-          updatedAt: new Date(),
-        },
-        $setOnInsert: {
-          id,
-          scope,
-        },
-      },
-      { upsert: true }
-    )
-  }
-
-  async createRun(run: CreateRunInput): Promise<string> {
-    const runId = crypto.randomUUID()
-
-    await this.agentRuns.insertOne({
-      _id: runId,
-      agentName: run.agentName,
-      threadId: run.threadId,
-      resourceId: run.resourceId,
-      status: run.status,
-      errorMessage: run.errorMessage ?? null,
-      suspendReason: run.suspendReason ?? null,
-      missingRpcs: run.missingRpcs ?? null,
-      usageInputTokens: run.usage.inputTokens,
-      usageOutputTokens: run.usage.outputTokens,
-      usageModel: run.usage.model,
-      createdAt: run.createdAt,
-      updatedAt: run.updatedAt,
-    })
-
-    if (run.pendingApprovals?.length) {
-      await this.setApprovals(runId, run.pendingApprovals)
-    }
-
-    return runId
-  }
-
-  async updateRun(
-    runId: string,
-    updates: Partial<AgentRunState>
-  ): Promise<void> {
-    const setValues: Record<string, any> = { updatedAt: new Date() }
-
-    if (updates.status !== undefined) {
-      setValues.status = updates.status
-    }
-    if (updates.errorMessage !== undefined) {
-      setValues.errorMessage = updates.errorMessage
-    }
-    if (updates.suspendReason !== undefined) {
-      setValues.suspendReason = updates.suspendReason
-    }
-    if (updates.missingRpcs !== undefined) {
-      setValues.missingRpcs = updates.missingRpcs
-    }
-    if (updates.usage !== undefined) {
-      setValues.usageInputTokens = updates.usage.inputTokens
-      setValues.usageOutputTokens = updates.usage.outputTokens
-      setValues.usageModel = updates.usage.model
-    }
-
-    await this.agentRuns.updateOne({ _id: runId }, { $set: setValues })
-
-    if (updates.pendingApprovals !== undefined) {
-      await this.toolCalls.updateMany(
-        { runId, approvalStatus: { $ne: null } },
-        {
-          $set: {
-            approvalStatus: null,
-            runId: null,
-            approvalType: null,
-            agentRunId: null,
-            displayToolName: null,
-            displayArgs: null,
-          },
-        }
-      )
-
-      if (updates.pendingApprovals.length) {
-        await this.setApprovals(runId, updates.pendingApprovals)
-      }
-    }
-  }
-
-  async getRun(runId: string): Promise<AgentRunState | null> {
-    const row = await this.agentRuns.findOne({ _id: runId })
-    if (!row) return null
-
-    const approvals = await this.toolCalls
-      .find({ runId, approvalStatus: 'pending' })
-      .toArray()
-
-    return this.mapRunRow(row, approvals)
-  }
-
-  async getRunsByThread(threadId: string): Promise<AgentRunState[]> {
-    const result = await this.agentRuns
-      .find({ threadId })
-      .sort({ createdAt: -1 })
-      .toArray()
-
-    const runs: AgentRunState[] = []
-    for (const row of result) {
-      const approvals = await this.toolCalls
-        .find({ runId: row._id, approvalStatus: 'pending' })
-        .toArray()
-      runs.push(this.mapRunRow(row, approvals))
-    }
-    return runs
-  }
-
-  private async setApprovals(
-    runId: string,
-    approvals: NonNullable<AgentRunState['pendingApprovals']>
-  ): Promise<void> {
-    for (const a of approvals) {
-      if (a.type === 'agent-call') {
-        await this.toolCalls.updateOne(
-          { _id: a.toolCallId },
-          {
-            $set: {
-              approvalStatus: 'pending',
-              runId,
-              approvalType: 'agent-call',
-              agentRunId: a.agentRunId,
-              displayToolName: a.displayToolName,
-              displayArgs: a.displayArgs,
-            },
-          }
-        )
-      } else {
-        await this.toolCalls.updateOne(
-          { _id: a.toolCallId },
-          {
-            $set: {
-              approvalStatus: 'pending',
-              runId,
-              approvalType: 'tool-call',
-            },
-          }
-        )
-      }
-    }
-  }
-
-  private mapRunRow(row: any, approvalRows?: any[]): AgentRunState {
-    const pendingApprovals = approvalRows?.length
-      ? approvalRows.map((a: any) => {
-          if (a.approvalType === 'agent-call') {
-            return {
-              type: 'agent-call' as const,
-              toolCallId: a._id as string,
-              agentName: a.toolName as string,
-              agentRunId: a.agentRunId as string,
-              displayToolName: a.displayToolName as string,
-              displayArgs: a.displayArgs as unknown,
-            }
-          }
-          return {
-            type: 'tool-call' as const,
-            toolCallId: a._id as string,
-            toolName: a.toolName as string,
-            args: a.args as unknown,
-          }
-        })
-      : undefined
-
-    return {
-      runId: row._id as string,
-      agentName: row.agentName as string,
-      threadId: row.threadId as string,
-      resourceId: row.resourceId as string,
-      status: row.status as AgentRunState['status'],
-      errorMessage: row.errorMessage ?? undefined,
-      suspendReason: row.suspendReason as AgentRunState['suspendReason'],
-      missingRpcs: row.missingRpcs ?? undefined,
-      pendingApprovals,
-      usage: {
-        inputTokens: Number(row.usageInputTokens),
-        outputTokens: Number(row.usageOutputTokens),
-        model: row.usageModel as string,
-      },
-      createdAt: new Date(row.createdAt as string),
-      updatedAt: new Date(row.updatedAt as string),
-    }
-  }
-
-  async findRunByToolCallId(toolCallId: string): Promise<{
-    run: AgentRunState
-    approval: NonNullable<AgentRunState['pendingApprovals']>[number]
-  } | null> {
-    const tc = await this.toolCalls.findOne({
-      _id: toolCallId,
-      approvalStatus: 'pending',
-    })
-
-    if (!tc || !tc.runId) return null
-
-    const run = await this.getRun(tc.runId)
-    if (!run) return null
-
-    let approval: NonNullable<AgentRunState['pendingApprovals']>[number]
-    if (tc.approvalType === 'agent-call') {
-      approval = {
-        type: 'agent-call',
-        toolCallId: tc._id,
-        agentName: tc.toolName,
-        agentRunId: tc.agentRunId!,
-        displayToolName: tc.displayToolName!,
-        displayArgs: tc.displayArgs as unknown,
-      }
-    } else {
-      approval = {
-        type: 'tool-call',
-        toolCallId: tc._id,
-        toolName: tc.toolName,
-        args: tc.args as unknown,
-      }
-    }
-
-    return { run, approval }
-  }
-
-  async resolveApproval(
-    toolCallId: string,
-    status: 'approved' | 'denied'
-  ): Promise<boolean> {
-    // Only the caller that moves the tool call off 'pending' has claimed it.
-    const result = await this.toolCalls.updateOne(
-      { _id: toolCallId, approvalStatus: 'pending' },
-      { $set: { approvalStatus: status } }
-    )
-    return result.modifiedCount > 0
-  }
-
-  async saveScore(score: SaveScoreInput): Promise<void> {
-    await this.agentRunScores.insertOne({
-      _id: crypto.randomUUID(),
-      runId: score.runId,
-      scorerName: score.scorerName,
-      score: score.score,
-      reason: score.reason ?? null,
-      metadata: score.metadata ?? null,
-      createdAt: new Date(),
-    })
-  }
-
-  async getScores(runId: string): Promise<AgentRunScore[]> {
-    const rows = await this.agentRunScores
-      .find({ runId })
-      .sort({ createdAt: 1 })
-      .toArray()
-    return rows.map((row) => ({
-      runId: row.runId,
-      scorerName: row.scorerName,
-      score: row.score,
-      reason: row.reason ?? undefined,
-      metadata: row.metadata ?? undefined,
-      createdAt: row.createdAt,
-    }))
-  }
-
-  public async close(): Promise<void> {}
-}