test_claude_code_service_chain.py

# Copyright (c) 2026 OceanBase.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import asyncio
import json
import os
import shutil
import socket
import subprocess
import sys
import threading
import time
from pathlib import Path

import httpx
import pytest
import uvicorn
from fastmcp import Client
from fastmcp.client.transports import StreamableHttpTransport
from pydantic import SecretStr

from powercontext.builtin.artifacts.handoff import HandoffDraft, HandoffGenerationRequest, HandoffStatement
from powercontext.builtin.persistence.sqlite import SQLiteConfig
from powercontext.builtin.runtime import InferenceConfig
from powercontext.server.factory import create_server_app
from powercontext.server.settings import AccessControlConfig, BearerAuthConfig, McpConfig, ServerSettings
from tests.e2e.atomic_memory_models import independent_atomic_memory_model

PROJECT_ROOT = Path(__file__).resolve().parents[2]
CLAUDE_PLUGIN = PROJECT_ROOT / "integrations" / "claude-code" / "plugins" / "powercontext"
CODEX_PLUGIN = PROJECT_ROOT / "integrations" / "codex" / "plugins" / "powercontext"
AUTH_TOKEN = "claude-code-e2e-token"  # noqa: S105 - non-secret test credential.
AUTHORIZATION = f"Bearer {AUTH_TOKEN}"


class _DeterministicHandoffPipeline:
    async def generate(self, request: HandoffGenerationRequest, /) -> HandoffDraft:
        citations = tuple(item.citation for item in request.evidence)
        return HandoffDraft(
            objective=request.objective,
            state=(
                HandoffStatement(text="Claude Code MCP exposes the explicit Handoff lifecycle.", citations=citations),
            ),
            disposition="continuable",
            next_action=HandoffStatement(text="Continue from the inspected Prepared Handoff.", citations=citations),
        )


@pytest.mark.parametrize("authentication_enabled", [False, True], ids=["public", "authenticated"])
def test_claude_sessions_and_codex_share_one_project_memory(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
    authentication_enabled: bool,
) -> None:
    monkeypatch.setattr(
        "pydantic_ai.models.infer_model",
        lambda _: independent_atomic_memory_model("Use PowerContext as the shared project context service."),
    )
    app = create_server_app(
        settings=ServerSettings(
            auth=BearerAuthConfig(
                token=SecretStr(AUTH_TOKEN) if authentication_enabled else None,
            ),
            access=AccessControlConfig(mode="enforced" if authentication_enabled else "disabled"),
            database=SQLiteConfig(url=f"sqlite+aiosqlite:///{tmp_path / 'runtime.db'}"),
            inference=InferenceConfig(generation_model="test"),
            mcp=McpConfig(enabled=True),
        )
    )

    listener = socket.socket()
    listener.bind(("127.0.0.1", 0))
    host, port = listener.getsockname()
    base_url = f"http://{host}:{port}"
    server = uvicorn.Server(uvicorn.Config(app, log_level="critical", lifespan="on"))
    thread = threading.Thread(target=server.run, kwargs={"sockets": [listener]}, daemon=True)
    thread.start()
    try:
        _wait_until_started(server, thread)
        codex_plugin = tmp_path / "codex-plugin"
        shutil.copytree(CODEX_PLUGIN, codex_plugin, ignore=shutil.ignore_patterns("__pycache__", ".venv"))
        mcp_configuration = json.loads((codex_plugin / ".mcp.json").read_text())
        mcp_configuration["mcpServers"]["powercontext"]["url"] = f"{base_url}/mcp"
        (codex_plugin / ".mcp.json").write_text(json.dumps(mcp_configuration))
        scope_id = _create_scope(
            base_url,
            authorization=AUTHORIZATION if authentication_enabled else None,
        )

        captured = _run_claude_hook(
            prompt="Remember the shared project context service.",
            session_id="claude-session-a",
            prompt_id="prompt-1",
            base_url=base_url,
            authorization=AUTHORIZATION if authentication_enabled else None,
            scope_id=scope_id,
        )
        assert captured.stdout == ""
        assert AUTH_TOKEN not in captured.stderr

        recalled_by_claude = _run_claude_hook(
            prompt="Which shared project context service should we use?",
            session_id="claude-session-b",
            prompt_id="prompt-2",
            base_url=base_url,
            authorization=AUTHORIZATION if authentication_enabled else None,
            scope_id=scope_id,
        )
        claude_context = json.loads(recalled_by_claude.stdout)["hookSpecificOutput"]["additionalContext"]
        claude_envelope = json.loads(claude_context.splitlines()[-2])
        assert claude_envelope["items"][0]["content"] == ("Use PowerContext as the shared project context service.")
        assert AUTH_TOKEN not in recalled_by_claude.stderr

        recalled_by_codex = _run_codex_hook(
            codex_plugin,
            prompt="Which shared project context service should we use?",
            authorization=AUTHORIZATION if authentication_enabled else None,
            scope_id=scope_id,
        )
        codex_context = json.loads(recalled_by_codex.stdout)["hookSpecificOutput"]["additionalContext"]
        codex_envelope = json.loads(codex_context.splitlines()[-2])
        assert codex_envelope["items"][0]["content"] == claude_envelope["items"][0]["content"]
        assert AUTH_TOKEN not in recalled_by_codex.stderr
    finally:
        server.should_exit = True
        thread.join(timeout=10)
        listener.close()
        assert not thread.is_alive()


@pytest.mark.parametrize("authentication_enabled", [False, True], ids=["public", "authenticated"])
def test_claude_plugin_mcp_supports_explicit_memory_and_handoff_workflows(
    tmp_path: Path,
    authentication_enabled: bool,
) -> None:
    app = create_server_app(
        settings=ServerSettings(
            auth=BearerAuthConfig(
                token=SecretStr(AUTH_TOKEN) if authentication_enabled else None,
            ),
            access=AccessControlConfig(mode="enforced" if authentication_enabled else "disabled"),
            database=SQLiteConfig(url=f"sqlite+aiosqlite:///{tmp_path / 'mcp.db'}"),
            mcp=McpConfig(enabled=True),
        ),
        handoff_pipeline=_DeterministicHandoffPipeline(),
    )

    listener = socket.socket()
    listener.bind(("127.0.0.1", 0))
    host, port = listener.getsockname()
    base_url = f"http://{host}:{port}"
    server = uvicorn.Server(uvicorn.Config(app, log_level="critical", lifespan="on"))
    thread = threading.Thread(target=server.run, kwargs={"sockets": [listener]}, daemon=True)
    thread.start()
    try:
        _wait_until_started(server, thread)
        endpoint, headers, helper_errors = _claude_mcp_connection(
            base_url,
            authorization=AUTHORIZATION if authentication_enabled else None,
        )
        result = asyncio.run(_exercise_explicit_mcp_workflows(endpoint, headers))
    finally:
        server.should_exit = True
        thread.join(timeout=10)
        listener.close()
        assert not thread.is_alive()

    assert AUTH_TOKEN not in helper_errors
    assert result == {
        "memory_state": "forgotten",
        "memory_text": "Use the Claude Code plugin MCP transport for explicit operations.",
        "temporary_selection": "prepared",
        "committed_family": "handoff",
        "latest_matches_commit": True,
    }


def _claude_mcp_connection(base_url: str, *, authorization: str | None) -> tuple[str, dict[str, str], str]:
    configuration = json.loads((CLAUDE_PLUGIN / ".mcp.json").read_text(encoding="utf-8"))["powercontext"]
    endpoint = configuration["url"].replace("${user_config.server_url}", base_url)
    header = configuration["headers"]["Authorization"].replace(
        "${POWERCONTEXT_CLAUDE_AUTHORIZATION:-}",
        authorization or "",
    )
    return endpoint, {"Authorization": header} if header else {}, ""


async def _exercise_explicit_mcp_workflows(endpoint: str, headers: dict[str, str]) -> dict[str, object]:
    async with Client(StreamableHttpTransport(endpoint, headers=headers)) as client:
        created_scope = await client.call_tool(
            "create_scope",
            {
                "title": "Claude Code MCP workflow",
                "summary": "Explicit memory and Handoff workflow acceptance.",
                "idempotency_key": "claude-code-mcp-workflow",
            },
        )
        scope_id = (created_scope.structured_content or {})["scope_id"]
        remembered_result = await client.call_tool(
            "remember_memory",
            {
                "scope_id": scope_id,
                "kind": "decision",
                "text": "Use Claude MCP for explicit operations.",
                "reason": "Phase 2 integration verification.",
            },
        )
        remembered = remembered_result.structured_content or {}
        memory = remembered["records"][0]["artifact"]
        target = {"scope_id": scope_id, "family": memory["family"], "artifact_id": memory["artifact_id"]}
        head_result = await client.call_tool("get_artifact", target)
        head = head_result.structured_content or {}
        revised_result = await client.call_tool(
            "replace_artifact",
            {
                **target,
                "If-Match": head["etag"],
                "content": {
                    "kind": "decision",
                    "text": "Use the Claude Code plugin MCP transport for explicit operations.",
                },
            },
        )
        revised = revised_result.structured_content or {}
        state_result = await client.call_tool(
            "get_atomic_memory_state", {"scope_id": scope_id, "artifact_id": memory["artifact_id"]}
        )
        state = state_result.structured_content or {}
        assert state["artifact"]["revision"] == revised["artifact"]["revision"]
        retired_result = await client.call_tool(
            "change_atomic_memory_lifecycle",
            {
                "scope_id": scope_id,
                "target": {"artifact": state["artifact"], "state_version": state["state_version"]},
                "state": "forgotten",
            },
        )
        retired = retired_result.structured_content or {}

        captured_result = await client.call_tool(
            "capture_content_source",
            {
                "scope_id": scope_id,
                "source_id": "claude-mcp-handoff-boundary",
                "content": "The Claude Code MCP integration completed its explicit workflow checks.",
            },
        )
        captured = captured_result.structured_content or {}
        activation_result = await client.call_tool(
            "activate_handoff",
            {
                "scope_id": scope_id,
                "boundary_source": captured["source"],
                "objective": "Transfer the verified Claude Code MCP integration state.",
            },
        )
        activation = activation_result.structured_content or {}
        prepared_result = await client.call_tool(
            "finalize_handoff",
            {"scope_id": scope_id, "draft": activation["draft"]},
        )
        prepared = prepared_result.structured_content or {}
        temporary_result = await client.call_tool(
            "continue_handoff",
            {"scope_id": scope_id, "selection": "prepared", "prepared": prepared},
        )
        temporary = temporary_result.structured_content or {}
        committed_result = await client.call_tool(
            "commit_handoff",
            {"scope_id": scope_id, "handoff": prepared},
        )
        committed = committed_result.structured_content or {}
        latest_result = await client.call_tool(
            "continue_handoff",
            {"scope_id": scope_id, "selection": "latest"},
        )
        latest = latest_result.structured_content or {}

    return {
        "memory_state": retired["records"][0]["state"],
        "memory_text": retired["records"][0]["text"],
        "temporary_selection": temporary["selection"],
        "committed_family": committed["reference"]["family"],
        "latest_matches_commit": latest["selected_revision"] == committed["reference"],
    }


def _run_claude_hook(
    *,
    prompt: str,
    session_id: str,
    prompt_id: str,
    base_url: str,
    authorization: str | None,
    scope_id: str,
) -> subprocess.CompletedProcess[str]:
    environment: dict[str, str] = {
        **os.environ,
        "POWERCONTEXT_CLAUDE_SERVER_URL": base_url,
        "POWERCONTEXT_CLAUDE_FLUSH_ON_CAPTURE": "true",
        "POWERCONTEXT_CLAUDE_HTTP_BUDGET_SECONDS": "10",
        "POWERCONTEXT_CLAUDE_REQUEST_TIMEOUT_SECONDS": "5",
        "POWERCONTEXT_CLAUDE_SCOPE_ID": scope_id,
    }
    environment.pop("POWERCONTEXT_CLAUDE_AUTHORIZATION", None)
    if authorization is not None:
        environment["POWERCONTEXT_CLAUDE_AUTHORIZATION"] = authorization
    return subprocess.run(
        [sys.executable, str(CLAUDE_PLUGIN / "hooks" / "user_prompt_submit.py")],
        cwd=PROJECT_ROOT,
        env=environment,
        input=json.dumps({
            "hook_event_name": "UserPromptSubmit",
            "cwd": str(PROJECT_ROOT),
            "prompt": prompt,
            "session_id": session_id,
            "prompt_id": prompt_id,
        }),
        text=True,
        capture_output=True,
        check=True,
        timeout=15,
    )


def _run_codex_hook(
    plugin: Path,
    *,
    prompt: str,
    authorization: str | None,
    scope_id: str,
) -> subprocess.CompletedProcess[str]:
    environment: dict[str, str] = {
        **os.environ,
        "POWERCONTEXT_CODEX_SCOPE_ID": scope_id,
        "POWERCONTEXT_CODEX_CAPTURE_PROMPTS": "false",
        "POWERCONTEXT_CODEX_HTTP_BUDGET_SECONDS": "10",
        "POWERCONTEXT_CODEX_REQUEST_TIMEOUT_SECONDS": "5",
    }
    environment.pop("POWERCONTEXT_CODEX_AUTHORIZATION", None)
    if authorization is not None:
        environment["POWERCONTEXT_CODEX_AUTHORIZATION"] = authorization
    return subprocess.run(
        [sys.executable, str(plugin / "hooks" / "recall.py")],
        cwd=PROJECT_ROOT,
        env=environment,
        input=json.dumps({
            "hook_event_name": "UserPromptSubmit",
            "cwd": str(PROJECT_ROOT),
            "prompt": prompt,
            "session_id": "codex-session",
            "turn_id": "turn-1",
        }),
        text=True,
        capture_output=True,
        check=True,
        timeout=15,
    )


def _create_scope(base_url: str, *, authorization: str | None) -> str:
    headers = {"Authorization": authorization} if authorization is not None else None
    response = httpx.post(
        f"{base_url}/v1/scopes",
        headers=headers,
        json={
            "title": "Shared agent work",
            "summary": "Shared memory for the Claude Code and Codex sessions.",
            "idempotency_key": "shared-agent-work",
        },
        timeout=5,
    )
    response.raise_for_status()
    return response.json()["scope_id"]


def _wait_until_started(server: uvicorn.Server, thread: threading.Thread) -> None:
    deadline = time.monotonic() + 10
    while thread.is_alive() and not server.started and time.monotonic() < deadline:
        time.sleep(0.01)
    assert server.started