fix: ensure close aiodocker.Docker() (#4251)

* fix: ensure close aiodocker.Docker()

* fix: code formatted

Author: fluidcat

astrbot/builtin_stars/python_interpreter/main.py

@@ -157,9 +157,8 @@ async def file_upload(self, file_path: str):
     async def is_docker_available(self) -> bool:
         """Check if docker is available"""
         try:
-            docker = aiodocker.Docker()
-            await docker.version()
-            await docker.close()
+            async with aiodocker.Docker() as docker:
+                await docker.version()
             return True
         except BaseException as e:
             logger.info(f"检查 Docker 可用性: {e}")
@@ -279,14 +278,14 @@ async def pi_mirror(self, event: AstrMessageEvent, url: str = ""):
     @pi.command("repull")
     async def pi_repull(self, event: AstrMessageEvent):
         """重新拉取沙箱镜像"""
-        docker = aiodocker.Docker()
-        image_name = await self.get_image_name()
-        try:
-            await docker.images.get(image_name)
-            await docker.images.delete(image_name, force=True)
-        except aiodocker.exceptions.DockerError:
-            pass
-        await docker.images.pull(image_name)
+        async with aiodocker.Docker() as docker:
+            image_name = await self.get_image_name()
+            try:
+                await docker.images.get(image_name)
+                await docker.images.delete(image_name, force=True)
+            except aiodocker.exceptions.DockerError:
+                pass
+            await docker.images.pull(image_name)
         yield event.plain_result("重新拉取沙箱镜像成功。")
 
     @pi.command("file")
@@ -371,137 +370,137 @@ async def python_interpreter(self, event: AstrMessageEvent):
         obs = ""
         n = 5
 
-        for i in range(n):
-            if i > 0:
-                logger.info(f"Try {i + 1}/{n}")
-
-            PROMPT_ = PROMPT.format(
-                prompt=plain_text,
-                extra_input=extra_inputs,
-                extra_prompt=obs,
-            )
-            provider = self.context.get_using_provider()
-            llm_response = await provider.text_chat(
-                prompt=PROMPT_,
-                session_id=f"{event.session_id}_{magic_code}_{i!s}",
-            )
-
-            logger.debug(
-                "code interpreter llm gened code:" + llm_response.completion_text,
-            )
-
-            # 整理代码并保存
-            code_clean = await self.tidy_code(llm_response.completion_text)
-            with open(os.path.join(workplace_path, "exec.py"), "w") as f:
-                f.write(code_clean)
-
-            # 启动容器
-            docker = aiodocker.Docker()
-
-            # 检查有没有image
-            image_name = await self.get_image_name()
-            try:
-                await docker.images.get(image_name)
-            except aiodocker.exceptions.DockerError:
-                # 拉取镜像
-                logger.info(f"未找到沙箱镜像，正在尝试拉取 {image_name}...")
-                await docker.images.pull(image_name)
+        async with aiodocker.Docker() as docker:
+            for i in range(n):
+                if i > 0:
+                    logger.info(f"Try {i + 1}/{n}")
 
-            yield event.plain_result(
-                f"使用沙箱执行代码中，请稍等...(尝试次数: {i + 1}/{n})",
-            )
+                PROMPT_ = PROMPT.format(
+                    prompt=plain_text,
+                    extra_input=extra_inputs,
+                    extra_prompt=obs,
+                )
+                provider = self.context.get_using_provider()
+                llm_response = await provider.text_chat(
+                    prompt=PROMPT_,
+                    session_id=f"{event.session_id}_{magic_code}_{i!s}",
+                )
 
-            self.docker_host_astrbot_abs_path = self.config.get(
-                "docker_host_astrbot_abs_path",
-                "",
-            )
-            if self.docker_host_astrbot_abs_path:
-                host_shared = os.path.join(
-                    self.docker_host_astrbot_abs_path,
-                    self.shared_path,
+                logger.debug(
+                    "code interpreter llm gened code:" + llm_response.completion_text,
                 )
-                host_output = os.path.join(
-                    self.docker_host_astrbot_abs_path,
-                    output_path,
+
+                # 整理代码并保存
+                code_clean = await self.tidy_code(llm_response.completion_text)
+                with open(os.path.join(workplace_path, "exec.py"), "w") as f:
+                    f.write(code_clean)
+
+                # 检查有没有image
+                image_name = await self.get_image_name()
+                try:
+                    await docker.images.get(image_name)
+                except aiodocker.exceptions.DockerError:
+                    # 拉取镜像
+                    logger.info(f"未找到沙箱镜像，正在尝试拉取 {image_name}...")
+                    await docker.images.pull(image_name)
+
+                yield event.plain_result(
+                    f"使用沙箱执行代码中，请稍等...(尝试次数: {i + 1}/{n})",
                 )
-                host_workplace = os.path.join(
-                    self.docker_host_astrbot_abs_path,
-                    workplace_path,
+
+                self.docker_host_astrbot_abs_path = self.config.get(
+                    "docker_host_astrbot_abs_path",
+                    "",
                 )
+                if self.docker_host_astrbot_abs_path:
+                    host_shared = os.path.join(
+                        self.docker_host_astrbot_abs_path,
+                        self.shared_path,
+                    )
+                    host_output = os.path.join(
+                        self.docker_host_astrbot_abs_path,
+                        output_path,
+                    )
+                    host_workplace = os.path.join(
+                        self.docker_host_astrbot_abs_path,
+                        workplace_path,
+                    )
 
-            else:
-                host_shared = os.path.abspath(self.shared_path)
-                host_output = os.path.abspath(output_path)
-                host_workplace = os.path.abspath(workplace_path)
+                else:
+                    host_shared = os.path.abspath(self.shared_path)
+                    host_output = os.path.abspath(output_path)
+                    host_workplace = os.path.abspath(workplace_path)
 
-            logger.debug(
-                f"host_shared: {host_shared}, host_output: {host_output}, host_workplace: {host_workplace}",
-            )
+                logger.debug(
+                    f"host_shared: {host_shared}, host_output: {host_output}, host_workplace: {host_workplace}",
+                )
 
-            container = await docker.containers.run(
-                {
-                    "Image": image_name,
-                    "Cmd": ["python", "exec.py"],
-                    "Memory": 512 * 1024 * 1024,
-                    "NanoCPUs": 1000000000,
-                    "HostConfig": {
-                        "Binds": [
-                            f"{host_shared}:/astrbot_sandbox/shared:ro",
-                            f"{host_output}:/astrbot_sandbox/output:rw",
-                            f"{host_workplace}:/astrbot_sandbox:rw",
-                        ],
+                container = await docker.containers.run(
+                    {
+                        "Image": image_name,
+                        "Cmd": ["python", "exec.py"],
+                        "Memory": 512 * 1024 * 1024,
+                        "NanoCPUs": 1000000000,
+                        "HostConfig": {
+                            "Binds": [
+                                f"{host_shared}:/astrbot_sandbox/shared:ro",
+                                f"{host_output}:/astrbot_sandbox/output:rw",
+                                f"{host_workplace}:/astrbot_sandbox:rw",
+                            ],
+                        },
+                        "Env": [f"MAGIC_CODE={magic_code}"],
+                        "AutoRemove": True,
                     },
-                    "Env": [f"MAGIC_CODE={magic_code}"],
-                    "AutoRemove": True,
-                },
-            )
+                )
 
-            logger.debug(f"Container {container.id} created.")
-            logs = await self.run_container(container)
-
-            logger.debug(f"Container {container.id} finished.")
-            logger.debug(f"Container {container.id} logs: {logs}")
-
-            # 发送结果
-            pattern = r"\[ASTRBOT_(TEXT|IMAGE|FILE)_OUTPUT#\w+\]: (.*)"
-            ok = False
-            traceback = ""
-            for idx, log in enumerate(logs):
-                match = re.match(pattern, log)
-                if match:
-                    ok = True
-                    if match.group(1) == "TEXT":
-                        yield event.plain_result(match.group(2))
-                    elif match.group(1) == "IMAGE":
-                        image_path = os.path.join(workplace_path, match.group(2))
-                        logger.debug(f"Sending image: {image_path}")
-                        yield event.image_result(image_path)
-                    elif match.group(1) == "FILE":
-                        file_path = os.path.join(workplace_path, match.group(2))
-                        # logger.debug(f"Sending file: {file_path}")
-                        # file_s3_url = await self.file_upload(file_path)
-                        # logger.info(f"文件上传到 AstrBot 云节点: {file_s3_url}")
-                        file_name = os.path.basename(file_path)
-                        chain: list[BaseMessageComponent] = [
-                            File(name=file_name, file=file_path)
-                        ]
-                        yield event.set_result(MessageEventResult(chain=chain))
-
-                elif "Traceback (most recent call last)" in log or "[Error]: " in log:
-                    traceback = "\n".join(logs[idx:])
-
-            if not ok:
-                if traceback:
-                    obs = f"## Observation \n When execute the code: ```python\n{code_clean}\n```\n\n Error occurred:\n\n{traceback}\n Need to improve/fix the code."
+                logger.debug(f"Container {container.id} created.")
+                logs = await self.run_container(container)
+
+                logger.debug(f"Container {container.id} finished.")
+                logger.debug(f"Container {container.id} logs: {logs}")
+
+                # 发送结果
+                pattern = r"\[ASTRBOT_(TEXT|IMAGE|FILE)_OUTPUT#\w+\]: (.*)"
+                ok = False
+                traceback = ""
+                for idx, log in enumerate(logs):
+                    match = re.match(pattern, log)
+                    if match:
+                        ok = True
+                        if match.group(1) == "TEXT":
+                            yield event.plain_result(match.group(2))
+                        elif match.group(1) == "IMAGE":
+                            image_path = os.path.join(workplace_path, match.group(2))
+                            logger.debug(f"Sending image: {image_path}")
+                            yield event.image_result(image_path)
+                        elif match.group(1) == "FILE":
+                            file_path = os.path.join(workplace_path, match.group(2))
+                            # logger.debug(f"Sending file: {file_path}")
+                            # file_s3_url = await self.file_upload(file_path)
+                            # logger.info(f"文件上传到 AstrBot 云节点: {file_s3_url}")
+                            file_name = os.path.basename(file_path)
+                            chain: list[BaseMessageComponent] = [
+                                File(name=file_name, file=file_path)
+                            ]
+                            yield event.set_result(MessageEventResult(chain=chain))
+
+                    elif (
+                        "Traceback (most recent call last)" in log or "[Error]: " in log
+                    ):
+                        traceback = "\n".join(logs[idx:])
+
+                if not ok:
+                    if traceback:
+                        obs = f"## Observation \n When execute the code: ```python\n{code_clean}\n```\n\n Error occurred:\n\n{traceback}\n Need to improve/fix the code."
+                    else:
+                        logger.warning(
+                            f"未从沙箱输出中捕获到合法的输出。沙箱输出日志: {logs}",
+                        )
+                        break
                 else:
-                    logger.warning(
-                        f"未从沙箱输出中捕获到合法的输出。沙箱输出日志: {logs}",
-                    )
-                    break
-            else:
-                # 成功了
-                self.user_file_msg_buffer.pop(event.get_session_id())
-                return
+                    # 成功了
+                    self.user_file_msg_buffer.pop(event.get_session_id())
+                    return
 
         yield event.plain_result(
             "经过多次尝试后，未从沙箱输出中捕获到合法的输出，请更换问法或者查看日志。",