From upstream: imporve oauth provider UI/UX (#2983) - #3
Conversation
* feat: imporve UI/UX * fix: stabilize provider enabled toggle and polish custom OAuth settings UX * fix: add access policy/message templates and persist advanced fields reliably * fix: move template fill actions below fields and keep advanced form flow cleaner
…service layer Restructure the task relay system for better separation of concerns: - Extract task billing into service/task_billing.go with unified settlement flow - Move task polling loop from controller to service/task_polling.go (supports Suno + video platforms) - Split RelayTask into fetch/submit paths with dedicated retry logic (taskSubmitWithRetry) - Add TaskDto, TaskResponse generics, and FetchReq to dto/task.go - Add taskcommon/helpers.go for shared task adaptor utilities - Remove controller/task_video.go (logic consolidated into service layer) - Update all task adaptors (ali, doubao, gemini, hailuo, jimeng, kling, sora, suno, vertex, vidu) - Simplify frontend task logs to use new TaskDto response format
Add three billing lifecycle methods to the TaskAdaptor interface: - EstimateBilling: compute OtherRatios from user request before pricing - AdjustBillingOnSubmit: adjust ratios from upstream submit response - AdjustBillingOnComplete: determine final quota at task terminal state Introduce BaseBilling as embeddable no-op default for adaptors without custom billing. Move Sora/Ali OtherRatios logic from shared validation into per-adaptor EstimateBilling implementations. Add TaskBillingContext to persist pricing params (model_price, group_ratio, other_ratios) in task private data for async polling settlement. Extract RecalculateTaskQuota as a general-purpose delta settlement function and unify polling billing via settleTaskBillingOnComplete (adaptor-first, then token-based fallback).
- Renamed RelayTask function to RelayTaskFetch for clarity. - Updated routing in relay-router.go and video-router.go to use RelayTaskFetch for fetch operations. - Enhanced error handling in RelayTaskFetch function. - Adjusted task data conversion in TaskAdaptor to include task ID.
- Updated the remix handling in ResolveOriginTask to prioritize extracting OtherRatios from the BillingContext of the original task if available. - Retained the previous logic for extracting seconds and size from task data as a fallback. - Improved clarity and maintainability of the remix logic by separating the new and old approaches.
- Enhanced the RelayTask function to utilize a locked channel when available, allowing for better reuse during retries. - Updated error handling to ensure proper context setup for the selected channel. - Clarified comments in ResolveOriginTask regarding channel locking and retry behavior. - Introduced a new field in TaskRelayInfo to store the locked channel object, improving type safety and reducing import cycles.
… conflicts Replace all bare task.Update() (DB.Save) calls with UpdateWithStatus(), which adds a WHERE status = ? guard to prevent concurrent processes from overwriting each other's state transitions. Key changes: model/task.go: - Add taskSnapshot struct with Equal() method for change detection - Add Snapshot() method to capture pre-update state - Add UpdateWithStatus(fromStatus) using DB.Where().Save() for CAS semantics with full-struct save (no explicit field listing needed) model/midjourney.go: - Add UpdateWithStatus(fromStatus string) with same CAS pattern service/task_polling.go (updateVideoSingleTask): - Snapshot before processing upstream response; skip DB write if unchanged - Terminal transitions (SUCCESS/FAILURE) use UpdateWithStatus CAS: billing/refund only executes if this process wins the transition - Non-terminal updates also use UpdateWithStatus to prevent overwriting a concurrent terminal transition back to IN_PROGRESS - Defer settleTaskBillingOnComplete to after CAS check (shouldSettle flag) relay/relay_task.go (tryRealtimeFetch): - Add snapshot + change detection; use UpdateWithStatus for CAS safety controller/midjourney.go (UpdateMidjourneyTaskBulk): - Capture preStatus before mutations; use UpdateWithStatus CAS - Gate refund (IncreaseUserQuota) on CAS success (won && shouldReturnQuota) This prevents the multi-instance race condition where: 1. Instance A reads task (IN_PROGRESS), fetches upstream (still IN_PROGRESS) 2. Instance B reads same task, fetches upstream (now SUCCESS), writes SUCCESS 3. Instance A's bare Save() overwrites SUCCESS back to IN_PROGRESS
…gration tests
- Updated UpdateWithStatus method to use Model().Select("*").Updates() for conditional updates, preventing GORM's INSERT fallback.
- Introduced comprehensive integration tests for UpdateWithStatus, covering scenarios for winning and losing CAS updates, as well as concurrent updates.
- Added task_cas_test.go to validate the new behavior and ensure data integrity during concurrent state transitions.
Summary of ChangesHello @Starry-Sky-World, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly enhances the custom OAuth provider functionality by introducing advanced access control policies, improving the configuration user experience, and adding support for OIDC Discovery. Users can now define granular access rules based on user information and customize denial messages. The UI for managing custom OAuth providers has been refined, including automatic endpoint discovery and better organization of advanced settings. Highlights
Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
There was a problem hiding this comment.
Code Review
This pull request significantly enhances the custom OAuth provider system, adding a flexible access policy engine and OIDC discovery support. However, the new OIDC discovery feature introduces a critical Server-Side Request Forgery (SSRF) vulnerability. While restricted to root administrators, it lacks proper validation of target URLs, potentially allowing requests to internal network resources. Additionally, I recommend optimizing the access policy message rendering by pre-compiling the regular expressions used for template variable replacement.
| httpReq.Header.Set("Accept", "application/json") | ||
|
|
||
| client := &http.Client{Timeout: 20 * time.Second} | ||
| resp, err := client.Do(httpReq) |
There was a problem hiding this comment.
The FetchCustomOAuthDiscovery function is vulnerable to Server-Side Request Forgery (SSRF). It allows a root administrator to specify an arbitrary URL (WellKnownURL or IssuerURL) which the server then fetches using http.Client.Do on line 172. There is no validation to ensure that the target URL does not point to internal network resources or cloud metadata services (e.g., 169.254.169.254). Although this endpoint is restricted to root administrators, it could be exploited via CSRF or if a root account is compromised to pivot into the internal network and access sensitive services.
| config *model.CustomOAuthProvider | ||
| } | ||
|
|
||
| type accessPolicy struct { |
There was a problem hiding this comment.
It is recommended to define the regular expressions used for access policy message rendering as global variables to avoid repeated compilation during the authentication flow.
var (
currentPattern = regexp.MustCompile(`\\{\\{current\\.([^}]+)\\}\\}`)
requiredPattern = regexp.MustCompile(`\\{\\{required\\.([^}]+)\\}\\}`)
)
type accessPolicy struct {| currentPattern := regexp.MustCompile(`\{\{current\.([^}]+)\}\}`) | ||
| message = currentPattern.ReplaceAllStringFunc(message, func(token string) string { |
| requiredPattern := regexp.MustCompile(`\{\{required\.([^}]+)\}\}`) | ||
| message = requiredPattern.ReplaceAllStringFunc(message, func(token string) string { |
…try fix for async tasks
1. Async task model redirection (aligned with sync tasks):
- Integrate ModelMappedHelper in RelayTaskSubmit after model name
determination, populating OriginModelName / UpstreamModelName on RelayInfo.
- All task adaptors now send UpstreamModelName to upstream providers:
- Gemini & Vertex: BuildRequestURL uses UpstreamModelName.
- Doubao & Ali: BuildRequestBody conditionally overwrites body.Model.
- Vidu, Kling, Hailuo, Jimeng: convertToRequestPayload accepts RelayInfo
and unconditionally uses info.UpstreamModelName.
- Sora: BuildRequestBody parses JSON and multipart bodies to replace
the "model" field with UpstreamModelName.
- Frontend log visibility: LogTaskConsumption and taskBillingOther now
emit is_model_mapped / upstream_model_name in the "other" JSON field.
- Billing safety: RecalculateTaskQuotaByTokens reads model name from
BillingContext.OriginModelName (via taskModelName) instead of
task.Data["model"], preventing billing leaks from upstream model names.
2. Per-call billing (TaskPricePatches lifecycle):
- Rename TaskBillingContext.ModelName → OriginModelName; add PerCallBilling
bool field, populated from TaskPricePatches at submission time.
- settleTaskBillingOnComplete short-circuits when PerCallBilling is true,
skipping both adaptor adjustments and token-based recalculation.
- Remove ModelName from TaskSubmitResult; use relayInfo.OriginModelName
consistently in controller/relay.go for billing context and logging.
3. Multipart retry boundary mismatch fix:
- Root cause: after Sora (or OpenAI audio) rebuilds a multipart body with a
new boundary and overwrites c.Request.Header["Content-Type"], subsequent
calls to ParseMultipartFormReusable on retry would parse the cached
original body with the wrong boundary, causing "NextPart: EOF".
- Fix: ParseMultipartFormReusable now caches the original Content-Type in
gin context key "_original_multipart_ct" on first call and reuses it for
all subsequent parses, making multipart parsing retry-safe globally.
- Sora adaptor reverted to the standard pattern (direct header set/get),
which is now safe thanks to the root fix.
4. Tests:
- task_billing_test.go: update makeTask to use OriginModelName; add
PerCallBilling settlement tests (skip adaptor adjust, skip token recalc);
add non-per-call adaptor adjustment test with refund verification.
…keys across multiple languages
refactor: async task
|
/gemini review |
|
@gemini-code-assist review |
There was a problem hiding this comment.
Code Review
This is an excellent and substantial pull request that significantly improves the OAuth provider functionality and refactors core components for asynchronous tasks. The introduction of access policies and OIDC discovery for custom OAuth providers is a major feature enhancement, and the implementation is robust, including backend validation and a polished user interface.
The refactoring of the task relaying and polling logic into dedicated services is a standout improvement. This greatly enhances code organization, maintainability, and correctness. Key highlights of the refactoring include:
- Improved Concurrency Control: The adoption of Compare-And-Swap (CAS) updates (
UpdateWithStatus) effectively prevents race conditions during concurrent task status updates. - Robust Billing Lifecycle: The new
TaskAdaptorinterface and the centralized billing logic inservice/task_billing.gocreate a clear and reliable lifecycle for pre-consumption, settlement, and refunds for async tasks. - Better Abstraction: Separating public-facing task IDs from upstream provider IDs and introducing helpers like
GetUpstreamTaskID()andGetResultURL()improves abstraction and future-proofs the data model. - Clearer Code Structure: Moving complex logic out of controllers and into dedicated service files (
task_polling.go,task_billing.go) and common helpers (taskcommon/helpers.go) follows best practices and makes the codebase much easier to navigate.
The frontend changes are equally impressive, providing a much-improved user experience for configuring custom OAuth providers, complete with presets, discovery, and access policy templates. The updates to the logging and task tables to reflect the new data model are also well-executed.
Overall, this is a high-quality contribution that demonstrates a deep understanding of the system's architecture and a commitment to improving its long-term health. The changes are well-designed, correctly implemented, and thoroughly tested. Great work!
|
@gemini-code-assist 这个修改了什么? |
Summary of ChangesThis pull request delivers a comprehensive overhaul of the custom OAuth provider system, significantly enhancing its configurability, security, and user experience. It introduces advanced access control capabilities, streamlines the setup process through OIDC discovery, and refactors the backend for asynchronous task management to ensure more reliable billing and status updates. The frontend has been updated to reflect these new features with an improved, more intuitive interface. Highlights
Changelog
Activity
|
feat: imporve UI/UX
fix: stabilize provider enabled toggle and polish custom OAuth settings UX
fix: add access policy/message templates and persist advanced fields reliably
fix: move template fill actions below fields and keep advanced form flow cleaner