| 1 | package session |
| 2 | |
| 3 | import ( |
| 4 | "bytes" |
| 5 | "context" |
| 6 | "crypto/sha256" |
| 7 | "encoding/base64" |
| 8 | "encoding/hex" |
| 9 | "encoding/json" |
| 10 | "os" |
| 11 | "path/filepath" |
| 12 | "strings" |
| 13 | "testing" |
| 14 | |
| 15 | "reasonix/internal/attachment" |
| 16 | "reasonix/internal/provider" |
| 17 | "reasonix/internal/sessioncontent" |
| 18 | "reasonix/internal/sessioninbox" |
| 19 | "reasonix/internal/store" |
| 20 | ) |
| 21 | |
| 22 | const attachmentReadPNG = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" |
| 23 | |
| 24 | func TestReadSessionAttachmentAuthorizesDigestOnly(t *testing.T) { |
| 25 | root := filepath.Join(t.TempDir(), "sessions-v4") |
| 26 | service, err := NewService("local", NewFilesystemPersistence(root)) |
| 27 | if err != nil { |
| 28 | t.Fatal(err) |
| 29 | } |
| 30 | t.Cleanup(func() { _ = service.CloseAll(context.Background()) }) |
| 31 | runtime, err := service.Create(t.Context(), CreateOptions{SessionID: "attach-read"}) |
| 32 | if err != nil { |
| 33 | t.Fatal(err) |
| 34 | } |
| 35 | raw, err := base64.StdEncoding.DecodeString(attachmentReadPNG) |
| 36 | if err != nil { |
| 37 | t.Fatal(err) |
| 38 | } |
| 39 | content := contentStoreForSessionDir(filepath.Join(root, runtime.Ref().SessionID)) |
| 40 | ref, err := content.Put(t.Context(), bytes.NewReader(raw), sessioncontent.Metadata{MediaType: "image/png", Name: "shot.png"}) |
| 41 | if err != nil { |
| 42 | t.Fatal(err) |
| 43 | } |
| 44 | msg := provider.Message{ |
| 45 | ID: "user-1", Role: provider.RoleUser, Content: "see this", |
| 46 | ImageInputs: []attachment.ImageInput{{ |
| 47 | Kind: attachment.KindAttachment, |
| 48 | Attachment: &attachment.AttachmentRef{ |
| 49 | Version: attachment.RefVersion, Content: ref, Width: 1, Height: 1, DisplayName: "shot.png", |
| 50 | }, |
| 51 | }}, |
| 52 | } |
| 53 | payload, err := json.Marshal(map[string]any{"message": msg}) |
| 54 | if err != nil { |
| 55 | t.Fatal(err) |
| 56 | } |
| 57 | if _, err := runtime.Session().AppendBatch(t.Context(), "user-1", []Event{{Kind: "message/complete", Payload: payload}}); err != nil { |
| 58 | t.Fatal(err) |
| 59 | } |
| 60 | if _, err := runtime.Session().Flush(t.Context()); err != nil { |
| 61 | t.Fatal(err) |
| 62 | } |
| 63 | _ = historyPageReady(t, service.Query(), runtime.Ref(), "", 8) |
| 64 | |
| 65 | got, total, err := service.Query().ReadSessionAttachment(t.Context(), runtime.Ref(), ref.Digest, 0, 1<<20) |
| 66 | if err != nil { |
| 67 | t.Fatal(err) |
| 68 | } |
| 69 | if total != int64(len(raw)) || !bytes.Equal(got, raw) { |
| 70 | t.Fatalf("read digest-only attachment total=%d len=%d", total, len(got)) |
| 71 | } |
| 72 | sum := sha256.Sum256(raw) |
| 73 | if ref.Digest != hex.EncodeToString(sum[:]) { |
| 74 | t.Fatal("fixture digest is not the original SHA-256") |
| 75 | } |
| 76 | foreign := strings.Repeat("0", 64) |
| 77 | if _, _, err := service.Query().ReadSessionAttachment(t.Context(), runtime.Ref(), foreign, 0, 1); err == nil { |
| 78 | t.Fatal("foreign digest was authorized") |
| 79 | } |
| 80 | if _, _, err := service.Query().ReadSessionAttachment(t.Context(), runtime.Ref(), ref.Digest, 0, (1<<20)+1); err == nil { |
| 81 | t.Fatal("oversized attachment range was accepted") |
| 82 | } |
| 83 | } |
| 84 | |
| 85 | func TestReadSessionAttachmentAuthorizesInboxRefs(t *testing.T) { |
| 86 | root := filepath.Join(t.TempDir(), "sessions-v4") |
| 87 | service, err := NewService("local", NewFilesystemPersistence(root)) |
| 88 | if err != nil { |
| 89 | t.Fatal(err) |
| 90 | } |
| 91 | t.Cleanup(func() { _ = service.CloseAll(context.Background()) }) |
| 92 | runtime, err := service.Create(t.Context(), CreateOptions{SessionID: "attach-inbox"}) |
| 93 | if err != nil { |
| 94 | t.Fatal(err) |
| 95 | } |
| 96 | raw, err := base64.StdEncoding.DecodeString(attachmentReadPNG) |
| 97 | if err != nil { |
| 98 | t.Fatal(err) |
| 99 | } |
| 100 | sessionDir := filepath.Join(root, runtime.Ref().SessionID) |
| 101 | content := contentStoreForSessionDir(sessionDir) |
| 102 | ref, err := content.Put(t.Context(), bytes.NewReader(raw), sessioncontent.Metadata{MediaType: "image/png", Name: "queued.png"}) |
| 103 | if err != nil { |
| 104 | t.Fatal(err) |
| 105 | } |
| 106 | inbox, err := sessioninbox.Open(sessionDir, sessioninbox.Limits{}) |
| 107 | if err != nil { |
| 108 | t.Fatal(err) |
| 109 | } |
| 110 | t.Cleanup(func() { inbox.Close() }) |
| 111 | if _, err := inbox.Enqueue(sessioninbox.EnqueueRequest{Envelope: sessioninbox.PromptEnvelope{ |
| 112 | SubmitText: "queued image", |
| 113 | ImageInputs: []attachment.ImageInput{{ |
| 114 | Kind: attachment.KindAttachment, |
| 115 | Attachment: &attachment.AttachmentRef{ |
| 116 | Version: attachment.RefVersion, Content: ref, Width: 1, Height: 1, DisplayName: "queued.png", |
| 117 | }, |
| 118 | }}, |
| 119 | }}); err != nil { |
| 120 | t.Fatal(err) |
| 121 | } |
| 122 | got, total, err := service.Query().ReadSessionAttachment(t.Context(), runtime.Ref(), ref.Digest, 0, 1<<20) |
| 123 | if err != nil { |
| 124 | t.Fatal(err) |
| 125 | } |
| 126 | if total != int64(len(raw)) || !bytes.Equal(got, raw) { |
| 127 | t.Fatalf("inbox attachment total=%d len=%d", total, len(got)) |
| 128 | } |
| 129 | if _, err := os.Stat(store.SessionInboxDir(sessionDir)); err != nil { |
| 130 | t.Fatal(err) |
| 131 | } |
| 132 | } |
| 133 | |
| 134 | func TestExportIncludesMessageAttachmentClosure(t *testing.T) { |
| 135 | root := filepath.Join(t.TempDir(), "sessions-v4") |
| 136 | service, err := NewService("local", NewFilesystemPersistence(root)) |
| 137 | if err != nil { |
| 138 | t.Fatal(err) |
| 139 | } |
| 140 | t.Cleanup(func() { _ = service.CloseAll(context.Background()) }) |
| 141 | runtime, err := service.Create(t.Context(), CreateOptions{SessionID: "export-images"}) |
| 142 | if err != nil { |
| 143 | t.Fatal(err) |
| 144 | } |
| 145 | raw, err := base64.StdEncoding.DecodeString(attachmentReadPNG) |
| 146 | if err != nil { |
| 147 | t.Fatal(err) |
| 148 | } |
| 149 | sessionDir := filepath.Join(root, runtime.Ref().SessionID) |
| 150 | content := contentStoreForSessionDir(sessionDir) |
| 151 | ref, err := content.Put(t.Context(), bytes.NewReader(raw), sessioncontent.Metadata{MediaType: "image/png", Name: "shot.png"}) |
| 152 | if err != nil { |
| 153 | t.Fatal(err) |
| 154 | } |
| 155 | message := provider.Message{ID: "user-image", Role: provider.RoleUser, Content: "see this", ImageInputs: []attachment.ImageInput{{ |
| 156 | Kind: attachment.KindAttachment, |
| 157 | Attachment: &attachment.AttachmentRef{ |
| 158 | Version: attachment.RefVersion, Content: ref, Width: 1, Height: 1, DisplayName: "shot.png", |
| 159 | }, |
| 160 | }}} |
| 161 | payload, err := json.Marshal(map[string]any{"message": message}) |
| 162 | if err != nil { |
| 163 | t.Fatal(err) |
| 164 | } |
| 165 | if _, err = runtime.Session().AppendBatch(t.Context(), "user-image", []Event{{Kind: "message/complete", Payload: payload}}); err != nil { |
| 166 | t.Fatal(err) |
| 167 | } |
| 168 | |
| 169 | exported := filepath.Join(t.TempDir(), "exported-session") |
| 170 | if err := service.Export(t.Context(), runtime.Ref(), exported); err != nil { |
| 171 | t.Fatal(err) |
| 172 | } |
| 173 | copied, err := sessioncontent.New(filepath.Join(exported, ".content-v1")).ReadRange(t.Context(), ref, 0, ref.Bytes) |
| 174 | if err != nil || !bytes.Equal(copied, raw) { |
| 175 | t.Fatalf("exported message attachment missing: %v", err) |
| 176 | } |
| 177 | } |
| 178 | |
| 179 | func TestExportRefusesMissingAttachmentObject(t *testing.T) { |
| 180 | root := filepath.Join(t.TempDir(), "sessions-v4") |
| 181 | service, err := NewService("local", NewFilesystemPersistence(root)) |
| 182 | if err != nil { |
| 183 | t.Fatal(err) |
| 184 | } |
| 185 | t.Cleanup(func() { _ = service.CloseAll(context.Background()) }) |
| 186 | runtime, err := service.Create(t.Context(), CreateOptions{SessionID: "export-missing"}) |
| 187 | if err != nil { |
| 188 | t.Fatal(err) |
| 189 | } |
| 190 | raw, err := base64.StdEncoding.DecodeString(attachmentReadPNG) |
| 191 | if err != nil { |
| 192 | t.Fatal(err) |
| 193 | } |
| 194 | sessionDir := filepath.Join(root, runtime.Ref().SessionID) |
| 195 | content := contentStoreForSessionDir(sessionDir) |
| 196 | ref, err := content.Put(t.Context(), bytes.NewReader(raw), sessioncontent.Metadata{MediaType: "image/png", Name: "shot.png"}) |
| 197 | if err != nil { |
| 198 | t.Fatal(err) |
| 199 | } |
| 200 | msg := provider.Message{ |
| 201 | ID: "user-1", Role: provider.RoleUser, Content: "see this", |
| 202 | ImageInputs: []attachment.ImageInput{{ |
| 203 | Kind: attachment.KindAttachment, |
| 204 | Attachment: &attachment.AttachmentRef{ |
| 205 | Version: attachment.RefVersion, Content: ref, Width: 1, Height: 1, DisplayName: "shot.png", |
| 206 | }, |
| 207 | }}, |
| 208 | } |
| 209 | payload, err := json.Marshal(map[string]any{"message": msg}) |
| 210 | if err != nil { |
| 211 | t.Fatal(err) |
| 212 | } |
| 213 | if _, err := runtime.Session().AppendBatch(t.Context(), "user-1", []Event{{Kind: "message/complete", Payload: payload}}); err != nil { |
| 214 | t.Fatal(err) |
| 215 | } |
| 216 | if _, err := runtime.Session().Flush(t.Context()); err != nil { |
| 217 | t.Fatal(err) |
| 218 | } |
| 219 | if err := os.RemoveAll(content.Root()); err != nil { |
| 220 | t.Fatal(err) |
| 221 | } |
| 222 | if err := service.Export(t.Context(), runtime.Ref(), filepath.Join(t.TempDir(), "broken-export")); err == nil { |
| 223 | t.Fatal("export accepted a missing attachment object") |
| 224 | } |
| 225 | } |
| 226 |