[BUG] skill_exec does not apply repository.skill_run_env() #359

import asyncio
import json
import os
import shlex
import sys
import tempfile
import uuid
from pathlib import Path

from trpc_agent_sdk.agents import BaseAgent
from trpc_agent_sdk.code_executors import create_local_workspace_runtime
from trpc_agent_sdk.context import AgentContext, InvocationContext
from trpc_agent_sdk.sessions import InMemorySessionService
from trpc_agent_sdk.skills import FsSkillRepository
from trpc_agent_sdk.skills.tools import SkillRunTool, SkillExecTool
from trpc_agent_sdk.skills.tools._skill_exec import PollSessionTool, KillSessionTool

KEY = "TEST_SKILL_ENV"

class NoopAgent(BaseAgent):
    async def _run_async_impl(self, ctx):
        if False:
            yield

class EnvRepository(FsSkillRepository):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.calls = []

    def skill_run_env(self, skill_name: str) -> dict[str, str]:
        self.calls.append(skill_name)
        return {KEY: "injected"}

async def main():
    previous = os.environ.pop(KEY, None)
    service = InMemorySessionService()
    agent = NoopAgent(name="env-repro")
    results = []
    try:
        with tempfile.TemporaryDirectory(prefix="skill-env-repro-") as root:
            skill = Path(root) / "skills" / "env-demo"
            skill.mkdir(parents=True)
            (skill / "SKILL.md").write_text("---\nname: env-demo\ndescription: Environment parity reproduction\n---\nPrint a variable.\n")
            (skill / "check.py").write_text("import os\nprint(os.environ.get('TEST_SKILL_ENV'))\n")
            runtime = create_local_workspace_runtime(work_root=str(Path(root) / "workspaces"))
            repo = EnvRepository(str(skill.parent), workspace_runtime=runtime)
            run = SkillRunTool(repository=repo)
            execute = SkillExecTool(run)
            poll = PollSessionTool(execute)
            kill = KillSessionTool(execute)
            command = f"{shlex.quote(sys.executable)} check.py"
            for case, explicit, host in [("repository_only", {}, None), ("explicit_override", {KEY: "explicit"}, None), ("host_override", {}, "host")]:
                if host is None:
                    os.environ.pop(KEY, None)
                else:
                    os.environ[KEY] = host
                for tool in [run, execute]:
                    session = await service.create_session(app_name="env-repro", user_id="repro-user", session_id=str(uuid.uuid4()))
                    ctx = InvocationContext(session_service=service, invocation_id=str(uuid.uuid4()), agent=agent, agent_context=AgentContext(), session=session)
                    repo.calls.clear()
                    args = {"skill": "env-demo", "command": command, "env": explicit}
                    if tool is execute:
                        args["yield_time_ms"] = 1000
                    out = await tool.run_async(tool_context=ctx, args=args)
                    if tool is execute:
                        sid = out["session_id"]
                        try:
                            for _ in range(10):
                                if out["status"] == "exited":
                                    break
                                out = await poll.run_async(tool_context=ctx, args={"session_id": sid, "yield_time_ms": 1000})
                            assert out["status"] == "exited", out
                            result = out["result"]
                        finally:
                            await kill.run_async(tool_context=ctx, args={"session_id": sid})
                    else:
                        result = out
                    assert result["exit_code"] == 0, result
                    row = {"case": case, "tool": tool.name, "stdout": result["stdout"].strip(), "exit_code": result["exit_code"], "hook_calls": list(repo.calls)}
                    results.append(row)
                    print(json.dumps(row), flush=True)
            observed = {(r["case"], r["tool"]): r["stdout"] for r in results}
            assert observed == {
                ("repository_only", "skill_run"): "injected",
                ("repository_only", "skill_exec"): "None",
                ("explicit_override", "skill_run"): "explicit",
                ("explicit_override", "skill_exec"): "explicit",
                ("host_override", "skill_run"): "host",
                ("host_override", "skill_exec"): "host",
            }, observed
            print("REPRO_CONFIRMED", flush=True)
    finally:
        if previous is None:
            os.environ.pop(KEY, None)
        else:
            os.environ[KEY] = previous

asyncio.run(main())