Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -218,12 +218,16 @@ async def _initialize_run(
await populate_mcp_server_credentials(self._context)
self._last_creds_refresh = time.monotonic()

# If the caller changed, destroy the worker and rebuild MCP servers +
# adapter so the new ClaudeSDKClient gets fresh mcp_servers config.
# The session ID is preserved — --resume works because each SDK client
# is a new CLI subprocess that spawns fresh MCP servers from os.environ.
# Rebuild MCP servers when credentials may have changed.
# On first run: always rebuild after credential fetch so credential-based servers
# (e.g. Jira via session endpoint) that were missed during _setup_platform().
# On user change: destroy worker so the new SDK client picks up fresh creds.
user_changed = current_user_id != prev_user
if user_changed and self._session_manager.get_existing(thread_id):
if self._first_run:
self._rebuild_mcp_servers()
self._rebuild_system_prompt()
self._adapter = None
elif user_changed and self._session_manager.get_existing(thread_id):
logger.info(
f"User changed for thread={thread_id}, "
"rebuilding MCP servers and adapter with new credentials"
Expand Down Expand Up @@ -719,6 +723,21 @@ def _rebuild_mcp_servers(self) -> None:
self._allowed_tools = build_allowed_tools(self._mcp_servers)
logger.info("Rebuilt MCP servers with updated credentials")

def _rebuild_system_prompt(self) -> None:
"""Rebuild the system prompt with current env vars.

Called on first run after credential refresh so the prompt accurately
reflects which integrations are configured (e.g. Jira via session
endpoint). The initial build in _setup_platform() may have run before
credentials were fully available.
"""
from ambient_runner.bridges.claude.prompts import build_sdk_system_prompt

self._system_prompt = build_sdk_system_prompt(
self._context.workspace_path, self._cwd_path
)
logger.info("Rebuilt system prompt with updated credentials")

# ------------------------------------------------------------------
# Private: adapter lifecycle
# ------------------------------------------------------------------
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -140,6 +140,24 @@ def build_mcp_servers(
"Added credential MCP servers: %s", list(credential_mcp_servers.keys())
)

# Fallback: add Jira MCP server from env vars when credentials are available but
# no credential binding exists
jira_server_name = _CREDENTIAL_MCP_REGISTRY["jira"]["server_name"]
if (
jira_server_name not in mcp_servers
and os.getenv("JIRA_URL", "").strip()
and os.getenv("JIRA_API_TOKEN", "").strip()
):
jira_entry = _CREDENTIAL_MCP_REGISTRY["jira"]
mcp_servers[jira_server_name] = {
"command": jira_entry["command"],
"args": list(jira_entry["args"]),
"env": {k: _expand_env_vars(v) for k, v in jira_entry["env"].items()},
}
logger.info(
"Added Jira MCP server (credentials available via session endpoint)"
)

# Gerrit MCP server (only if credentials are configured)
gerrit_config = os.environ.get("GERRIT_CONFIG_PATH", "")
if gerrit_config and Path(gerrit_config).exists():
Expand Down Expand Up @@ -243,11 +261,15 @@ def _build_sidecar_mcp_servers(credential_mcp_urls_raw: str) -> dict:
try:
credential_mcp_urls = json.loads(credential_mcp_urls_raw)
except (json.JSONDecodeError, TypeError):
logger.warning("Failed to parse CREDENTIAL_MCP_URLS — skipping credential MCP servers")
logger.warning(
"Failed to parse CREDENTIAL_MCP_URLS — skipping credential MCP servers"
)
return {}

if not isinstance(credential_mcp_urls, dict):
logger.warning("CREDENTIAL_MCP_URLS is not a JSON object — skipping credential MCP servers")
logger.warning(
"CREDENTIAL_MCP_URLS is not a JSON object — skipping credential MCP servers"
)
return {}

servers: dict = {}
Expand Down Expand Up @@ -298,7 +320,11 @@ def _wait_for_sidecar_readiness(
if not endpoints:
return

logger.info("Waiting for %d credential sidecar(s) to become ready (timeout=%ds)", len(endpoints), int(timeout))
logger.info(
"Waiting for %d credential sidecar(s) to become ready (timeout=%ds)",
len(endpoints),
int(timeout),
)
deadline = time.monotonic() + timeout
pending = list(endpoints)

Expand All @@ -307,7 +333,9 @@ def _wait_for_sidecar_readiness(
for name, host, port in pending:
try:
with socket.create_connection((host, port), timeout=1.0):
logger.info("Credential sidecar %s ready at %s:%d", name, host, port)
logger.info(
"Credential sidecar %s ready at %s:%d", name, host, port
)
except (ConnectionRefusedError, OSError, socket.timeout):
still_pending.append((name, host, port))
pending = still_pending
Expand All @@ -316,7 +344,9 @@ def _wait_for_sidecar_readiness(

if pending:
names = [p[0] for p in pending]
logger.warning("Credential sidecar(s) not ready after %ds: %s", int(timeout), names)
logger.warning(
"Credential sidecar(s) not ready after %ds: %s", int(timeout), names
)


def _build_subprocess_mcp_servers() -> dict:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -117,16 +117,31 @@ def _build_system_prompt(cwd_path: str) -> str:
"""Build the full system.md content string."""
from ambient_runner.platform.config import get_repos_config, load_ambient_config
from ambient_runner.platform.prompts import (
GITHUB_TOKEN_PROMPT,
GITLAB_TOKEN_PROMPT,
GIT_PUSH_INSTRUCTIONS_BODY,
GIT_PUSH_INSTRUCTIONS_HEADER,
GIT_PUSH_MCP_STEPS,
GIT_PUSH_STEPS,
MCP_INTEGRATIONS_PROMPT,
WORKSPACE_FIXED_PATHS_PROMPT,
_build_integrations_prompt,
_detect_github,
)
from ambient_runner.platform.utils import derive_workflow_name

# Detect GitHub mode once so the git-push step selection is consistent
# with what _build_integrations_prompt() will tell the model to use.
_cmu: dict = {}
_cmu_raw = os.getenv("CREDENTIAL_MCP_URLS", "").strip()
if _cmu_raw:
try:
import json as _json

_parsed = _json.loads(_cmu_raw)
if isinstance(_parsed, dict):
_cmu = _parsed
except (ValueError, TypeError):
pass
github_mode = _detect_github(_cmu)

# Pull in Gemini's dynamically-built default sections via variable substitution.
# These are expanded at runtime by the CLI — no static text to maintain.
sections = [
Expand Down Expand Up @@ -173,7 +188,8 @@ def _build_system_prompt(cwd_path: str) -> str:
sections.append(GIT_PUSH_INSTRUCTIONS_BODY)
for r in auto_push:
sections.append(f"- **repos/{r.get('name', 'unknown')}/**")
sections.append(GIT_PUSH_STEPS.format(branch=branch))
push_steps = GIT_PUSH_MCP_STEPS if github_mode == "mcp" else GIT_PUSH_STEPS
sections.append(push_steps.format(branch=branch))

# ---- Workflow directory ----
if active_workflow_url:
Expand All @@ -200,14 +216,8 @@ def _build_system_prompt(cwd_path: str) -> str:
except Exception as exc:
logger.warning("Could not list uploaded files in %s: %s", uploads, exc)

# ---- MCP integration hints ----
sections.append(MCP_INTEGRATIONS_PROMPT)

# ---- Token visibility ----
if os.getenv("GITHUB_TOKEN"):
sections.append(GITHUB_TOKEN_PROMPT)
if os.getenv("GITLAB_TOKEN"):
sections.append(GITLAB_TOKEN_PROMPT)
# ---- Integration status — conditional on actual credential state ----
sections.append(_build_integrations_prompt())

Comment thread
squizzi marked this conversation as resolved.
# ---- Workflow custom instructions ----
ambient_config: dict = {}
Expand Down
Loading