session.go 4.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150
  1. package session
  2. import (
  3. "context"
  4. "database/sql"
  5. "github.com/google/uuid"
  6. "github.com/opencode-ai/opencode/internal/db"
  7. "github.com/opencode-ai/opencode/internal/pubsub"
  8. )
  9. type Session struct {
  10. ID string
  11. ParentSessionID string
  12. Title string
  13. MessageCount int64
  14. PromptTokens int64
  15. CompletionTokens int64
  16. Cost float64
  17. CreatedAt int64
  18. UpdatedAt int64
  19. }
  20. type Service interface {
  21. pubsub.Suscriber[Session]
  22. Create(ctx context.Context, title string) (Session, error)
  23. CreateTitleSession(ctx context.Context, parentSessionID string) (Session, error)
  24. CreateTaskSession(ctx context.Context, toolCallID, parentSessionID, title string) (Session, error)
  25. Get(ctx context.Context, id string) (Session, error)
  26. List(ctx context.Context) ([]Session, error)
  27. Save(ctx context.Context, session Session) (Session, error)
  28. Delete(ctx context.Context, id string) error
  29. }
  30. type service struct {
  31. *pubsub.Broker[Session]
  32. q db.Querier
  33. }
  34. func (s *service) Create(ctx context.Context, title string) (Session, error) {
  35. dbSession, err := s.q.CreateSession(ctx, db.CreateSessionParams{
  36. ID: uuid.New().String(),
  37. Title: title,
  38. })
  39. if err != nil {
  40. return Session{}, err
  41. }
  42. session := s.fromDBItem(dbSession)
  43. s.Publish(pubsub.CreatedEvent, session)
  44. return session, nil
  45. }
  46. func (s *service) CreateTaskSession(ctx context.Context, toolCallID, parentSessionID, title string) (Session, error) {
  47. dbSession, err := s.q.CreateSession(ctx, db.CreateSessionParams{
  48. ID: toolCallID,
  49. ParentSessionID: sql.NullString{String: parentSessionID, Valid: true},
  50. Title: title,
  51. })
  52. if err != nil {
  53. return Session{}, err
  54. }
  55. session := s.fromDBItem(dbSession)
  56. s.Publish(pubsub.CreatedEvent, session)
  57. return session, nil
  58. }
  59. func (s *service) CreateTitleSession(ctx context.Context, parentSessionID string) (Session, error) {
  60. dbSession, err := s.q.CreateSession(ctx, db.CreateSessionParams{
  61. ID: "title-" + parentSessionID,
  62. ParentSessionID: sql.NullString{String: parentSessionID, Valid: true},
  63. Title: "Generate a title",
  64. })
  65. if err != nil {
  66. return Session{}, err
  67. }
  68. session := s.fromDBItem(dbSession)
  69. s.Publish(pubsub.CreatedEvent, session)
  70. return session, nil
  71. }
  72. func (s *service) Delete(ctx context.Context, id string) error {
  73. session, err := s.Get(ctx, id)
  74. if err != nil {
  75. return err
  76. }
  77. err = s.q.DeleteSession(ctx, session.ID)
  78. if err != nil {
  79. return err
  80. }
  81. s.Publish(pubsub.DeletedEvent, session)
  82. return nil
  83. }
  84. func (s *service) Get(ctx context.Context, id string) (Session, error) {
  85. dbSession, err := s.q.GetSessionByID(ctx, id)
  86. if err != nil {
  87. return Session{}, err
  88. }
  89. return s.fromDBItem(dbSession), nil
  90. }
  91. func (s *service) Save(ctx context.Context, session Session) (Session, error) {
  92. dbSession, err := s.q.UpdateSession(ctx, db.UpdateSessionParams{
  93. ID: session.ID,
  94. Title: session.Title,
  95. PromptTokens: session.PromptTokens,
  96. CompletionTokens: session.CompletionTokens,
  97. Cost: session.Cost,
  98. })
  99. if err != nil {
  100. return Session{}, err
  101. }
  102. session = s.fromDBItem(dbSession)
  103. s.Publish(pubsub.UpdatedEvent, session)
  104. return session, nil
  105. }
  106. func (s *service) List(ctx context.Context) ([]Session, error) {
  107. dbSessions, err := s.q.ListSessions(ctx)
  108. if err != nil {
  109. return nil, err
  110. }
  111. sessions := make([]Session, len(dbSessions))
  112. for i, dbSession := range dbSessions {
  113. sessions[i] = s.fromDBItem(dbSession)
  114. }
  115. return sessions, nil
  116. }
  117. func (s service) fromDBItem(item db.Session) Session {
  118. return Session{
  119. ID: item.ID,
  120. ParentSessionID: item.ParentSessionID.String,
  121. Title: item.Title,
  122. MessageCount: item.MessageCount,
  123. PromptTokens: item.PromptTokens,
  124. CompletionTokens: item.CompletionTokens,
  125. Cost: item.Cost,
  126. CreatedAt: item.CreatedAt,
  127. UpdatedAt: item.UpdatedAt,
  128. }
  129. }
  130. func NewService(q db.Querier) Service {
  131. broker := pubsub.NewBroker[Session]()
  132. return &service{
  133. broker,
  134. q,
  135. }
  136. }