deepseek-harness/packages/agent-loop/tests/agent.spec.ts

import { describe, expect, it, vi } from 'vitest'
import { Context } from 'cordis'
import { AgentId } from '@deepseek-ai/dsh-agent'
import LlmService from '@deepseek-ai/dsh-llm'
import SessionStore from '@deepseek-ai/dsh-session'
import SystemPrompt from '@deepseek-ai/dsh-system-prompt'
import ToolRegistry from '@deepseek-ai/dsh-tools'
import AgentRegistry from '@deepseek-ai/dsh-agent'
import AgentLoop, { LoopAgent } from '@deepseek-ai/dsh-agent-loop'
import { MockAdapter, textResponse } from './mock-adapter.ts'

async function harness(adapter: MockAdapter) {
  const ctx = new Context()
  await ctx.plugin(LlmService)
  await ctx.plugin(SessionStore)
  await ctx.plugin(SystemPrompt)
  await ctx.plugin(ToolRegistry)
  await ctx.plugin(AgentRegistry)
  await ctx.plugin(AgentLoop, { agents: [] })
  ctx.llm.registerAdapter(['mock'], adapter)
  return ctx
}

function waitForIdle(ctx: Context, agent: LoopAgent): Promise<void> {
  return new Promise((resolve) => {
    const dispose = ctx.on('agent/status', (subject, status) => {
      if (subject === agent && status === 'idle') {
        dispose()
        resolve()
      }
    })
  })
}

function send(agent: LoopAgent, text: string) {
  agent.send([{ type: 'text', text }])
}

describe('LoopAgent', () => {
  it('send() throws after disposal', async () => {
    const adapter = new MockAdapter(['hang'])
    const ctx = await harness(adapter)
    let agent!: LoopAgent
    const fiber = await ctx.plugin(Object.assign((inner: Context) => {
      agent = inner.agentLoop.create('scoped', { model: 'mock' })
    }, { inject: ['agentLoop'] }))
    send(agent, 'go')
    await new Promise(r => setTimeout(r, 30))
    await fiber.dispose()
    await agent.done

    expect(() => { agent.send([{ type: 'text', text: 'too late' }]) }).toThrow('disposed')
  })

  it('steer() throws after disposal', async () => {
    const adapter = new MockAdapter(['hang'])
    const ctx = await harness(adapter)
    let agent!: LoopAgent
    const fiber = await ctx.plugin(Object.assign((inner: Context) => {
      agent = inner.agentLoop.create('scoped', { model: 'mock' })
    }, { inject: ['agentLoop'] }))
    send(agent, 'go')
    await new Promise(r => setTimeout(r, 30))
    await fiber.dispose()
    await agent.done

    expect(() => { agent.steer([{ type: 'text', text: 'too late' }]) }).toThrow('disposed')
  })

  it('inject() throws after disposal', async () => {
    const adapter = new MockAdapter(['hang'])
    const ctx = await harness(adapter)
    let agent!: LoopAgent
    const fiber = await ctx.plugin(Object.assign((inner: Context) => {
      agent = inner.agentLoop.create('scoped', { model: 'mock' })
    }, { inject: ['agentLoop'] }))
    send(agent, 'go')
    await new Promise(r => setTimeout(r, 30))
    await fiber.dispose()
    await agent.done

    expect(() => { agent.inject([{ type: 'text', text: 'too late' }]) }).toThrow('disposed')
  })

  it('inject() decides enclosure from the LOG (open turn), not agent status', async () => {
    const adapter = new MockAdapter([textResponse('ok')])
    const ctx = await harness(adapter)
    const agent = ctx.agentLoop.create('a1', { model: 'mock' })

    // Simulate an OPEN turn in the log while the agent is idle (status is not a
    // reliable open-turn signal). inject must append into that open turn, NOT
    // wrap a new one.
    agent.session.append('turn/start', { turn: 1, trigger: { kind: 'message', source: { kind: 'user' } } })
    agent.inject([{ type: 'text', text: 'mid' }], { source: { kind: 'plugin', plugin: 'p' } })
    expect(agent.session.events.filter(e => e.type === 'turn/start')).toHaveLength(1)
    expect(agent.session.events.at(-1)!.type).toBe('context/message')

    // Close the turn; now inject must wrap its own one-shot injection turn.
    agent.session.append('turn/end', { turn: 1, reason: { kind: 'completed' } })
    agent.inject([{ type: 'text', text: 'after' }], { source: { kind: 'plugin', plugin: 'p' } })
    const starts = agent.session.events.filter(e => e.type === 'turn/start')
    expect(starts).toHaveLength(2)
    const last = starts[1]!
    expect(last.type === 'turn/start' && last.data.trigger.kind).toBe('injection')
    expect(agent.session.events.at(-1)!.type).toBe('turn/end') // turn-enclosed
  })

  it('idle inject() contains a failing flush (logs, does not throw into the caller)', async () => {
    const adapter = new MockAdapter([textResponse('ok')])
    const ctx = await harness(adapter)
    // A persistence-like listener whose flush rejects.
    ctx.on('session/flush', () => { throw new Error('disk gone') })
    const warn = vi.spyOn(ctx.logger, 'warn').mockImplementation(() => undefined)
    const agent = ctx.agentLoop.create('a1', { model: 'mock' })

    // inject() is synchronous and fires a fire-and-forget flush; a rejecting
    // flush must be contained (logged), never thrown into the caller.
    expect(() => { agent.inject([{ type: 'text', text: 'notice' }], { source: { kind: 'plugin', plugin: 'p' } }) }).not.toThrow()
    await new Promise(r => setTimeout(r, 20)) // let the contained flush settle
    expect(warn).toHaveBeenCalledWith(expect.stringContaining('flush after idle injection failed'))
    warn.mockRestore()
  })

  it('idle inject() closes its one-shot turn AND still checkpoints even if the append throws', async () => {
    const adapter = new MockAdapter([textResponse('ok')])
    const ctx = await harness(adapter)
    const agent = ctx.agentLoop.create('a1', { model: 'mock' })
    let flushes = 0
    ctx.on('session/flush', () => { flushes += 1 })

    // Non-serializable injected content makes Session.append throw AFTER
    // turn/start was recorded. The turn/end must still be appended (finally),
    // AND the durability checkpoint must still fire — the balanced turn is in
    // memory and a crash before the next turn/dispose would otherwise lose it.
    expect(() => {
      agent.inject([{ type: 'text', text: 'x', bad: 1n } as never], { source: { kind: 'plugin', plugin: 'p' } })
    }).toThrow(/non-JSON-serializable/)
    const types = agent.session.events.map(e => e.type)
    expect(types).toEqual(['turn/start', 'turn/end']) // balanced, no open turn
    await new Promise(r => setTimeout(r, 10)) // let the fire-and-forget flush run
    expect(flushes).toBe(1) // checkpoint fired despite the throw
  })

  it('idle inject() still checkpoints when a listener throws on the synthetic turn/end', async () => {
    const adapter = new MockAdapter([textResponse('ok')])
    const ctx = await harness(adapter)
    const agent = ctx.agentLoop.create('a1', { model: 'mock' })
    let flushes = 0
    ctx.on('session/flush', () => { flushes += 1 })
    // A session/event listener that throws on the synthetic turn/end. Append
    // pushes before notifying, so turn/end is in the log (turn balanced) but the
    // throw must NOT skip the durability checkpoint — the flush decision is made
    // from the log, not a flag set after the (throwing) append.
    let threw = false
    ctx.on('session/event', (_s, event) => {
      if (!threw && event.type === 'turn/end') { threw = true; throw new Error('boom turn/end') }
    })

    expect(() => { agent.inject([{ type: 'text', text: 'notice' }], { source: { kind: 'plugin', plugin: 'p' } }) }).not.toThrow()
    const types = agent.session.events.map(e => e.type)
    expect(types).toEqual(['turn/start', 'context/message', 'turn/end']) // balanced
    await new Promise(r => setTimeout(r, 10))
    expect(flushes).toBe(1) // checkpoint fired despite the throwing turn/end listener
  })

  it('idle inject() reports a failing flush via agent/error (step 0) AND the logger', async () => {
    const adapter = new MockAdapter([textResponse('ok')])
    const ctx = await harness(adapter)
    // A non-Error rejection exercises the String() normalization branch.
    ctx.on('session/flush', () => { throw 'disk gone' })
    const warn = vi.spyOn(ctx.logger, 'warn').mockImplementation(() => undefined)
    const agent = ctx.agentLoop.create('a1', { model: 'mock' })
    const errors: { turn: number; step: number; message: string }[] = []
    ctx.on('agent/error', (_a, turn, step, error) => void errors.push({ turn, step, message: error.message }))

    agent.inject([{ type: 'text', text: 'notice' }], { source: { kind: 'plugin', plugin: 'p' } })
    await new Promise(r => setTimeout(r, 20)) // let the contained flush settle

    // Reported via agent/error (step 0 — the idle-injection convention) so
    // plugins monitoring agent/error see idle-injection persistence failures,
    // mirroring the loop's post-turn/end flush path. A non-Error throw is
    // normalized to an Error.
    expect(errors).toEqual([{ turn: 1, step: 0, message: 'disk gone' }])
    expect(warn).toHaveBeenCalledWith(expect.stringContaining('flush after idle injection failed'))
    warn.mockRestore()
  })

  it('idle inject() with a non-serializable source opens no turn (nothing to close)', async () => {
    const adapter = new MockAdapter([textResponse('ok')])
    const ctx = await harness(adapter)
    const agent = ctx.agentLoop.create('a1', { model: 'mock' })

    // A non-serializable source makes the turn/start append throw BEFORE the
    // event is pushed (Session.append validates before push), so NO turn opens.
    // The finally's isTurnOpen() guard sees no open turn and appends nothing —
    // the log stays empty, not left with a dangling turn/start.
    expect(() => {
      agent.inject([{ type: 'text', text: 'x' }], { source: { kind: 'plugin', plugin: 'p', bad: 1n } as never })
    }).toThrow(/non-JSON-serializable/)
    expect(agent.session.events).toHaveLength(0)
  })

  it('steer() when idle falls through to send() and starts a turn', async () => {
    const adapter = new MockAdapter([textResponse('ok')])
    const ctx = await harness(adapter)
    const agent = ctx.agentLoop.create('a1', { model: 'mock' })

    // steer while idle delegates to send
    agent.steer([{ type: 'text', text: 'steer idle' }], { source: { kind: 'plugin', plugin: 'test' } })
    await waitForIdle(ctx, agent)

    // The message was recorded as a user-level message (send path)
    expect(agent.session.events.some(e => e.type === 'user/message')).toBe(true)
    expect(adapter.requests).toHaveLength(1)
  })

  it('disposer is idempotent (double-stop)', async () => {
    // Create a bare LoopAgent and call start() directly to get the disposer.
    // Then call it twice — the second call hits the early-return branch.
    const ctx = new Context()
    await ctx.plugin(SessionStore)
    const session = ctx.sessions.create('test')
    const agent = new LoopAgent(ctx, AgentId('bare'), { model: 'mock' }, session)

    // Start the loop to get the disposer; the agent waits for messages
    // (idle, never-resolving cancel), so it will stay idle.
    const dispose = agent.start()

    // First dispose
    dispose()
    expect(agent.status).toBe('disposed')

    // Second dispose — idempotent, no throw
    expect(() => { dispose() }).not.toThrow()
    expect(agent.status).toBe('disposed')
  })

  it('setting the same status does not emit agent/status again', async () => {
    const adapter = new MockAdapter([textResponse('ok')])
    const ctx = await harness(adapter)
    const agent = ctx.agentLoop.create('a1', { model: 'mock' })

    const statuses: string[] = []
    ctx.on('agent/status', (subject, status) => {
      if (subject === agent) statuses.push(status)
    })

    send(agent, 'hi')
    await waitForIdle(ctx, agent)

    // After the turn, agent is idle. Send again to trigger another attempt
    // to go idle — but it's already idle, so no emission.
    const idleTransitionCount = statuses.filter(s => s === 'idle').length
    expect(idleTransitionCount).toBe(1) // only the final transition from running
  })

  it('whenIdle() resolves immediately when the agent is not running', async () => {
    const adapter = new MockAdapter([textResponse('ok')])
    const ctx = await harness(adapter)
    const agent = ctx.agentLoop.create('a1', { model: 'mock' })

    // Fresh agent is idle — whenIdle() takes the not-running fast path and
    // resolves without subscribing. await must not hang.
    await agent.whenIdle()
    expect(agent.status).not.toBe('running')
  })

  it('whenIdle() awaits the running→idle transition, ignoring other subjects/running events', async () => {
    const adapter = new MockAdapter([textResponse('ok'), textResponse('ok')])
    const ctx = await harness(adapter)
    const agent = ctx.agentLoop.create('a1', { model: 'mock' })
    const other = ctx.agentLoop.create('a2', { model: 'mock' })

    // Drive `agent` into `running`, then await whenIdle() — it subscribes to
    // agent/status and resolves on the first transition out of running.
    const running = new Promise<void>((resolve) => {
      const dispose = ctx.on('agent/status', (subject, status) => {
        if (subject === agent && status === 'running') { dispose(); resolve() }
      })
    })
    send(agent, 'go')
    await running
    expect(agent.status).toBe('running')

    // While `agent`'s whenIdle is pending, churn `other` through running→idle:
    // every status event it emits hits whenIdle's guard with `subject !== this`,
    // so the wait must ignore them and only resolve on `agent`'s own idle.
    send(other, 'go')

    await agent.whenIdle()
    expect(agent.status).toBe('idle')
  })

  it('whenIdle() subscribed while running resolves via done when the agent is then disposed', async () => {
    // Covers the waiter's disposed arm: whenIdle() queues an internal waiter
    // while running (not the fast path), then the disposer settles it and chains
    // `done` (loop exit), not an eager resolve. A bare LoopAgent + direct
    // start() disposer keeps the emit synchronous.
    const ctx = new Context()
    await ctx.plugin(LlmService)
    await ctx.plugin(SessionStore)
    await ctx.plugin(SystemPrompt)
    await ctx.plugin(ToolRegistry)
    await ctx.plugin(AgentRegistry)
    const adapter = new MockAdapter(['hang'])
    ctx.llm.registerAdapter(['mock'], adapter)
    const session = ctx.sessions.create('bare')
    const agent = new LoopAgent(ctx, AgentId('bare'), { model: 'mock' }, session)
    const dispose = agent.start()
    agent.send([{ type: 'text', text: 'go' }])
    await new Promise(r => setTimeout(r, 30))
    expect(agent.status).toBe('running')

    const idle = agent.whenIdle() // queues an internal waiter (running)
    dispose() // settles the waiter synchronously; whenIdle chains done
    await idle
    expect(agent.status).toBe('disposed')
    await agent.done
  })

  it('whenIdle() subscribed while running survives a FIBER dispose (no hung promise)', async () => {
    // The waiter is internal agent state, NOT an effect-scoped ctx.on listener:
    // disposing the OWNING fiber runs the agent's listener disposers, which would
    // have dropped a ctx.on-based waiter before the 'disposed' transition and
    // hung the promise. With internal waiters, the fiber disposer still settles
    // it. Regression for the round-3 whenIdle finding.
    const adapter = new MockAdapter(['hang'])
    const ctx = await harness(adapter)
    let agent!: LoopAgent
    const fiber = await ctx.plugin(Object.assign((inner: Context) => {
      agent = inner.agentLoop.create('scoped', { model: 'mock' })
    }, { inject: ['agentLoop'] }))
    send(agent, 'go')
    await new Promise(r => setTimeout(r, 30))
    expect(agent.status).toBe('running')

    const idle = agent.whenIdle() // queued while running
    await fiber.dispose()         // tears the fiber down (drops agent listeners)
    await idle                    // must resolve, not hang
    expect(agent.status).toBe('disposed')
  })

  it('whenIdle() on a disposed agent awaits the loop exit (done), not just the status flip', async () => {
    // The disposer emits agent/status('disposed') BEFORE the driver loop
    // unwinds, so whenIdle() must chain `done` (true quiescence) on the
    // disposed path. Dispose a running agent, then assert whenIdle() resolves
    // only after `done` — i.e. the loop has actually exited.
    const adapter = new MockAdapter(['hang'])
    const ctx = await harness(adapter)
    let agent!: LoopAgent
    const fiber = await ctx.plugin(Object.assign((inner: Context) => {
      agent = inner.agentLoop.create('scoped', { model: 'mock' })
    }, { inject: ['agentLoop'] }))
    send(agent, 'go')
    await new Promise(r => setTimeout(r, 30))

    let doneResolved = false
    void agent.done.then(() => { doneResolved = true })
    await fiber.dispose() // sets status disposed, aborts, drains the loop
    expect(agent.status).toBe('disposed')

    // whenIdle() must not resolve before `done` has — chaining `done` is the
    // quiescence guarantee. By here dispose() awaited the loop, so done is
    // settled; whenIdle resolves and done is observed resolved.
    await agent.whenIdle()
    expect(doneResolved).toBe(true)
  })

  it('abort() resolves reason to "aborted" when no reason provided', async () => {
    const adapter = new MockAdapter(['hang'])
    const ctx = await harness(adapter)
    const agent = ctx.agentLoop.create('a1', { model: 'mock' })

    const reasons: { kind: string; reason?: string }[] = []
    ctx.on('agent/turn-end', (_agent, _turn, reason) => void reasons.push(reason))

    send(agent, 'go')
    await new Promise(r => setTimeout(r, 30))
    agent.abort() // no reason string
    await waitForIdle(ctx, agent)

    expect(reasons[0]).toMatchObject({ kind: 'aborted', reason: 'aborted' })
  })
})