Merge branch 'staging' into feat/byok

Author: kah-seng

.gitignore

@@ -50,3 +50,6 @@ overleaf.kubeconfig
 # coverage report
 coverage.out
 coverage.html
+
+# claude code
+CLAUDE.md

go.sum

@@ -55,6 +55,8 @@ github.com/golang/snappy v1.0.0/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEW
 github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
 github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
 github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
+github.com/google/subcommands v1.2.0 h1:vWQspBTo2nEqTUFita5/KeEWlUL8kQObDFbub/EN9oE=
+github.com/google/subcommands v1.2.0/go.mod h1:ZjhPrFU+Olkh9WazFPsl27BQ4UPiG37m3yTrtFlrHVk=
 github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
 github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/google/wire v0.7.0 h1:JxUKI6+CVBgCO2WToKy/nQk0sS+amI9z9EjVmdaocj4=
@@ -160,6 +162,8 @@ golang.org/x/crypto v0.41.0/go.mod h1:pO5AFd7FA68rFak7rOAGVuygIISepHftHnr8dr6+sU
 golang.org/x/exp v0.0.0-20250819193227-8b4c13bb791b h1:DXr+pvt3nC887026GRP39Ej11UATqWDmWuS99x26cD0=
 golang.org/x/exp v0.0.0-20250819193227-8b4c13bb791b/go.mod h1:4QTo5u+SEIbbKW1RacMZq1YEfOBqeXa19JeshGi+zc4=
 golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4=
+golang.org/x/mod v0.27.0 h1:kb+q2PyFnEADO2IEF935ehFUXlWiNjJWtRNgBLSfbxQ=
+golang.org/x/mod v0.27.0/go.mod h1:rWI627Fq0DEoudcK+MBkNkCe0EetEaDSwJJkCcjpazc=
 golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
 golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
 golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c=
@@ -188,6 +192,8 @@ golang.org/x/text v0.28.0/go.mod h1:U8nCwOR8jO/marOQ0QbDiOngZVEBB7MAiitBuMjXiNU=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
 golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc=
+golang.org/x/tools v0.36.0 h1:kWS0uv/zsvHEle1LbV5LE8QujrxB3wfQyxHfhOk0Qkg=
+golang.org/x/tools v0.36.0/go.mod h1:WBDiHKJK8YgLHlcQPYQzNCkUxUypCaa5ZegCVutKm+s=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 gonum.org/v1/gonum v0.16.0 h1:5+ul4Swaf3ESvrOnidPp4GZbzf0mxVQpDCYUQE7OJfk=
 gonum.org/v1/gonum v0.16.0/go.mod h1:fef3am4MQ93R2HHpKnLk4/Tbh/s0+wqD5nfa6Pnwy4E=

helm-chart/values.yaml

@@ -14,7 +14,7 @@ paperdebuggerMcpServer:
   image: ghcr.io/paperdebugger/paperdebugger-mcp-server:main-14409c5
 
 paperdebuggerXtraMcpServer:
-  image: ghcr.io/paperdebugger/xtragpt-mcp-server:sha-95f70ce
+  image: ghcr.io/paperdebugger/xtragpt-mcp-server:sha-510bc96
 
 mongo:
   in_cluster: true

internal/api/chat/create_conversation_message_stream_v2.go

@@ -303,7 +303,7 @@ func (s *ChatServerV2) CreateConversationMessageStream(
 		}
 	}
 
-	openaiChatHistory, inappChatHistory, err := s.aiClientV2.ChatCompletionStreamV2(ctx, stream, conversation.ID.Hex(), modelSlug, conversation.OpenaiChatHistoryCompletion, llmProvider)
+	openaiChatHistory, inappChatHistory, err := s.aiClientV2.ChatCompletionStreamV2(ctx, stream, conversation.UserID, conversation.ID.Hex(), modelSlug, conversation.OpenaiChatHistoryCompletion, llmProvider)
 	if err != nil {
 		return s.sendStreamError(stream, err)
 	}
@@ -329,7 +329,7 @@ func (s *ChatServerV2) CreateConversationMessageStream(
 			for i, bsonMsg := range conversation.InappChatHistory {
 				protoMessages[i] = mapper.BSONToChatMessageV2(bsonMsg)
 			}
-			title, err := s.aiClientV2.GetConversationTitleV2(ctx, protoMessages, llmProvider, modelSlug)
+			title, err := s.aiClientV2.GetConversationTitleV2(ctx, conversation.UserID, protoMessages, llmProvider, modelSlug)
 			if err != nil {
 				s.logger.Error("Failed to get conversation title", "error", err, "conversationID", conversation.ID.Hex())
 				return

internal/api/chat/get_citation_keys.go

@@ -0,0 +1,51 @@
+package chat
+
+import (
+	"context"
+
+	"paperdebugger/internal/libs/contextutil"
+	"paperdebugger/internal/libs/shared"
+	"paperdebugger/internal/models"
+	chatv2 "paperdebugger/pkg/gen/api/chat/v2"
+)
+
+func (s *ChatServerV2) GetCitationKeys(
+	ctx context.Context,
+	req *chatv2.GetCitationKeysRequest,
+) (*chatv2.GetCitationKeysResponse, error) {
+	if req.GetSentence() == "" {
+		return nil, shared.ErrBadRequest("sentence is required")
+	}
+	if req.GetProjectId() == "" {
+		return nil, shared.ErrBadRequest("project_id is required")
+	}
+
+	actor, err := contextutil.GetActor(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	settings, err := s.userService.GetUserSettings(ctx, actor.ID)
+	if err != nil {
+		return nil, err
+	}
+
+	llmProvider := &models.LLMProviderConfig{
+		APIKey: settings.OpenAIAPIKey,
+	}
+
+	citationKeys, err := s.aiClientV2.GetCitationKeys(
+		ctx,
+		req.GetSentence(),
+		actor.ID,
+		req.GetProjectId(),
+		llmProvider,
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	return &chatv2.GetCitationKeysResponse{
+		CitationKeys: citationKeys,
+	}, nil
+}

internal/api/grpc.go

@@ -15,6 +15,7 @@ import (
 	chatv2 "paperdebugger/pkg/gen/api/chat/v2"
 	commentv1 "paperdebugger/pkg/gen/api/comment/v1"
 	projectv1 "paperdebugger/pkg/gen/api/project/v1"
+	usagev1 "paperdebugger/pkg/gen/api/usage/v1"
 	userv1 "paperdebugger/pkg/gen/api/user/v1"
 
 	// "github.com/grpc-ecosystem/go-grpc-middleware"
@@ -106,6 +107,7 @@ func NewGrpcServer(
 	userServer userv1.UserServiceServer,
 	projectServer projectv1.ProjectServiceServer,
 	commentServer commentv1.CommentServiceServer,
+	usageServer usagev1.UsageServiceServer,
 ) *GrpcServer {
 	grpcServer := &GrpcServer{}
 	grpcServer.userService = userService
@@ -121,5 +123,6 @@ func NewGrpcServer(
 	userv1.RegisterUserServiceServer(grpcServer.Server, userServer)
 	projectv1.RegisterProjectServiceServer(grpcServer.Server, projectServer)
 	commentv1.RegisterCommentServiceServer(grpcServer.Server, commentServer)
+	usagev1.RegisterUsageServiceServer(grpcServer.Server, usageServer)
 	return grpcServer
 }

internal/api/mapper/user.go

@@ -35,7 +35,7 @@ func MapProtoSettingsToModel(settings *userv1.Settings) *models.Settings {
 	return &models.Settings{
 		ShowShortcutsAfterSelection:  settings.ShowShortcutsAfterSelection,
 		FullWidthPaperDebuggerButton: settings.FullWidthPaperDebuggerButton,
-		EnableCompletion:             settings.EnableCompletion,
+		EnableCitationSuggestion:             settings.EnableCitationSuggestion,
 		FullDocumentRag:              settings.FullDocumentRag,
 		ShowedOnboarding:             settings.ShowedOnboarding,
 		OpenAIAPIKey:                 settings.OpenaiApiKey,
@@ -62,7 +62,7 @@ func MapModelSettingsToProto(settings *models.Settings) *userv1.Settings {
 	return &userv1.Settings{
 		ShowShortcutsAfterSelection:  settings.ShowShortcutsAfterSelection,
 		FullWidthPaperDebuggerButton: settings.FullWidthPaperDebuggerButton,
-		EnableCompletion:             settings.EnableCompletion,
+		EnableCitationSuggestion:             settings.EnableCitationSuggestion,
 		FullDocumentRag:              settings.FullDocumentRag,
 		ShowedOnboarding:             settings.ShowedOnboarding,
 		OpenaiApiKey:                 settings.OpenAIAPIKey,

internal/api/server.go

@@ -6,17 +6,23 @@ import (
 	"fmt"
 	"net"
 	"net/http"
+	"os"
+	"os/signal"
 	"strings"
+	"syscall"
 
 	"paperdebugger/internal/libs/logger"
 	"paperdebugger/internal/libs/metadatautil"
 	"paperdebugger/internal/libs/shared"
+	"paperdebugger/internal/services"
+	aiclient "paperdebugger/internal/services/toolkit/client"
 	authv1 "paperdebugger/pkg/gen/api/auth/v1"
 	chatv1 "paperdebugger/pkg/gen/api/chat/v1"
 	chatv2 "paperdebugger/pkg/gen/api/chat/v2"
 	commentv1 "paperdebugger/pkg/gen/api/comment/v1"
 	projectv1 "paperdebugger/pkg/gen/api/project/v1"
 	sharedv1 "paperdebugger/pkg/gen/api/shared/v1"
+	usagev1 "paperdebugger/pkg/gen/api/usage/v1"
 	userv1 "paperdebugger/pkg/gen/api/user/v1"
 
 	"github.com/gin-gonic/gin"
@@ -30,25 +36,37 @@ import (
 )
 
 type Server struct {
-	grpcServer *GrpcServer
-	ginServer  *GinServer
+	grpcServer     *GrpcServer
+	ginServer      *GinServer
+	pricingService *services.PricingService
+	aiClientV2     *aiclient.AIClientV2
 
 	logger *logger.Logger
 }
 
 func NewServer(
 	grpcServer *GrpcServer,
 	ginServer *GinServer,
+	pricingService *services.PricingService,
+	aiClientV2 *aiclient.AIClientV2,
 	logger *logger.Logger,
 ) *Server {
 	return &Server{
-		grpcServer: grpcServer,
-		ginServer:  ginServer,
-		logger:     logger,
+		grpcServer:     grpcServer,
+		ginServer:      ginServer,
+		pricingService: pricingService,
+		aiClientV2:     aiClientV2,
+		logger:         logger,
 	}
 }
 
 func (s *Server) Run(addr string) {
+	// Start the pricing updater in the background
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	s.pricingService.StartPriceUpdater(ctx)
+
 	listener, err := net.Listen("tcp", ":0")
 	if err != nil {
 		s.logger.Fatalf("failed to start grpc server listener: %v", err)
@@ -105,6 +123,22 @@ func (s *Server) Run(addr string) {
 		s.logger.Fatalf("failed to register comment service grpc gateway: %v", err)
 		return
 	}
+	err = usagev1.RegisterUsageServiceHandler(context.Background(), mux, client)
+	if err != nil {
+		s.logger.Fatalf("failed to register usage service grpc gateway: %v", err)
+		return
+	}
+
+	// Set up signal handling for graceful shutdown
+	sigChan := make(chan os.Signal, 1)
+	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
+
+	go func() {
+		<-sigChan
+		s.logger.Info("[PAPERDEBUGGER] received shutdown signal, shutting down gracefully...")
+		s.Shutdown()
+		os.Exit(0)
+	}()
 
 	s.logger.Infof("[PAPERDEBUGGER] http server listening on %s", addr)
 	s.ginServer.Any("/_pd/api/*path", func(c *gin.Context) { mux.ServeHTTP(c.Writer, c.Request) })
@@ -114,6 +148,16 @@ func (s *Server) Run(addr string) {
 	}
 }
 
+// Shutdown gracefully shuts down all server components.
+func (s *Server) Shutdown() {
+	s.logger.Info("[PAPERDEBUGGER] shutting down AI client (draining usage records)...")
+	s.aiClientV2.Shutdown()
+	s.logger.Info("[PAPERDEBUGGER] AI client shutdown complete")
+
+	s.grpcServer.GracefulStop()
+	s.logger.Info("[PAPERDEBUGGER] gRPC server shutdown complete")
+}
+
 func (s *Server) metadataAnnotator() func(ctx context.Context, req *http.Request) metadata.MD {
 	return func(ctx context.Context, req *http.Request) metadata.MD {
 		md := metadata.New(map[string]string{})

internal/api/usage/get_session_usage.go

@@ -0,0 +1,40 @@
+package usage
+
+import (
+	"context"
+
+	"paperdebugger/internal/libs/contextutil"
+	usagev1 "paperdebugger/pkg/gen/api/usage/v1"
+
+	"google.golang.org/protobuf/types/known/timestamppb"
+)
+
+func (s *UsageServer) GetSessionUsage(
+	ctx context.Context,
+	req *usagev1.GetSessionUsageRequest,
+) (*usagev1.GetSessionUsageResponse, error) {
+	actor, err := contextutil.GetActor(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	// Get session with costs already calculated by the service layer
+	session, err := s.usageService.GetActiveSessionWithCosts(ctx, actor.ID)
+	if err != nil {
+		return nil, err
+	}
+
+	if session == nil {
+		return &usagev1.GetSessionUsageResponse{
+			Session: nil,
+		}, nil
+	}
+
+	return &usagev1.GetSessionUsageResponse{
+		Session: &usagev1.SessionUsage{
+			SessionExpiry: timestamppb.New(session.SessionExpiry),
+			Models:        convertModelsToProto(session.Models),
+			TotalCostUsd:  session.TotalCostUSD,
+		},
+	}, nil
+}

internal/api/usage/get_weekly_usage.go

@@ -0,0 +1,32 @@
+package usage
+
+import (
+	"context"
+
+	"paperdebugger/internal/libs/contextutil"
+	usagev1 "paperdebugger/pkg/gen/api/usage/v1"
+)
+
+func (s *UsageServer) GetWeeklyUsage(
+	ctx context.Context,
+	req *usagev1.GetWeeklyUsageRequest,
+) (*usagev1.GetWeeklyUsageResponse, error) {
+	actor, err := contextutil.GetActor(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	// Get weekly stats with costs already calculated by the service layer
+	stats, err := s.usageService.GetWeeklyUsageWithCosts(ctx, actor.ID)
+	if err != nil {
+		return nil, err
+	}
+
+	return &usagev1.GetWeeklyUsageResponse{
+		Usage: &usagev1.WeeklyUsage{
+			Models:       convertModelsToProto(stats.Models),
+			SessionCount: stats.SessionCount,
+			TotalCostUsd: stats.TotalCostUSD,
+		},
+	}, nil
+}