| 1 | package sessionexport |
| 2 | |
| 3 | import ( |
| 4 | "bytes" |
| 5 | "context" |
| 6 | "encoding/base64" |
| 7 | "encoding/json" |
| 8 | "fmt" |
| 9 | "os" |
| 10 | "path/filepath" |
| 11 | "reasonix/internal/attachment" |
| 12 | "reasonix/internal/provider" |
| 13 | "reasonix/internal/session" |
| 14 | "reasonix/internal/sessioncontent" |
| 15 | "strings" |
| 16 | "testing" |
| 17 | ) |
| 18 | |
| 19 | const exportAttachmentPNG = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" |
| 20 | |
| 21 | func TestBuildStagesAndRendersAuthorizedImageInputs(t *testing.T) { |
| 22 | root := filepath.Join(t.TempDir(), "sessions") |
| 23 | service, err := session.NewService("local", session.NewFilesystemPersistence(root)) |
| 24 | if err != nil { |
| 25 | t.Fatal(err) |
| 26 | } |
| 27 | defer service.CloseAll(context.Background()) |
| 28 | runtime, err := service.Create(t.Context(), session.CreateOptions{SessionID: "images"}) |
| 29 | if err != nil { |
| 30 | t.Fatal(err) |
| 31 | } |
| 32 | raw, err := base64.StdEncoding.DecodeString(exportAttachmentPNG) |
| 33 | if err != nil { |
| 34 | t.Fatal(err) |
| 35 | } |
| 36 | contentRef, err := runtime.Session().ContentStore().Put(t.Context(), bytes.NewReader(raw), sessioncontent.Metadata{MediaType: "image/png", Name: "shot.png"}) |
| 37 | if err != nil { |
| 38 | t.Fatal(err) |
| 39 | } |
| 40 | message := provider.Message{ID: "user-image", Role: provider.RoleUser, Content: "inspect", Images: []string{"https://example.test/legacy.png"}, ImageInputs: []attachment.ImageInput{{ |
| 41 | Kind: attachment.KindAttachment, |
| 42 | Attachment: &attachment.AttachmentRef{Version: attachment.RefVersion, Content: contentRef, Width: 1, Height: 1, DisplayName: "shot.png"}, |
| 43 | }}} |
| 44 | payload, _ := json.Marshal(map[string]any{"message": message}) |
| 45 | if _, err = runtime.Session().AppendBatch(t.Context(), "user-image", []session.Event{{Kind: "message/complete", Payload: payload}}); err != nil { |
| 46 | t.Fatal(err) |
| 47 | } |
| 48 | snapshot, err := service.Query().CaptureExportSnapshot(t.Context(), runtime.Ref()) |
| 49 | if err != nil { |
| 50 | t.Fatal(err) |
| 51 | } |
| 52 | dir := t.TempDir() |
| 53 | if _, err = Build(t.Context(), service.Query(), snapshot, dir, nil); err != nil { |
| 54 | t.Fatal(err) |
| 55 | } |
| 56 | staged, err := os.ReadFile(filepath.Join(dir, "attachments", contentRef.Digest)) |
| 57 | if err != nil || !bytes.Equal(staged, raw) { |
| 58 | t.Fatalf("staged attachment mismatch: %v", err) |
| 59 | } |
| 60 | markdown, err := os.ReadFile(filepath.Join(dir, "markdown")) |
| 61 | if err != nil { |
| 62 | t.Fatal(err) |
| 63 | } |
| 64 | wantDataURL := attachment.DataURL("image/png", raw) |
| 65 | if !strings.Contains(string(markdown), "") || !strings.Contains(string(markdown), "![attachment]("+wantDataURL+")") { |
| 66 | t.Fatalf("markdown did not preserve legacy and durable images: %s", markdown) |
| 67 | } |
| 68 | var document struct { |
| 69 | Items []Item `json:"items"` |
| 70 | } |
| 71 | jsonBody, err := os.ReadFile(filepath.Join(dir, "json")) |
| 72 | if err != nil || json.Unmarshal(jsonBody, &document) != nil { |
| 73 | t.Fatalf("read json export: %v", err) |
| 74 | } |
| 75 | if len(document.Items) != 1 { |
| 76 | t.Fatalf("items = %d", len(document.Items)) |
| 77 | } |
| 78 | images, ok := document.Items[0]["images"].([]any) |
| 79 | if !ok || len(images) != 2 || images[0] != "https://example.test/legacy.png" || images[1] != wantDataURL { |
| 80 | t.Fatalf("export images = %#v", document.Items[0]["images"]) |
| 81 | } |
| 82 | } |
| 83 | |
| 84 | func TestBuildFullSnapshotAcrossPagesAndLargeTools(t *testing.T) { |
| 85 | service, err := session.NewService("local", session.NewFilesystemPersistence(filepath.Join(t.TempDir(), "sessions"))) |
| 86 | if err != nil { |
| 87 | t.Fatal(err) |
| 88 | } |
| 89 | defer service.CloseAll(context.Background()) |
| 90 | runtime, err := service.Create(t.Context(), session.CreateOptions{SessionID: "complete"}) |
| 91 | if err != nil { |
| 92 | t.Fatal(err) |
| 93 | } |
| 94 | turnID := "" |
| 95 | appendMessage := func(m provider.Message) { |
| 96 | t.Helper() |
| 97 | payload, _ := json.Marshal(map[string]any{"message": m}) |
| 98 | if _, err := runtime.Session().Append(t.Context(), session.Batch{OperationID: m.ID, TurnID: turnID, Events: []session.Event{{Kind: "message/complete", Payload: payload}}}); err != nil { |
| 99 | t.Fatal(err) |
| 100 | } |
| 101 | } |
| 102 | for i := range 110 { |
| 103 | appendMessage(provider.Message{ID: fmt.Sprintf("user-%d", i), Role: provider.RoleUser, Content: fmt.Sprintf("QUESTION-%03d", i), Origin: provider.MessageOrigin("user")}) |
| 104 | } |
| 105 | output := " prefix\n\n\n" + strings.Repeat("中文✓ ", 200000) + "\n``````\n suffix \n" |
| 106 | |
| 107 | turnID = "incident" |
| 108 | appendEvents := func(id string, events ...session.Event) { |
| 109 | t.Helper() |
| 110 | if _, err := runtime.Session().Append(t.Context(), session.Batch{OperationID: id, TurnID: turnID, Events: events}); err != nil { |
| 111 | t.Fatal(err) |
| 112 | } |
| 113 | } |
| 114 | appendEvents("turn-start", session.Event{Kind: "turn/start", Payload: json.RawMessage(`{"status":"in_progress"}`)}) |
| 115 | appendEvents("mcp-notice", session.Event{Kind: "diagnostic", Optional: true, Payload: json.RawMessage(`{"type":"display-notice-v1","displayRecord":{"id":"mcp-notice","role":"notice","content":"MCP tools/list","code":"mcp_tools_list","detail":"{\"source\":\"shared_host\",\"network_call\":true}"}}`)}) |
| 116 | callIndex := 0 |
| 117 | for sample := range 15 { |
| 118 | attempt, _ := json.Marshal(map[string]any{"id": fmt.Sprintf("attempt-%d", sample), "action": "begin"}) |
| 119 | appendEvents(fmt.Sprintf("sampling-%d", sample), session.Event{Kind: "assistant/attempt", Payload: attempt}) |
| 120 | count := 1 |
| 121 | if sample < 8 { |
| 122 | count = 2 |
| 123 | } |
| 124 | calls := make([]provider.ToolCall, 0, count) |
| 125 | for range count { |
| 126 | id := fmt.Sprintf("call-%d", callIndex) |
| 127 | callIndex++ |
| 128 | calls = append(calls, provider.ToolCall{ID: id, Name: "bash", Arguments: `{"command":"echo test"}`}) |
| 129 | } |
| 130 | appendMessage(provider.Message{ID: fmt.Sprintf("assistant-%d", sample), Role: provider.RoleAssistant, ReasoningContent: fmt.Sprintf("reasoning-%d", sample), ToolCalls: calls}) |
| 131 | for _, call := range calls { |
| 132 | payload, _ := json.Marshal(map[string]string{"id": call.ID, "name": "bash"}) |
| 133 | appendEvents("dispatch-"+call.ID, session.Event{Kind: "tool/call", Payload: payload}) |
| 134 | appendMessage(provider.Message{ID: "result-" + call.ID, Role: provider.RoleTool, ToolCallID: call.ID, Content: output, ToolRunState: provider.ToolRunCompleted}) |
| 135 | } |
| 136 | } |
| 137 | appendMessage(provider.Message{ID: "final", Role: provider.RoleAssistant, Content: "FINAL-ANSWER"}) |
| 138 | appendEvents("turn-end", session.Event{Kind: "turn/end", Payload: json.RawMessage(`{"status":"completed"}`)}) |
| 139 | |
| 140 | snapshot, err := service.Query().CaptureExportSnapshot(t.Context(), runtime.Ref()) |
| 141 | if err != nil { |
| 142 | t.Fatal(err) |
| 143 | } |
| 144 | appendMessage(provider.Message{ID: "later", Role: provider.RoleUser, Content: "MUST-NOT-APPEAR"}) |
| 145 | dir := t.TempDir() |
| 146 | doc, err := Build(t.Context(), service.Query(), snapshot, dir, nil) |
| 147 | if err != nil { |
| 148 | t.Fatal(err) |
| 149 | } |
| 150 | if doc.Records != 150 { |
| 151 | t.Fatalf("records=%d", doc.Records) |
| 152 | } |
| 153 | file, err := os.Open(filepath.Join(dir, "json")) |
| 154 | if err != nil { |
| 155 | t.Fatal(err) |
| 156 | } |
| 157 | defer file.Close() |
| 158 | var data struct { |
| 159 | MCP map[string]int `json:"mcpList"` |
| 160 | Items []Item `json:"items"` |
| 161 | Metadata map[string]any `json:"exportMetadata"` |
| 162 | } |
| 163 | if err = json.NewDecoder(file).Decode(&data); err != nil { |
| 164 | t.Fatal(err) |
| 165 | } |
| 166 | if data.MCP["sharedHost"] != 1 || data.MCP["networkCalls"] != 1 { |
| 167 | t.Fatalf("lost full-history MCP attribution: %+v", data.MCP) |
| 168 | } |
| 169 | count := 0 |
| 170 | for _, item := range data.Items { |
| 171 | if text(item, "text") == "FINAL-ANSWER" && (item["samplingCount"] != float64(15) || item["toolCount"] != float64(23) || item["turnFinal"] != true) { |
| 172 | t.Fatalf("lost incident telemetry: %+v", item) |
| 173 | } |
| 174 | if text(item, "kind") == "tool" { |
| 175 | count++ |
| 176 | if text(item, "output") != output || text(item, "status") != "done" { |
| 177 | t.Fatal("tool output/state lost") |
| 178 | } |
| 179 | } |
| 180 | } |
| 181 | if count != 23 { |
| 182 | t.Fatalf("tools=%d", count) |
| 183 | } |
| 184 | md, err := os.ReadFile(filepath.Join(dir, "markdown")) |
| 185 | if err != nil { |
| 186 | t.Fatal(err) |
| 187 | } |
| 188 | for _, expected := range []string{"QUESTION-000", "FINAL-ANSWER", output, "```````"} { |
| 189 | if !strings.Contains(string(md), expected) { |
| 190 | t.Fatal("missing markdown content") |
| 191 | } |
| 192 | } |
| 193 | if strings.Contains(string(md), "MUST-NOT-APPEAR") { |
| 194 | t.Fatal("snapshot leaked later message") |
| 195 | } |
| 196 | } |
| 197 |