Skip to content

Commit f8b7566

Browse files
priyanshujainpriyanshujain
authored
Merge pull request #110 from 73ai/refactor-file-formats
Migrate storage formats: SQLite → JSONL/Markdown
2 parents a94bae9 + e1e7073 commit f8b7566

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

48 files changed

+1619
-1736
lines changed

‎agent/audit/audit.go‎

Lines changed: 49 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
11
package audit
22

33
import (
4+
"encoding/json"
45
"log/slog"
56
"os"
67
"path/filepath"
8+
"sync"
79
"time"
8-
9-
"github.com/73ai/openbotkit/store"
1010
)
1111

1212
// Entry represents a single audit log record.
@@ -20,44 +20,51 @@ type Entry struct {
2020
Error string
2121
}
2222

23-
// Logger writes audit entries to a database.
24-
type Logger struct {
25-
db *store.DB
23+
type jsonEntry struct {
24+
Timestamp string `json:"timestamp"`
25+
Context string `json:"context"`
26+
ToolName string `json:"tool_name"`
27+
InputSummary string `json:"input_summary"`
28+
OutputSummary string `json:"output_summary"`
29+
ApprovalStatus string `json:"approval_status"`
30+
Error string `json:"error,omitempty"`
2631
}
2732

28-
// NewLogger creates an audit logger backed by the given database.
29-
func NewLogger(db *store.DB) *Logger {
30-
return &Logger{db: db}
33+
// Logger writes audit entries to a JSONL file.
34+
type Logger struct {
35+
mu sync.Mutex
36+
file *os.File
37+
enc *json.Encoder
3138
}
3239

33-
// OpenDefault opens (or creates) the audit database at dbPath,
34-
// runs migrations, and returns a ready Logger.
40+
// OpenDefault opens (or creates) the audit JSONL file at path,
41+
// creating parent directories as needed.
3542
// Returns nil if any step fails (errors are logged via slog).
36-
func OpenDefault(dbPath string) *Logger {
37-
dir := filepath.Dir(dbPath)
43+
func OpenDefault(path string) *Logger {
44+
dir := filepath.Dir(path)
3845
if err := os.MkdirAll(dir, 0700); err != nil {
3946
slog.Debug("audit: cannot create dir", "error", err)
4047
return nil
4148
}
42-
db, err := store.Open(store.SQLiteConfig(dbPath))
49+
f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0600)
4350
if err != nil {
44-
slog.Debug("audit: open db failed", "error", err)
51+
slog.Debug("audit: open file failed", "error", err)
4552
return nil
4653
}
47-
if err := Migrate(db); err != nil {
48-
db.Close()
49-
slog.Debug("audit: migrate failed", "error", err)
50-
return nil
51-
}
52-
return NewLogger(db)
54+
return &Logger{file: f, enc: json.NewEncoder(f)}
5355
}
5456

55-
// Close closes the underlying database connection.
57+
// Close closes the underlying file.
5658
func (l *Logger) Close() error {
57-
if l == nil || l.db == nil {
59+
if l == nil || l.file == nil {
5860
return nil
5961
}
60-
return l.db.Close()
62+
l.mu.Lock()
63+
defer l.mu.Unlock()
64+
err := l.file.Close()
65+
l.file = nil
66+
l.enc = nil
67+
return err
6168
}
6269

6370
const maxSummaryLen = 200
@@ -69,53 +76,36 @@ func truncate(s string, max int) string {
6976
return s[:max] + "..."
7077
}
7178

72-
// Log writes an audit entry. It never returns an error to the caller;
73-
// failures are logged via slog.
79+
// Log writes an audit entry as a JSON line. It never returns an error
80+
// to the caller; failures are logged via slog.
7481
func (l *Logger) Log(e Entry) {
75-
if l == nil || l.db == nil {
82+
if l == nil || l.file == nil {
7683
return
7784
}
7885
ts := e.Timestamp
7986
if ts.IsZero() {
8087
ts = time.Now().UTC()
8188
}
82-
inputSum := truncate(e.InputSummary, maxSummaryLen)
83-
outputSum := truncate(e.OutputSummary, maxSummaryLen)
8489
if e.ApprovalStatus == "" {
8590
e.ApprovalStatus = "n/a"
8691
}
8792

88-
query := l.db.Rebind(`INSERT INTO audit_log (timestamp, context, tool_name, input_summary, output_summary, approval_status, error)
89-
VALUES (?, ?, ?, ?, ?, ?, ?)`)
90-
_, err := l.db.Exec(query,
91-
ts.Format(time.RFC3339),
92-
e.Context,
93-
e.ToolName,
94-
inputSum,
95-
outputSum,
96-
e.ApprovalStatus,
97-
e.Error,
98-
)
99-
if err != nil {
100-
slog.Error("audit log write failed", "tool", e.ToolName, "error", err)
93+
je := jsonEntry{
94+
Timestamp: ts.Format(time.RFC3339),
95+
Context: e.Context,
96+
ToolName: e.ToolName,
97+
InputSummary: truncate(e.InputSummary, maxSummaryLen),
98+
OutputSummary: truncate(e.OutputSummary, maxSummaryLen),
99+
ApprovalStatus: e.ApprovalStatus,
100+
Error: e.Error,
101101
}
102-
}
103102

104-
// Migrate creates the audit_log table if it doesn't exist.
105-
func Migrate(db *store.DB) error {
106-
_, err := db.Exec(`
107-
CREATE TABLE IF NOT EXISTS audit_log (
108-
id INTEGER PRIMARY KEY AUTOINCREMENT,
109-
timestamp TEXT NOT NULL DEFAULT (datetime('now')),
110-
context TEXT NOT NULL,
111-
tool_name TEXT NOT NULL,
112-
input_summary TEXT,
113-
output_summary TEXT,
114-
approval_status TEXT DEFAULT 'n/a',
115-
error TEXT
116-
);
117-
CREATE INDEX IF NOT EXISTS idx_audit_timestamp ON audit_log(timestamp);
118-
CREATE INDEX IF NOT EXISTS idx_audit_tool ON audit_log(tool_name);
119-
`)
120-
return err
103+
l.mu.Lock()
104+
defer l.mu.Unlock()
105+
if l.enc == nil {
106+
return
107+
}
108+
if err := l.enc.Encode(je); err != nil {
109+
slog.Error("audit log write failed", "tool", e.ToolName, "error", err)
110+
}
121111
}

‎agent/audit/audit_test.go‎

Lines changed: 89 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -1,42 +1,22 @@
11
package audit
22

33
import (
4+
"bufio"
5+
"encoding/json"
6+
"os"
47
"path/filepath"
58
"testing"
69
"time"
7-
8-
"github.com/73ai/openbotkit/store"
910
)
1011

11-
func openTestDB(t *testing.T) *store.DB {
12-
t.Helper()
13-
path := filepath.Join(t.TempDir(), "audit_test.db")
14-
db, err := store.Open(store.SQLiteConfig(path))
15-
if err != nil {
16-
t.Fatalf("open test db: %v", err)
17-
}
18-
t.Cleanup(func() { db.Close() })
19-
return db
20-
}
21-
22-
func TestMigrate(t *testing.T) {
23-
db := openTestDB(t)
24-
if err := Migrate(db); err != nil {
25-
t.Fatalf("Migrate: %v", err)
26-
}
27-
// Idempotent.
28-
if err := Migrate(db); err != nil {
29-
t.Fatalf("Migrate (2nd call): %v", err)
30-
}
31-
}
32-
33-
func TestLogger_Log(t *testing.T) {
34-
db := openTestDB(t)
35-
if err := Migrate(db); err != nil {
36-
t.Fatalf("Migrate: %v", err)
12+
func TestLog(t *testing.T) {
13+
path := filepath.Join(t.TempDir(), "audit.jsonl")
14+
l := OpenDefault(path)
15+
if l == nil {
16+
t.Fatal("OpenDefault returned nil")
3717
}
18+
defer l.Close()
3819

39-
l := NewLogger(db)
4020
l.Log(Entry{
4121
Timestamp: time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC),
4222
Context: "cli",
@@ -46,34 +26,32 @@ func TestLogger_Log(t *testing.T) {
4626
ApprovalStatus: "n/a",
4727
})
4828

49-
var count int
50-
if err := db.QueryRow("SELECT COUNT(*) FROM audit_log").Scan(&count); err != nil {
51-
t.Fatalf("query: %v", err)
29+
entries := readJSONL(t, path)
30+
if len(entries) != 1 {
31+
t.Fatalf("got %d entries, want 1", len(entries))
5232
}
53-
if count != 1 {
54-
t.Errorf("count = %d, want 1", count)
33+
if entries[0].ToolName != "bash" {
34+
t.Errorf("tool_name = %q, want %q", entries[0].ToolName, "bash")
5535
}
56-
57-
var toolName, ctx string
58-
err := db.QueryRow("SELECT tool_name, context FROM audit_log WHERE id=1").Scan(&toolName, &ctx)
59-
if err != nil {
60-
t.Fatalf("query row: %v", err)
36+
if entries[0].Context != "cli" {
37+
t.Errorf("context = %q, want %q", entries[0].Context, "cli")
6138
}
62-
if toolName != "bash" {
63-
t.Errorf("tool_name = %q, want %q", toolName, "bash")
39+
if entries[0].Timestamp != "2026-01-01T00:00:00Z" {
40+
t.Errorf("timestamp = %q, want %q", entries[0].Timestamp, "2026-01-01T00:00:00Z")
6441
}
65-
if ctx != "cli" {
66-
t.Errorf("context = %q, want %q", ctx, "cli")
42+
if entries[0].Error != "" {
43+
t.Errorf("error should be omitted, got %q", entries[0].Error)
6744
}
6845
}
6946

70-
func TestLogger_Truncation(t *testing.T) {
71-
db := openTestDB(t)
72-
if err := Migrate(db); err != nil {
73-
t.Fatalf("Migrate: %v", err)
47+
func TestTruncation(t *testing.T) {
48+
path := filepath.Join(t.TempDir(), "audit.jsonl")
49+
l := OpenDefault(path)
50+
if l == nil {
51+
t.Fatal("OpenDefault returned nil")
7452
}
53+
defer l.Close()
7554

76-
l := NewLogger(db)
7755
longInput := make([]byte, 500)
7856
for i := range longInput {
7957
longInput[i] = 'x'
@@ -84,62 +62,96 @@ func TestLogger_Truncation(t *testing.T) {
8462
InputSummary: string(longInput),
8563
})
8664

87-
var inputSum string
88-
err := db.QueryRow("SELECT input_summary FROM audit_log WHERE id=1").Scan(&inputSum)
89-
if err != nil {
90-
t.Fatalf("query: %v", err)
65+
entries := readJSONL(t, path)
66+
if len(entries) != 1 {
67+
t.Fatalf("got %d entries, want 1", len(entries))
9168
}
92-
if len(inputSum) > maxSummaryLen+10 {
93-
t.Errorf("input_summary len = %d, expected truncated to ~%d", len(inputSum), maxSummaryLen)
69+
if len(entries[0].InputSummary) > maxSummaryLen+10 {
70+
t.Errorf("input_summary len = %d, expected truncated to ~%d", len(entries[0].InputSummary), maxSummaryLen)
9471
}
9572
}
9673

97-
func TestLogger_NilSafe(t *testing.T) {
74+
func TestNilSafe(t *testing.T) {
9875
var l *Logger
99-
// Should not panic.
10076
l.Log(Entry{ToolName: "bash"})
10177
if err := l.Close(); err != nil {
10278
t.Errorf("nil Close: %v", err)
10379
}
10480
}
10581

106-
func TestOpenDefault(t *testing.T) {
107-
dbPath := filepath.Join(t.TempDir(), "audit", "data.db")
108-
l := OpenDefault(dbPath)
82+
func TestClose(t *testing.T) {
83+
path := filepath.Join(t.TempDir(), "audit.jsonl")
84+
l := OpenDefault(path)
10985
if l == nil {
11086
t.Fatal("OpenDefault returned nil")
11187
}
112-
defer l.Close()
113-
l.Log(Entry{Context: "test", ToolName: "bash", InputSummary: "echo hi"})
88+
if err := l.Close(); err != nil {
89+
t.Fatalf("Close: %v", err)
90+
}
91+
// Log after close should not panic.
92+
l.Log(Entry{ToolName: "bash", Context: "test"})
93+
}
11494

115-
var count int
116-
if err := l.db.QueryRow("SELECT COUNT(*) FROM audit_log").Scan(&count); err != nil {
117-
t.Fatalf("query: %v", err)
95+
func TestOpenDefault_CreatesDir(t *testing.T) {
96+
path := filepath.Join(t.TempDir(), "sub", "dir", "audit.jsonl")
97+
l := OpenDefault(path)
98+
if l == nil {
99+
t.Fatal("OpenDefault returned nil")
118100
}
119-
if count != 1 {
120-
t.Errorf("count = %d, want 1", count)
101+
defer l.Close()
102+
103+
l.Log(Entry{Context: "test", ToolName: "bash", InputSummary: "echo hi"})
104+
entries := readJSONL(t, path)
105+
if len(entries) != 1 {
106+
t.Fatalf("got %d entries, want 1", len(entries))
121107
}
122108
}
123109

124110
func TestOpenDefault_BadPath(t *testing.T) {
125-
// A null byte in the path is invalid on all platforms.
126-
l := OpenDefault("/bad\x00path/data.db")
111+
l := OpenDefault("/bad\x00path/audit.jsonl")
127112
if l != nil {
128113
l.Close()
129114
t.Error("expected nil for bad path")
130115
}
131116
}
132117

133-
func TestLogger_Close(t *testing.T) {
134-
path := filepath.Join(t.TempDir(), "close_test.db")
135-
db, err := store.Open(store.SQLiteConfig(path))
118+
func TestMultipleEntries(t *testing.T) {
119+
path := filepath.Join(t.TempDir(), "audit.jsonl")
120+
l := OpenDefault(path)
121+
if l == nil {
122+
t.Fatal("OpenDefault returned nil")
123+
}
124+
defer l.Close()
125+
126+
for i := 0; i < 5; i++ {
127+
l.Log(Entry{Context: "cli", ToolName: "bash"})
128+
}
129+
130+
entries := readJSONL(t, path)
131+
if len(entries) != 5 {
132+
t.Fatalf("got %d entries, want 5", len(entries))
133+
}
134+
}
135+
136+
func readJSONL(t *testing.T, path string) []jsonEntry {
137+
t.Helper()
138+
f, err := os.Open(path)
136139
if err != nil {
137-
t.Fatalf("open db: %v", err)
140+
t.Fatalf("open %s: %v", path, err)
138141
}
139-
l := NewLogger(db)
140-
if err := l.Close(); err != nil {
141-
t.Fatalf("Close: %v", err)
142+
defer f.Close()
143+
144+
var entries []jsonEntry
145+
scanner := bufio.NewScanner(f)
146+
for scanner.Scan() {
147+
var e jsonEntry
148+
if err := json.Unmarshal(scanner.Bytes(), &e); err != nil {
149+
t.Fatalf("parse JSON line: %v", err)
150+
}
151+
entries = append(entries, e)
142152
}
143-
// Log after close should not panic (fire-and-forget).
144-
l.Log(Entry{ToolName: "bash", Context: "test"})
153+
if err := scanner.Err(); err != nil {
154+
t.Fatalf("scan: %v", err)
155+
}
156+
return entries
145157
}

0 commit comments

Comments
 (0)