From f228d3def9667387215f3c37e799039c59c3ff70 Mon Sep 17 00:00:00 2001 From: locnguyen1986 Date: Fri, 30 Jan 2026 13:57:57 +0700 Subject: [PATCH 1/2] ad usage and --- apps/web/src/components/form/login.tsx | 23 ++++- apps/web/src/constants/ui.ts | 1 + apps/web/src/routes/__root.tsx | 92 ++++++++++++++++++- apps/web/src/routes/auth/callback.tsx | 17 +++- apps/web/src/routes/login.tsx | 52 +++++++++++ services/llm-api/cmd/server/wire_gen.go | 2 +- .../messageshandler/messages_handler.go | 61 +++++++++++- 7 files changed, 236 insertions(+), 12 deletions(-) create mode 100644 apps/web/src/routes/login.tsx diff --git a/apps/web/src/components/form/login.tsx b/apps/web/src/components/form/login.tsx index fbcb210b..fae77e52 100644 --- a/apps/web/src/components/form/login.tsx +++ b/apps/web/src/components/form/login.tsx @@ -21,7 +21,7 @@ export function LoginForm({ className, onSuccess, ...props -}: React.ComponentProps<"div"> & { onSuccess?: () => void }) { +}: React.ComponentProps<"div"> & { onSuccess?: (redirectUrl?: string) => void }) { const [isGoogleLoading, setIsGoogleLoading] = useState(false); const [isPasswordLoading, setIsPasswordLoading] = useState(false); const [email, setEmail] = useState(""); @@ -30,13 +30,30 @@ export function LoginForm({ const { loginWithOAuth } = useAuth(); const router = useRouter(); + const isAllowedExternalRedirect = (value: string) => { + // Allow localhost with any port for development + return /^http:\/\/localhost:\d+/.test(value); + }; + + const getRedirectUrl = () => { + const url = new URL(window.location.href); + const redirectParam = url.searchParams.get(URL_PARAM.REDIRECT); + if (redirectParam && (redirectParam.startsWith("/") || isAllowedExternalRedirect(redirectParam))) { + return redirectParam; + } + if (url.pathname === "/login") { + return "/"; + } + return url.pathname + url.search; + }; + const handleGoogleLogin = async () => { try { setIsGoogleLoading(true); setError(null); // Store the current URL to redirect back after OAuth - const currentUrl = window.location.pathname + window.location.search; + const currentUrl = getRedirectUrl(); // Build Keycloak authorization URL with Google IdP const authUrl = await buildGoogleAuthUrl(currentUrl); @@ -89,7 +106,7 @@ export function LoginForm({ const tokens: OAuthTokenResponse = await response.json(); loginWithOAuth(tokens); - onSuccess?.(); + onSuccess?.(getRedirectUrl()); } catch (error) { console.error("Password login error:", error); setError( diff --git a/apps/web/src/constants/ui.ts b/apps/web/src/constants/ui.ts index c4c41c23..66e4427c 100644 --- a/apps/web/src/constants/ui.ts +++ b/apps/web/src/constants/ui.ts @@ -22,6 +22,7 @@ export const URL_PARAM = { SETTING: "setting", PROJECTS: "projects", SEARCH: "search", + REDIRECT: "redirect", } as const; // URL param values diff --git a/apps/web/src/routes/__root.tsx b/apps/web/src/routes/__root.tsx index ad55c5a1..1de09423 100644 --- a/apps/web/src/routes/__root.tsx +++ b/apps/web/src/routes/__root.tsx @@ -32,6 +32,7 @@ function RootLayout() { const guestLogin = useAuth((state) => state.guestLogin); const hasAttemptedGuestLogin = useRef(false); const isAuthenticated = useAuth((state) => state.isAuthenticated); + const isGuest = useAuth((state) => state.isGuest); const clearRightSidebar = useRightSidebarStore( (state) => state.clearSelection ); @@ -56,21 +57,76 @@ function RootLayout() { setSidebarOpen(false); }, [location.pathname]); - // Auto guest login if no token exists + const isProtectedPath = (path: string) => { + if (path === "/login" || path === "/") { + return false; + } + + const publicPrefixes = ["/auth", "/share", "/docs"]; + if (publicPrefixes.some((prefix) => path.startsWith(prefix))) { + return false; + } + + const protectedPrefixes = [ + "/dashboard", + "/admin", + "/profile", + "/projects", + "/connectors", + "/artifacts", + ]; + + return protectedPrefixes.some((prefix) => path.startsWith(prefix)); + }; + + const getRedirectTarget = () => { + const url = new URL(window.location.href); + const redirectParam = url.searchParams.get(URL_PARAM.REDIRECT); + if (redirectParam && redirectParam.startsWith("/")) { + return redirectParam; + } + return url.pathname + url.search; + }; + + const redirectToLogin = () => { + const redirectUrl = getRedirectTarget(); + const params = new URLSearchParams(); + params.set(URL_PARAM.REDIRECT, redirectUrl); + router.navigate({ to: `/login?${params.toString()}` }); + }; + + // Auto guest login if no token exists (skip for protected pages) useEffect(() => { - if (!accessToken && !hasAttemptedGuestLogin.current) { + if ( + !accessToken && + !hasAttemptedGuestLogin.current && + !isProtectedPath(location.pathname) && + location.pathname !== "/login" + ) { hasAttemptedGuestLogin.current = true; guestLogin().catch((error) => { console.error("Auto guest login failed:", error); hasAttemptedGuestLogin.current = false; // Allow retry on failure }); } - }, [accessToken, guestLogin]); + }, [accessToken, guestLogin, location.pathname]); + + // Redirect unauthenticated/guest users to login for protected pages + useEffect(() => { + if (!isProtectedPath(location.pathname)) { + return; + } + + if (!isAuthenticated || isGuest) { + redirectToLogin(); + } + }, [isAuthenticated, isGuest, location.pathname]); // Check if modals should be shown via search params const searchParams = new URLSearchParams(location.search); + const isLoginRoute = location.pathname === "/login"; const isLoginModal = - searchParams.get(URL_PARAM.MODAL) === URL_PARAM_VALUE.LOGIN; + isLoginRoute || searchParams.get(URL_PARAM.MODAL) === URL_PARAM_VALUE.LOGIN; const isRegisterModal = searchParams.get(URL_PARAM.MODAL) === URL_PARAM_VALUE.REGISTER; const settingSection = searchParams.get(URL_PARAM.SETTING); @@ -80,10 +136,36 @@ function RootLayout() { const searchSection = searchParams.get(URL_PARAM.SEARCH); const isSearchOpen = !!searchSection; - const handleCloseModal = () => { + const isAllowedExternalRedirect = (value: string) => { + // Allow localhost with any port for development + return /^http:\/\/localhost:\d+/.test(value); + }; + + const handleCloseModal = (redirectUrl?: string) => { + // Handle external redirect (e.g., http://localhost:19999) + if (redirectUrl && isAllowedExternalRedirect(redirectUrl)) { + const target = new URL(redirectUrl); + target.searchParams.set("token", `Bearer ${accessToken}`); + window.location.href = target.toString(); + return; + } + + if (location.pathname === "/login") { + if (redirectUrl && redirectUrl.startsWith("/")) { + router.navigate({ to: redirectUrl }); + } + return; + } + + if (redirectUrl && redirectUrl.startsWith("/")) { + router.navigate({ to: redirectUrl }); + return; + } + // Remove the modal search param by navigating without it const url = new URL(window.location.href); url.searchParams.delete(URL_PARAM.MODAL); + url.searchParams.delete(URL_PARAM.REDIRECT); router.navigate({ to: url.pathname + url.search }); }; diff --git a/apps/web/src/routes/auth/callback.tsx b/apps/web/src/routes/auth/callback.tsx index 88fccfb0..06efa9ec 100644 --- a/apps/web/src/routes/auth/callback.tsx +++ b/apps/web/src/routes/auth/callback.tsx @@ -55,7 +55,22 @@ function OAuthCallbackPage() { loginWithOAuth(tokens); console.log("OAuth login successful"); // Navigate to the original URL or home - const redirectUrl = "/"; + const isAllowedExternalRedirect = (value: string) => + /^http:\/\/localhost:\d+/.test(value); + + if (oauthData.redirectUrl && isAllowedExternalRedirect(oauthData.redirectUrl)) { + const bearerToken = `Bearer ${tokens.access_token}`; + const encodedToken = btoa(bearerToken); + const target = new URL(oauthData.redirectUrl); + target.searchParams.set("token", encodedToken); + window.location.href = target.toString(); + return; + } + + const redirectUrl = + oauthData.redirectUrl && oauthData.redirectUrl.startsWith("/") + ? oauthData.redirectUrl + : "/"; navigate({ to: redirectUrl }); } catch (err) { diff --git a/apps/web/src/routes/login.tsx b/apps/web/src/routes/login.tsx new file mode 100644 index 00000000..cf0ad203 --- /dev/null +++ b/apps/web/src/routes/login.tsx @@ -0,0 +1,52 @@ +import { createFileRoute, useNavigate } from "@tanstack/react-router"; +import { useEffect } from "react"; +import { useAuth } from "@/stores/auth-store"; +import { URL_PARAM } from "@/constants"; + +export const Route = createFileRoute("/login" as "/")({ + component: LoginRoute, +}); + +function LoginRoute() { + const navigate = useNavigate(); + const isAuthenticated = useAuth((state) => state.isAuthenticated); + const isGuest = useAuth((state) => state.isGuest); + const accessToken = useAuth((state) => state.accessToken); + + const isAllowedExternalRedirect = (value: string) => { + // Allow localhost with any port for development + return /^http:\/\/localhost:\d+/.test(value); + }; + + useEffect(() => { + if (!isAuthenticated || isGuest) { + return; + } + + const url = new URL(window.location.href); + const redirectUrl = url.searchParams.get(URL_PARAM.REDIRECT); + + // Handle external redirect (e.g., http://localhost:19999) + if (redirectUrl && isAllowedExternalRedirect(redirectUrl)) { + if (!accessToken) { + return; + } + const bearerToken = `Bearer ${accessToken}`; + const encodedToken = btoa(bearerToken); + const target = new URL(redirectUrl); + target.searchParams.set("token", encodedToken); + window.location.href = target.toString(); + return; + } + + // Handle internal redirect (e.g., /dashboard) + if (redirectUrl && redirectUrl.startsWith("/")) { + navigate({ to: redirectUrl }); + return; + } + + navigate({ to: "/" }); + }, [accessToken, isAuthenticated, isGuest, navigate]); + + return null; +} diff --git a/services/llm-api/cmd/server/wire_gen.go b/services/llm-api/cmd/server/wire_gen.go index aa4d4d7c..d0d6dadc 100644 --- a/services/llm-api/cmd/server/wire_gen.go +++ b/services/llm-api/cmd/server/wire_gen.go @@ -168,7 +168,7 @@ func CreateApplication() (*Application, error) { shareHandler := sharehandler.NewShareHandler(shareService, conversationHandler, config) shareRoute := share2.NewShareRoute(shareHandler, authHandler, conversationHandler) publicShareRoute := public.NewPublicShareRoute(shareHandler) - messagesHandler := messageshandler.NewMessagesHandler(inferenceProvider, providerHandler, conversationService) + messagesHandler := messageshandler.NewMessagesHandler(inferenceProvider, providerHandler, conversationService, tokenusageService) messagesRoute := messages.NewMessagesRoute(messagesHandler, authHandler) usageRoute := usage.NewUsageRoute(usageHandler, authHandler) // Document OCR service and routes diff --git a/services/llm-api/internal/interfaces/httpserver/handlers/messageshandler/messages_handler.go b/services/llm-api/internal/interfaces/httpserver/handlers/messageshandler/messages_handler.go index ab2c840a..7425ac50 100644 --- a/services/llm-api/internal/interfaces/httpserver/handlers/messageshandler/messages_handler.go +++ b/services/llm-api/internal/interfaces/httpserver/handlers/messageshandler/messages_handler.go @@ -15,7 +15,9 @@ import ( "go.opentelemetry.io/otel/codes" "jan-server/services/llm-api/internal/domain/conversation" + "jan-server/services/llm-api/internal/domain/tokenusage" "jan-server/services/llm-api/internal/infrastructure/inference" + "jan-server/services/llm-api/internal/infrastructure/metrics" "jan-server/services/llm-api/internal/infrastructure/observability" "jan-server/services/llm-api/internal/interfaces/httpserver/handlers/modelhandler" messagesrequests "jan-server/services/llm-api/internal/interfaces/httpserver/requests/messages" @@ -35,6 +37,7 @@ type MessagesHandler struct { inferenceProvider *inference.InferenceProvider providerHandler *modelhandler.ProviderHandler conversationService *conversation.ConversationService + tokenUsageService *tokenusage.Service } // NewMessagesHandler creates a new messages handler @@ -42,11 +45,13 @@ func NewMessagesHandler( inferenceProvider *inference.InferenceProvider, providerHandler *modelhandler.ProviderHandler, conversationService *conversation.ConversationService, + tokenUsageService *tokenusage.Service, ) *MessagesHandler { return &MessagesHandler{ inferenceProvider: inferenceProvider, providerHandler: providerHandler, conversationService: conversationService, + tokenUsageService: tokenUsageService, } } @@ -121,9 +126,9 @@ func (h *MessagesHandler) CreateMessage(ctx context.Context, reqCtx *gin.Context llmStartTime := time.Now() if request.Stream { - err = h.streamCompletion(ctx, reqCtx, chatClient, llmRequest, request.Model, conv, conversationID) + err = h.streamCompletion(ctx, reqCtx, chatClient, llmRequest, request.Model, selectedProvider.DisplayName, conv, conversationID, userID) } else { - err = h.callCompletion(ctx, reqCtx, chatClient, llmRequest, request.Model, conv, conversationID, userID, request) + err = h.callCompletion(ctx, reqCtx, chatClient, llmRequest, request.Model, selectedProvider.DisplayName, conv, conversationID, userID, request) } llmDuration := time.Since(llmStartTime) @@ -150,6 +155,7 @@ func (h *MessagesHandler) callCompletion( chatClient *chat.ChatCompletionClient, request chat.CompletionRequest, originalModel string, + providerName string, conv *conversation.Conversation, conversationID string, userID uint, @@ -184,6 +190,29 @@ func (h *MessagesHandler) callCompletion( h.storeToConversation(ctx, conv, anthropicReq, response) } + // Record token usage for dashboard + if response != nil && response.Usage.TotalTokens > 0 && h.tokenUsageService != nil { + metrics.RecordTokens(originalModel, providerName, response.Usage.PromptTokens, response.Usage.CompletionTokens) + usage := &tokenusage.TokenUsage{ + UserID: fmt.Sprintf("%d", userID), + Model: originalModel, + Provider: providerName, + PromptTokens: response.Usage.PromptTokens, + CompletionTokens: response.Usage.CompletionTokens, + TotalTokens: response.Usage.TotalTokens, + Stream: false, + } + if conversationID != "" { + usage.ConversationID = &conversationID + } + if conv != nil && conv.ProjectPublicID != nil { + usage.ProjectID = conv.ProjectPublicID + } + go func(ctx context.Context, u *tokenusage.TokenUsage) { + _ = h.tokenUsageService.RecordUsage(ctx, u) + }(context.Background(), usage) + } + reqCtx.JSON(http.StatusOK, anthropicResponse) return nil } @@ -195,8 +224,10 @@ func (h *MessagesHandler) streamCompletion( chatClient *chat.ChatCompletionClient, request chat.CompletionRequest, originalModel string, + providerName string, conv *conversation.Conversation, conversationID string, + userID uint, ) error { // Ensure streaming is enabled request.Stream = true @@ -262,6 +293,32 @@ func (h *MessagesHandler) streamCompletion( return platformerrors.AsError(ctx, platformerrors.LayerHandler, err, "streaming error") } + // Record token usage for dashboard (streaming) + if h.tokenUsageService != nil && (state.InputTokens > 0 || state.OutputTokens > 0) { + totalTokens := state.InputTokens + state.OutputTokens + if totalTokens > 0 { + metrics.RecordTokens(originalModel, providerName, state.InputTokens, state.OutputTokens) + usage := &tokenusage.TokenUsage{ + UserID: fmt.Sprintf("%d", userID), + Model: originalModel, + Provider: providerName, + PromptTokens: state.InputTokens, + CompletionTokens: state.OutputTokens, + TotalTokens: totalTokens, + Stream: true, + } + if conversationID != "" { + usage.ConversationID = &conversationID + } + if conv != nil && conv.ProjectPublicID != nil { + usage.ProjectID = conv.ProjectPublicID + } + go func(ctx context.Context, u *tokenusage.TokenUsage) { + _ = h.tokenUsageService.RecordUsage(ctx, u) + }(context.Background(), usage) + } + } + return nil } From 9632f163133a639aa6542403752e9ed46fe8f845 Mon Sep 17 00:00:00 2001 From: locnguyen1986 Date: Fri, 30 Jan 2026 14:51:39 +0700 Subject: [PATCH 2/2] allow API Key access inactive models --- .../handlers/chathandler/chat_handler.go | 10 ++++++++-- .../messageshandler/messages_handler.go | 20 +++++++++++++++++-- 2 files changed, 26 insertions(+), 4 deletions(-) diff --git a/services/llm-api/internal/interfaces/httpserver/handlers/chathandler/chat_handler.go b/services/llm-api/internal/interfaces/httpserver/handlers/chathandler/chat_handler.go index 06db1fe2..191b4026 100644 --- a/services/llm-api/internal/interfaces/httpserver/handlers/chathandler/chat_handler.go +++ b/services/llm-api/internal/interfaces/httpserver/handlers/chathandler/chat_handler.go @@ -192,7 +192,14 @@ func (h *ChatHandler) CreateChatCompletion( // Get provider based on the requested model observability.AddSpanEvent(ctx, "selecting_provider") - selectedProviderModel, selectedProvider, err := h.providerHandler.SelectProviderModelForModelPublicID(ctx, request.Model) + isAPIKeyAuth := strings.EqualFold(reqCtx.GetHeader("X-Auth-Method"), "apikey") + var selectedProviderModel *domainmodel.ProviderModel + var selectedProvider *domainmodel.Provider + if isAPIKeyAuth { + selectedProviderModel, selectedProvider, err = h.providerHandler.SelectProviderModelForModelPublicIDIncludingInactive(ctx, request.Model) + } else { + selectedProviderModel, selectedProvider, err = h.providerHandler.SelectProviderModelForModelPublicID(ctx, request.Model) + } if err != nil { observability.RecordError(ctx, err) return nil, platformerrors.AsError(ctx, platformerrors.LayerHandler, err, "failed to select provider model") @@ -220,7 +227,6 @@ func (h *ChatHandler) CreateChatCompletion( // Check if we should use the instruct model instead // This happens when enable_thinking is explicitly false and the model has an instruct model configured // Skip instruct fallback for API key authentication (API users should get the model they requested) - isAPIKeyAuth := strings.EqualFold(reqCtx.GetHeader("X-Auth-Method"), "apikey") if !isAPIKeyAuth && request.EnableThinking != nil && !*request.EnableThinking && selectedProviderModel.InstructModelID != nil && !imageRequested { instructModel, instructProvider, err := h.providerHandler.GetProviderModelByID(ctx, *selectedProviderModel.InstructModelID) if err == nil && instructModel != nil && instructProvider != nil { diff --git a/services/llm-api/internal/interfaces/httpserver/handlers/messageshandler/messages_handler.go b/services/llm-api/internal/interfaces/httpserver/handlers/messageshandler/messages_handler.go index 7425ac50..3c8de7c1 100644 --- a/services/llm-api/internal/interfaces/httpserver/handlers/messageshandler/messages_handler.go +++ b/services/llm-api/internal/interfaces/httpserver/handlers/messageshandler/messages_handler.go @@ -15,6 +15,7 @@ import ( "go.opentelemetry.io/otel/codes" "jan-server/services/llm-api/internal/domain/conversation" + domainmodel "jan-server/services/llm-api/internal/domain/model" "jan-server/services/llm-api/internal/domain/tokenusage" "jan-server/services/llm-api/internal/infrastructure/inference" "jan-server/services/llm-api/internal/infrastructure/metrics" @@ -74,7 +75,15 @@ func (h *MessagesHandler) CreateMessage(ctx context.Context, reqCtx *gin.Context // Get provider for the requested model observability.AddSpanEvent(ctx, "selecting_provider") - selectedProviderModel, selectedProvider, err := h.providerHandler.SelectProviderModelForModelPublicID(ctx, request.Model) + isAPIKeyAuth := strings.EqualFold(reqCtx.GetHeader("X-Auth-Method"), "apikey") + var selectedProviderModel *domainmodel.ProviderModel + var selectedProvider *domainmodel.Provider + var err error + if isAPIKeyAuth { + selectedProviderModel, selectedProvider, err = h.providerHandler.SelectProviderModelForModelPublicIDIncludingInactive(ctx, request.Model) + } else { + selectedProviderModel, selectedProvider, err = h.providerHandler.SelectProviderModelForModelPublicID(ctx, request.Model) + } if err != nil { observability.RecordError(ctx, err) return h.writeErrorResponse(reqCtx, http.StatusBadRequest, "invalid_request_error", fmt.Sprintf("Model not found: %s", request.Model)) @@ -450,7 +459,14 @@ func (h *MessagesHandler) CountTokens(ctx context.Context, reqCtx *gin.Context, ) // Get provider for the requested model to validate it exists - selectedProviderModel, _, err := h.providerHandler.SelectProviderModelForModelPublicID(ctx, request.Model) + isAPIKeyAuth := strings.EqualFold(reqCtx.GetHeader("X-Auth-Method"), "apikey") + var selectedProviderModel *domainmodel.ProviderModel + var err error + if isAPIKeyAuth { + selectedProviderModel, _, err = h.providerHandler.SelectProviderModelForModelPublicIDIncludingInactive(ctx, request.Model) + } else { + selectedProviderModel, _, err = h.providerHandler.SelectProviderModelForModelPublicID(ctx, request.Model) + } if err != nil || selectedProviderModel == nil { return h.writeErrorResponse(reqCtx, http.StatusBadRequest, "invalid_request_error", fmt.Sprintf("Model not found: %s", request.Model)) }