595 lines
21 KiB
Go
595 lines
21 KiB
Go
|
// Package bot — Milestone 3 payload assembly tests.
|
|||
|
|
//
|
|||
|
|
// These tests cover the four payload-assembly helpers backend-eng adds
|
|||
|
|
// in M3 (per plan §7 mapping table):
|
|||
|
|
//
|
|||
|
|
// - Conversation history: last RecentTurns completed turns × 2 messages,
|
|||
|
|
// ascending ordinal.
|
|||
|
|
// - Scope payload assembly: stored doc ids ∪ folder-expanded doc ids,
|
|||
|
|
// with the original folder ids preserved in the payload.
|
|||
|
|
// - Cached-results stripping on updated_session_state.
|
|||
|
|
// - 64KB session-state cap (chatbot.StateMaxBytes).
|
|||
|
|
// - Metadata summarization (sum input/output tokens across the
|
|||
|
|
// metadata.agents map; wall-clock fallback when metadata is absent).
|
|||
|
|
//
|
|||
|
|
// SQL-driven helpers (conversation history, scope expansion) use the
|
|||
|
|
// real Postgres testcontainer. Pure functions (strip, cap, summarize)
|
|||
|
|
// run in-memory. No httptest.NewServer here — that lives in
|
|||
|
|
// fastapi_test.go for transport-layer cases only.
|
|||
|
|
//
|
|||
|
|
// Compile contract: this file references bot.BuildConversationHistory,
|
|||
|
|
// bot.BuildChatScope, bot.StripCachedResults, bot.SummarizeMetadata,
|
|||
|
|
// bot.SessionStateOversize, and the previously-declared types from
|
|||
|
|
// fastapi_test.go (ChatScope, ConversationMessage, Metadata,
|
|||
|
|
// AgentTelemetry). Backend-eng adds payload.go / state_strip.go /
|
|||
|
|
// metadata_summarise.go and the package compiles.
|
|||
|
|
package bot_test
|
|||
|
|
|
|||
|
|
import (
|
|||
|
|
"context"
|
|||
|
|
"encoding/json"
|
|||
|
|
"os"
|
|||
|
|
"strings"
|
|||
|
|
"testing"
|
|||
|
|
"time"
|
|||
|
|
|
|||
|
|
"queryorchestration/internal/bot"
|
|||
|
|
"queryorchestration/internal/database/repository"
|
|||
|
|
"queryorchestration/internal/serviceconfig"
|
|||
|
|
"queryorchestration/internal/serviceconfig/chatbot"
|
|||
|
|
"queryorchestration/internal/test"
|
|||
|
|
|
|||
|
|
"github.com/google/uuid"
|
|||
|
|
"github.com/stretchr/testify/assert"
|
|||
|
|
"github.com/stretchr/testify/require"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
// payloadDefaultCfg returns a chatbot.ChatbotConfig with the plan §8
|
|||
|
|
// defaults so the test does not have to repeat the literal values. The
|
|||
|
|
// SERVICE_URL/API_KEY values are placeholders; payload tests do not
|
|||
|
|
// dial out.
|
|||
|
|
func payloadDefaultCfg() chatbot.ChatbotConfig {
|
|||
|
|
return chatbot.ChatbotConfig{
|
|||
|
|
ServiceURL: "http://chatbot.test:8000",
|
|||
|
|
APIKey: "test-api-key",
|
|||
|
|
RequestTimeoutSeconds: 60,
|
|||
|
|
MaxTurnChars: 2000,
|
|||
|
|
RecentTurns: 5,
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// resetClientForPayload mirrors resetClientForBotTests but lives in the
|
|||
|
|
// bot package's test space so this file does not pull in queryapi_test
|
|||
|
|
// helpers. The shared testcontainer Postgres reuses volumes across runs,
|
|||
|
|
// so leftover bot rows must be wiped before CreateClient.
|
|||
|
|
func resetClientForPayload(t *testing.T, ctx context.Context, cfg *serviceconfig.BaseConfig, clientID, name string) {
|
|||
|
|
t.Helper()
|
|||
|
|
pool := cfg.GetDBPool()
|
|||
|
|
stmts := []string{
|
|||
|
|
`DELETE FROM document_custom_metadata WHERE document_id IN (SELECT id FROM documents WHERE clientId = $1)`,
|
|||
|
|
`UPDATE documents SET custom_schema_id = NULL WHERE clientId = $1`,
|
|||
|
|
`DELETE FROM client_metadata_schemas WHERE client_id = $1`,
|
|||
|
|
`DELETE FROM bot_session_documents WHERE document_id IN (SELECT id FROM documents WHERE clientId = $1)`,
|
|||
|
|
`DELETE FROM bot_session_folders WHERE folder_id IN (SELECT id FROM folders WHERE clientId = $1)`,
|
|||
|
|
`DELETE FROM bot_sessions WHERE client_id = $1`,
|
|||
|
|
`DELETE FROM documents WHERE clientId = $1`,
|
|||
|
|
`DELETE FROM folders WHERE clientId = $1`,
|
|||
|
|
`DELETE FROM clients WHERE clientId = $1`,
|
|||
|
|
}
|
|||
|
|
for _, stmt := range stmts {
|
|||
|
|
_, err := pool.Exec(ctx, stmt, clientID)
|
|||
|
|
require.NoErrorf(t, err, "reset stmt failed: %s", stmt)
|
|||
|
|
}
|
|||
|
|
err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
|
|||
|
|
Name: name,
|
|||
|
|
Clientid: clientID,
|
|||
|
|
})
|
|||
|
|
require.NoError(t, err)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// seedPayloadSession inserts a bot_sessions row owned by (clientID, actor).
|
|||
|
|
func seedPayloadSession(t *testing.T, ctx context.Context, cfg *serviceconfig.BaseConfig, clientID, actor string) uuid.UUID {
|
|||
|
|
t.Helper()
|
|||
|
|
var id uuid.UUID
|
|||
|
|
err := cfg.GetDBPool().QueryRow(ctx, `
|
|||
|
|
INSERT INTO bot_sessions (client_id, created_by) VALUES ($1, $2) RETURNING id
|
|||
|
|
`, clientID, actor).Scan(&id)
|
|||
|
|
require.NoError(t, err)
|
|||
|
|
return id
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// seedCompletedTurn inserts one completed bot_turns row at the given
|
|||
|
|
// ordinal, with prompt and completion populated.
|
|||
|
|
func seedCompletedTurn(t *testing.T, ctx context.Context, cfg *serviceconfig.BaseConfig, sessionID uuid.UUID, ordinal int, prompt, completion string, createdAt time.Time) {
|
|||
|
|
t.Helper()
|
|||
|
|
_, err := cfg.GetDBPool().Exec(ctx, `
|
|||
|
|
INSERT INTO bot_turns (session_id, ordinal, prompt, status, attempt_id, completion, attempt_started_at, created_at)
|
|||
|
|
VALUES ($1, $2, $3, 'completed', $4, $5, $6, $6)
|
|||
|
|
`, sessionID, ordinal, prompt, uuid.New(), completion, createdAt)
|
|||
|
|
require.NoError(t, err)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// ---------- conversation history ----------
|
|||
|
|
|
|||
|
|
// TestConversationHistoryConversion: given the last RecentTurns
|
|||
|
|
// completed turns, BuildConversationHistory produces N×2 messages
|
|||
|
|
// (user prompt + assistant completion) in ascending ordinal order.
|
|||
|
|
// Plan §7 mapping table.
|
|||
|
|
func TestConversationHistoryConversion(t *testing.T) {
|
|||
|
|
if testing.Short() {
|
|||
|
|
t.SkipNow()
|
|||
|
|
}
|
|||
|
|
ctx := t.Context()
|
|||
|
|
|
|||
|
|
dbCfg := &serviceconfig.BaseConfig{}
|
|||
|
|
test.CreateDB(t, dbCfg)
|
|||
|
|
|
|||
|
|
const (
|
|||
|
|
clientID = "BOT_M3_CONV"
|
|||
|
|
actor = "history@example.com"
|
|||
|
|
)
|
|||
|
|
resetClientForPayload(t, ctx, dbCfg, clientID, "bot-m3-conv")
|
|||
|
|
sessionID := seedPayloadSession(t, ctx, dbCfg, clientID, actor)
|
|||
|
|
|
|||
|
|
// Seed 7 completed turns. With cfg.RecentTurns=5 the helper must
|
|||
|
|
// pick the last 5 ordered ascending: ordinals 3..7. createdAt is
|
|||
|
|
// staggered by one second so any future ordering by created_at
|
|||
|
|
// (rather than ordinal) is also deterministic.
|
|||
|
|
now := time.Now().UTC().Truncate(time.Second)
|
|||
|
|
for ord := 1; ord <= 7; ord++ {
|
|||
|
|
seedCompletedTurn(t, ctx, dbCfg, sessionID, ord,
|
|||
|
|
"prompt-"+string(rune('0'+ord)),
|
|||
|
|
"answer-"+string(rune('0'+ord)),
|
|||
|
|
now.Add(time.Duration(ord)*time.Second))
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
svc, err := bot.New(payloadDefaultCfg(), dbCfg)
|
|||
|
|
require.NoError(t, err)
|
|||
|
|
|
|||
|
|
history, err := svc.BuildConversationHistory(ctx, sessionID)
|
|||
|
|
require.NoError(t, err)
|
|||
|
|
require.Len(t, history, 5*2,
|
|||
|
|
"plan §7: RecentTurns=5 -> 5 turns x 2 messages = 10")
|
|||
|
|
|
|||
|
|
// Expect ordinals 3,4,5,6,7 in ascending order. Each turn produces
|
|||
|
|
// (user prompt, assistant completion) in that order.
|
|||
|
|
for i := 0; i < 5; i++ {
|
|||
|
|
ord := i + 3
|
|||
|
|
userMsg := history[i*2]
|
|||
|
|
assistantMsg := history[i*2+1]
|
|||
|
|
|
|||
|
|
assert.Equal(t, "user", userMsg.Role,
|
|||
|
|
"index %d must be user; ordinal=%d", i*2, ord)
|
|||
|
|
assert.Equal(t, "prompt-"+string(rune('0'+ord)), userMsg.Content,
|
|||
|
|
"user content for ordinal %d", ord)
|
|||
|
|
assert.Equal(t, "assistant", assistantMsg.Role,
|
|||
|
|
"index %d must be assistant; ordinal=%d", i*2+1, ord)
|
|||
|
|
assert.Equal(t, "answer-"+string(rune('0'+ord)), assistantMsg.Content,
|
|||
|
|
"assistant content for ordinal %d", ord)
|
|||
|
|
assert.Falsef(t, userMsg.Timestamp.IsZero(),
|
|||
|
|
"user message ordinal %d must carry a non-zero timestamp", ord)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// TestConversationHistoryConversion_EmptySession: a session with zero
|
|||
|
|
// completed turns yields a zero-length slice (NOT nil), so the FastAPI
|
|||
|
|
// client emits the JSON `"conversation_history": []` per the swagger
|
|||
|
|
// requirement that the field be a non-null array.
|
|||
|
|
func TestConversationHistoryConversion_EmptySession(t *testing.T) {
|
|||
|
|
if testing.Short() {
|
|||
|
|
t.SkipNow()
|
|||
|
|
}
|
|||
|
|
ctx := t.Context()
|
|||
|
|
|
|||
|
|
dbCfg := &serviceconfig.BaseConfig{}
|
|||
|
|
test.CreateDB(t, dbCfg)
|
|||
|
|
|
|||
|
|
const (
|
|||
|
|
clientID = "BOT_M3_CONV_E"
|
|||
|
|
actor = "empty@example.com"
|
|||
|
|
)
|
|||
|
|
resetClientForPayload(t, ctx, dbCfg, clientID, "bot-m3-conv-e")
|
|||
|
|
sessionID := seedPayloadSession(t, ctx, dbCfg, clientID, actor)
|
|||
|
|
|
|||
|
|
svc, err := bot.New(payloadDefaultCfg(), dbCfg)
|
|||
|
|
require.NoError(t, err)
|
|||
|
|
|
|||
|
|
history, err := svc.BuildConversationHistory(ctx, sessionID)
|
|||
|
|
require.NoError(t, err)
|
|||
|
|
require.NotNil(t, history,
|
|||
|
|
"empty result must be a non-nil slice so json.Marshal emits []")
|
|||
|
|
assert.Empty(t, history)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// ---------- scope payload assembly ----------
|
|||
|
|
|
|||
|
|
// TestScopePayloadAssembly_FoldersExpanded: stored scope has 2 doc ids
|
|||
|
|
// + 1 folder id; folder contains 3 unrelated docs. Resulting payload
|
|||
|
|
// has documents = the union of 5 ids; folders = the original 1 folder.
|
|||
|
|
// Plan §7 mapping table: "scope": "Stored document ids union expanded
|
|||
|
|
// folder document ids; stored folder ids preserved".
|
|||
|
|
func TestScopePayloadAssembly_FoldersExpanded(t *testing.T) {
|
|||
|
|
if testing.Short() {
|
|||
|
|
t.SkipNow()
|
|||
|
|
}
|
|||
|
|
ctx := t.Context()
|
|||
|
|
|
|||
|
|
dbCfg := &serviceconfig.BaseConfig{}
|
|||
|
|
test.CreateDB(t, dbCfg)
|
|||
|
|
pool := dbCfg.GetDBPool()
|
|||
|
|
|
|||
|
|
const (
|
|||
|
|
clientID = "BOT_M3_SCOPE_FX"
|
|||
|
|
actor = "scope@example.com"
|
|||
|
|
)
|
|||
|
|
resetClientForPayload(t, ctx, dbCfg, clientID, "bot-m3-scope-fx")
|
|||
|
|
sessionID := seedPayloadSession(t, ctx, dbCfg, clientID, actor)
|
|||
|
|
|
|||
|
|
queries := dbCfg.GetDBQueries()
|
|||
|
|
|
|||
|
|
// 2 stored documents (no folder).
|
|||
|
|
docA, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{Clientid: clientID, Hash: "scope-fx-a"})
|
|||
|
|
require.NoError(t, err)
|
|||
|
|
docB, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{Clientid: clientID, Hash: "scope-fx-b"})
|
|||
|
|
require.NoError(t, err)
|
|||
|
|
|
|||
|
|
// 1 folder containing 3 documents (different from the stored docs).
|
|||
|
|
folder, err := queries.CreateFolder(ctx, &repository.CreateFolderParams{
|
|||
|
|
Path: "/scope-fx-folder", Parentid: nil, Clientid: clientID, Createdby: "tester",
|
|||
|
|
})
|
|||
|
|
require.NoError(t, err)
|
|||
|
|
folderDocs := make([]uuid.UUID, 3)
|
|||
|
|
for i := 0; i < 3; i++ {
|
|||
|
|
var id uuid.UUID
|
|||
|
|
err := pool.QueryRow(ctx, `
|
|||
|
|
INSERT INTO documents (clientId, hash, folderId) VALUES ($1, $2, $3) RETURNING id
|
|||
|
|
`, clientID, "scope-fx-folderdoc-"+string(rune('0'+i)), folder.ID).Scan(&id)
|
|||
|
|
require.NoError(t, err)
|
|||
|
|
folderDocs[i] = id
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// Persist the scope: 2 stored docs + 1 folder.
|
|||
|
|
for _, d := range []uuid.UUID{docA, docB} {
|
|||
|
|
_, err := pool.Exec(ctx,
|
|||
|
|
`INSERT INTO bot_session_documents (session_id, document_id) VALUES ($1, $2)`,
|
|||
|
|
sessionID, d)
|
|||
|
|
require.NoError(t, err)
|
|||
|
|
}
|
|||
|
|
_, err = pool.Exec(ctx,
|
|||
|
|
`INSERT INTO bot_session_folders (session_id, folder_id) VALUES ($1, $2)`,
|
|||
|
|
sessionID, folder.ID)
|
|||
|
|
require.NoError(t, err)
|
|||
|
|
|
|||
|
|
svc, err := bot.New(payloadDefaultCfg(), dbCfg)
|
|||
|
|
require.NoError(t, err)
|
|||
|
|
|
|||
|
|
scope, err := svc.BuildChatScope(ctx, sessionID, clientID)
|
|||
|
|
require.NoError(t, err)
|
|||
|
|
|
|||
|
|
wantDocs := []string{
|
|||
|
|
docA.String(), docB.String(),
|
|||
|
|
folderDocs[0].String(), folderDocs[1].String(), folderDocs[2].String(),
|
|||
|
|
}
|
|||
|
|
assert.ElementsMatch(t, wantDocs, scope.DocumentIDs,
|
|||
|
|
"plan §7: documents = stored docs ∪ folder-expanded docs")
|
|||
|
|
assert.ElementsMatch(t, []string{folder.ID.String()}, scope.FolderIDs,
|
|||
|
|
"plan §7: stored folder ids preserved")
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// TestScopePayloadAssembly_DocsAndFolderOverlap: stored doc id is also
|
|||
|
|
// inside the folder being expanded. Payload documents must dedup so the
|
|||
|
|
// shared id appears exactly once. Folder id preserved.
|
|||
|
|
func TestScopePayloadAssembly_DocsAndFolderOverlap(t *testing.T) {
|
|||
|
|
if testing.Short() {
|
|||
|
|
t.SkipNow()
|
|||
|
|
}
|
|||
|
|
ctx := t.Context()
|
|||
|
|
|
|||
|
|
dbCfg := &serviceconfig.BaseConfig{}
|
|||
|
|
test.CreateDB(t, dbCfg)
|
|||
|
|
pool := dbCfg.GetDBPool()
|
|||
|
|
|
|||
|
|
const (
|
|||
|
|
clientID = "BOT_M3_SCOPE_OL"
|
|||
|
|
actor = "overlap@example.com"
|
|||
|
|
)
|
|||
|
|
resetClientForPayload(t, ctx, dbCfg, clientID, "bot-m3-scope-ol")
|
|||
|
|
sessionID := seedPayloadSession(t, ctx, dbCfg, clientID, actor)
|
|||
|
|
|
|||
|
|
queries := dbCfg.GetDBQueries()
|
|||
|
|
|
|||
|
|
folder, err := queries.CreateFolder(ctx, &repository.CreateFolderParams{
|
|||
|
|
Path: "/scope-ol-folder", Parentid: nil, Clientid: clientID, Createdby: "tester",
|
|||
|
|
})
|
|||
|
|
require.NoError(t, err)
|
|||
|
|
|
|||
|
|
// docA is the overlap: stored AND inside the folder.
|
|||
|
|
var docA uuid.UUID
|
|||
|
|
err = pool.QueryRow(ctx, `
|
|||
|
|
INSERT INTO documents (clientId, hash, folderId) VALUES ($1, $2, $3) RETURNING id
|
|||
|
|
`, clientID, "scope-ol-overlap", folder.ID).Scan(&docA)
|
|||
|
|
require.NoError(t, err)
|
|||
|
|
|
|||
|
|
// docB is folder-only.
|
|||
|
|
var docB uuid.UUID
|
|||
|
|
err = pool.QueryRow(ctx, `
|
|||
|
|
INSERT INTO documents (clientId, hash, folderId) VALUES ($1, $2, $3) RETURNING id
|
|||
|
|
`, clientID, "scope-ol-folder-only", folder.ID).Scan(&docB)
|
|||
|
|
require.NoError(t, err)
|
|||
|
|
|
|||
|
|
// Stored scope: docA (also in folder) + folder.
|
|||
|
|
_, err = pool.Exec(ctx,
|
|||
|
|
`INSERT INTO bot_session_documents (session_id, document_id) VALUES ($1, $2)`,
|
|||
|
|
sessionID, docA)
|
|||
|
|
require.NoError(t, err)
|
|||
|
|
_, err = pool.Exec(ctx,
|
|||
|
|
`INSERT INTO bot_session_folders (session_id, folder_id) VALUES ($1, $2)`,
|
|||
|
|
sessionID, folder.ID)
|
|||
|
|
require.NoError(t, err)
|
|||
|
|
|
|||
|
|
svc, err := bot.New(payloadDefaultCfg(), dbCfg)
|
|||
|
|
require.NoError(t, err)
|
|||
|
|
|
|||
|
|
scope, err := svc.BuildChatScope(ctx, sessionID, clientID)
|
|||
|
|
require.NoError(t, err)
|
|||
|
|
|
|||
|
|
wantDocs := []string{docA.String(), docB.String()}
|
|||
|
|
assert.ElementsMatch(t, wantDocs, scope.DocumentIDs,
|
|||
|
|
"overlap document must appear exactly once after dedup")
|
|||
|
|
assert.ElementsMatch(t, []string{folder.ID.String()}, scope.FolderIDs)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// TestScopePayloadAssembly_EmptyScopeReturnsAllClientDocs: with no
|
|||
|
|
// stored doc/folder rows, the payload documents fall back to "all
|
|||
|
|
// documents for the client" per plan §4 ("Empty scope tables mean all
|
|||
|
|
// documents for the client"). Folders array is empty.
|
|||
|
|
func TestScopePayloadAssembly_EmptyScopeReturnsAllClientDocs(t *testing.T) {
|
|||
|
|
if testing.Short() {
|
|||
|
|
t.SkipNow()
|
|||
|
|
}
|
|||
|
|
ctx := t.Context()
|
|||
|
|
|
|||
|
|
dbCfg := &serviceconfig.BaseConfig{}
|
|||
|
|
test.CreateDB(t, dbCfg)
|
|||
|
|
|
|||
|
|
const (
|
|||
|
|
clientID = "BOT_M3_SCOPE_E"
|
|||
|
|
actor = "empty-scope@example.com"
|
|||
|
|
)
|
|||
|
|
resetClientForPayload(t, ctx, dbCfg, clientID, "bot-m3-scope-e")
|
|||
|
|
sessionID := seedPayloadSession(t, ctx, dbCfg, clientID, actor)
|
|||
|
|
|
|||
|
|
queries := dbCfg.GetDBQueries()
|
|||
|
|
wantDocs := []string{}
|
|||
|
|
for i := 0; i < 4; i++ {
|
|||
|
|
id, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{
|
|||
|
|
Clientid: clientID, Hash: "scope-e-" + string(rune('0'+i)),
|
|||
|
|
})
|
|||
|
|
require.NoError(t, err)
|
|||
|
|
wantDocs = append(wantDocs, id.String())
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
svc, err := bot.New(payloadDefaultCfg(), dbCfg)
|
|||
|
|
require.NoError(t, err)
|
|||
|
|
|
|||
|
|
scope, err := svc.BuildChatScope(ctx, sessionID, clientID)
|
|||
|
|
require.NoError(t, err)
|
|||
|
|
assert.ElementsMatch(t, wantDocs, scope.DocumentIDs,
|
|||
|
|
"empty scope tables -> all documents for the client")
|
|||
|
|
assert.Empty(t, scope.FolderIDs)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// ---------- cached_results stripping ----------
|
|||
|
|
|
|||
|
|
// TestStripCachedResults_RemovesKey: input state with a cached_results
|
|||
|
|
// entry returns the same object minus that one key. Other keys are
|
|||
|
|
// preserved verbatim.
|
|||
|
|
func TestStripCachedResults_RemovesKey(t *testing.T) {
|
|||
|
|
t.Parallel()
|
|||
|
|
|
|||
|
|
in := json.RawMessage(`{"foo":1,"cached_results":{"abc":{"k":"v"}},"bar":"hello"}`)
|
|||
|
|
out, err := bot.StripCachedResults(in)
|
|||
|
|
require.NoError(t, err)
|
|||
|
|
|
|||
|
|
var parsed map[string]interface{}
|
|||
|
|
require.NoError(t, json.Unmarshal(out, &parsed))
|
|||
|
|
_, hasCached := parsed["cached_results"]
|
|||
|
|
assert.False(t, hasCached, "cached_results must be removed")
|
|||
|
|
assert.EqualValues(t, 1, parsed["foo"], "other keys must be preserved")
|
|||
|
|
assert.Equal(t, "hello", parsed["bar"])
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// TestStripCachedResults_NoOpWhenAbsent: input without cached_results is
|
|||
|
|
// returned with the same logical shape.
|
|||
|
|
func TestStripCachedResults_NoOpWhenAbsent(t *testing.T) {
|
|||
|
|
t.Parallel()
|
|||
|
|
|
|||
|
|
in := json.RawMessage(`{"foo":1,"bar":"hello"}`)
|
|||
|
|
out, err := bot.StripCachedResults(in)
|
|||
|
|
require.NoError(t, err)
|
|||
|
|
|
|||
|
|
var got map[string]interface{}
|
|||
|
|
require.NoError(t, json.Unmarshal(out, &got))
|
|||
|
|
|
|||
|
|
var want map[string]interface{}
|
|||
|
|
require.NoError(t, json.Unmarshal(in, &want))
|
|||
|
|
|
|||
|
|
assert.Equal(t, want, got)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// TestStripCachedResults_NullInput: a nil/empty input passes through
|
|||
|
|
// unchanged so the caller can hand it null state without a special case.
|
|||
|
|
func TestStripCachedResults_NullInput(t *testing.T) {
|
|||
|
|
t.Parallel()
|
|||
|
|
|
|||
|
|
out, err := bot.StripCachedResults(nil)
|
|||
|
|
require.NoError(t, err)
|
|||
|
|
assert.Nil(t, out, "nil input must round-trip as nil")
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// ---------- 64KB session-state cap ----------
|
|||
|
|
|
|||
|
|
// TestSessionStateSizeCap_64KB_Boundary: the cap is plan §8's
|
|||
|
|
// chatbot.StateMaxBytes. A stripped state at exactly the cap reports
|
|||
|
|
// as not-oversize; one byte over reports oversize. The helper returns
|
|||
|
|
// (oversize bool) without mutating the input.
|
|||
|
|
func TestSessionStateSizeCap_64KB_Boundary(t *testing.T) {
|
|||
|
|
t.Parallel()
|
|||
|
|
|
|||
|
|
// Build a JSON object whose length is exactly chatbot.StateMaxBytes.
|
|||
|
|
// We use a single-key string-valued object and pad the value so the
|
|||
|
|
// total raw byte length is exact.
|
|||
|
|
const head = `{"k":"`
|
|||
|
|
const tail = `"}`
|
|||
|
|
overhead := len(head) + len(tail)
|
|||
|
|
require.Greater(t, chatbot.StateMaxBytes, overhead)
|
|||
|
|
value := strings.Repeat("a", chatbot.StateMaxBytes-overhead)
|
|||
|
|
stateAtCap := json.RawMessage(head + value + tail)
|
|||
|
|
require.Equal(t, chatbot.StateMaxBytes, len(stateAtCap),
|
|||
|
|
"test setup: stateAtCap must equal the cap byte-for-byte")
|
|||
|
|
|
|||
|
|
assert.False(t, bot.SessionStateOversize(stateAtCap),
|
|||
|
|
"plan §8: state at exactly StateMaxBytes is not oversize")
|
|||
|
|
|
|||
|
|
// One byte over the cap.
|
|||
|
|
overValue := strings.Repeat("a", chatbot.StateMaxBytes-overhead+1)
|
|||
|
|
stateOver := json.RawMessage(head + overValue + tail)
|
|||
|
|
require.Equal(t, chatbot.StateMaxBytes+1, len(stateOver))
|
|||
|
|
|
|||
|
|
assert.True(t, bot.SessionStateOversize(stateOver),
|
|||
|
|
"plan §8: state above StateMaxBytes is oversize")
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// ---------- metadata summarization ----------
|
|||
|
|
|
|||
|
|
// TestMetadataSummarization: response metadata.agents = {a: {input:10,
|
|||
|
|
// output:20}, b: {input:5, output:30}} -> tokens_in=15, tokens_out=50,
|
|||
|
|
// latency = metadata.total_latency_ms (NOT wall clock when metadata is
|
|||
|
|
// present).
|
|||
|
|
func TestMetadataSummarization(t *testing.T) {
|
|||
|
|
t.Parallel()
|
|||
|
|
|
|||
|
|
sentAt := time.Date(2026, 5, 4, 12, 0, 0, 0, time.UTC)
|
|||
|
|
receivedAt := sentAt.Add(7 * time.Second)
|
|||
|
|
|
|||
|
|
metadata := &bot.Metadata{
|
|||
|
|
TotalLatencyMs: 6000,
|
|||
|
|
RetryCount: 0,
|
|||
|
|
Agents: map[string]bot.AgentTelemetry{
|
|||
|
|
"a": {LatencyMs: 1000, InputTokens: 10, OutputTokens: 20},
|
|||
|
|
"b": {LatencyMs: 5000, InputTokens: 5, OutputTokens: 30},
|
|||
|
|
},
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
latencyMs, tokensIn, tokensOut := bot.SummarizeMetadata(metadata, sentAt, receivedAt)
|
|||
|
|
assert.Equal(t, 6000, latencyMs,
|
|||
|
|
"plan §7: latency = metadata.total_latency_ms when metadata is present")
|
|||
|
|
assert.Equal(t, 15, tokensIn,
|
|||
|
|
"plan §7: tokens_in = sum of agents[].input_tokens")
|
|||
|
|
assert.Equal(t, 50, tokensOut,
|
|||
|
|
"plan §7: tokens_out = sum of agents[].output_tokens")
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// TestMetadataSummarization_NilMetadataUsesWallClock: nil metadata ->
|
|||
|
|
// latency = ReceivedAt - SentAt in ms, tokens default to 0.
|
|||
|
|
func TestMetadataSummarization_NilMetadataUsesWallClock(t *testing.T) {
|
|||
|
|
t.Parallel()
|
|||
|
|
|
|||
|
|
sentAt := time.Date(2026, 5, 4, 12, 0, 0, 0, time.UTC)
|
|||
|
|
receivedAt := sentAt.Add(2500 * time.Millisecond)
|
|||
|
|
|
|||
|
|
latencyMs, tokensIn, tokensOut := bot.SummarizeMetadata(nil, sentAt, receivedAt)
|
|||
|
|
assert.Equal(t, 2500, latencyMs,
|
|||
|
|
"nil metadata -> wall-clock latency in milliseconds")
|
|||
|
|
assert.Equal(t, 0, tokensIn)
|
|||
|
|
assert.Equal(t, 0, tokensOut)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// TestMetadataSummarization_EmptyAgentsMap: metadata present but agents
|
|||
|
|
// map is empty -> tokens default to 0; latency uses metadata.total_latency_ms.
|
|||
|
|
func TestMetadataSummarization_EmptyAgentsMap(t *testing.T) {
|
|||
|
|
t.Parallel()
|
|||
|
|
|
|||
|
|
sentAt := time.Date(2026, 5, 4, 12, 0, 0, 0, time.UTC)
|
|||
|
|
receivedAt := sentAt.Add(5 * time.Second)
|
|||
|
|
|
|||
|
|
metadata := &bot.Metadata{
|
|||
|
|
TotalLatencyMs: 4500,
|
|||
|
|
RetryCount: 2,
|
|||
|
|
Agents: map[string]bot.AgentTelemetry{},
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
latencyMs, tokensIn, tokensOut := bot.SummarizeMetadata(metadata, sentAt, receivedAt)
|
|||
|
|
assert.Equal(t, 4500, latencyMs)
|
|||
|
|
assert.Equal(t, 0, tokensIn)
|
|||
|
|
assert.Equal(t, 0, tokensOut)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// TestStripCachedResults_OnRealAariaSample exercises StripCachedResults
|
|||
|
|
// against the authoritative Aaria session_state shape Q committed to
|
|||
|
|
// plans/chatbot.stuff/sample.response.json. The sample carries the full
|
|||
|
|
// 10-key session-state envelope including a populated cached_results
|
|||
|
|
// map keyed by result hash; stripping must leave the other 9 keys
|
|||
|
|
// intact and the resulting bytes must fit comfortably under
|
|||
|
|
// chatbot.StateMaxBytes.
|
|||
|
|
//
|
|||
|
|
// The other-key preservation check uses json.Unmarshal->reflect.DeepEqual
|
|||
|
|
// rather than byte equality because Go's encoding/json does not
|
|||
|
|
// guarantee map-key emission order: a sound implementation may re-emit
|
|||
|
|
// keys in a different order than the source file, but the parsed value
|
|||
|
|
// must be identical. Plan §7 mapping table, plan §8 64KB cap.
|
|||
|
|
func TestStripCachedResults_OnRealAariaSample(t *testing.T) {
|
|||
|
|
t.Parallel()
|
|||
|
|
|
|||
|
|
const samplePath = "../../plans/chatbot.stuff/sample.response.json"
|
|||
|
|
raw, err := os.ReadFile(samplePath)
|
|||
|
|
require.NoErrorf(t, err, "must read fixture at %s", samplePath)
|
|||
|
|
|
|||
|
|
// Sanity-precondition: the fixture has the 10 documented top-level
|
|||
|
|
// keys including cached_results. If Q reshapes the fixture, this
|
|||
|
|
// assertion fails first with a clear pointer at the source of drift.
|
|||
|
|
var input map[string]interface{}
|
|||
|
|
require.NoError(t, json.Unmarshal(raw, &input),
|
|||
|
|
"fixture must be a JSON object")
|
|||
|
|
require.Len(t, input, 10,
|
|||
|
|
"fixture must have 10 top-level keys before stripping; got %d", len(input))
|
|||
|
|
_, hadCached := input["cached_results"]
|
|||
|
|
require.True(t, hadCached, "fixture must include cached_results")
|
|||
|
|
|
|||
|
|
stripped, err := bot.StripCachedResults(json.RawMessage(raw))
|
|||
|
|
require.NoError(t, err)
|
|||
|
|
require.NotNil(t, stripped, "non-empty input must produce non-nil output")
|
|||
|
|
|
|||
|
|
var out map[string]interface{}
|
|||
|
|
require.NoError(t, json.Unmarshal(stripped, &out),
|
|||
|
|
"stripped output must be a valid JSON object")
|
|||
|
|
|
|||
|
|
require.Len(t, out, 9,
|
|||
|
|
"stripping cached_results must leave exactly 9 keys; got %d", len(out))
|
|||
|
|
_, hasCached := out["cached_results"]
|
|||
|
|
assert.False(t, hasCached, "cached_results must be absent after strip")
|
|||
|
|
|
|||
|
|
// Every other key must round-trip with the original value. Plan §7
|
|||
|
|
// is explicit that ONLY cached_results is touched.
|
|||
|
|
expectedKeys := []string{
|
|||
|
|
"turn",
|
|||
|
|
"resolved_entities",
|
|||
|
|
"last_scope",
|
|||
|
|
"last_intent",
|
|||
|
|
"last_query_sql",
|
|||
|
|
"last_query_plan",
|
|||
|
|
"last_result_hash",
|
|||
|
|
"last_scope_signature",
|
|||
|
|
"validator_retry_count",
|
|||
|
|
}
|
|||
|
|
for _, k := range expectedKeys {
|
|||
|
|
require.Containsf(t, out, k,
|
|||
|
|
"key %q must be preserved after strip", k)
|
|||
|
|
assert.Equalf(t, input[k], out[k],
|
|||
|
|
"key %q value must round-trip unchanged", k)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
assert.Falsef(t, bot.SessionStateOversize(stripped),
|
|||
|
|
"plan §8: real Aaria sample is well under 64 KiB after strip; got %d bytes", len(stripped))
|
|||
|
|
}
|