Include Attached AI Sessions in Case Reports

agent/modules/export/case.go

@@ -58,6 +58,30 @@ func (export *Export) getCaseDetailsFromServer(caseId string) (*CaseTemplateInpu
 		}).WithError(err).Error("failed to get case observables")
 	}
 
+	for _, attachment := range caseTemplateInput.Attachments {
+		if attachment.ArtifactType == "assistant_chat" {
+			currSessionId := attachment.Value
+			sessionExists := false
+			for _, session := range caseTemplateInput.AssistantSessions {
+				if session.Session.SessionId == currSessionId {
+					sessionExists = true
+					break
+				}
+			}
+			if !sessionExists {
+				sessionDetails := &model.AssistantSessionDetails{}
+				_, err := export.agent.Client.SendAuthorizedObject("GET", fmt.Sprintf("/api/assistant/sessions/%s", currSessionId), nil, sessionDetails)
+				if err != nil {
+					log.WithFields(log.Fields{
+						"sessionId": currSessionId,
+					}).WithError(err).Error("failed to get assistant session for attachment")
+				} else {
+					caseTemplateInput.AssistantSessions = append(caseTemplateInput.AssistantSessions, sessionDetails)
+				}
+			}
+		}
+	}
+
 	// Get related events
 	_, err = export.agent.Client.SendAuthorizedObject("GET", fmt.Sprintf("/api/case/events/%s", caseId), nil, &caseTemplateInput.RelatedEvents)
 	if err != nil {
@@ -119,14 +143,15 @@ func (export *Export) getCaseDetailsFromServer(caseId string) (*CaseTemplateInpu
 }
 
 type CaseTemplateInput struct {
-	Case          *model.Case
-	Comments      []*model.Comment
-	Attachments   []*model.Artifact
-	Observables   []*model.Artifact
-	Detections    []*model.Detection
-	RelatedEvents []*model.RelatedEvent
-	History       []*model.Auditable
-	TotalHours    float64
+	Case              *model.Case
+	Comments          []*model.Comment
+	Attachments       []*model.Artifact
+	AssistantSessions []*model.AssistantSessionDetails
+	Observables       []*model.Artifact
+	Detections        []*model.Detection
+	RelatedEvents     []*model.RelatedEvent
+	History           []*model.Auditable
+	TotalHours        float64
 }
 
 func (export *Export) generateCaseReport(caseId string) ([]byte, error) {

agent/modules/export/export.go

@@ -270,23 +270,24 @@ func (export *Export) populateTemplatesCache() {
 		"templatePath": export.templatePath,
 	}).Info("Refreshing templates cache for export module")
 	master := template.New("export").Funcs(template.FuncMap{
-		"getUserDetail":         export.getUserDetail,
-		"formatDateTime":        export.formatDateTime,
-		"join":                  strings.Join,
-		"lower":                 strings.ToLower,
-		"upper":                 strings.ToUpper,
-		"sortHistory":           export.sortHistory,
-		"sortComments":          export.sortComments,
-		"sortRelatedEvents":     export.sortRelatedEvents,
-		"sortArtifacts":         export.sortArtifacts,
-		"sortDetections":        export.sortDetections,
-		"sortMetrics":           export.sortMetrics,
-		"sortAssistantMessages": export.sortAssistantMessages,
-		"formatNumber":          export.formatNumber,
-		"parseJSON":             export.parseJSON,
-		"toJSON":                export.toJSON,
-		"add":                   export.add,
-		"stripEmoji":            export.stripEmoji,
+		"getUserDetail":               export.getUserDetail,
+		"formatDateTime":              export.formatDateTime,
+		"join":                        strings.Join,
+		"lower":                       strings.ToLower,
+		"upper":                       strings.ToUpper,
+		"sortHistory":                 export.sortHistory,
+		"sortComments":                export.sortComments,
+		"sortRelatedEvents":           export.sortRelatedEvents,
+		"sortArtifacts":               export.sortArtifacts,
+		"sortDetections":              export.sortDetections,
+		"sortMetrics":                 export.sortMetrics,
+		"sortAssistantMessages":       export.sortAssistantMessages,
+		"sortAssistantSessionDetails": export.SortAssistantSessionDetails,
+		"formatNumber":                export.formatNumber,
+		"parseJSON":                   export.parseJSON,
+		"toJSON":                      export.toJSON,
+		"add":                         export.add,
+		"stripEmoji":                  export.stripEmoji,
 	})
 
 	var err error

agent/modules/export/export_test.go

@@ -191,6 +191,178 @@ func TestGetCaseDetailsFromServer(t *testing.T) {
 	assert.NotNil(t, templateInput)
 }
 
+func TestGetCaseDetailsFromServer_SessionAttachments(t *testing.T) {
+	export := NewExport(agent.NewAgent(&config.AgentConfig{}, "test-version"))
+	export.agent.Client = web.NewClient("http://localhost:8080", true)
+	export.agent.Client.Auth = FakeClientAuth{}
+
+	caseId := "test-case-with-sessions"
+
+	// Mock attachments JSON with assistant_chat artifacts
+	attachmentsWithSessionsJson := `[{
+		"id":"attachment1",
+		"createTime": "2025-07-01T16:41:09.698562704-04:00",
+		"updateTime": "2025-07-01T16:41:09.698562704-04:00",
+		"description": "This is a test file",
+		"artifactType": "file"
+	},{
+		"id":"attachment2",
+		"createTime": "2025-07-01T16:42:09.698562704-04:00",
+		"updateTime": "2025-07-01T16:42:09.698562704-04:00",
+		"description": "Assistant chat session 1",
+		"artifactType": "assistant_chat",
+		"value": "chat_session_123"
+	},{
+		"id":"attachment3",
+		"createTime": "2025-07-01T16:43:09.698562704-04:00",
+		"updateTime": "2025-07-01T16:43:09.698562704-04:00",
+		"description": "Assistant chat session 2",
+		"artifactType": "assistant_chat",
+		"value": "chat_session_456"
+	},{
+		"id":"attachment4",
+		"createTime": "2025-07-01T16:44:09.698562704-04:00",
+		"updateTime": "2025-07-01T16:44:09.698562704-04:00",
+		"description": "Duplicate session reference",
+		"artifactType": "assistant_chat",
+		"value": "chat_session_123"
+	}]`
+
+	sessionDetails1Json := `{
+		"session": {
+			"id": "session1",
+			"createTime": "2025-07-01T16:41:09.698562704-04:00",
+			"updateTime": "2025-07-01T16:41:09.698562704-04:00",
+			"sessionId": "chat_session_123",
+			"title": "First Assistant Session",
+			"tags": ["investigation"]
+		},
+		"history": [
+			{
+				"id": "msg1",
+				"createTime": "2025-07-01T16:41:09.698562704-04:00",
+				"updateTime": "2025-07-01T16:41:09.698562704-04:00",
+				"sessionId": "chat_session_123",
+				"message": {
+					"id": "msg1",
+					"role": "user",
+					"contentStr": "What is this alert about?"
+				}
+			}
+		]
+	}`
+
+	sessionDetails2Json := `{
+		"session": {
+			"id": "session2",
+			"createTime": "2025-07-01T16:42:09.698562704-04:00",
+			"updateTime": "2025-07-01T16:42:09.698562704-04:00",
+			"sessionId": "chat_session_456",
+			"title": "Second Assistant Session",
+			"tags": ["analysis"]
+		},
+		"history": [
+			{
+				"id": "msg2",
+				"createTime": "2025-07-01T16:42:09.698562704-04:00",
+				"updateTime": "2025-07-01T16:42:09.698562704-04:00",
+				"sessionId": "chat_session_456",
+				"message": {
+					"id": "msg2",
+					"role": "user",
+					"contentStr": "Can you help analyze this?"
+				}
+			}
+		]
+	}`
+
+	// Mock successful responses - order matters based on the sequence of API calls in getCaseDetailsFromServer
+	export.agent.Client.MockStringResponse(caseJson, 200, nil)                    // Get Case
+	export.agent.Client.MockStringResponse(commentJson, 200, nil)                 // Get Case Comments
+	export.agent.Client.MockStringResponse(attachmentsWithSessionsJson, 200, nil) // Get Case Attachments with sessions
+	export.agent.Client.MockStringResponse(observablesJson, 200, nil)             // Get Case Observables
+	export.agent.Client.MockStringResponse(sessionDetails1Json, 200, nil)         // Get Assistant Session 1 (chat_session_123)
+	export.agent.Client.MockStringResponse(sessionDetails2Json, 200, nil)         // Get Assistant Session 2 (chat_session_456) - duplicate chat_session_123 should not trigger another call
+	export.agent.Client.MockStringResponse(eventsJson, 200, nil)                  // Get Case Related Events
+	export.agent.Client.MockStringResponse(detectionJson, 200, nil)               // Get Detection for Event
+	export.agent.Client.MockStringResponse(historyJson, 200, nil)                 // Get Case History
+
+	templateInput, err := export.getCaseDetailsFromServer(caseId)
+	assert.Nil(t, err)
+	assert.NotNil(t, templateInput)
+
+	// Verify basic case data
+	assert.Equal(t, "case1", templateInput.Case.Id)
+	assert.Len(t, templateInput.Comments, 1)
+	assert.Len(t, templateInput.Attachments, 4) // 1 file + 3 assistant_chat attachments
+
+	// Verify assistant sessions were fetched
+	assert.Len(t, templateInput.AssistantSessions, 2) // Only 2 unique sessions despite 3 assistant_chat attachments
+
+	// Verify session details
+	sessionIds := make(map[string]bool)
+	for _, session := range templateInput.AssistantSessions {
+		sessionIds[session.Session.SessionId] = true
+	}
+	assert.True(t, sessionIds["chat_session_123"])
+	assert.True(t, sessionIds["chat_session_456"])
+
+	// Verify session titles
+	var session1, session2 *model.AssistantSessionDetails
+	for _, session := range templateInput.AssistantSessions {
+		if session.Session.SessionId == "chat_session_123" {
+			session1 = session
+		} else if session.Session.SessionId == "chat_session_456" {
+			session2 = session
+		}
+	}
+	assert.NotNil(t, session1)
+	assert.NotNil(t, session2)
+	assert.Equal(t, "First Assistant Session", session1.Session.Title)
+	assert.Equal(t, "Second Assistant Session", session2.Session.Title)
+	assert.Len(t, session1.History, 1)
+	assert.Len(t, session2.History, 1)
+}
+
+func TestGetCaseDetailsFromServer_SessionAttachmentError(t *testing.T) {
+	export := NewExport(agent.NewAgent(&config.AgentConfig{}, "test-version"))
+	export.agent.Client = web.NewClient("http://localhost:8080", true)
+	export.agent.Client.Auth = FakeClientAuth{}
+
+	caseId := "test-case-session-error"
+
+	// Mock attachments JSON with assistant_chat artifact
+	attachmentsWithSessionJson := `[{
+		"id":"attachment1",
+		"createTime": "2025-07-01T16:41:09.698562704-04:00",
+		"updateTime": "2025-07-01T16:41:09.698562704-04:00",
+		"description": "Assistant chat session",
+		"artifactType": "assistant_chat",
+		"value": "chat_session_error"
+	}]`
+
+	// Mock successful responses for case data, but error for session
+	export.agent.Client.MockStringResponse(caseJson, 200, nil)                   // Get Case
+	export.agent.Client.MockStringResponse(commentJson, 200, nil)                // Get Case Comments
+	export.agent.Client.MockStringResponse(attachmentsWithSessionJson, 200, nil) // Get Case Attachments
+	export.agent.Client.MockStringResponse("", 500, assert.AnError)              // Get Assistant Session fails
+	export.agent.Client.MockStringResponse(observablesJson, 200, nil)            // Get Case Observables
+	export.agent.Client.MockStringResponse(eventsJson, 200, nil)                 // Get Case Related Events
+	export.agent.Client.MockStringResponse(detectionJson, 200, nil)              // Get Detection for Event
+	export.agent.Client.MockStringResponse(historyJson, 200, nil)                // Get Case History
+
+	templateInput, err := export.getCaseDetailsFromServer(caseId)
+	assert.Nil(t, err) // Error is logged but not returned
+	assert.NotNil(t, templateInput)
+
+	// Verify case data is still returned
+	assert.Equal(t, "case1", templateInput.Case.Id)
+	assert.Len(t, templateInput.Attachments, 1)
+
+	// Verify no assistant sessions were added due to error
+	assert.Len(t, templateInput.AssistantSessions, 0)
+}
+
 func TestPrerequisiteModules(t *testing.T) {
 	export := NewExport(nil)
 	assert.Nil(t, export.PrerequisiteModules())

agent/modules/export/sort.go

@@ -246,6 +246,42 @@ func (export *Export) sortAssistantMessages(field string, dir string, list []*mo
 	return list
 }
 
+func (export *Export) SortAssistantSessionDetails(field string, dir string, list []*model.AssistantSessionDetails) []*model.AssistantSessionDetails {
+	sort.Slice(list, func(i, j int) bool {
+		a := list[i]
+		b := list[j]
+		if a == nil || b == nil || a.Session == nil || b.Session == nil {
+			return false
+		}
+		switch strings.ToLower(field) {
+		case "id":
+			return export.compareWithDir(a.Session.Id, b.Session.Id, dir)
+		case "createtime":
+			return export.compareWithDir(a.Session.CreateTime, b.Session.CreateTime, dir)
+		case "updatetime":
+			return export.compareWithDir(a.Session.UpdateTime, b.Session.UpdateTime, dir)
+		case "deletetime":
+			if a.Session.DeleteTime == nil || b.Session.DeleteTime == nil {
+				return false
+			}
+			return export.compareWithDir(*a.Session.DeleteTime, *b.Session.DeleteTime, dir)
+		case "userid":
+			return export.compareWithDir(a.Session.UserId, b.Session.UserId, dir)
+		case "kind":
+			return export.compareWithDir(a.Session.Kind, b.Session.Kind, dir)
+		case "operation":
+			return export.compareWithDir(a.Session.Operation, b.Session.Operation, dir)
+		case "title":
+			return export.compareWithDir(a.Session.Title, b.Session.Title, dir)
+		case "sessionid":
+			return export.compareWithDir(a.Session.SessionId, b.Session.SessionId, dir)
+		default:
+			return false
+		}
+	})
+	return list
+}
+
 func (export *Export) sortMetrics(field string, dir string, list []*model.EventMetric) []*model.EventMetric {
 	sort.Slice(list, func(i, j int) bool {
 		a := list[i]

agent/modules/export/sort_test.go

@@ -363,3 +363,70 @@ func TestSortAssistantMessages(t *testing.T) {
 		assert.Equal(t, "msg1", sortedList[2].Id)
 	})
 }
+
+func TestSortAssistantSessionDetails(t *testing.T) {
+	export := NewExport(nil)
+
+	now := time.Now()
+	time1 := now.Add(-3 * time.Hour)
+	time2 := now.Add(-2 * time.Hour)
+	time3 := now.Add(-1 * time.Hour)
+
+	createSession := func(id, sessionId, title string, createTime *time.Time) *model.AssistantSessionDetails {
+		return &model.AssistantSessionDetails{
+			Session: &model.AssistantSession{
+				Auditable: model.Auditable{Id: id, CreateTime: createTime},
+				SessionId: sessionId,
+				Title:     title,
+			},
+		}
+	}
+
+	t.Run("SortByIdAsc", func(t *testing.T) {
+		list := []*model.AssistantSessionDetails{
+			createSession("session3", "chat_3", "C", &time1),
+			createSession("session1", "chat_1", "A", &time2),
+			createSession("session2", "chat_2", "B", &time3),
+		}
+		sorted := export.SortAssistantSessionDetails("id", "asc", list)
+		assert.Equal(t, "session1", sorted[0].Session.Id)
+		assert.Equal(t, "session2", sorted[1].Session.Id)
+		assert.Equal(t, "session3", sorted[2].Session.Id)
+	})
+
+	t.Run("SortByTitleDesc", func(t *testing.T) {
+		list := []*model.AssistantSessionDetails{
+			createSession("s1", "c1", "Alpha", &time1),
+			createSession("s2", "c2", "Charlie", &time2),
+			createSession("s3", "c3", "Bravo", &time3),
+		}
+		sorted := export.SortAssistantSessionDetails("title", "desc", list)
+		assert.Equal(t, "Charlie", sorted[0].Session.Title)
+		assert.Equal(t, "Bravo", sorted[1].Session.Title)
+		assert.Equal(t, "Alpha", sorted[2].Session.Title)
+	})
+
+	t.Run("SortByCreateTime", func(t *testing.T) {
+		list := []*model.AssistantSessionDetails{
+			createSession("s1", "c1", "A", &time2),
+			createSession("s2", "c2", "B", &time3),
+			createSession("s3", "c3", "C", &time1),
+		}
+		sorted := export.SortAssistantSessionDetails("createtime", "asc", list)
+		assert.Equal(t, "s3", sorted[0].Session.Id)
+		assert.Equal(t, "s1", sorted[1].Session.Id)
+		assert.Equal(t, "s2", sorted[2].Session.Id)
+	})
+
+	t.Run("HandleNilValues", func(t *testing.T) {
+		list := []*model.AssistantSessionDetails{
+			createSession("s1", "c1", "A", &time1),
+			nil,
+			{Session: nil},
+		}
+		sorted := export.SortAssistantSessionDetails("id", "asc", list)
+		assert.NotNil(t, sorted[0].Session)
+		assert.Nil(t, sorted[1])
+		assert.Nil(t, sorted[2].Session)
+	})
+}