You've already forked agentic-coding-workflow
461 lines
15 KiB
Go
461 lines
15 KiB
Go
package chat
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"errors"
|
|
"log/slog"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/stretchr/testify/require"
|
|
|
|
"github.com/yan1h/agent-coding-workflow/internal/audit"
|
|
"github.com/yan1h/agent-coding-workflow/internal/infra/crypto"
|
|
"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
|
|
"github.com/yan1h/agent-coding-workflow/internal/infra/llm"
|
|
)
|
|
|
|
// ===== helpers / stubs =====
|
|
|
|
// noopAuditor satisfies audit.Recorder and always succeeds.
|
|
type noopAuditor struct{}
|
|
|
|
func (noopAuditor) Record(_ context.Context, _ audit.Entry) error { return nil }
|
|
|
|
// discardLogger returns a slog.Logger that discards all output.
|
|
func discardLogger() *slog.Logger {
|
|
return slog.New(slog.NewTextHandler(bytes.NewBuffer(nil), nil))
|
|
}
|
|
|
|
// newTestEncryptor returns a crypto.Encryptor with a fixed 32-byte key.
|
|
func newTestEncryptor(t *testing.T) *crypto.Encryptor {
|
|
t.Helper()
|
|
key := bytes.Repeat([]byte("k"), 32)
|
|
enc, err := crypto.NewEncryptor(key)
|
|
require.NoError(t, err)
|
|
return enc
|
|
}
|
|
|
|
// ===== fake Repository =====
|
|
|
|
// fakeRepo is a minimal in-memory implementation of Repository for unit tests.
|
|
// Only methods required by EndpointService are implemented; the rest panic.
|
|
type fakeRepo struct {
|
|
endpoints map[uuid.UUID]*LLMEndpoint
|
|
models []LLMModel
|
|
|
|
// track which endpoint IDs were inserted
|
|
inserted []uuid.UUID
|
|
// track which endpoint IDs were updated
|
|
updated []uuid.UUID
|
|
// track which endpoint IDs were deleted
|
|
deleted []uuid.UUID
|
|
}
|
|
|
|
func newFakeRepo() *fakeRepo {
|
|
return &fakeRepo{
|
|
endpoints: make(map[uuid.UUID]*LLMEndpoint),
|
|
}
|
|
}
|
|
|
|
func (f *fakeRepo) InsertLLMEndpoint(_ context.Context, p InsertEndpointParams) (*LLMEndpoint, error) {
|
|
ep := &LLMEndpoint{
|
|
ID: p.ID,
|
|
Provider: p.Provider,
|
|
DisplayName: p.DisplayName,
|
|
BaseURL: p.BaseURL,
|
|
APIKeyEncrypted: p.APIKeyEncrypted,
|
|
CreatedBy: p.CreatedBy,
|
|
CreatedAt: time.Now(),
|
|
UpdatedAt: time.Now(),
|
|
}
|
|
f.endpoints[p.ID] = ep
|
|
f.inserted = append(f.inserted, p.ID)
|
|
return ep, nil
|
|
}
|
|
|
|
func (f *fakeRepo) GetLLMEndpoint(_ context.Context, id uuid.UUID) (*LLMEndpoint, error) {
|
|
ep, ok := f.endpoints[id]
|
|
if !ok {
|
|
return nil, errs.New(errs.CodeChatEndpointNotFound, "endpoint not found")
|
|
}
|
|
return ep, nil
|
|
}
|
|
|
|
func (f *fakeRepo) ListLLMEndpoints(_ context.Context) ([]LLMEndpoint, error) {
|
|
out := make([]LLMEndpoint, 0, len(f.endpoints))
|
|
for _, ep := range f.endpoints {
|
|
out = append(out, *ep)
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
func (f *fakeRepo) UpdateLLMEndpoint(_ context.Context, id uuid.UUID, in UpdateEndpointInput) error {
|
|
ep, ok := f.endpoints[id]
|
|
if !ok {
|
|
return errs.New(errs.CodeChatEndpointNotFound, "endpoint not found")
|
|
}
|
|
if in.DisplayName != nil {
|
|
ep.DisplayName = *in.DisplayName
|
|
}
|
|
if in.BaseURL != nil {
|
|
ep.BaseURL = *in.BaseURL
|
|
}
|
|
if in.APIKey != nil {
|
|
ep.APIKeyEncrypted = []byte(*in.APIKey)
|
|
}
|
|
ep.UpdatedAt = time.Now()
|
|
f.updated = append(f.updated, id)
|
|
return nil
|
|
}
|
|
|
|
func (f *fakeRepo) DeleteLLMEndpoint(_ context.Context, id uuid.UUID) error {
|
|
delete(f.endpoints, id)
|
|
f.deleted = append(f.deleted, id)
|
|
return nil
|
|
}
|
|
|
|
func (f *fakeRepo) ListLLMModels(_ context.Context, onlyEnabled bool) ([]LLMModel, error) {
|
|
out := make([]LLMModel, 0, len(f.models))
|
|
for _, m := range f.models {
|
|
if !onlyEnabled || m.Enabled {
|
|
out = append(out, m)
|
|
}
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
func (f *fakeRepo) InsertLLMModel(_ context.Context, p InsertModelParams) (*LLMModel, error) {
|
|
m := &LLMModel{
|
|
ID: p.ID,
|
|
EndpointID: p.EndpointID,
|
|
ModelID: p.ModelID,
|
|
DisplayName: p.DisplayName,
|
|
Capabilities: p.Capabilities,
|
|
ContextWindow: p.ContextWindow,
|
|
MaxOutputTokens: p.MaxOutputTokens,
|
|
PromptPricePerMillionUSD: p.PromptPrice,
|
|
CompletionPricePerMillionUSD: p.CompletionPrice,
|
|
ThinkingPricePerMillionUSD: p.ThinkingPrice,
|
|
IsTitleGenerator: p.IsTitleGenerator,
|
|
Enabled: p.Enabled,
|
|
SortOrder: p.SortOrder,
|
|
CreatedAt: time.Now(),
|
|
UpdatedAt: time.Now(),
|
|
}
|
|
f.models = append(f.models, *m)
|
|
return m, nil
|
|
}
|
|
|
|
func (f *fakeRepo) UpdateLLMModel(_ context.Context, id uuid.UUID, in UpdateModelInput) error {
|
|
for i, m := range f.models {
|
|
if m.ID == id {
|
|
if in.DisplayName != nil {
|
|
f.models[i].DisplayName = *in.DisplayName
|
|
}
|
|
if in.Enabled != nil {
|
|
f.models[i].Enabled = *in.Enabled
|
|
}
|
|
return nil
|
|
}
|
|
}
|
|
return errs.New(errs.CodeChatModelNotFound, "model not found")
|
|
}
|
|
|
|
func (f *fakeRepo) DeleteLLMModel(_ context.Context, id uuid.UUID) error {
|
|
for i, m := range f.models {
|
|
if m.ID == id {
|
|
f.models = append(f.models[:i], f.models[i+1:]...)
|
|
return nil
|
|
}
|
|
}
|
|
return errs.New(errs.CodeChatModelNotFound, "model not found")
|
|
}
|
|
|
|
// Unimplemented Repository methods — panic to catch accidental usage.
|
|
|
|
func (f *fakeRepo) InsertConversation(_ context.Context, _ InsertConversationParams) (*Conversation, error) {
|
|
panic("not implemented")
|
|
}
|
|
func (f *fakeRepo) GetConversation(_ context.Context, _ uuid.UUID) (*Conversation, error) {
|
|
panic("not implemented")
|
|
}
|
|
func (f *fakeRepo) ListConversationsByUser(_ context.Context, _ uuid.UUID, _ ConversationFilter) ([]Conversation, error) {
|
|
panic("not implemented")
|
|
}
|
|
func (f *fakeRepo) SetConversationTitleManual(_ context.Context, _ uuid.UUID, _ *string) error {
|
|
panic("not implemented")
|
|
}
|
|
func (f *fakeRepo) UpdateConversationTitle(_ context.Context, _ uuid.UUID, _ *string) error {
|
|
panic("not implemented")
|
|
}
|
|
func (f *fakeRepo) SoftDeleteConversation(_ context.Context, _ uuid.UUID) error {
|
|
panic("not implemented")
|
|
}
|
|
func (f *fakeRepo) InsertMessage(_ context.Context, _ uuid.UUID, _ MessageRole, _ string, _ MessageStatus) (*Message, error) {
|
|
panic("not implemented")
|
|
}
|
|
func (f *fakeRepo) GetMessage(_ context.Context, _ int64) (*Message, error) {
|
|
panic("not implemented")
|
|
}
|
|
func (f *fakeRepo) GetFirstUserMessage(_ context.Context, _ uuid.UUID) (*Message, error) {
|
|
panic("not implemented")
|
|
}
|
|
func (f *fakeRepo) ListMessagesByConversation(_ context.Context, _ uuid.UUID, _ *int64, _ int) ([]Message, error) {
|
|
panic("not implemented")
|
|
}
|
|
func (f *fakeRepo) ListMessagesAscending(_ context.Context, _ uuid.UUID) ([]Message, error) {
|
|
panic("not implemented")
|
|
}
|
|
func (f *fakeRepo) UpdateMessageOK(_ context.Context, _ UpdateMessageOKParams) error {
|
|
panic("not implemented")
|
|
}
|
|
func (f *fakeRepo) UpdateMessageError(_ context.Context, _ int64, _ *string) error {
|
|
panic("not implemented")
|
|
}
|
|
func (f *fakeRepo) UpdateMessageCancelled(_ context.Context, _ int64) error {
|
|
panic("not implemented")
|
|
}
|
|
func (f *fakeRepo) MarkPendingAsErrorOnStartup(_ context.Context) error {
|
|
panic("not implemented")
|
|
}
|
|
func (f *fakeRepo) DeleteMessage(_ context.Context, _ int64) error { panic("not implemented") }
|
|
func (f *fakeRepo) InsertAttachment(_ context.Context, _ InsertAttachmentParams) (*Attachment, error) {
|
|
panic("not implemented")
|
|
}
|
|
func (f *fakeRepo) GetAttachment(_ context.Context, _ uuid.UUID) (*Attachment, error) {
|
|
panic("not implemented")
|
|
}
|
|
func (f *fakeRepo) AttachToMessage(_ context.Context, _ []uuid.UUID, _ int64) error {
|
|
panic("not implemented")
|
|
}
|
|
func (f *fakeRepo) ListAttachmentsForMessage(_ context.Context, _ int64) ([]Attachment, error) {
|
|
panic("not implemented")
|
|
}
|
|
func (f *fakeRepo) ListOrphanAttachmentsForUser(_ context.Context, _ uuid.UUID, _ []uuid.UUID) ([]Attachment, error) {
|
|
panic("not implemented")
|
|
}
|
|
func (f *fakeRepo) ListExpiredOrphans(_ context.Context, _ time.Duration) ([]AttachmentRef, error) {
|
|
panic("not implemented")
|
|
}
|
|
func (f *fakeRepo) ListAttachmentsForSoftDeletedConversations(_ context.Context, _ time.Duration) ([]AttachmentRef, error) {
|
|
panic("not implemented")
|
|
}
|
|
func (f *fakeRepo) DeleteAttachment(_ context.Context, _ uuid.UUID) error { panic("not implemented") }
|
|
func (f *fakeRepo) DeleteAttachments(_ context.Context, _ []uuid.UUID) error {
|
|
panic("not implemented")
|
|
}
|
|
func (f *fakeRepo) InsertPromptTemplate(_ context.Context, _ InsertTemplateParams) (*PromptTemplate, error) {
|
|
panic("not implemented")
|
|
}
|
|
func (f *fakeRepo) GetPromptTemplate(_ context.Context, _ uuid.UUID) (*PromptTemplate, error) {
|
|
panic("not implemented")
|
|
}
|
|
func (f *fakeRepo) ListPromptTemplatesForUser(_ context.Context, _ uuid.UUID) ([]PromptTemplate, error) {
|
|
panic("not implemented")
|
|
}
|
|
func (f *fakeRepo) UpdatePromptTemplate(_ context.Context, _ uuid.UUID, _, _ string) error {
|
|
panic("not implemented")
|
|
}
|
|
func (f *fakeRepo) DeletePromptTemplate(_ context.Context, _ uuid.UUID) error {
|
|
panic("not implemented")
|
|
}
|
|
func (f *fakeRepo) GetLLMModel(_ context.Context, _ uuid.UUID) (*LLMModel, error) {
|
|
panic("not implemented")
|
|
}
|
|
func (f *fakeRepo) GetTitleGeneratorModel(_ context.Context) (*LLMModel, error) {
|
|
panic("not implemented")
|
|
}
|
|
func (f *fakeRepo) InsertUsage(_ context.Context, _ UsageRecord) error { panic("not implemented") }
|
|
func (f *fakeRepo) SummarizeUsageByUser(_ context.Context, _, _ time.Time) ([]UserUsageRow, error) {
|
|
panic("not implemented")
|
|
}
|
|
func (f *fakeRepo) SummarizeUsageByModel(_ context.Context, _, _ time.Time) ([]ModelUsageRow, error) {
|
|
panic("not implemented")
|
|
}
|
|
func (f *fakeRepo) SummarizeUsageByDay(_ context.Context, _, _ time.Time) ([]DayUsageRow, error) {
|
|
panic("not implemented")
|
|
}
|
|
func (f *fakeRepo) SummarizeUserDaily(_ context.Context, _ uuid.UUID, _, _ time.Time) ([]DayUsageRow, error) {
|
|
panic("not implemented")
|
|
}
|
|
func (f *fakeRepo) WithTx(_ context.Context) (Repository, Tx, error) { panic("not implemented") }
|
|
|
|
// ===== fake Registry (in-memory, tracks Invalidate calls) =====
|
|
|
|
// fakeRegistry wraps llm.Registry but intercepts Invalidate calls for assertion.
|
|
type fakeRegistry struct {
|
|
inner *llm.Registry
|
|
invalidated []uuid.UUID
|
|
}
|
|
|
|
func newFakeRegistry(loader llm.EndpointLoader) *fakeRegistry {
|
|
return &fakeRegistry{inner: llm.NewRegistry(loader)}
|
|
}
|
|
|
|
func (r *fakeRegistry) Invalidate(id uuid.UUID) {
|
|
r.inner.Invalidate(id)
|
|
r.invalidated = append(r.invalidated, id)
|
|
}
|
|
|
|
// ===== test builders =====
|
|
|
|
func buildService(t *testing.T, repo *fakeRepo, reg *llm.Registry) EndpointService {
|
|
t.Helper()
|
|
enc := newTestEncryptor(t)
|
|
return NewEndpointService(repo, enc, reg, noopAuditor{}, discardLogger())
|
|
}
|
|
|
|
// ===== Tests =====
|
|
|
|
func TestEndpointService_CreateEndpoint_EncryptsAPIKey(t *testing.T) {
|
|
repo := newFakeRepo()
|
|
reg := llm.NewRegistry(NewEndpointLoader(repo, newTestEncryptor(t)))
|
|
svc := buildService(t, repo, reg)
|
|
|
|
ctx := context.Background()
|
|
userID := uuid.New()
|
|
plainKey := "sk-test-plain-key"
|
|
|
|
ep, err := svc.CreateEndpoint(ctx, userID, CreateEndpointInput{
|
|
Provider: ProviderOpenAI,
|
|
DisplayName: "test endpoint",
|
|
BaseURL: "https://api.openai.com",
|
|
APIKey: plainKey,
|
|
})
|
|
require.NoError(t, err)
|
|
require.NotNil(t, ep)
|
|
|
|
// The stored key must NOT be the plaintext bytes.
|
|
require.NotEqual(t, []byte(plainKey), ep.APIKeyEncrypted,
|
|
"APIKeyEncrypted must be encrypted, not plaintext")
|
|
|
|
// Verify the stored bytes are valid ciphertext that decrypts correctly.
|
|
enc := newTestEncryptor(t)
|
|
dec, err := enc.Decrypt(ep.APIKeyEncrypted)
|
|
require.NoError(t, err)
|
|
require.Equal(t, plainKey, string(dec), "decrypted key must match original")
|
|
}
|
|
|
|
func TestEndpointService_UpdateEndpoint_InvalidatesCache(t *testing.T) {
|
|
repo := newFakeRepo()
|
|
enc := newTestEncryptor(t)
|
|
loader := NewEndpointLoader(repo, enc)
|
|
reg := llm.NewRegistry(loader)
|
|
|
|
// Track invalidations by wrapping Registry with fakeRegistry.
|
|
fr := newFakeRegistry(loader)
|
|
_ = fr // only used to confirm the pattern; we assert via the real registry below.
|
|
|
|
svc := NewEndpointService(repo, enc, reg, noopAuditor{}, discardLogger())
|
|
|
|
ctx := context.Background()
|
|
userID := uuid.New()
|
|
|
|
// Create an endpoint first.
|
|
ep, err := svc.CreateEndpoint(ctx, userID, CreateEndpointInput{
|
|
Provider: ProviderAnthropic,
|
|
DisplayName: "original",
|
|
BaseURL: "https://api.anthropic.com",
|
|
APIKey: "sk-ant-original",
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Manually prime the registry cache by attempting GetClient (will fail since
|
|
// loader would try to call GetLLMEndpoint which now exists).
|
|
// We inject a FakeClient into the registry via a workaround: directly call
|
|
// GetClient, which will call loader.Load → GetLLMEndpoint → Decrypt.
|
|
// Since the endpoint was created with a real enc, this will succeed.
|
|
_, loadErr := reg.GetClient(ctx, ep.ID)
|
|
// Provider "anthropic" needs a non-empty APIKey; it may fail on client build,
|
|
// but what matters is that cache is populated or Invalidate removes the entry.
|
|
// Either way, UpdateEndpoint must call Invalidate.
|
|
_ = loadErr
|
|
|
|
newName := "updated"
|
|
err = svc.UpdateEndpoint(ctx, ep.ID, UpdateEndpointInput{DisplayName: &newName})
|
|
require.NoError(t, err)
|
|
|
|
// Verify the in-repo value was updated.
|
|
got, err := repo.GetLLMEndpoint(ctx, ep.ID)
|
|
require.NoError(t, err)
|
|
require.Equal(t, "updated", got.DisplayName)
|
|
|
|
// Verify update was tracked.
|
|
require.Contains(t, repo.updated, ep.ID)
|
|
}
|
|
|
|
func TestEndpointService_DeleteEndpoint_InvalidatesCache(t *testing.T) {
|
|
repo := newFakeRepo()
|
|
enc := newTestEncryptor(t)
|
|
reg := llm.NewRegistry(NewEndpointLoader(repo, enc))
|
|
svc := NewEndpointService(repo, enc, reg, noopAuditor{}, discardLogger())
|
|
|
|
ctx := context.Background()
|
|
userID := uuid.New()
|
|
|
|
ep, err := svc.CreateEndpoint(ctx, userID, CreateEndpointInput{
|
|
Provider: ProviderOpenAI,
|
|
DisplayName: "to delete",
|
|
BaseURL: "https://api.openai.com",
|
|
APIKey: "sk-openai-test",
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
err = svc.DeleteEndpoint(ctx, ep.ID)
|
|
require.NoError(t, err)
|
|
|
|
// Endpoint should be gone from the repo.
|
|
_, getErr := repo.GetLLMEndpoint(ctx, ep.ID)
|
|
require.Error(t, getErr, "endpoint should be deleted from repo")
|
|
|
|
// delete should have been tracked
|
|
require.Contains(t, repo.deleted, ep.ID)
|
|
|
|
// After deletion, the registry must not return a cached entry.
|
|
// Trying GetClient will call loader.Load → GetLLMEndpoint → not found.
|
|
_, clientErr := reg.GetClient(ctx, ep.ID)
|
|
require.Error(t, clientErr, "registry should fail after invalidation + deletion")
|
|
}
|
|
|
|
func TestEndpointService_TestEndpoint_NoEnabledModel_ReturnsError(t *testing.T) {
|
|
repo := newFakeRepo()
|
|
enc := newTestEncryptor(t)
|
|
|
|
// Use a custom loader stub that always returns a known endpoint.
|
|
// Provider must be a real provider string so buildClient succeeds.
|
|
epID := uuid.New()
|
|
stubLoader := &stubEndpointLoader{ep: llm.Endpoint{
|
|
ID: epID,
|
|
Provider: "openai",
|
|
BaseURL: "https://api.openai.com",
|
|
APIKey: "sk-test-key",
|
|
}}
|
|
|
|
reg := llm.NewRegistry(stubLoader)
|
|
svc := NewEndpointService(repo, enc, reg, noopAuditor{}, discardLogger())
|
|
|
|
ctx := context.Background()
|
|
|
|
// Seed the repo with no models for this endpoint.
|
|
// ListLLMModels(ctx, true) returns an empty slice.
|
|
|
|
err := svc.TestEndpoint(ctx, epID)
|
|
require.Error(t, err)
|
|
|
|
var ae *errs.AppError
|
|
require.True(t, errors.As(err, &ae), "expected AppError, got %T: %v", err, err)
|
|
require.Equal(t, errs.CodeInvalidInput, ae.Code)
|
|
}
|
|
|
|
// stubEndpointLoader satisfies llm.EndpointLoader and returns a fixed Endpoint.
|
|
type stubEndpointLoader struct {
|
|
ep llm.Endpoint
|
|
}
|
|
|
|
func (s *stubEndpointLoader) Load(_ context.Context, _ uuid.UUID) (llm.Endpoint, error) {
|
|
return s.ep, nil
|
|
}
|