JohnCosta27/Haystack
1
0
Fork 0
Code Issues 7 Pull Requests Actions Packages Projects 1 Releases Wiki Activity

fix: notification system

Browse Source
This commit is contained in:
John Costa
2025-10-05 12:10:06 +01:00
parent 649cfe0b02
commit 980b42aa44
12 changed files with 191 additions and 124 deletions
Download Patch File Download Diff File Expand all files Collapse all files

63
backend/events.go
View File

@@ -5,77 +5,22 @@ import (
"fmt"
"net/http"
"screenmark/screenmark/middleware"
"screenmark/screenmark/notifications"
"strconv"
"github.com/google/uuid"
)
const (
IMAGE_TYPE = "image"
LIST_TYPE = "list"
)
type imageNotification struct {
Type string
ImageID uuid.UUID
ImageName string
Status string
}
type listNotification struct {
Type string
ListID uuid.UUID
Name string
Status string
}
type Notification struct {
image *imageNotification
list *listNotification
}
func getImageNotification(image imageNotification) Notification {
return Notification{
image: &image,
}
}
func getListNotification(list listNotification) Notification {
return Notification{
list: &list,
}
}
func (n Notification) MarshalJSON() ([]byte, error) {
if n.image != nil {
return json.Marshal(n.image)
}
if n.list != nil {
return json.Marshal(n.list)
}
return nil, fmt.Errorf("no image or list present")
}
func (n *Notification) UnmarshalJSON(data []byte) error {
return fmt.Errorf("unimplemented")
}
/*
* TODO: We have channels open every a user sends an image.
* We never close these channels.
*
* What is a reasonable default? Close the channel after 1 minute of inactivity?
*/
func CreateEventsHandler(notifier *Notifier[Notification]) http.HandlerFunc {
func CreateEventsHandler(notifier *notifications.Notifier[notifications.Notification]) http.HandlerFunc {
counter := 0
userSplitters := make(map[string]*ChannelSplitter[Notification])
userSplitters := make(map[string]*notifications.ChannelSplitter[notifications.Notification])
return func(w http.ResponseWriter, r *http.Request) {
_userId := r.Context().Value(middleware.USER_ID).(uuid.UUID)
@@ -98,7 +43,7 @@ func CreateEventsHandler(notifier *Notifier[Notification]) http.HandlerFunc {
userNotifications := notifier.Listeners[userId]
if _, exists := userSplitters[userId]; !exists {
splitter := NewChannelSplitter(userNotifications)
splitter := notifications.NewChannelSplitter(userNotifications)
userSplitters[userId] = &splitter
splitter.Listen()

6
backend/images/handler.go
View File

@@ -154,7 +154,11 @@ func (h *ImageHandler) uploadImage(w http.ResponseWriter, r *http.Request) {
h.logger.Info("About to add image")
h.processor.Add(newImage)
w.WriteHeader(http.StatusOK)
// We nullify the image's data, so we're not transferring all that
// data back to the frontend.
newImage.Image = nil
middleware.WriteJsonOrError(h.logger, newImage, w)
}
func (h *ImageHandler) deleteImage(w http.ResponseWriter, r *http.Request) {

6
backend/integration_test.go
View File

@@ -181,7 +181,11 @@ func setupTestContext(t *testing.T) *TestContext {
}
jwtManager := middleware.NewJwtManager([]byte("test-jwt-secret"))
router := setupRouter(db, jwtManager)
router, err := setupRouter(db, jwtManager)
if err != nil {
panic(err)
}
server := httptest.NewServer(router)
tc.db = db

5
backend/main.go
View File

@@ -28,7 +28,10 @@ func main() {
panic(err)
}
router := setupRouter(db, jwtManager)
router, err := setupRouter(db, jwtManager)
if err != nil {
panic(err)
}
port, exists := os.LookupEnv("PORT")
if !exists {

38
backend/notifications/channel_splitter.go Normal file
View File

@@ -0,0 +1,38 @@
package notifications
type ChannelSplitter[TNotification any] struct {
ch chan TNotification
Listeners map[string]chan TNotification
}
func (s *ChannelSplitter[TNotification]) Listen() {
go func() {
for {
select {
case msg := <-s.ch:
for _, v := range s.Listeners {
v <- msg
}
}
}
}()
}
func (s *ChannelSplitter[TNotification]) Add(id string) chan TNotification {
ch := make(chan TNotification)
s.Listeners[id] = ch
return ch
}
func (s *ChannelSplitter[TNotification]) Remove(id string) {
delete(s.Listeners, id)
}
func NewChannelSplitter[TNotification any](ch chan TNotification) ChannelSplitter[TNotification] {
return ChannelSplitter[TNotification]{
ch: ch,
Listeners: make(map[string]chan TNotification),
}
}

64
backend/notifications/image_notification.go Normal file
View File

@@ -0,0 +1,64 @@
package notifications
import (
"encoding/json"
"fmt"
"github.com/google/uuid"
)
const (
IMAGE_TYPE = "image"
LIST_TYPE = "list"
)
type ImageNotification struct {
Type string
ImageID uuid.UUID
ImageName string
Status string
}
type ListNotification struct {
Type string
ListID uuid.UUID
Name string
Status string
}
type Notification struct {
image *ImageNotification
list *ListNotification
}
func GetImageNotification(image ImageNotification) Notification {
return Notification{
image: &image,
}
}
func GetListNotification(list ListNotification) Notification {
return Notification{
list: &list,
}
}
func (n Notification) MarshalJSON() ([]byte, error) {
if n.image != nil {
return json.Marshal(n.image)
}
if n.list != nil {
return json.Marshal(n.list)
}
return nil, fmt.Errorf("no image or list present")
}
func (n *Notification) UnmarshalJSON(data []byte) error {
return fmt.Errorf("unimplemented")
}

41
backend/notifications.go → backend/notifications/notifications.go
View File

@@ -1,4 +1,4 @@
package main
package notifications
import (
"errors"
@@ -56,42 +56,3 @@ func NewNotifier[TNotification any](bufferSize int) Notifier[TNotification] {
Listeners: make(map[string]chan TNotification),
}
}
// ----------------------------------
type ChannelSplitter[TNotification any] struct {
ch chan TNotification
Listeners map[string]chan TNotification
}
func (s *ChannelSplitter[TNotification]) Listen() {
go func() {
for {
select {
case msg := <-s.ch:
for _, v := range s.Listeners {
v <- msg
}
}
}
}()
}
func (s *ChannelSplitter[TNotification]) Add(id string) chan TNotification {
ch := make(chan TNotification)
s.Listeners[id] = ch
return ch
}
func (s *ChannelSplitter[TNotification]) Remove(id string) {
delete(s.Listeners, id)
}
func NewChannelSplitter[TNotification any](ch chan TNotification) ChannelSplitter[TNotification] {
return ChannelSplitter[TNotification]{
ch: ch,
Listeners: make(map[string]chan TNotification),
}
}

2
backend/notifications_test.go → backend/notifications/notifications_test.go
View File

@@ -1,4 +1,4 @@
package main
package notifications
import (
"testing"

48
backend/processor/image.go
View File

@@ -2,11 +2,13 @@ package processor
import (
"context"
"fmt"
"screenmark/screenmark/.gen/haystack/haystack/model"
"screenmark/screenmark/agents"
"screenmark/screenmark/agents/client"
"screenmark/screenmark/limits"
"screenmark/screenmark/models"
"screenmark/screenmark/notifications"
"sync"
"github.com/charmbracelet/log"
@@ -24,6 +26,8 @@ type ImageProcessor struct {
// TODO: add the notifier here
Processor *Processor[model.Image]
notifier *notifications.Notifier[notifications.Notification]
}
func (p *ImageProcessor) setImageToProcess(ctx context.Context, image model.Image) {
@@ -71,6 +75,19 @@ func (p *ImageProcessor) processImage(image model.Image) {
var wg sync.WaitGroup
wg.Add(2)
imageNotification := notifications.GetImageNotification(notifications.ImageNotification{
Type: notifications.IMAGE_TYPE,
ImageID: image.ID,
ImageName: image.ImageName,
Status: string(model.Progress_InProgress),
})
err := p.notifier.SendAndCreate(image.UserID.String(), imageNotification)
if err != nil {
p.logger.Error("sending in progress notification", "err", err)
return
}
go func() {
p.describe(ctx, image)
wg.Done()
@@ -82,9 +99,34 @@ func (p *ImageProcessor) processImage(image model.Image) {
}()
wg.Wait()
// TODO: there is some repeated code here. The ergonomicts of the notifications,
// isn't the best.
imageNotification = notifications.GetImageNotification(notifications.ImageNotification{
Type: notifications.IMAGE_TYPE,
ImageID: image.ID,
ImageName: image.ImageName,
Status: string(model.Progress_Complete),
})
err = p.notifier.SendAndCreate(image.UserID.String(), imageNotification)
if err != nil {
p.logger.Error("sending done notification", "err", err)
return
}
}
func NewImageProcessor(logger *log.Logger, imageModel models.ImageModel, listModel models.StackModel, limitsManager limits.LimitsManagerMethods) ImageProcessor {
func NewImageProcessor(
logger *log.Logger,
imageModel models.ImageModel,
listModel models.StackModel,
limitsManager limits.LimitsManagerMethods,
notifier *notifications.Notifier[notifications.Notification],
) (ImageProcessor, error) {
if notifier == nil {
return ImageProcessor{}, fmt.Errorf("notifier is nil")
}
descriptionAgent := agents.NewDescriptionAgent(logger, imageModel)
stackAgent := agents.NewListAgent(logger, listModel, limitsManager)
@@ -93,9 +135,11 @@ func NewImageProcessor(logger *log.Logger, imageModel models.ImageModel, listMod
logger: logger,
descriptionAgent: descriptionAgent,
stackAgent: stackAgent,
notifier: notifier,
}
imageProcessor.Processor = NewProcessor(int(IMAGE_PROCESS_AT_A_TIME), imageProcessor.processImage)
return imageProcessor
return imageProcessor, nil
}

16
backend/router.go
View File

@@ -2,12 +2,14 @@ package main
import (
"database/sql"
"fmt"
"os"
"screenmark/screenmark/agents/client"
"screenmark/screenmark/auth"
"screenmark/screenmark/images"
"screenmark/screenmark/limits"
"screenmark/screenmark/models"
"screenmark/screenmark/notifications"
"screenmark/screenmark/processor"
"screenmark/screenmark/stacks"
@@ -25,22 +27,26 @@ func (client TestAiClient) GetImageInfo(imageName string, imageData []byte) (cli
return client.ImageInfo, nil
}
func setupRouter(db *sql.DB, jwtManager *ourmiddleware.JwtManager) chi.Router {
func setupRouter(db *sql.DB, jwtManager *ourmiddleware.JwtManager) (chi.Router, error) {
limitsManager := limits.CreateLimitsManager(db)
imageModel := models.NewImageModel(db)
stackModel := models.NewStackModel(db)
notifier := notifications.NewNotifier[notifications.Notification](10)
imageProcessorLogger := createLogger("Image Processor", os.Stdout)
imageProcessor := processor.NewImageProcessor(imageProcessorLogger, imageModel, stackModel, limitsManager)
imageProcessor, err := processor.NewImageProcessor(imageProcessorLogger, imageModel, stackModel, limitsManager, &notifier)
if err != nil {
return nil, fmt.Errorf("processor: %w", err)
}
go imageProcessor.Processor.Work()
stackHandler := stacks.CreateStackHandler(db, limitsManager, jwtManager)
authHandler := auth.CreateAuthHandler(db, jwtManager)
imageHandler := images.CreateImageHandler(db, limitsManager, jwtManager, imageProcessor.Processor)
notifier := NewNotifier[Notification](10)
r := chi.NewRouter()
r.Use(middleware.Logger)
@@ -56,5 +62,5 @@ func setupRouter(db *sql.DB, jwtManager *ourmiddleware.JwtManager) chi.Router {
r.Get("/", CreateEventsHandler(&notifier))
})
return r
return r, nil
}