diff --git a/src/agentpool_server/acp_server/acp_agent.py b/src/agentpool_server/acp_server/acp_agent.py index 8cc5ffaa6..53131e831 100644 --- a/src/agentpool_server/acp_server/acp_agent.py +++ b/src/agentpool_server/acp_server/acp_agent.py @@ -77,6 +77,30 @@ logger = get_logger(__name__) +_LOCALE_MAX_LENGTH = 15 + + +def _extract_locale(field_meta: dict[str, Any] | None) -> str | None: + """Extract and validate locale from ACP client metadata. + + Only allows alphanumeric characters, hyphens, and underscores to + prevent prompt injection via crafted locale strings. + """ + if field_meta is None: + return None + raw = field_meta.get("locale") + if not isinstance(raw, str): + return None + locale = raw.strip() + if not locale or len(locale) > _LOCALE_MAX_LENGTH: + logger.warning("Invalid client locale ignored", locale=raw) + return None + if not all(c.isascii() and (c.isalnum() or c in "-_") for c in locale): + logger.warning("Invalid client locale ignored", locale=raw) + return None + logger.info("Client locale", locale=locale) + return locale + async def get_session_model_state( agent: BaseAgent, @@ -254,6 +278,7 @@ def __post_init__(self) -> None: """Initialize derived attributes and setup after field assignment.""" self.client_capabilities: ClientCapabilities | None = None self.client_info: Implementation | None = None + self.client_locale: str | None = None pool = self.agent_pool if pool is None: msg = "Default agent has no associated pool" @@ -399,6 +424,9 @@ async def initialize(self, params: InitializeRequest) -> InitializeResponse: version = min(params.protocol_version, self.PROTOCOL_VERSION) self.client_capabilities = params.client_capabilities self.client_info = params.client_info + # Extract locale from client metadata (field_meta / _meta in JSON-RPC). + # Validate format to prevent prompt injection via crafted locale strings. + self.client_locale = _extract_locale(params.field_meta) logger.info("Client info", request=params.model_dump_json()) self._initialized = True # Forward client capabilities to the SessionPool protocol handler so @@ -449,6 +477,7 @@ async def new_session(self, params: NewSessionRequest) -> NewSessionResponse: mcp_servers=params.mcp_servers, client_capabilities=self.client_capabilities, client_info=self.client_info, + client_locale=self.client_locale, subagent_display_mode=self.subagent_display_mode, raw_input_mode=self.raw_input_mode, connection_id=self._get_connection_id(), @@ -529,6 +558,7 @@ async def load_session(self, params: LoadSessionRequest) -> LoadSessionResponse: mcp_servers=params.mcp_servers, client_capabilities=self.client_capabilities, client_info=self.client_info, + client_locale=self.client_locale, subagent_display_mode=self.subagent_display_mode, connection_id=self._get_connection_id(), ) @@ -623,6 +653,7 @@ async def fork_session(self, params: ForkSessionRequest) -> ForkSessionResponse: mcp_servers=params.mcp_servers, client_capabilities=self.client_capabilities, client_info=self.client_info, + client_locale=self.client_locale, subagent_display_mode=self.subagent_display_mode, raw_input_mode=self.raw_input_mode, connection_id=self._get_connection_id(), @@ -652,6 +683,7 @@ async def resume_session(self, params: ResumeSessionRequest) -> ResumeSessionRes mcp_servers=params.mcp_servers, client_capabilities=self.client_capabilities, client_info=self.client_info, + client_locale=self.client_locale, subagent_display_mode=self.subagent_display_mode, connection_id=self._get_connection_id(), ) diff --git a/src/agentpool_server/acp_server/session.py b/src/agentpool_server/acp_server/session.py index c8451c75e..18c3aba36 100644 --- a/src/agentpool_server/acp_server/session.py +++ b/src/agentpool_server/acp_server/session.py @@ -187,6 +187,9 @@ class ACPSession: client_info: Implementation | None = None """Client implementation info (name, version, title)""" + client_locale: str | None = None + """Client locale for i18n (e.g. "en", "zh-CN"). Injected into agent prompts.""" + manager: ACPSessionManager | None = None """Session manager for managing sessions. Used for session management commands.""" @@ -250,6 +253,8 @@ def __post_init__(self) -> None: # to AcpMcpConnectionManager.cleanup_session(), leaking per-session # ACP stream pairs and reverse-index entries. self.agent.mcp._acp_mcp_manager = self.acp_agent._mcp_manager + if self.client_locale: + self.agent.sys_prompts.prompts.append(self.get_locale_prompt) # pyright: ignore[reportArgumentType] # ty: ignore[invalid-argument-type] if isinstance(self.agent, ACPAgent): async def permission_callback( @@ -585,6 +590,16 @@ def get_cwd_context(self) -> str: """Get current working directory context for prompts.""" return f"Working directory: {self.cwd}" if self.cwd else "" + def get_locale_prompt(self) -> str: + """Get locale directive for the agent. + + Returns a prompt instructing the agent to respond in the client's + preferred locale. Called by pydantic-ai on each model invocation. + """ + if not self.client_locale: + return "" + return f"Language: You MUST respond in {self.client_locale}." + async def switch_active_agent(self, agent_name: str) -> None: """Switch to a different agent in the pool. @@ -605,6 +620,11 @@ async def switch_active_agent(self, agent_name: str) -> None: # Remove session-specific mutations from old agent before switching if isinstance(self.agent, Agent) and self.get_cwd_context in self.agent.sys_prompts.prompts: self.agent.sys_prompts.prompts.remove(self.get_cwd_context) # pyright: ignore[reportArgumentType] # ty: ignore[invalid-argument-type] + if ( + isinstance(self.agent, Agent) + and self.get_locale_prompt in self.agent.sys_prompts.prompts + ): + self.agent.sys_prompts.prompts.remove(self.get_locale_prompt) # pyright: ignore[reportArgumentType] # ty: ignore[invalid-argument-type] # Create new session agent via SessionPool (pool-level agents removed) pool = self.agent_pool @@ -623,6 +643,8 @@ async def switch_active_agent(self, agent_name: str) -> None: self.agent._input_provider = self.input_provider if isinstance(self.agent, Agent): self.agent.sys_prompts.prompts.append(self.get_cwd_context) # pyright: ignore[reportArgumentType] # ty: ignore[invalid-argument-type] + if self.client_locale: + self.agent.sys_prompts.prompts.append(self.get_locale_prompt) # pyright: ignore[reportArgumentType] # ty: ignore[invalid-argument-type] # Reconnect signal with suppress(Exception): @@ -830,6 +852,10 @@ async def close(self) -> None: self.get_cwd_context in self.agent.sys_prompts.prompts ): self.agent.sys_prompts.prompts.remove(self.get_cwd_context) # pyright: ignore[reportArgumentType] # ty: ignore[invalid-argument-type] + if isinstance(self.agent, Agent) and ( + self.get_locale_prompt in self.agent.sys_prompts.prompts + ): + self.agent.sys_prompts.prompts.remove(self.get_locale_prompt) # pyright: ignore[reportArgumentType] # ty: ignore[invalid-argument-type] # Unregister skill command callback to prevent memory leak if hasattr(self, "_skill_command_callback"): diff --git a/src/agentpool_server/acp_server/session_manager.py b/src/agentpool_server/acp_server/session_manager.py index a9bcf5000..e8cf116d9 100644 --- a/src/agentpool_server/acp_server/session_manager.py +++ b/src/agentpool_server/acp_server/session_manager.py @@ -76,6 +76,7 @@ async def create_session( session_id: str | None = None, client_capabilities: ClientCapabilities | None = None, client_info: Implementation | None = None, + client_locale: str | None = None, subagent_display_mode: Literal["legacy", "zed", "qwen"] = "legacy", raw_input_mode: Literal["dict", "skip", "json_str"] = "dict", parent_session_id: str | None = None, @@ -92,6 +93,7 @@ async def create_session( session_id: Optional specific session ID (generated if None) client_capabilities: Client capabilities for tool registration client_info: Client implementation info (name, version) + client_locale: Client locale for i18n (e.g. "en", "zh-CN") subagent_display_mode: Display mode for subagent outputs raw_input_mode: How to emit tool call raw_input parent_session_id: Optional parent session ID for child sessions. @@ -198,6 +200,7 @@ async def create_session( acp_agent=acp_agent, client_capabilities=client_capabilities or ClientCapabilities(), client_info=client_info, + client_locale=client_locale, manager=self, subagent_display_mode=subagent_display_mode, raw_input_mode=raw_input_mode, @@ -229,6 +232,7 @@ async def resume_session( acp_agent: AgentPoolACPAgent, client_capabilities: ClientCapabilities | None = None, client_info: Implementation | None = None, + client_locale: str | None = None, subagent_display_mode: Literal["legacy", "zed", "qwen"] = "legacy", raw_input_mode: Literal["dict", "skip", "json_str"] = "dict", mcp_servers: Sequence[McpServer] | None = None, @@ -242,6 +246,7 @@ async def resume_session( acp_agent: ACP agent instance client_capabilities: Client capabilities client_info: Client implementation info (name, version) + client_locale: Client locale for i18n (e.g. "en", "zh-CN") subagent_display_mode: Display mode for subagent outputs raw_input_mode: How to emit tool call raw_input mcp_servers: MCP server configurations to (re-)initialize @@ -321,6 +326,7 @@ async def resume_session( acp_agent=acp_agent, client_capabilities=client_capabilities or ClientCapabilities(), client_info=client_info, + client_locale=client_locale, manager=self, subagent_display_mode=subagent_display_mode, raw_input_mode=raw_input_mode, diff --git a/tests/agents/native_agent/test_get_agentlet_capabilities.py b/tests/agents/native_agent/test_get_agentlet_capabilities.py index fe3c7098b..d80360fed 100644 --- a/tests/agents/native_agent/test_get_agentlet_capabilities.py +++ b/tests/agents/native_agent/test_get_agentlet_capabilities.py @@ -639,6 +639,7 @@ async def test_from_config_capabilities_not_duplicated() -> None: This test calls from_config() with a config containing a capability, then calls get_agentlet() and verifies the capability appears exactly once. """ + from llmling_models_config import TestModelConfig from pydantic_ai.capabilities import Instrumentation from agentpool.models.agents import NativeAgentConfig @@ -650,7 +651,7 @@ async def test_from_config_capabilities_not_duplicated() -> None: ) config = NativeAgentConfig( name="test_dedup_agent", - model="openai:gpt-4o-mini", + model=TestModelConfig(custom_output_text="test"), system_prompt=["Be helpful."], capabilities=[cap_config], ) diff --git a/tests/servers/acp_server/test_acp_session_resume.py b/tests/servers/acp_server/test_acp_session_resume.py index 6633345bd..850ddcab3 100644 --- a/tests/servers/acp_server/test_acp_session_resume.py +++ b/tests/servers/acp_server/test_acp_session_resume.py @@ -536,6 +536,7 @@ async def test_resume_passes_through_to_session_manager_correctly( mcp_servers=resume_session_request.mcp_servers, client_capabilities=mock_acp_agent.client_capabilities, client_info=mock_acp_agent.client_info, + client_locale=mock_acp_agent.client_locale, subagent_display_mode=mock_acp_agent.subagent_display_mode, connection_id=mock_acp_agent._get_connection_id(), )