context-paging/tests/ContextPagingTest.php

<?php

declare(strict_types=1);

namespace ContextPaging\Tests;

use ContextPaging\ContextPaging;
use ContextPaging\TokenCounter;
use GuzzleHttp\Psr7\Response;
use GuzzleHttp\Psr7\ServerRequest;
use PHPUnit\Framework\TestCase;

class ContextPagingTest extends TestCase
{
    private ContextPaging $contextPaging;

    protected function setUp(): void
    {
        $this->contextPaging = new ContextPaging();
    }

    /**
     * Test basic fit with a small payload that doesn't need summarization.
     */
    public function testFitWithSmallPayload(): void
    {
        $messages = [
            ['role' => 'user', 'content' => 'Hello, how are you?'],
        ];

        $request = $this->createRequest($messages);

        $fitted = $this->contextPaging->fit($request);

        $this->assertTrue($fitted->getAttribute('context_fitted'));

        $fittedMessages = $fitted->getParsedBody()['messages'];
        $this->assertCount(1, $fittedMessages);
        $this->assertEquals('Hello, how are you?', $fittedMessages[0]['content']);
    }

    /**
     * Test fit with a larger payload that exceeds context limit.
     */
    public function testFitWithLargePayloadTriggersSummarization(): void
    {
        // Set a limit low enough to force summarization but high enough for last message
        $this->contextPaging->setMaxContextTokens(100)->setResponseReserve(20);

        $messages = [
            ['role' => 'user', 'content' => str_repeat('This is a long message that should be summarized. ', 50)],
            ['role' => 'assistant', 'content' => 'I understand your message.'],
            ['role' => 'user', 'content' => 'Short question'],
        ];

        $request = $this->createRequest($messages);

        $fitted = $this->contextPaging->fit($request);

        $fittedMessages = $fitted->getParsedBody()['messages'];

        // First message should be summarized
        $this->assertTrue($fittedMessages[0]['_summarized'] ?? false);
        $this->assertStringContainsString('[md5:', $fittedMessages[0]['content']);

        // Last message should NOT be summarized
        $lastIndex = count($fittedMessages) - 1;
        $this->assertFalse($fittedMessages[$lastIndex]['_summarized'] ?? false);

        // Should be under budget
        $tokens = $fitted->getAttribute('context_tokens');
        $this->assertLessThanOrEqual(80, $tokens);
    }

    /**
     * Test execute with no tool calls returns response as-is.
     */
    public function testExecuteWithNoToolCalls(): void
    {
        $messages = [
            ['role' => 'user', 'content' => 'Hello!'],
        ];

        $request = $this->createRequest($messages);

        $response = $this->contextPaging->execute($request, function (array $msgs, $req) {
            return new Response(200, [], json_encode(['choices' => [[
                'message' => ['role' => 'assistant', 'content' => 'Hi there!'],
            ]]]));
        });

        $body = json_decode($response->getBody()->getContents(), true);
        $this->assertEquals('Hi there!', $body['choices'][0]['message']['content']);
    }

    /**
     * Test that original messages are stored for dereferencing.
     */
    public function testOriginalMessagesStoredForDereferencing(): void
    {
        $longContent = str_repeat('This is a long message. ', 100);
        $messages = [
            ['role' => 'user', 'content' => $longContent],
            ['role' => 'user', 'content' => 'Short question'],
        ];

        $this->contextPaging->setMaxContextTokens(100)->setResponseReserve(20);

        $request = $this->createRequest($messages);
        $fitted = $this->contextPaging->fit($request);

        // The MD5 hash in the summarized message should reference the original
        $fittedMessages = $fitted->getParsedBody()['messages'];
        $this->assertMatchesRegularExpression('/\[md5:([a-f0-9]{32})\]/', $fittedMessages[0]['content']);

        // Extract MD5
        preg_match('/\[md5:([a-f0-9]{32})\]/', $fittedMessages[0]['content'], $matches);
        $md5 = $matches[1];

        // Verify the message store has the original
        // (We'd need to expose this or use reflection in a real test)
        $this->assertNotEmpty($md5);
    }

    /**
     * Test that last message is never summarized.
     */
    public function testLastMessageNeverSummarized(): void
    {
        $this->contextPaging->setMaxContextTokens(80)->setResponseReserve(15);

        $messages = [
            ['role' => 'user', 'content' => 'First message that is quite long and should be summarized'],
            ['role' => 'user', 'content' => 'Second message that is also quite long'],
            ['role' => 'user', 'content' => 'Short'],
        ];

        $request = $this->createRequest($messages);
        $fitted = $this->contextPaging->fit($request);

        $fittedMessages = $fitted->getParsedBody()['messages'];
        $lastIndex = count($fittedMessages) - 1;

        // Last message must not be summarized
        $this->assertFalse($fittedMessages[$lastIndex]['_summarized'] ?? false);
        $this->assertEquals('Short', $fittedMessages[$lastIndex]['content']);
    }

    /**
     * Test error when last message itself exceeds context.
     */
    public function testErrorWhenLastMessageTooLarge(): void
    {
        $this->contextPaging->setMaxContextTokens(20)->setResponseReserve(5);

        // Single message that's way too big
        $messages = [
            ['role' => 'user', 'content' => str_repeat('This is a massive message that will never fit. ', 100)],
        ];

        $request = $this->createRequest($messages);

        $this->expectException(\RuntimeException::class);
        $this->expectExceptionMessage('too large');

        $this->contextPaging->fit($request);
    }

    /**
     * Test that max_tokens in request is used for budget calculation.
     */
    public function testMaxTokensUsedForBudgetCalculation(): void
    {
        $this->contextPaging->setMaxContextTokens(65536);

        $messages = [
            ['role' => 'user', 'content' => 'Hello'],
        ];

        // Request with max_tokens: 8000
        $request = $this->createRequest($messages, ['max_tokens' => 8000]);

        $fitted = $this->contextPaging->fit($request);

        // Budget should be 65536 - 8000 = 57536
        $this->assertEquals(57536, $fitted->getAttribute('context_budget'));
    }

    /**
     * Test fallback to responseReserve when max_tokens not provided.
     */
    public function testFallbackToResponseReserveWhenNoMaxTokens(): void
    {
        $this->contextPaging->setMaxContextTokens(65536)->setResponseReserve(4096);

        $messages = [
            ['role' => 'user', 'content' => 'Hello'],
        ];

        // Request WITHOUT max_tokens
        $request = $this->createRequest($messages);

        $fitted = $this->contextPaging->fit($request);

        // Budget should be 65536 - 4096 = 61440
        $this->assertEquals(61440, $fitted->getAttribute('context_budget'));
    }

    // -----------------------------------------------------------------
    // Helpers
    // -----------------------------------------------------------------

    private function createRequest(array $messages, array $extraBody = []): ServerRequest
    {
        $body = array_merge(['messages' => $messages], $extraBody);

        return new ServerRequest(
            method: 'POST',
            uri: 'test://localhost',
            headers: ['Content-Type' => 'application/json'],
            body: json_encode($body),
            version: '1.1',
            serverParams: []
        )->withParsedBody($body);
    }
}