import { describe, it, expect, vi } from 'vitest';
import { createTokenUsageMiddleware } from '../../context-manager/modes/agentic/agentic-context.js';
import { Message, TextBlock } from '../../types/messages.js';
import { Agent } from '../../agent/agent.js';
import { MockMessageModel } from '../../__fixtures__/mock-message-model.js';
function createMockModel(contextWindowLimit) {
    return {
        getConfig: () => ({ contextWindowLimit }),
    };
}
function createContext(overrides = {}) {
    return {
        agent: {},
        messages: [new Message({ role: 'user', content: [new TextBlock('hello')] })],
        toolSpecs: [],
        invocationState: {},
        ...overrides,
    };
}
describe('createTokenUsageMiddleware', () => {
    it('returns context unchanged when projectedInputTokens is not set', async () => {
        const middleware = createTokenUsageMiddleware(createMockModel(200_000));
        const context = createContext();
        const result = await middleware(context);
        expect(result).toBe(context);
    });
    it('returns context unchanged when messages are empty', async () => {
        const middleware = createTokenUsageMiddleware(createMockModel(200_000));
        const context = createContext({
            messages: [],
            projectedInputTokens: 50_000,
        });
        const result = await middleware(context);
        expect(result).toBe(context);
    });
    it('appends context status to last message content', async () => {
        const middleware = createTokenUsageMiddleware(createMockModel(200_000));
        const context = createContext({
            projectedInputTokens: 50_000,
        });
        const result = await middleware(context);
        expect(result.messages.length).toBe(1);
        const lastMsg = result.messages[0];
        expect(lastMsg.content.length).toBe(2);
        const statusBlock = lastMsg.content[1];
        expect(statusBlock.text).toContain('<context-status>');
        expect(statusBlock.text).toContain('25.0%');
        expect(statusBlock.text).toContain('<remaining>');
        expect(statusBlock.text).toContain('</context-status>');
    });
    it('does not mutate the original messages array', async () => {
        const middleware = createTokenUsageMiddleware(createMockModel(200_000));
        const originalMessage = new Message({ role: 'user', content: [new TextBlock('hello')] });
        const context = createContext({
            messages: [originalMessage],
            projectedInputTokens: 50_000,
        });
        const result = await middleware(context);
        expect(result.messages).not.toBe(context.messages);
        expect(originalMessage.content.length).toBe(1);
    });
    it('uses DEFAULT_CONTEXT_WINDOW_LIMIT when model config has no limit', async () => {
        const middleware = createTokenUsageMiddleware(createMockModel(undefined));
        const context = createContext({
            projectedInputTokens: 100_000,
        });
        const result = await middleware(context);
        const statusBlock = result.messages[0].content[1];
        expect(statusBlock.text).toContain('<context-status>');
        expect(statusBlock.text).toContain('50.0%');
        expect(statusBlock.text).toContain('200,000');
    });
    it('preserves message metadata', async () => {
        const middleware = createTokenUsageMiddleware(createMockModel(200_000));
        const metadata = { usage: { inputTokens: 10, outputTokens: 5, totalTokens: 15 } };
        const context = createContext({
            messages: [new Message({ role: 'user', content: [new TextBlock('hello')], metadata })],
            projectedInputTokens: 50_000,
        });
        const result = await middleware(context);
        expect(result.messages[0].metadata).toEqual(metadata);
    });
    it('reports correct remaining tokens', async () => {
        const middleware = createTokenUsageMiddleware(createMockModel(100_000));
        const context = createContext({
            projectedInputTokens: 80_000,
        });
        const result = await middleware(context);
        const statusBlock = result.messages[0].content[1];
        expect(statusBlock.text).toContain('80.0%');
        expect(statusBlock.text).toContain('<remaining>~20,000 tokens</remaining>');
    });
});
describe('token usage middleware integration', () => {
    it('injects context status into messages sent to model on second call', async () => {
        const model = new MockMessageModel()
            .addTurn({ type: 'textBlock', text: 'First response' }, { usage: { inputTokens: 1000, outputTokens: 200, totalTokens: 1200 } })
            .addTurn({ type: 'textBlock', text: 'Second response' }, { usage: { inputTokens: 2000, outputTokens: 300, totalTokens: 2300 } });
        model.updateConfig({ contextWindowLimit: 100_000 });
        const agent = new Agent({ model, contextManager: 'agentic', printer: false });
        const streamSpy = vi.spyOn(model, 'stream');
        await agent.invoke('First message');
        await agent.invoke('Second message');
        const secondCallMessages = streamSpy.mock.calls[1][0];
        const lastMessage = secondCallMessages[secondCallMessages.length - 1];
        const lastBlock = lastMessage.content[lastMessage.content.length - 1];
        expect(lastBlock.text).toContain('<context-status>');
        expect(lastBlock.text).toContain('100,000');
        expect(lastBlock.text).toContain('<remaining>');
    });
    it('injects on first call when token estimation succeeds', async () => {
        const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'First response' }, { usage: { inputTokens: 1000, outputTokens: 200, totalTokens: 1200 } });
        model.updateConfig({ contextWindowLimit: 100_000 });
        const agent = new Agent({ model, contextManager: 'agentic', printer: false });
        const streamSpy = vi.spyOn(model, 'stream');
        await agent.invoke('First message');
        const firstCallMessages = streamSpy.mock.calls[0][0];
        const lastMessage = firstCallMessages[firstCallMessages.length - 1];
        const lastBlock = lastMessage.content[lastMessage.content.length - 1];
        expect(lastBlock.text).toContain('<context-status>');
        expect(lastBlock.text).toContain('100,000');
    });
    it('does not register middleware when contextManager is not agentic', async () => {
        const model = new MockMessageModel()
            .addTurn({ type: 'textBlock', text: 'First response' }, { usage: { inputTokens: 1000, outputTokens: 200, totalTokens: 1200 } })
            .addTurn({ type: 'textBlock', text: 'Second response' }, { usage: { inputTokens: 2000, outputTokens: 300, totalTokens: 2300 } });
        model.updateConfig({ contextWindowLimit: 100_000 });
        const agent = new Agent({ model, printer: false });
        const streamSpy = vi.spyOn(model, 'stream');
        await agent.invoke('First message');
        await agent.invoke('Second message');
        const secondCallMessages = streamSpy.mock.calls[1][0];
        const lastMessage = secondCallMessages[secondCallMessages.length - 1];
        const lastBlock = lastMessage.content[lastMessage.content.length - 1];
        expect(lastBlock.text).not.toContain('<context-status>');
    });
});
//# sourceMappingURL=token-usage-middleware.test.js.map