opencode/picoclaw
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat(memory): add kernel contract, DAG snapshot persistence, and map_ops schema

Browse source 
pkg/memory/kernel_contract.go
- Defines the MemoryKernelContract interface for active context projection:
  ActiveContext(), ProjectSession(), SyncProjection(), RecoverSession()
- Provides a deterministic contract for lossless session continuity with
  DAG-based compression and projection pointer tracking

pkg/memory/dag/store.go
- DAGStore: persistent snapshot storage for DAG nodes and edges in SQLite
- Stores compression snapshots with provenance metadata (origin session,
  message range, token counts)
- Supports node/edge insertion, snapshot listing, and node lookup by ID

pkg/memory/dag/backfill.go + backfill_test.go
- BackfillDAGFromSession(): builds a DAG snapshot from an existing session
  history for sessions that pre-date DAG compression
- Idempotent: skips sessions that already have a snapshot

pkg/memory/migrations/011_dag_tables.go
- SQL migration adding dag_snapshots, dag_nodes, and dag_edges tables
- Includes indexes for snapshot_id, session_key, and node type lookups

pkg/memory/migrations/012_map_operator_runs.go
- SQL migration adding map_operator_runs and map_operator_items tables
- Tracks batch map operation runs with per-item status and FlatBuffers
  serialized payloads

pkg/memory/sqlc/queries/dag.sql + dag.sql.go
- SQLC queries for DAG snapshot CRUD operations

pkg/memory/sqlc/queries/map_ops.sql + map_ops.sql.go
- SQLC queries for map operator run and item tracking

pkg/memory/store/retrieval_policy.go
- RetrievalPolicy: configures token budget, recency bias, and semantic
  score thresholds for memory retrieval operations
- Used by the kernel contract to bound context window usage

pkg/memory/delegate/sqlite_dag_test.go
- Integration tests for DAG table creation, snapshot insertion, and
  node/edge persistence through the SQLite delegate
This commit is contained in:
ZanzyTHEbar 2026-02-21 18:57:46 +00:00
parent db4d98a1d4
commit 1a433dc596
14 changed files with 3026 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

236
pkg/memory/dag/backfill.go Normal file
View file

@ -0,0 +1,236 @@
package dag
import (
"context"
"database/sql"
"errors"
"sort"
"strings"
"time"
jsonv2 "github.com/go-json-experiment/json"
"github.com/ZanzyTHEbar/dragonscale/pkg/memory"
memsqlc "github.com/ZanzyTHEbar/dragonscale/pkg/memory/sqlc"
)
const dagBackfillStatusKVKey = "migration:dag_backfill:v1"
// BackfillOptions controls DAG backfill behavior.
type BackfillOptions struct {
PageSize int
MaxSessions int
Force bool
}
func DefaultBackfillOptions() BackfillOptions {
return BackfillOptions{
PageSize: 500,
MaxSessions: 0,
Force: false,
}
}
// BackfillStatus tracks one DAG backfill pass over existing sessions.
type BackfillStatus struct {
Version int `json:"version"`
SessionsScanned int `json:"sessions_scanned"`
SnapshotsCreated int `json:"snapshots_created"`
SkippedExisting int `json:"skipped_existing"`
Failures int `json:"failures"`
CompletedAt time.Time `json:"completed_at"`
Skipped bool `json:"skipped,omitempty"`
}
// BackfillMissingSessionDAGs creates DAG snapshots for sessions that have
// session-message history but no persisted DAG snapshot yet.
//
// The pass is one-shot by default and records status in KV under
// dagBackfillStatusKVKey. Use opts.Force=true to run again.
func BackfillMissingSessionDAGs(
ctx context.Context,
delegate memory.MemoryDelegate,
queries *memsqlc.Queries,
agentID string,
opts BackfillOptions,
) (*BackfillStatus, error) {
if delegate == nil || queries == nil {
return nil, nil
}
persister, ok := delegate.(DAGPersister)
if !ok {
return nil, nil
}
if opts.PageSize <= 0 {
opts.PageSize = 500
}
if !opts.Force {
if raw, err := delegate.GetKV(ctx, agentID, dagBackfillStatusKVKey); err == nil && strings.TrimSpace(raw) != "" {
var status BackfillStatus
if uErr := jsonv2.Unmarshal([]byte(raw), &status); uErr == nil {
status.Skipped = true
return &status, nil
}
return &BackfillStatus{Version: 1, Skipped: true}, nil
}
}
sessionKeys, err := collectSessionKeysForBackfill(ctx, delegate, agentID, opts.PageSize)
if err != nil {
return nil, err
}
if opts.MaxSessions > 0 && len(sessionKeys) > opts.MaxSessions {
sessionKeys = sessionKeys[:opts.MaxSessions]
}
status := &BackfillStatus{
Version: 1,
SessionsScanned: len(sessionKeys),
}
compressor := NewCompressor(DefaultCompressorConfig())
for _, sessionKey := range sessionKeys {
if ctx.Err() != nil {
return status, ctx.Err()
}
_, err := queries.GetLatestDAGSnapshotBySession(ctx, memsqlc.GetLatestDAGSnapshotBySessionParams{
AgentID: agentID,
SessionKey: sessionKey,
})
if err == nil {
status.SkippedExisting++
continue
}
if !errors.Is(err, sql.ErrNoRows) {
status.Failures++
continue
}
msgs, err := loadSessionMessagesForBackfill(ctx, delegate, agentID, sessionKey, opts.PageSize)
if err != nil {
status.Failures++
continue
}
if len(msgs) == 0 {
continue
}
d := compressor.Compress(msgs)
if d == nil || len(d.Nodes) == 0 {
continue
}
if err := persister.PersistDAG(ctx, agentID, sessionKey, &PersistSnapshot{
FromMsgIdx: 0,
ToMsgIdx: len(msgs),
MsgCount: len(msgs),
DAG: d,
}); err != nil {
status.Failures++
continue
}
status.SnapshotsCreated++
}
status.CompletedAt = time.Now().UTC()
data, err := jsonv2.Marshal(status)
if err != nil {
return status, err
}
if err := delegate.UpsertKV(ctx, agentID, dagBackfillStatusKVKey, string(data)); err != nil {
return status, err
}
return status, nil
}
func collectSessionKeysForBackfill(ctx context.Context, delegate memory.MemoryDelegate, agentID string, pageSize int) ([]string, error) {
keys := make(map[string]struct{})
offset := 0
for {
items, err := delegate.ListRecallItems(ctx, agentID, "", pageSize, offset)
if err != nil {
return nil, err
}
if len(items) == 0 {
break
}
for _, item := range items {
if !strings.Contains(item.Tags, "session-message") {
continue
}
keys[item.SessionKey] = struct{}{}
}
if len(items) < pageSize {
break
}
offset += len(items)
}
out := make([]string, 0, len(keys))
for k := range keys {
out = append(out, k)
}
sort.Strings(out)
return out, nil
}
func loadSessionMessagesForBackfill(
ctx context.Context,
delegate memory.MemoryDelegate,
agentID string,
sessionKey string,
pageSize int,
) ([]Message, error) {
type sessionLister interface {
ListSessionMessages(ctx context.Context, agentID, sessionKey, role string, limit int) ([]*memory.RecallItem, error)
}
if lister, ok := delegate.(sessionLister); ok {
count, err := delegate.CountRecallItems(ctx, agentID, sessionKey)
if err != nil {
count = 0
}
limit := count + 32
if limit < 32 {
limit = 32
}
rows, err := lister.ListSessionMessages(ctx, agentID, sessionKey, "", limit)
if err == nil {
msgs := make([]Message, 0, len(rows))
for _, row := range rows {
msgs = append(msgs, Message{Role: row.Role, Content: row.Content})
}
return msgs, nil
}
}
var all []*memory.RecallItem
offset := 0
for {
items, err := delegate.ListRecallItems(ctx, agentID, sessionKey, pageSize, offset)
if err != nil {
return nil, err
}
if len(items) == 0 {
break
}
all = append(all, items...)
if len(items) < pageSize {
break
}
offset += len(items)
}
msgs := make([]Message, 0, len(all))
for i := len(all) - 1; i >= 0; i-- {
item := all[i]
if !strings.Contains(item.Tags, "session-message") {
continue
}
msgs = append(msgs, Message{Role: item.Role, Content: item.Content})
}
return msgs, nil
}

79
pkg/memory/dag/backfill_test.go Normal file
View file

@ -0,0 +1,79 @@
package dag_test
import (
"context"
"fmt"
"testing"
"time"
jsonv2 "github.com/go-json-experiment/json"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/ZanzyTHEbar/dragonscale/pkg/ids"
"github.com/ZanzyTHEbar/dragonscale/pkg/memory"
"github.com/ZanzyTHEbar/dragonscale/pkg/memory/dag"
"github.com/ZanzyTHEbar/dragonscale/pkg/memory/delegate"
memsqlc "github.com/ZanzyTHEbar/dragonscale/pkg/memory/sqlc"
)
func TestBackfillMissingSessionDAGs_CreatesSnapshotsAndPersistsStatus(t *testing.T) {
ctx := context.Background()
d, err := delegate.NewLibSQLInMemory()
require.NoError(t, err)
require.NoError(t, d.Init(ctx))
defer d.Close()
agentID := "agent-backfill"
sessionKey := "legacy-session-1"
for i := 0; i < 12; i++ {
role := "user"
if i%2 == 1 {
role = "assistant"
}
require.NoError(t, d.InsertRecallItem(ctx, &memory.RecallItem{
ID: ids.New(),
AgentID: agentID,
SessionKey: sessionKey,
Role: role,
Sector: memory.SectorEpisodic,
Importance: 0.5,
Salience: 0.5,
DecayRate: 0.01,
Content: fmt.Sprintf("legacy message %d", i),
Tags: "session-message",
CreatedAt: time.Now().UTC(),
UpdatedAt: time.Now().UTC(),
}))
}
status, err := dag.BackfillMissingSessionDAGs(ctx, d, d.Queries(), agentID, dag.DefaultBackfillOptions())
require.NoError(t, err)
require.NotNil(t, status)
assert.Equal(t, 1, status.SnapshotsCreated)
assert.Equal(t, 0, status.Failures)
row, err := d.Queries().GetLatestDAGSnapshotBySession(ctx, memsqlc.GetLatestDAGSnapshotBySessionParams{
AgentID: agentID,
SessionKey: sessionKey,
})
require.NoError(t, err)
assert.Equal(t, int64(12), row.MsgCount)
kv, err := d.ListKVByPrefix(ctx, agentID, "migration:dag_backfill", 10)
require.NoError(t, err)
require.NotEmpty(t, kv)
var stored dag.BackfillStatus
for _, raw := range kv {
require.NoError(t, jsonv2.Unmarshal([]byte(raw), &stored))
break
}
assert.Equal(t, status.SnapshotsCreated, stored.SnapshotsCreated)
assert.False(t, stored.CompletedAt.IsZero())
status2, err := dag.BackfillMissingSessionDAGs(ctx, d, d.Queries(), agentID, dag.DefaultBackfillOptions())
require.NoError(t, err)
require.NotNil(t, status2)
assert.True(t, status2.Skipped)
}

166
pkg/memory/dag/store.go Normal file
View file

@ -0,0 +1,166 @@
package dag
import (
"context"
"crypto/sha256"
"database/sql"
"encoding/hex"
"encoding/json"
"fmt"
"sort"
"github.com/ZanzyTHEbar/dragonscale/pkg/ids"
memsqlc "github.com/ZanzyTHEbar/dragonscale/pkg/memory/sqlc"
)
// PersistSnapshot is the serializable shape for persisting a DAG snapshot.
type PersistSnapshot struct {
FromMsgIdx int
ToMsgIdx int
MsgCount int
DAG *DAG
}
// contentHash returns a deterministic hash of the DAG structure for deduplication.
func contentHash(d *DAG) string {
if d == nil || len(d.Nodes) == 0 {
return ""
}
h := sha256.New()
roots, _ := json.Marshal(d.Roots)
h.Write(roots)
nodeIDs := make([]string, 0, len(d.Nodes))
for id := range d.Nodes {
nodeIDs = append(nodeIDs, id)
}
sort.Strings(nodeIDs)
for _, id := range nodeIDs {
n := d.Nodes[id]
h.Write([]byte(id))
h.Write([]byte(n.Summary))
h.Write([]byte(fmt.Sprintf("%d:%d:%d", n.StartIdx, n.EndIdx, n.Tokens)))
if len(n.Children) > 0 {
children := append([]string(nil), n.Children...)
sort.Strings(children)
for _, child := range children {
h.Write([]byte(child))
}
}
}
return hex.EncodeToString(h.Sum(nil))
}
// DAGPersister persists DAG snapshots to storage.
// Implemented by delegates that support DAG persistence.
type DAGPersister interface {
PersistDAG(ctx context.Context, agentID, sessionKey string, snap *PersistSnapshot) error
}
// PersistDAG persists a DAG snapshot via sqlc in a transaction.
func PersistDAG(ctx context.Context, db *sql.DB, q *memsqlc.Queries, agentID, sessionKey string, snap *PersistSnapshot) error {
if snap == nil || snap.DAG == nil || len(snap.DAG.Nodes) == 0 {
return nil
}
hash := contentHash(snap.DAG)
snapshotID := ids.New()
rootsJSON, err := json.Marshal(snap.DAG.Roots)
if err != nil {
return fmt.Errorf("marshal roots: %w", err)
}
tx, err := db.BeginTx(ctx, nil)
if err != nil {
return fmt.Errorf("begin tx: %w", err)
}
defer tx.Rollback()
qTx := q.WithTx(tx)
_, err = qTx.InsertDAGSnapshot(ctx, memsqlc.InsertDAGSnapshotParams{
ID: snapshotID,
AgentID: agentID,
SessionKey: sessionKey,
FromMsgIdx: int64(snap.FromMsgIdx),
ToMsgIdx: int64(snap.ToMsgIdx),
MsgCount: int64(snap.MsgCount),
RootsJson: string(rootsJSON),
ContentHash: hash,
})
if err != nil {
return fmt.Errorf("insert snapshot: %w", err)
}
nodeIDsOrdered := make([]string, 0, len(snap.DAG.Nodes))
for nodeID := range snap.DAG.Nodes {
nodeIDsOrdered = append(nodeIDsOrdered, nodeID)
}
sort.Strings(nodeIDsOrdered)
nodeIDs := make(map[string]ids.UUID)
for _, nodeID := range nodeIDsOrdered {
node := snap.DAG.Nodes[nodeID]
nodeUUID := ids.New()
nodeIDs[nodeID] = nodeUUID
nodeHash := sha256.Sum256([]byte(fmt.Sprintf("%s|%s|%d|%d|%d", node.ID, node.Summary, node.StartIdx, node.EndIdx, node.Tokens)))
metricsJSONBytes, _ := json.Marshal(map[string]any{
"tokens": node.Tokens,
"span": node.Span(),
"children_count": len(node.Children),
})
metadataJSONBytes, _ := json.Marshal(map[string]any{
"level": node.Level.String(),
"children": node.Children,
})
_, err = qTx.InsertDAGNode(ctx, memsqlc.InsertDAGNodeParams{
ID: nodeUUID,
SnapshotID: snapshotID,
NodeID: node.ID,
Level: int64(node.Level),
Summary: node.Summary,
Tokens: int64(node.Tokens),
StartIdx: int64(node.StartIdx),
EndIdx: int64(node.EndIdx),
Span: int64(node.Span()),
ContentHash: hex.EncodeToString(nodeHash[:]),
MetricsJson: metricsJSONBytes,
MetadataJson: metadataJSONBytes,
})
if err != nil {
return fmt.Errorf("insert node %s: %w", node.ID, err)
}
}
for _, nodeID := range nodeIDsOrdered {
node := snap.DAG.Nodes[nodeID]
for edgeIdx, childID := range node.Children {
childUUID, ok := nodeIDs[childID]
if !ok {
continue
}
parentUUID := nodeIDs[node.ID]
edgeMetadata, _ := json.Marshal(map[string]any{
"parent_node_id": node.ID,
"child_node_id": childID,
})
err = qTx.InsertDAGEdge(ctx, memsqlc.InsertDAGEdgeParams{
ID: ids.New(),
SnapshotID: snapshotID,
ParentNodeID: parentUUID,
ChildNodeID: childUUID,
EdgeIndex: int64(edgeIdx),
MetadataJson: edgeMetadata,
})
if err != nil {
return fmt.Errorf("insert edge %s->%s: %w", node.ID, childID, err)
}
}
}
return tx.Commit()
}

50
pkg/memory/delegate/sqlite_dag_test.go Normal file
View file

@ -0,0 +1,50 @@
package delegate
import (
"context"
"testing"
"github.com/ZanzyTHEbar/dragonscale/pkg/memory/dag"
memsqlc "github.com/ZanzyTHEbar/dragonscale/pkg/memory/sqlc"
"github.com/stretchr/testify/require"
)
func TestLibSQLDelegate_PersistDAG(t *testing.T) {
d := newTestDelegate(t)
ctx := context.Background()
compressor := dag.NewCompressor(dag.DefaultCompressorConfig())
msgs := make([]dag.Message, 16)
for i := range msgs {
role := "user"
if i%2 == 1 {
role = "assistant"
}
msgs[i] = dag.Message{Role: role, Content: "msg " + string(rune('A'+i%26))}
}
dagOut := compressor.Compress(msgs)
require.NotEmpty(t, dagOut.Nodes)
snap := &dag.PersistSnapshot{
FromMsgIdx: 0,
ToMsgIdx: 16,
MsgCount: 16,
DAG: dagOut,
}
require.NoError(t, d.PersistDAG(ctx, "agent1", "session1", snap))
row, err := d.Queries().GetLatestDAGSnapshotBySession(ctx, memsqlc.GetLatestDAGSnapshotBySessionParams{
AgentID: "agent1",
SessionKey: "session1",
})
require.NoError(t, err)
require.Equal(t, "agent1", row.AgentID)
require.Equal(t, "session1", row.SessionKey)
require.Equal(t, int64(16), row.MsgCount)
nodes, err := d.Queries().ListDAGNodesBySnapshotID(ctx, memsqlc.ListDAGNodesBySnapshotIDParams{
SnapshotID: row.ID,
})
require.NoError(t, err)
require.NotEmpty(t, nodes)
}

95
pkg/memory/kernel_contract.go Normal file
View file

@ -0,0 +1,95 @@
package memory
import (
"context"
"time"
"github.com/ZanzyTHEbar/dragonscale/pkg/ids"
)
// ProjectionSegmentKind classifies what role a segment plays in active context.
type ProjectionSegmentKind string
const (
ProjectionSegmentSystem ProjectionSegmentKind = "system"
ProjectionSegmentRecent ProjectionSegmentKind = "recent"
ProjectionSegmentDAG ProjectionSegmentKind = "dag"
ProjectionSegmentRecall ProjectionSegmentKind = "recall"
ProjectionSegmentArchival ProjectionSegmentKind = "archival"
ProjectionSegmentTool ProjectionSegmentKind = "tool_result"
)
// ImmutableSpanRef is a lossless reference back into immutable session history.
type ImmutableSpanRef struct {
SessionKey string `json:"session_key"`
StartIdx int `json:"start_idx"` // inclusive
EndIdx int `json:"end_idx"` // exclusive
FirstID ids.UUID `json:"first_id"`
LastID ids.UUID `json:"last_id"`
FromTime time.Time `json:"from_time"`
ToTime time.Time `json:"to_time"`
}
// ProjectionSegment is one unit included in active context.
// Every segment must carry a lossless reference back to immutable history.
type ProjectionSegment struct {
Kind ProjectionSegmentKind `json:"kind"`
Source string `json:"source"`
Text string `json:"text"`
Tokens int `json:"tokens"`
Ref ImmutableSpanRef `json:"ref"`
}
// ActiveContextProjection is the assembled view injected into the model for a turn.
// It is a materialized projection, not a source of truth.
type ActiveContextProjection struct {
AgentID string `json:"agent_id"`
SessionKey string `json:"session_key"`
BudgetTokens int `json:"budget_tokens"`
GeneratedAt time.Time `json:"generated_at"`
ProjectionRef string `json:"projection_ref"`
Segments []ProjectionSegment `json:"segments"`
}
func (p *ActiveContextProjection) TotalTokens() int {
if p == nil {
return 0
}
total := 0
for _, seg := range p.Segments {
total += seg.Tokens
}
return total
}
// HasLosslessRefs verifies the projection preserves deterministic pointers to
// immutable history for every segment.
func (p *ActiveContextProjection) HasLosslessRefs() bool {
if p == nil {
return false
}
for _, seg := range p.Segments {
if seg.Ref.SessionKey == "" || seg.Ref.EndIdx < seg.Ref.StartIdx {
return false
}
}
return true
}
type ProjectionRequest struct {
AgentID string
SessionKey string
MaxTokens int
IncludeTools bool
}
// ImmutableHistoryReader resolves lossless references to original messages.
type ImmutableHistoryReader interface {
ListRecallItems(ctx context.Context, agentID, sessionKey string, limit, offset int) ([]*RecallItem, error)
}
// ActiveContextBuilder materializes the turn-time projection from immutable
// storage, retrieval tiers, and DAG summaries.
type ActiveContextBuilder interface {
BuildActiveContext(ctx context.Context, req ProjectionRequest) (*ActiveContextProjection, error)
}

47
pkg/memory/kernel_contract_test.go Normal file
View file

@ -0,0 +1,47 @@
package memory
import (
"testing"
"time"
"github.com/ZanzyTHEbar/dragonscale/pkg/ids"
"github.com/stretchr/testify/assert"
)
func TestActiveContextProjection_TotalTokens(t *testing.T) {
p := &ActiveContextProjection{
Segments: []ProjectionSegment{
{Tokens: 120},
{Tokens: 80},
{Tokens: 35},
},
}
assert.Equal(t, 235, p.TotalTokens())
}
func TestActiveContextProjection_HasLosslessRefs(t *testing.T) {
now := time.Now().UTC()
p := &ActiveContextProjection{
Segments: []ProjectionSegment{
{
Kind: ProjectionSegmentRecent,
Source: "session-tail",
Text: "latest messages",
Tokens: 64,
Ref: ImmutableSpanRef{
SessionKey: "s1",
StartIdx: 10,
EndIdx: 15,
FirstID: ids.New(),
LastID: ids.New(),
FromTime: now,
ToTime: now,
},
},
},
}
assert.True(t, p.HasLosslessRefs())
p.Segments[0].Ref.SessionKey = ""
assert.False(t, p.HasLosslessRefs())
}

86
pkg/memory/migrations/011_dag_tables.go Normal file
View file

@ -0,0 +1,86 @@
package migrations
import (
"context"
"database/sql"
"fmt"
"github.com/pressly/goose/v3"
)
func init() {
goose.AddMigrationContext(up011DAGTables, down011DAGTables)
}
func up011DAGTables(ctx context.Context, tx *sql.Tx) error {
stmts := []string{
`CREATE TABLE IF NOT EXISTS dag_snapshots (
id BLOB PRIMARY KEY,
agent_id TEXT NOT NULL,
session_key TEXT NOT NULL DEFAULT '',
from_msg_idx INTEGER NOT NULL DEFAULT 0,
to_msg_idx INTEGER NOT NULL DEFAULT 0,
msg_count INTEGER NOT NULL DEFAULT 0,
roots_json TEXT NOT NULL DEFAULT '[]',
content_hash TEXT NOT NULL DEFAULT '',
created_at DATETIME NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
updated_at DATETIME NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now'))
)`,
`CREATE INDEX IF NOT EXISTS idx_dag_snapshots_agent_session ON dag_snapshots(agent_id, session_key)`,
`CREATE INDEX IF NOT EXISTS idx_dag_snapshots_created_at ON dag_snapshots(created_at DESC)`,
`CREATE TABLE IF NOT EXISTS dag_nodes (
id BLOB PRIMARY KEY,
snapshot_id BLOB NOT NULL REFERENCES dag_snapshots(id) ON DELETE CASCADE,
node_id TEXT NOT NULL,
level INTEGER NOT NULL DEFAULT 1,
summary TEXT NOT NULL DEFAULT '',
tokens INTEGER NOT NULL DEFAULT 0,
start_idx INTEGER NOT NULL DEFAULT 0,
end_idx INTEGER NOT NULL DEFAULT 0,
span INTEGER NOT NULL DEFAULT 0,
content_hash TEXT NOT NULL DEFAULT '',
metrics_json JSON NOT NULL DEFAULT '{}',
metadata_json JSON NOT NULL DEFAULT '{}',
created_at DATETIME NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
updated_at DATETIME NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now'))
)`,
`CREATE INDEX IF NOT EXISTS idx_dag_nodes_snapshot ON dag_nodes(snapshot_id)`,
`CREATE INDEX IF NOT EXISTS idx_dag_nodes_node_id ON dag_nodes(snapshot_id, node_id)`,
`CREATE INDEX IF NOT EXISTS idx_dag_nodes_level_start ON dag_nodes(snapshot_id, level, start_idx)`,
`CREATE TABLE IF NOT EXISTS dag_edges (
id BLOB PRIMARY KEY,
snapshot_id BLOB NOT NULL REFERENCES dag_snapshots(id) ON DELETE CASCADE,
parent_node_id BLOB NOT NULL REFERENCES dag_nodes(id) ON DELETE CASCADE,
child_node_id BLOB NOT NULL REFERENCES dag_nodes(id) ON DELETE CASCADE,
edge_index INTEGER NOT NULL DEFAULT 0,
metadata_json JSON NOT NULL DEFAULT '{}',
created_at DATETIME NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
updated_at DATETIME NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now'))
)`,
`CREATE INDEX IF NOT EXISTS idx_dag_edges_snapshot ON dag_edges(snapshot_id)`,
`CREATE INDEX IF NOT EXISTS idx_dag_edges_parent ON dag_edges(parent_node_id)`,
`CREATE INDEX IF NOT EXISTS idx_dag_edges_child ON dag_edges(child_node_id)`,
}
for _, s := range stmts {
if _, err := tx.ExecContext(ctx, s); err != nil {
return fmt.Errorf("011_dag_tables up: %w\nSQL: %s", err, s)
}
}
return nil
}
func down011DAGTables(ctx context.Context, tx *sql.Tx) error {
stmts := []string{
`DROP TABLE IF EXISTS dag_edges`,
`DROP TABLE IF EXISTS dag_nodes`,
`DROP TABLE IF EXISTS dag_snapshots`,
}
for _, s := range stmts {
if _, err := tx.ExecContext(ctx, s); err != nil {
return fmt.Errorf("011_dag_tables down: %w\nSQL: %s", err, s)
}
}
return nil
}

77
pkg/memory/migrations/012_map_operator_runs.go Normal file
View file

@ -0,0 +1,77 @@
package migrations
import (
"context"
"database/sql"
"fmt"
"github.com/pressly/goose/v3"
)
func init() {
goose.AddMigrationContext(up012MapOperatorRuns, down012MapOperatorRuns)
}
func up012MapOperatorRuns(ctx context.Context, tx *sql.Tx) error {
stmts := []string{
`CREATE TABLE IF NOT EXISTS map_runs (
id BLOB PRIMARY KEY,
agent_id TEXT NOT NULL,
session_key TEXT NOT NULL,
operator_kind TEXT NOT NULL,
idempotency_key TEXT,
status TEXT NOT NULL DEFAULT 'queued',
total_items INTEGER NOT NULL DEFAULT 0,
queued_items INTEGER NOT NULL DEFAULT 0,
running_items INTEGER NOT NULL DEFAULT 0,
succeeded_items INTEGER NOT NULL DEFAULT 0,
failed_items INTEGER NOT NULL DEFAULT 0,
spec_fb BLOB NOT NULL,
last_error TEXT,
created_at DATETIME NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
updated_at DATETIME NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
completed_at DATETIME
)`,
`CREATE INDEX IF NOT EXISTS idx_map_runs_agent_session_created_at ON map_runs(agent_id, session_key, created_at DESC)`,
`CREATE INDEX IF NOT EXISTS idx_map_runs_status_updated_at ON map_runs(status, updated_at DESC)`,
`CREATE UNIQUE INDEX IF NOT EXISTS idx_map_runs_dedupe ON map_runs(agent_id, session_key, operator_kind, idempotency_key) WHERE idempotency_key IS NOT NULL`,
`CREATE TABLE IF NOT EXISTS map_items (
id BLOB PRIMARY KEY,
run_id BLOB NOT NULL REFERENCES map_runs(id) ON DELETE CASCADE,
item_index INTEGER NOT NULL,
status TEXT NOT NULL DEFAULT 'queued',
attempts INTEGER NOT NULL DEFAULT 0,
last_error TEXT,
input_fb BLOB NOT NULL,
output_fb BLOB,
input_hash TEXT,
output_hash TEXT,
created_at DATETIME NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
updated_at DATETIME NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
completed_at DATETIME,
UNIQUE(run_id, item_index)
)`,
`CREATE INDEX IF NOT EXISTS idx_map_items_run_id_status_item_index ON map_items(run_id, status, item_index)`,
`CREATE INDEX IF NOT EXISTS idx_map_items_run_id_item_index ON map_items(run_id, item_index)`,
}
for _, s := range stmts {
if _, err := tx.ExecContext(ctx, s); err != nil {
return fmt.Errorf("012_map_operator_runs up: %w\nSQL: %s", err, s)
}
}
return nil
}
func down012MapOperatorRuns(ctx context.Context, tx *sql.Tx) error {
stmts := []string{
`DROP TABLE IF EXISTS map_items`,
`DROP TABLE IF EXISTS map_runs`,
}
for _, s := range stmts {
if _, err := tx.ExecContext(ctx, s); err != nil {
return fmt.Errorf("012_map_operator_runs down: %w\nSQL: %s", err, s)
}
}
return nil
}

638
pkg/memory/sqlc/dag.sql.go Normal file
View file

@ -0,0 +1,638 @@
// Code generated by sqlc. DO NOT EDIT.
// versions:
// sqlc v1.30.0
// source: dag.sql
package sqlc
import (
"context"
"encoding/json"
"github.com/ZanzyTHEbar/dragonscale/pkg/ids"
)
const GetDAGNodeBySnapshotAndNodeID = `-- name: GetDAGNodeBySnapshotAndNodeID :one
SELECT id,
snapshot_id,
node_id,
level,
summary,
tokens,
start_idx,
end_idx,
span,
content_hash,
metrics_json,
metadata_json,
created_at,
updated_at
FROM dag_nodes
WHERE snapshot_id = ?1
AND node_id = ?2
LIMIT 1
`
type GetDAGNodeBySnapshotAndNodeIDParams struct {
SnapshotID ids.UUID `db:"snapshot_id" json:"snapshot_id"`
NodeID string `db:"node_id" json:"node_id"`
}
// GetDAGNodeBySnapshotAndNodeID
//
// SELECT id,
// snapshot_id,
// node_id,
// level,
// summary,
// tokens,
// start_idx,
// end_idx,
// span,
// content_hash,
// metrics_json,
// metadata_json,
// created_at,
// updated_at
// FROM dag_nodes
// WHERE snapshot_id = ?1
// AND node_id = ?2
// LIMIT 1
func (q *Queries) GetDAGNodeBySnapshotAndNodeID(ctx context.Context, arg GetDAGNodeBySnapshotAndNodeIDParams) (DagNode, error) {
row := q.db.QueryRowContext(ctx, GetDAGNodeBySnapshotAndNodeID, arg.SnapshotID, arg.NodeID)
var i DagNode
err := row.Scan(
&i.ID,
&i.SnapshotID,
&i.NodeID,
&i.Level,
&i.Summary,
&i.Tokens,
&i.StartIdx,
&i.EndIdx,
&i.Span,
&i.ContentHash,
&i.MetricsJson,
&i.MetadataJson,
&i.CreatedAt,
&i.UpdatedAt,
)
return i, err
}
const GetDAGSnapshotByID = `-- name: GetDAGSnapshotByID :one
SELECT id,
agent_id,
session_key,
from_msg_idx,
to_msg_idx,
msg_count,
roots_json,
content_hash,
created_at,
updated_at
FROM dag_snapshots
WHERE id = ?1
LIMIT 1
`
type GetDAGSnapshotByIDParams struct {
ID ids.UUID `db:"id" json:"id"`
}
// GetDAGSnapshotByID
//
// SELECT id,
// agent_id,
// session_key,
// from_msg_idx,
// to_msg_idx,
// msg_count,
// roots_json,
// content_hash,
// created_at,
// updated_at
// FROM dag_snapshots
// WHERE id = ?1
// LIMIT 1
func (q *Queries) GetDAGSnapshotByID(ctx context.Context, arg GetDAGSnapshotByIDParams) (DagSnapshot, error) {
row := q.db.QueryRowContext(ctx, GetDAGSnapshotByID, arg.ID)
var i DagSnapshot
err := row.Scan(
&i.ID,
&i.AgentID,
&i.SessionKey,
&i.FromMsgIdx,
&i.ToMsgIdx,
&i.MsgCount,
&i.RootsJson,
&i.ContentHash,
&i.CreatedAt,
&i.UpdatedAt,
)
return i, err
}
const GetLatestDAGSnapshotBySession = `-- name: GetLatestDAGSnapshotBySession :one
SELECT id,
agent_id,
session_key,
from_msg_idx,
to_msg_idx,
msg_count,
roots_json,
content_hash,
created_at,
updated_at
FROM dag_snapshots
WHERE agent_id = ?1
AND session_key = ?2
ORDER BY created_at DESC
LIMIT 1
`
type GetLatestDAGSnapshotBySessionParams struct {
AgentID string `db:"agent_id" json:"agent_id"`
SessionKey string `db:"session_key" json:"session_key"`
}
// GetLatestDAGSnapshotBySession
//
// SELECT id,
// agent_id,
// session_key,
// from_msg_idx,
// to_msg_idx,
// msg_count,
// roots_json,
// content_hash,
// created_at,
// updated_at
// FROM dag_snapshots
// WHERE agent_id = ?1
// AND session_key = ?2
// ORDER BY created_at DESC
// LIMIT 1
func (q *Queries) GetLatestDAGSnapshotBySession(ctx context.Context, arg GetLatestDAGSnapshotBySessionParams) (DagSnapshot, error) {
row := q.db.QueryRowContext(ctx, GetLatestDAGSnapshotBySession, arg.AgentID, arg.SessionKey)
var i DagSnapshot
err := row.Scan(
&i.ID,
&i.AgentID,
&i.SessionKey,
&i.FromMsgIdx,
&i.ToMsgIdx,
&i.MsgCount,
&i.RootsJson,
&i.ContentHash,
&i.CreatedAt,
&i.UpdatedAt,
)
return i, err
}
const InsertDAGEdge = `-- name: InsertDAGEdge :exec
INSERT INTO dag_edges (
id,
snapshot_id,
parent_node_id,
child_node_id,
edge_index,
metadata_json
)
VALUES (
?1,
?2,
?3,
?4,
?5,
?6
)
`
type InsertDAGEdgeParams struct {
ID ids.UUID `db:"id" json:"id"`
SnapshotID ids.UUID `db:"snapshot_id" json:"snapshot_id"`
ParentNodeID ids.UUID `db:"parent_node_id" json:"parent_node_id"`
ChildNodeID ids.UUID `db:"child_node_id" json:"child_node_id"`
EdgeIndex int64 `db:"edge_index" json:"edge_index"`
MetadataJson json.RawMessage `db:"metadata_json" json:"metadata_json"`
}
// InsertDAGEdge
//
// INSERT INTO dag_edges (
// id,
// snapshot_id,
// parent_node_id,
// child_node_id,
// edge_index,
// metadata_json
// )
// VALUES (
// ?1,
// ?2,
// ?3,
// ?4,
// ?5,
// ?6
// )
func (q *Queries) InsertDAGEdge(ctx context.Context, arg InsertDAGEdgeParams) error {
_, err := q.db.ExecContext(ctx, InsertDAGEdge,
arg.ID,
arg.SnapshotID,
arg.ParentNodeID,
arg.ChildNodeID,
arg.EdgeIndex,
arg.MetadataJson,
)
return err
}
const InsertDAGNode = `-- name: InsertDAGNode :one
INSERT INTO dag_nodes (
id,
snapshot_id,
node_id,
level,
summary,
tokens,
start_idx,
end_idx,
span,
content_hash,
metrics_json,
metadata_json
)
VALUES (
?1,
?2,
?3,
?4,
?5,
?6,
?7,
?8,
?9,
?10,
?11,
?12
)
RETURNING id,
snapshot_id,
node_id,
level,
summary,
tokens,
start_idx,
end_idx,
span,
content_hash,
metrics_json,
metadata_json,
created_at,
updated_at
`
type InsertDAGNodeParams struct {
ID ids.UUID `db:"id" json:"id"`
SnapshotID ids.UUID `db:"snapshot_id" json:"snapshot_id"`
NodeID string `db:"node_id" json:"node_id"`
Level int64 `db:"level" json:"level"`
Summary string `db:"summary" json:"summary"`
Tokens int64 `db:"tokens" json:"tokens"`
StartIdx int64 `db:"start_idx" json:"start_idx"`
EndIdx int64 `db:"end_idx" json:"end_idx"`
Span int64 `db:"span" json:"span"`
ContentHash string `db:"content_hash" json:"content_hash"`
MetricsJson json.RawMessage `db:"metrics_json" json:"metrics_json"`
MetadataJson json.RawMessage `db:"metadata_json" json:"metadata_json"`
}
// InsertDAGNode
//
// INSERT INTO dag_nodes (
// id,
// snapshot_id,
// node_id,
// level,
// summary,
// tokens,
// start_idx,
// end_idx,
// span,
// content_hash,
// metrics_json,
// metadata_json
// )
// VALUES (
// ?1,
// ?2,
// ?3,
// ?4,
// ?5,
// ?6,
// ?7,
// ?8,
// ?9,
// ?10,
// ?11,
// ?12
// )
// RETURNING id,
// snapshot_id,
// node_id,
// level,
// summary,
// tokens,
// start_idx,
// end_idx,
// span,
// content_hash,
// metrics_json,
// metadata_json,
// created_at,
// updated_at
func (q *Queries) InsertDAGNode(ctx context.Context, arg InsertDAGNodeParams) (DagNode, error) {
row := q.db.QueryRowContext(ctx, InsertDAGNode,
arg.ID,
arg.SnapshotID,
arg.NodeID,
arg.Level,
arg.Summary,
arg.Tokens,
arg.StartIdx,
arg.EndIdx,
arg.Span,
arg.ContentHash,
arg.MetricsJson,
arg.MetadataJson,
)
var i DagNode
err := row.Scan(
&i.ID,
&i.SnapshotID,
&i.NodeID,
&i.Level,
&i.Summary,
&i.Tokens,
&i.StartIdx,
&i.EndIdx,
&i.Span,
&i.ContentHash,
&i.MetricsJson,
&i.MetadataJson,
&i.CreatedAt,
&i.UpdatedAt,
)
return i, err
}
const InsertDAGSnapshot = `-- name: InsertDAGSnapshot :one
INSERT INTO dag_snapshots (
id,
agent_id,
session_key,
from_msg_idx,
to_msg_idx,
msg_count,
roots_json,
content_hash
)
VALUES (
?1,
?2,
?3,
?4,
?5,
?6,
?7,
?8
)
RETURNING id,
agent_id,
session_key,
from_msg_idx,
to_msg_idx,
msg_count,
roots_json,
content_hash,
created_at,
updated_at
`
type InsertDAGSnapshotParams struct {
ID ids.UUID `db:"id" json:"id"`
AgentID string `db:"agent_id" json:"agent_id"`
SessionKey string `db:"session_key" json:"session_key"`
FromMsgIdx int64 `db:"from_msg_idx" json:"from_msg_idx"`
ToMsgIdx int64 `db:"to_msg_idx" json:"to_msg_idx"`
MsgCount int64 `db:"msg_count" json:"msg_count"`
RootsJson string `db:"roots_json" json:"roots_json"`
ContentHash string `db:"content_hash" json:"content_hash"`
}
// DAG persistence queries
//
// INSERT INTO dag_snapshots (
// id,
// agent_id,
// session_key,
// from_msg_idx,
// to_msg_idx,
// msg_count,
// roots_json,
// content_hash
// )
// VALUES (
// ?1,
// ?2,
// ?3,
// ?4,
// ?5,
// ?6,
// ?7,
// ?8
// )
// RETURNING id,
// agent_id,
// session_key,
// from_msg_idx,
// to_msg_idx,
// msg_count,
// roots_json,
// content_hash,
// created_at,
// updated_at
func (q *Queries) InsertDAGSnapshot(ctx context.Context, arg InsertDAGSnapshotParams) (DagSnapshot, error) {
row := q.db.QueryRowContext(ctx, InsertDAGSnapshot,
arg.ID,
arg.AgentID,
arg.SessionKey,
arg.FromMsgIdx,
arg.ToMsgIdx,
arg.MsgCount,
arg.RootsJson,
arg.ContentHash,
)
var i DagSnapshot
err := row.Scan(
&i.ID,
&i.AgentID,
&i.SessionKey,
&i.FromMsgIdx,
&i.ToMsgIdx,
&i.MsgCount,
&i.RootsJson,
&i.ContentHash,
&i.CreatedAt,
&i.UpdatedAt,
)
return i, err
}
const ListDAGEdgesBySnapshotID = `-- name: ListDAGEdgesBySnapshotID :many
SELECT id,
snapshot_id,
parent_node_id,
child_node_id,
edge_index,
metadata_json,
created_at,
updated_at
FROM dag_edges
WHERE snapshot_id = ?1
ORDER BY edge_index ASC,
created_at ASC
`
type ListDAGEdgesBySnapshotIDParams struct {
SnapshotID ids.UUID `db:"snapshot_id" json:"snapshot_id"`
}
// ListDAGEdgesBySnapshotID
//
// SELECT id,
// snapshot_id,
// parent_node_id,
// child_node_id,
// edge_index,
// metadata_json,
// created_at,
// updated_at
// FROM dag_edges
// WHERE snapshot_id = ?1
// ORDER BY edge_index ASC,
// created_at ASC
func (q *Queries) ListDAGEdgesBySnapshotID(ctx context.Context, arg ListDAGEdgesBySnapshotIDParams) ([]DagEdge, error) {
rows, err := q.db.QueryContext(ctx, ListDAGEdgesBySnapshotID, arg.SnapshotID)
if err != nil {
return nil, err
}
defer rows.Close()
items := []DagEdge{}
for rows.Next() {
var i DagEdge
if err := rows.Scan(
&i.ID,
&i.SnapshotID,
&i.ParentNodeID,
&i.ChildNodeID,
&i.EdgeIndex,
&i.MetadataJson,
&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 ListDAGNodesBySnapshotID = `-- name: ListDAGNodesBySnapshotID :many
SELECT id,
snapshot_id,
node_id,
level,
summary,
tokens,
start_idx,
end_idx,
span,
content_hash,
metrics_json,
metadata_json,
created_at,
updated_at
FROM dag_nodes
WHERE snapshot_id = ?1
ORDER BY start_idx ASC
`
type ListDAGNodesBySnapshotIDParams struct {
SnapshotID ids.UUID `db:"snapshot_id" json:"snapshot_id"`
}
// ListDAGNodesBySnapshotID
//
// SELECT id,
// snapshot_id,
// node_id,
// level,
// summary,
// tokens,
// start_idx,
// end_idx,
// span,
// content_hash,
// metrics_json,
// metadata_json,
// created_at,
// updated_at
// FROM dag_nodes
// WHERE snapshot_id = ?1
// ORDER BY start_idx ASC
func (q *Queries) ListDAGNodesBySnapshotID(ctx context.Context, arg ListDAGNodesBySnapshotIDParams) ([]DagNode, error) {
rows, err := q.db.QueryContext(ctx, ListDAGNodesBySnapshotID, arg.SnapshotID)
if err != nil {
return nil, err
}
defer rows.Close()
items := []DagNode{}
for rows.Next() {
var i DagNode
if err := rows.Scan(
&i.ID,
&i.SnapshotID,
&i.NodeID,
&i.Level,
&i.Summary,
&i.Tokens,
&i.StartIdx,
&i.EndIdx,
&i.Span,
&i.ContentHash,
&i.MetricsJson,
&i.MetadataJson,
&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
}

815
pkg/memory/sqlc/map_ops.sql.go Normal file
View file

@ -0,0 +1,815 @@
// Code generated by sqlc. DO NOT EDIT.
// versions:
// sqlc v1.30.0
// source: map_ops.sql
package sqlc
import (
"context"
"time"
"github.com/ZanzyTHEbar/dragonscale/pkg/ids"
)
const CountMapItemsByRun = `-- name: CountMapItemsByRun :one
SELECT count(*) AS count
FROM map_items
WHERE run_id = ?1
`
type CountMapItemsByRunParams struct {
RunID ids.UUID `db:"run_id" json:"run_id"`
}
// CountMapItemsByRun
//
// SELECT count(*) AS count
// FROM map_items
// WHERE run_id = ?1
func (q *Queries) CountMapItemsByRun(ctx context.Context, arg CountMapItemsByRunParams) (int64, error) {
row := q.db.QueryRowContext(ctx, CountMapItemsByRun, arg.RunID)
var count int64
err := row.Scan(&count)
return count, err
}
const CountMapItemsByRunAndStatus = `-- name: CountMapItemsByRunAndStatus :many
SELECT status,
count(*) AS count
FROM map_items
WHERE run_id = ?1
GROUP BY status
`
type CountMapItemsByRunAndStatusParams struct {
RunID ids.UUID `db:"run_id" json:"run_id"`
}
type CountMapItemsByRunAndStatusRow struct {
Status string `db:"status" json:"status"`
Count int64 `db:"count" json:"count"`
}
// CountMapItemsByRunAndStatus
//
// SELECT status,
// count(*) AS count
// FROM map_items
// WHERE run_id = ?1
// GROUP BY status
func (q *Queries) CountMapItemsByRunAndStatus(ctx context.Context, arg CountMapItemsByRunAndStatusParams) ([]CountMapItemsByRunAndStatusRow, error) {
rows, err := q.db.QueryContext(ctx, CountMapItemsByRunAndStatus, arg.RunID)
if err != nil {
return nil, err
}
defer rows.Close()
items := []CountMapItemsByRunAndStatusRow{}
for rows.Next() {
var i CountMapItemsByRunAndStatusRow
if err := rows.Scan(&i.Status, &i.Count); 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 GetMapItemByID = `-- name: GetMapItemByID :one
SELECT id, run_id, item_index, status, attempts, last_error, input_fb, output_fb, input_hash, output_hash, created_at, updated_at, completed_at
FROM map_items
WHERE id = ?1
LIMIT 1
`
type GetMapItemByIDParams struct {
ID ids.UUID `db:"id" json:"id"`
}
// GetMapItemByID
//
// SELECT id, run_id, item_index, status, attempts, last_error, input_fb, output_fb, input_hash, output_hash, created_at, updated_at, completed_at
// FROM map_items
// WHERE id = ?1
// LIMIT 1
func (q *Queries) GetMapItemByID(ctx context.Context, arg GetMapItemByIDParams) (MapItem, error) {
row := q.db.QueryRowContext(ctx, GetMapItemByID, arg.ID)
var i MapItem
err := row.Scan(
&i.ID,
&i.RunID,
&i.ItemIndex,
&i.Status,
&i.Attempts,
&i.LastError,
&i.InputFb,
&i.OutputFb,
&i.InputHash,
&i.OutputHash,
&i.CreatedAt,
&i.UpdatedAt,
&i.CompletedAt,
)
return i, err
}
const GetMapItemByRunAndIndex = `-- name: GetMapItemByRunAndIndex :one
SELECT id, run_id, item_index, status, attempts, last_error, input_fb, output_fb, input_hash, output_hash, created_at, updated_at, completed_at
FROM map_items
WHERE run_id = ?1
AND item_index = ?2
LIMIT 1
`
type GetMapItemByRunAndIndexParams struct {
RunID ids.UUID `db:"run_id" json:"run_id"`
ItemIndex int64 `db:"item_index" json:"item_index"`
}
// GetMapItemByRunAndIndex
//
// SELECT id, run_id, item_index, status, attempts, last_error, input_fb, output_fb, input_hash, output_hash, created_at, updated_at, completed_at
// FROM map_items
// WHERE run_id = ?1
// AND item_index = ?2
// LIMIT 1
func (q *Queries) GetMapItemByRunAndIndex(ctx context.Context, arg GetMapItemByRunAndIndexParams) (MapItem, error) {
row := q.db.QueryRowContext(ctx, GetMapItemByRunAndIndex, arg.RunID, arg.ItemIndex)
var i MapItem
err := row.Scan(
&i.ID,
&i.RunID,
&i.ItemIndex,
&i.Status,
&i.Attempts,
&i.LastError,
&i.InputFb,
&i.OutputFb,
&i.InputHash,
&i.OutputHash,
&i.CreatedAt,
&i.UpdatedAt,
&i.CompletedAt,
)
return i, err
}
const GetMapRunByID = `-- name: GetMapRunByID :one
SELECT id, agent_id, session_key, operator_kind, idempotency_key, status, total_items, queued_items, running_items, succeeded_items, failed_items, spec_fb, last_error, created_at, updated_at, completed_at
FROM map_runs
WHERE id = ?1
LIMIT 1
`
type GetMapRunByIDParams struct {
ID ids.UUID `db:"id" json:"id"`
}
// GetMapRunByID
//
// SELECT id, agent_id, session_key, operator_kind, idempotency_key, status, total_items, queued_items, running_items, succeeded_items, failed_items, spec_fb, last_error, created_at, updated_at, completed_at
// FROM map_runs
// WHERE id = ?1
// LIMIT 1
func (q *Queries) GetMapRunByID(ctx context.Context, arg GetMapRunByIDParams) (MapRun, error) {
row := q.db.QueryRowContext(ctx, GetMapRunByID, arg.ID)
var i MapRun
err := row.Scan(
&i.ID,
&i.AgentID,
&i.SessionKey,
&i.OperatorKind,
&i.IdempotencyKey,
&i.Status,
&i.TotalItems,
&i.QueuedItems,
&i.RunningItems,
&i.SucceededItems,
&i.FailedItems,
&i.SpecFb,
&i.LastError,
&i.CreatedAt,
&i.UpdatedAt,
&i.CompletedAt,
)
return i, err
}
const GetMapRunByIdempotencyKey = `-- name: GetMapRunByIdempotencyKey :one
SELECT id, agent_id, session_key, operator_kind, idempotency_key, status, total_items, queued_items, running_items, succeeded_items, failed_items, spec_fb, last_error, created_at, updated_at, completed_at
FROM map_runs
WHERE agent_id = ?1
AND session_key = ?2
AND operator_kind = ?3
AND idempotency_key = ?4
LIMIT 1
`
type GetMapRunByIdempotencyKeyParams struct {
AgentID string `db:"agent_id" json:"agent_id"`
SessionKey string `db:"session_key" json:"session_key"`
OperatorKind string `db:"operator_kind" json:"operator_kind"`
IdempotencyKey *string `db:"idempotency_key" json:"idempotency_key"`
}
// GetMapRunByIdempotencyKey
//
// SELECT id, agent_id, session_key, operator_kind, idempotency_key, status, total_items, queued_items, running_items, succeeded_items, failed_items, spec_fb, last_error, created_at, updated_at, completed_at
// FROM map_runs
// WHERE agent_id = ?1
// AND session_key = ?2
// AND operator_kind = ?3
// AND idempotency_key = ?4
// LIMIT 1
func (q *Queries) GetMapRunByIdempotencyKey(ctx context.Context, arg GetMapRunByIdempotencyKeyParams) (MapRun, error) {
row := q.db.QueryRowContext(ctx, GetMapRunByIdempotencyKey,
arg.AgentID,
arg.SessionKey,
arg.OperatorKind,
arg.IdempotencyKey,
)
var i MapRun
err := row.Scan(
&i.ID,
&i.AgentID,
&i.SessionKey,
&i.OperatorKind,
&i.IdempotencyKey,
&i.Status,
&i.TotalItems,
&i.QueuedItems,
&i.RunningItems,
&i.SucceededItems,
&i.FailedItems,
&i.SpecFb,
&i.LastError,
&i.CreatedAt,
&i.UpdatedAt,
&i.CompletedAt,
)
return i, err
}
const InsertMapItem = `-- name: InsertMapItem :one
INSERT INTO map_items (
id,
run_id,
item_index,
status,
attempts,
last_error,
input_fb,
output_fb,
input_hash,
output_hash
)
VALUES (
?1,
?2,
?3,
?4,
?5,
?6,
?7,
?8,
?9,
?10
)
RETURNING id, run_id, item_index, status, attempts, last_error, input_fb, output_fb, input_hash, output_hash, created_at, updated_at, completed_at
`
type InsertMapItemParams struct {
ID ids.UUID `db:"id" json:"id"`
RunID ids.UUID `db:"run_id" json:"run_id"`
ItemIndex int64 `db:"item_index" json:"item_index"`
Status string `db:"status" json:"status"`
Attempts int64 `db:"attempts" json:"attempts"`
LastError *string `db:"last_error" json:"last_error"`
InputFb []byte `db:"input_fb" json:"input_fb"`
OutputFb []byte `db:"output_fb" json:"output_fb"`
InputHash *string `db:"input_hash" json:"input_hash"`
OutputHash *string `db:"output_hash" json:"output_hash"`
}
// InsertMapItem
//
// INSERT INTO map_items (
// id,
// run_id,
// item_index,
// status,
// attempts,
// last_error,
// input_fb,
// output_fb,
// input_hash,
// output_hash
// )
// VALUES (
// ?1,
// ?2,
// ?3,
// ?4,
// ?5,
// ?6,
// ?7,
// ?8,
// ?9,
// ?10
// )
// RETURNING id, run_id, item_index, status, attempts, last_error, input_fb, output_fb, input_hash, output_hash, created_at, updated_at, completed_at
func (q *Queries) InsertMapItem(ctx context.Context, arg InsertMapItemParams) (MapItem, error) {
row := q.db.QueryRowContext(ctx, InsertMapItem,
arg.ID,
arg.RunID,
arg.ItemIndex,
arg.Status,
arg.Attempts,
arg.LastError,
arg.InputFb,
arg.OutputFb,
arg.InputHash,
arg.OutputHash,
)
var i MapItem
err := row.Scan(
&i.ID,
&i.RunID,
&i.ItemIndex,
&i.Status,
&i.Attempts,
&i.LastError,
&i.InputFb,
&i.OutputFb,
&i.InputHash,
&i.OutputHash,
&i.CreatedAt,
&i.UpdatedAt,
&i.CompletedAt,
)
return i, err
}
const InsertMapRun = `-- name: InsertMapRun :one
INSERT INTO map_runs (
id,
agent_id,
session_key,
operator_kind,
idempotency_key,
status,
total_items,
queued_items,
running_items,
succeeded_items,
failed_items,
spec_fb,
last_error
)
VALUES (
?1,
?2,
?3,
?4,
?5,
?6,
?7,
?8,
?9,
?10,
?11,
?12,
?13
)
RETURNING id, agent_id, session_key, operator_kind, idempotency_key, status, total_items, queued_items, running_items, succeeded_items, failed_items, spec_fb, last_error, created_at, updated_at, completed_at
`
type InsertMapRunParams struct {
ID ids.UUID `db:"id" json:"id"`
AgentID string `db:"agent_id" json:"agent_id"`
SessionKey string `db:"session_key" json:"session_key"`
OperatorKind string `db:"operator_kind" json:"operator_kind"`
IdempotencyKey *string `db:"idempotency_key" json:"idempotency_key"`
Status string `db:"status" json:"status"`
TotalItems int64 `db:"total_items" json:"total_items"`
QueuedItems int64 `db:"queued_items" json:"queued_items"`
RunningItems int64 `db:"running_items" json:"running_items"`
SucceededItems int64 `db:"succeeded_items" json:"succeeded_items"`
FailedItems int64 `db:"failed_items" json:"failed_items"`
SpecFb []byte `db:"spec_fb" json:"spec_fb"`
LastError *string `db:"last_error" json:"last_error"`
}
// InsertMapRun
//
// INSERT INTO map_runs (
// id,
// agent_id,
// session_key,
// operator_kind,
// idempotency_key,
// status,
// total_items,
// queued_items,
// running_items,
// succeeded_items,
// failed_items,
// spec_fb,
// last_error
// )
// VALUES (
// ?1,
// ?2,
// ?3,
// ?4,
// ?5,
// ?6,
// ?7,
// ?8,
// ?9,
// ?10,
// ?11,
// ?12,
// ?13
// )
// RETURNING id, agent_id, session_key, operator_kind, idempotency_key, status, total_items, queued_items, running_items, succeeded_items, failed_items, spec_fb, last_error, created_at, updated_at, completed_at
func (q *Queries) InsertMapRun(ctx context.Context, arg InsertMapRunParams) (MapRun, error) {
row := q.db.QueryRowContext(ctx, InsertMapRun,
arg.ID,
arg.AgentID,
arg.SessionKey,
arg.OperatorKind,
arg.IdempotencyKey,
arg.Status,
arg.TotalItems,
arg.QueuedItems,
arg.RunningItems,
arg.SucceededItems,
arg.FailedItems,
arg.SpecFb,
arg.LastError,
)
var i MapRun
err := row.Scan(
&i.ID,
&i.AgentID,
&i.SessionKey,
&i.OperatorKind,
&i.IdempotencyKey,
&i.Status,
&i.TotalItems,
&i.QueuedItems,
&i.RunningItems,
&i.SucceededItems,
&i.FailedItems,
&i.SpecFb,
&i.LastError,
&i.CreatedAt,
&i.UpdatedAt,
&i.CompletedAt,
)
return i, err
}
const ListMapItemsByRunPaged = `-- name: ListMapItemsByRunPaged :many
SELECT id, run_id, item_index, status, attempts, last_error, input_fb, output_fb, input_hash, output_hash, created_at, updated_at, completed_at
FROM map_items
WHERE run_id = ?1
ORDER BY item_index ASC
LIMIT ?3 OFFSET ?2
`
type ListMapItemsByRunPagedParams struct {
RunID ids.UUID `db:"run_id" json:"run_id"`
Off int64 `db:"off" json:"off"`
Lim int64 `db:"lim" json:"lim"`
}
// ListMapItemsByRunPaged
//
// SELECT id, run_id, item_index, status, attempts, last_error, input_fb, output_fb, input_hash, output_hash, created_at, updated_at, completed_at
// FROM map_items
// WHERE run_id = ?1
// ORDER BY item_index ASC
// LIMIT ?3 OFFSET ?2
func (q *Queries) ListMapItemsByRunPaged(ctx context.Context, arg ListMapItemsByRunPagedParams) ([]MapItem, error) {
rows, err := q.db.QueryContext(ctx, ListMapItemsByRunPaged, arg.RunID, arg.Off, arg.Lim)
if err != nil {
return nil, err
}
defer rows.Close()
items := []MapItem{}
for rows.Next() {
var i MapItem
if err := rows.Scan(
&i.ID,
&i.RunID,
&i.ItemIndex,
&i.Status,
&i.Attempts,
&i.LastError,
&i.InputFb,
&i.OutputFb,
&i.InputHash,
&i.OutputHash,
&i.CreatedAt,
&i.UpdatedAt,
&i.CompletedAt,
); 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 ListMapRunsBySession = `-- name: ListMapRunsBySession :many
SELECT id, agent_id, session_key, operator_kind, idempotency_key, status, total_items, queued_items, running_items, succeeded_items, failed_items, spec_fb, last_error, created_at, updated_at, completed_at
FROM map_runs
WHERE agent_id = ?1
AND session_key = ?2
ORDER BY created_at DESC
LIMIT ?4 OFFSET ?3
`
type ListMapRunsBySessionParams struct {
AgentID string `db:"agent_id" json:"agent_id"`
SessionKey string `db:"session_key" json:"session_key"`
Off int64 `db:"off" json:"off"`
Lim int64 `db:"lim" json:"lim"`
}
// ListMapRunsBySession
//
// SELECT id, agent_id, session_key, operator_kind, idempotency_key, status, total_items, queued_items, running_items, succeeded_items, failed_items, spec_fb, last_error, created_at, updated_at, completed_at
// FROM map_runs
// WHERE agent_id = ?1
// AND session_key = ?2
// ORDER BY created_at DESC
// LIMIT ?4 OFFSET ?3
func (q *Queries) ListMapRunsBySession(ctx context.Context, arg ListMapRunsBySessionParams) ([]MapRun, error) {
rows, err := q.db.QueryContext(ctx, ListMapRunsBySession,
arg.AgentID,
arg.SessionKey,
arg.Off,
arg.Lim,
)
if err != nil {
return nil, err
}
defer rows.Close()
items := []MapRun{}
for rows.Next() {
var i MapRun
if err := rows.Scan(
&i.ID,
&i.AgentID,
&i.SessionKey,
&i.OperatorKind,
&i.IdempotencyKey,
&i.Status,
&i.TotalItems,
&i.QueuedItems,
&i.RunningItems,
&i.SucceededItems,
&i.FailedItems,
&i.SpecFb,
&i.LastError,
&i.CreatedAt,
&i.UpdatedAt,
&i.CompletedAt,
); 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 MarkMapItemFailed = `-- name: MarkMapItemFailed :one
UPDATE map_items
SET status = 'failed',
last_error = ?1,
updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now'),
completed_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now')
WHERE id = ?2
RETURNING id, run_id, item_index, status, attempts, last_error, input_fb, output_fb, input_hash, output_hash, created_at, updated_at, completed_at
`
type MarkMapItemFailedParams struct {
LastError *string `db:"last_error" json:"last_error"`
ID ids.UUID `db:"id" json:"id"`
}
// MarkMapItemFailed
//
// UPDATE map_items
// SET status = 'failed',
// last_error = ?1,
// updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now'),
// completed_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now')
// WHERE id = ?2
// RETURNING id, run_id, item_index, status, attempts, last_error, input_fb, output_fb, input_hash, output_hash, created_at, updated_at, completed_at
func (q *Queries) MarkMapItemFailed(ctx context.Context, arg MarkMapItemFailedParams) (MapItem, error) {
row := q.db.QueryRowContext(ctx, MarkMapItemFailed, arg.LastError, arg.ID)
var i MapItem
err := row.Scan(
&i.ID,
&i.RunID,
&i.ItemIndex,
&i.Status,
&i.Attempts,
&i.LastError,
&i.InputFb,
&i.OutputFb,
&i.InputHash,
&i.OutputHash,
&i.CreatedAt,
&i.UpdatedAt,
&i.CompletedAt,
)
return i, err
}
const MarkMapItemRunning = `-- name: MarkMapItemRunning :one
UPDATE map_items
SET status = 'running',
attempts = attempts + 1,
last_error = NULL,
updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now')
WHERE id = ?1
RETURNING id, run_id, item_index, status, attempts, last_error, input_fb, output_fb, input_hash, output_hash, created_at, updated_at, completed_at
`
type MarkMapItemRunningParams struct {
ID ids.UUID `db:"id" json:"id"`
}
// MarkMapItemRunning
//
// UPDATE map_items
// SET status = 'running',
// attempts = attempts + 1,
// last_error = NULL,
// updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now')
// WHERE id = ?1
// RETURNING id, run_id, item_index, status, attempts, last_error, input_fb, output_fb, input_hash, output_hash, created_at, updated_at, completed_at
func (q *Queries) MarkMapItemRunning(ctx context.Context, arg MarkMapItemRunningParams) (MapItem, error) {
row := q.db.QueryRowContext(ctx, MarkMapItemRunning, arg.ID)
var i MapItem
err := row.Scan(
&i.ID,
&i.RunID,
&i.ItemIndex,
&i.Status,
&i.Attempts,
&i.LastError,
&i.InputFb,
&i.OutputFb,
&i.InputHash,
&i.OutputHash,
&i.CreatedAt,
&i.UpdatedAt,
&i.CompletedAt,
)
return i, err
}
const MarkMapItemSucceeded = `-- name: MarkMapItemSucceeded :one
UPDATE map_items
SET status = 'succeeded',
output_fb = ?1,
output_hash = ?2,
last_error = NULL,
updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now'),
completed_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now')
WHERE id = ?3
RETURNING id, run_id, item_index, status, attempts, last_error, input_fb, output_fb, input_hash, output_hash, created_at, updated_at, completed_at
`
type MarkMapItemSucceededParams struct {
OutputFb []byte `db:"output_fb" json:"output_fb"`
OutputHash *string `db:"output_hash" json:"output_hash"`
ID ids.UUID `db:"id" json:"id"`
}
// MarkMapItemSucceeded
//
// UPDATE map_items
// SET status = 'succeeded',
// output_fb = ?1,
// output_hash = ?2,
// last_error = NULL,
// updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now'),
// completed_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now')
// WHERE id = ?3
// RETURNING id, run_id, item_index, status, attempts, last_error, input_fb, output_fb, input_hash, output_hash, created_at, updated_at, completed_at
func (q *Queries) MarkMapItemSucceeded(ctx context.Context, arg MarkMapItemSucceededParams) (MapItem, error) {
row := q.db.QueryRowContext(ctx, MarkMapItemSucceeded, arg.OutputFb, arg.OutputHash, arg.ID)
var i MapItem
err := row.Scan(
&i.ID,
&i.RunID,
&i.ItemIndex,
&i.Status,
&i.Attempts,
&i.LastError,
&i.InputFb,
&i.OutputFb,
&i.InputHash,
&i.OutputHash,
&i.CreatedAt,
&i.UpdatedAt,
&i.CompletedAt,
)
return i, err
}
const UpdateMapRunProgress = `-- name: UpdateMapRunProgress :one
UPDATE map_runs
SET status = ?1,
queued_items = ?2,
running_items = ?3,
succeeded_items = ?4,
failed_items = ?5,
last_error = ?6,
updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now'),
completed_at = ?7
WHERE id = ?8
RETURNING id, agent_id, session_key, operator_kind, idempotency_key, status, total_items, queued_items, running_items, succeeded_items, failed_items, spec_fb, last_error, created_at, updated_at, completed_at
`
type UpdateMapRunProgressParams struct {
Status string `db:"status" json:"status"`
QueuedItems int64 `db:"queued_items" json:"queued_items"`
RunningItems int64 `db:"running_items" json:"running_items"`
SucceededItems int64 `db:"succeeded_items" json:"succeeded_items"`
FailedItems int64 `db:"failed_items" json:"failed_items"`
LastError *string `db:"last_error" json:"last_error"`
CompletedAt *time.Time `db:"completed_at" json:"completed_at"`
ID ids.UUID `db:"id" json:"id"`
}
// UpdateMapRunProgress
//
// UPDATE map_runs
// SET status = ?1,
// queued_items = ?2,
// running_items = ?3,
// succeeded_items = ?4,
// failed_items = ?5,
// last_error = ?6,
// updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now'),
// completed_at = ?7
// WHERE id = ?8
// RETURNING id, agent_id, session_key, operator_kind, idempotency_key, status, total_items, queued_items, running_items, succeeded_items, failed_items, spec_fb, last_error, created_at, updated_at, completed_at
func (q *Queries) UpdateMapRunProgress(ctx context.Context, arg UpdateMapRunProgressParams) (MapRun, error) {
row := q.db.QueryRowContext(ctx, UpdateMapRunProgress,
arg.Status,
arg.QueuedItems,
arg.RunningItems,
arg.SucceededItems,
arg.FailedItems,
arg.LastError,
arg.CompletedAt,
arg.ID,
)
var i MapRun
err := row.Scan(
&i.ID,
&i.AgentID,
&i.SessionKey,
&i.OperatorKind,
&i.IdempotencyKey,
&i.Status,
&i.TotalItems,
&i.QueuedItems,
&i.RunningItems,
&i.SucceededItems,
&i.FailedItems,
&i.SpecFb,
&i.LastError,
&i.CreatedAt,
&i.UpdatedAt,
&i.CompletedAt,
)
return i, err
}

172
pkg/memory/sqlc/queries/dag.sql Normal file
View file

@ -0,0 +1,172 @@
-- DAG persistence queries
-- name: InsertDAGSnapshot :one
INSERT INTO dag_snapshots (
id,
agent_id,
session_key,
from_msg_idx,
to_msg_idx,
msg_count,
roots_json,
content_hash
)
VALUES (
sqlc.arg(id),
sqlc.arg(agent_id),
sqlc.arg(session_key),
sqlc.arg(from_msg_idx),
sqlc.arg(to_msg_idx),
sqlc.arg(msg_count),
sqlc.arg(roots_json),
sqlc.arg(content_hash)
)
RETURNING id,
agent_id,
session_key,
from_msg_idx,
to_msg_idx,
msg_count,
roots_json,
content_hash,
created_at,
updated_at;
-- name: InsertDAGNode :one
INSERT INTO dag_nodes (
id,
snapshot_id,
node_id,
level,
summary,
tokens,
start_idx,
end_idx,
span,
content_hash,
metrics_json,
metadata_json
)
VALUES (
sqlc.arg(id),
sqlc.arg(snapshot_id),
sqlc.arg(node_id),
sqlc.arg(level),
sqlc.arg(summary),
sqlc.arg(tokens),
sqlc.arg(start_idx),
sqlc.arg(end_idx),
sqlc.arg(span),
sqlc.arg(content_hash),
sqlc.arg(metrics_json),
sqlc.arg(metadata_json)
)
RETURNING id,
snapshot_id,
node_id,
level,
summary,
tokens,
start_idx,
end_idx,
span,
content_hash,
metrics_json,
metadata_json,
created_at,
updated_at;
-- name: InsertDAGEdge :exec
INSERT INTO dag_edges (
id,
snapshot_id,
parent_node_id,
child_node_id,
edge_index,
metadata_json
)
VALUES (
sqlc.arg(id),
sqlc.arg(snapshot_id),
sqlc.arg(parent_node_id),
sqlc.arg(child_node_id),
sqlc.arg(edge_index),
sqlc.arg(metadata_json)
);
-- name: GetLatestDAGSnapshotBySession :one
SELECT id,
agent_id,
session_key,
from_msg_idx,
to_msg_idx,
msg_count,
roots_json,
content_hash,
created_at,
updated_at
FROM dag_snapshots
WHERE agent_id = sqlc.arg(agent_id)
AND session_key = sqlc.arg(session_key)
ORDER BY created_at DESC
LIMIT 1;
-- name: GetDAGSnapshotByID :one
SELECT id,
agent_id,
session_key,
from_msg_idx,
to_msg_idx,
msg_count,
roots_json,
content_hash,
created_at,
updated_at
FROM dag_snapshots
WHERE id = sqlc.arg(id)
LIMIT 1;
-- name: ListDAGNodesBySnapshotID :many
SELECT id,
snapshot_id,
node_id,
level,
summary,
tokens,
start_idx,
end_idx,
span,
content_hash,
metrics_json,
metadata_json,
created_at,
updated_at
FROM dag_nodes
WHERE snapshot_id = sqlc.arg(snapshot_id)
ORDER BY start_idx ASC;
-- name: GetDAGNodeBySnapshotAndNodeID :one
SELECT id,
snapshot_id,
node_id,
level,
summary,
tokens,
start_idx,
end_idx,
span,
content_hash,
metrics_json,
metadata_json,
created_at,
updated_at
FROM dag_nodes
WHERE snapshot_id = sqlc.arg(snapshot_id)
AND node_id = sqlc.arg(node_id)
LIMIT 1;
-- name: ListDAGEdgesBySnapshotID :many
SELECT id,
snapshot_id,
parent_node_id,
child_node_id,
edge_index,
metadata_json,
created_at,
updated_at
FROM dag_edges
WHERE snapshot_id = sqlc.arg(snapshot_id)
ORDER BY edge_index ASC,
created_at ASC;

143
pkg/memory/sqlc/queries/map_ops.sql Normal file
View file

@ -0,0 +1,143 @@
-- name: InsertMapRun :one
INSERT INTO map_runs (
id,
agent_id,
session_key,
operator_kind,
idempotency_key,
status,
total_items,
queued_items,
running_items,
succeeded_items,
failed_items,
spec_fb,
last_error
)
VALUES (
sqlc.arg(id),
sqlc.arg(agent_id),
sqlc.arg(session_key),
sqlc.arg(operator_kind),
sqlc.arg(idempotency_key),
sqlc.arg(status),
sqlc.arg(total_items),
sqlc.arg(queued_items),
sqlc.arg(running_items),
sqlc.arg(succeeded_items),
sqlc.arg(failed_items),
sqlc.arg(spec_fb),
sqlc.arg(last_error)
)
RETURNING *;
-- name: GetMapRunByIdempotencyKey :one
SELECT *
FROM map_runs
WHERE agent_id = sqlc.arg(agent_id)
AND session_key = sqlc.arg(session_key)
AND operator_kind = sqlc.arg(operator_kind)
AND idempotency_key = sqlc.arg(idempotency_key)
LIMIT 1;
-- name: GetMapRunByID :one
SELECT *
FROM map_runs
WHERE id = sqlc.arg(id)
LIMIT 1;
-- name: ListMapRunsBySession :many
SELECT *
FROM map_runs
WHERE agent_id = sqlc.arg(agent_id)
AND session_key = sqlc.arg(session_key)
ORDER BY created_at DESC
LIMIT sqlc.arg(lim) OFFSET sqlc.arg(off);
-- name: UpdateMapRunProgress :one
UPDATE map_runs
SET status = sqlc.arg(status),
queued_items = sqlc.arg(queued_items),
running_items = sqlc.arg(running_items),
succeeded_items = sqlc.arg(succeeded_items),
failed_items = sqlc.arg(failed_items),
last_error = sqlc.arg(last_error),
updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now'),
completed_at = sqlc.arg(completed_at)
WHERE id = sqlc.arg(id)
RETURNING *;
-- name: InsertMapItem :one
INSERT INTO map_items (
id,
run_id,
item_index,
status,
attempts,
last_error,
input_fb,
output_fb,
input_hash,
output_hash
)
VALUES (
sqlc.arg(id),
sqlc.arg(run_id),
sqlc.arg(item_index),
sqlc.arg(status),
sqlc.arg(attempts),
sqlc.arg(last_error),
sqlc.arg(input_fb),
sqlc.arg(output_fb),
sqlc.arg(input_hash),
sqlc.arg(output_hash)
)
RETURNING *;
-- name: GetMapItemByID :one
SELECT *
FROM map_items
WHERE id = sqlc.arg(id)
LIMIT 1;
-- name: GetMapItemByRunAndIndex :one
SELECT *
FROM map_items
WHERE run_id = sqlc.arg(run_id)
AND item_index = sqlc.arg(item_index)
LIMIT 1;
-- name: ListMapItemsByRunPaged :many
SELECT *
FROM map_items
WHERE run_id = sqlc.arg(run_id)
ORDER BY item_index ASC
LIMIT sqlc.arg(lim) OFFSET sqlc.arg(off);
-- name: MarkMapItemRunning :one
UPDATE map_items
SET status = 'running',
attempts = attempts + 1,
last_error = NULL,
updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now')
WHERE id = sqlc.arg(id)
RETURNING *;
-- name: MarkMapItemSucceeded :one
UPDATE map_items
SET status = 'succeeded',
output_fb = sqlc.arg(output_fb),
output_hash = sqlc.arg(output_hash),
last_error = NULL,
updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now'),
completed_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now')
WHERE id = sqlc.arg(id)
RETURNING *;
-- name: MarkMapItemFailed :one
UPDATE map_items
SET status = 'failed',
last_error = sqlc.arg(last_error),
updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now'),
completed_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now')
WHERE id = sqlc.arg(id)
RETURNING *;
-- name: CountMapItemsByRun :one
SELECT count(*) AS count
FROM map_items
WHERE run_id = sqlc.arg(run_id);
-- name: CountMapItemsByRunAndStatus :many
SELECT status,
count(*) AS count
FROM map_items
WHERE run_id = sqlc.arg(run_id)
GROUP BY status;

404
pkg/memory/store/retrieval_policy.go Normal file
View file

@ -0,0 +1,404 @@
package store
import (
"context"
"encoding/json"
"fmt"
"time"
"github.com/ZanzyTHEbar/dragonscale/pkg/ids"
"github.com/ZanzyTHEbar/dragonscale/pkg/logger"
"github.com/ZanzyTHEbar/dragonscale/pkg/memory"
)
type retrievalMode string
const (
retrievalModeShadow retrievalMode = "shadow"
retrievalModePromoted retrievalMode = "promoted"
retrievalModeRollback retrievalMode = "rollback"
)
const (
retrievalPolicyStateKey = "memory:retrieval:policy_state"
retrievalPolicyGatesKey = "memory:retrieval:promotion_gates"
retrievalPolicyMetricsKey = "memory:retrieval:shadow_metrics"
)
type retrievalPolicyState struct {
Mode retrievalMode `json:"mode"`
Reason string `json:"reason,omitempty"`
UpdatedAt time.Time `json:"updated_at"`
}
type retrievalPromotionGates struct {
MinSamples int `json:"min_samples"`
MinAugmentedSamples int `json:"min_augmented_samples"`
MinTop1Parity float64 `json:"min_top1_parity"`
MinOverlapAtK float64 `json:"min_overlap_at_k"`
MinPromotedSamples int `json:"min_promoted_samples"`
RollbackTop1Parity float64 `json:"rollback_top1_parity"`
RollbackOverlapAtK float64 `json:"rollback_overlap_at_k"`
MaxNoResultRate float64 `json:"max_no_result_rate"`
PromotedNoResultRate float64 `json:"promoted_no_result_rate"`
}
type retrievalShadowMetrics struct {
TotalQueries int `json:"total_queries"`
NoResultQueries int `json:"no_result_queries"`
AugmentedQueries int `json:"augmented_queries"`
Top1ParityHits int `json:"top1_parity_hits"`
OverlapAtKTotal float64 `json:"overlap_at_k_total"`
PromotedQueries int `json:"promoted_queries"`
PromotedTop1Hits int `json:"promoted_top1_hits"`
PromotedOverlap float64 `json:"promoted_overlap_total"`
PromotedNoResultQueries int `json:"promoted_no_result_queries"`
UpdatedAt time.Time `json:"updated_at"`
}
type retrievalParity struct {
Top1Match bool
OverlapAtK float64
NoResultPair bool
}
func defaultRetrievalPolicyState() retrievalPolicyState {
return retrievalPolicyState{
Mode: retrievalModeShadow,
Reason: "default_shadow_bootstrap",
UpdatedAt: time.Now().UTC(),
}
}
func defaultRetrievalPromotionGates() retrievalPromotionGates {
return retrievalPromotionGates{
MinSamples: 25,
MinAugmentedSamples: 10,
MinTop1Parity: 0.65,
MinOverlapAtK: 0.60,
MinPromotedSamples: 10,
RollbackTop1Parity: 0.45,
RollbackOverlapAtK: 0.35,
MaxNoResultRate: 0.90,
PromotedNoResultRate: 0.95,
}
}
func defaultRetrievalShadowMetrics() retrievalShadowMetrics {
return retrievalShadowMetrics{
UpdatedAt: time.Now().UTC(),
}
}
func (m *MemoryStore) loadRetrievalPolicyState(ctx context.Context) retrievalPolicyState {
raw, err := m.delegate.GetKV(ctx, m.agentID, retrievalPolicyStateKey)
if err != nil || raw == "" {
state := defaultRetrievalPolicyState()
if perr := m.persistRetrievalPolicyState(ctx, state); perr != nil {
logger.WarnCF("memory", "failed to persist default retrieval policy state",
map[string]interface{}{"error": perr.Error()})
}
return state
}
var state retrievalPolicyState
if err := json.Unmarshal([]byte(raw), &state); err != nil {
state = defaultRetrievalPolicyState()
if perr := m.persistRetrievalPolicyState(ctx, state); perr != nil {
logger.WarnCF("memory", "failed to persist repaired retrieval policy state",
map[string]interface{}{"error": perr.Error()})
}
return state
}
if !isValidRetrievalMode(state.Mode) {
state = defaultRetrievalPolicyState()
if perr := m.persistRetrievalPolicyState(ctx, state); perr != nil {
logger.WarnCF("memory", "failed to persist healed retrieval policy state",
map[string]interface{}{"error": perr.Error()})
}
}
return state
}
func (m *MemoryStore) loadRetrievalPromotionGates(ctx context.Context) retrievalPromotionGates {
raw, err := m.delegate.GetKV(ctx, m.agentID, retrievalPolicyGatesKey)
if err != nil || raw == "" {
gates := defaultRetrievalPromotionGates()
if perr := m.persistRetrievalPromotionGates(ctx, gates); perr != nil {
logger.WarnCF("memory", "failed to persist default retrieval promotion gates",
map[string]interface{}{"error": perr.Error()})
}
return gates
}
var gates retrievalPromotionGates
if err := json.Unmarshal([]byte(raw), &gates); err != nil {
gates = defaultRetrievalPromotionGates()
if perr := m.persistRetrievalPromotionGates(ctx, gates); perr != nil {
logger.WarnCF("memory", "failed to persist repaired retrieval promotion gates",
map[string]interface{}{"error": perr.Error()})
}
return gates
}
// Guard-clauses for malformed/partial gate records.
if gates.MinSamples <= 0 {
gates.MinSamples = defaultRetrievalPromotionGates().MinSamples
}
if gates.MinPromotedSamples <= 0 {
gates.MinPromotedSamples = defaultRetrievalPromotionGates().MinPromotedSamples
}
if gates.MinAugmentedSamples <= 0 {
gates.MinAugmentedSamples = defaultRetrievalPromotionGates().MinAugmentedSamples
}
if gates.MinTop1Parity <= 0 {
gates.MinTop1Parity = defaultRetrievalPromotionGates().MinTop1Parity
}
if gates.MinOverlapAtK <= 0 {
gates.MinOverlapAtK = defaultRetrievalPromotionGates().MinOverlapAtK
}
if gates.RollbackTop1Parity <= 0 {
gates.RollbackTop1Parity = defaultRetrievalPromotionGates().RollbackTop1Parity
}
if gates.RollbackOverlapAtK <= 0 {
gates.RollbackOverlapAtK = defaultRetrievalPromotionGates().RollbackOverlapAtK
}
if gates.MaxNoResultRate <= 0 {
gates.MaxNoResultRate = defaultRetrievalPromotionGates().MaxNoResultRate
}
if gates.PromotedNoResultRate <= 0 {
gates.PromotedNoResultRate = defaultRetrievalPromotionGates().PromotedNoResultRate
}
return gates
}
func (m *MemoryStore) loadRetrievalShadowMetrics(ctx context.Context) retrievalShadowMetrics {
raw, err := m.delegate.GetKV(ctx, m.agentID, retrievalPolicyMetricsKey)
if err != nil || raw == "" {
metrics := defaultRetrievalShadowMetrics()
if perr := m.persistRetrievalShadowMetrics(ctx, metrics); perr != nil {
logger.WarnCF("memory", "failed to persist default retrieval shadow metrics",
map[string]interface{}{"error": perr.Error()})
}
return metrics
}
var metrics retrievalShadowMetrics
if err := json.Unmarshal([]byte(raw), &metrics); err != nil {
metrics = defaultRetrievalShadowMetrics()
if perr := m.persistRetrievalShadowMetrics(ctx, metrics); perr != nil {
logger.WarnCF("memory", "failed to persist repaired retrieval shadow metrics",
map[string]interface{}{"error": perr.Error()})
}
return metrics
}
return metrics
}
func (m *MemoryStore) persistRetrievalPolicyState(ctx context.Context, state retrievalPolicyState) error {
payload, err := json.Marshal(state)
if err != nil {
return fmt.Errorf("marshal retrieval policy state: %w", err)
}
if err := m.delegate.UpsertKV(ctx, m.agentID, retrievalPolicyStateKey, string(payload)); err != nil {
return fmt.Errorf("upsert retrieval policy state: %w", err)
}
return nil
}
func (m *MemoryStore) persistRetrievalPromotionGates(ctx context.Context, gates retrievalPromotionGates) error {
payload, err := json.Marshal(gates)
if err != nil {
return fmt.Errorf("marshal retrieval promotion gates: %w", err)
}
if err := m.delegate.UpsertKV(ctx, m.agentID, retrievalPolicyGatesKey, string(payload)); err != nil {
return fmt.Errorf("upsert retrieval promotion gates: %w", err)
}
return nil
}
func (m *MemoryStore) persistRetrievalShadowMetrics(ctx context.Context, metrics retrievalShadowMetrics) error {
payload, err := json.Marshal(metrics)
if err != nil {
return fmt.Errorf("marshal retrieval shadow metrics: %w", err)
}
if err := m.delegate.UpsertKV(ctx, m.agentID, retrievalPolicyMetricsKey, string(payload)); err != nil {
return fmt.Errorf("upsert retrieval shadow metrics: %w", err)
}
return nil
}
func computeRetrievalParity(baseline, augmented []memory.SearchResult, k int) retrievalParity {
if k <= 0 {
k = 5
}
parity := retrievalParity{
NoResultPair: len(baseline) == 0 && len(augmented) == 0,
}
if len(baseline) > 0 && len(augmented) > 0 {
parity.Top1Match = baseline[0].ID == augmented[0].ID
}
baseTop := topKResultIDs(baseline, k)
augTop := topKResultIDs(augmented, k)
if len(baseTop) == 0 && len(augTop) == 0 {
parity.OverlapAtK = 1.0
return parity
}
union := make(map[ids.UUID]bool, len(baseTop)+len(augTop))
intersection := 0
for id := range baseTop {
union[id] = true
if augTop[id] {
intersection++
}
}
for id := range augTop {
union[id] = true
}
if len(union) > 0 {
parity.OverlapAtK = float64(intersection) / float64(len(union))
}
return parity
}
func topKResultIDs(results []memory.SearchResult, k int) map[ids.UUID]bool {
if k <= 0 {
return nil
}
if len(results) < k {
k = len(results)
}
out := make(map[ids.UUID]bool, k)
for i := 0; i < k; i++ {
out[results[i].ID] = true
}
return out
}
func (m *MemoryStore) updateRetrievalPolicy(
ctx context.Context,
state retrievalPolicyState,
gates retrievalPromotionGates,
metrics retrievalShadowMetrics,
parity retrievalParity,
augmentedUsed bool,
baseline []memory.SearchResult,
) retrievalPolicyState {
prevMode := state.Mode
metrics.TotalQueries++
if len(baseline) == 0 {
metrics.NoResultQueries++
}
if augmentedUsed {
metrics.AugmentedQueries++
}
if parity.Top1Match {
metrics.Top1ParityHits++
}
metrics.OverlapAtKTotal += parity.OverlapAtK
metrics.UpdatedAt = time.Now().UTC()
switch state.Mode {
case retrievalModeShadow:
top1Rate := safeRate(metrics.Top1ParityHits, metrics.TotalQueries)
overlapAvg := safeAvg(metrics.OverlapAtKTotal, metrics.TotalQueries)
noResultRate := safeRate(metrics.NoResultQueries, metrics.TotalQueries)
if metrics.TotalQueries >= gates.MinSamples &&
metrics.AugmentedQueries >= gates.MinAugmentedSamples &&
top1Rate >= gates.MinTop1Parity &&
overlapAvg >= gates.MinOverlapAtK &&
noResultRate <= gates.MaxNoResultRate {
state.Mode = retrievalModePromoted
state.Reason = fmt.Sprintf(
"promotion_gates_passed(top1=%.3f overlap=%.3f samples=%d)",
top1Rate,
overlapAvg,
metrics.TotalQueries,
)
state.UpdatedAt = time.Now().UTC()
metrics.PromotedQueries = 0
metrics.PromotedTop1Hits = 0
metrics.PromotedOverlap = 0
metrics.PromotedNoResultQueries = 0
}
case retrievalModePromoted:
metrics.PromotedQueries++
if len(baseline) == 0 {
metrics.PromotedNoResultQueries++
}
if parity.Top1Match {
metrics.PromotedTop1Hits++
}
metrics.PromotedOverlap += parity.OverlapAtK
if metrics.PromotedQueries >= gates.MinPromotedSamples {
promotedTop1 := safeRate(metrics.PromotedTop1Hits, metrics.PromotedQueries)
promotedOverlap := safeAvg(metrics.PromotedOverlap, metrics.PromotedQueries)
promotedNoResultRate := safeRate(metrics.PromotedNoResultQueries, metrics.PromotedQueries)
if promotedTop1 < gates.RollbackTop1Parity ||
promotedOverlap < gates.RollbackOverlapAtK ||
promotedNoResultRate > gates.PromotedNoResultRate {
state.Mode = retrievalModeRollback
state.Reason = fmt.Sprintf(
"rollback_triggered(top1=%.3f overlap=%.3f no_result_rate=%.3f promoted_samples=%d)",
promotedTop1,
promotedOverlap,
promotedNoResultRate,
metrics.PromotedQueries,
)
state.UpdatedAt = time.Now().UTC()
}
}
}
if prevMode != state.Mode {
logger.InfoCF("memory", "retrieval policy mode transition",
map[string]interface{}{
"from": prevMode,
"to": state.Mode,
"reason": state.Reason,
"total_queries": metrics.TotalQueries,
"augmented_queries": metrics.AugmentedQueries,
"promoted_queries": metrics.PromotedQueries,
})
}
if err := m.persistRetrievalShadowMetrics(ctx, metrics); err != nil {
logger.WarnCF("memory", "failed to persist retrieval shadow metrics",
map[string]interface{}{"mode": state.Mode, "error": err.Error()})
}
if err := m.persistRetrievalPolicyState(ctx, state); err != nil {
logger.WarnCF("memory", "failed to persist retrieval policy state",
map[string]interface{}{"mode": state.Mode, "error": err.Error()})
}
return state
}
func safeRate(numerator, denominator int) float64 {
if denominator <= 0 {
return 0
}
return float64(numerator) / float64(denominator)
}
func safeAvg(total float64, count int) float64 {
if count <= 0 {
return 0
}
return total / float64(count)
}
func isValidRetrievalMode(mode retrievalMode) bool {
switch mode {
case retrievalModeShadow, retrievalModePromoted, retrievalModeRollback:
return true
default:
return false
}
}

18
pkg/memory/store/token_test_helpers_test.go Normal file
View file

@ -0,0 +1,18 @@
package store
import (
"fmt"
"strings"
)
func contentAtLeastTokens(target int) string {
if target <= 0 {
return ""
}
var b strings.Builder
for i := 0; estimateTokens(b.String()) < target; i++ {
b.WriteString(fmt.Sprintf("token-%d ", i))
}
return b.String()
}