Conversation
- Fix parser factories not initialized in tests: add create_test_context_with_parsers() helper - Refactor handlers.rs: replace nested match with early returns using let-else pattern - Fix race condition in reserve_id_for_url: use atomic entry().or_default() API - Extract verbose WorkerInfo creation: add WorkerInfo::pending() constructor - Reduce worker_id parsing duplication: add parse_worker_id_or_error() helper
Summary of ChangesHello @slin1237, 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 focuses on minor code cleanup and reliability improvements across the 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
|
There was a problem hiding this comment.
Code Review
This pull request introduces a series of valuable code cleanups and fixes that improve code quality and maintainability. The refactoring in handlers.rs to use the let-else pattern significantly improves readability. The race condition fix in worker_registry.rs is crucial for correctness. The extraction of helper functions and constructors like WorkerInfo::pending() and parse_worker_id_or_error() effectively reduces code duplication in the application logic.
The changes also include fixes to tests to properly initialize parser factories, making the test suite more robust.
I have one suggestion regarding code duplication in the new test helper create_test_context_with_parsers(). Addressing this would further improve the maintainability of the test code.
| /// Helper function to create AppContext for tests with parser factories initialized | ||
| pub async fn create_test_context_with_parsers(config: RouterConfig) -> Arc<AppContext> { | ||
| let client = reqwest::Client::new(); | ||
|
|
||
| // Initialize rate limiter | ||
| let rate_limiter = match config.max_concurrent_requests { | ||
| n if n <= 0 => None, | ||
| n => { | ||
| let rate_limit_tokens = config | ||
| .rate_limit_tokens_per_second | ||
| .filter(|&t| t > 0) | ||
| .unwrap_or(n); | ||
| Some(Arc::new(TokenBucket::new( | ||
| n as usize, | ||
| rate_limit_tokens as usize, | ||
| ))) | ||
| } | ||
| }; | ||
|
|
||
| // Initialize registries | ||
| let worker_registry = Arc::new(WorkerRegistry::new()); | ||
| let policy_registry = Arc::new(PolicyRegistry::new(config.policy.clone())); | ||
|
|
||
| // Initialize storage backends (Memory for tests) | ||
| let response_storage = Arc::new(MemoryResponseStorage::new()); | ||
| let conversation_storage = Arc::new(MemoryConversationStorage::new()); | ||
| let conversation_item_storage = Arc::new(MemoryConversationItemStorage::new()); | ||
|
|
||
| // Initialize load monitor | ||
| let load_monitor = Some(Arc::new(LoadMonitor::new( | ||
| worker_registry.clone(), | ||
| policy_registry.clone(), | ||
| client.clone(), | ||
| config.worker_startup_check_interval_secs, | ||
| ))); | ||
|
|
||
| // Create empty OnceLock for worker job queue, workflow engine, and mcp manager | ||
| let worker_job_queue = Arc::new(OnceLock::new()); | ||
| let workflow_engine = Arc::new(OnceLock::new()); | ||
| let mcp_manager_lock = Arc::new(OnceLock::new()); | ||
|
|
||
| // Initialize parser factories | ||
| let reasoning_parser_factory = Some(ReasoningParserFactory::new()); | ||
| let tool_parser_factory = Some(ToolParserFactory::new()); | ||
|
|
||
| let app_context = Arc::new( | ||
| AppContext::builder() | ||
| .router_config(config.clone()) | ||
| .client(client) | ||
| .rate_limiter(rate_limiter) | ||
| .tokenizer(None) // tokenizer | ||
| .reasoning_parser_factory(reasoning_parser_factory) | ||
| .tool_parser_factory(tool_parser_factory) | ||
| .worker_registry(worker_registry) | ||
| .policy_registry(policy_registry) | ||
| .response_storage(response_storage) | ||
| .conversation_storage(conversation_storage) | ||
| .conversation_item_storage(conversation_item_storage) | ||
| .load_monitor(load_monitor) | ||
| .worker_job_queue(worker_job_queue) | ||
| .workflow_engine(workflow_engine) | ||
| .mcp_manager(mcp_manager_lock) | ||
| .build() | ||
| .unwrap(), | ||
| ); | ||
|
|
||
| // Initialize JobQueue after AppContext is created | ||
| let weak_context = Arc::downgrade(&app_context); | ||
| let job_queue = sgl_model_gateway::core::JobQueue::new( | ||
| sgl_model_gateway::core::JobQueueConfig::default(), | ||
| weak_context, | ||
| ); | ||
| app_context | ||
| .worker_job_queue | ||
| .set(job_queue) | ||
| .expect("JobQueue should only be initialized once"); | ||
|
|
||
| // Initialize WorkflowEngine and register workflows | ||
| use sgl_model_gateway::{ | ||
| core::steps::{create_worker_registration_workflow, create_worker_removal_workflow}, | ||
| workflow::WorkflowEngine, | ||
| }; | ||
| let engine = Arc::new(WorkflowEngine::new()); | ||
| engine | ||
| .register_workflow(create_worker_registration_workflow(&config)) | ||
| .expect("worker_registration workflow should be valid"); | ||
| engine | ||
| .register_workflow(create_worker_removal_workflow()) | ||
| .expect("worker_removal workflow should be valid"); | ||
| app_context | ||
| .workflow_engine | ||
| .set(engine) | ||
| .expect("WorkflowEngine should only be initialized once"); | ||
|
|
||
| // Register external workers for OpenAI mode | ||
| if let RoutingMode::OpenAI { worker_urls, .. } = &config.mode { | ||
| for url in worker_urls { | ||
| // Create a worker that supports common test models | ||
| let models = vec![ | ||
| ModelCard::new("mock-model"), | ||
| ModelCard::new("gpt-4"), | ||
| ModelCard::new("gpt-3.5-turbo"), | ||
| ]; | ||
| let worker: Arc<dyn Worker> = Arc::new( | ||
| BasicWorkerBuilder::new(url) | ||
| .worker_type(WorkerType::Regular) | ||
| .runtime_type(RuntimeType::External) | ||
| .models(models) | ||
| .build(), | ||
| ); | ||
| app_context.worker_registry.register(worker); | ||
| } | ||
| } | ||
|
|
||
| // Initialize MCP manager with empty config | ||
| use sgl_model_gateway::mcp::{McpConfig, McpManager}; | ||
| let empty_config = McpConfig { | ||
| servers: vec![], | ||
| pool: Default::default(), | ||
| proxy: None, | ||
| warmup: vec![], | ||
| inventory: Default::default(), | ||
| }; | ||
| let mcp_manager = McpManager::with_defaults(empty_config) | ||
| .await | ||
| .expect("Failed to create MCP manager"); | ||
| app_context | ||
| .mcp_manager | ||
| .set(Arc::new(mcp_manager)) | ||
| .ok() | ||
| .expect("McpManager should only be initialized once"); | ||
|
|
||
| app_context | ||
| } |
There was a problem hiding this comment.
This new helper function create_test_context_with_parsers almost entirely duplicates the existing create_test_context function. This introduces a significant maintenance burden, as any changes to the test context setup will need to be manually synchronized across both functions.
To improve maintainability, this duplicated code should be refactored. A good approach would be to extract the common logic into a private helper function that can be configured.
For example, you could introduce a private create_test_context_internal(config: RouterConfig, with_parsers: bool) function. Then, create_test_context could call it with with_parsers: false, and create_test_context_with_parsers would call it with with_parsers: true. This would eliminate over 100 lines of duplicated code and make the test setup easier to maintain.
Checklist