From 0c3e25a07b39d006d47984184bb438bf9bb7d0c4 Mon Sep 17 00:00:00 2001
From: Manoj Aggarwal <manoj6543@gmail.com>
Date: Tue, 13 Jan 2026 13:57:04 -0800
Subject: [PATCH] fix linting issues with ruff

---
 ..._allow_multiple_connectors_with_unique_.py | 10 ++++-----
 .../app/agents/new_chat/tools/mcp_client.py   | 21 ++++++++++---------
 .../app/agents/new_chat/tools/mcp_tool.py     |  3 ++-
 .../app/utils/connector_naming.py             |  6 +++---
 4 files changed, 20 insertions(+), 20 deletions(-)

diff --git a/surfsense_backend/alembic/versions/5263aa4e7f94_allow_multiple_connectors_with_unique_.py b/surfsense_backend/alembic/versions/5263aa4e7f94_allow_multiple_connectors_with_unique_.py
index 954a864ab..5a61af0cb 100644
--- a/surfsense_backend/alembic/versions/5263aa4e7f94_allow_multiple_connectors_with_unique_.py
+++ b/surfsense_backend/alembic/versions/5263aa4e7f94_allow_multiple_connectors_with_unique_.py
@@ -5,17 +5,15 @@ Revises: ffd7445eb90a
 Create Date: 2026-01-13 12:23:31.481643
 
 """
-from typing import Sequence, Union
+from collections.abc import Sequence
 
 from alembic import op
-import sqlalchemy as sa
-
 
 # revision identifiers, used by Alembic.
 revision: str = '5263aa4e7f94'
-down_revision: Union[str, None] = 'ffd7445eb90a'
-branch_labels: Union[str, Sequence[str], None] = None
-depends_on: Union[str, Sequence[str], None] = None
+down_revision: str | None = 'ffd7445eb90a'
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
 
 
 def upgrade() -> None:
diff --git a/surfsense_backend/app/agents/new_chat/tools/mcp_client.py b/surfsense_backend/app/agents/new_chat/tools/mcp_client.py
index 62ad258b9..181cbb642 100644
--- a/surfsense_backend/app/agents/new_chat/tools/mcp_client.py
+++ b/surfsense_backend/app/agents/new_chat/tools/mcp_client.py
@@ -4,7 +4,6 @@ This module provides a client for communicating with MCP servers via stdio trans
 It handles server lifecycle management, tool discovery, and tool execution.
 """
 
-import asyncio
 import logging
 import os
 from contextlib import asynccontextmanager
@@ -54,15 +53,17 @@ class MCPClient:
             )
             
             # Spawn server process and create session
-            async with stdio_client(server=server_params) as (read, write):
-                async with ClientSession(read, write) as session:
-                    # Initialize the connection
-                    await session.initialize()
-                    self.session = session
-                    logger.info(
-                        f"Connected to MCP server: {self.command} {' '.join(self.args)}"
-                    )
-                    yield session
+            async with (
+                stdio_client(server=server_params) as (read, write),
+                ClientSession(read, write) as session,
+            ):
+                # Initialize the connection
+                await session.initialize()
+                self.session = session
+                logger.info(
+                    f"Connected to MCP server: {self.command} {' '.join(self.args)}"
+                )
+                yield session
 
         except Exception as e:
             logger.error(f"Failed to connect to MCP server: {e!s}", exc_info=True)
diff --git a/surfsense_backend/app/agents/new_chat/tools/mcp_tool.py b/surfsense_backend/app/agents/new_chat/tools/mcp_tool.py
index 47f4a4cd4..f7ba27a3d 100644
--- a/surfsense_backend/app/agents/new_chat/tools/mcp_tool.py
+++ b/surfsense_backend/app/agents/new_chat/tools/mcp_tool.py
@@ -59,7 +59,7 @@ def _normalize_gemini_params(params: dict[str, Any], mcp_schema: dict[str, Any])
                             normalized_item[item_key] = item_value
                     
                     # Add missing required fields with empty defaults if needed
-                    for required_field in items_properties.keys():
+                    for required_field in items_properties:
                         if required_field not in normalized_item:
                             # For arrays like observations, default to empty array
                             if items_properties[required_field].get("type") == "array":
@@ -102,6 +102,7 @@ def _create_dynamic_input_model_from_schema(
         # Use Any type for complex schemas to preserve structure
         # This allows the MCP server to do its own validation
         from typing import Any as AnyType
+
         from pydantic import Field
         
         if is_required:
diff --git a/surfsense_backend/app/utils/connector_naming.py b/surfsense_backend/app/utils/connector_naming.py
index bfc9decdd..a2b748a3a 100644
--- a/surfsense_backend/app/utils/connector_naming.py
+++ b/surfsense_backend/app/utils/connector_naming.py
@@ -8,9 +8,9 @@ from typing import Any
 from urllib.parse import urlparse
 from uuid import UUID
 
-from sqlalchemy import func
+from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.future import select
+from sqlalchemy.sql import func
 
 from app.db import SearchSourceConnector, SearchSourceConnectorType
 
@@ -76,7 +76,7 @@ def extract_identifier_from_credentials(
                 if ".atlassian.net" in hostname:
                     return hostname.replace(".atlassian.net", "")
                 return hostname
-            except Exception:
+            except (ValueError, TypeError, AttributeError):
                 pass
         return None