feat(auth): add health check endpoint for auth file models #1208
Conversation
Summary of ChangesHello @myuan19, 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! 此拉取请求引入了一项重要的健康检查功能，允许用户验证其配置的认证文件中的模型是否正常工作。通过模拟真实的API请求并检查响应，系统能够提供每个模型的详细健康状态，包括延迟和潜在错误信息。这极大地提升了系统的可观测性和故障排除能力，确保用户可以快速识别并解决模型配置或连接问题。 Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. 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
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
此拉取请求添加了一个新的健康检查端点，用于验证与身份验证文件关联的模型的可用性。主要逻辑在 internal/api/handlers/management/auth_files.go 中实现，通过向每个模型发送一个简短的流式请求并检查初始响应来工作。sdk/cliproxy/auth/conductor.go 中的新辅助方法支持此功能，允许直接针对特定的身份验证凭据执行请求，绕过了通常的负载均衡。
代码整体结构良好，但存在一些可以改进的地方。conductor.go 中的 ExecuteStreamWithAuth 和 ExecuteWithAuth 函数有大量重复代码，可以重构为一个共享的辅助函数。此外，在 auth_files.go 中，并发和顺序健康检查的执行逻辑也是重复的，可以合并以提高可维护性。最后，在处理失败的健康检查时，启动一个goroutine来排空流是不必要的，因为在调用 cancel() 之后，流的读取端会收到一个错误，可以安全地退出。
| // Drain remaining chunks | ||
| go func() { | ||
| for range stream { | ||
| } | ||
| }() |
There was a problem hiding this comment.
在处理失败的健康检查时，在调用 cancel() 之后，没有必要再启动一个 goroutine 来排空流。cancel() 会导致 ExecuteStreamWithAuth 中的上下文被取消，这通常会使流的读取端（即 <-stream）立即出错或关闭。因此，select 语句中的 case chunk, ok := <-stream: 将会处理这个关闭或错误，而不需要额外的 goroutine 来排空。
| if isConcurrent { | ||
| // Concurrent execution | ||
| for _, model := range models { | ||
| wg.Add(1) | ||
| go checkModel(model) | ||
| } | ||
| } else { | ||
| // Sequential execution | ||
| for _, model := range models { | ||
| wg.Add(1) | ||
| checkModel(model) | ||
| } | ||
| } |
There was a problem hiding this comment.
并发和顺序执行的逻辑几乎完全相同，唯一的区别在于是否使用 go 关键字。这种重复可以通过将循环提取出来并有条件地决定是并发还是同步调用 checkModel 来避免，从而使代码更简洁。
for _, model := range models {
wg.Add(1)
if isConcurrent {
go checkModel(model)
} else {
checkModel(model)
}
}| // ExecuteStreamWithAuth performs a streaming execution using a specific auth directly. | ||
| // Unlike ExecuteStream which performs load balancing across auths of the same provider, | ||
| // this method targets a specific auth entry - useful for health checks and diagnostics. | ||
| // The method handles model rewriting, proxy/RoundTripper setup, and proper request translation. | ||
| func (m *Manager) ExecuteStreamWithAuth(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (<-chan cliproxyexecutor.StreamChunk, error) { | ||
| if m == nil { | ||
| return nil, &Error{Code: "manager_nil", Message: "auth manager is nil"} | ||
| } | ||
| if auth == nil { | ||
| return nil, &Error{Code: "auth_not_found", Message: "auth is nil"} | ||
| } | ||
|
|
||
| // Determine provider key from auth | ||
| providerKey := executorKeyFromAuth(auth) | ||
| if providerKey == "" { | ||
| return nil, &Error{Code: "provider_not_found", Message: "auth provider is empty"} | ||
| } | ||
|
|
||
| // Get the executor for this provider | ||
| executor := m.executorFor(providerKey) | ||
| if executor == nil { | ||
| return nil, &Error{Code: "provider_not_found", Message: "executor not registered for provider: " + providerKey} | ||
| } | ||
|
|
||
| // Setup execution context with RoundTripper (proxy support) | ||
| execCtx := ctx | ||
| if rt := m.roundTripperFor(auth); rt != nil { | ||
| execCtx = context.WithValue(execCtx, roundTripperContextKey{}, rt) | ||
| execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt) | ||
| } | ||
|
|
||
| // Apply model rewriting if configured | ||
| routeModel := req.Model | ||
| execReq := req | ||
| execReq.Model = rewriteModelForAuth(routeModel, auth) | ||
| execReq.Model = m.applyOAuthModelAlias(auth, execReq.Model) | ||
| execReq.Model = m.applyAPIKeyModelAlias(auth, execReq.Model) | ||
|
|
||
| // Ensure requested model metadata is set | ||
| opts = ensureRequestedModelMetadata(opts, routeModel) | ||
|
|
||
| // Execute the stream request directly with the specified auth | ||
| return executor.ExecuteStream(execCtx, auth, execReq, opts) | ||
| } | ||
|
|
||
| // ExecuteWithAuth performs a non-streaming execution using a specific auth directly. | ||
| // This is the non-streaming counterpart to ExecuteStreamWithAuth. | ||
| func (m *Manager) ExecuteWithAuth(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { | ||
| if m == nil { | ||
| return cliproxyexecutor.Response{}, &Error{Code: "manager_nil", Message: "auth manager is nil"} | ||
| } | ||
| if auth == nil { | ||
| return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "auth is nil"} | ||
| } | ||
|
|
||
| // Determine provider key from auth | ||
| providerKey := executorKeyFromAuth(auth) | ||
| if providerKey == "" { | ||
| return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "auth provider is empty"} | ||
| } | ||
|
|
||
| // Get the executor for this provider | ||
| executor := m.executorFor(providerKey) | ||
| if executor == nil { | ||
| return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "executor not registered for provider: " + providerKey} | ||
| } | ||
|
|
||
| // Setup execution context with RoundTripper (proxy support) | ||
| execCtx := ctx | ||
| if rt := m.roundTripperFor(auth); rt != nil { | ||
| execCtx = context.WithValue(execCtx, roundTripperContextKey{}, rt) | ||
| execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt) | ||
| } | ||
|
|
||
| // Apply model rewriting if configured | ||
| routeModel := req.Model | ||
| execReq := req | ||
| execReq.Model = rewriteModelForAuth(routeModel, auth) | ||
| execReq.Model = m.applyOAuthModelAlias(auth, execReq.Model) | ||
| execReq.Model = m.applyAPIKeyModelAlias(auth, execReq.Model) | ||
|
|
||
| // Ensure requested model metadata is set | ||
| opts = ensureRequestedModelMetadata(opts, routeModel) | ||
|
|
||
| // Execute the request directly with the specified auth | ||
| return executor.Execute(execCtx, auth, execReq, opts) | ||
| } |
There was a problem hiding this comment.
ExecuteStreamWithAuth 和 ExecuteWithAuth 两个函数之间存在大量重复代码。它们都执行相同的设置步骤：空值检查、获取执行器、设置执行上下文和重写模型。这些通用逻辑可以提取到一个私有的辅助函数中，以减少重复并提高可维护性。这个辅助函数可以处理所有共享的准备工作，并返回执行器和准备好的请求/选项，然后每个公共函数只需调用相应的执行器方法（ExecuteStream 或 Execute）。
2 participants
添加健康检查功能endpoints