245 lines
11 KiB
Go
245 lines
11 KiB
Go
|
package memory
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"fmt"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"github.com/Tencent/WeKnora/internal/types"
|
||
|
|
"github.com/stretchr/testify/require"
|
||
|
|
)
|
||
|
|
|
||
|
|
func conflictingMemoryProposal(t *testing.T, svc *Service, ctx context.Context) (*types.MemoryItem, *types.MemoryItem) {
|
||
|
|
t.Helper()
|
||
|
|
accepted, err := svc.Remember(ctx, types.MemoryItem{
|
||
|
|
Kind: types.MemoryKindProfile, Topic: "database", Content: "My database is MySQL",
|
||
|
|
Origin: types.MemoryOriginExplicit, Importance: 4,
|
||
|
|
})
|
||
|
|
require.NoError(t, err)
|
||
|
|
proposal, err := svc.Remember(ctx, types.MemoryItem{
|
||
|
|
Kind: types.MemoryKindProfile, Topic: "database", Content: "My database is PostgreSQL",
|
||
|
|
Origin: types.MemoryOriginExtracted, Inferred: true, Importance: 4,
|
||
|
|
})
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, types.MemoryStatusPending, proposal.Status)
|
||
|
|
return accepted, proposal
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestPendingConflictPreservesAcceptedMemoryWhenRejected(t *testing.T) {
|
||
|
|
svc, _, tenants := newMemoryHarness(t)
|
||
|
|
ctx := enabledCtx(t, tenants, 1, "alice")
|
||
|
|
accepted, proposal := conflictingMemoryProposal(t, svc, ctx)
|
||
|
|
scope := scopeFor(t, ctx)
|
||
|
|
|
||
|
|
prior, err := svc.repo.GetItem(ctx, scope, accepted.ID)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, types.MemoryStatusActive, prior.Status)
|
||
|
|
require.Nil(t, prior.InvalidAt)
|
||
|
|
require.Empty(t, prior.SupersededBy)
|
||
|
|
recall := svc.Recall(ctx, "database")
|
||
|
|
require.Contains(t, recall.Prompt, accepted.Content)
|
||
|
|
require.NotContains(t, recall.Prompt, proposal.Content)
|
||
|
|
require.Contains(t, svc.RetrievalContextFor(ctx).Background, accepted.Content)
|
||
|
|
require.NotContains(t, svc.RetrievalContextFor(ctx).Background, proposal.Content)
|
||
|
|
|
||
|
|
require.NoError(t, svc.RejectItem(ctx, proposal.ID))
|
||
|
|
deleted, err := svc.repo.GetItem(ctx, scope, proposal.ID)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Nil(t, deleted)
|
||
|
|
prior, err = svc.repo.GetItem(ctx, scope, accepted.ID)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, types.MemoryStatusActive, prior.Status)
|
||
|
|
require.Nil(t, prior.InvalidAt)
|
||
|
|
require.Empty(t, prior.SupersededBy)
|
||
|
|
require.Contains(t, svc.Recall(ctx, "database").Prompt, accepted.Content)
|
||
|
|
|
||
|
|
_, err = svc.Remember(ctx, types.MemoryItem{
|
||
|
|
Kind: proposal.Kind, Topic: proposal.Topic, Content: proposal.Content,
|
||
|
|
Origin: types.MemoryOriginExtracted, Inferred: true,
|
||
|
|
})
|
||
|
|
require.ErrorIs(t, err, ErrPreviouslyForgotten)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestConfirmPendingConflictSupersedesOnlyAcceptedConflict(t *testing.T) {
|
||
|
|
svc, _, tenants := newMemoryHarness(t)
|
||
|
|
ctx := enabledCtx(t, tenants, 1, "alice")
|
||
|
|
accepted, proposal := conflictingMemoryProposal(t, svc, ctx)
|
||
|
|
unrelated, err := svc.Remember(ctx, types.MemoryItem{
|
||
|
|
Kind: types.MemoryKindProfile, Topic: "role", Content: "I work as a data engineer",
|
||
|
|
})
|
||
|
|
require.NoError(t, err)
|
||
|
|
confirmed, err := svc.ConfirmItem(ctx, proposal.ID)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, types.MemoryStatusActive, confirmed.Status)
|
||
|
|
scope := scopeFor(t, ctx)
|
||
|
|
prior, err := svc.repo.GetItem(ctx, scope, accepted.ID)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, types.MemoryStatusSuperseded, prior.Status)
|
||
|
|
require.Equal(t, confirmed.ID, prior.SupersededBy)
|
||
|
|
require.NotNil(t, prior.InvalidAt)
|
||
|
|
require.Equal(t, accepted.Content, prior.Content)
|
||
|
|
other, err := svc.repo.GetItem(ctx, scope, unrelated.ID)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, types.MemoryStatusActive, other.Status)
|
||
|
|
recall := svc.Recall(ctx, "database")
|
||
|
|
require.Contains(t, recall.Prompt, proposal.Content)
|
||
|
|
require.NotContains(t, recall.Prompt, accepted.Content)
|
||
|
|
|
||
|
|
_, err = svc.ConfirmItem(ctx, accepted.ID)
|
||
|
|
require.ErrorIs(t, err, ErrItemNotPending, "historical memory must never be reactivated by confirmation")
|
||
|
|
_, err = svc.ConfirmItem(ctx, confirmed.ID)
|
||
|
|
require.ErrorIs(t, err, ErrItemNotPending)
|
||
|
|
require.ErrorIs(t, svc.RejectItem(ctx, confirmed.ID), ErrItemNotPending,
|
||
|
|
"a stale reject action must not delete a confirmed item")
|
||
|
|
stillConfirmed, err := svc.repo.GetItem(ctx, scope, confirmed.ID)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, types.MemoryStatusActive, stillConfirmed.Status)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestConfirmPendingExtensionSupersedesContainedAcceptedMemory(t *testing.T) {
|
||
|
|
svc, _, tenants := newMemoryHarness(t)
|
||
|
|
ctx := enabledCtx(t, tenants, 1, "alice")
|
||
|
|
accepted, err := svc.Remember(ctx, types.MemoryItem{
|
||
|
|
Kind: types.MemoryKindProfile, Topic: "role", Content: "I maintain data pipelines",
|
||
|
|
})
|
||
|
|
require.NoError(t, err)
|
||
|
|
proposal, err := svc.Remember(ctx, types.MemoryItem{
|
||
|
|
Kind: types.MemoryKindProfile, Topic: "engineering", Content: "I maintain data pipelines in Go",
|
||
|
|
Origin: types.MemoryOriginExtracted, Inferred: true,
|
||
|
|
})
|
||
|
|
require.NoError(t, err)
|
||
|
|
prior, err := svc.repo.GetItem(ctx, scopeFor(t, ctx), accepted.ID)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, types.MemoryStatusActive, prior.Status)
|
||
|
|
_, err = svc.ConfirmItem(ctx, proposal.ID)
|
||
|
|
require.NoError(t, err)
|
||
|
|
prior, err = svc.repo.GetItem(ctx, scopeFor(t, ctx), accepted.ID)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, types.MemoryStatusSuperseded, prior.Status)
|
||
|
|
require.Equal(t, proposal.ID, prior.SupersededBy)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestMemoryReviewRejectsWrongStatusAndScope(t *testing.T) {
|
||
|
|
svc, db, tenants := newMemoryHarness(t)
|
||
|
|
ctx := enabledCtx(t, tenants, 1, "alice")
|
||
|
|
_, proposal := conflictingMemoryProposal(t, svc, ctx)
|
||
|
|
for _, foreignCtx := range []context.Context{
|
||
|
|
enabledCtx(t, tenants, 1, "bob"),
|
||
|
|
enabledCtx(t, tenants, 2, "alice"),
|
||
|
|
} {
|
||
|
|
foreignAccepted, foreignProposal := conflictingMemoryProposal(t, svc, foreignCtx)
|
||
|
|
_, err := svc.ConfirmItem(foreignCtx, proposal.ID)
|
||
|
|
require.ErrorIs(t, err, ErrItemNotFound)
|
||
|
|
require.ErrorIs(t, svc.RejectItem(foreignCtx, proposal.ID), ErrItemNotFound)
|
||
|
|
_, err = svc.ConfirmItem(foreignCtx, foreignProposal.ID)
|
||
|
|
require.NoError(t, err)
|
||
|
|
prior, err := svc.repo.GetItem(foreignCtx, scopeFor(t, foreignCtx), foreignAccepted.ID)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, types.MemoryStatusSuperseded, prior.Status)
|
||
|
|
}
|
||
|
|
stillPending, err := svc.repo.GetItem(ctx, scopeFor(t, ctx), proposal.ID)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, types.MemoryStatusPending, stillPending.Status)
|
||
|
|
|
||
|
|
for _, status := range []string{types.MemoryStatusActive, types.MemoryStatusArchived, types.MemoryStatusSuperseded} {
|
||
|
|
require.NoError(t, db.Model(&types.MemoryItem{}).Where("id = ?", proposal.ID).Update("status", status).Error)
|
||
|
|
_, err := svc.ConfirmItem(ctx, proposal.ID)
|
||
|
|
require.ErrorIs(t, err, ErrItemNotPending)
|
||
|
|
require.ErrorIs(t, svc.RejectItem(ctx, proposal.ID), ErrItemNotPending)
|
||
|
|
unchanged, err := svc.repo.GetItem(ctx, scopeFor(t, ctx), proposal.ID)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, status, unchanged.Status)
|
||
|
|
}
|
||
|
|
_, err = svc.ConfirmItem(ctx, "missing")
|
||
|
|
require.ErrorIs(t, err, ErrItemNotFound)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestPendingConfirmationRollsBackOnWriteFailureAndRetries(t *testing.T) {
|
||
|
|
for _, failingStatus := range []string{types.MemoryStatusSuperseded, types.MemoryStatusActive} {
|
||
|
|
t.Run(failingStatus, func(t *testing.T) {
|
||
|
|
svc, db, tenants := newMemoryHarness(t)
|
||
|
|
ctx := enabledCtx(t, tenants, 1, "alice")
|
||
|
|
accepted, proposal := conflictingMemoryProposal(t, svc, ctx)
|
||
|
|
require.NoError(t, db.Exec(fmt.Sprintf(`CREATE TRIGGER fail_memory_review
|
||
|
|
BEFORE UPDATE OF status ON memory_items WHEN NEW.status = '%s'
|
||
|
|
BEGIN SELECT RAISE(ABORT, 'simulated review write failure'); END`, failingStatus)).Error)
|
||
|
|
confirmed, err := svc.ConfirmItem(ctx, proposal.ID)
|
||
|
|
require.ErrorContains(t, err, "simulated review write failure")
|
||
|
|
require.Nil(t, confirmed)
|
||
|
|
scope := scopeFor(t, ctx)
|
||
|
|
prior, err := svc.repo.GetItem(ctx, scope, accepted.ID)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, types.MemoryStatusActive, prior.Status)
|
||
|
|
require.Nil(t, prior.InvalidAt)
|
||
|
|
require.Empty(t, prior.SupersededBy)
|
||
|
|
pending, err := svc.repo.GetItem(ctx, scope, proposal.ID)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, types.MemoryStatusPending, pending.Status)
|
||
|
|
require.NoError(t, db.Exec("DROP TRIGGER fail_memory_review").Error)
|
||
|
|
|
||
|
|
confirmed, err = svc.ConfirmItem(ctx, proposal.ID)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, types.MemoryStatusActive, confirmed.Status)
|
||
|
|
prior, err = svc.repo.GetItem(ctx, scope, accepted.ID)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, types.MemoryStatusSuperseded, prior.Status)
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestPendingRejectionRollsBackTombstoneAndRetries(t *testing.T) {
|
||
|
|
svc, db, tenants := newMemoryHarness(t)
|
||
|
|
ctx := enabledCtx(t, tenants, 1, "alice")
|
||
|
|
accepted, proposal := conflictingMemoryProposal(t, svc, ctx)
|
||
|
|
require.NoError(t, db.Exec(`CREATE TRIGGER fail_memory_rejection
|
||
|
|
BEFORE DELETE ON memory_items WHEN OLD.status = 'pending'
|
||
|
|
BEGIN SELECT RAISE(ABORT, 'simulated review delete failure'); END`).Error)
|
||
|
|
require.ErrorContains(t, svc.RejectItem(ctx, proposal.ID), "simulated review delete failure")
|
||
|
|
scope := scopeFor(t, ctx)
|
||
|
|
forgotten, err := svc.repo.HasTombstone(ctx, scope, types.MemoryFingerprint(proposal.Content))
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.False(t, forgotten, "failed rejection must not prevent a later review")
|
||
|
|
pending, err := svc.repo.GetItem(ctx, scope, proposal.ID)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, types.MemoryStatusPending, pending.Status)
|
||
|
|
require.NoError(t, db.Exec("DROP TRIGGER fail_memory_rejection").Error)
|
||
|
|
require.NoError(t, svc.RejectItem(ctx, proposal.ID))
|
||
|
|
forgotten, err = svc.repo.HasTombstone(ctx, scope, types.MemoryFingerprint(proposal.Content))
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.True(t, forgotten)
|
||
|
|
prior, err := svc.repo.GetItem(ctx, scope, accepted.ID)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, types.MemoryStatusActive, prior.Status)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestPendingWriteFailurePreservesAcceptedMemoryAndRetries(t *testing.T) {
|
||
|
|
svc, db, tenants := newMemoryHarness(t)
|
||
|
|
ctx := enabledCtx(t, tenants, 1, "alice")
|
||
|
|
accepted, err := svc.Remember(ctx, types.MemoryItem{
|
||
|
|
Kind: types.MemoryKindProfile, Topic: "database", Content: "My database is MySQL",
|
||
|
|
})
|
||
|
|
require.NoError(t, err)
|
||
|
|
proposal := types.MemoryItem{
|
||
|
|
Kind: types.MemoryKindProfile, Topic: "database", Content: "My database is PostgreSQL",
|
||
|
|
Origin: types.MemoryOriginExtracted, Inferred: true,
|
||
|
|
}
|
||
|
|
require.NoError(t, db.Exec(`CREATE TRIGGER fail_memory_proposal
|
||
|
|
BEFORE INSERT ON memory_items WHEN NEW.status = 'pending'
|
||
|
|
BEGIN SELECT RAISE(ABORT, 'simulated proposal insert failure'); END`).Error)
|
||
|
|
_, err = svc.Remember(ctx, proposal)
|
||
|
|
require.ErrorContains(t, err, "simulated proposal insert failure")
|
||
|
|
prior, err := svc.repo.GetItem(ctx, scopeFor(t, ctx), accepted.ID)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, types.MemoryStatusActive, prior.Status)
|
||
|
|
require.NoError(t, db.Exec("DROP TRIGGER fail_memory_proposal").Error)
|
||
|
|
stored, err := svc.Remember(ctx, proposal)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, types.MemoryStatusPending, stored.Status)
|
||
|
|
_, pendingCount, err := svc.ListItems(ctx, types.MemoryStatusPending, 10, 0)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, int64(1), pendingCount)
|
||
|
|
prior, err = svc.repo.GetItem(ctx, scopeFor(t, ctx), accepted.ID)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, types.MemoryStatusActive, prior.Status)
|
||
|
|
}
|