1  import { mkdtempSync, rmSync } from 'node:fs'
  2  import { tmpdir } from 'node:os'
  3  import path from 'node:path'
  4  
  5  import { afterEach, beforeEach, describe, expect, it } from 'vitest'
  6  
  7  import { ensureSession } from '@/server/storage/chat-store'
  8  import {
  9    appendAgentTaskEvent,
 10    cancelAgentTask,
 11    getPendingWorkflowInput,
 12    listAgentTaskEvents,
 13    listAgentTasks,
 14    recordWorkflowInput,
 15    resolveWorkflowInput,
 16    upsertAgentTask,
 17    upsertWorkflowRun,
 18  } from '@/server/storage/agent-runtime'
 19  
 20  const originalEnv = { ...process.env }
 21  
 22  let tempDataDir = ''
 23  
 24  describe('agent runtime storage', () => {
 25    beforeEach(() => {
 26      tempDataDir = mkdtempSync(path.join(tmpdir(), 'helper-agent-runtime-'))
 27      resetDbConnection()
 28      process.env = {
 29        ...originalEnv,
 30        APP_DATA_DIR: tempDataDir,
 31      }
 32    })
 33  
 34    afterEach(() => {
 35      resetDbConnection()
 36      rmSync(tempDataDir, { recursive: true, force: true })
 37      process.env = originalEnv
 38    })
 39  
 40    it('persists task lifecycle and timeline events', () => {
 41      ensureSession('session-agent-runtime')
 42  
 43      const task = upsertAgentTask({
 44        id: 'task-1',
 45        sessionId: 'session-agent-runtime',
 46        runId: 'run-1',
 47        title: 'Agent run',
 48        kind: 'agent_run',
 49        status: 'running',
 50        source: 'chat-stream',
 51      })
 52      expect(task.status).toBe('running')
 53  
 54      appendAgentTaskEvent({
 55        taskId: task.id,
 56        type: 'agent.started',
 57        text: 'Started run.',
 58      })
 59      appendAgentTaskEvent({
 60        taskId: task.id,
 61        type: 'agent.progress',
 62        text: 'Working...',
 63      })
 64  
 65      const events = listAgentTaskEvents({
 66        taskId: task.id,
 67        limit: 10,
 68      })
 69      expect(events).toHaveLength(2)
 70      expect(events[0]?.type).toBe('agent.started')
 71      expect(events[1]?.type).toBe('agent.progress')
 72  
 73      const cancelled = cancelAgentTask(task.id)
 74      expect(cancelled?.status).toBe('cancelled')
 75      expect(
 76        listAgentTasks({
 77          scope: 'active',
 78        }).some((entry) => entry.id === task.id),
 79      ).toBe(false)
 80    })
 81  
 82    it('records workflow input and resolves it', () => {
 83      ensureSession('session-workflow-runtime')
 84  
 85      upsertAgentTask({
 86        id: 'task-workflow',
 87        sessionId: 'session-workflow-runtime',
 88        title: 'Workflow task',
 89        kind: 'workflow',
 90        status: 'waiting_input',
 91        source: 'workflow',
 92      })
 93      upsertWorkflowRun({
 94        id: 'wf-1',
 95        taskId: 'task-workflow',
 96        sessionId: 'session-workflow-runtime',
 97        commandId: 'plan',
 98        status: 'waiting_input',
 99        awaitingInput: true,
100        resumeToken: 'resume-1',
101      })
102  
103      const request = recordWorkflowInput({
104        workflowId: 'wf-1',
105        taskId: 'task-workflow',
106        commandId: 'plan',
107        prompt: 'Pick a repo',
108        fields: [
109          {
110            id: 'repo',
111            label: 'Repo',
112            type: 'text',
113          },
114        ],
115        resumeToken: 'resume-1',
116      })
117      const pending = getPendingWorkflowInput({
118        workflowId: request.workflowId,
119        resumeToken: request.resumeToken,
120      })
121      expect(pending?.prompt).toBe('Pick a repo')
122  
123      const resolved = resolveWorkflowInput({
124        workflowId: request.workflowId,
125        resumeToken: request.resumeToken,
126        values: {
127          repo: 'owner/repo',
128        },
129      })
130      expect(resolved?.status).toBe('running')
131      expect(resolved?.awaitingInput).toBe(false)
132    })
133  })
134  
135  function resetDbConnection(): void {
136    const globalWithDb = globalThis as typeof globalThis & {
137      __helperDb__?: { close: () => void }
138    }
139    globalWithDb.__helperDb__?.close()
140    globalWithDb.__helperDb__ = undefined
141  }