266 lines
10 KiB
Go
266 lines
10 KiB
Go
package handlers
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"net/http"
|
||
"strings"
|
||
"testing"
|
||
|
||
"chat-switchboard/compaction"
|
||
"chat-switchboard/database"
|
||
"chat-switchboard/models"
|
||
"chat-switchboard/roles"
|
||
"chat-switchboard/store"
|
||
postgres "chat-switchboard/store/postgres"
|
||
sqlite "chat-switchboard/store/sqlite"
|
||
"chat-switchboard/treepath"
|
||
)
|
||
|
||
// ═══════════════════════════════════════════
|
||
// Live Compaction Tests
|
||
// ═══════════════════════════════════════════
|
||
// Requires: TEST_DATABASE_URL + PROVIDER_KEY (or VENICE_API_KEY)
|
||
//
|
||
// Tests the full Compact() pipeline end-to-end:
|
||
// seed messages → call utility model → verify summary tree node
|
||
// ═══════════════════════════════════════════
|
||
|
||
// dialectStores returns the correct store bundle for the active dialect.
|
||
func dialectStores() store.Stores {
|
||
if database.IsSQLite() {
|
||
return sqlite.NewStores(database.TestDB)
|
||
}
|
||
return postgres.NewStores(database.TestDB)
|
||
}
|
||
|
||
func TestLive_CompactionFullPipeline(t *testing.T) {
|
||
h := setupHarness(t)
|
||
pc := requireLiveProvider(t)
|
||
userID, adminToken := h.createAdminUser("compactadmin", "compact@test.com")
|
||
|
||
// Set up provider + enable first model
|
||
configID, modelID := setupProviderWithModel(t, h, adminToken, pc)
|
||
|
||
// Configure utility role
|
||
w := h.request("PUT", "/api/v1/admin/roles/utility", adminToken, map[string]interface{}{
|
||
"primary": map[string]interface{}{
|
||
"provider_config_id": configID,
|
||
"model_id": modelID,
|
||
},
|
||
})
|
||
if w.Code != http.StatusOK {
|
||
t.Fatalf("configure utility role: %d: %s", w.Code, w.Body.String())
|
||
}
|
||
t.Log(" ✓ Utility role configured with", modelID)
|
||
|
||
// Create a channel with enough messages to summarize
|
||
w = h.request("POST", "/api/v1/channels", adminToken, map[string]interface{}{
|
||
"title": "Compaction Test", "type": "direct",
|
||
})
|
||
if w.Code != http.StatusCreated {
|
||
t.Fatalf("create channel: %d: %s", w.Code, w.Body.String())
|
||
}
|
||
var ch map[string]interface{}
|
||
decode(w, &ch)
|
||
channelID := ch["id"].(string)
|
||
|
||
// Seed messages manually (not via completion — faster and cheaper)
|
||
msgs := []struct{ role, content string }{
|
||
{"user", "I'm planning a trip to Japan next month. What cities should I visit?"},
|
||
{"assistant", "For a Japan trip, I'd recommend Tokyo for modern culture, Kyoto for temples and traditional culture, Osaka for food, and Hiroshima for history. How long is your trip?"},
|
||
{"user", "About 10 days. I'm most interested in food and temples."},
|
||
{"assistant", "With 10 days focusing on food and temples, I'd suggest: 3 days in Kyoto for Fushimi Inari, Kinkaku-ji, and Arashiyama. 2 days in Osaka for Dotonbori street food, Kuromon Market, and Namba. 2 days in Tokyo for Tsukiji Outer Market, Senso-ji, and Meiji Shrine. Then day trips to Nara (deer park and Todai-ji) and Kamakura (Great Buddha)."},
|
||
{"user", "That sounds great! What about budget? I'm trying to keep it under $3000 for the whole trip."},
|
||
{"assistant", "A $3000 budget for 10 days in Japan is doable. Rough breakdown: Flights ~$800-1000, Accommodation ~$50-80/night in hostels or budget hotels ($500-800), JR Pass 14-day ~$380, Food ~$30-50/day ($300-500), Activities and temples ~$200. Total: $2180-2880. Tips: eat at konbini (convenience stores) for cheap meals, visit free shrines, and book accommodation early."},
|
||
{"user", "Should I get a JR Pass or just buy individual tickets?"},
|
||
{"assistant", "For your itinerary covering Tokyo, Kyoto, Osaka, Nara, and Kamakura, a 7-day JR Pass ($200) would cover most of your intercity travel. The Tokyo-Kyoto shinkansen alone costs $120 one-way. I'd recommend the 7-day pass starting when you leave Tokyo for Kyoto, then use IC cards (Suica/Pasmo) for local transit."},
|
||
{"user", "Perfect. One more question — what's the best time to visit temples to avoid crowds?"},
|
||
{"assistant", "Early morning is best for temples. Fushimi Inari at sunrise (5-6am) is magical and nearly empty. Kinkaku-ji opens at 9am — arrive right at opening. Arashiyama bamboo grove is best before 8am. For Senso-ji in Tokyo, go at dawn for beautiful photos. Weekdays are always less crowded than weekends."},
|
||
}
|
||
|
||
ph := database.PH
|
||
var lastMsgID string
|
||
for _, m := range msgs {
|
||
var parentPtr *string
|
||
if lastMsgID != "" {
|
||
parentPtr = &lastMsgID
|
||
}
|
||
err := database.TestDB.QueryRow(
|
||
"INSERT INTO messages (channel_id, parent_id, role, content, sibling_index) VALUES ("+ph(1)+", "+ph(2)+", "+ph(3)+", "+ph(4)+", 0) RETURNING id",
|
||
channelID, parentPtr, m.role, m.content).Scan(&lastMsgID)
|
||
if err != nil {
|
||
t.Fatalf("seed message: %v", err)
|
||
}
|
||
}
|
||
|
||
// Set cursor to last message
|
||
database.TestDB.Exec(
|
||
"INSERT INTO channel_cursors (channel_id, user_id, active_leaf_id) VALUES ("+ph(1)+", "+ph(2)+", "+ph(3)+") ON CONFLICT (channel_id, user_id) DO UPDATE SET active_leaf_id = excluded.active_leaf_id",
|
||
channelID, userID, lastMsgID)
|
||
|
||
t.Logf(" ✓ Seeded %d messages in channel %s", len(msgs), channelID)
|
||
|
||
// ── Run compaction ──
|
||
stores := dialectStores()
|
||
resolver := roles.NewResolver(stores, nil)
|
||
svc := compaction.NewService(stores, resolver)
|
||
|
||
result, err := svc.Compact(context.Background(), compaction.CompactRequest{
|
||
ChannelID: channelID,
|
||
UserID: userID,
|
||
TeamID: nil,
|
||
Trigger: "manual",
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("Compact() failed: %v", err)
|
||
}
|
||
|
||
t.Logf(" ✓ Compact result: %d messages summarized, model=%s, %d chars",
|
||
result.SummarizedCount, result.Model, len(result.Content))
|
||
|
||
// Verify summary was created
|
||
if result.SummarizedCount != len(msgs) {
|
||
t.Errorf("summarized_count = %d, want %d", result.SummarizedCount, len(msgs))
|
||
}
|
||
if result.Content == "" {
|
||
t.Fatal("summary content should not be empty")
|
||
}
|
||
if result.Model != modelID {
|
||
t.Errorf("model = %q, want %q", result.Model, modelID)
|
||
}
|
||
|
||
// Verify summary message exists in the tree
|
||
var summaryContent, summaryRole string
|
||
var metadataRaw []byte
|
||
err = database.TestDB.QueryRow(
|
||
"SELECT role, content, metadata FROM messages WHERE id = "+ph(1),
|
||
result.SummaryID).Scan(&summaryRole, &summaryContent, &metadataRaw)
|
||
if err != nil {
|
||
t.Fatalf("query summary message: %v", err)
|
||
}
|
||
if summaryRole != "assistant" {
|
||
t.Errorf("summary role = %q, want assistant", summaryRole)
|
||
}
|
||
|
||
var metadata models.JSONMap
|
||
json.Unmarshal(metadataRaw, &metadata)
|
||
if metadata["type"] != "summary" {
|
||
t.Errorf("metadata.type = %q, want summary", metadata["type"])
|
||
}
|
||
if metadata["trigger"] != "manual" {
|
||
t.Errorf("metadata.trigger = %q, want manual", metadata["trigger"])
|
||
}
|
||
t.Logf(" ✓ Summary message verified: id=%s, metadata=%v", result.SummaryID, metadata)
|
||
|
||
// Verify cursor was updated to point to summary
|
||
path, err := treepath.GetActivePath(channelID, userID)
|
||
if err != nil {
|
||
t.Fatalf("GetActivePath after compaction: %v", err)
|
||
}
|
||
|
||
if len(path) == 0 {
|
||
t.Fatal("path should not be empty after compaction")
|
||
}
|
||
lastPath := path[len(path)-1]
|
||
if lastPath.ID != result.SummaryID {
|
||
t.Errorf("active path leaf = %s, want summary %s", lastPath.ID, result.SummaryID)
|
||
}
|
||
t.Logf(" ✓ Cursor updated: path has %d messages, leaf is summary", len(path))
|
||
|
||
// Verify usage was logged
|
||
var usageCount int
|
||
database.TestDB.QueryRow(
|
||
"SELECT COUNT(*) FROM usage_log WHERE channel_id = "+ph(1)+" AND role = 'utility'",
|
||
channelID).Scan(&usageCount)
|
||
if usageCount == 0 {
|
||
t.Error("expected usage_log entry for utility role")
|
||
}
|
||
t.Logf(" ✓ Usage logged: %d entries", usageCount)
|
||
}
|
||
|
||
// TestLive_CompactionContextBudgetGuardRail verifies that Compact()
|
||
// rejects input that exceeds the utility model's context window.
|
||
func TestLive_CompactionContextBudgetGuardRail(t *testing.T) {
|
||
h := setupHarness(t)
|
||
pc := requireLiveProvider(t)
|
||
_, adminToken := h.createAdminUser("guardrail_admin", "guardrail@test.com")
|
||
userID := database.SeedTestUser(t, "guardrail_user", "gruser@test.com")
|
||
|
||
// Set up provider + enable first model
|
||
configID, modelID := setupProviderWithModel(t, h, adminToken, pc)
|
||
|
||
ph := database.PH
|
||
|
||
// Set max_context to 32000 in catalog (in case sync didn't populate it)
|
||
if database.IsSQLite() {
|
||
database.TestDB.Exec(
|
||
"UPDATE model_catalog SET capabilities = json_set(COALESCE(capabilities,'{}'), '$.max_context', 32000) WHERE provider_config_id = "+ph(1)+" AND model_id = "+ph(2),
|
||
configID, modelID)
|
||
} else {
|
||
database.TestDB.Exec(
|
||
"UPDATE model_catalog SET capabilities = capabilities || '{\"max_context\": 32000}'::jsonb WHERE provider_config_id = "+ph(1)+" AND model_id = "+ph(2),
|
||
configID, modelID)
|
||
}
|
||
|
||
// Configure utility role
|
||
h.request("PUT", "/api/v1/admin/roles/utility", adminToken, map[string]interface{}{
|
||
"primary": map[string]interface{}{
|
||
"provider_config_id": configID,
|
||
"model_id": modelID,
|
||
},
|
||
})
|
||
|
||
// Create channel with huge messages (>32K context worth)
|
||
w := h.request("POST", "/api/v1/channels", adminToken, map[string]interface{}{
|
||
"title": "Guard Rail Test", "type": "direct",
|
||
})
|
||
var ch map[string]interface{}
|
||
decode(w, &ch)
|
||
channelID := ch["id"].(string)
|
||
|
||
// Seed 20 messages × 8000 chars each = 160K chars ≈ 40K tokens
|
||
bigContent := make([]byte, 8000)
|
||
for i := range bigContent {
|
||
bigContent[i] = 'a'
|
||
}
|
||
var lastMsgID string
|
||
for i := 0; i < 20; i++ {
|
||
var parentPtr *string
|
||
if lastMsgID != "" {
|
||
parentPtr = &lastMsgID
|
||
}
|
||
role := "user"
|
||
if i%2 == 1 {
|
||
role = "assistant"
|
||
}
|
||
database.TestDB.QueryRow(
|
||
"INSERT INTO messages (channel_id, parent_id, role, content, sibling_index) VALUES ("+ph(1)+", "+ph(2)+", "+ph(3)+", "+ph(4)+", 0) RETURNING id",
|
||
channelID, parentPtr, role, string(bigContent)).Scan(&lastMsgID)
|
||
}
|
||
database.TestDB.Exec(
|
||
"INSERT INTO channel_cursors (channel_id, user_id, active_leaf_id) VALUES ("+ph(1)+", "+ph(2)+", "+ph(3)+") ON CONFLICT (channel_id, user_id) DO UPDATE SET active_leaf_id = excluded.active_leaf_id",
|
||
channelID, userID, lastMsgID)
|
||
|
||
// Run compaction — should fail with context budget error
|
||
stores := dialectStores()
|
||
resolver := roles.NewResolver(stores, nil)
|
||
svc := compaction.NewService(stores, resolver)
|
||
|
||
_, err := svc.Compact(context.Background(), compaction.CompactRequest{
|
||
ChannelID: channelID,
|
||
UserID: userID,
|
||
TeamID: nil,
|
||
Trigger: "auto",
|
||
})
|
||
if err == nil {
|
||
t.Fatal("Compact() should fail when content exceeds utility model context window")
|
||
}
|
||
|
||
t.Logf(" ✓ Guard rail triggered: %v", err)
|
||
|
||
if !strings.Contains(err.Error(), "context window") {
|
||
t.Errorf("expected context window error, got: %v", err)
|
||
}
|
||
}
|