From aaaafb723d1e045a467b767af040eede1ca6fb9d Mon Sep 17 00:00:00 2001 From: ZanzyTHEbar Date: Wed, 18 Feb 2026 15:52:59 +0000 Subject: [PATCH] feat(memory): add KV store, documents, and audit log subsystems Extend the memory delegate with three new storage subsystems: - agent_kv: key-value store for agent state and preferences - agent_documents: structured document storage with metadata - agent_audit_log: append-only audit trail for agent actions Includes goose migrations (004-006), sqlc query generation, delegate methods, and comprehensive test coverage. --- pkg/memory/delegate/sqlite.go | 266 +++++++++++++ pkg/memory/delegate/sqlite_audit_test.go | 281 +++++++++++++ pkg/memory/delegate/sqlite_doc_test.go | 330 ++++++++++++++++ .../delegate/sqlite_integration_test.go | 227 +++++++++++ pkg/memory/delegate/sqlite_kv_test.go | 330 ++++++++++++++++ .../delegate/sqlite_session_msg_test.go | 269 +++++++++++++ pkg/memory/delegate/sqlite_test.go | 138 +++++++ pkg/memory/memory.go | 47 +++ pkg/memory/migrations/004_agent_kv.go | 35 ++ pkg/memory/migrations/005_agent_documents.go | 47 +++ pkg/memory/migrations/006_agent_audit_log.go | 46 +++ pkg/memory/sqlc/agent_audit_log.sql.go | 368 ++++++++++++++++++ pkg/memory/sqlc/agent_documents.sql.go | 305 +++++++++++++++ pkg/memory/sqlc/agent_kv.sql.go | 160 ++++++++ pkg/memory/sqlc/models.go | 31 ++ pkg/memory/sqlc/querier.go | 285 ++++++++++++++ pkg/memory/sqlc/queries/agent_audit_log.sql | 81 ++++ pkg/memory/sqlc/queries/agent_documents.sql | 77 ++++ pkg/memory/sqlc/queries/agent_kv.sql | 34 ++ pkg/memory/sqlc/queries/recall.sql | 60 ++- pkg/memory/sqlc/recall.sql.go | 204 ++++++++++ pkg/memory/sqlc/schema.sql | 39 +- pkg/memory/sqlc/sqlc.yaml | 8 + pkg/memory/store/memory_store.go | 10 + 24 files changed, 3676 insertions(+), 2 deletions(-) create mode 100644 pkg/memory/delegate/sqlite_audit_test.go create mode 100644 pkg/memory/delegate/sqlite_doc_test.go create mode 100644 pkg/memory/delegate/sqlite_integration_test.go create mode 100644 pkg/memory/delegate/sqlite_kv_test.go create mode 100644 pkg/memory/delegate/sqlite_session_msg_test.go create mode 100644 pkg/memory/migrations/004_agent_kv.go create mode 100644 pkg/memory/migrations/005_agent_documents.go create mode 100644 pkg/memory/migrations/006_agent_audit_log.go create mode 100644 pkg/memory/sqlc/agent_audit_log.sql.go create mode 100644 pkg/memory/sqlc/agent_documents.sql.go create mode 100644 pkg/memory/sqlc/agent_kv.sql.go create mode 100644 pkg/memory/sqlc/queries/agent_audit_log.sql create mode 100644 pkg/memory/sqlc/queries/agent_documents.sql create mode 100644 pkg/memory/sqlc/queries/agent_kv.sql diff --git a/pkg/memory/delegate/sqlite.go b/pkg/memory/delegate/sqlite.go index 22487fd68..800c198a8 100644 --- a/pkg/memory/delegate/sqlite.go +++ b/pkg/memory/delegate/sqlite.go @@ -411,6 +411,235 @@ func (d *LibSQLDelegate) CountArchivalChunks(ctx context.Context) (int, error) { return int(count), err } +// --- Key-Value Store --- + +func (d *LibSQLDelegate) GetKV(ctx context.Context, agentID, key string) (string, error) { + row, err := d.queries.GetKV(ctx, memsqlc.GetKVParams{ + AgentID: agentID, + Key: key, + }) + if err == sql.ErrNoRows { + return "", nil + } + if err != nil { + return "", err + } + return row.Value, nil +} + +func (d *LibSQLDelegate) UpsertKV(ctx context.Context, agentID, key, value string) error { + return d.queries.UpsertKV(ctx, memsqlc.UpsertKVParams{ + AgentID: agentID, + Key: key, + Value: value, + }) +} + +func (d *LibSQLDelegate) DeleteKV(ctx context.Context, agentID, key string) error { + return d.queries.DeleteKV(ctx, memsqlc.DeleteKVParams{ + AgentID: agentID, + Key: key, + }) +} + +func (d *LibSQLDelegate) ListKVByPrefix(ctx context.Context, agentID, prefix string, limit int) (map[string]string, error) { + rows, err := d.queries.ListKVByPrefix(ctx, memsqlc.ListKVByPrefixParams{ + AgentID: agentID, + Prefix: &prefix, + Lim: int64(limit), + }) + if err != nil { + return nil, err + } + result := make(map[string]string, len(rows)) + for _, row := range rows { + result[row.Key] = row.Value + } + return result, nil +} + +// --- Documents --- + +func (d *LibSQLDelegate) GetDocument(ctx context.Context, agentID, name string) (*memory.AgentDocument, error) { + row, err := d.queries.GetDocument(ctx, memsqlc.GetDocumentParams{ + AgentID: agentID, + Name: name, + }) + if err == sql.ErrNoRows { + return nil, nil + } + if err != nil { + return nil, err + } + return sqlcDocToMemory(row), nil +} + +func (d *LibSQLDelegate) UpsertDocument(ctx context.Context, doc *memory.AgentDocument) error { + return d.queries.UpsertDocument(ctx, memsqlc.UpsertDocumentParams{ + ID: doc.ID, + AgentID: doc.AgentID, + Name: doc.Name, + Category: doc.Category, + Content: doc.Content, + }) +} + +func (d *LibSQLDelegate) DeleteDocument(ctx context.Context, agentID, name string) error { + return d.queries.DeleteDocument(ctx, memsqlc.DeleteDocumentParams{ + AgentID: agentID, + Name: name, + }) +} + +func (d *LibSQLDelegate) ListDocumentsByCategory(ctx context.Context, agentID, category string) ([]*memory.AgentDocument, error) { + rows, err := d.queries.ListDocumentsByCategory(ctx, memsqlc.ListDocumentsByCategoryParams{ + AgentID: agentID, + Category: category, + }) + if err != nil { + return nil, err + } + docs := make([]*memory.AgentDocument, len(rows)) + for i, row := range rows { + docs[i] = sqlcDocToMemory(row) + } + return docs, nil +} + +func (d *LibSQLDelegate) ListAllDocuments(ctx context.Context, agentID string) ([]*memory.AgentDocument, error) { + rows, err := d.queries.ListAllDocuments(ctx, memsqlc.ListAllDocumentsParams{ + AgentID: agentID, + }) + if err != nil { + return nil, err + } + docs := make([]*memory.AgentDocument, len(rows)) + for i, row := range rows { + docs[i] = sqlcDocToMemory(row) + } + return docs, nil +} + +// --- Session Messages --- + +func (d *LibSQLDelegate) InsertSessionMessage(ctx context.Context, agentID, sessionKey, role, content string) error { + return d.queries.InsertSessionMessage(ctx, memsqlc.InsertSessionMessageParams{ + ID: ids.New(), + AgentID: agentID, + SessionKey: sessionKey, + Role: role, + Content: content, + }) +} + +func (d *LibSQLDelegate) ListSessionMessages(ctx context.Context, agentID, sessionKey, role string, limit int) ([]*memory.RecallItem, error) { + rows, err := d.queries.ListSessionMessages(ctx, memsqlc.ListSessionMessagesParams{ + AgentID: agentID, + SessionKey: sessionKey, + Role: role, + Lim: int64(limit), + }) + if err != nil { + return nil, err + } + items := make([]*memory.RecallItem, len(rows)) + for i, row := range rows { + items[i] = sqlcRecallToMemory(row) + } + return items, nil +} + +func (d *LibSQLDelegate) CountSessionMessages(ctx context.Context, agentID, sessionKey string) (int64, error) { + return d.queries.CountSessionMessages(ctx, memsqlc.CountSessionMessagesParams{ + AgentID: agentID, + SessionKey: sessionKey, + }) +} + +// --- Audit Log --- + +func (d *LibSQLDelegate) InsertAuditEntry(ctx context.Context, entry *memory.AuditEntry) error { + return d.queries.InsertAuditEntry(ctx, memsqlc.InsertAuditEntryParams{ + ID: entry.ID, + AgentID: entry.AgentID, + SessionKey: entry.SessionKey, + Action: entry.Action, + Target: entry.Target, + Input: &entry.Input, + Output: &entry.Output, + DurationMs: ptrInt64(int64(entry.DurationMS)), + }) +} + +func (d *LibSQLDelegate) ListAuditEntries(ctx context.Context, agentID string, limit int) ([]*memory.AuditEntry, error) { + rows, err := d.queries.ListAuditEntries(ctx, memsqlc.ListAuditEntriesParams{ + AgentID: agentID, + Lim: int64(limit), + }) + if err != nil { + return nil, err + } + entries := make([]*memory.AuditEntry, len(rows)) + for i, row := range rows { + entries[i] = sqlcAuditToMemory(row) + } + return entries, nil +} + +func (d *LibSQLDelegate) ListAuditEntriesByAction(ctx context.Context, agentID, action string, limit int) ([]*memory.AuditEntry, error) { + rows, err := d.queries.ListAuditEntriesByAction(ctx, memsqlc.ListAuditEntriesByActionParams{ + AgentID: agentID, + Action: action, + Lim: int64(limit), + }) + if err != nil { + return nil, err + } + entries := make([]*memory.AuditEntry, len(rows)) + for i, row := range rows { + entries[i] = sqlcAuditToMemory(row) + } + return entries, nil +} + +func (d *LibSQLDelegate) CountAuditEntries(ctx context.Context, agentID string) (int, error) { + count, err := d.queries.CountAuditEntries(ctx, memsqlc.CountAuditEntriesParams{ + AgentID: agentID, + }) + return int(count), err +} + +func (d *LibSQLDelegate) ListAuditEntriesBySession(ctx context.Context, agentID, sessionKey string, limit int) ([]*memory.AuditEntry, error) { + rows, err := d.queries.ListAuditEntriesBySession(ctx, memsqlc.ListAuditEntriesBySessionParams{ + AgentID: agentID, + SessionKey: sessionKey, + Lim: int64(limit), + }) + if err != nil { + return nil, err + } + entries := make([]*memory.AuditEntry, len(rows)) + for i, row := range rows { + entries[i] = sqlcAuditToMemory(row) + } + return entries, nil +} + +func (d *LibSQLDelegate) PruneOldAuditEntries(ctx context.Context, agentID string, before time.Time) error { + return d.queries.PruneOldAuditEntries(ctx, memsqlc.PruneOldAuditEntriesParams{ + AgentID: agentID, + Before: before, + }) +} + +func (d *LibSQLDelegate) CountAuditEntriesByAction(ctx context.Context, agentID, action string) (int, error) { + count, err := d.queries.CountAuditEntriesByAction(ctx, memsqlc.CountAuditEntriesByActionParams{ + AgentID: agentID, + Action: action, + }) + return int(count), err +} + // --- Conversion helpers --- func sqlcRecallToMemory(row memsqlc.RecallItem) *memory.RecallItem { @@ -442,3 +671,40 @@ func sqlcChunkToMemory(row memsqlc.ArchivalChunk) *memory.ArchivalChunk { CreatedAt: row.CreatedAt, } } + +func sqlcDocToMemory(row memsqlc.AgentDocument) *memory.AgentDocument { + return &memory.AgentDocument{ + ID: row.ID, + AgentID: row.AgentID, + Name: row.Name, + Category: row.Category, + Content: row.Content, + Version: int(row.Version), + IsActive: row.IsActive, + CreatedAt: row.CreatedAt, + UpdatedAt: row.UpdatedAt, + } +} + +func sqlcAuditToMemory(row memsqlc.AgentAuditLog) *memory.AuditEntry { + entry := &memory.AuditEntry{ + ID: row.ID, + AgentID: row.AgentID, + SessionKey: row.SessionKey, + Action: row.Action, + Target: row.Target, + CreatedAt: row.CreatedAt, + } + if row.Input != nil { + entry.Input = *row.Input + } + if row.Output != nil { + entry.Output = *row.Output + } + if row.DurationMs != nil { + entry.DurationMS = int(*row.DurationMs) + } + return entry +} + +func ptrInt64(v int64) *int64 { return &v } diff --git a/pkg/memory/delegate/sqlite_audit_test.go b/pkg/memory/delegate/sqlite_audit_test.go new file mode 100644 index 000000000..125c1b115 --- /dev/null +++ b/pkg/memory/delegate/sqlite_audit_test.go @@ -0,0 +1,281 @@ +package delegate + +import ( + "context" + "testing" + "time" + + "github.com/sipeed/picoclaw/pkg/ids" + "github.com/sipeed/picoclaw/pkg/memory" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func makeAuditEntry(agentID, sessionKey, action, target string) *memory.AuditEntry { + return &memory.AuditEntry{ + ID: ids.New(), + AgentID: agentID, + SessionKey: sessionKey, + Action: action, + Target: target, + Input: `{"arg":"val"}`, + Output: `{"result":"ok"}`, + DurationMS: 42, + } +} + +func TestLibSQLDelegate_InsertAuditEntry(t *testing.T) { + tests := []struct { + name string + entry *memory.AuditEntry + }{ + { + name: "insert tool_call entry", + entry: makeAuditEntry("a1", "sess-1", "tool_call", "exec"), + }, + { + name: "insert memory_write entry", + entry: makeAuditEntry("a1", "sess-1", "memory_write", "working_context"), + }, + { + name: "insert entry with empty optional fields", + entry: &memory.AuditEntry{ + ID: ids.New(), + AgentID: "a1", + SessionKey: "sess-1", + Action: "state_change", + Target: "", + Input: "", + Output: "", + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + d := newTestDelegate(t) + ctx := context.Background() + require.NoError(t, d.InsertAuditEntry(ctx, tt.entry)) + + count, err := d.CountAuditEntries(ctx, tt.entry.AgentID) + require.NoError(t, err) + assert.Equal(t, 1, count) + }) + } +} + +func TestLibSQLDelegate_ListAuditEntries(t *testing.T) { + tests := []struct { + name string + setup func(t *testing.T, d *LibSQLDelegate, ctx context.Context) + agentID string + limit int + wantLen int + }{ + { + name: "empty returns empty", + agentID: "a1", + limit: 10, + wantLen: 0, + }, + { + name: "returns all for agent", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + require.NoError(t, d.InsertAuditEntry(ctx, makeAuditEntry("a1", "s1", "tool_call", "exec"))) + require.NoError(t, d.InsertAuditEntry(ctx, makeAuditEntry("a1", "s1", "memory_write", "kv"))) + }, + agentID: "a1", + limit: 10, + wantLen: 2, + }, + { + name: "agent isolation", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + require.NoError(t, d.InsertAuditEntry(ctx, makeAuditEntry("a1", "s1", "tool_call", "t1"))) + require.NoError(t, d.InsertAuditEntry(ctx, makeAuditEntry("a2", "s1", "tool_call", "t2"))) + }, + agentID: "a1", + limit: 10, + wantLen: 1, + }, + { + name: "respects limit", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + for i := 0; i < 10; i++ { + require.NoError(t, d.InsertAuditEntry(ctx, makeAuditEntry("a1", "s1", "tool_call", kvKey("t", i)))) + } + }, + agentID: "a1", + limit: 3, + wantLen: 3, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + d := newTestDelegate(t) + ctx := context.Background() + if tt.setup != nil { + tt.setup(t, d, ctx) + } + entries, err := d.ListAuditEntries(ctx, tt.agentID, tt.limit) + require.NoError(t, err) + assert.Len(t, entries, tt.wantLen) + }) + } +} + +func TestLibSQLDelegate_ListAuditEntriesByAction(t *testing.T) { + tests := []struct { + name string + setup func(t *testing.T, d *LibSQLDelegate, ctx context.Context) + agentID string + action string + limit int + wantLen int + wantAction string + }{ + { + name: "no matches returns empty", + agentID: "a1", + action: "tool_call", + limit: 10, + wantLen: 0, + }, + { + name: "filters by action", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + require.NoError(t, d.InsertAuditEntry(ctx, makeAuditEntry("a1", "s1", "tool_call", "exec"))) + require.NoError(t, d.InsertAuditEntry(ctx, makeAuditEntry("a1", "s1", "memory_write", "kv"))) + require.NoError(t, d.InsertAuditEntry(ctx, makeAuditEntry("a1", "s1", "tool_call", "read"))) + }, + agentID: "a1", + action: "tool_call", + limit: 10, + wantLen: 2, + wantAction: "tool_call", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + d := newTestDelegate(t) + ctx := context.Background() + if tt.setup != nil { + tt.setup(t, d, ctx) + } + entries, err := d.ListAuditEntriesByAction(ctx, tt.agentID, tt.action, tt.limit) + require.NoError(t, err) + assert.Len(t, entries, tt.wantLen) + if tt.wantAction != "" { + for _, e := range entries { + assert.Equal(t, tt.wantAction, e.Action) + } + } + }) + } +} + +func TestLibSQLDelegate_ListAuditEntriesBySession(t *testing.T) { + tests := []struct { + name string + setup func(t *testing.T, d *LibSQLDelegate, ctx context.Context) + agentID string + sessionKey string + limit int + wantLen int + }{ + { + name: "empty returns empty", + agentID: "a1", + sessionKey: "sess-1", + limit: 10, + wantLen: 0, + }, + { + name: "filters by session", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + require.NoError(t, d.InsertAuditEntry(ctx, makeAuditEntry("a1", "sess-A", "tool_call", "t1"))) + require.NoError(t, d.InsertAuditEntry(ctx, makeAuditEntry("a1", "sess-A", "tool_call", "t2"))) + require.NoError(t, d.InsertAuditEntry(ctx, makeAuditEntry("a1", "sess-B", "tool_call", "t3"))) + }, + agentID: "a1", + sessionKey: "sess-A", + limit: 10, + wantLen: 2, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + d := newTestDelegate(t) + ctx := context.Background() + if tt.setup != nil { + tt.setup(t, d, ctx) + } + entries, err := d.ListAuditEntriesBySession(ctx, tt.agentID, tt.sessionKey, tt.limit) + require.NoError(t, err) + assert.Len(t, entries, tt.wantLen) + }) + } +} + +func TestLibSQLDelegate_PruneOldAuditEntries(t *testing.T) { + d := newTestDelegate(t) + ctx := context.Background() + + require.NoError(t, d.InsertAuditEntry(ctx, makeAuditEntry("a1", "s1", "tool_call", "t1"))) + require.NoError(t, d.InsertAuditEntry(ctx, makeAuditEntry("a1", "s1", "tool_call", "t2"))) + + count, err := d.CountAuditEntries(ctx, "a1") + require.NoError(t, err) + assert.Equal(t, 2, count) + + // Prune entries created before "now + 1 minute" (should remove all) + require.NoError(t, d.PruneOldAuditEntries(ctx, "a1", time.Now().Add(time.Minute))) + + count, err = d.CountAuditEntries(ctx, "a1") + require.NoError(t, err) + assert.Equal(t, 0, count) +} + +func TestLibSQLDelegate_CountAuditEntriesByAction(t *testing.T) { + tests := []struct { + name string + setup func(t *testing.T, d *LibSQLDelegate, ctx context.Context) + agentID string + action string + want int + }{ + { + name: "no entries returns zero", + agentID: "a1", + action: "tool_call", + want: 0, + }, + { + name: "counts only matching action", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + require.NoError(t, d.InsertAuditEntry(ctx, makeAuditEntry("a1", "s1", "tool_call", "t1"))) + require.NoError(t, d.InsertAuditEntry(ctx, makeAuditEntry("a1", "s1", "memory_write", "kv"))) + require.NoError(t, d.InsertAuditEntry(ctx, makeAuditEntry("a1", "s1", "tool_call", "t2"))) + }, + agentID: "a1", + action: "tool_call", + want: 2, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + d := newTestDelegate(t) + ctx := context.Background() + if tt.setup != nil { + tt.setup(t, d, ctx) + } + count, err := d.CountAuditEntriesByAction(ctx, tt.agentID, tt.action) + require.NoError(t, err) + assert.Equal(t, tt.want, count) + }) + } +} diff --git a/pkg/memory/delegate/sqlite_doc_test.go b/pkg/memory/delegate/sqlite_doc_test.go new file mode 100644 index 000000000..46811dbd5 --- /dev/null +++ b/pkg/memory/delegate/sqlite_doc_test.go @@ -0,0 +1,330 @@ +package delegate + +import ( + "context" + "testing" + + "github.com/sipeed/picoclaw/pkg/ids" + "github.com/sipeed/picoclaw/pkg/memory" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func makeDoc(agentID, name, category, content string) *memory.AgentDocument { + return &memory.AgentDocument{ + ID: ids.New(), + AgentID: agentID, + Name: name, + Category: category, + Content: content, + } +} + +func TestLibSQLDelegate_GetDocument(t *testing.T) { + tests := []struct { + name string + setup func(t *testing.T, d *LibSQLDelegate, ctx context.Context) + agentID string + docName string + wantNil bool + wantContent string + }{ + { + name: "missing document returns nil", + agentID: "a1", + docName: "nonexistent", + wantNil: true, + }, + { + name: "returns stored document", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + require.NoError(t, d.UpsertDocument(ctx, makeDoc("a1", "README", "core", "hello world"))) + }, + agentID: "a1", + docName: "README", + wantContent: "hello world", + }, + { + name: "agent isolation — different agent sees nil", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + require.NoError(t, d.UpsertDocument(ctx, makeDoc("a1", "SECRET", "core", "mine"))) + }, + agentID: "a2", + docName: "SECRET", + wantNil: true, + }, + { + name: "document with special characters in content", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + require.NoError(t, d.UpsertDocument(ctx, makeDoc("a1", "config.json", "config", `{"key":"val","nested":{"a":1}}`))) + }, + agentID: "a1", + docName: "config.json", + wantContent: `{"key":"val","nested":{"a":1}}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + d := newTestDelegate(t) + ctx := context.Background() + if tt.setup != nil { + tt.setup(t, d, ctx) + } + doc, err := d.GetDocument(ctx, tt.agentID, tt.docName) + require.NoError(t, err) + if tt.wantNil { + assert.Nil(t, doc) + return + } + require.NotNil(t, doc) + assert.Equal(t, tt.wantContent, doc.Content) + assert.Equal(t, tt.agentID, doc.AgentID) + assert.Equal(t, tt.docName, doc.Name) + }) + } +} + +func TestLibSQLDelegate_UpsertDocument(t *testing.T) { + tests := []struct { + name string + ops []*memory.AgentDocument + agentID string + docName string + wantContent string + }{ + { + name: "insert new document", + ops: []*memory.AgentDocument{makeDoc("a1", "doc1", "cat", "first")}, + agentID: "a1", + docName: "doc1", + wantContent: "first", + }, + { + name: "overwrite existing document content", + ops: []*memory.AgentDocument{ + makeDoc("a1", "doc1", "cat", "original"), + {ID: ids.New(), AgentID: "a1", Name: "doc1", Category: "cat", Content: "updated"}, + }, + agentID: "a1", + docName: "doc1", + wantContent: "updated", + }, + { + name: "same name different agents are independent", + ops: []*memory.AgentDocument{ + makeDoc("a1", "shared", "cat", "from-a1"), + makeDoc("a2", "shared", "cat", "from-a2"), + }, + agentID: "a1", + docName: "shared", + wantContent: "from-a1", + }, + { + name: "empty content is valid", + ops: []*memory.AgentDocument{makeDoc("a1", "empty", "cat", "")}, + agentID: "a1", + docName: "empty", + wantContent: "", + }, + { + name: "large content roundtrip", + ops: []*memory.AgentDocument{makeDoc("a1", "big", "cat", longValue(8192))}, + agentID: "a1", + docName: "big", + wantContent: longValue(8192), + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + d := newTestDelegate(t) + ctx := context.Background() + for _, doc := range tt.ops { + require.NoError(t, d.UpsertDocument(ctx, doc)) + } + got, err := d.GetDocument(ctx, tt.agentID, tt.docName) + require.NoError(t, err) + require.NotNil(t, got) + assert.Equal(t, tt.wantContent, got.Content) + }) + } +} + +func TestLibSQLDelegate_DeleteDocument(t *testing.T) { + tests := []struct { + name string + setup func(t *testing.T, d *LibSQLDelegate, ctx context.Context) + agentID string + docName string + }{ + { + name: "delete nonexistent is idempotent", + agentID: "a1", + docName: "nope", + }, + { + name: "delete existing document", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + require.NoError(t, d.UpsertDocument(ctx, makeDoc("a1", "doomed", "cat", "bye"))) + }, + agentID: "a1", + docName: "doomed", + }, + { + name: "delete only affects target agent", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + require.NoError(t, d.UpsertDocument(ctx, makeDoc("a1", "shared", "cat", "a1-doc"))) + require.NoError(t, d.UpsertDocument(ctx, makeDoc("a2", "shared", "cat", "a2-doc"))) + }, + agentID: "a1", + docName: "shared", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + d := newTestDelegate(t) + ctx := context.Background() + if tt.setup != nil { + tt.setup(t, d, ctx) + } + require.NoError(t, d.DeleteDocument(ctx, tt.agentID, tt.docName)) + + got, err := d.GetDocument(ctx, tt.agentID, tt.docName) + require.NoError(t, err) + assert.Nil(t, got, "document should be gone after delete") + + if tt.name == "delete only affects target agent" { + other, err := d.GetDocument(ctx, "a2", "shared") + require.NoError(t, err) + require.NotNil(t, other, "other agent's document must survive") + assert.Equal(t, "a2-doc", other.Content) + } + }) + } +} + +func TestLibSQLDelegate_ListDocumentsByCategory(t *testing.T) { + tests := []struct { + name string + setup func(t *testing.T, d *LibSQLDelegate, ctx context.Context) + agentID string + category string + wantLen int + wantName string // first result name, if any + }{ + { + name: "empty store returns empty slice", + agentID: "a1", + category: "core", + wantLen: 0, + }, + { + name: "filters by category", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + require.NoError(t, d.UpsertDocument(ctx, makeDoc("a1", "d1", "core", "one"))) + require.NoError(t, d.UpsertDocument(ctx, makeDoc("a1", "d2", "core", "two"))) + require.NoError(t, d.UpsertDocument(ctx, makeDoc("a1", "d3", "config", "three"))) + }, + agentID: "a1", + category: "core", + wantLen: 2, + }, + { + name: "agent isolation in category listing", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + require.NoError(t, d.UpsertDocument(ctx, makeDoc("a1", "d1", "core", "a1-val"))) + require.NoError(t, d.UpsertDocument(ctx, makeDoc("a2", "d2", "core", "a2-val"))) + }, + agentID: "a1", + category: "core", + wantLen: 1, + wantName: "d1", + }, + { + name: "no match returns empty slice", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + require.NoError(t, d.UpsertDocument(ctx, makeDoc("a1", "d1", "core", "val"))) + }, + agentID: "a1", + category: "nonexistent", + wantLen: 0, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + d := newTestDelegate(t) + ctx := context.Background() + if tt.setup != nil { + tt.setup(t, d, ctx) + } + docs, err := d.ListDocumentsByCategory(ctx, tt.agentID, tt.category) + require.NoError(t, err) + assert.Len(t, docs, tt.wantLen) + if tt.wantName != "" && len(docs) > 0 { + assert.Equal(t, tt.wantName, docs[0].Name) + } + }) + } +} + +func TestLibSQLDelegate_ListAllDocuments(t *testing.T) { + tests := []struct { + name string + setup func(t *testing.T, d *LibSQLDelegate, ctx context.Context) + agentID string + wantLen int + }{ + { + name: "empty store returns empty slice", + agentID: "a1", + wantLen: 0, + }, + { + name: "returns all documents for agent", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + require.NoError(t, d.UpsertDocument(ctx, makeDoc("a1", "d1", "core", "one"))) + require.NoError(t, d.UpsertDocument(ctx, makeDoc("a1", "d2", "config", "two"))) + require.NoError(t, d.UpsertDocument(ctx, makeDoc("a1", "d3", "skills", "three"))) + }, + agentID: "a1", + wantLen: 3, + }, + { + name: "agent isolation — counts only own docs", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + require.NoError(t, d.UpsertDocument(ctx, makeDoc("a1", "d1", "core", "a1"))) + require.NoError(t, d.UpsertDocument(ctx, makeDoc("a2", "d2", "core", "a2"))) + require.NoError(t, d.UpsertDocument(ctx, makeDoc("a2", "d3", "core", "a2b"))) + }, + agentID: "a1", + wantLen: 1, + }, + { + name: "multiple categories all returned", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + cats := []string{"core", "config", "skills", "templates"} + for i, c := range cats { + require.NoError(t, d.UpsertDocument(ctx, makeDoc("a1", kvKey("doc-", i), c, "val"))) + } + }, + agentID: "a1", + wantLen: 4, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + d := newTestDelegate(t) + ctx := context.Background() + if tt.setup != nil { + tt.setup(t, d, ctx) + } + docs, err := d.ListAllDocuments(ctx, tt.agentID) + require.NoError(t, err) + assert.Len(t, docs, tt.wantLen) + }) + } +} diff --git a/pkg/memory/delegate/sqlite_integration_test.go b/pkg/memory/delegate/sqlite_integration_test.go new file mode 100644 index 000000000..98c2b59ef --- /dev/null +++ b/pkg/memory/delegate/sqlite_integration_test.go @@ -0,0 +1,227 @@ +package delegate + +import ( + "context" + "encoding/json" + "testing" + + "github.com/sipeed/picoclaw/pkg/ids" + "github.com/sipeed/picoclaw/pkg/memory" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestCronKVBackend_Roundtrip(t *testing.T) { + d := newTestDelegate(t) + ctx := context.Background() + agentID := "picoclaw" + kvKey := "cron:store" + + type cronStore struct { + Version int `json:"version"` + Jobs []struct { + ID string `json:"id"` + Name string `json:"name"` + Enabled bool `json:"enabled"` + } `json:"jobs"` + } + + store := cronStore{ + Version: 1, + Jobs: []struct { + ID string `json:"id"` + Name string `json:"name"` + Enabled bool `json:"enabled"` + }{ + {ID: "job-1", Name: "daily report", Enabled: true}, + {ID: "job-2", Name: "weekly backup", Enabled: false}, + }, + } + + data, err := json.Marshal(store) + require.NoError(t, err) + + require.NoError(t, d.UpsertKV(ctx, agentID, kvKey, string(data))) + + raw, err := d.GetKV(ctx, agentID, kvKey) + require.NoError(t, err) + require.NotEmpty(t, raw) + + var loaded cronStore + require.NoError(t, json.Unmarshal([]byte(raw), &loaded)) + + assert.Equal(t, 1, loaded.Version) + assert.Len(t, loaded.Jobs, 2) + assert.Equal(t, "daily report", loaded.Jobs[0].Name) + assert.True(t, loaded.Jobs[0].Enabled) + assert.Equal(t, "weekly backup", loaded.Jobs[1].Name) + assert.False(t, loaded.Jobs[1].Enabled) +} + +func TestCronKVBackend_UpdatePreservesShape(t *testing.T) { + d := newTestDelegate(t) + ctx := context.Background() + agentID := "picoclaw" + kvKey := "cron:store" + + v1 := `{"version":1,"jobs":[{"id":"j1","name":"test","enabled":true}]}` + require.NoError(t, d.UpsertKV(ctx, agentID, kvKey, v1)) + + v2 := `{"version":1,"jobs":[{"id":"j1","name":"test","enabled":true},{"id":"j2","name":"new","enabled":true}]}` + require.NoError(t, d.UpsertKV(ctx, agentID, kvKey, v2)) + + raw, err := d.GetKV(ctx, agentID, kvKey) + require.NoError(t, err) + assert.Equal(t, v2, raw) +} + +func TestCronKVBackend_PrefixScan(t *testing.T) { + d := newTestDelegate(t) + ctx := context.Background() + agentID := "picoclaw" + + require.NoError(t, d.UpsertKV(ctx, agentID, "cron:store", "{}")) + require.NoError(t, d.UpsertKV(ctx, agentID, "cron:lock", "held")) + require.NoError(t, d.UpsertKV(ctx, agentID, "focus:sess1", "{}")) + + result, err := d.ListKVByPrefix(ctx, agentID, "cron:", 10) + require.NoError(t, err) + assert.Len(t, result, 2) + assert.Contains(t, result, "cron:store") + assert.Contains(t, result, "cron:lock") +} + +func TestEndToEnd_SessionAndAuditFlow(t *testing.T) { + d := newTestDelegate(t) + ctx := context.Background() + agentID := "a1" + sessionKey := "sess-integration" + + // 1. Insert session messages + require.NoError(t, d.InsertSessionMessage(ctx, agentID, sessionKey, "user", "What is the weather?")) + require.NoError(t, d.InsertSessionMessage(ctx, agentID, sessionKey, "assistant", "Let me check...")) + require.NoError(t, d.InsertSessionMessage(ctx, agentID, sessionKey, "tool", `{"temp":72,"unit":"F"}`)) + require.NoError(t, d.InsertSessionMessage(ctx, agentID, sessionKey, "assistant", "It's 72F.")) + + // 2. Verify session message count + count, err := d.CountSessionMessages(ctx, agentID, sessionKey) + require.NoError(t, err) + assert.Equal(t, int64(4), count) + + // 3. Verify message ordering (ASC) + msgs, err := d.ListSessionMessages(ctx, agentID, sessionKey, "", 50) + require.NoError(t, err) + require.Len(t, msgs, 4) + assert.Equal(t, "user", msgs[0].Role) + assert.Equal(t, "What is the weather?", msgs[0].Content) + assert.Equal(t, "assistant", msgs[3].Role) + + // 4. Insert audit entries for the tool call + auditEntry := &memory.AuditEntry{ + ID: ids.New(), + AgentID: agentID, + SessionKey: sessionKey, + Action: "tool_call", + Target: "weather_api", + Input: `{"location":"here"}`, + Output: `{"temp":72}`, + DurationMS: 150, + } + require.NoError(t, d.InsertAuditEntry(ctx, auditEntry)) + + // 5. Verify audit log by session + auditEntries, err := d.ListAuditEntriesBySession(ctx, agentID, sessionKey, 10) + require.NoError(t, err) + require.Len(t, auditEntries, 1) + assert.Equal(t, "tool_call", auditEntries[0].Action) + assert.Equal(t, "weather_api", auditEntries[0].Target) + + // 6. Store KV state (e.g. focus checkpoint) + require.NoError(t, d.UpsertKV(ctx, agentID, "focus:"+sessionKey, `{"topic":"weather query","checkpoint_index":2}`)) + + kvVal, err := d.GetKV(ctx, agentID, "focus:"+sessionKey) + require.NoError(t, err) + assert.Contains(t, kvVal, "weather query") + + // 7. Store a document + doc := &memory.AgentDocument{ + ID: ids.New(), + AgentID: agentID, + Name: "AGENT.md", + Category: "core", + Content: "# Agent Identity\nI am picoclaw.", + } + require.NoError(t, d.UpsertDocument(ctx, doc)) + + loadedDoc, err := d.GetDocument(ctx, agentID, "AGENT.md") + require.NoError(t, err) + require.NotNil(t, loadedDoc) + assert.Equal(t, "# Agent Identity\nI am picoclaw.", loadedDoc.Content) + + // 8. Verify cross-table isolation: different session sees nothing + otherMsgs, err := d.ListSessionMessages(ctx, agentID, "sess-other", "", 50) + require.NoError(t, err) + assert.Empty(t, otherMsgs) + + otherAudit, err := d.ListAuditEntriesBySession(ctx, agentID, "sess-other", 10) + require.NoError(t, err) + assert.Empty(t, otherAudit) +} + +func TestEndToEnd_WorkingContextAndRecallRoundtrip(t *testing.T) { + d := newTestDelegate(t) + ctx := context.Background() + agentID := "a1" + sessionKey := "sess-wc" + + // 1. Upsert working context + require.NoError(t, d.UpsertWorkingContext(ctx, agentID, sessionKey, "Initial system prompt state")) + + wc, err := d.GetWorkingContext(ctx, agentID, sessionKey) + require.NoError(t, err) + require.NotNil(t, wc) + assert.Equal(t, "Initial system prompt state", wc.Content) + + // 2. Insert recall items + item := &memory.RecallItem{ + ID: ids.New(), + AgentID: agentID, + SessionKey: sessionKey, + Role: "user", + Sector: memory.SectorEpisodic, + Importance: 0.8, + Salience: 0.7, + DecayRate: 0.01, + Content: "Important user preference", + Tags: "preference", + } + require.NoError(t, d.InsertRecallItem(ctx, item)) + + // 3. Count recall items + recallCount, err := d.CountRecallItems(ctx, agentID, sessionKey) + require.NoError(t, err) + assert.Equal(t, 1, recallCount) + + // 4. Update working context + require.NoError(t, d.UpsertWorkingContext(ctx, agentID, sessionKey, "Updated with preference awareness")) + + wc, err = d.GetWorkingContext(ctx, agentID, sessionKey) + require.NoError(t, err) + assert.Equal(t, "Updated with preference awareness", wc.Content) + + // 5. Insert a summary + summary := &memory.MemorySummary{ + ID: ids.New(), + AgentID: agentID, + SessionKey: sessionKey, + Content: "User discussed preferences. Key info captured.", + FromMsgIdx: 0, + ToMsgIdx: 5, + } + require.NoError(t, d.InsertSummary(ctx, summary)) + + summaries, err := d.ListSummaries(ctx, agentID, sessionKey, 10) + require.NoError(t, err) + require.Len(t, summaries, 1) + assert.Equal(t, "User discussed preferences. Key info captured.", summaries[0].Content) +} diff --git a/pkg/memory/delegate/sqlite_kv_test.go b/pkg/memory/delegate/sqlite_kv_test.go new file mode 100644 index 000000000..b112f341a --- /dev/null +++ b/pkg/memory/delegate/sqlite_kv_test.go @@ -0,0 +1,330 @@ +package delegate + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestLibSQLDelegate_GetKV(t *testing.T) { + tests := []struct { + name string + setup func(t *testing.T, d *LibSQLDelegate, ctx context.Context) + agentID string + key string + wantValue string + wantErr bool + }{ + { + name: "missing key returns empty string", + agentID: "agent-1", + key: "nonexistent", + wantValue: "", + }, + { + name: "returns stored value", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + require.NoError(t, d.UpsertKV(ctx, "agent-1", "greeting", "hello")) + }, + agentID: "agent-1", + key: "greeting", + wantValue: "hello", + }, + { + name: "agent isolation — different agent sees empty", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + require.NoError(t, d.UpsertKV(ctx, "agent-A", "secret", "mine")) + }, + agentID: "agent-B", + key: "secret", + wantValue: "", + }, + { + name: "empty value is valid", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + require.NoError(t, d.UpsertKV(ctx, "agent-1", "blank", "")) + }, + agentID: "agent-1", + key: "blank", + wantValue: "", + }, + { + name: "value with special characters", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + require.NoError(t, d.UpsertKV(ctx, "agent-1", "json", `{"key":"val","n":42}`)) + }, + agentID: "agent-1", + key: "json", + wantValue: `{"key":"val","n":42}`, + }, + { + name: "key with colons and slashes", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + require.NoError(t, d.UpsertKV(ctx, "agent-1", "focus:session/abc", "data")) + }, + agentID: "agent-1", + key: "focus:session/abc", + wantValue: "data", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + d := newTestDelegate(t) + ctx := context.Background() + if tt.setup != nil { + tt.setup(t, d, ctx) + } + got, err := d.GetKV(ctx, tt.agentID, tt.key) + if tt.wantErr { + assert.Error(t, err) + return + } + require.NoError(t, err) + assert.Equal(t, tt.wantValue, got) + }) + } +} + +func TestLibSQLDelegate_UpsertKV(t *testing.T) { + tests := []struct { + name string + ops []kvOp + agentID string + key string + want string + }{ + { + name: "insert new key", + ops: []kvOp{{agent: "a1", key: "k1", val: "v1"}}, + agentID: "a1", + key: "k1", + want: "v1", + }, + { + name: "overwrite existing key", + ops: []kvOp{ + {agent: "a1", key: "k1", val: "first"}, + {agent: "a1", key: "k1", val: "second"}, + }, + agentID: "a1", + key: "k1", + want: "second", + }, + { + name: "overwrite with empty", + ops: []kvOp{ + {agent: "a1", key: "k1", val: "nonempty"}, + {agent: "a1", key: "k1", val: ""}, + }, + agentID: "a1", + key: "k1", + want: "", + }, + { + name: "same key different agents are independent", + ops: []kvOp{ + {agent: "a1", key: "shared", val: "from-a1"}, + {agent: "a2", key: "shared", val: "from-a2"}, + }, + agentID: "a1", + key: "shared", + want: "from-a1", + }, + { + name: "large value roundtrip", + ops: []kvOp{ + {agent: "a1", key: "big", val: longValue(4096)}, + }, + agentID: "a1", + key: "big", + want: longValue(4096), + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + d := newTestDelegate(t) + ctx := context.Background() + for _, op := range tt.ops { + require.NoError(t, d.UpsertKV(ctx, op.agent, op.key, op.val)) + } + got, err := d.GetKV(ctx, tt.agentID, tt.key) + require.NoError(t, err) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestLibSQLDelegate_DeleteKV(t *testing.T) { + tests := []struct { + name string + setup func(t *testing.T, d *LibSQLDelegate, ctx context.Context) + agentID string + key string + wantErr bool + }{ + { + name: "delete nonexistent key is idempotent", + agentID: "a1", + key: "nope", + }, + { + name: "delete existing key", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + require.NoError(t, d.UpsertKV(ctx, "a1", "doomed", "bye")) + }, + agentID: "a1", + key: "doomed", + }, + { + name: "delete only affects target agent", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + require.NoError(t, d.UpsertKV(ctx, "a1", "shared", "a1-val")) + require.NoError(t, d.UpsertKV(ctx, "a2", "shared", "a2-val")) + }, + agentID: "a1", + key: "shared", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + d := newTestDelegate(t) + ctx := context.Background() + if tt.setup != nil { + tt.setup(t, d, ctx) + } + err := d.DeleteKV(ctx, tt.agentID, tt.key) + if tt.wantErr { + assert.Error(t, err) + return + } + require.NoError(t, err) + + got, err := d.GetKV(ctx, tt.agentID, tt.key) + require.NoError(t, err) + assert.Empty(t, got, "key should be gone after delete") + + if tt.name == "delete only affects target agent" { + other, err := d.GetKV(ctx, "a2", "shared") + require.NoError(t, err) + assert.Equal(t, "a2-val", other, "other agent's key must survive") + } + }) + } +} + +func TestLibSQLDelegate_ListKVByPrefix(t *testing.T) { + tests := []struct { + name string + setup func(t *testing.T, d *LibSQLDelegate, ctx context.Context) + agentID string + prefix string + limit int + want map[string]string + wantErr bool + }{ + { + name: "empty store returns empty map", + agentID: "a1", + prefix: "obs:", + limit: 10, + want: map[string]string{}, + }, + { + name: "matches prefix exactly", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + require.NoError(t, d.UpsertKV(ctx, "a1", "obs:1", "one")) + require.NoError(t, d.UpsertKV(ctx, "a1", "obs:2", "two")) + require.NoError(t, d.UpsertKV(ctx, "a1", "focus:x", "other")) + }, + agentID: "a1", + prefix: "obs:", + limit: 10, + want: map[string]string{"obs:1": "one", "obs:2": "two"}, + }, + { + name: "respects limit", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + for i := 0; i < 5; i++ { + require.NoError(t, d.UpsertKV(ctx, "a1", kvKey("item:", i), kvVal(i))) + } + }, + agentID: "a1", + prefix: "item:", + limit: 3, + want: nil, // checked via length only + }, + { + name: "agent isolation in prefix scan", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + require.NoError(t, d.UpsertKV(ctx, "a1", "ns:k1", "a1-val")) + require.NoError(t, d.UpsertKV(ctx, "a2", "ns:k2", "a2-val")) + }, + agentID: "a1", + prefix: "ns:", + limit: 10, + want: map[string]string{"ns:k1": "a1-val"}, + }, + { + name: "prefix with no trailing separator", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + require.NoError(t, d.UpsertKV(ctx, "a1", "abc", "one")) + require.NoError(t, d.UpsertKV(ctx, "a1", "abd", "two")) + require.NoError(t, d.UpsertKV(ctx, "a1", "xyz", "three")) + }, + agentID: "a1", + prefix: "ab", + limit: 10, + want: map[string]string{"abc": "one", "abd": "two"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + d := newTestDelegate(t) + ctx := context.Background() + if tt.setup != nil { + tt.setup(t, d, ctx) + } + got, err := d.ListKVByPrefix(ctx, tt.agentID, tt.prefix, tt.limit) + if tt.wantErr { + assert.Error(t, err) + return + } + require.NoError(t, err) + + if tt.name == "respects limit" { + assert.LessOrEqual(t, len(got), tt.limit) + return + } + + assert.Equal(t, tt.want, got) + }) + } +} + +// --- helpers --- + +type kvOp struct { + agent, key, val string +} + +func longValue(n int) string { + b := make([]byte, n) + for i := range b { + b[i] = 'A' + byte(i%26) + } + return string(b) +} + +func kvKey(prefix string, i int) string { + return prefix + string(rune('0'+i)) +} + +func kvVal(i int) string { + return string(rune('a'+i)) + "-val" +} diff --git a/pkg/memory/delegate/sqlite_session_msg_test.go b/pkg/memory/delegate/sqlite_session_msg_test.go new file mode 100644 index 000000000..6d5bc4be0 --- /dev/null +++ b/pkg/memory/delegate/sqlite_session_msg_test.go @@ -0,0 +1,269 @@ +package delegate + +import ( + "context" + "testing" + "time" + + "github.com/sipeed/picoclaw/pkg/ids" + "github.com/sipeed/picoclaw/pkg/memory" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func makeRecallItem(agentID, sessionKey, role, content, tags string) *memory.RecallItem { + now := time.Now() + return &memory.RecallItem{ + ID: ids.New(), + AgentID: agentID, + SessionKey: sessionKey, + Role: role, + Sector: memory.SectorEpisodic, + Importance: 0.5, + Salience: 0.5, + DecayRate: 0.01, + Content: content, + Tags: tags, + CreatedAt: now, + UpdatedAt: now, + } +} + +func TestLibSQLDelegate_InsertSessionMessage(t *testing.T) { + tests := []struct { + name string + agentID string + sessionKey string + role string + content string + }{ + { + name: "insert user message", + agentID: "a1", + sessionKey: "sess-1", + role: "user", + content: "Hello, agent!", + }, + { + name: "insert assistant message", + agentID: "a1", + sessionKey: "sess-1", + role: "assistant", + content: "Hello, human!", + }, + { + name: "insert tool message", + agentID: "a1", + sessionKey: "sess-1", + role: "tool", + content: `{"result":"ok"}`, + }, + { + name: "empty content is valid", + agentID: "a1", + sessionKey: "sess-1", + role: "system", + content: "", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + d := newTestDelegate(t) + ctx := context.Background() + + require.NoError(t, d.InsertSessionMessage(ctx, tt.agentID, tt.sessionKey, tt.role, tt.content)) + + count, err := d.CountSessionMessages(ctx, tt.agentID, tt.sessionKey) + require.NoError(t, err) + assert.Equal(t, int64(1), count) + }) + } +} + +func TestLibSQLDelegate_ListSessionMessages(t *testing.T) { + tests := []struct { + name string + setup func(t *testing.T, d *LibSQLDelegate, ctx context.Context) + agentID string + sessionKey string + role string + limit int + wantLen int + wantRole string + }{ + { + name: "empty session returns empty", + agentID: "a1", + sessionKey: "sess-empty", + role: "", + limit: 50, + wantLen: 0, + }, + { + name: "list all roles (role=empty string)", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-1", "user", "msg1")) + require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-1", "assistant", "msg2")) + require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-1", "tool", "msg3")) + }, + agentID: "a1", + sessionKey: "sess-1", + role: "", + limit: 50, + wantLen: 3, + }, + { + name: "filter by role=user", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-1", "user", "u1")) + require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-1", "assistant", "a1")) + require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-1", "user", "u2")) + }, + agentID: "a1", + sessionKey: "sess-1", + role: "user", + limit: 50, + wantLen: 2, + wantRole: "user", + }, + { + name: "session isolation", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-A", "user", "msgA")) + require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-B", "user", "msgB")) + }, + agentID: "a1", + sessionKey: "sess-A", + role: "", + limit: 50, + wantLen: 1, + }, + { + name: "agent isolation", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-1", "user", "from-a1")) + require.NoError(t, d.InsertSessionMessage(ctx, "a2", "sess-1", "user", "from-a2")) + }, + agentID: "a1", + sessionKey: "sess-1", + role: "", + limit: 50, + wantLen: 1, + }, + { + name: "respects limit", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + for i := 0; i < 10; i++ { + require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-1", "user", kvVal(i))) + } + }, + agentID: "a1", + sessionKey: "sess-1", + role: "", + limit: 5, + wantLen: 5, + }, + { + name: "ordered by created_at ASC", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-1", "user", "first")) + require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-1", "assistant", "second")) + require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-1", "user", "third")) + }, + agentID: "a1", + sessionKey: "sess-1", + role: "", + limit: 50, + wantLen: 3, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + d := newTestDelegate(t) + ctx := context.Background() + if tt.setup != nil { + tt.setup(t, d, ctx) + } + msgs, err := d.ListSessionMessages(ctx, tt.agentID, tt.sessionKey, tt.role, tt.limit) + require.NoError(t, err) + assert.Len(t, msgs, tt.wantLen) + + if tt.wantRole != "" { + for _, m := range msgs { + assert.Equal(t, tt.wantRole, m.Role) + } + } + + if tt.name == "ordered by created_at ASC" && len(msgs) == 3 { + assert.Equal(t, "first", msgs[0].Content) + assert.Equal(t, "second", msgs[1].Content) + assert.Equal(t, "third", msgs[2].Content) + } + }) + } +} + +func TestLibSQLDelegate_CountSessionMessages(t *testing.T) { + tests := []struct { + name string + setup func(t *testing.T, d *LibSQLDelegate, ctx context.Context) + agentID string + sessionKey string + want int64 + }{ + { + name: "empty session returns zero", + agentID: "a1", + sessionKey: "sess-empty", + want: 0, + }, + { + name: "counts session messages only", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-1", "user", "m1")) + require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-1", "assistant", "m2")) + }, + agentID: "a1", + sessionKey: "sess-1", + want: 2, + }, + { + name: "excludes non-session recall items", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-1", "user", "session-msg")) + // Insert a generic recall item via the standard method + item := makeRecallItem("a1", "sess-1", "user", "generic-item", "other-tag") + require.NoError(t, d.InsertRecallItem(ctx, item)) + }, + agentID: "a1", + sessionKey: "sess-1", + want: 1, + }, + { + name: "session isolation in count", + setup: func(t *testing.T, d *LibSQLDelegate, ctx context.Context) { + require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-A", "user", "m1")) + require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-A", "user", "m2")) + require.NoError(t, d.InsertSessionMessage(ctx, "a1", "sess-B", "user", "m3")) + }, + agentID: "a1", + sessionKey: "sess-A", + want: 2, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + d := newTestDelegate(t) + ctx := context.Background() + if tt.setup != nil { + tt.setup(t, d, ctx) + } + count, err := d.CountSessionMessages(ctx, tt.agentID, tt.sessionKey) + require.NoError(t, err) + assert.Equal(t, tt.want, count) + }) + } +} diff --git a/pkg/memory/delegate/sqlite_test.go b/pkg/memory/delegate/sqlite_test.go index 9957cfbd1..712ea60ff 100644 --- a/pkg/memory/delegate/sqlite_test.go +++ b/pkg/memory/delegate/sqlite_test.go @@ -564,3 +564,141 @@ func TestEmbeddingValueScanRoundTrip(t *testing.T) { } } } + +func TestIntegration_FullStackNoDisk(t *testing.T) { + d := newTestDelegate(t) + ctx := context.Background() + agentID := "integration-agent" + + t.Run("KV_Store", func(t *testing.T) { + if err := d.UpsertKV(ctx, agentID, "last_channel", "telegram:42"); err != nil { + t.Fatalf("UpsertKV: %v", err) + } + val, err := d.GetKV(ctx, agentID, "last_channel") + if err != nil { + t.Fatalf("GetKV: %v", err) + } + if val != "telegram:42" { + t.Errorf("expected telegram:42, got %q", val) + } + + if err := d.UpsertKV(ctx, agentID, "cron:store", `{"version":1,"jobs":[]}`); err != nil { + t.Fatalf("UpsertKV cron: %v", err) + } + kvs, err := d.ListKVByPrefix(ctx, agentID, "cron:", 10) + if err != nil { + t.Fatalf("ListKVByPrefix: %v", err) + } + if len(kvs) != 1 { + t.Errorf("expected 1 cron KV, got %d", len(kvs)) + } + }) + + t.Run("Documents", func(t *testing.T) { + doc := &memory.AgentDocument{ + ID: ids.New(), + AgentID: agentID, + Name: "AGENTS.md", + Category: "bootstrap", + Content: "# Agent Config\nBootstrap content", + Version: 1, + IsActive: true, + } + if err := d.UpsertDocument(ctx, doc); err != nil { + t.Fatalf("UpsertDocument: %v", err) + } + + got, err := d.GetDocument(ctx, agentID, "AGENTS.md") + if err != nil { + t.Fatalf("GetDocument: %v", err) + } + if got.Content != doc.Content { + t.Errorf("content mismatch: %q vs %q", got.Content, doc.Content) + } + + docs, err := d.ListDocumentsByCategory(ctx, agentID, "bootstrap") + if err != nil { + t.Fatalf("ListDocumentsByCategory: %v", err) + } + if len(docs) != 1 { + t.Errorf("expected 1 bootstrap doc, got %d", len(docs)) + } + }) + + t.Run("Sessions_via_RecallItems", func(t *testing.T) { + for i, msg := range []struct{ role, content string }{ + {"user", "hello agent"}, + {"assistant", "hi there!"}, + {"user", "how are you?"}, + } { + item := &memory.RecallItem{ + ID: ids.New(), + AgentID: agentID, + SessionKey: "test-session", + Role: msg.role, + Sector: memory.SectorEpisodic, + Importance: 0.5, + Salience: 0.5, + Content: msg.content, + Tags: "session-message", + } + if err := d.InsertRecallItem(ctx, item); err != nil { + t.Fatalf("InsertRecallItem[%d]: %v", i, err) + } + } + + items, err := d.ListRecallItems(ctx, agentID, "test-session", 100, 0) + if err != nil { + t.Fatalf("ListRecallItems: %v", err) + } + if len(items) != 3 { + t.Errorf("expected 3 session items, got %d", len(items)) + } + }) + + t.Run("AuditLog", func(t *testing.T) { + entry := &memory.AuditEntry{ + ID: ids.New(), + AgentID: agentID, + SessionKey: "test-session", + Action: "tool_call", + Target: "exec", + Input: `{"command":"ls"}`, + } + if err := d.InsertAuditEntry(ctx, entry); err != nil { + t.Fatalf("InsertAuditEntry: %v", err) + } + + entries, err := d.ListAuditEntries(ctx, agentID, 10) + if err != nil { + t.Fatalf("ListAuditEntries: %v", err) + } + if len(entries) != 1 { + t.Errorf("expected 1 audit entry, got %d", len(entries)) + } + if entries[0].Target != "exec" { + t.Errorf("expected target 'exec', got %q", entries[0].Target) + } + + count, err := d.CountAuditEntries(ctx, agentID) + if err != nil { + t.Fatalf("CountAuditEntries: %v", err) + } + if count != 1 { + t.Errorf("expected count 1, got %d", count) + } + }) + + t.Run("WorkingContext", func(t *testing.T) { + if err := d.UpsertWorkingContext(ctx, agentID, "sess-1", "agent memory contents"); err != nil { + t.Fatalf("UpsertWorkingContext: %v", err) + } + wc, err := d.GetWorkingContext(ctx, agentID, "sess-1") + if err != nil { + t.Fatalf("GetWorkingContext: %v", err) + } + if wc == nil || wc.Content != "agent memory contents" { + t.Errorf("unexpected working context: %v", wc) + } + }) +} diff --git a/pkg/memory/memory.go b/pkg/memory/memory.go index 21ec1a381..df604c305 100644 --- a/pkg/memory/memory.go +++ b/pkg/memory/memory.go @@ -211,6 +211,34 @@ const ( PressureFlush PressureLevel = "flush" // > 85% ) +// --- Additional domain types --- + +// AgentDocument is a versioned named document stored in the database. +type AgentDocument struct { + ID ids.UUID + AgentID string + Name string + Category string + Content string + Version int + IsActive bool + CreatedAt time.Time + UpdatedAt time.Time +} + +// AuditEntry records a single agent action for the audit log. +type AuditEntry struct { + ID ids.UUID + AgentID string + SessionKey string + Action string // "tool_call", "memory_write", "doc_update", "state_change" + Target string // tool name, doc name, key name + Input string + Output string + DurationMS int + CreatedAt time.Time +} + // --- Delegate interface (backend) --- // MemoryDelegate is the pure storage backend for the memory system. @@ -260,6 +288,25 @@ type MemoryDelegate interface { CountRecallItems(ctx context.Context, agentID, sessionKey string) (int, error) CountArchivalChunks(ctx context.Context) (int, error) + // --- Key-Value Store --- + GetKV(ctx context.Context, agentID, key string) (string, error) + UpsertKV(ctx context.Context, agentID, key, value string) error + DeleteKV(ctx context.Context, agentID, key string) error + ListKVByPrefix(ctx context.Context, agentID, prefix string, limit int) (map[string]string, error) + + // --- Documents --- + GetDocument(ctx context.Context, agentID, name string) (*AgentDocument, error) + UpsertDocument(ctx context.Context, doc *AgentDocument) error + DeleteDocument(ctx context.Context, agentID, name string) error + ListDocumentsByCategory(ctx context.Context, agentID, category string) ([]*AgentDocument, error) + ListAllDocuments(ctx context.Context, agentID string) ([]*AgentDocument, error) + + // --- Audit Log --- + InsertAuditEntry(ctx context.Context, entry *AuditEntry) error + ListAuditEntries(ctx context.Context, agentID string, limit int) ([]*AuditEntry, error) + ListAuditEntriesByAction(ctx context.Context, agentID, action string, limit int) ([]*AuditEntry, error) + CountAuditEntries(ctx context.Context, agentID string) (int, error) + // --- Capability Detection --- HasVectorSearch() bool HasFTS() bool diff --git a/pkg/memory/migrations/004_agent_kv.go b/pkg/memory/migrations/004_agent_kv.go new file mode 100644 index 000000000..73bc2ac14 --- /dev/null +++ b/pkg/memory/migrations/004_agent_kv.go @@ -0,0 +1,35 @@ +package migrations + +import ( + "context" + "database/sql" + "fmt" + + "github.com/pressly/goose/v3" +) + +func init() { + goose.AddMigrationContext(up004AgentKV, down004AgentKV) +} + +func up004AgentKV(ctx context.Context, tx *sql.Tx) error { + _, err := tx.ExecContext(ctx, `CREATE TABLE IF NOT EXISTS agent_kv ( + agent_id TEXT NOT NULL, + key TEXT NOT NULL, + value TEXT NOT NULL DEFAULT '', + updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + PRIMARY KEY (agent_id, key) +)`) + if err != nil { + return fmt.Errorf("004_agent_kv up: %w", err) + } + return nil +} + +func down004AgentKV(ctx context.Context, tx *sql.Tx) error { + _, err := tx.ExecContext(ctx, `DROP TABLE IF EXISTS agent_kv`) + if err != nil { + return fmt.Errorf("004_agent_kv down: %w", err) + } + return nil +} diff --git a/pkg/memory/migrations/005_agent_documents.go b/pkg/memory/migrations/005_agent_documents.go new file mode 100644 index 000000000..9181cd071 --- /dev/null +++ b/pkg/memory/migrations/005_agent_documents.go @@ -0,0 +1,47 @@ +package migrations + +import ( + "context" + "database/sql" + "fmt" + + "github.com/pressly/goose/v3" +) + +func init() { + goose.AddMigrationContext(up005AgentDocuments, down005AgentDocuments) +} + +func up005AgentDocuments(ctx context.Context, tx *sql.Tx) error { + stmts := []string{ + `CREATE TABLE IF NOT EXISTS agent_documents ( + id BLOB PRIMARY KEY, + agent_id TEXT NOT NULL, + name TEXT NOT NULL, + category TEXT NOT NULL DEFAULT 'bootstrap', + content TEXT NOT NULL, + version INTEGER NOT NULL DEFAULT 1, + is_active INTEGER NOT NULL DEFAULT 1, + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + UNIQUE(agent_id, name) +)`, + `CREATE INDEX IF NOT EXISTS idx_docs_agent_cat ON agent_documents(agent_id, category)`, + `CREATE INDEX IF NOT EXISTS idx_docs_name ON agent_documents(name)`, + } + + for _, s := range stmts { + if _, err := tx.ExecContext(ctx, s); err != nil { + return fmt.Errorf("005_agent_documents up: %w\nSQL: %s", err, s) + } + } + return nil +} + +func down005AgentDocuments(ctx context.Context, tx *sql.Tx) error { + _, err := tx.ExecContext(ctx, `DROP TABLE IF EXISTS agent_documents`) + if err != nil { + return fmt.Errorf("005_agent_documents down: %w", err) + } + return nil +} diff --git a/pkg/memory/migrations/006_agent_audit_log.go b/pkg/memory/migrations/006_agent_audit_log.go new file mode 100644 index 000000000..ebc1be164 --- /dev/null +++ b/pkg/memory/migrations/006_agent_audit_log.go @@ -0,0 +1,46 @@ +package migrations + +import ( + "context" + "database/sql" + "fmt" + + "github.com/pressly/goose/v3" +) + +func init() { + goose.AddMigrationContext(up006AgentAuditLog, down006AgentAuditLog) +} + +func up006AgentAuditLog(ctx context.Context, tx *sql.Tx) error { + stmts := []string{ + `CREATE TABLE IF NOT EXISTS agent_audit_log ( + id BLOB PRIMARY KEY, + agent_id TEXT NOT NULL, + session_key TEXT NOT NULL DEFAULT '', + action TEXT NOT NULL, + target TEXT NOT NULL DEFAULT '', + input TEXT, + output TEXT, + duration_ms INTEGER, + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP +)`, + `CREATE INDEX IF NOT EXISTS idx_audit_agent_time ON agent_audit_log(agent_id, created_at DESC)`, + `CREATE INDEX IF NOT EXISTS idx_audit_action ON agent_audit_log(action)`, + } + + for _, s := range stmts { + if _, err := tx.ExecContext(ctx, s); err != nil { + return fmt.Errorf("006_agent_audit_log up: %w\nSQL: %s", err, s) + } + } + return nil +} + +func down006AgentAuditLog(ctx context.Context, tx *sql.Tx) error { + _, err := tx.ExecContext(ctx, `DROP TABLE IF EXISTS agent_audit_log`) + if err != nil { + return fmt.Errorf("006_agent_audit_log down: %w", err) + } + return nil +} diff --git a/pkg/memory/sqlc/agent_audit_log.sql.go b/pkg/memory/sqlc/agent_audit_log.sql.go new file mode 100644 index 000000000..ae8fccc0b --- /dev/null +++ b/pkg/memory/sqlc/agent_audit_log.sql.go @@ -0,0 +1,368 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.30.0 +// source: agent_audit_log.sql + +package sqlc + +import ( + "context" + "time" + + "github.com/sipeed/picoclaw/pkg/ids" +) + +const CountAuditEntries = `-- name: CountAuditEntries :one +SELECT COUNT(*) +FROM agent_audit_log +WHERE agent_id = ?1 +` + +type CountAuditEntriesParams struct { + AgentID string `json:"agent_id"` +} + +// CountAuditEntries +// +// SELECT COUNT(*) +// FROM agent_audit_log +// WHERE agent_id = ?1 +func (q *Queries) CountAuditEntries(ctx context.Context, arg CountAuditEntriesParams) (int64, error) { + row := q.db.QueryRowContext(ctx, CountAuditEntries, arg.AgentID) + var count int64 + err := row.Scan(&count) + return count, err +} + +const CountAuditEntriesByAction = `-- name: CountAuditEntriesByAction :one +SELECT COUNT(*) +FROM agent_audit_log +WHERE agent_id = ?1 + AND action = ?2 +` + +type CountAuditEntriesByActionParams struct { + AgentID string `json:"agent_id"` + Action string `json:"action"` +} + +// CountAuditEntriesByAction +// +// SELECT COUNT(*) +// FROM agent_audit_log +// WHERE agent_id = ?1 +// AND action = ?2 +func (q *Queries) CountAuditEntriesByAction(ctx context.Context, arg CountAuditEntriesByActionParams) (int64, error) { + row := q.db.QueryRowContext(ctx, CountAuditEntriesByAction, arg.AgentID, arg.Action) + var count int64 + err := row.Scan(&count) + return count, err +} + +const InsertAuditEntry = `-- name: InsertAuditEntry :exec +INSERT INTO agent_audit_log ( + id, + agent_id, + session_key, + action, + target, + input, + output, + duration_ms, + created_at + ) +VALUES ( + ?1, + ?2, + ?3, + ?4, + ?5, + ?6, + ?7, + ?8, + datetime('now') + ) +` + +type InsertAuditEntryParams struct { + ID ids.UUID `json:"id"` + AgentID string `json:"agent_id"` + SessionKey string `json:"session_key"` + Action string `json:"action"` + Target string `json:"target"` + Input *string `json:"input"` + Output *string `json:"output"` + DurationMs *int64 `json:"duration_ms"` +} + +// Agent Audit Log queries +// +// INSERT INTO agent_audit_log ( +// id, +// agent_id, +// session_key, +// action, +// target, +// input, +// output, +// duration_ms, +// created_at +// ) +// VALUES ( +// ?1, +// ?2, +// ?3, +// ?4, +// ?5, +// ?6, +// ?7, +// ?8, +// datetime('now') +// ) +func (q *Queries) InsertAuditEntry(ctx context.Context, arg InsertAuditEntryParams) error { + _, err := q.db.ExecContext(ctx, InsertAuditEntry, + arg.ID, + arg.AgentID, + arg.SessionKey, + arg.Action, + arg.Target, + arg.Input, + arg.Output, + arg.DurationMs, + ) + return err +} + +const ListAuditEntries = `-- name: ListAuditEntries :many +SELECT id, + agent_id, + session_key, + action, + target, + input, + output, + duration_ms, + created_at +FROM agent_audit_log +WHERE agent_id = ?1 +ORDER BY created_at DESC +LIMIT ?2 +` + +type ListAuditEntriesParams struct { + AgentID string `json:"agent_id"` + Lim int64 `json:"lim"` +} + +// ListAuditEntries +// +// SELECT id, +// agent_id, +// session_key, +// action, +// target, +// input, +// output, +// duration_ms, +// created_at +// FROM agent_audit_log +// WHERE agent_id = ?1 +// ORDER BY created_at DESC +// LIMIT ?2 +func (q *Queries) ListAuditEntries(ctx context.Context, arg ListAuditEntriesParams) ([]AgentAuditLog, error) { + rows, err := q.db.QueryContext(ctx, ListAuditEntries, arg.AgentID, arg.Lim) + if err != nil { + return nil, err + } + defer rows.Close() + items := []AgentAuditLog{} + for rows.Next() { + var i AgentAuditLog + if err := rows.Scan( + &i.ID, + &i.AgentID, + &i.SessionKey, + &i.Action, + &i.Target, + &i.Input, + &i.Output, + &i.DurationMs, + &i.CreatedAt, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const ListAuditEntriesByAction = `-- name: ListAuditEntriesByAction :many +SELECT id, + agent_id, + session_key, + action, + target, + input, + output, + duration_ms, + created_at +FROM agent_audit_log +WHERE agent_id = ?1 + AND action = ?2 +ORDER BY created_at DESC +LIMIT ?3 +` + +type ListAuditEntriesByActionParams struct { + AgentID string `json:"agent_id"` + Action string `json:"action"` + Lim int64 `json:"lim"` +} + +// ListAuditEntriesByAction +// +// SELECT id, +// agent_id, +// session_key, +// action, +// target, +// input, +// output, +// duration_ms, +// created_at +// FROM agent_audit_log +// WHERE agent_id = ?1 +// AND action = ?2 +// ORDER BY created_at DESC +// LIMIT ?3 +func (q *Queries) ListAuditEntriesByAction(ctx context.Context, arg ListAuditEntriesByActionParams) ([]AgentAuditLog, error) { + rows, err := q.db.QueryContext(ctx, ListAuditEntriesByAction, arg.AgentID, arg.Action, arg.Lim) + if err != nil { + return nil, err + } + defer rows.Close() + items := []AgentAuditLog{} + for rows.Next() { + var i AgentAuditLog + if err := rows.Scan( + &i.ID, + &i.AgentID, + &i.SessionKey, + &i.Action, + &i.Target, + &i.Input, + &i.Output, + &i.DurationMs, + &i.CreatedAt, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const ListAuditEntriesBySession = `-- name: ListAuditEntriesBySession :many +SELECT id, + agent_id, + session_key, + action, + target, + input, + output, + duration_ms, + created_at +FROM agent_audit_log +WHERE agent_id = ?1 + AND session_key = ?2 +ORDER BY created_at DESC +LIMIT ?3 +` + +type ListAuditEntriesBySessionParams struct { + AgentID string `json:"agent_id"` + SessionKey string `json:"session_key"` + Lim int64 `json:"lim"` +} + +// ListAuditEntriesBySession +// +// SELECT id, +// agent_id, +// session_key, +// action, +// target, +// input, +// output, +// duration_ms, +// created_at +// FROM agent_audit_log +// WHERE agent_id = ?1 +// AND session_key = ?2 +// ORDER BY created_at DESC +// LIMIT ?3 +func (q *Queries) ListAuditEntriesBySession(ctx context.Context, arg ListAuditEntriesBySessionParams) ([]AgentAuditLog, error) { + rows, err := q.db.QueryContext(ctx, ListAuditEntriesBySession, arg.AgentID, arg.SessionKey, arg.Lim) + if err != nil { + return nil, err + } + defer rows.Close() + items := []AgentAuditLog{} + for rows.Next() { + var i AgentAuditLog + if err := rows.Scan( + &i.ID, + &i.AgentID, + &i.SessionKey, + &i.Action, + &i.Target, + &i.Input, + &i.Output, + &i.DurationMs, + &i.CreatedAt, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const PruneOldAuditEntries = `-- name: PruneOldAuditEntries :exec +DELETE FROM agent_audit_log +WHERE agent_id = ?1 + AND created_at < ?2 +` + +type PruneOldAuditEntriesParams struct { + AgentID string `json:"agent_id"` + Before time.Time `json:"before"` +} + +// PruneOldAuditEntries +// +// DELETE FROM agent_audit_log +// WHERE agent_id = ?1 +// AND created_at < ?2 +func (q *Queries) PruneOldAuditEntries(ctx context.Context, arg PruneOldAuditEntriesParams) error { + _, err := q.db.ExecContext(ctx, PruneOldAuditEntries, arg.AgentID, arg.Before) + return err +} diff --git a/pkg/memory/sqlc/agent_documents.sql.go b/pkg/memory/sqlc/agent_documents.sql.go new file mode 100644 index 000000000..43530f112 --- /dev/null +++ b/pkg/memory/sqlc/agent_documents.sql.go @@ -0,0 +1,305 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.30.0 +// source: agent_documents.sql + +package sqlc + +import ( + "context" + + "github.com/sipeed/picoclaw/pkg/ids" +) + +const DeleteDocument = `-- name: DeleteDocument :exec +DELETE FROM agent_documents +WHERE agent_id = ?1 + AND name = ?2 +` + +type DeleteDocumentParams struct { + AgentID string `json:"agent_id"` + Name string `json:"name"` +} + +// DeleteDocument +// +// DELETE FROM agent_documents +// WHERE agent_id = ?1 +// AND name = ?2 +func (q *Queries) DeleteDocument(ctx context.Context, arg DeleteDocumentParams) error { + _, err := q.db.ExecContext(ctx, DeleteDocument, arg.AgentID, arg.Name) + return err +} + +const GetDocument = `-- name: GetDocument :one +SELECT id, + agent_id, + name, + category, + content, + version, + is_active, + created_at, + updated_at +FROM agent_documents +WHERE agent_id = ?1 + AND name = ?2 +` + +type GetDocumentParams struct { + AgentID string `json:"agent_id"` + Name string `json:"name"` +} + +// Agent Documents queries +// +// SELECT id, +// agent_id, +// name, +// category, +// content, +// version, +// is_active, +// created_at, +// updated_at +// FROM agent_documents +// WHERE agent_id = ?1 +// AND name = ?2 +func (q *Queries) GetDocument(ctx context.Context, arg GetDocumentParams) (AgentDocument, error) { + row := q.db.QueryRowContext(ctx, GetDocument, arg.AgentID, arg.Name) + var i AgentDocument + err := row.Scan( + &i.ID, + &i.AgentID, + &i.Name, + &i.Category, + &i.Content, + &i.Version, + &i.IsActive, + &i.CreatedAt, + &i.UpdatedAt, + ) + return i, err +} + +const ListAllDocuments = `-- name: ListAllDocuments :many +SELECT id, + agent_id, + name, + category, + content, + version, + is_active, + created_at, + updated_at +FROM agent_documents +WHERE agent_id = ?1 + AND is_active = 1 +ORDER BY category, + name +` + +type ListAllDocumentsParams struct { + AgentID string `json:"agent_id"` +} + +// ListAllDocuments +// +// SELECT id, +// agent_id, +// name, +// category, +// content, +// version, +// is_active, +// created_at, +// updated_at +// FROM agent_documents +// WHERE agent_id = ?1 +// AND is_active = 1 +// ORDER BY category, +// name +func (q *Queries) ListAllDocuments(ctx context.Context, arg ListAllDocumentsParams) ([]AgentDocument, error) { + rows, err := q.db.QueryContext(ctx, ListAllDocuments, arg.AgentID) + if err != nil { + return nil, err + } + defer rows.Close() + items := []AgentDocument{} + for rows.Next() { + var i AgentDocument + if err := rows.Scan( + &i.ID, + &i.AgentID, + &i.Name, + &i.Category, + &i.Content, + &i.Version, + &i.IsActive, + &i.CreatedAt, + &i.UpdatedAt, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const ListDocumentsByCategory = `-- name: ListDocumentsByCategory :many +SELECT id, + agent_id, + name, + category, + content, + version, + is_active, + created_at, + updated_at +FROM agent_documents +WHERE agent_id = ?1 + AND category = ?2 + AND is_active = 1 +ORDER BY name +` + +type ListDocumentsByCategoryParams struct { + AgentID string `json:"agent_id"` + Category string `json:"category"` +} + +// ListDocumentsByCategory +// +// SELECT id, +// agent_id, +// name, +// category, +// content, +// version, +// is_active, +// created_at, +// updated_at +// FROM agent_documents +// WHERE agent_id = ?1 +// AND category = ?2 +// AND is_active = 1 +// ORDER BY name +func (q *Queries) ListDocumentsByCategory(ctx context.Context, arg ListDocumentsByCategoryParams) ([]AgentDocument, error) { + rows, err := q.db.QueryContext(ctx, ListDocumentsByCategory, arg.AgentID, arg.Category) + if err != nil { + return nil, err + } + defer rows.Close() + items := []AgentDocument{} + for rows.Next() { + var i AgentDocument + if err := rows.Scan( + &i.ID, + &i.AgentID, + &i.Name, + &i.Category, + &i.Content, + &i.Version, + &i.IsActive, + &i.CreatedAt, + &i.UpdatedAt, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const UpsertDocument = `-- name: UpsertDocument :exec +INSERT INTO agent_documents ( + id, + agent_id, + name, + category, + content, + version, + is_active, + created_at, + updated_at + ) +VALUES ( + ?1, + ?2, + ?3, + ?4, + ?5, + 1, + 1, + datetime('now'), + datetime('now') + ) ON CONFLICT (agent_id, name) DO +UPDATE +SET content = excluded.content, + category = excluded.category, + version = agent_documents.version + 1, + is_active = 1, + updated_at = datetime('now') +` + +type UpsertDocumentParams struct { + ID ids.UUID `json:"id"` + AgentID string `json:"agent_id"` + Name string `json:"name"` + Category string `json:"category"` + Content string `json:"content"` +} + +// UpsertDocument +// +// INSERT INTO agent_documents ( +// id, +// agent_id, +// name, +// category, +// content, +// version, +// is_active, +// created_at, +// updated_at +// ) +// VALUES ( +// ?1, +// ?2, +// ?3, +// ?4, +// ?5, +// 1, +// 1, +// datetime('now'), +// datetime('now') +// ) ON CONFLICT (agent_id, name) DO +// UPDATE +// SET content = excluded.content, +// category = excluded.category, +// version = agent_documents.version + 1, +// is_active = 1, +// updated_at = datetime('now') +func (q *Queries) UpsertDocument(ctx context.Context, arg UpsertDocumentParams) error { + _, err := q.db.ExecContext(ctx, UpsertDocument, + arg.ID, + arg.AgentID, + arg.Name, + arg.Category, + arg.Content, + ) + return err +} diff --git a/pkg/memory/sqlc/agent_kv.sql.go b/pkg/memory/sqlc/agent_kv.sql.go new file mode 100644 index 000000000..4edb3521c --- /dev/null +++ b/pkg/memory/sqlc/agent_kv.sql.go @@ -0,0 +1,160 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.30.0 +// source: agent_kv.sql + +package sqlc + +import ( + "context" +) + +const DeleteKV = `-- name: DeleteKV :exec +DELETE FROM agent_kv +WHERE agent_id = ?1 + AND key = ?2 +` + +type DeleteKVParams struct { + AgentID string `json:"agent_id"` + Key string `json:"key"` +} + +// DeleteKV +// +// DELETE FROM agent_kv +// WHERE agent_id = ?1 +// AND key = ?2 +func (q *Queries) DeleteKV(ctx context.Context, arg DeleteKVParams) error { + _, err := q.db.ExecContext(ctx, DeleteKV, arg.AgentID, arg.Key) + return err +} + +const GetKV = `-- name: GetKV :one +SELECT agent_id, + key, + value, + updated_at +FROM agent_kv +WHERE agent_id = ?1 + AND key = ?2 +` + +type GetKVParams struct { + AgentID string `json:"agent_id"` + Key string `json:"key"` +} + +// Agent KV Store queries +// +// SELECT agent_id, +// key, +// value, +// updated_at +// FROM agent_kv +// WHERE agent_id = ?1 +// AND key = ?2 +func (q *Queries) GetKV(ctx context.Context, arg GetKVParams) (AgentKv, error) { + row := q.db.QueryRowContext(ctx, GetKV, arg.AgentID, arg.Key) + var i AgentKv + err := row.Scan( + &i.AgentID, + &i.Key, + &i.Value, + &i.UpdatedAt, + ) + return i, err +} + +const ListKVByPrefix = `-- name: ListKVByPrefix :many +SELECT agent_id, + key, + value, + updated_at +FROM agent_kv +WHERE agent_id = ?1 + AND key LIKE ?2 || '%' +ORDER BY key +LIMIT ?3 +` + +type ListKVByPrefixParams struct { + AgentID string `json:"agent_id"` + Prefix *string `json:"prefix"` + Lim int64 `json:"lim"` +} + +// ListKVByPrefix +// +// SELECT agent_id, +// key, +// value, +// updated_at +// FROM agent_kv +// WHERE agent_id = ?1 +// AND key LIKE ?2 || '%' +// ORDER BY key +// LIMIT ?3 +func (q *Queries) ListKVByPrefix(ctx context.Context, arg ListKVByPrefixParams) ([]AgentKv, error) { + rows, err := q.db.QueryContext(ctx, ListKVByPrefix, arg.AgentID, arg.Prefix, arg.Lim) + if err != nil { + return nil, err + } + defer rows.Close() + items := []AgentKv{} + for rows.Next() { + var i AgentKv + if err := rows.Scan( + &i.AgentID, + &i.Key, + &i.Value, + &i.UpdatedAt, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const UpsertKV = `-- name: UpsertKV :exec +INSERT INTO agent_kv (agent_id, key, value, updated_at) +VALUES ( + ?1, + ?2, + ?3, + datetime('now') + ) ON CONFLICT (agent_id, key) DO +UPDATE +SET value = excluded.value, + updated_at = excluded.updated_at +` + +type UpsertKVParams struct { + AgentID string `json:"agent_id"` + Key string `json:"key"` + Value string `json:"value"` +} + +// UpsertKV +// +// INSERT INTO agent_kv (agent_id, key, value, updated_at) +// VALUES ( +// ?1, +// ?2, +// ?3, +// datetime('now') +// ) ON CONFLICT (agent_id, key) DO +// UPDATE +// SET value = excluded.value, +// updated_at = excluded.updated_at +func (q *Queries) UpsertKV(ctx context.Context, arg UpsertKVParams) error { + _, err := q.db.ExecContext(ctx, UpsertKV, arg.AgentID, arg.Key, arg.Value) + return err +} diff --git a/pkg/memory/sqlc/models.go b/pkg/memory/sqlc/models.go index 471faf9f1..bd9198046 100644 --- a/pkg/memory/sqlc/models.go +++ b/pkg/memory/sqlc/models.go @@ -11,6 +11,37 @@ import ( "github.com/sipeed/picoclaw/pkg/memory" ) +type AgentAuditLog struct { + ID ids.UUID `json:"id"` + AgentID string `json:"agent_id"` + SessionKey string `json:"session_key"` + Action string `json:"action"` + Target string `json:"target"` + Input *string `json:"input"` + Output *string `json:"output"` + DurationMs *int64 `json:"duration_ms"` + CreatedAt time.Time `json:"created_at"` +} + +type AgentDocument struct { + ID ids.UUID `json:"id"` + AgentID string `json:"agent_id"` + Name string `json:"name"` + Category string `json:"category"` + Content string `json:"content"` + Version int64 `json:"version"` + IsActive bool `json:"is_active"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` +} + +type AgentKv struct { + AgentID string `json:"agent_id"` + Key string `json:"key"` + Value string `json:"value"` + UpdatedAt time.Time `json:"updated_at"` +} + type ArchivalChunk struct { ID ids.UUID `json:"id"` RecallID ids.UUID `json:"recall_id"` diff --git a/pkg/memory/sqlc/querier.go b/pkg/memory/sqlc/querier.go index defb86eff..a8ac5293c 100644 --- a/pkg/memory/sqlc/querier.go +++ b/pkg/memory/sqlc/querier.go @@ -14,6 +14,19 @@ type Querier interface { // SELECT COUNT(*) // FROM archival_chunks CountArchivalChunks(ctx context.Context) (int64, error) + //CountAuditEntries + // + // SELECT COUNT(*) + // FROM agent_audit_log + // WHERE agent_id = ?1 + CountAuditEntries(ctx context.Context, arg CountAuditEntriesParams) (int64, error) + //CountAuditEntriesByAction + // + // SELECT COUNT(*) + // FROM agent_audit_log + // WHERE agent_id = ?1 + // AND action = ?2 + CountAuditEntriesByAction(ctx context.Context, arg CountAuditEntriesByActionParams) (int64, error) //CountRecallItems // // SELECT COUNT(*) @@ -24,11 +37,31 @@ type Querier interface { // OR ?2 = '' // ) CountRecallItems(ctx context.Context, arg CountRecallItemsParams) (int64, error) + //CountSessionMessages + // + // SELECT COUNT(*) + // FROM recall_items + // WHERE agent_id = ?1 + // AND session_key = ?2 + // AND tags = 'session-message' + CountSessionMessages(ctx context.Context, arg CountSessionMessagesParams) (int64, error) //DeleteArchivalChunksByRecall // // DELETE FROM archival_chunks // WHERE recall_id = ?1 DeleteArchivalChunksByRecall(ctx context.Context, arg DeleteArchivalChunksByRecallParams) error + //DeleteDocument + // + // DELETE FROM agent_documents + // WHERE agent_id = ?1 + // AND name = ?2 + DeleteDocument(ctx context.Context, arg DeleteDocumentParams) error + //DeleteKV + // + // DELETE FROM agent_kv + // WHERE agent_id = ?1 + // AND key = ?2 + DeleteKV(ctx context.Context, arg DeleteKVParams) error //DeleteRecallItem // // DELETE FROM recall_items @@ -60,6 +93,31 @@ type Querier interface { // FROM archival_chunks // WHERE id IN (/*SLICE:ids*/?) GetArchivalChunksByIDs(ctx context.Context, arg GetArchivalChunksByIDsParams) ([]ArchivalChunk, error) + // Agent Documents queries + // + // SELECT id, + // agent_id, + // name, + // category, + // content, + // version, + // is_active, + // created_at, + // updated_at + // FROM agent_documents + // WHERE agent_id = ?1 + // AND name = ?2 + GetDocument(ctx context.Context, arg GetDocumentParams) (AgentDocument, error) + // Agent KV Store queries + // + // SELECT agent_id, + // key, + // value, + // updated_at + // FROM agent_kv + // WHERE agent_id = ?1 + // AND key = ?2 + GetKV(ctx context.Context, arg GetKVParams) (AgentKv, error) //GetRecallItem // // SELECT id, @@ -127,6 +185,31 @@ type Querier interface { // datetime('now') // ) InsertArchivalChunk(ctx context.Context, arg InsertArchivalChunkParams) error + // Agent Audit Log queries + // + // INSERT INTO agent_audit_log ( + // id, + // agent_id, + // session_key, + // action, + // target, + // input, + // output, + // duration_ms, + // created_at + // ) + // VALUES ( + // ?1, + // ?2, + // ?3, + // ?4, + // ?5, + // ?6, + // ?7, + // ?8, + // datetime('now') + // ) + InsertAuditEntry(ctx context.Context, arg InsertAuditEntryParams) error // Recall Item queries // // INSERT INTO recall_items ( @@ -158,6 +241,37 @@ type Querier interface { // datetime('now') // ) InsertRecallItem(ctx context.Context, arg InsertRecallItemParams) error + //InsertSessionMessage + // + // INSERT INTO recall_items ( + // id, + // agent_id, + // session_key, + // role, + // sector, + // importance, + // salience, + // decay_rate, + // content, + // tags, + // created_at, + // updated_at + // ) + // VALUES ( + // ?1, + // ?2, + // ?3, + // ?4, + // 'episodic', + // 0.5, + // 0.5, + // 0.01, + // ?5, + // 'session-message', + // datetime('now'), + // datetime('now') + // ) + InsertSessionMessage(ctx context.Context, arg InsertSessionMessageParams) error // Memory Summary queries // // INSERT INTO memory_summaries ( @@ -193,6 +307,23 @@ type Querier interface { // ORDER BY created_at DESC // LIMIT ?2 OFFSET ?1 ListAllArchivalChunks(ctx context.Context, arg ListAllArchivalChunksParams) ([]ArchivalChunk, error) + //ListAllDocuments + // + // SELECT id, + // agent_id, + // name, + // category, + // content, + // version, + // is_active, + // created_at, + // updated_at + // FROM agent_documents + // WHERE agent_id = ?1 + // AND is_active = 1 + // ORDER BY category, + // name + ListAllDocuments(ctx context.Context, arg ListAllDocumentsParams) ([]AgentDocument, error) //ListArchivalChunks // // SELECT id, @@ -207,6 +338,85 @@ type Querier interface { // WHERE recall_id = ?1 // ORDER BY chunk_index ListArchivalChunks(ctx context.Context, arg ListArchivalChunksParams) ([]ArchivalChunk, error) + //ListAuditEntries + // + // SELECT id, + // agent_id, + // session_key, + // action, + // target, + // input, + // output, + // duration_ms, + // created_at + // FROM agent_audit_log + // WHERE agent_id = ?1 + // ORDER BY created_at DESC + // LIMIT ?2 + ListAuditEntries(ctx context.Context, arg ListAuditEntriesParams) ([]AgentAuditLog, error) + //ListAuditEntriesByAction + // + // SELECT id, + // agent_id, + // session_key, + // action, + // target, + // input, + // output, + // duration_ms, + // created_at + // FROM agent_audit_log + // WHERE agent_id = ?1 + // AND action = ?2 + // ORDER BY created_at DESC + // LIMIT ?3 + ListAuditEntriesByAction(ctx context.Context, arg ListAuditEntriesByActionParams) ([]AgentAuditLog, error) + //ListAuditEntriesBySession + // + // SELECT id, + // agent_id, + // session_key, + // action, + // target, + // input, + // output, + // duration_ms, + // created_at + // FROM agent_audit_log + // WHERE agent_id = ?1 + // AND session_key = ?2 + // ORDER BY created_at DESC + // LIMIT ?3 + ListAuditEntriesBySession(ctx context.Context, arg ListAuditEntriesBySessionParams) ([]AgentAuditLog, error) + //ListDocumentsByCategory + // + // SELECT id, + // agent_id, + // name, + // category, + // content, + // version, + // is_active, + // created_at, + // updated_at + // FROM agent_documents + // WHERE agent_id = ?1 + // AND category = ?2 + // AND is_active = 1 + // ORDER BY name + ListDocumentsByCategory(ctx context.Context, arg ListDocumentsByCategoryParams) ([]AgentDocument, error) + //ListKVByPrefix + // + // SELECT agent_id, + // key, + // value, + // updated_at + // FROM agent_kv + // WHERE agent_id = ?1 + // AND key LIKE ?2 || '%' + // ORDER BY key + // LIMIT ?3 + ListKVByPrefix(ctx context.Context, arg ListKVByPrefixParams) ([]AgentKv, error) //ListRecallItems // // SELECT id, @@ -230,6 +440,31 @@ type Querier interface { // ORDER BY created_at DESC // LIMIT ?4 OFFSET ?3 ListRecallItems(ctx context.Context, arg ListRecallItemsParams) ([]RecallItem, error) + //ListSessionMessages + // + // SELECT id, + // agent_id, + // session_key, + // role, + // sector, + // importance, + // salience, + // decay_rate, + // content, + // tags, + // created_at, + // updated_at + // FROM recall_items + // WHERE agent_id = ?1 + // AND session_key = ?2 + // AND tags = 'session-message' + // AND ( + // role = ?3 + // OR ?3 = '' + // ) + // ORDER BY created_at ASC + // LIMIT ?4 + ListSessionMessages(ctx context.Context, arg ListSessionMessagesParams) ([]RecallItem, error) //ListSummaries // // SELECT id, @@ -248,6 +483,12 @@ type Querier interface { // ORDER BY created_at DESC // LIMIT ?3 ListSummaries(ctx context.Context, arg ListSummariesParams) ([]MemorySummary, error) + //PruneOldAuditEntries + // + // DELETE FROM agent_audit_log + // WHERE agent_id = ?1 + // AND created_at < ?2 + PruneOldAuditEntries(ctx context.Context, arg PruneOldAuditEntriesParams) error //SearchRecallByKeyword // // SELECT ri.id, @@ -281,6 +522,50 @@ type Querier interface { // updated_at = datetime('now') // WHERE id = ?8 UpdateRecallItem(ctx context.Context, arg UpdateRecallItemParams) error + //UpsertDocument + // + // INSERT INTO agent_documents ( + // id, + // agent_id, + // name, + // category, + // content, + // version, + // is_active, + // created_at, + // updated_at + // ) + // VALUES ( + // ?1, + // ?2, + // ?3, + // ?4, + // ?5, + // 1, + // 1, + // datetime('now'), + // datetime('now') + // ) ON CONFLICT (agent_id, name) DO + // UPDATE + // SET content = excluded.content, + // category = excluded.category, + // version = agent_documents.version + 1, + // is_active = 1, + // updated_at = datetime('now') + UpsertDocument(ctx context.Context, arg UpsertDocumentParams) error + //UpsertKV + // + // INSERT INTO agent_kv (agent_id, key, value, updated_at) + // VALUES ( + // ?1, + // ?2, + // ?3, + // datetime('now') + // ) ON CONFLICT (agent_id, key) DO + // UPDATE + // SET value = excluded.value, + // updated_at = excluded.updated_at + UpsertKV(ctx context.Context, arg UpsertKVParams) error //UpsertWorkingContext // // INSERT INTO working_context (agent_id, session_key, content, updated_at) diff --git a/pkg/memory/sqlc/queries/agent_audit_log.sql b/pkg/memory/sqlc/queries/agent_audit_log.sql new file mode 100644 index 000000000..dd06ad225 --- /dev/null +++ b/pkg/memory/sqlc/queries/agent_audit_log.sql @@ -0,0 +1,81 @@ +-- Agent Audit Log queries +-- name: InsertAuditEntry :exec +INSERT INTO agent_audit_log ( + id, + agent_id, + session_key, + action, + target, + input, + output, + duration_ms, + created_at + ) +VALUES ( + sqlc.arg(id), + sqlc.arg(agent_id), + sqlc.arg(session_key), + sqlc.arg(action), + sqlc.arg(target), + sqlc.arg(input), + sqlc.arg(output), + sqlc.arg(duration_ms), + datetime('now') + ); +-- name: ListAuditEntries :many +SELECT id, + agent_id, + session_key, + action, + target, + input, + output, + duration_ms, + created_at +FROM agent_audit_log +WHERE agent_id = sqlc.arg(agent_id) +ORDER BY created_at DESC +LIMIT sqlc.arg(lim); +-- name: ListAuditEntriesByAction :many +SELECT id, + agent_id, + session_key, + action, + target, + input, + output, + duration_ms, + created_at +FROM agent_audit_log +WHERE agent_id = sqlc.arg(agent_id) + AND action = sqlc.arg(action) +ORDER BY created_at DESC +LIMIT sqlc.arg(lim); +-- name: CountAuditEntries :one +SELECT COUNT(*) +FROM agent_audit_log +WHERE agent_id = sqlc.arg(agent_id); +-- name: ListAuditEntriesBySession :many +SELECT id, + agent_id, + session_key, + action, + target, + input, + output, + duration_ms, + created_at +FROM agent_audit_log +WHERE agent_id = sqlc.arg(agent_id) + AND session_key = sqlc.arg(session_key) +ORDER BY created_at DESC +LIMIT sqlc.arg(lim); +-- name: PruneOldAuditEntries :exec +DELETE FROM agent_audit_log +WHERE agent_id = sqlc.arg(agent_id) + AND created_at < sqlc.arg(before); +-- name: CountAuditEntriesByAction :one +SELECT COUNT(*) +FROM agent_audit_log +WHERE agent_id = sqlc.arg(agent_id) + AND action = sqlc.arg(action); \ No newline at end of file diff --git a/pkg/memory/sqlc/queries/agent_documents.sql b/pkg/memory/sqlc/queries/agent_documents.sql new file mode 100644 index 000000000..7ba1540c4 --- /dev/null +++ b/pkg/memory/sqlc/queries/agent_documents.sql @@ -0,0 +1,77 @@ +-- Agent Documents queries +-- name: GetDocument :one +SELECT id, + agent_id, + name, + category, + content, + version, + is_active, + created_at, + updated_at +FROM agent_documents +WHERE agent_id = sqlc.arg(agent_id) + AND name = sqlc.arg(name); +-- name: UpsertDocument :exec +INSERT INTO agent_documents ( + id, + agent_id, + name, + category, + content, + version, + is_active, + created_at, + updated_at + ) +VALUES ( + sqlc.arg(id), + sqlc.arg(agent_id), + sqlc.arg(name), + sqlc.arg(category), + sqlc.arg(content), + 1, + 1, + datetime('now'), + datetime('now') + ) ON CONFLICT (agent_id, name) DO +UPDATE +SET content = excluded.content, + category = excluded.category, + version = agent_documents.version + 1, + is_active = 1, + updated_at = datetime('now'); +-- name: ListDocumentsByCategory :many +SELECT id, + agent_id, + name, + category, + content, + version, + is_active, + created_at, + updated_at +FROM agent_documents +WHERE agent_id = sqlc.arg(agent_id) + AND category = sqlc.arg(category) + AND is_active = 1 +ORDER BY name; +-- name: DeleteDocument :exec +DELETE FROM agent_documents +WHERE agent_id = sqlc.arg(agent_id) + AND name = sqlc.arg(name); +-- name: ListAllDocuments :many +SELECT id, + agent_id, + name, + category, + content, + version, + is_active, + created_at, + updated_at +FROM agent_documents +WHERE agent_id = sqlc.arg(agent_id) + AND is_active = 1 +ORDER BY category, + name; \ No newline at end of file diff --git a/pkg/memory/sqlc/queries/agent_kv.sql b/pkg/memory/sqlc/queries/agent_kv.sql new file mode 100644 index 000000000..5b3bdaee2 --- /dev/null +++ b/pkg/memory/sqlc/queries/agent_kv.sql @@ -0,0 +1,34 @@ +-- Agent KV Store queries +-- name: GetKV :one +SELECT agent_id, + key, + value, + updated_at +FROM agent_kv +WHERE agent_id = sqlc.arg(agent_id) + AND key = sqlc.arg(key); +-- name: UpsertKV :exec +INSERT INTO agent_kv (agent_id, key, value, updated_at) +VALUES ( + sqlc.arg(agent_id), + sqlc.arg(key), + sqlc.arg(value), + datetime('now') + ) ON CONFLICT (agent_id, key) DO +UPDATE +SET value = excluded.value, + updated_at = excluded.updated_at; +-- name: DeleteKV :exec +DELETE FROM agent_kv +WHERE agent_id = sqlc.arg(agent_id) + AND key = sqlc.arg(key); +-- name: ListKVByPrefix :many +SELECT agent_id, + key, + value, + updated_at +FROM agent_kv +WHERE agent_id = sqlc.arg(agent_id) + AND key LIKE sqlc.arg(prefix) || '%' +ORDER BY key +LIMIT sqlc.arg(lim); \ No newline at end of file diff --git a/pkg/memory/sqlc/queries/recall.sql b/pkg/memory/sqlc/queries/recall.sql index 66542c25d..73b9e535e 100644 --- a/pkg/memory/sqlc/queries/recall.sql +++ b/pkg/memory/sqlc/queries/recall.sql @@ -118,4 +118,62 @@ SELECT id, created_at, updated_at FROM recall_items -WHERE id IN (sqlc.slice('ids')); \ No newline at end of file +WHERE id IN (sqlc.slice('ids')); +-- name: InsertSessionMessage :exec +INSERT INTO recall_items ( + id, + agent_id, + session_key, + role, + sector, + importance, + salience, + decay_rate, + content, + tags, + created_at, + updated_at + ) +VALUES ( + sqlc.arg(id), + sqlc.arg(agent_id), + sqlc.arg(session_key), + sqlc.arg(role), + 'episodic', + 0.5, + 0.5, + 0.01, + sqlc.arg(content), + 'session-message', + datetime('now'), + datetime('now') + ); +-- name: ListSessionMessages :many +SELECT id, + agent_id, + session_key, + role, + sector, + importance, + salience, + decay_rate, + content, + tags, + created_at, + updated_at +FROM recall_items +WHERE agent_id = sqlc.arg(agent_id) + AND session_key = sqlc.arg(session_key) + AND tags = 'session-message' + AND ( + role = sqlc.arg(role) + OR sqlc.arg(role) = '' + ) +ORDER BY created_at ASC +LIMIT sqlc.arg(lim); +-- name: CountSessionMessages :one +SELECT COUNT(*) +FROM recall_items +WHERE agent_id = sqlc.arg(agent_id) + AND session_key = sqlc.arg(session_key) + AND tags = 'session-message'; \ No newline at end of file diff --git a/pkg/memory/sqlc/recall.sql.go b/pkg/memory/sqlc/recall.sql.go index 756f6a629..32ac3e1ef 100644 --- a/pkg/memory/sqlc/recall.sql.go +++ b/pkg/memory/sqlc/recall.sql.go @@ -44,6 +44,33 @@ func (q *Queries) CountRecallItems(ctx context.Context, arg CountRecallItemsPara return count, err } +const CountSessionMessages = `-- name: CountSessionMessages :one +SELECT COUNT(*) +FROM recall_items +WHERE agent_id = ?1 + AND session_key = ?2 + AND tags = 'session-message' +` + +type CountSessionMessagesParams struct { + AgentID string `json:"agent_id"` + SessionKey string `json:"session_key"` +} + +// CountSessionMessages +// +// SELECT COUNT(*) +// FROM recall_items +// WHERE agent_id = ?1 +// AND session_key = ?2 +// AND tags = 'session-message' +func (q *Queries) CountSessionMessages(ctx context.Context, arg CountSessionMessagesParams) (int64, error) { + row := q.db.QueryRowContext(ctx, CountSessionMessages, arg.AgentID, arg.SessionKey) + var count int64 + err := row.Scan(&count) + return count, err +} + const DeleteRecallItem = `-- name: DeleteRecallItem :exec DELETE FROM recall_items WHERE id = ?1 @@ -292,6 +319,86 @@ func (q *Queries) InsertRecallItem(ctx context.Context, arg InsertRecallItemPara return err } +const InsertSessionMessage = `-- name: InsertSessionMessage :exec +INSERT INTO recall_items ( + id, + agent_id, + session_key, + role, + sector, + importance, + salience, + decay_rate, + content, + tags, + created_at, + updated_at + ) +VALUES ( + ?1, + ?2, + ?3, + ?4, + 'episodic', + 0.5, + 0.5, + 0.01, + ?5, + 'session-message', + datetime('now'), + datetime('now') + ) +` + +type InsertSessionMessageParams struct { + ID ids.UUID `json:"id"` + AgentID string `json:"agent_id"` + SessionKey string `json:"session_key"` + Role string `json:"role"` + Content string `json:"content"` +} + +// InsertSessionMessage +// +// INSERT INTO recall_items ( +// id, +// agent_id, +// session_key, +// role, +// sector, +// importance, +// salience, +// decay_rate, +// content, +// tags, +// created_at, +// updated_at +// ) +// VALUES ( +// ?1, +// ?2, +// ?3, +// ?4, +// 'episodic', +// 0.5, +// 0.5, +// 0.01, +// ?5, +// 'session-message', +// datetime('now'), +// datetime('now') +// ) +func (q *Queries) InsertSessionMessage(ctx context.Context, arg InsertSessionMessageParams) error { + _, err := q.db.ExecContext(ctx, InsertSessionMessage, + arg.ID, + arg.AgentID, + arg.SessionKey, + arg.Role, + arg.Content, + ) + return err +} + const ListRecallItems = `-- name: ListRecallItems :many SELECT id, agent_id, @@ -385,6 +492,103 @@ func (q *Queries) ListRecallItems(ctx context.Context, arg ListRecallItemsParams return items, nil } +const ListSessionMessages = `-- name: ListSessionMessages :many +SELECT id, + agent_id, + session_key, + role, + sector, + importance, + salience, + decay_rate, + content, + tags, + created_at, + updated_at +FROM recall_items +WHERE agent_id = ?1 + AND session_key = ?2 + AND tags = 'session-message' + AND ( + role = ?3 + OR ?3 = '' + ) +ORDER BY created_at ASC +LIMIT ?4 +` + +type ListSessionMessagesParams struct { + AgentID string `json:"agent_id"` + SessionKey string `json:"session_key"` + Role string `json:"role"` + Lim int64 `json:"lim"` +} + +// ListSessionMessages +// +// SELECT id, +// agent_id, +// session_key, +// role, +// sector, +// importance, +// salience, +// decay_rate, +// content, +// tags, +// created_at, +// updated_at +// FROM recall_items +// WHERE agent_id = ?1 +// AND session_key = ?2 +// AND tags = 'session-message' +// AND ( +// role = ?3 +// OR ?3 = '' +// ) +// ORDER BY created_at ASC +// LIMIT ?4 +func (q *Queries) ListSessionMessages(ctx context.Context, arg ListSessionMessagesParams) ([]RecallItem, error) { + rows, err := q.db.QueryContext(ctx, ListSessionMessages, + arg.AgentID, + arg.SessionKey, + arg.Role, + arg.Lim, + ) + if err != nil { + return nil, err + } + defer rows.Close() + items := []RecallItem{} + for rows.Next() { + var i RecallItem + if err := rows.Scan( + &i.ID, + &i.AgentID, + &i.SessionKey, + &i.Role, + &i.Sector, + &i.Importance, + &i.Salience, + &i.DecayRate, + &i.Content, + &i.Tags, + &i.CreatedAt, + &i.UpdatedAt, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + const SearchRecallByKeyword = `-- name: SearchRecallByKeyword :many SELECT ri.id, ri.agent_id, diff --git a/pkg/memory/sqlc/schema.sql b/pkg/memory/sqlc/schema.sql index 4b4a03415..5eacba709 100644 --- a/pkg/memory/sqlc/schema.sql +++ b/pkg/memory/sqlc/schema.sql @@ -58,4 +58,41 @@ CREATE TABLE IF NOT EXISTS memory_summaries ( to_msg_idx INTEGER NOT NULL DEFAULT 0, created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP ); -CREATE INDEX IF NOT EXISTS idx_summaries_agent_session ON memory_summaries(agent_id, session_key); \ No newline at end of file +CREATE INDEX IF NOT EXISTS idx_summaries_agent_session ON memory_summaries(agent_id, session_key); +-- Agent KV store: generic key-value pairs for agent state, preferences, config +CREATE TABLE IF NOT EXISTS agent_kv ( + agent_id TEXT NOT NULL, + key TEXT NOT NULL, + value TEXT NOT NULL DEFAULT '', + updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + PRIMARY KEY (agent_id, key) +); +-- Agent documents: versioned named documents (bootstrap files, identity, etc.) +CREATE TABLE IF NOT EXISTS agent_documents ( + id BLOB PRIMARY KEY, + agent_id TEXT NOT NULL, + name TEXT NOT NULL, + category TEXT NOT NULL DEFAULT 'bootstrap', + content TEXT NOT NULL, + version INTEGER NOT NULL DEFAULT 1, + is_active BOOLEAN NOT NULL DEFAULT 1, + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + UNIQUE(agent_id, name) +); +CREATE INDEX IF NOT EXISTS idx_docs_agent_cat ON agent_documents(agent_id, category); +CREATE INDEX IF NOT EXISTS idx_docs_name ON agent_documents(name); +-- Agent audit log: append-only record of tool calls, state changes, etc. +CREATE TABLE IF NOT EXISTS agent_audit_log ( + id BLOB PRIMARY KEY, + agent_id TEXT NOT NULL, + session_key TEXT NOT NULL DEFAULT '', + action TEXT NOT NULL, + target TEXT NOT NULL DEFAULT '', + input TEXT, + output TEXT, + duration_ms INTEGER, + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP +); +CREATE INDEX IF NOT EXISTS idx_audit_agent_time ON agent_audit_log(agent_id, created_at DESC); +CREATE INDEX IF NOT EXISTS idx_audit_action ON agent_audit_log(action); \ No newline at end of file diff --git a/pkg/memory/sqlc/sqlc.yaml b/pkg/memory/sqlc/sqlc.yaml index 625e38c10..17c52018b 100644 --- a/pkg/memory/sqlc/sqlc.yaml +++ b/pkg/memory/sqlc/sqlc.yaml @@ -40,6 +40,14 @@ sql: go_type: import: "github.com/sipeed/picoclaw/pkg/ids" type: "UUID" + - column: "agent_documents.id" + go_type: + import: "github.com/sipeed/picoclaw/pkg/ids" + type: "UUID" + - column: "agent_audit_log.id" + go_type: + import: "github.com/sipeed/picoclaw/pkg/ids" + type: "UUID" # Domain type: recall_items.sector → memory.Sector - column: "recall_items.sector" go_type: diff --git a/pkg/memory/store/memory_store.go b/pkg/memory/store/memory_store.go index 66d4a2767..ed0f39c8d 100644 --- a/pkg/memory/store/memory_store.go +++ b/pkg/memory/store/memory_store.go @@ -204,6 +204,16 @@ func (m *MemoryStore) RetrieveArchival(ctx context.Context, id ids.UUID) (string return string(buf), nil } +// ReadByID loads the content of a memory entry by its UUID string. +// Tries archival chunks first, then falls back to recall item. +func (m *MemoryStore) ReadByID(ctx context.Context, agentID, idStr string) (string, error) { + id, err := ids.Parse(idStr) + if err != nil { + return "", fmt.Errorf("invalid ID: %w", err) + } + return m.RetrieveArchival(ctx, id) +} + // --- Retrieval pipeline --- func (m *MemoryStore) Search(ctx context.Context, query string, opts memory.SearchOptions) ([]memory.SearchResult, error) {