Memory v2 fixup I1+I4: expires_at validation + audit JSON marshal

Two small Important findings from self-review, bundled because both
are <20 line changes touching the same file.

I1: expires_at silent drop
  - mcp_tools_memory_v2.go:130 had `if t, err := ...; err == nil { ... }`
    which dropped malformed timestamps without telling the agent.
    Agent passes `expires_at: "tomorrow"`, gets a 200, and the memory
    has no TTL.
  - Now returns a clear error: "invalid expires_at: must be RFC3339"
  - Test renamed: TestCommitMemoryV2_BadExpiresIsIgnored (which
    codified the bug) → TestCommitMemoryV2_BadExpiresReturnsError
    (which pins the fix).

I4: audit log JSON via Sprintf-%q
  - auditOrgWrite was building activity_logs.metadata via fmt.Sprintf
    with %q. Go-quoted strings happen to coincide with JSON-quoted
    for ASCII (and today's values are pure ASCII: UUID + hex digest)
    so the bug was latent.
  - Replaced with json.Marshal of map[string]string. Same wire shape
    today, but won't silently produce invalid JSON if metadata grows
    to include arbitrary content snippets.
  - New test TestAuditOrgWrite_MetadataIsValidJSON uses a custom
    sqlmock.Argument matcher (jsonValidMatcher) that fails the test
    if the metadata column isn't parseable JSON. The test runs
    auditOrgWrite with a content string containing quotes,
    backslashes, and a control byte — values where %q would diverge
    from JSON-quote.

Both pre-existing tests (TestCommitMemoryV2_AuditsOrgWrites etc.)
remain green.
This commit is contained in:
Hongming Wang 2026-05-04 08:57:58 -07:00
parent 7cffff844b
commit d48693144b
2 changed files with 79 additions and 12 deletions

View File

@ -127,9 +127,11 @@ func (h *MCPHandler) toolCommitMemoryV2(ctx context.Context, workspaceID string,
Source: contract.MemorySourceAgent,
}
if exp, ok := args["expires_at"].(string); ok && exp != "" {
if t, err := time.Parse(time.RFC3339, exp); err == nil {
body.ExpiresAt = &t
t, err := time.Parse(time.RFC3339, exp)
if err != nil {
return "", fmt.Errorf("invalid expires_at: must be RFC3339 (got %q): %w", exp, err)
}
body.ExpiresAt = &t
}
if pin, ok := args["pin"].(bool); ok {
body.Pin = pin
@ -331,10 +333,23 @@ func (h *MCPHandler) toolForgetMemory(ctx context.Context, workspaceID string, a
func (h *MCPHandler) auditOrgWrite(ctx context.Context, workspaceID, ns, content, memID string) error {
hash := sha256.Sum256([]byte(content))
hashHex := hex.EncodeToString(hash[:])
_, err := h.database.ExecContext(ctx, `
// json.Marshal, not Sprintf-%q. %q produces Go-quoted strings,
// which are NOT valid JSON for non-ASCII inputs (Go's escapes
// like \xNN aren't part of the JSON spec). Today's values are
// pure-ASCII so the bug was latent; if metadata grows to include
// arbitrary content snippets it would silently produce invalid
// JSON in activity_logs.
metadata, err := json.Marshal(map[string]string{
"memory_id": memID,
"sha256": hashHex,
})
if err != nil {
return fmt.Errorf("audit metadata marshal: %w", err)
}
_, err = h.database.ExecContext(ctx, `
INSERT INTO activity_logs (workspace_id, action, target, metadata, created_at)
VALUES ($1, 'memory.org_write', $2, $3, now())
`, workspaceID, ns, fmt.Sprintf(`{"memory_id":%q,"sha256":%q}`, memID, hashHex))
`, workspaceID, ns, string(metadata))
if err != nil && err != sql.ErrNoRows {
return err
}

View File

@ -3,6 +3,7 @@ package handlers
import (
"context"
"database/sql"
"database/sql/driver"
"encoding/json"
"errors"
"strings"
@ -342,13 +343,19 @@ func TestCommitMemoryV2_AcceptsExpiresAndPin(t *testing.T) {
}
}
func TestCommitMemoryV2_BadExpiresIsIgnored(t *testing.T) {
// TestCommitMemoryV2_BadExpiresReturnsError pins the I1 fix: malformed
// expires_at must surface as an error, not silently drop (which would
// leave the agent thinking it set a TTL when it didn't).
//
// Replaces TestCommitMemoryV2_BadExpiresIsIgnored which incorrectly
// codified silent-drop as a feature.
func TestCommitMemoryV2_BadExpiresReturnsError(t *testing.T) {
db, _, _ := sqlmock.New()
defer db.Close()
gotExp := (*time.Time)(nil)
pluginCalled := false
h := newV2Handler(t, db, &stubMemoryPlugin{
commitFn: func(_ context.Context, _ string, body contract.MemoryWrite) (*contract.MemoryWriteResponse, error) {
gotExp = body.ExpiresAt
commitFn: func(_ context.Context, _ string, _ contract.MemoryWrite) (*contract.MemoryWriteResponse, error) {
pluginCalled = true
return &contract.MemoryWriteResponse{ID: "mem-1", Namespace: "workspace:root-1"}, nil
},
}, rootNamespaceResolver())
@ -356,12 +363,57 @@ func TestCommitMemoryV2_BadExpiresIsIgnored(t *testing.T) {
"content": "x",
"expires_at": "tomorrow at noon",
})
if err != nil {
t.Fatalf("err: %v", err)
if err == nil {
t.Fatalf("expected error for malformed expires_at, got nil")
}
if gotExp != nil {
t.Errorf("malformed expires must be ignored, got %v", gotExp)
if !strings.Contains(err.Error(), "invalid expires_at") {
t.Errorf("err = %v, want substring 'invalid expires_at'", err)
}
if pluginCalled {
t.Errorf("plugin must NOT be called when expires_at fails to parse")
}
}
// TestAuditOrgWrite_MetadataIsValidJSON pins the I4 fix: audit metadata
// is built via json.Marshal, not Sprintf-%q. This test exercises
// auditOrgWrite directly with a content string containing characters
// where Go-quote would diverge from JSON-quote, and asserts the
// metadata column receives valid JSON.
func TestAuditOrgWrite_MetadataIsValidJSON(t *testing.T) {
db, mock, _ := sqlmock.New()
defer db.Close()
// jsonValidArg is a sqlmock.Argument that asserts its input
// parses as JSON. Used as the metadata-arg matcher so the test
// fails loudly if a future refactor regresses to Sprintf-%q.
matcher := jsonValidMatcher{}
mock.ExpectExec("INSERT INTO activity_logs").
WithArgs("ws-1", "org:abc", matcher).
WillReturnResult(sqlmock.NewResult(0, 1))
h := &MCPHandler{database: db}
if err := h.auditOrgWrite(context.Background(),
"ws-1", "org:abc",
"content with \"quotes\" \\backslash and \x01 control",
"mem-uuid-1"); err != nil {
t.Fatalf("auditOrgWrite: %v", err)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Errorf("expectations: %v", err)
}
}
// jsonValidMatcher is a sqlmock.Argument that passes only when the
// driver-encoded value parses as JSON. Lets the I4 test fail loudly
// if metadata regresses to non-JSON output.
type jsonValidMatcher struct{}
func (jsonValidMatcher) Match(v driver.Value) bool {
s, ok := v.(string)
if !ok {
return false
}
var out map[string]interface{}
return json.Unmarshal([]byte(s), &out) == nil
}
// --- search_memory ---