284 lines
8.9 KiB
TypeScript
284 lines
8.9 KiB
TypeScript
import { beforeEach, describe, expect, it, vi } from 'vitest'
|
|
|
|
const addMessageMock = vi.fn()
|
|
const updateSessionStatsMock = vi.fn()
|
|
const updateUsageMock = vi.fn()
|
|
const calcAndUpdateUsageMock = vi.fn()
|
|
const buildDbHistoryMock = vi.fn()
|
|
const buildSnapshotAwareHistoryMock = vi.fn()
|
|
const estimateUsageTokensFromMessagesMock = vi.fn()
|
|
|
|
vi.mock('../../packages/server/src/db/hermes/session-store', () => ({
|
|
addMessage: addMessageMock,
|
|
createSession: vi.fn(),
|
|
getSession: vi.fn(() => ({ id: 'session-resume', profile: 'default', model: 'gpt-test', provider: 'openai' })),
|
|
updateSession: vi.fn(),
|
|
updateSessionStats: updateSessionStatsMock,
|
|
}))
|
|
|
|
vi.mock('../../packages/server/src/db/hermes/usage-store', () => ({
|
|
updateUsage: updateUsageMock,
|
|
}))
|
|
|
|
vi.mock('../../packages/server/src/services/logger', () => ({
|
|
logger: { info: vi.fn(), warn: vi.fn(), error: vi.fn(), debug: vi.fn() },
|
|
bridgeLogger: { info: vi.fn(), warn: vi.fn(), error: vi.fn(), debug: vi.fn() },
|
|
}))
|
|
|
|
vi.mock('../../packages/server/src/lib/llm-prompt', () => ({
|
|
getSystemPrompt: vi.fn(() => 'system prompt'),
|
|
}))
|
|
|
|
vi.mock('../../packages/server/src/lib/context-compressor', () => ({
|
|
countTokens: vi.fn(() => 1),
|
|
SUMMARY_PREFIX: '[Summary] ',
|
|
}))
|
|
|
|
vi.mock('../../packages/server/src/db/hermes/compression-snapshot', () => ({
|
|
getCompressionSnapshot: vi.fn(),
|
|
}))
|
|
|
|
vi.mock('../../packages/server/src/services/hermes/run-chat/compression', async () => {
|
|
const actual = await vi.importActual<any>('../../packages/server/src/services/hermes/run-chat/compression')
|
|
return {
|
|
...actual,
|
|
buildDbHistory: buildDbHistoryMock,
|
|
buildSnapshotAwareHistory: buildSnapshotAwareHistoryMock,
|
|
buildCompressedHistory: vi.fn(),
|
|
forceCompressBridgeHistory: vi.fn(),
|
|
}
|
|
})
|
|
|
|
vi.mock('../../packages/server/src/services/hermes/run-chat/usage', () => ({
|
|
calcAndUpdateUsage: calcAndUpdateUsageMock,
|
|
contextTokensWithCachedOverhead: vi.fn((_state, tokens) => tokens),
|
|
estimateUsageTokensFromMessages: estimateUsageTokensFromMessagesMock,
|
|
getCachedBridgeContextOverhead: vi.fn(() => undefined),
|
|
updateMessageContextTokenUsage: vi.fn((_sid, state, _emit, tokens) => {
|
|
state.contextTokens = tokens
|
|
return tokens
|
|
}),
|
|
}))
|
|
|
|
function createNamespace() {
|
|
const emitted: Array<{ event: string; payload: any }> = []
|
|
return {
|
|
emitted,
|
|
nsp: {
|
|
adapter: { rooms: { get: vi.fn(() => new Set(['socket-1'])) } },
|
|
to: vi.fn(() => ({
|
|
emit: vi.fn((event: string, payload: any) => emitted.push({ event, payload })),
|
|
})),
|
|
},
|
|
}
|
|
}
|
|
|
|
describe('resumeBridgeRun', () => {
|
|
beforeEach(() => {
|
|
vi.clearAllMocks()
|
|
addMessageMock.mockReturnValue(42)
|
|
calcAndUpdateUsageMock.mockResolvedValue({ inputTokens: 3, outputTokens: 2 })
|
|
buildDbHistoryMock.mockResolvedValue([
|
|
{ role: 'user', content: 'hello' },
|
|
{ role: 'assistant', content: 'Hello world' },
|
|
])
|
|
buildSnapshotAwareHistoryMock.mockResolvedValue([
|
|
{ role: 'user', content: 'hello' },
|
|
{ role: 'assistant', content: 'Hello world' },
|
|
])
|
|
estimateUsageTokensFromMessagesMock.mockReturnValue({ inputTokens: 3, outputTokens: 2 })
|
|
})
|
|
|
|
it('continues polling an existing bridge run after web server state is recreated', async () => {
|
|
const { resumeBridgeRun } = await import('../../packages/server/src/services/hermes/run-chat/handle-bridge-run')
|
|
const { nsp, emitted } = createNamespace()
|
|
const socket = { id: 'socket-1', connected: true, emit: vi.fn() }
|
|
const sessionMap = new Map<string, any>()
|
|
sessionMap.set('session-resume', {
|
|
messages: [
|
|
{ id: 1, session_id: 'session-resume', role: 'user', content: 'hello', timestamp: 1 },
|
|
{ id: 2, session_id: 'session-resume', role: 'assistant', content: 'Hello', timestamp: 2 },
|
|
],
|
|
isWorking: true,
|
|
events: [],
|
|
queue: [],
|
|
})
|
|
|
|
const bridge = {
|
|
getResult: vi.fn(async () => ({
|
|
ok: true,
|
|
run_id: 'run-resume',
|
|
session_id: 'session-resume',
|
|
status: 'running',
|
|
output: 'Hello',
|
|
deltas: ['Hello'],
|
|
events: [{
|
|
event: 'model.usage',
|
|
api_request_id: 'request-before-resume',
|
|
turn_id: 'turn-1',
|
|
api_call_count: 1,
|
|
model: 'gpt-test-response',
|
|
provider: 'openai',
|
|
usage: {
|
|
input_tokens: 80,
|
|
output_tokens: 12,
|
|
cache_read_tokens: 50,
|
|
cache_write_tokens: 3,
|
|
reasoning_tokens: 4,
|
|
},
|
|
}],
|
|
})),
|
|
getOutput: vi.fn(async () => ({
|
|
ok: true,
|
|
run_id: 'run-resume',
|
|
session_id: 'session-resume',
|
|
status: 'complete',
|
|
delta: ' world',
|
|
cursor: 2,
|
|
output: 'Hello world',
|
|
done: true,
|
|
result: { final_response: 'Hello world' },
|
|
error: null,
|
|
events: [],
|
|
event_cursor: 0,
|
|
})),
|
|
contextEstimate: vi.fn(async () => ({
|
|
ok: true,
|
|
session_id: 'session-resume',
|
|
fixed_context_tokens: 0,
|
|
system_prompt_tokens: 0,
|
|
tool_tokens: 0,
|
|
message_count: 0,
|
|
tool_count: 0,
|
|
system_prompt_chars: 0,
|
|
})),
|
|
goalEvaluate: vi.fn(async () => ({
|
|
ok: true,
|
|
session_id: 'session-resume',
|
|
handled: true,
|
|
should_continue: false,
|
|
})),
|
|
}
|
|
|
|
await resumeBridgeRun(
|
|
nsp as any,
|
|
socket as any,
|
|
{
|
|
sessionId: 'session-resume',
|
|
runId: 'run-resume',
|
|
profile: 'default',
|
|
instructions: 'system prompt',
|
|
model: 'gpt-test',
|
|
provider: 'openai',
|
|
},
|
|
sessionMap,
|
|
bridge as any,
|
|
vi.fn(),
|
|
)
|
|
|
|
expect(bridge.getResult).toHaveBeenCalledWith('run-resume')
|
|
expect(bridge.getOutput).toHaveBeenCalledWith('run-resume', 1, 1)
|
|
expect(updateUsageMock).toHaveBeenCalledWith('session-resume', expect.objectContaining({
|
|
runId: 'run-resume:api:request-before-resume',
|
|
source: 'hermes',
|
|
usageScope: 'model_call',
|
|
apiCalls: 1,
|
|
inputTokens: 80,
|
|
outputTokens: 12,
|
|
cacheReadTokens: 50,
|
|
cacheWriteTokens: 3,
|
|
reasoningTokens: 4,
|
|
model: 'gpt-test-response',
|
|
provider: 'openai',
|
|
isEstimated: false,
|
|
}))
|
|
expect(updateUsageMock).toHaveBeenCalledTimes(1)
|
|
expect(emitted).toEqual(expect.arrayContaining([
|
|
expect.objectContaining({
|
|
event: 'message.delta',
|
|
payload: expect.objectContaining({ delta: ' world', session_id: 'session-resume' }),
|
|
}),
|
|
expect.objectContaining({
|
|
event: 'run.completed',
|
|
payload: expect.objectContaining({ output: 'Hello world', session_id: 'session-resume' }),
|
|
}),
|
|
]))
|
|
expect(sessionMap.get('session-resume').isWorking).toBe(false)
|
|
})
|
|
|
|
it('completes a timed-out abort when the resumed bridge run reaches a terminal state', async () => {
|
|
const { resumeBridgeRun } = await import('../../packages/server/src/services/hermes/run-chat/handle-bridge-run')
|
|
const { nsp, emitted } = createNamespace()
|
|
const socket = { id: 'socket-1', connected: true, emit: vi.fn() }
|
|
const sessionMap = new Map<string, any>()
|
|
sessionMap.set('session-resume', {
|
|
messages: [
|
|
{ id: 1, session_id: 'session-resume', role: 'user', content: 'hello', timestamp: 1 },
|
|
{ id: 2, session_id: 'session-resume', role: 'assistant', content: 'Hello', timestamp: 2, isStreaming: true },
|
|
],
|
|
isWorking: true,
|
|
isAborting: true,
|
|
runId: 'run-resume',
|
|
profile: 'default',
|
|
source: 'cli',
|
|
events: [],
|
|
queue: [],
|
|
})
|
|
|
|
const bridge = {
|
|
getResult: vi.fn(async () => ({
|
|
ok: true,
|
|
run_id: 'run-resume',
|
|
session_id: 'session-resume',
|
|
status: 'running',
|
|
output: 'Hello',
|
|
deltas: ['Hello'],
|
|
events: [],
|
|
})),
|
|
getOutput: vi.fn(async () => ({
|
|
ok: true,
|
|
run_id: 'run-resume',
|
|
session_id: 'session-resume',
|
|
status: 'interrupted',
|
|
delta: '',
|
|
cursor: 1,
|
|
output: 'Hello',
|
|
done: true,
|
|
result: { interrupted: true, completed: false, final_response: 'Operation interrupted' },
|
|
error: null,
|
|
events: [],
|
|
event_cursor: 0,
|
|
})),
|
|
}
|
|
|
|
await resumeBridgeRun(
|
|
nsp as any,
|
|
socket as any,
|
|
{
|
|
sessionId: 'session-resume',
|
|
runId: 'run-resume',
|
|
profile: 'default',
|
|
instructions: 'system prompt',
|
|
model: 'gpt-test',
|
|
provider: 'openai',
|
|
},
|
|
sessionMap,
|
|
bridge as any,
|
|
vi.fn(),
|
|
)
|
|
|
|
expect(emitted).toEqual(expect.arrayContaining([
|
|
expect.objectContaining({
|
|
event: 'abort.completed',
|
|
payload: expect.objectContaining({ session_id: 'session-resume', run_id: 'run-resume', synced: true }),
|
|
}),
|
|
]))
|
|
expect(emitted.some(item => item.event === 'run.failed')).toBe(false)
|
|
expect(sessionMap.get('session-resume')).toEqual(expect.objectContaining({
|
|
isWorking: false,
|
|
isAborting: false,
|
|
runId: undefined,
|
|
}))
|
|
})
|
|
})
|