deepseek-harness/packages/context/time-context/tests/time-context.spec.ts

import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
import { Context } from 'cordis'
import Loader from '@cordisjs/plugin-loader'
import { CallId, LlmAdapter } from '@deepseek-ai/dsh-llm'
import type { GenerateOptions, StreamChunk } from '@deepseek-ai/dsh-llm'
import { Session, SessionId, type SessionEvent } from '@deepseek-ai/dsh-session'
import AgentRegistry, { agentEvents, AgentMessageId, type Agent } from '@deepseek-ai/dsh-agent'
import { defineContentToolFixture } from '@deepseek-ai/dsh-tools'
import AgentLoop from '@deepseek-ai/dsh-agent-loop'
import { mountAgentLoopTestDependencies } from '@deepseek-ai/dsh-agent-loop-testkit'
import * as timeContext from '@deepseek-ai/dsh-time-context'
import type { Config } from '@deepseek-ai/dsh-time-context'

const BASE = Date.parse('2026-07-14T00:00:00.000Z')
const ORIGINAL_TIME_ZONE = process.env['TZ']
const SIGNAL = new AbortController().signal

beforeEach(() => {
  process.env['TZ'] = 'UTC'
  vi.useFakeTimers()
  vi.setSystemTime(BASE)
})

afterEach(() => {
  vi.restoreAllMocks()
  vi.useRealTimers()
  if (ORIGINAL_TIME_ZONE === undefined) delete process.env['TZ']
  else process.env['TZ'] = ORIGINAL_TIME_ZONE
})

async function mount(config: Config = {}) {
  const ctx = new Context()
  await ctx.plugin(AgentRegistry)
  const fiber = await ctx.plugin(timeContext, config)
  return { ctx, fiber }
}

function sessionAgent(session: Session, id = 'agent'): Agent {
  return {
    id: SessionId(id),
    options: {},
    session,
    status: 'running',
    acceptsNextStep: true,
    ctx: new Context(),
    followup: () => AgentMessageId('stub'),
    steer: () => AgentMessageId('stub'),
    inject(input) {
      session.append('user/message', input, { surfaceOp: 'append' })
      return AgentMessageId('stub')
    },
    send: () => AgentMessageId('stub'),
    cancel() {},
    whenIdle: () => Promise.resolve(),
  }
}

function openMessageTurn(session: Session, turn: number): void {
  session.append('turn/start', { turn, trigger: { kind: 'message', source: { kind: 'user' } } })
  session.append('user/message', {
    content: [{ type: 'text', text: `turn ${turn}` }],
    source: { kind: 'user' },
  }, { surfaceOp: 'append' })
}

function contextTexts(session: Session): string[] {
  const texts: string[] = []
  for (const event of session.events) {
    if (event.type === 'user/message'
      && event.data.source.kind === 'plugin'
      && event.data.source.plugin === 'time-context') {
      texts.push(event.data.content.find(block => block.type === 'text')?.text ?? '')
    }
  }
  return texts
}

async function fire(
  ctx: Context,
  agent: Agent,
  turn: number,
  step: number,
  signal: AbortSignal = SIGNAL,
): Promise<void> {
  await agentEvents(ctx, agent).serial('agent/step', turn, step, signal)
}

function textResponse(text: string): StreamChunk[] {
  return [
    { type: 'block-start', index: 0, blockType: 'text' },
    { type: 'block-end', index: 0, block: { type: 'text', text } },
    { type: 'finish', reason: { kind: 'stop' } },
  ]
}

function toolCallResponse(): StreamChunk[] {
  return [
    { type: 'block-start', index: 0, blockType: 'tool-call' },
    {
      type: 'block-end',
      index: 0,
      block: { type: 'tool-call', id: CallId('tick-1'), name: 'tick', arguments: '{}' },
    },
    { type: 'finish', reason: { kind: 'tool-calls' } },
  ]
}

class ScriptedAdapter extends LlmAdapter {
  readonly requests: GenerateOptions[] = []

  constructor(private readonly script: StreamChunk[][]) {
    super()
  }

  override async * stream(options: GenerateOptions): AsyncIterable<StreamChunk> {
    this.requests.push(options)
    const chunks = this.script.shift()
    if (chunks === undefined) throw new Error('ScriptedAdapter: script exhausted')
    for (const chunk of chunks) yield chunk
  }
}

async function loopHarness(adapter: ScriptedAdapter, config: Config = {}): Promise<Context> {
  const ctx = new Context()
  await mountAgentLoopTestDependencies(ctx)
  await ctx.plugin(AgentLoop, { agents: [] })
  await ctx.plugin(timeContext, config)
  ctx.llm.registerAdapter(['mock'], adapter)
  return ctx
}

function requestText(request: GenerateOptions): string {
  return request.messages
    .flatMap(message => message.content)
    .filter(block => block.type === 'text')
    .map(block => block.text)
    .join('\n')
}

describe('durable step context', () => {
  it('records turn, step, zoned time, and the preceding model-visible message baseline', async () => {
    const { ctx } = await mount({ timeZone: 'Asia/Shanghai' })
    const session = new Session(SessionId('first'))
    openMessageTurn(session, 1)
    vi.setSystemTime(BASE + 90_061_000)

    await fire(ctx, sessionAgent(session), 1, 1)

    expect(contextTexts(session)).toEqual([
      'Time sampled while preparing turn 1, step 1: 2026-07-15T09:01:01+08:00[Asia/Shanghai]\n'
      + 'Elapsed since the preceding model-visible message: 1d 1h 1m 1s.',
    ])
    const event = session.events.at(-1)
    expect(event?.type).toBe('user/message')
    if (event?.type !== 'user/message') throw new Error('missing time context')
    expect(event.data.source).toEqual({ kind: 'plugin', plugin: 'time-context' })
    expect(event.surfaceOp).toBe('append')
  })

  it('reports an unavailable first-step baseline when no model-visible message precedes it', async () => {
    const { ctx } = await mount()
    const session = new Session(SessionId('unavailable'))
    session.append('turn/start', { turn: 1, trigger: { kind: 'message', source: { kind: 'user' } } })

    await fire(ctx, sessionAgent(session), 1, 1)

    expect(contextTexts(session)[0]).toContain(
      'Elapsed since the preceding model-visible message: unavailable.',
    )
  })

  it.each([
    ['omitted interval', {}],
    ['zero interval', { refreshIntervalMs: 0 }],
  ] as const)('uses the preceding durable step-context timestamp after step one with %s', async (_label, config) => {
    const { ctx } = await mount(config)
    const session = new Session(SessionId('later-step'))
    const agent = sessionAgent(session)
    openMessageTurn(session, 3)
    await fire(ctx, agent, 3, 1)
    vi.setSystemTime(BASE + 61_000)

    await fire(ctx, agent, 3, 2)

    expect(contextTexts(session)[1]).toBe(
      'Time sampled while preparing turn 3, step 2: 2026-07-14T00:01:01+00:00[UTC]\n'
      + 'Elapsed since the preceding step context: 1m 1s.',
    )
  })

  it('reports an unavailable later-step baseline at the matching turn boundary', async () => {
    const { ctx } = await mount()
    const session = new Session(SessionId('later-step-boundary'))
    openMessageTurn(session, 4)

    await fire(ctx, sessionAgent(session), 4, 2)

    expect(contextTexts(session)[0]).toContain(
      'Elapsed since the preceding step context: unavailable.',
    )
  })

  it('reports an unavailable later-step baseline when event lookup is exhausted', async () => {
    const { ctx } = await mount()
    const session = new Session(SessionId('later-step-exhausted'))

    await fire(ctx, sessionAgent(session), 1, 2)

    expect(contextTexts(session)[0]).toContain(
      'Elapsed since the preceding step context: unavailable.',
    )
  })

  it('injects after backward wall-clock movement and clamps elapsed time to zero', async () => {
    const { ctx } = await mount({ refreshIntervalMs: 60_000 })
    const session = new Session(SessionId('backward'))
    const agent = sessionAgent(session)
    openMessageTurn(session, 1)
    await fire(ctx, agent, 1, 1)
    vi.setSystemTime(BASE - 5_000)

    await fire(ctx, agent, 1, 2)

    expect(contextTexts(session)).toHaveLength(2)
    expect(contextTexts(session)[1]).toContain('Elapsed since the preceding step context: 0s.')
  })

  it('uses a shadowed durable injection after resume and injects at the exact threshold', async () => {
    const { ctx } = await mount({ refreshIntervalMs: 1_000 })
    const original = new Session(SessionId('seed-source'))
    openMessageTurn(original, 1)
    await fire(ctx, sessionAgent(original), 1, 1)
    const user = original.events.find(event => event.type === 'user/message' && event.data.source.kind === 'user')
    const reading = original.events.find(event => event.type === 'user/message' && event.data.source.kind === 'plugin')
    if (user === undefined || reading === undefined) throw new Error('missing source surface events')
    original.append('user/message', {
      content: [{ type: 'text', text: 'compacted history' }],
      source: { kind: 'plugin', plugin: 'compact-basic' },
    }, {
      surfaceOp: { op: 'replace', start: user.seq, end: reading.seq },
      sourceEventSeqs: [user.seq, reading.seq],
    })
    original.append('turn/end', { turn: 1, reason: { kind: 'completed' } })
    expect(JSON.stringify(original.deriveMessages())).not.toContain('Time sampled while preparing')

    const resumed = new Session(SessionId('resumed'), [...original.events])
    const resumedAgent = sessionAgent(resumed)
    vi.setSystemTime(BASE + 999)
    openMessageTurn(resumed, 2)
    const beforeSkip = resumed.events.length

    await fire(ctx, resumedAgent, 2, 1)

    expect(resumed.events).toHaveLength(beforeSkip)
    expect(contextTexts(resumed)).toHaveLength(1)

    vi.setSystemTime(BASE + 1_000)
    await fire(ctx, resumedAgent, 2, 2)

    expect(contextTexts(resumed)).toHaveLength(2)
    expect(contextTexts(resumed)[1]).toContain(
      'Elapsed since the preceding step context: unavailable.',
    )
  })

  it('applies a positive interval across turns without sharing state between sessions', async () => {
    const { ctx } = await mount({ refreshIntervalMs: 1_000 })
    const first = new Session(SessionId('interval-first'))
    const firstAgent = sessionAgent(first, 'first-agent')
    openMessageTurn(first, 1)
    await fire(ctx, firstAgent, 1, 1)
    first.append('turn/end', { turn: 1, reason: { kind: 'completed' } })

    vi.setSystemTime(BASE + 500)
    openMessageTurn(first, 2)
    const beforeSkip = first.events.length
    await fire(ctx, firstAgent, 2, 1)

    const independent = new Session(SessionId('interval-independent'))
    openMessageTurn(independent, 1)
    await fire(ctx, sessionAgent(independent, 'independent-agent'), 1, 1)

    expect(first.events).toHaveLength(beforeSkip)
    expect(contextTexts(first)).toHaveLength(1)
    expect(contextTexts(independent)).toHaveLength(1)
  })

  it('runs before ordinary pre-step listeners and skips an already-aborted step', async () => {
    const { ctx } = await mount()
    const session = new Session(SessionId('ordering'))
    const agent = sessionAgent(session)
    openMessageTurn(session, 1)
    let ordinarySawContext = false
    ctx.on('agent/step', (subject) => {
      ordinarySawContext = subject.session.events.some(event => event.type === 'user/message')
    })

    await fire(ctx, agent, 1, 1)
    const abort = new AbortController()
    abort.abort()
    await fire(ctx, agent, 1, 2, abort.signal)

    expect(ordinarySawContext).toBe(true)
    expect(contextTexts(session)).toHaveLength(1)
  })
})

describe('configuration and lifecycle', () => {
  it('defaults to the process system zone and retains the zone resolved at plugin load', async () => {
    process.env['TZ'] = 'Asia/Shanghai'
    const { ctx } = await mount()
    process.env['TZ'] = 'America/New_York'
    const session = new Session(SessionId('system-zone'))
    openMessageTurn(session, 1)

    await fire(ctx, sessionAgent(session), 1, 1)

    expect(contextTexts(session)[0]).toContain('2026-07-14T08:00:00+08:00[Asia/Shanghai]')
  })

  it('fails loud for an invalid explicit zone or an unavailable process zone', async () => {
    const invalid = new Context()
    await invalid.plugin(AgentRegistry)
    await expect(invalid.plugin(timeContext, { timeZone: 'Not/A_Real_Zone' })).rejects.toThrow(
      /invalid IANA timeZone/,
    )

    vi.spyOn(Intl, 'DateTimeFormat').mockImplementationOnce(() => {
      throw new RangeError('system zone unavailable')
    })
    const unresolved = new Context()
    await unresolved.plugin(AgentRegistry)
    await expect(unresolved.plugin(timeContext, {})).rejects.toThrow(/failed to resolve the system time zone/)
  })

  it('rejects invalid refresh intervals at plugin load with one diagnostic', async () => {
    const invalid = [-1, 0.5, Number.MAX_SAFE_INTEGER + 1, Number.POSITIVE_INFINITY, Number.NaN]
    for (const refreshIntervalMs of invalid) {
      await expect(mount({ refreshIntervalMs })).rejects.toThrow(
        'time-context: refreshIntervalMs must be a non-negative safe integer',
      )
    }
  })

  it('removes its listener when the plugin fiber disposes', async () => {
    const { ctx, fiber } = await mount()
    const session = new Session(SessionId('dispose'))
    const agent = sessionAgent(session)
    openMessageTurn(session, 1)
    await fire(ctx, agent, 1, 1)

    await fiber.dispose()
    await fire(ctx, agent, 1, 2)

    expect(contextTexts(session)).toHaveLength(1)
  })
})

describe('real agent-loop request history', () => {
  it.each([
    ['throws', 'error'],
    ['cancels', 'aborted'],
  ] as const)('discards the pending preparation reading when a later step listener %s', async (mode, reasonKind) => {
    const adapter = new ScriptedAdapter([textResponse('unused')])
    const ctx = await loopHarness(adapter)
    let laterSawReading = false
    ctx.on('agent/step', (subject) => {
      laterSawReading = contextTexts(subject.session).length === 1
      if (mode === 'throws') throw new Error('later pre-step failure')
      subject.cancel({ kind: 'user' })
    })
    const agent = ctx.agentLoop.create(SessionId(`late-${mode}`), { provider: 'mock', model: 'mock' })

    agent.followup({ content: [{ type: 'text', text: 'start' }], source: { kind: 'user' } })
    await agent.whenIdle()

    expect(laterSawReading).toBe(false)
    expect(contextTexts(agent.session)).toHaveLength(0)
    expect(adapter.requests).toHaveLength(0)
    expect(agent.session.events.some(event => event.type === 'step/start')).toBe(false)
    const turnEnd = agent.session.events.findLast(event => event.type === 'turn/end')
    expect(turnEnd?.type === 'turn/end' && turnEnd.data.reason.kind).toBe(reasonKind)
    await ctx.fiber.dispose()
  })

  it('persists one ordered context per request, accumulates readings, and leaves system headers unchanged', async () => {
    const adapter = new ScriptedAdapter([toolCallResponse(), textResponse('done')])
    const ctx = await loopHarness(adapter)
    ctx.tools.register(defineContentToolFixture({
      name: 'tick',
      description: 'advance fake time',
      parameters: {},
      async execute() {
        vi.setSystemTime(BASE + 61_000)
        return [{ type: 'text' as const, text: 'advanced' }]
      },
    }))
    const agent = ctx.agentLoop.create(SessionId('loop'), { provider: 'mock', model: 'mock' })

    agent.followup({ content: [{ type: 'text', text: 'start' }], source: { kind: 'user' } })
    await agent.whenIdle()

    expect(adapter.requests).toHaveLength(2)
    const contexts = agent.session.events.filter(
      (event): event is SessionEvent<'user/message'> => event.type === 'user/message' && event.data.source.kind === 'plugin')
    const starts = agent.session.events.filter(event => event.type === 'step/start')
    expect(contexts).toHaveLength(adapter.requests.length)
    expect(starts).toHaveLength(adapter.requests.length)
    for (let index = 0; index < contexts.length; index += 1) {
      expect(contexts[index]!.seq).toBeLessThan(starts[index]!.seq)
    }
    expect(contexts.every(event => event.data.source.kind === 'plugin'
      && event.data.source.plugin === 'time-context'
      && event.surfaceOp === 'append')).toBe(true)

    const firstRequestText = requestText(adapter.requests[0]!)
    const secondRequestText = requestText(adapter.requests[1]!)
    expect(firstRequestText).toContain('Time sampled while preparing turn 1, step 1:')
    expect(firstRequestText).toContain('Elapsed since the preceding model-visible message: 0s.')
    expect(firstRequestText).not.toContain('Time sampled while preparing turn 1, step 2:')
    expect(secondRequestText).toContain('Time sampled while preparing turn 1, step 1:')
    expect(secondRequestText).toContain('Time sampled while preparing turn 1, step 2:')
    expect(secondRequestText).toContain('Elapsed since the preceding step context: 1m 1s.')

    for (const request of adapter.requests) expect(request.system).not.toContain('Time sampled while preparing')
    const headers = agent.session.events.filter(event => event.type === 'request/header')
    expect(JSON.stringify(headers)).not.toContain('Time sampled while preparing')
    await ctx.fiber.dispose()
  })
})

describe('real Loader export path', () => {
  it('keeps namespace metadata and boots the agent listener through unwrapExports', async () => {
    expect('default' in timeContext).toBe(false)
    const loader = Object.create(Loader.prototype) as Loader
    const unwrapped = loader.unwrapExports(timeContext) as Record<string, unknown>
    expect(unwrapped).toBe(timeContext)
    expect(unwrapped.name).toBe('time-context')
    expect(unwrapped.inject).toEqual(['agents'])
    expect(unwrapped.Config).toBeDefined()
    expect(typeof unwrapped.apply).toBe('function')

    const ctx = new Context()
    await ctx.plugin(AgentRegistry)
    const plugin = loader.unwrapExports(timeContext) as Parameters<Context['plugin']>[0]
    await ctx.plugin(plugin)
    const session = new Session(SessionId('loader'))
    openMessageTurn(session, 1)
    await fire(ctx, sessionAgent(session), 1, 1)
    expect(contextTexts(session)[0]).toContain('Time sampled while preparing turn 1, step 1:')
  })
})