oidc: remove actor from final extension (#20475)oidc-1.2.5-devoidc-1.2.5

16 files changed, 190 insertions, 176 deletions

diff --git a/ydb/mvp/oidc_proxy/extension.cpp b/ydb/mvp/oidc_proxy/extension.cpp
index a5b572c6e75..8a950cb2b87 100644
--- a/ydb/mvp/oidc_proxy/extension.cpp
+++ b/ydb/mvp/oidc_proxy/extension.cpp
@@ -2,34 +2,36 @@
 
 namespace NMVP::NOIDC {
 
-void TExtension::Bootstrap() {
-    Become(&TExtension::StateWork);
+void TExtensionContext::Reply(NHttp::THttpOutgoingResponsePtr httpResponse) {
+    NActors::TActivationContext::Send(Sender, std::make_unique<NHttp::TEvHttpProxy::TEvHttpOutgoingResponse>(std::move(httpResponse)));
 }
 
-void TExtension::ReplyAndPassAway(NHttp::THttpOutgoingResponsePtr httpResponse) {
-    Send(Context->Sender, new NHttp::TEvHttpProxy::TEvHttpOutgoingResponse(std::move(httpResponse)));
-    PassAway();
-}
-
-void TExtension::ReplyAndPassAway() {
-    auto& params = Context->Params;
-    if (params->StatusOverride) {
-        return ReplyAndPassAway(params->Request->CreateResponse(params->StatusOverride, params->MessageOverride, *params->HeadersOverride, params->BodyOverride));
+void TExtensionContext::Reply() {
+    if (Params->StatusOverride) {
+        return Reply(Params->Request->CreateResponse(Params->StatusOverride, Params->MessageOverride, *Params->HeadersOverride, Params->BodyOverride));
     } else {
         static constexpr size_t MAX_LOGGED_SIZE = 1024;
-        BLOG_D("Can not process request to protected resource:\n" << Context->Params->Request->GetObfuscatedData().substr(0, MAX_LOGGED_SIZE));
-        return ReplyAndPassAway(CreateResponseForNotExistingResponseFromProtectedResource(Context->Params->Request, Context->Params->ResponseError));
+        BLOG_D("Can not process request to protected resource:\n" << Params->Request->GetObfuscatedData().substr(0, MAX_LOGGED_SIZE));
+        return Reply(CreateResponseForNotExistingResponseFromProtectedResource(Params->Request, Params->ResponseError));
     }
 }
 
-void TExtension::ContinueAndPassAway() {
-    if (!Context->Route.empty()) {
-        const auto route = Context->Route.Next();
-        Send(route, new TEvPrivate::TEvExtensionRequest(std::move(Context)));
-        PassAway();
+void TExtensionContext::Continue() {
+    const auto step = Steps.Next();
+    if (step) {
+        step->Execute(this);
     } else {
-        ReplyAndPassAway();
+        Reply();
+    }
+}
+
+std::unique_ptr<IExtension> TExtensionsSteps::Next() {
+    if (empty()) {
+        return nullptr;
     }
+    auto target = std::move(front());
+    pop();
+    return target;
 }
 
 } // NMVP::NOIDC
diff --git a/ydb/mvp/oidc_proxy/extension.h b/ydb/mvp/oidc_proxy/extension.h
index 5a551fea106..ad5f6c63fcb 100644
--- a/ydb/mvp/oidc_proxy/extension.h
+++ b/ydb/mvp/oidc_proxy/extension.h
@@ -1,30 +1,48 @@
 #pragma once
 
+#include "cracked_page.h"
 #include "openid_connect.h"
 
+#include <ydb/library/actors/core/actorid.h>
+#include <ydb/library/actors/http/http_proxy.h>
+
+#include <util/generic/strbuf.h>
+#include <util/generic/string.h>
+#include <util/generic/queue.h>
+
 namespace NMVP::NOIDC {
 
-class TExtension : public NActors::TActorBootstrapped<TExtension> {
-protected:
-    const TOpenIdConnectSettings Settings;
-    TIntrusivePtr<TExtensionContext> Context;
+class IExtension;
 
-public:
-    TExtension(const TOpenIdConnectSettings& settings)
-        : Settings(settings)
-    {}
-    virtual void Bootstrap();
-    virtual void Handle(TEvPrivate::TEvExtensionRequest::TPtr event) = 0;
-    void ReplyAndPassAway(NHttp::THttpOutgoingResponsePtr httpResponse);
-    void ReplyAndPassAway();
-    void ContinueAndPassAway();
-
-    STFUNC(StateWork) {
-        switch (ev->GetTypeRewrite()) {
-            hFunc(TEvPrivate::TEvExtensionRequest, Handle);
-        }
-    }
+struct TProxiedResponseParams {
+    NHttp::THttpIncomingRequestPtr Request;
+    THolder<TCrackedPage> ProtectedPage;
+    TString ResponseError;
 
+    TString StatusOverride;
+    TString MessageOverride;
+    TString BodyOverride;
+    THolder<NHttp::THeadersBuilder> HeadersOverride;
+};
+
+struct TExtensionsSteps : public TQueue<std::unique_ptr<IExtension>> {
+    std::unique_ptr<IExtension> Next();
+};
+
+struct TExtensionContext : public TThrRefBase {
+    TActorId Sender;
+    TExtensionsSteps Steps;
+    THolder<TProxiedResponseParams> Params;
+
+    void Reply(NHttp::THttpOutgoingResponsePtr httpResponse);
+    void Reply();
+    void Continue();
+};
+
+class IExtension {
+public:
+    virtual ~IExtension() = default;
+    virtual void Execute(TIntrusivePtr<TExtensionContext> ctx) = 0;
 };
 
 } // NMVP::NOIDC
diff --git a/ydb/mvp/oidc_proxy/extension_context.h b/ydb/mvp/oidc_proxy/extension_context.h
deleted file mode 100644
index 4ca520e0fb9..00000000000
--- a/ydb/mvp/oidc_proxy/extension_context.h
+++ /dev/null
@@ -1,42 +0,0 @@
-#pragma once
-
-#include "cracked_page.h"
-
-#include <ydb/library/actors/core/actorid.h>
-#include <ydb/library/actors/http/http_proxy.h>
-
-#include <util/generic/strbuf.h>
-#include <util/generic/string.h>
-#include <util/generic/queue.h>
-
-namespace NMVP::NOIDC {
-
-struct TProxiedResponseParams {
-    NHttp::THttpIncomingRequestPtr Request;
-    THolder<TCrackedPage> ProtectedPage;
-    TString ResponseError;
-
-    TString StatusOverride;
-    TString MessageOverride;
-    TString BodyOverride;
-    THolder<NHttp::THeadersBuilder> HeadersOverride;
-};
-
-struct TExtensionContext : public TThrRefBase {
-    struct TRoute: public TQueue<NActors::TActorId> {
-        TActorId Next() {
-            if (empty()) {
-                return TActorId();
-            }
-            TActorId target = front();
-            pop();
-            return target;
-        }
-    };
-
-    TActorId Sender;
-    TRoute Route;
-    THolder<TProxiedResponseParams> Params;
-};
-
-} // NMVP::NOIDC
diff --git a/ydb/mvp/oidc_proxy/extension_final.cpp b/ydb/mvp/oidc_proxy/extension_final.cpp
index af9205d2346..e38c3db50b1 100644
--- a/ydb/mvp/oidc_proxy/extension_final.cpp
+++ b/ydb/mvp/oidc_proxy/extension_final.cpp
@@ -81,14 +81,14 @@ TString TExtensionFinal::GetFixedLocationHeader(TStringBuf location) {
     return TString(location);
 }
 
-void TExtensionFinal::Handle(TEvPrivate::TEvExtensionRequest::TPtr ev) {
-    Context = std::move(ev->Get()->Context);
+void TExtensionFinal::Execute(TIntrusivePtr<TExtensionContext> ctx) {
+    Context = std::move(ctx);
     if (Context->Params->StatusOverride) {
         SetProxyResponseHeaders();
         SetProxyResponseBody();
     }
 
-    ContinueAndPassAway();
+    Context->Continue();
 }
 
 } // NMVP::NOIDC
diff --git a/ydb/mvp/oidc_proxy/extension_final.h b/ydb/mvp/oidc_proxy/extension_final.h
index 2755b72660f..dccda091111 100644
--- a/ydb/mvp/oidc_proxy/extension_final.h
+++ b/ydb/mvp/oidc_proxy/extension_final.h
@@ -4,15 +4,16 @@
 
 namespace NMVP::NOIDC {
 
-class TExtensionFinal : public TExtension {
+class TExtensionFinal : public IExtension {
 private:
-    using TBase = TExtension;
+    const TOpenIdConnectSettings Settings;
+    TIntrusivePtr<TExtensionContext> Context;
 
 public:
     TExtensionFinal(const TOpenIdConnectSettings& settings)
-        : TBase(settings)
+        : Settings(settings)
     {}
-    void Handle(TEvPrivate::TEvExtensionRequest::TPtr event) override;
+    void Execute(TIntrusivePtr<TExtensionContext> ctx) override;
 
 private:
     TString FixReferenceInHtml(TStringBuf html, TStringBuf host, TStringBuf findStr);
diff --git a/ydb/mvp/oidc_proxy/extension_manager.cpp b/ydb/mvp/oidc_proxy/extension_manager.cpp
index 08eacc814d9..bc63125bdee 100644
--- a/ydb/mvp/oidc_proxy/extension_manager.cpp
+++ b/ydb/mvp/oidc_proxy/extension_manager.cpp
@@ -16,6 +16,11 @@ TExtensionManager::TExtensionManager(const TActorId sender,
     ExtensionCtx->Params = MakeHolder<TProxiedResponseParams>();
     ExtensionCtx->Params->ProtectedPage = MakeHolder<TCrackedPage>(protectedPage);
     ExtensionCtx->Sender = sender;
+    Timeout = settings.DefaultRequestTimeout;
+}
+
+void TExtensionManager::SetExtensionTimeout(TDuration timeout) {
+    Timeout = timeout;
 }
 
 void TExtensionManager::SetRequest(NHttp::THttpIncomingRequestPtr request) {
@@ -24,9 +29,9 @@ void TExtensionManager::SetRequest(NHttp::THttpIncomingRequestPtr request) {
 
 void TExtensionManager::SetOverrideResponse(NHttp::TEvHttpProxy::TEvHttpIncomingResponse::TPtr event) {
     ExtensionCtx->Params->HeadersOverride = MakeHolder<NHttp::THeadersBuilder>();
-    ExtensionCtx->Params->ResponseError = event ? event->Get()->GetError() : "Timeout while waiting info";
+    ExtensionCtx->Params->ResponseError = event->Get()->GetError();
 
-    if (!event || !event->Get()->Response)
+    if (!event->Get()->Response)
         return;
 
     auto& response = event->Get()->Response;
@@ -40,32 +45,28 @@ void TExtensionManager::SetOverrideResponse(NHttp::TEvHttpProxy::TEvHttpIncoming
 }
 
 void TExtensionManager::AddExtensionWhoami() {
-    EnrichmentExtension = true;
-    AddExtension(NActors::TActivationContext::ActorSystem()->Register(new TExtensionWhoami(Settings, AuthHeader)));
+    auto ext = std::make_unique<TExtensionWhoami>(Settings, AuthHeader, Timeout);
+    AddExtension(std::move(ext));
 }
 
 void TExtensionManager::AddExtensionFinal() {
-    AddExtension(NActors::TActivationContext::ActorSystem()->Register(new TExtensionFinal(Settings)));
+    auto ext = std::make_unique<TExtensionFinal>(Settings);
+    AddExtension(std::move(ext));
 }
 
-void TExtensionManager::AddExtension(const NActors::TActorId& stage) {
-    ExtensionCtx->Route.push(stage);
+void TExtensionManager::AddExtension(std::unique_ptr<IExtension> ext) {
+    ExtensionCtx->Steps.push(std::move(ext));
 }
 
 bool TExtensionManager::NeedExtensionWhoami(const NHttp::THttpIncomingRequestPtr& request) const {
-    if (Settings.AccessServiceType == NMvp::yandex_v2) {
-        return false; // does not support whoami extension
-    }
-
-    if (request->Method == "OPTIONS" || Settings.WhoamiExtendedInfoEndpoint.empty()) {
+    if (!Settings.EnabledExtensionWhoami() || request->Method == "OPTIONS") {
         return false;
     }
 
     TCrackedPage page(request);
     auto path = TStringBuf(page.Url).Before('?');
 
-    static constexpr TStringBuf WHOAMI_PATHS[] = { "/viewer/json/whoami", "/viewer/whoami" };
-    for (const auto& whoamiPath : WHOAMI_PATHS) {
+    for (const auto& whoamiPath : Settings.WHOAMI_PATHS) {
         if (path.EndsWith(whoamiPath)) {
             return true;
         }
@@ -80,17 +81,13 @@ void TExtensionManager::ArrangeExtensions(const NHttp::THttpIncomingRequestPtr&
     AddExtensionFinal();
 }
 
-bool TExtensionManager::HasEnrichmentExtension() {
-    return EnrichmentExtension;
-}
-
 void TExtensionManager::StartExtensionProcess(NHttp::THttpIncomingRequestPtr request,
                                               NHttp::TEvHttpProxy::TEvHttpIncomingResponse::TPtr event) {
     SetRequest(std::move(request));
     SetOverrideResponse(std::move(event));
 
-    const auto route = ExtensionCtx->Route.Next();
-    NActors::TActivationContext::ActorSystem()->Send(route, new TEvPrivate::TEvExtensionRequest(std::move(ExtensionCtx)));
+    const auto step = ExtensionCtx->Steps.Next();
+    step->Execute(std::move(ExtensionCtx));
 }
 
 } // NMVP::NOIDC
diff --git a/ydb/mvp/oidc_proxy/extension_manager.h b/ydb/mvp/oidc_proxy/extension_manager.h
index ee4d541ee6b..7f838fe2d9d 100644
--- a/ydb/mvp/oidc_proxy/extension_manager.h
+++ b/ydb/mvp/oidc_proxy/extension_manager.h
@@ -1,6 +1,6 @@
 #pragma once
 
-#include "extension_context.h"
+#include "extension.h"
 
 namespace NMVP::NOIDC {
 
@@ -8,15 +8,15 @@ struct TExtensionManager {
     TIntrusivePtr<TExtensionContext> ExtensionCtx;
     const TOpenIdConnectSettings Settings;
     TString AuthHeader;
-    bool EnrichmentExtension = false;
+    TDuration Timeout;
 
 public:
     TExtensionManager(const TActorId sender,
                       const TOpenIdConnectSettings& settings,
                       const TCrackedPage& protectedPage,
                       const TString authHeader);
+    void SetExtensionTimeout(TDuration timeout);
     void ArrangeExtensions(const NHttp::THttpIncomingRequestPtr& request);
-    bool HasEnrichmentExtension();
     void StartExtensionProcess(NHttp::THttpIncomingRequestPtr request,
                                NHttp::TEvHttpProxy::TEvHttpIncomingResponse::TPtr event = nullptr);
 
@@ -26,7 +26,7 @@ private:
     bool NeedExtensionWhoami(const NHttp::THttpIncomingRequestPtr& request) const;
     void AddExtensionWhoami();
     void AddExtensionFinal();
-    void AddExtension(const NActors::TActorId& stage);
+    void AddExtension(std::unique_ptr<IExtension> ext);
 };
 
 } // NMVP::NOIDC
diff --git a/ydb/mvp/oidc_proxy/extension_whoami.cpp b/ydb/mvp/oidc_proxy/extension_whoami.cpp
index aa65ea215cd..400eb6ec7a6 100644
--- a/ydb/mvp/oidc_proxy/extension_whoami.cpp
+++ b/ydb/mvp/oidc_proxy/extension_whoami.cpp
@@ -5,7 +5,7 @@
 
 namespace NMVP::NOIDC {
 
-void TExtensionWhoami::Bootstrap() {
+void TExtensionWhoamiWorker::Bootstrap() {
     auto connection = CreateGRpcServiceConnection<TProfileService>(Settings.WhoamiExtendedInfoEndpoint);
 
     nebius::iam::v1::GetProfileRequest request;
@@ -22,25 +22,25 @@ void TExtensionWhoami::Bootstrap() {
 
     NYdbGrpc::TCallMeta meta;
     SetHeader(meta, "authorization", AuthHeader);
-    meta.Timeout = TDuration::MilliSeconds(Settings.EnrichmentProcessTimeoutMs);
+    meta.Timeout = Timeout;
 
     connection->DoRequest(request, std::move(responseCb), &nebius::iam::v1::ProfileService::Stub::AsyncGet, meta);
-    Become(&TExtensionWhoami::StateWork);
+    Become(&TExtensionWhoamiWorker::StateWork);
 }
 
-void TExtensionWhoami::Handle(TEvPrivate::TEvGetProfileResponse::TPtr event) {
-    BLOG_D("Whoami Extention Info: OK");
+void TExtensionWhoamiWorker::Handle(TEvPrivate::TEvGetProfileResponse::TPtr event) {
+    BLOG_D("Whoami Extension Info: OK");
     IamResponse = std::move(event);
     ApplyIfReady();
 }
 
-void TExtensionWhoami::Handle(TEvPrivate::TEvErrorResponse::TPtr event) {
-    BLOG_D("Whoami Extention Info " << event->Get()->Status << ": " << event->Get()->Message << ", " << event->Get()->Details);
+void TExtensionWhoamiWorker::Handle(TEvPrivate::TEvErrorResponse::TPtr event) {
+    BLOG_D("Whoami Extension Info " << event->Get()->Status << ": " << event->Get()->Message << ", " << event->Get()->Details);
     IamError = std::move(event);
     ApplyIfReady();
 }
 
-void TExtensionWhoami::PatchResponse(NJson::TJsonValue& json, NJson::TJsonValue& errorJson) {
+void TExtensionWhoamiWorker::PatchResponse(NJson::TJsonValue& json, NJson::TJsonValue& errorJson) {
     TString statusOverride;
     TString messageOverride;
     NJson::TJsonValue* outJson = nullptr;
@@ -79,7 +79,7 @@ void TExtensionWhoami::PatchResponse(NJson::TJsonValue& json, NJson::TJsonValue&
     params->BodyOverride = content.Str();
 }
 
-void TExtensionWhoami::Handle(TEvPrivate::TEvExtensionRequest::TPtr ev) {
+void TExtensionWhoamiWorker::Handle(TEvPrivate::TEvExtensionRequest::TPtr ev) {
     Context = std::move(ev->Get()->Context);
     if (Context->Params->StatusOverride.StartsWith("3") || Context->Params->StatusOverride == "404") {
         ContinueAndPassAway();
@@ -87,22 +87,22 @@ void TExtensionWhoami::Handle(TEvPrivate::TEvExtensionRequest::TPtr ev) {
     ApplyIfReady();
 }
 
-void TExtensionWhoami::SetExtendedError(NJson::TJsonValue& root, const TStringBuf section, const TStringBuf key, const TStringBuf value) {
+void TExtensionWhoamiWorker::SetExtendedError(NJson::TJsonValue& root, const TStringBuf section, const TStringBuf key, const TStringBuf value) {
     if (!value.empty()) {
         root[EXTENDED_ERRORS][section][key] = value;
     }
 }
 
-void TExtensionWhoami::ApplyIfReady() {
+void TExtensionWhoamiWorker::ApplyIfReady() {
     if (!Context) {
         return;
     }
-    if (IamResponse.has_value() || IamError.has_value() || Timeout) {
+    if (IamResponse.has_value() || IamError.has_value()) {
         ApplyExtension();
     }
 }
 
-void TExtensionWhoami::ApplyExtension() {
+void TExtensionWhoamiWorker::ApplyExtension() {
     NJson::TJsonValue json;
     NJson::TJsonValue errorJson;
     NHttp::THttpIncomingResponsePtr response;
@@ -117,6 +117,9 @@ void TExtensionWhoami::ApplyExtension() {
         }
     } else {
         TString& error = params->ResponseError;
+        if (!error) {
+            error = "Can not process request to protected resource";
+        }
         BLOG_D("Incoming client error for protected resource: " << error);
         SetExtendedError(errorJson, "Ydb", "ClientError", error);
     }
@@ -158,4 +161,18 @@ void TExtensionWhoami::ApplyExtension() {
     ContinueAndPassAway();
 }
 
+void TExtensionWhoamiWorker::ContinueAndPassAway() {
+    Context->Continue();
+    PassAway();
+}
+
+TExtensionWhoami::TExtensionWhoami(const TOpenIdConnectSettings& settings, const TString& authHeader, const TDuration timeout)
+{
+    WhoamiHandlerId = NActors::TActivationContext::ActorSystem()->Register(new TExtensionWhoamiWorker(settings, authHeader, timeout));
+}
+
+void TExtensionWhoami::Execute(TIntrusivePtr<TExtensionContext> ctx) {
+    NActors::TActivationContext::ActorSystem()->Send(WhoamiHandlerId, new TEvPrivate::TEvExtensionRequest(std::move(ctx)));
+}
+
 } // NMVP::NOIDC
diff --git a/ydb/mvp/oidc_proxy/extension_whoami.h b/ydb/mvp/oidc_proxy/extension_whoami.h
index 3544fc28382..dd49964bc2a 100644
--- a/ydb/mvp/oidc_proxy/extension_whoami.h
+++ b/ydb/mvp/oidc_proxy/extension_whoami.h
@@ -6,25 +6,27 @@
 
 namespace NMVP::NOIDC {
 
-class TExtensionWhoami : public TExtension {
-private:
-    using TBase = TExtension;
+class TExtensionWhoamiWorker : public NActors::TActorBootstrapped<TExtensionWhoamiWorker> {
+    using TBase = IExtension;
     using TProfileService = nebius::iam::v1::ProfileService;
 
-protected:
     const TString AuthHeader;
-    bool Timeout = false;
+
+    const TOpenIdConnectSettings Settings;
+    TIntrusivePtr<TExtensionContext> Context;
 
     std::optional<TEvPrivate::TEvGetProfileResponse::TPtr> IamResponse;
     std::optional<TEvPrivate::TEvErrorResponse::TPtr> IamError;
+    TDuration Timeout;
 
 public:
-    TExtensionWhoami(const TOpenIdConnectSettings& settings, const TString& authHeader)
-        : TBase(settings)
-        , AuthHeader(authHeader)
+    TExtensionWhoamiWorker(const TOpenIdConnectSettings& settings, const TString& authHeader, const TDuration timeout)
+        : AuthHeader(authHeader)
+        , Settings(settings)
+        , Timeout(timeout)
     {}
-    void Bootstrap() override;
-    void Handle(TEvPrivate::TEvExtensionRequest::TPtr event) override;
+    void Bootstrap();
+    void Handle(TEvPrivate::TEvExtensionRequest::TPtr event);
     void Handle(TEvPrivate::TEvGetProfileResponse::TPtr event);
     void Handle(TEvPrivate::TEvErrorResponse::TPtr event);
 
@@ -32,9 +34,7 @@ public:
         switch (ev->GetTypeRewrite()) {
             hFunc(TEvPrivate::TEvGetProfileResponse, Handle);
             hFunc(TEvPrivate::TEvErrorResponse, Handle);
-            default:
-                TBase::StateWork(ev);
-                break;
+            hFunc(TEvPrivate::TEvExtensionRequest, Handle);
         }
     }
 
@@ -43,6 +43,16 @@ private:
     void ApplyIfReady();
     void ApplyExtension();
     void SetExtendedError(NJson::TJsonValue& root, const TStringBuf section, const TStringBuf key, const TStringBuf value);
+    void ContinueAndPassAway();
+};
+
+class TExtensionWhoami : public IExtension {
+private:
+    TActorId WhoamiHandlerId;
+
+public:
+    TExtensionWhoami(const TOpenIdConnectSettings& settings, const TString& authHeader, const TDuration timeout);
+    void Execute(TIntrusivePtr<TExtensionContext> ctx) override;
 };
 
 } // NMVP::NOIDC
diff --git a/ydb/mvp/oidc_proxy/mvp.cpp b/ydb/mvp/oidc_proxy/mvp.cpp
index f94bf2fc726..10fbe455790 100644
--- a/ydb/mvp/oidc_proxy/mvp.cpp
+++ b/ydb/mvp/oidc_proxy/mvp.cpp
@@ -294,6 +294,7 @@ void TMVP::TryGetGenericOptionsFromConfig(
             ythrow yexception() << "Unknown access_service_type value: " << accessServiceTypeStr;
         }
     }
+    OpenIdConnectSettings.InitRequestTimeoutsByPath();
 }
 
 THolder<NActors::TActorSystemSetup> TMVP::BuildActorSystemSetup(int argc, char** argv) {
diff --git a/ydb/mvp/oidc_proxy/oidc_protected_page.cpp b/ydb/mvp/oidc_proxy/oidc_protected_page.cpp
index 5a634842198..7e0956a2e0f 100644
--- a/ydb/mvp/oidc_proxy/oidc_protected_page.cpp
+++ b/ydb/mvp/oidc_proxy/oidc_protected_page.cpp
@@ -45,11 +45,6 @@ void THandlerSessionServiceCheck::HandleProxy(NHttp::TEvHttpProxy::TEvHttpIncomi
     PassAway();
 }
 
-void THandlerSessionServiceCheck::HandleEnrichmentTimeout() {
-    ExtensionManager->StartExtensionProcess(std::move(Request));
-    PassAway();
-}
-
 void THandlerSessionServiceCheck::HandleIncompleteProxy(NHttp::TEvHttpProxy::TEvHttpIncompleteIncomingResponse::TPtr event) {
     if (event->Get()->Response != nullptr) {
         NHttp::THttpIncomingResponsePtr response = std::move(event->Get()->Response);
@@ -110,14 +105,29 @@ bool THandlerSessionServiceCheck::IsAuthorizedRequest(TStringBuf authHeader) {
     return to_lower(ToString(authHeader)).StartsWith(IAM_TOKEN_SCHEME_LOWER);
 }
 
+TDuration THandlerSessionServiceCheck::GetRequestTimeout() const {
+    auto path = TStringBuf(ProtectedPage.Url).Before('?');
+    for (const auto& [suffix, timeout] : Settings.RequestTimeoutsByPath) {
+        if (path.EndsWith(suffix)) {
+            return timeout;
+        }
+    }
+    return Settings.DefaultRequestTimeout;
+}
+
 void THandlerSessionServiceCheck::ForwardUserRequest(TStringBuf authHeader, bool secure) {
     BLOG_D("Forward user request bypass OIDC");
 
+    auto timeout = GetRequestTimeout();
+    ExtensionManager = MakeHolder<TExtensionManager>(Sender, Settings, ProtectedPage, TString(authHeader));
+    ExtensionManager->SetExtensionTimeout(timeout);
+    ExtensionManager->ArrangeExtensions(Request);
+
     TProxiedRequestParams params(Request, authHeader, secure, ProtectedPage, Settings);
     auto httpRequest = CreateProxiedRequest(params);
 
     auto requestEvent = std::make_unique<NHttp::TEvHttpProxy::TEvHttpOutgoingRequest>(httpRequest);
-    requestEvent->Timeout = TDuration::Seconds(120);
+    requestEvent->Timeout = timeout;
     requestEvent->AllowConnectionReuse = !Request->IsConnectionClose();
     requestEvent->StreamContentTypes = {
         "multipart/x-mixed-replace",
@@ -126,12 +136,6 @@ void THandlerSessionServiceCheck::ForwardUserRequest(TStringBuf authHeader, bool
     };
 
     Send(HttpProxyId, requestEvent.release());
-
-    ExtensionManager = MakeHolder<TExtensionManager>(Sender, Settings, ProtectedPage, TString(authHeader));
-    ExtensionManager->ArrangeExtensions(Request);
-    if (ExtensionManager->HasEnrichmentExtension()) {
-        Schedule(TDuration::MilliSeconds(Settings.EnrichmentProcessTimeoutMs), new NActors::TEvents::TEvWakeup());
-    }
 }
 
 void THandlerSessionServiceCheck::SendSecureHttpRequest(const NHttp::THttpIncomingResponsePtr& response) {
diff --git a/ydb/mvp/oidc_proxy/oidc_protected_page.h b/ydb/mvp/oidc_proxy/oidc_protected_page.h
index 4e9d7928d36..7e387fe4e33 100644
--- a/ydb/mvp/oidc_proxy/oidc_protected_page.h
+++ b/ydb/mvp/oidc_proxy/oidc_protected_page.h
@@ -37,7 +37,6 @@ public:
     void HandleDataChunk(NHttp::TEvHttpProxy::TEvHttpIncomingDataChunk::TPtr event);
     void HandleUndelivered(NActors::TEvents::TEvUndelivered::TPtr event);
     void HandleCancelled();
-    void HandleEnrichmentTimeout();
 
     STFUNC(StateWork) {
         switch (ev->GetTypeRewrite()) {
@@ -46,12 +45,12 @@ public:
             hFunc(NHttp::TEvHttpProxy::TEvHttpIncomingDataChunk, HandleDataChunk);
             cFunc(NHttp::TEvHttpProxy::EvRequestCancelled, HandleCancelled);
             hFunc(NActors::TEvents::TEvUndelivered, HandleUndelivered);
-            cFunc(NActors::TEvents::TEvWakeup::EventType, HandleEnrichmentTimeout);
         }
     }
 
 protected:
     virtual void StartOidcProcess(const NActors::TActorContext& ctx) = 0;
+    TDuration GetRequestTimeout() const;
     virtual void ForwardUserRequest(TStringBuf authHeader, bool secure = false);
     virtual bool NeedSendSecureHttpRequest(const NHttp::THttpIncomingResponsePtr& response) const = 0;
     void ReplyAndPassAway(NHttp::THttpOutgoingResponsePtr httpResponse);
diff --git a/ydb/mvp/oidc_proxy/oidc_proxy_ut.cpp b/ydb/mvp/oidc_proxy/oidc_proxy_ut.cpp
index 1034f29a30c..f775f8d4f97 100644
--- a/ydb/mvp/oidc_proxy/oidc_proxy_ut.cpp
+++ b/ydb/mvp/oidc_proxy/oidc_proxy_ut.cpp
@@ -1515,7 +1515,6 @@ Y_UNIT_TEST_SUITE(Mvp) {
         TString WhoamiResponse;
         TString ExpectedStatus;
         NMvp::EAccessServiceType AccessServiceType = NMvp::nebius_v1;
-        bool YdbTimeout = false;
 
         TWhoamiContext(TStringBuf token, TString whoamiResponse, TStringBuf expectedStatus)
             : Token(token)
@@ -1609,14 +1608,10 @@ Y_UNIT_TEST_SUITE(Mvp) {
         UNIT_ASSERT_STRINGS_EQUAL(outgoingRequestEv->Request->Host, allowedProxyHost);
         UNIT_ASSERT_STRINGS_EQUAL(outgoingRequestEv->Request->URL, "/viewer/whoami");
 
-        if (context.YdbTimeout) {
-            runtime.Send(new IEventHandle(handle->Sender, edge, new NActors::TEvents::TEvWakeup()));
-        } else {
-            // viewer returns response to oidc
-            NHttp::THttpIncomingResponsePtr incomingResponse = new NHttp::THttpIncomingResponse(outgoingRequestEv->Request);
-            EatWholeString(incomingResponse, context.WhoamiResponse);
-            runtime.Send(new IEventHandle(handle->Sender, edge, new NHttp::TEvHttpProxy::TEvHttpIncomingResponse(outgoingRequestEv->Request, incomingResponse)));
-        }
+        // viewer returns response to oidc
+        NHttp::THttpIncomingResponsePtr incomingResponse = new NHttp::THttpIncomingResponse(outgoingRequestEv->Request);
+        EatWholeString(incomingResponse, context.WhoamiResponse);
+        runtime.Send(new IEventHandle(handle->Sender, edge, new NHttp::TEvHttpProxy::TEvHttpIncomingResponse(outgoingRequestEv->Request, incomingResponse)));
 
         // oidc final response
         auto* outgoing = runtime.GrabEdgeEvent<NHttp::TEvHttpProxy::TEvHttpOutgoingResponse>(handle);
@@ -1726,19 +1721,7 @@ Y_UNIT_TEST_SUITE(Mvp) {
         UNIT_ASSERT(!json.Has(EXTENDED_ERRORS));
     }
 
-    Y_UNIT_TEST(OidcYdbTimeout200) {
-        TWhoamiContext ctx(TProfileServiceMock::VALID_USER_TOKEN, TWhoamiContext::GetViewerResponse200(), "200");
-        ctx.YdbTimeout = true;
-        auto json = OidcWhoamiExtendedInfoTest(ctx);
-        UNIT_ASSERT_VALUES_EQUAL(json[USER_SID], TProfileServiceMock::USER_ACCOUNT_ID);
-        UNIT_ASSERT_VALUES_EQUAL(json[ORIGINAL_USER_TOKEN], TProfileServiceMock::VALID_USER_TOKEN);
-        UNIT_ASSERT(json.Has(EXTENDED_INFO));
-        UNIT_ASSERT(json[EXTENDED_ERRORS].Has("Ydb"));
-        UNIT_ASSERT_VALUES_EQUAL(json[EXTENDED_ERRORS]["Ydb"]["ClientError"], "Timeout while waiting info");
-        UNIT_ASSERT(!json[EXTENDED_ERRORS].Has("Iam"));
-    }
-
-    Y_UNIT_TEST(OidcYandexIgnoresWhoamiExtention) {
+    Y_UNIT_TEST(OidcYandexIgnoresWhoamiExtension) {
         TWhoamiContext ctx(TProfileServiceMock::VALID_USER_TOKEN, TWhoamiContext::GetViewerResponse200(), "200");
         ctx.AccessServiceType = NMvp::yandex_v2;
         auto json = OidcWhoamiExtendedInfoTest(ctx);
diff --git a/ydb/mvp/oidc_proxy/oidc_settings.cpp b/ydb/mvp/oidc_proxy/oidc_settings.cpp
index 9a5d0efce1e..6a759541e79 100644
--- a/ydb/mvp/oidc_proxy/oidc_settings.cpp
+++ b/ydb/mvp/oidc_proxy/oidc_settings.cpp
@@ -33,6 +33,21 @@ const TVector<TStringBuf> TOpenIdConnectSettings::RESPONSE_HEADERS_WHITE_LIST =
     "traceresponse"
 };
 
+bool TOpenIdConnectSettings::EnabledExtensionWhoami() const {
+    return AccessServiceType == NMvp::nebius_v1 && !WhoamiExtendedInfoEndpoint.empty();
+}
+
+void TOpenIdConnectSettings::InitRequestTimeoutsByPath() {
+    RequestTimeoutsByPath.clear();
+
+    // For requests with enrichment (e.g. extended whoami), set an explicit timeout
+    // to avoid long waits in case YDB is slow or unresponsive
+    if (EnabledExtensionWhoami()) {
+        for (auto path : WHOAMI_PATHS) {
+            RequestTimeoutsByPath[path] = TDuration::Seconds(10);
+        }
+    }
+}
 
 TString TOpenIdConnectSettings::GetAuthorizationString() const {
     return "Basic " + Base64Encode(ClientId + ":" + ClientSecret);
diff --git a/ydb/mvp/oidc_proxy/oidc_settings.h b/ydb/mvp/oidc_proxy/oidc_settings.h
index aa004433579..34eb9019cf9 100644
--- a/ydb/mvp/oidc_proxy/oidc_settings.h
+++ b/ydb/mvp/oidc_proxy/oidc_settings.h
@@ -1,6 +1,8 @@
 #pragma once
 
 #include <ydb/mvp/core/protos/mvp.pb.h>
+#include <util/datetime/base.h>
+#include <util/generic/hash.h>
 #include <util/generic/string.h>
 #include <util/generic/vector.h>
 
@@ -17,10 +19,12 @@ struct TOpenIdConnectSettings {
     static const inline TString DEFAULT_EXCHANGE_URL_PATH = "/oauth2/session/exchange";
     static const inline TString DEFAULT_IMPERSONATE_URL_PATH = "/oauth2/impersonation/impersonate";
 
-    static const inline ui32 DEFAULT_ENRICHMENT_PROCESS_TIMEOUT_MS = 10000;
+    static constexpr inline TDuration DEFAULT_REQUEST_TIMEOUT = TDuration::Seconds(120);
 
     static const TVector<TStringBuf> REQUEST_HEADERS_WHITE_LIST;
     static const TVector<TStringBuf> RESPONSE_HEADERS_WHITE_LIST;
+    static constexpr TStringBuf WHOAMI_PATHS[] =  { "/viewer/json/whoami", "/viewer/whoami" };
+
     TString ClientId = DEFAULT_CLIENT_ID;
     TString SessionServiceEndpoint;
     TString SessionServiceTokenName;
@@ -28,7 +32,8 @@ struct TOpenIdConnectSettings {
     TString ClientSecret;
     std::vector<TString> AllowedProxyHosts;
     TString WhoamiExtendedInfoEndpoint;
-    ui32 EnrichmentProcessTimeoutMs = DEFAULT_ENRICHMENT_PROCESS_TIMEOUT_MS;
+    TDuration DefaultRequestTimeout = DEFAULT_REQUEST_TIMEOUT;
+    THashMap<TStringBuf, TDuration> RequestTimeoutsByPath;
 
     NMvp::EAccessServiceType AccessServiceType = NMvp::yandex_v2;
     TString AuthUrlPath = DEFAULT_AUTH_URL_PATH;
@@ -36,6 +41,9 @@ struct TOpenIdConnectSettings {
     TString ExchangeUrlPath = DEFAULT_EXCHANGE_URL_PATH;
     TString ImpersonateUrlPath = DEFAULT_IMPERSONATE_URL_PATH;
 
+    bool EnabledExtensionWhoami() const;
+    void InitRequestTimeoutsByPath();
+
     TString GetAuthorizationString() const;
     TString GetAuthEndpointURL() const;
     TString GetTokenEndpointURL() const;
diff --git a/ydb/mvp/oidc_proxy/openid_connect.h b/ydb/mvp/oidc_proxy/openid_connect.h
index a4b3f958832..3ffe58290d4 100644
--- a/ydb/mvp/oidc_proxy/openid_connect.h
+++ b/ydb/mvp/oidc_proxy/openid_connect.h
@@ -1,6 +1,6 @@
 #pragma once
 #include "cracked_page.h"
-#include "extension_context.h"
+#include "extension.h"
 #include "context.h"
 #include "oidc_settings.h"
 #include <ydb/library/actors/core/events.h>
@@ -16,6 +16,7 @@
 namespace NMVP::NOIDC {
 
 struct TOpenIdConnectSettings;
+struct TExtensionContext;
 
 constexpr TStringBuf IAM_TOKEN_SCHEME = "Bearer ";
 constexpr TStringBuf IAM_TOKEN_SCHEME_LOWER = "bearer ";