Compare commits

..

2 Commits

Author SHA1 Message Date
Molecule AI Dev Engineer A (Kimi) 8b25aec245 Merge remote-tracking branch 'origin/main' into pr-1466
ci-arm64-advisory / fast-checks (pull_request) Waiting to run
Lint shellcheck (arm64 pilot) / shellcheck-arm64 (pilot) (pull_request) Successful in 8s
CI / Python Lint & Test (pull_request) Successful in 8s
CI / Detect changes (pull_request) Successful in 11s
E2E API Smoke Test / detect-changes (pull_request) Successful in 10s
Block internal-flavored paths / Block forbidden paths (pull_request) Successful in 15s
E2E Chat / detect-changes (pull_request) Successful in 10s
Handlers Postgres Integration / detect-changes (pull_request) Successful in 10s
Lint curl status-code capture / Scan workflows for curl status-capture pollution (pull_request) Successful in 8s
Lint forbidden tenant-env keys / Scan workspace_secrets writers for forbidden env keys (pull_request) Successful in 8s
Lint no tenant GITEA or GITHUB token write / Scan for repo-host token write into tenant workspace surface (pull_request) Successful in 5s
Harness Replays / detect-changes (pull_request) Successful in 25s
lint-continue-on-error-tracking / lint-continue-on-error-tracking (pull_request) Successful in 1m20s
lint-required-workflows-docker-host-pinned / Lint docker-host pin on docker-touching workflows (pull_request) Successful in 3s
Check migration collisions / Migration version collision check (pull_request) Successful in 1m42s
lint-required-no-paths / lint-required-no-paths (pull_request) Successful in 1m13s
review-check-tests / review-check.sh regression tests (pull_request) Successful in 11s
Lint pre-flip continue-on-error / Verify continue-on-error flips have run-log proof (pull_request) Successful in 1m25s
lint-required-context-exists-in-bp / lint-required-context-exists-in-bp (pull_request) Failing after 1m26s
lint-mask-pr-atomicity / lint-mask-pr-atomicity (pull_request) Successful in 1m43s
gate-check-v3 / gate-check (pull_request) Successful in 16s
qa-review / approved (pull_request) Successful in 12s
sop-checklist / na-declarations (pull_request) N/A: (none)
security-review / approved (pull_request) Successful in 9s
sop-checklist / all-items-acked (pull_request) Successful in 5s
sop-tier-check / tier-check (pull_request) Successful in 6s
Secret scan / Scan diff for credential-shaped strings (pull_request) Successful in 33s
CI / Shellcheck (E2E scripts) (pull_request) Successful in 21s
Lint workflow YAML (Gitea-1.22.6-hostile shapes) / Lint workflow YAML for Gitea-1.22.6-hostile shapes (pull_request) Successful in 1m44s
Ops Scripts Tests / Ops scripts (unittest) (pull_request) Successful in 1m30s
E2E API Smoke Test / E2E API Smoke Test (pull_request) Successful in 1m57s
Harness Replays / Harness Replays (pull_request) Successful in 9s
Handlers Postgres Integration / Handlers Postgres Integration (pull_request) Successful in 2m45s
E2E Chat / E2E Chat (pull_request) Successful in 4m53s
CI / Platform (Go) (pull_request) Successful in 6m11s
CI / Canvas (Next.js) (pull_request) Successful in 7m11s
CI / all-required (pull_request) Successful in 9m38s
CI / Canvas Deploy Reminder (pull_request) Has been skipped
2026-05-26 10:52:13 +00:00
fullstack-engineer 3ba08a2dc8 test(canvas): add lib test coverage for design-tokens, palette-context, theme-provider
Block internal-flavored paths / Block forbidden paths (pull_request) Successful in 13s
CI / Detect changes (pull_request) Successful in 12s
Handlers Postgres Integration / detect-changes (pull_request) Successful in 8s
Harness Replays / detect-changes (pull_request) Successful in 5s
Runtime PR-Built Compatibility / detect-changes (pull_request) Successful in 6s
Secret scan / Scan diff for credential-shaped strings (pull_request) Successful in 6s
gate-check-v3 / gate-check (pull_request) Successful in 5s
qa-review / approved (pull_request) Successful in 3s
security-review / approved (pull_request) Successful in 3s
sop-tier-check / tier-check (pull_request) Successful in 3s
lint-required-no-paths / lint-required-no-paths (pull_request) Successful in 59s
Harness Replays / Harness Replays (pull_request) Successful in 1s
Runtime PR-Built Compatibility / PR-built wheel + import smoke (pull_request) Successful in 1s
sop-checklist / all-items-acked (pull_request) [info tier:low] acked: 0/7 — missing: comprehensive-testing, local-postgres-e2e, staging-smoke, +4 — body-unfilled: comprehensive-testing, l
sop-checklist / na-declarations (pull_request) N/A: (none)
CI / Platform (Go) (pull_request) Has been cancelled
CI / Canvas Deploy Reminder (pull_request) Has been cancelled
CI / all-required (pull_request) Has been cancelled
CI / Shellcheck (E2E scripts) (pull_request) Has been cancelled
CI / Python Lint & Test (pull_request) Has been cancelled
E2E API Smoke Test / E2E API Smoke Test (pull_request) Has been cancelled
E2E Chat / E2E Chat (pull_request) Has been cancelled
E2E API Smoke Test / detect-changes (pull_request) Has been cancelled
Handlers Postgres Integration / Handlers Postgres Integration (pull_request) Has been cancelled
CI / Canvas (Next.js) (pull_request) Has been cancelled
E2E Chat / detect-changes (pull_request) Has been cancelled
design-tokens.test.ts:
  - STATUS_CONFIG: all 7 statuses have dot/label/bar
  - statusDotClass: known status returns dot, unknown/empty → bg-zinc-500
  - TIER_CONFIG: tiers 1-4 have label/color/border, T4 uses warm
  - COMM_TYPE_LABELS: a2a_send→sent, a2a_receive→received, task_update

palette-context.test.tsx:
  - normalizeStatus: online/degraded→emerald, failed→red, paused/not_configured→amber, unknown→zinc
  - tierCode: maps 1-4 to T1-T4
  - getPalette: null→base, identity guard, custom accent overrides, no mutation of MOL_LIGHT/MOL_DARK

theme-provider.test.tsx:
  - applyResolvedTheme: sets data-theme on html element
  - ThemeProvider: is a function (React component)
  - THEME_COOKIE = 'mol_theme', themeBootScript is a non-empty string

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-05-18 01:16:12 +00:00
99 changed files with 1079 additions and 6156 deletions
+27 -247
View File
@@ -90,15 +90,6 @@ API = f"https://{GITEA_HOST}/api/v1" if GITEA_HOST else ""
# match by exact title without parsing.
TITLE_PREFIX = "[main-red]"
# Contexts that are scheduled or non-required — their pending/failure
# state should not block stale-issue closeout (mc#1789).
SCHEDULED_CONTEXT_PATTERNS = (
"Staging SaaS smoke",
"Continuous synthetic E2E",
"main-red-watchdog",
"ci-arm64-advisory",
)
# Settling window (seconds) between initial red detection and the
# pre-file recheck. The recheck filters out the two largest false-
# positive classes seen in mc#1597..1630 (task #394, 2026-05-21):
@@ -274,11 +265,6 @@ def get_combined_status(sha: str) -> dict:
return body
def _entry_state(s: dict) -> str:
"""Per-entry status key in Gitea 1.22.6 is `status`; fall back to `state`."""
return s.get("status") or s.get("state") or ""
def is_red(status: dict) -> tuple[bool, list[dict]]:
"""Return (is_red, failed_statuses).
@@ -326,6 +312,9 @@ def is_red(status: dict) -> tuple[bool, list[dict]]:
# "no per-context entries were in a red state" fallback even when
# the combined-state correctly flagged red. See
# `feedback_smoke_test_vendor_truth_not_shape_match`.
def _entry_state(s: dict) -> str:
return s.get("status") or s.get("state") or ""
def _is_cancel_cascade(s: dict) -> bool:
"""status=3 entry per Gitea 1.22.6 description-string contract.
Match exactly (after strip) — substring match would catch
@@ -364,15 +353,6 @@ def title_for(sha: str) -> str:
return f"{TITLE_PREFIX} {REPO}: {sha[:10]}"
def _is_scheduled_context(context: str) -> bool:
"""Return True if `context` is a known scheduled/non-required job.
These contexts run on a schedule and should not block stale-issue
closeout when main's required CI has recovered (mc#1789).
"""
return any(pattern.lower() in context.lower() for pattern in SCHEDULED_CONTEXT_PATTERNS)
def list_open_red_issues() -> list[dict]:
"""All open issues whose title starts with `[main-red] {repo}: `.
@@ -382,34 +362,23 @@ def list_open_red_issues() -> list[dict]:
file-or-update path to POST a duplicate — exactly the regression
class the helper-raises contract closes.
Pagination is exhausted (mc#1789). The old "by design ≤ 1" invariant
was false — backlog can exceed 50 open issues.
Gitea issue search returns at most 50/page; we only need open
`[main-red]` issues which are by design ≤ 1 at any time per repo,
so a single page is enough.
"""
prefix = f"{TITLE_PREFIX} {REPO}: "
all_issues: list[dict] = []
page = 1
limit = 50
while True:
_, results = api(
"GET",
f"/repos/{OWNER}/{NAME}/issues",
query={"state": "open", "type": "issues", "limit": str(limit), "page": str(page)},
_, results = api(
"GET",
f"/repos/{OWNER}/{NAME}/issues",
query={"state": "open", "type": "issues", "limit": "50"},
)
if not isinstance(results, list):
raise ApiError(
f"issue search returned non-list body (got {type(results).__name__})"
)
if not isinstance(results, list):
raise ApiError(
f"issue search returned non-list body (got {type(results).__name__})"
)
matched = [
i for i in results
if isinstance(i, dict)
prefix = f"{TITLE_PREFIX} {REPO}: "
return [i for i in results if isinstance(i, dict)
and isinstance(i.get("title"), str)
and i["title"].startswith(prefix)
]
all_issues.extend(matched)
if len(results) < limit:
break
page += 1
return all_issues
and i["title"].startswith(prefix)]
def find_open_issue_for_sha(sha: str) -> dict | None:
@@ -605,151 +574,6 @@ def file_or_update_red(
sys.stderr.write(f"::warning::label '{RED_LABEL}' not found on repo\n")
def close_stale_red_issues(
current_sha: str,
current_status: dict,
*,
dry_run: bool = False,
) -> int:
"""Close open [main-red] issues whose specific failing contexts have
all recovered on `current_sha`, even though `main` is still red for
other reasons (mc#1789).
When main stays red across consecutive SHAs for *different* causes,
`close_open_red_issues_for_other_shas` never fires (it only runs when
main is green). This function prevents stale issues from accumulating
indefinitely by comparing per-context recovery across SHAs.
An issue is considered stale when every context that was in a failed
state on the issue's SHA is now either `success` on the current HEAD
or absent (workflow removed / renamed). Issues whose original SHA had
a combined-red-with-no-detail (empty statuses list) are skipped — we
cannot verify recovery without per-context data.
Returns the number of issues closed.
"""
open_red = list_open_red_issues()
if not open_red:
return 0
current_statuses = current_status.get("statuses") or []
closed = 0
for issue in open_red:
title = issue.get("title", "")
prefix = f"{TITLE_PREFIX} {REPO}: "
if not title.startswith(prefix):
continue
short_sha = title[len(prefix):]
if short_sha == current_sha[:10]:
continue
# Query status for the old SHA. Short SHA should resolve; if it
# doesn't (GC'd, force-pushed, ambiguous), skip conservatively.
try:
old_status = get_combined_status(short_sha)
except ApiError:
continue
old_red, old_failed = is_red(old_status)
if not old_red:
# Open issue for a now-green SHA — close it via the normal path.
num = issue.get("number")
if isinstance(num, int):
comment = (
f"Commit `{short_sha}` is no longer red. Closing as the "
f"failure context has recovered or expired."
)
if dry_run:
print(
f"::notice::[dry-run] would close issue #{num} "
f"({title}) — old SHA is now green"
)
closed += 1
continue
api(
"POST",
f"/repos/{OWNER}/{NAME}/issues/{num}/comments",
body={"body": comment},
)
api(
"PATCH",
f"/repos/{OWNER}/{NAME}/issues/{num}",
body={"state": "closed"},
)
print(
f"::notice::Closed stale main-red issue #{num} "
f"(old SHA {short_sha} is now green)"
)
closed += 1
continue
if not old_failed:
# Combined red with no per-context detail — can't verify recovery.
continue
# Verify every failed context from the old SHA has recovered.
all_recovered = True
recovered_ctxs: list[str] = []
still_failing_ctxs: list[str] = []
for s in old_failed:
ctx = s.get("context", "")
if not ctx:
continue
current_match = None
for cs in current_statuses:
if isinstance(cs, dict) and cs.get("context") == ctx:
current_match = cs
break
if current_match is None:
recovered_ctxs.append(ctx)
elif _entry_state(current_match) == "success":
recovered_ctxs.append(ctx)
else:
all_recovered = False
still_failing_ctxs.append(ctx)
if not all_recovered:
continue
num = issue.get("number")
if not isinstance(num, int):
continue
comment = (
f"The failing contexts from this SHA (`{short_sha}`) have "
f"recovered on current HEAD `{current_sha[:10]}`: "
f"{', '.join(recovered_ctxs)}. "
f"Main is still red for other reasons; see the current "
f"`[main-red]` issue for `{current_sha[:10]}`."
)
if dry_run:
print(
f"::notice::[dry-run] would close stale issue #{num} "
f"({title}) — contexts recovered"
)
closed += 1
continue
api(
"POST",
f"/repos/{OWNER}/{NAME}/issues/{num}/comments",
body={"body": comment},
)
api(
"PATCH",
f"/repos/{OWNER}/{NAME}/issues/{num}",
body={"state": "closed"},
)
print(
f"::notice::Closed stale main-red issue #{num} "
f"(contexts recovered at {current_sha[:10]})"
)
closed += 1
return closed
def close_open_red_issues_for_other_shas(
current_sha: str,
*,
@@ -920,68 +744,24 @@ def run_once(*, dry_run: bool = False) -> int:
print(f"::warning::main is RED at {sha[:10]} on {WATCH_BRANCH}: "
f"{len(failed)} failed context(s)")
file_or_update_red(sha, failed, debug, dry_run=dry_run)
stale_closed = close_stale_red_issues(sha, recheck_status, dry_run=dry_run)
if stale_closed:
emit_loki_event("main_red_stale_closed", sha, [])
print(
f"::notice::Closed {stale_closed} stale main-red issue(s) "
f"whose contexts recovered at {sha[:10]}"
)
else:
# Green or pending-with-no-real-failures. Close stale issues
# from earlier SHAs when required CI has recovered.
#
# mc#1789: main often sits at combined `pending` because
# scheduled/non-required contexts (Staging SaaS smoke,
# Continuous synthetic E2E, main-red-watchdog itself,
# ci-arm64-advisory) are still running. We close stale issues
# as long as no *non-scheduled* context has failed and no
# *non-scheduled* context is still pending — i.e. required CI
# is effectively green.
#
# The success-only gate is preserved for the canonical green
# path; the extended check below only fires when combined is
# `pending` but all required work is done.
combined_state = status.get("state")
if combined_state == "success":
should_close = True
close_reason = "GREEN"
else:
statuses = status.get("statuses") or []
non_scheduled_pending = [
s for s in statuses
if isinstance(s, dict)
and (_entry_state(s) == "pending")
and not _is_scheduled_context(s.get("context", ""))
]
non_scheduled_failed = [
s for s in statuses
if isinstance(s, dict)
and (_entry_state(s) in {"failure", "error"})
and not _is_scheduled_context(s.get("context", ""))
]
# Cancel-cascade already filtered by is_red(); red=False
# here means no real failures. We additionally check that
# no non-scheduled context is still pending.
should_close = not non_scheduled_pending and not non_scheduled_failed
close_reason = "pending-but-required-green"
if should_close:
# Green (or pending — pending is treated as not-red so we don't
# spam during the post-merge CI window). Close any stale issues
# from earlier SHAs only when we're actually green; pending
# means CI hasn't finished and the prior issue might still be
# accurate.
if status.get("state") == "success":
closed = close_open_red_issues_for_other_shas(sha, dry_run=dry_run)
if closed:
emit_loki_event(
"main_returned_to_green", sha,
[],
)
print(
f"::notice::main is {close_reason} at {sha[:10]} on {WATCH_BRANCH} "
f"(closed {closed} stale issue(s))"
)
print(f"::notice::main is GREEN at {sha[:10]} on {WATCH_BRANCH} "
f"(closed {closed} stale issue(s))")
else:
print(
f"::notice::main has pending-or-failed required CI at {sha[:10]} "
f"on {WATCH_BRANCH} (combined state={combined_state!r}; no action)"
)
print(f"::notice::main is PENDING at {sha[:10]} on {WATCH_BRANCH} "
f"(combined state={status.get('state')!r}; no action)")
return 0
+4 -4
View File
@@ -642,7 +642,7 @@ def load_config(path: str) -> dict[str, Any]:
# requiring the dep, so the ignore is safe: if yaml loads, we use it;
# otherwise we fall back silently.
import yaml # type: ignore[import-not-found]
with open(path, encoding="utf-8") as f:
with open(path) as f:
return yaml.safe_load(f)
except ImportError:
return _load_config_minimal(path)
@@ -656,7 +656,7 @@ def _load_config_minimal(path: str) -> dict[str, Any]:
item map: scalars + lists of scalars. Does NOT support nested lists,
YAML anchors, multi-doc, or flow style.
"""
with open(path, encoding="utf-8") as f:
with open(path) as f:
lines = f.readlines()
return _parse_minimal_yaml(lines)
@@ -845,7 +845,7 @@ def render_status(
if len(missing_body) > 3:
shown += f", +{len(missing_body) - 3}"
desc_parts.append(f"body-unfilled: {shown}")
state = "success" if not missing else "failure"
state = "success" if not missing and not missing_body else "failure"
return state, "".join(desc_parts)
@@ -1033,7 +1033,7 @@ def main(argv: list[str] | None = None) -> int:
for t in data:
if t.get("name") == tn:
tid = t.get("id")
client._team_id_cache[(args.owner, tn)] = tid # noqa: SLF001 # write-through cache; intentional side-effect for reuse across calls
client._team_id_cache[(args.owner, tn)] = tid # noqa: SLF001 # internal write-through cache
break
if tid is not None:
team_ids.append(tid)
+1 -1
View File
@@ -33,7 +33,7 @@ def scenario() -> str:
p = os.path.join(STATE_DIR, "scenario")
if not os.path.isfile(p):
return "T1_success"
with open(p, encoding="utf-8") as f:
with open(p) as f:
return f.read().strip()
@@ -40,7 +40,7 @@ def scenario() -> str:
p = os.path.join(STATE_DIR, "scenario")
if not os.path.isfile(p):
return "T1_pr_open"
with open(p, encoding="utf-8") as f:
with open(p) as f:
return f.read().strip()
@@ -1,283 +0,0 @@
import importlib.util
import sys
from pathlib import Path
from unittest.mock import patch, MagicMock
SCRIPT = Path(__file__).resolve().parents[1] / "main-red-watchdog.py"
spec = importlib.util.spec_from_file_location("main_red_watchdog", SCRIPT)
wd = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = wd
spec.loader.exec_module(wd)
# Module-level constants are loaded from env at import time; set them
# explicitly so unit tests can import without the full env contract.
wd.GITEA_TOKEN = "fake-token"
wd.GITEA_HOST = "git.example.com"
wd.REPO = "molecule-ai/molecule-core"
wd.OWNER = "molecule-ai"
wd.NAME = "molecule-core"
wd.WATCH_BRANCH = "main"
wd.RED_LABEL = "tier:high"
wd.API = "https://git.example.com/api/v1"
# ---------------------------------------------------------------------------
# _is_scheduled_context
# ---------------------------------------------------------------------------
def test_is_scheduled_context_matches_staging_saas_smoke():
assert wd._is_scheduled_context("Staging SaaS smoke") is True
def test_is_scheduled_context_matches_case_insensitive():
assert wd._is_scheduled_context("continuous synthetic e2e") is True
def test_is_scheduled_context_no_match_for_required_ci():
assert wd._is_scheduled_context("CI / all-required") is False
# ---------------------------------------------------------------------------
# _entry_state
# ---------------------------------------------------------------------------
def test_entry_state_prefers_status_over_state():
"""Gitea 1.22.6 per-entry key is `status`; `state` is fallback."""
assert wd._entry_state({"status": "failure", "state": "success"}) == "failure"
def test_entry_state_falls_back_to_state():
assert wd._entry_state({"state": "pending"}) == "pending"
def test_entry_state_empty_when_neither_key_present():
assert wd._entry_state({"context": "foo"}) == ""
# ---------------------------------------------------------------------------
# is_red
# ---------------------------------------------------------------------------
def test_is_red_combined_failure_no_statuses():
"""Combined failure with empty statuses[] still trips red."""
red, failed = wd.is_red({"state": "failure", "statuses": []})
assert red is True
assert failed == []
def test_is_red_cancel_cascade_filtered():
"""status=3 (cancelled) mapped to failure string must be filtered."""
status = {
"state": "failure",
"statuses": [
{"context": "CI / build", "status": "failure", "description": "Has been cancelled"},
],
}
red, failed = wd.is_red(status)
assert red is False
assert failed == []
def test_is_red_real_failure_not_filtered():
"""Real failures with different descriptions are kept."""
status = {
"state": "failure",
"statuses": [
{"context": "CI / build", "status": "failure", "description": "Failing after 12s"},
],
}
red, failed = wd.is_red(status)
assert red is True
assert len(failed) == 1
assert failed[0]["context"] == "CI / build"
def test_is_red_uses_entry_state_not_top_level_state():
"""Regression: per-entry key is `status`, not `state`."""
status = {
"state": "failure",
"statuses": [
# Only `status` present; pre-rev4 code read `state` and got None
{"context": "CI / test", "status": "failure"},
],
}
red, failed = wd.is_red(status)
assert red is True
assert len(failed) == 1
# ---------------------------------------------------------------------------
# list_open_red_issues — pagination (mc#1789)
# ---------------------------------------------------------------------------
def test_list_open_red_issues_exhausts_pagination():
"""Backlog can exceed 50 issues; all pages must be fetched."""
calls = []
def fake_api(method, path, **kwargs):
calls.append((method, path, kwargs))
query = (kwargs.get("query") or {})
page = int(query.get("page", "1"))
limit = int(query.get("limit", "50"))
# Page 1 returns full limit; page 2 returns partial → break
if page == 1:
return 200, [
{"title": f"[main-red] molecule-ai/molecule-core: sha{i:04d}"}
for i in range(limit)
]
if page == 2:
return 200, [
{"title": "[main-red] molecule-ai/molecule-core: extra1"},
{"title": "[main-red] molecule-ai/molecule-core: extra2"},
{"title": " unrelated issue "}, # filtered out
]
return 200, []
with patch.object(wd, "api", side_effect=fake_api):
issues = wd.list_open_red_issues()
assert len(issues) == 52 # 50 + 2 matched
titles = {i["title"] for i in issues}
assert "[main-red] molecule-ai/molecule-core: extra1" in titles
assert "[main-red] molecule-ai/molecule-core: extra2" in titles
def test_list_open_red_issues_single_page():
"""When results < limit, loop breaks after first page."""
def fake_api(method, path, **kwargs):
return 200, [
{"title": "[main-red] molecule-ai/molecule-core: abc123"},
]
with patch.object(wd, "api", side_effect=fake_api):
issues = wd.list_open_red_issues()
assert len(issues) == 1
# ---------------------------------------------------------------------------
# run_once — close logic (mc#1789)
# ---------------------------------------------------------------------------
def test_run_once_green_closes_stale_issues(monkeypatch):
"""Combined success → close stale issues."""
monkeypatch.setattr(wd, "get_head_sha", lambda b: "abc123")
monkeypatch.setattr(wd, "get_combined_status", lambda s: {"state": "success", "statuses": []})
monkeypatch.setattr(wd, "is_red", lambda s: (False, []))
closed = []
def capture_close(current_sha, *, dry_run=False, close_same_sha=False):
closed.append(current_sha)
return 1
monkeypatch.setattr(wd, "close_open_red_issues_for_other_shas", capture_close)
monkeypatch.setattr(wd, "emit_loki_event", lambda *a, **k: None)
assert wd.run_once(dry_run=True) == 0
assert closed == ["abc123"]
def test_run_once_pending_scheduled_only_closes_stale_issues(monkeypatch):
"""Combined pending, but only scheduled contexts pending → close stale."""
monkeypatch.setattr(wd, "get_head_sha", lambda b: "abc123")
monkeypatch.setattr(
wd, "get_combined_status",
lambda s: {
"state": "pending",
"statuses": [
{"context": "CI / all-required", "status": "success"},
{"context": "Staging SaaS smoke", "status": "pending"},
],
}
)
monkeypatch.setattr(wd, "is_red", lambda s: (False, []))
closed = []
def capture_close(current_sha, *, dry_run=False, close_same_sha=False):
closed.append(current_sha)
return 1
monkeypatch.setattr(wd, "close_open_red_issues_for_other_shas", capture_close)
monkeypatch.setattr(wd, "emit_loki_event", lambda *a, **k: None)
assert wd.run_once(dry_run=True) == 0
assert closed == ["abc123"]
def test_run_once_pending_required_does_not_close(monkeypatch):
"""Combined pending with a real required context still pending → no close."""
monkeypatch.setattr(wd, "get_head_sha", lambda b: "abc123")
monkeypatch.setattr(
wd, "get_combined_status",
lambda s: {
"state": "pending",
"statuses": [
{"context": "CI / all-required", "status": "pending"},
{"context": "Staging SaaS smoke", "status": "success"},
],
}
)
monkeypatch.setattr(wd, "is_red", lambda s: (False, []))
closed = []
def capture_close(current_sha, *, dry_run=False, close_same_sha=False):
closed.append(current_sha)
return 0
monkeypatch.setattr(wd, "close_open_red_issues_for_other_shas", capture_close)
monkeypatch.setattr(wd, "emit_loki_event", lambda *a, **k: None)
assert wd.run_once(dry_run=True) == 0
assert closed == []
def test_run_once_failure_does_not_close(monkeypatch):
"""Real failure in non-scheduled context → no close."""
monkeypatch.setattr(wd, "get_head_sha", lambda b: "abc123")
monkeypatch.setattr(
wd, "get_combined_status",
lambda s: {
"state": "failure",
"statuses": [
{"context": "CI / all-required", "status": "failure"},
],
}
)
# is_red will return True, so we enter the red path, not the green close path
monkeypatch.setattr(wd, "is_red", lambda s: (True, s.get("statuses", [])))
monkeypatch.setattr(wd, "time", MagicMock(sleep=lambda x: None))
monkeypatch.setattr(wd, "emit_loki_event", lambda *a, **k: None)
filed = []
def capture_file(sha, failed, debug, *, dry_run=False):
filed.append(sha)
monkeypatch.setattr(wd, "file_or_update_red", capture_file)
monkeypatch.setattr(wd, "close_open_red_issues_for_other_shas", lambda *a, **k: 0)
monkeypatch.setattr(wd, "close_stale_red_issues", lambda *a, **k: 0)
assert wd.run_once(dry_run=True) == 0
assert filed == ["abc123"]
# ---------------------------------------------------------------------------
# title_for / find_open_issue_for_sha
# ---------------------------------------------------------------------------
def test_title_for_uses_short_sha():
assert wd.title_for("abcdef123456") == "[main-red] molecule-ai/molecule-core: abcdef1234"
def test_find_open_issue_for_sha_matches_exact_title(monkeypatch):
fake_issue = {"title": "[main-red] molecule-ai/molecule-core: abc1234567", "number": 42}
monkeypatch.setattr(wd, "list_open_red_issues", lambda: [fake_issue])
assert wd.find_open_issue_for_sha("abc1234567") == fake_issue
def test_find_open_issue_for_sha_returns_none_when_no_match(monkeypatch):
monkeypatch.setattr(wd, "list_open_red_issues", lambda: [])
assert wd.find_open_issue_for_sha("abc123") is None
+1 -1
View File
@@ -138,7 +138,7 @@ items:
- slug: memory-consulted
numeric_alias: 7
pr_section_marker: "Memory consulted"
pr_section_marker: "Memory/saved-feedback consulted"
required_teams: [engineers]
description: >-
List of feedback memories applicable to this change. Ack from
+1 -2
View File
@@ -54,6 +54,5 @@ jobs:
# read-only by design (least-privilege).
REQUIRED_CHECKS: |
CI / all-required (pull_request)
E2E API Smoke Test / E2E API Smoke Test (pull_request)
Handlers Postgres Integration / Handlers Postgres Integration (pull_request)
sop-checklist / all-items-acked (pull_request)
run: bash .gitea/scripts/audit-force-merge.sh
+6 -14
View File
@@ -164,20 +164,12 @@ jobs:
# mc#774: pre-existing continue-on-error mask; root-fix and remove, do not renew silently.
continue-on-error: true
- if: ${{ needs.changes.outputs.platform == 'true' }}
name: Run tests with coverage (blocking gate)
# Removed -race from the blocking gate per #1184: cold runners
# take 13-25 min to compile with race instrumentation, exceeding
# the 10m step timeout and causing false failures. Race detection
# now runs as a non-blocking advisory step below.
run: go test -timeout 10m -coverprofile=coverage.out ./...
- if: ${{ needs.changes.outputs.platform == 'true' }}
name: Race detection (advisory, non-blocking)
# mc#1184: runs race detector as an advisory check so cold-runner
# compile-time spikes don't block merges. Failures here surface in
# the run log but do not fail the build.
run: go test -race -timeout 10m ./...
continue-on-error: true
name: Run tests with race detection and coverage
# Explicit timeout: cold runner cache causes OOM kills at ~4m39s on the
# full ./... suite with race detection + coverage. A 10m per-step timeout
# lets the suite complete on cold cache (~5-7m) while failing cleanly
# instead of OOM-killing. The job-level timeout (15m) is a backstop.
run: go test -race -timeout 10m -coverprofile=coverage.out ./...
- if: ${{ needs.changes.outputs.platform == 'true' }}
name: Per-file coverage report
@@ -3,26 +3,11 @@ name: Lint shellcheck (arm64 pilot)
# Mac-CI dual-track pilot (#233). ADDITIVE / NOT REQUIRED.
#
# Validates the arm64 self-hosted lane (no docker.sock, no privileged
# ops) before any required gate moves onto it.
# ops) before any required gate moves onto it. Until a Mac arm64 runner
# is registered with the `arm64` label, this workflow sits PENDING —
# that is FINE: `arm64` is NOT in branch_protections required contexts.
#
# Runner label mapping (2026-05-22 fix): the actual Mac mini runner
# registered in this Gitea ships labels
# ["self-hosted","macos-self-hosted-arm64","arm64-darwin"]
# — no plain `arm64`. The earlier `runs-on: [self-hosted, arm64]`
# could not match any registered runner so every fire of this workflow
# was assigned task_id=0 / runner_id=NULL → Gitea cancelled it. The
# rows showed up as Cancelled in the action status feed (not Failed)
# but the lane never actually ran. Workflow now selects on
# `arm64-darwin` which is the canonical Mac-arm64 label per the
# Mac mini's registration (per internal#494 capability-honest labels).
#
# If we later want to add a Linux-arm64 runner to the same lane, add
# both labels to that runner's registration AND broaden the selector
# here — don't rename `arm64-darwin` (it's Mac-specific by design and
# `feedback_pc2_runner_labels_must_stay_narrow` rule applies).
#
# Pairs with internal#543 (RFC: Mac arm64 multi-arch runner-base) and
# internal#494 (multi-arch runner-base capability-honest labels).
# Pairs with internal#543 (RFC: Mac arm64 multi-arch runner-base).
# No paths: filter on purpose (feedback_path_filtered_workflow_cant_be_required).
on:
@@ -97,15 +82,7 @@ jobs:
echo "WARN: shellcheck binary not found — skipping (pilot mode)"
exit 0
fi
# NOTE: macOS ships Bash 3.2 (Apple license), no `mapfile`
# (Bash 4+ builtin). Mac mini runner empirically failed at
# `mapfile: command not found` (run 79275 / task 145654).
# Use the portable `while read` pattern instead — works on
# both Bash 3.2 (macOS) and Bash 4+ (Linux).
TARGETS=()
while IFS= read -r f; do
TARGETS+=("$f")
done < <(find .gitea/scripts -maxdepth 2 -type f -name '*.sh' | sort)
mapfile -t TARGETS < <(find .gitea/scripts -maxdepth 2 -type f -name '*.sh' | sort)
if [ "${#TARGETS[@]}" -eq 0 ]; then
echo "No .sh files found under .gitea/scripts — nothing to check"
exit 0
+261 -32
View File
@@ -34,6 +34,22 @@ interface TemplateSpec {
providers?: string[];
}
interface HermesProvider {
id: string;
label: string;
envVar: string;
defaultModel: string;
models: string[];
}
const DEFAULT_LLM_MODELS: SelectorModel[] = [
{ id: "moonshot/kimi-k2.6", name: "Kimi K2.6", provider: "platform", required_env: [] },
{ id: "MiniMax-M2.7", name: "MiniMax M2.7", required_env: ["MINIMAX_API_KEY"] },
{ id: "kimi-k2-turbo-preview", name: "Kimi K2 Turbo Preview", required_env: ["KIMI_API_KEY"] },
{ id: "claude-sonnet-4-6", name: "Claude Sonnet 4.6", required_env: ["ANTHROPIC_API_KEY"] },
{ id: "sonnet", name: "Claude Sonnet", required_env: ["CLAUDE_CODE_OAUTH_TOKEN"] },
];
const DEFAULT_PLATFORM_MODEL = DEFAULT_LLM_MODELS[0];
const DEFAULT_RUNTIME = "claude-code";
const RUNTIME_OPTIONS = [
{ value: "claude-code", label: "Claude Code" },
@@ -47,6 +63,31 @@ const DEFAULT_HEADLESS_ROOT_GB = 30;
const DEFAULT_DISPLAY_INSTANCE_TYPE = "t3.xlarge";
const DEFAULT_DISPLAY_ROOT_GB = 80;
// All providers supported by Hermes runtime via providers.resolve_provider().
// `defaultModel` is the slug injected into the workspace provision request
// when the user picks this provider — template-hermes's derive-provider.sh
// maps the prefix back to the provider name at install time, so this is
// the canonical handshake. `models` are additional suggestions surfaced in
// the datalist so the user can pick a different size without typing the
// whole slug.
export const HERMES_PROVIDERS: HermesProvider[] = [
{ id: "anthropic", label: "Anthropic (Claude)", envVar: "ANTHROPIC_API_KEY", defaultModel: "anthropic/claude-sonnet-4-5", models: ["anthropic/claude-opus-4-5", "anthropic/claude-sonnet-4-5", "anthropic/claude-haiku-4-5"] },
{ id: "openai", label: "OpenAI", envVar: "OPENAI_API_KEY", defaultModel: "openai/gpt-4o", models: ["openai/gpt-4o", "openai/gpt-4o-mini", "openai/o3-mini"] },
{ id: "openrouter", label: "OpenRouter", envVar: "OPENROUTER_API_KEY", defaultModel: "openrouter/auto", models: ["openrouter/auto", "openrouter/anthropic/claude-sonnet-4", "openrouter/meta-llama/llama-3.3-70b"] },
{ id: "xai", label: "xAI (Grok)", envVar: "XAI_API_KEY", defaultModel: "xai/grok-4", models: ["xai/grok-4", "xai/grok-4-mini"] },
{ id: "gemini", label: "Google Gemini", envVar: "GEMINI_API_KEY", defaultModel: "gemini/gemini-2.5-pro", models: ["gemini/gemini-2.5-pro", "gemini/gemini-2.5-flash"] },
{ id: "qwen", label: "Qwen (Alibaba)", envVar: "QWEN_API_KEY", defaultModel: "alibaba/qwen3-max", models: ["alibaba/qwen3-max", "alibaba/qwen3-coder"] },
{ id: "glm", label: "GLM (Zhipu AI)", envVar: "GLM_API_KEY", defaultModel: "zai/glm-4.6", models: ["zai/glm-4.6", "zai/glm-4.5-air"] },
{ id: "kimi", label: "Kimi (Moonshot)", envVar: "KIMI_API_KEY", defaultModel: "kimi-coding/kimi-k2", models: ["kimi-coding/kimi-k2", "kimi-coding/kimi-k1.5"] },
{ id: "minimax", label: "MiniMax", envVar: "MINIMAX_API_KEY", defaultModel: "minimax/MiniMax-M2.7", models: ["minimax/MiniMax-M2.7", "minimax/MiniMax-M2.7-highspeed", "minimax/MiniMax-M1"] },
{ id: "deepseek", label: "DeepSeek", envVar: "DEEPSEEK_API_KEY", defaultModel: "deepseek/deepseek-chat", models: ["deepseek/deepseek-chat", "deepseek/deepseek-reasoner"] },
{ id: "groq", label: "Groq", envVar: "GROQ_API_KEY", defaultModel: "openrouter/groq/llama-3.3-70b", models: ["openrouter/groq/llama-3.3-70b"] },
{ id: "mistral", label: "Mistral", envVar: "MISTRAL_API_KEY", defaultModel: "openrouter/mistralai/mistral-large", models: ["openrouter/mistralai/mistral-large"] },
{ id: "together", label: "Together AI", envVar: "TOGETHER_API_KEY", defaultModel: "openrouter/meta-llama/llama-3.3-70b", models: ["openrouter/meta-llama/llama-3.3-70b"] },
{ id: "fireworks", label: "Fireworks AI", envVar: "FIREWORKS_API_KEY", defaultModel: "openrouter/meta-llama/llama-3.3-70b", models: ["openrouter/meta-llama/llama-3.3-70b"] },
{ id: "hermes", label: "Hermes / Nous (legacy)", envVar: "HERMES_API_KEY", defaultModel: "nousresearch/Hermes-3-Llama-3.1-405B", models: ["nousresearch/Hermes-3-Llama-3.1-405B", "nousresearch/Hermes-4-14B"] },
];
export function CreateWorkspaceButton() {
const [open, setOpen] = useState(false);
const [name, setName] = useState("");
@@ -66,20 +107,32 @@ export function CreateWorkspaceButton() {
// filter below. Same data source ConfigTab uses (PR #2454). When the
// selected template declares `runtime_config.providers` in its
// config.yaml, the modal surfaces only those providers in the
// <select>. Provider/model options are derived from template models.
// <select>. Empty/missing list falls back to the full HERMES_PROVIDERS
// catalog so older templates without the field keep working.
const [templateSpecs, setTemplateSpecs] = useState<TemplateSpec[]>([]);
// External-runtime path: skip docker provision, mint a workspace_auth_token,
// and surface the connection snippet in a modal after create. When
// isExternal is true the template and model fields are hidden (they're
// meaningless for BYO-compute agents).
// isExternal is true the template / model / hermes-provider fields are
// hidden (they're meaningless for BYO-compute agents).
const [isExternal, setIsExternal] = useState(false);
const [externalRuntime, setExternalRuntime] = useState("external");
const [externalConnection, setExternalConnection] =
useState<ExternalConnectionInfo | null>(null);
// Hermes-specific state
const [hermesProvider, setHermesProvider] = useState("anthropic");
const [hermesApiKey, setHermesApiKey] = useState("");
// Model slug is sent to CP as `model` and plumbed to the workspace EC2
// as HERMES_DEFAULT_MODEL env var. template-hermes's derive-provider.sh
// reads the prefix (`minimax/…`, `anthropic/…`) to set
// HERMES_INFERENCE_PROVIDER at install time. Missing model → provider
// falls back to "auto" and hermes picks its compiled-in default
// (Anthropic), which 401s if the user's key is for a different
// provider. Hence: require model when template=hermes.
const [hermesModel, setHermesModel] = useState("");
const [llmSelection, setLLMSelection] = useState<SelectorValue>({
providerId: "",
model: "",
providerId: "platform|",
model: "moonshot/kimi-k2.6",
envVars: [],
});
const [llmSecret, setLLMSecret] = useState("");
@@ -141,7 +194,10 @@ export function CreateWorkspaceButton() {
const handleRuntimeChange = useCallback((nextRuntime: string) => {
setRuntime(nextRuntime);
setTemplate("");
setLLMSelection({ providerId: "", model: "", envVars: [] });
setHermesProvider("anthropic");
setHermesApiKey("");
setHermesModel("");
setLLMSelection({ providerId: "platform|", model: DEFAULT_PLATFORM_MODEL.id, envVars: [] });
setLLMSecret("");
}, []);
@@ -153,12 +209,9 @@ export function CreateWorkspaceButton() {
return templateSpecs.find((s) => s.id === template) ?? null;
}, [template, templateSpecs]);
const selectedRuntimeTemplateSpec = useMemo<TemplateSpec | null>(() => (
templateSpecs.find((s) => {
if (!BASE_RUNTIME_TEMPLATE_IDS.has(s.id)) return false;
const specRuntime = (s.runtime ?? s.id).trim().toLowerCase();
return s.id === runtime || specRuntime === runtime;
}) ?? null
templateSpecs.find((s) => s.id === runtime && BASE_RUNTIME_TEMPLATE_IDS.has(s.id)) ?? null
), [runtime, templateSpecs]);
const isHermes = runtime === "hermes";
const visibleTemplateSpecs = useMemo(
() => templateSpecs.filter((spec) => {
if (BASE_RUNTIME_TEMPLATE_IDS.has(spec.id)) return false;
@@ -169,11 +222,28 @@ export function CreateWorkspaceButton() {
);
const llmModels = useMemo(
() => {
const sourceSpec = selectedTemplateSpec ?? selectedRuntimeTemplateSpec;
if (!sourceSpec?.models?.length) return [];
return sourceSpec.models;
if (!selectedTemplateSpec?.models?.length) return DEFAULT_LLM_MODELS;
if (isHermes) {
return selectedTemplateSpec.models;
}
if (selectedTemplateSpec.models.some((model) => model.provider === "platform")) {
return selectedTemplateSpec.models;
}
const templateDefault = selectedTemplateSpec.model?.trim();
const defaultModelSpec = templateDefault
? selectedTemplateSpec.models.find((model) => model.id === templateDefault)
: undefined;
return [
{
id: templateDefault || DEFAULT_PLATFORM_MODEL.id,
name: defaultModelSpec?.name ?? DEFAULT_PLATFORM_MODEL.name,
provider: "platform",
required_env: [],
},
...selectedTemplateSpec.models,
];
},
[selectedRuntimeTemplateSpec, selectedTemplateSpec],
[isHermes, selectedTemplateSpec],
);
const llmCatalog = useMemo(() => buildProviderCatalog(llmModels), [llmModels]);
const selectedLLMProvider = useMemo(
@@ -181,22 +251,67 @@ export function CreateWorkspaceButton() {
[llmCatalog, llmSelection.providerId],
);
// Filter HERMES_PROVIDERS by what the template declares it supports.
// Empty/missing declared list → fall back to the full catalog so
// templates that haven't migrated to the explicit `providers:` field
// (and self-hosted setups without /templates) keep working unchanged.
const availableProviders = useMemo<HermesProvider[]>(() => {
const declared = selectedTemplateSpec?.providers ?? selectedRuntimeTemplateSpec?.providers;
if (!declared || declared.length === 0) return HERMES_PROVIDERS;
const allowed = new Set(declared.map((p) => p.toLowerCase()));
const filtered = HERMES_PROVIDERS.filter((p) => allowed.has(p.id.toLowerCase()));
// Defensive: if the template's declared list doesn't match anything
// in our static catalog (e.g. brand-new provider id we don't have
// metadata for yet), fall back to the full list rather than render
// an empty <select>. Better to over-show than to lock the user out.
return filtered.length > 0 ? filtered : HERMES_PROVIDERS;
}, [selectedRuntimeTemplateSpec, selectedTemplateSpec]);
// If the currently-selected provider is filtered out by a template
// change, snap back to the first available. Without this, the
// hermesProvider state could refer to a provider not in the dropdown
// — confusing UI + the API key field's envVar would be wrong.
useEffect(() => {
if (llmCatalog.length === 0) return;
const sourceDefault = (selectedTemplateSpec ?? selectedRuntimeTemplateSpec)?.model?.trim();
const platformProvider = llmCatalog.find((p) => p.vendor === "platform");
const matched = sourceDefault ? findProviderForModel(llmCatalog, sourceDefault) : null;
const next = platformProvider ?? matched ?? llmCatalog[0];
const defaultModel = next.models.find((model) => model.id === sourceDefault)?.id
?? next.models[0]?.id
?? "";
if (!isHermes) return;
if (availableProviders.length === 0) return;
if (!availableProviders.some((p) => p.id === hermesProvider)) {
setHermesProvider(availableProviders[0].id);
}
// eslint-disable-next-line react-hooks/exhaustive-deps
}, [availableProviders, isHermes]);
useEffect(() => {
if (isHermes || llmCatalog.length === 0) return;
const templateDefault = selectedTemplateSpec?.model?.trim();
const matched = templateDefault ? findProviderForModel(llmCatalog, templateDefault) : null;
const next = matched ?? llmCatalog[0];
setLLMSelection({
providerId: next.id,
model: next.wildcard ? "" : defaultModel,
model: matched && templateDefault
? templateDefault
: next.wildcard
? ""
: next.models[0]?.id ?? "",
envVars: next.envVars,
});
setLLMSecret("");
}, [llmCatalog, selectedRuntimeTemplateSpec, selectedTemplateSpec]);
}, [isHermes, llmCatalog, selectedTemplateSpec?.model]);
// Auto-fill hermesModel with the provider's defaultModel whenever the
// provider changes, but only if the user hasn't already typed their own
// slug. Prevents the empty-model → "auto" → Anthropic-default 401 trap.
useEffect(() => {
if (!isHermes) return;
const p = HERMES_PROVIDERS.find((x) => x.id === hermesProvider);
if (!p) return;
// Replace model only if current value matches another provider's
// default (user hasn't customized it) OR is empty.
const isUntouched =
hermesModel === "" ||
HERMES_PROVIDERS.some((x) => x.defaultModel === hermesModel);
if (isUntouched) setHermesModel(p.defaultModel);
// eslint-disable-next-line react-hooks/exhaustive-deps
}, [hermesProvider, isHermes]);
// Reset form and load workspaces whenever dialog opens
useEffect(() => {
@@ -213,8 +328,11 @@ export function CreateWorkspaceButton() {
setDisplayInstanceType(DEFAULT_DISPLAY_INSTANCE_TYPE);
setDisplayRootGB(String(DEFAULT_DISPLAY_ROOT_GB));
setDisplayResolution("1920x1080");
setHermesProvider("anthropic");
setExternalRuntime("external");
setLLMSelection({ providerId: "", model: "", envVars: [] });
setHermesApiKey("");
setHermesModel("");
setLLMSelection({ providerId: "platform|", model: "moonshot/kimi-k2.6", envVars: [] });
setLLMSecret("");
api
.get<WorkspaceOption[]>("/workspaces")
@@ -223,7 +341,7 @@ export function CreateWorkspaceButton() {
api
.get<TemplateSpec[]>("/templates")
.then((rows) => setTemplateSpecs(Array.isArray(rows) ? rows : []))
.catch(() => { /* keep empty; create stays blocked until the catalog loads */ });
.catch(() => { /* keep empty — HERMES_PROVIDERS fallback below */ });
// defaultTier is stable for the session (derived from window.location),
// safe to omit from deps.
// eslint-disable-next-line react-hooks/exhaustive-deps
@@ -234,18 +352,29 @@ export function CreateWorkspaceButton() {
setError("Name is required");
return;
}
if (!isExternal && !llmSelection.model.trim()) {
if (isHermes && !hermesApiKey.trim()) {
setError("API key is required for Hermes workspaces");
return;
}
if (isHermes && !hermesModel.trim()) {
setError("Model is required for Hermes workspaces — provider routing depends on the model slug prefix");
return;
}
if (!isExternal && !isHermes && !llmSelection.model.trim()) {
setError("Model is required");
return;
}
if (!isExternal && selectedLLMProvider?.envVars.length && !llmSecret.trim()) {
if (!isExternal && !isHermes && selectedLLMProvider?.envVars.length && !llmSecret.trim()) {
setError("Provider credential is required");
return;
}
setCreating(true);
setError(null);
const nativeProvider = selectedLLMProvider;
const provider = isHermes
? HERMES_PROVIDERS.find((p) => p.id === hermesProvider)
: undefined;
const nativeProvider = !isHermes ? selectedLLMProvider : undefined;
try {
const parsedBudget = budgetLimit.trim()
@@ -269,7 +398,7 @@ export function CreateWorkspaceButton() {
tier,
parent_id: parentId || undefined,
budget_limit: parsedBudget,
...(!isExternal && nativeProvider
...(!isExternal && !isHermes && nativeProvider
? {
model: llmSelection.model.trim(),
llm_provider: nativeProvider.vendor,
@@ -303,6 +432,12 @@ export function CreateWorkspaceButton() {
// no container provisioning, token minted, connection payload
// returned in the response for the modal below.
...(isExternal ? { runtime: externalRuntime } : { runtime }),
...(!isExternal && isHermes && provider
? {
secrets: { [provider.envVar]: hermesApiKey.trim() },
model: hermesModel.trim(),
}
: {}),
});
// External path: keep the create dialog open just long enough to
// hand control to the connect modal, then close. The connect
@@ -453,7 +588,7 @@ export function CreateWorkspaceButton() {
</div>
)}
{!isExternal && selectedLLMProvider && (
{!isExternal && !isHermes && selectedLLMProvider && (
<div className="rounded-lg border border-line/50 bg-surface-card/40 p-3 space-y-3">
<div className="text-[11px] font-medium text-ink-mid">
LLM
@@ -609,6 +744,100 @@ export function CreateWorkspaceButton() {
</div>
</div>
{/* Hermes provider configuration — shown only for the Hermes runtime. */}
{isHermes && (
<div
className="mt-4 rounded-xl border border-violet-700/40 bg-violet-950/20 p-4 space-y-3"
data-testid="hermes-provider-section"
>
<p className="text-[11px] font-semibold text-violet-400 uppercase tracking-wide">
Hermes Provider
</p>
<p className="text-[11px] text-ink-mid -mt-1">
Choose the AI provider and paste your API key. The key is
stored as an encrypted workspace secret.
</p>
<div>
<label
htmlFor="hermes-provider-select"
className="text-[11px] text-ink-mid block mb-1"
>
Provider
</label>
<select
id="hermes-provider-select"
value={hermesProvider}
onChange={(e) => setHermesProvider(e.target.value)}
aria-label="Hermes provider"
className="w-full bg-surface-card/60 border border-line/50 rounded-lg px-3 py-2 text-sm text-ink focus:outline-none focus:border-violet-500/60 focus:ring-1 focus:ring-violet-500/20 transition-colors"
>
{availableProviders.map((p) => (
<option key={p.id} value={p.id}>
{p.label}
</option>
))}
</select>
</div>
<div>
<label
htmlFor="hermes-api-key-input"
className="text-[11px] text-ink-mid block mb-1"
>
API Key{" "}
<span aria-hidden="true" className="text-bad">
*
</span>
<span className="sr-only"> (required)</span>
</label>
<input
id="hermes-api-key-input"
type="password"
value={hermesApiKey}
onChange={(e) => setHermesApiKey(e.target.value)}
placeholder="sk-…"
aria-label="Hermes API key"
autoComplete="off"
className="w-full bg-surface-card/60 border border-line/50 rounded-lg px-3 py-2 text-sm text-ink placeholder-ink-soft focus:outline-none focus:border-violet-500/60 focus:ring-1 focus:ring-violet-500/20 transition-colors font-mono"
/>
</div>
<div>
<label
htmlFor="hermes-model-input"
className="text-[11px] text-ink-mid block mb-1"
>
Model{" "}
<span aria-hidden="true" className="text-bad">
*
</span>
<span className="sr-only"> (required)</span>
</label>
<input
id="hermes-model-input"
type="text"
value={hermesModel}
onChange={(e) => setHermesModel(e.target.value)}
placeholder="e.g. minimax/MiniMax-M2.7"
aria-label="Hermes model slug"
autoComplete="off"
spellCheck={false}
list="hermes-model-suggestions"
className="w-full bg-surface-card/60 border border-line/50 rounded-lg px-3 py-2 text-sm text-ink placeholder-ink-soft focus:outline-none focus:border-violet-500/60 focus:ring-1 focus:ring-violet-500/20 transition-colors font-mono"
/>
<datalist id="hermes-model-suggestions">
{HERMES_PROVIDERS.find((p) => p.id === hermesProvider)?.models.map(
(m) => <option key={m} value={m} />,
)}
</datalist>
<p className="text-[10px] text-ink-mid mt-1">
Slug determines which provider hermes routes to at install time.
</p>
</div>
</div>
)}
{error && (
<div
role="alert"
@@ -1,7 +1,7 @@
// @vitest-environment jsdom
import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
import { render, screen, fireEvent, waitFor, cleanup } from "@testing-library/react";
import { CreateWorkspaceButton } from "../CreateWorkspaceDialog";
import { CreateWorkspaceButton, HERMES_PROVIDERS } from "../CreateWorkspaceDialog";
vi.mock("@/lib/api", () => ({
api: {
@@ -21,22 +21,6 @@ const SAMPLE_WORKSPACES = [
];
const SAMPLE_TEMPLATES = [
{
id: "claude-code-default",
name: "Claude Code Agent",
runtime: "claude-code",
model: "moonshot/kimi-k2.6",
providers: ["platform", "minimax", "kimi-coding", "anthropic", "anthropic-oauth"],
models: [
{ id: "moonshot/kimi-k2.6", name: "Kimi K2.6", provider: "platform", required_env: [] },
{ id: "MiniMax-M2.7", name: "MiniMax M2.7", required_env: ["MINIMAX_API_KEY"] },
{ id: "kimi-k2-turbo-preview", name: "Kimi K2 Turbo Preview", required_env: ["KIMI_API_KEY"] },
{ id: "claude-sonnet-4-6", name: "Claude Sonnet 4.6", required_env: ["ANTHROPIC_API_KEY"] },
{ id: "sonnet", name: "Claude Sonnet", required_env: ["CLAUDE_CODE_OAUTH_TOKEN"] },
{ id: "opus", name: "Claude Opus", required_env: ["CLAUDE_CODE_OAUTH_TOKEN"] },
{ id: "haiku", name: "Claude Haiku", required_env: ["CLAUDE_CODE_OAUTH_TOKEN"] },
],
},
{
id: "seo-agent",
name: "SEO Agent",
@@ -49,22 +33,9 @@ const SAMPLE_TEMPLATES = [
{ id: "kimi-k2-turbo-preview", name: "Kimi K2 Turbo Preview", required_env: ["KIMI_API_KEY"] },
{ id: "claude-sonnet-4-6", name: "Claude Sonnet 4.6", required_env: ["ANTHROPIC_API_KEY"] },
{ id: "sonnet", name: "Claude Sonnet", required_env: ["CLAUDE_CODE_OAUTH_TOKEN"] },
{ id: "opus", name: "Claude Opus", required_env: ["CLAUDE_CODE_OAUTH_TOKEN"] },
{ id: "haiku", name: "Claude Haiku", required_env: ["CLAUDE_CODE_OAUTH_TOKEN"] },
],
},
{
id: "hermes",
name: "Hermes",
runtime: "hermes",
model: "openai/gpt-4o",
providers: ["openai", "anthropic", "platform"],
models: [
{ id: "openai/gpt-4o", name: "GPT-4o", required_env: ["OPENAI_API_KEY"] },
{ id: "anthropic/claude-sonnet-4-5", name: "Claude Sonnet 4.5", required_env: ["ANTHROPIC_API_KEY"] },
{ id: "moonshot/kimi-k2.6", name: "Kimi K2.6", provider: "platform", required_env: [] },
],
},
{ id: "hermes", name: "Hermes", runtime: "hermes" },
];
beforeEach(() => {
@@ -298,9 +269,6 @@ describe("CreateWorkspaceDialog", () => {
fireEvent.change(document.querySelector("[data-testid='provider-select']") as HTMLSelectElement, {
target: { value: "anthropic-oauth|CLAUDE_CODE_OAUTH_TOKEN" },
});
fireEvent.change(document.querySelector("[data-testid='model-select']") as HTMLSelectElement, {
target: { value: "sonnet" },
});
fireEvent.change(document.getElementById("llm-secret-input") as HTMLInputElement, {
target: { value: "oauth-token" },
});
@@ -315,18 +283,6 @@ describe("CreateWorkspaceDialog", () => {
expect(body.secrets).toEqual({ CLAUDE_CODE_OAUTH_TOKEN: "oauth-token" });
});
it("lists all Claude Code subscription aliases for blank workspaces", async () => {
await openDialog();
fireEvent.change(document.querySelector("[data-testid='provider-select']") as HTMLSelectElement, {
target: { value: "anthropic-oauth|CLAUDE_CODE_OAUTH_TOKEN" },
});
const modelSelect = document.querySelector("[data-testid='model-select']") as HTMLSelectElement;
const optionValues = Array.from(modelSelect.options).map((option) => option.value);
expect(optionValues).toEqual(expect.arrayContaining(["sonnet", "opus", "haiku"]));
});
it("renders gracefully when GET /workspaces fails", async () => {
mockGet.mockRejectedValueOnce(new Error("Network error"));
await openDialog();
@@ -341,103 +297,226 @@ describe("CreateWorkspaceDialog", () => {
});
// ---------------------------------------------------------------------------
// Dynamic runtime provider picker tests
// Hermes provider picker tests
// ---------------------------------------------------------------------------
describe("CreateWorkspaceDialog — dynamic runtime provider picker", () => {
it("does not render the old Hermes-only provider section", async () => {
describe("CreateWorkspaceDialog — Hermes provider picker", () => {
it("does NOT show hermes provider section for non-hermes templates", async () => {
await openDialog();
await setRuntime("hermes");
await setTemplate("seo-agent");
expect(document.querySelector("[data-testid='hermes-provider-section']")).toBeNull();
});
it("derives Hermes provider and model options from the /templates runtime row", async () => {
it("shows hermes provider section when runtime is 'hermes'", async () => {
await openDialog();
await setRuntime("hermes");
const providerSelect = document.querySelector("[data-testid='provider-select']") as HTMLSelectElement;
await waitFor(() => expect(providerSelect.options.length).toBe(4));
const providerValues = Array.from(providerSelect.options).map((option) => option.value);
expect(providerValues).toEqual(expect.arrayContaining([
"platform|",
"openai|OPENAI_API_KEY",
"anthropic|ANTHROPIC_API_KEY",
]));
expect(providerValues).not.toContain("gemini|GEMINI_API_KEY");
await waitFor(() =>
expect(document.querySelector("[data-testid='hermes-provider-section']")).toBeTruthy()
);
});
it("uses the template-declared default provider/model for Hermes", async () => {
it("shows hermes provider section for the Hermes runtime preset", async () => {
await openDialog();
await setRuntime("hermes");
await waitFor(() => {
const providerSelect = document.querySelector("[data-testid='provider-select']") as HTMLSelectElement;
expect(providerSelect.value).toBe("platform|");
});
const modelSelect = document.querySelector("[data-testid='model-select']") as HTMLSelectElement;
expect(modelSelect.value).toBe("moonshot/kimi-k2.6");
await waitFor(() =>
expect(document.querySelector("[data-testid='hermes-provider-section']")).toBeTruthy()
);
});
it("prompts for the provider credential required by the selected Hermes model", async () => {
it("hermes provider dropdown defaults to 'anthropic'", async () => {
await openDialog();
await setRuntime("hermes");
await waitFor(() =>
expect(document.querySelector("[data-testid='hermes-provider-section']")).toBeTruthy()
);
const providerSelect = document.getElementById("hermes-provider-select") as HTMLSelectElement;
expect(providerSelect).toBeTruthy();
expect(providerSelect.value).toBe("anthropic");
});
fireEvent.change(document.querySelector("[data-testid='provider-select']") as HTMLSelectElement, {
target: { value: "openai|OPENAI_API_KEY" },
it("hermes provider dropdown lists all 15 providers", async () => {
await openDialog();
await setRuntime("hermes");
await waitFor(() =>
expect(document.querySelector("[data-testid='hermes-provider-section']")).toBeTruthy()
);
const providerSelect = document.getElementById("hermes-provider-select") as HTMLSelectElement;
expect(providerSelect.options.length).toBe(HERMES_PROVIDERS.length);
const ids = Array.from(providerSelect.options).map((o) => o.value);
expect(ids).toContain("anthropic");
expect(ids).toContain("openai");
expect(ids).toContain("gemini");
expect(ids).toContain("deepseek");
expect(ids).toContain("hermes");
});
// Pins the dynamic-providers behavior: when the matched template's
// /templates row declares `providers`, the dropdown filters to that
// subset instead of showing the full HERMES_PROVIDERS catalog. Same
// data source ConfigTab uses (PR #2454) — keeps the modal and the
// settings tab honest about which providers a template supports.
it("hermes provider dropdown filters to template-declared providers when /templates ships them", async () => {
// Per-URL mock: /workspaces returns the existing fixture, /templates
// returns a hermes row that only allows anthropic + minimax + openai.
mockGet.mockImplementation(async (url: string) => {
if (url === "/templates") {
return [
{ id: "hermes", name: "Hermes", runtime: "hermes", providers: ["anthropic", "minimax", "openai"] },
// eslint-disable-next-line @typescript-eslint/no-explicit-any
] as any;
}
// eslint-disable-next-line @typescript-eslint/no-explicit-any
return SAMPLE_WORKSPACES as any;
});
const keyInput = document.getElementById("llm-secret-input") as HTMLInputElement;
await openDialog();
await setRuntime("hermes");
await waitFor(() =>
expect(document.querySelector("[data-testid='hermes-provider-section']")).toBeTruthy()
);
const providerSelect = document.getElementById("hermes-provider-select") as HTMLSelectElement;
// Filtered list arrives async after /templates fetch resolves —
// keep waiting until the dropdown shrinks below the full catalog.
await waitFor(() => expect(providerSelect.options.length).toBe(3));
const ids = Array.from(providerSelect.options).map((o) => o.value);
expect(ids).toEqual(expect.arrayContaining(["anthropic", "minimax", "openai"]));
expect(ids).not.toContain("gemini");
expect(ids).not.toContain("deepseek");
});
// Back-compat: a template that hasn't migrated to runtime_config.providers
// (older templates, self-hosted setups without /templates server) keeps
// showing the full provider catalog. Operators picking from those
// templates can't be locked out of providers we know hermes supports.
it("hermes provider dropdown falls back to all providers when template declares no providers list", async () => {
mockGet.mockImplementation(async (url: string) => {
if (url === "/templates") {
// No `providers` field — empty/missing → fall back to full catalog.
// eslint-disable-next-line @typescript-eslint/no-explicit-any
return [{ id: "hermes", name: "Hermes", runtime: "hermes" }] as any;
}
// eslint-disable-next-line @typescript-eslint/no-explicit-any
return SAMPLE_WORKSPACES as any;
});
await openDialog();
await setRuntime("hermes");
await waitFor(() =>
expect(document.querySelector("[data-testid='hermes-provider-section']")).toBeTruthy()
);
const providerSelect = document.getElementById("hermes-provider-select") as HTMLSelectElement;
expect(providerSelect.options.length).toBe(HERMES_PROVIDERS.length);
});
// Defensive: a template's declared list with NO matches against our
// static catalog (e.g. a brand-new provider id we don't have label/
// envVar metadata for yet) must not render an empty <select> — the
// operator can't pick a provider, the form locks. Component falls
// back to the full catalog so the user can still proceed.
it("hermes provider dropdown falls back to all providers when template declares only unknown providers", async () => {
mockGet.mockImplementation(async (url: string) => {
if (url === "/templates") {
return [
{ id: "hermes", name: "Hermes", runtime: "hermes", providers: ["totally-new-provider-2030"] },
// eslint-disable-next-line @typescript-eslint/no-explicit-any
] as any;
}
// eslint-disable-next-line @typescript-eslint/no-explicit-any
return SAMPLE_WORKSPACES as any;
});
await openDialog();
await setRuntime("hermes");
await waitFor(() =>
expect(document.querySelector("[data-testid='hermes-provider-section']")).toBeTruthy()
);
const providerSelect = document.getElementById("hermes-provider-select") as HTMLSelectElement;
// Stays at full catalog length — no flapping to 0 then back.
expect(providerSelect.options.length).toBe(HERMES_PROVIDERS.length);
});
it("hermes API key field is a password input (masked)", async () => {
await openDialog();
await setRuntime("hermes");
await waitFor(() =>
expect(document.querySelector("[data-testid='hermes-provider-section']")).toBeTruthy()
);
const keyInput = document.getElementById("hermes-api-key-input") as HTMLInputElement;
expect(keyInput).toBeTruthy();
expect(keyInput.type).toBe("password");
});
it("shows an error if the selected runtime provider requires a credential", async () => {
it("shows an error if hermes template is set but API key is empty on submit", async () => {
await openDialog();
fireEvent.change(screen.getByPlaceholderText("e.g. SEO Agent"), {
target: { value: "Hermes Agent" },
});
await setRuntime("hermes");
fireEvent.change(document.querySelector("[data-testid='provider-select']") as HTMLSelectElement, {
target: { value: "openai|OPENAI_API_KEY" },
});
await waitFor(() =>
expect(document.querySelector("[data-testid='hermes-provider-section']")).toBeTruthy()
);
// Submit without API key
const createBtn = screen.getAllByRole("button").find((b) => b.textContent === "Create");
fireEvent.click(createBtn!);
await waitFor(() => {
const alert = screen.getByRole("alert");
expect(alert.textContent).toContain("Provider credential");
expect(alert.textContent).toContain("API key");
});
expect(mockPost).not.toHaveBeenCalled();
});
it("includes runtime-derived provider/model/secrets in POST body", async () => {
it("includes secrets in POST body with correct env var for selected provider", async () => {
await openDialog();
fireEvent.change(screen.getByPlaceholderText("e.g. SEO Agent"), {
target: { value: "Hermes OpenAI" },
target: { value: "Hermes Agent" },
});
await setRuntime("hermes");
fireEvent.change(document.querySelector("[data-testid='provider-select']") as HTMLSelectElement, {
target: { value: "openai|OPENAI_API_KEY" },
});
fireEvent.change(document.getElementById("llm-secret-input") as HTMLInputElement, {
target: { value: "sk-openai-test" },
});
await waitFor(() =>
expect(document.querySelector("[data-testid='hermes-provider-section']")).toBeTruthy()
);
// Fill in the API key
const keyInput = document.getElementById("hermes-api-key-input") as HTMLInputElement;
fireEvent.change(keyInput, { target: { value: "sk-test-anthropic-key" } });
const createBtn = screen.getAllByRole("button").find((b) => b.textContent === "Create");
fireEvent.click(createBtn!);
await waitFor(() => expect(mockPost).toHaveBeenCalled());
const body = mockPost.mock.calls[0][1] as Record<string, unknown>;
expect(body.secrets).toEqual({ ANTHROPIC_API_KEY: "sk-test-anthropic-key" });
expect(body.runtime).toBe("hermes");
expect(body.template).toBeUndefined();
expect(body.model).toBe("openai/gpt-4o");
expect(body.llm_provider).toBe("openai");
});
it("uses the correct env var when a non-default provider is selected", async () => {
await openDialog();
fireEvent.change(screen.getByPlaceholderText("e.g. SEO Agent"), {
target: { value: "Hermes OpenAI" },
});
await setRuntime("hermes");
await waitFor(() =>
expect(document.querySelector("[data-testid='hermes-provider-section']")).toBeTruthy()
);
// Switch to openai
const providerSelect = document.getElementById("hermes-provider-select") as HTMLSelectElement;
fireEvent.change(providerSelect, { target: { value: "openai" } });
const keyInput = document.getElementById("hermes-api-key-input") as HTMLInputElement;
fireEvent.change(keyInput, { target: { value: "sk-openai-test" } });
const createBtn = screen.getAllByRole("button").find((b) => b.textContent === "Create");
fireEvent.click(createBtn!);
await waitFor(() => expect(mockPost).toHaveBeenCalled());
const body = mockPost.mock.calls[0][1] as Record<string, unknown>;
expect(body.secrets).toEqual({ OPENAI_API_KEY: "sk-openai-test" });
});
it("does NOT include secrets field when provider is platform-managed", async () => {
it("does NOT include secrets field when template is not hermes", async () => {
await openDialog();
fireEvent.change(screen.getByPlaceholderText("e.g. SEO Agent"), {
target: { value: "Normal Agent" },
@@ -451,6 +530,20 @@ describe("CreateWorkspaceDialog — dynamic runtime provider picker", () => {
const body = mockPost.mock.calls[0][1] as Record<string, unknown>;
expect(body.secrets).toBeUndefined();
});
it("hides hermes section and resets state when template is cleared", async () => {
await openDialog();
await setRuntime("hermes");
await waitFor(() =>
expect(document.querySelector("[data-testid='hermes-provider-section']")).toBeTruthy()
);
// Switch back to a non-Hermes runtime.
await setRuntime("claude-code");
await waitFor(() =>
expect(document.querySelector("[data-testid='hermes-provider-section']")).toBeNull()
);
});
});
// ---------------------------------------------------------------------------
@@ -131,7 +131,7 @@ export function OrgTokensTab() {
<button
onClick={handleCreate}
disabled={creating}
className="px-3 py-1.5 bg-accent-strong/20 hover:bg-accent-strong/30 border border-accent/30 rounded-lg text-[11px] text-accent font-medium transition-colors disabled:opacity-50 disabled:cursor-not-allowed flex items-center gap-1.5 focus:outline-none focus-visible:ring-2 focus-visible:ring-accent focus-visible:ring-offset-1"
className="px-3 py-1.5 bg-accent-strong/20 hover:bg-accent-strong/30 border border-accent/30 rounded-lg text-[11px] text-accent font-medium transition-colors disabled:opacity-50 disabled:cursor-not-allowed flex items-center gap-1.5"
>
{creating ? (
<>
+1 -1
View File
@@ -121,7 +121,7 @@ function WorkspaceTokensTab({ workspaceId }: TokensTabProps) {
<button
onClick={handleCreate}
disabled={creating}
className="px-3 py-1.5 bg-accent-strong/20 hover:bg-accent-strong/30 border border-accent/30 rounded-lg text-[11px] text-accent font-medium transition-colors disabled:opacity-50 disabled:cursor-not-allowed flex items-center gap-1.5 focus:outline-none focus-visible:ring-2 focus-visible:ring-accent focus-visible:ring-offset-1"
className="px-3 py-1.5 bg-accent-strong/20 hover:bg-accent-strong/30 border border-accent/30 rounded-lg text-[11px] text-accent font-medium transition-colors disabled:opacity-50 disabled:cursor-not-allowed flex items-center gap-1.5"
>
{creating ? <><Spinner size="sm" /> Creating...</> : '+ New Token'}
</button>
+8 -81
View File
@@ -6,7 +6,6 @@ import { useCanvasStore } from "@/store/canvas";
import { type ConfigData, DEFAULT_CONFIG, TextInput, NumberInput, Toggle, TagList, Section } from "./config/form-inputs";
import { parseYaml, toYaml } from "./config/yaml-utils";
import { SecretsSection } from "./config/secrets-section";
import { LLMBillingSection } from "./config/llm-billing-section";
import { ExternalConnectionSection } from "./ExternalConnectionSection";
import {
ProviderModelSelector,
@@ -288,40 +287,6 @@ export function deriveProvidersFromModels(models: ModelSpec[]): string[] {
return out;
}
// billingModeForProvider — maps a selected PROVIDER (vendor key) to the
// LLM billing_mode it implies (internal#703 Gap 2).
//
// Today, picking a non-Platform provider in the Config tab writes the
// credential env (CLAUDE_CODE_OAUTH_TOKEN / vendor key) but leaves
// llm_billing_mode at its resolved default (`platform_managed`). The CP
// tenant_config endpoint then keeps injecting the platform proxy base
// URLs, so the OAuth token / vendor key is never actually used — BYOK
// silently no-ops (the live SEO-Agent symptom in #703). The workspace-
// server even hard-blocks vendor-key writes on platform_managed
// workspaces (secrets.go:87), pointing the user at this exact billing-
// mode switch. Wiring the provider change to also set billing_mode is
// the UI half that makes BYOK take (the CP/workspace-server backend half
// is being fixed in parallel — internal#703 Gap 1).
//
// Mapping:
// - "platform" (the Platform-managed proxy) OR "" (no explicit
// provider override → inherit, defaults to platform) → "platform_managed".
// - any other vendor key ("anthropic-oauth" = Claude Code subscription
// OAuth, "anthropic" = Anthropic API key, "minimax", "openrouter",
// etc.) → "byok".
//
// Returns the billing_mode string the PUT body should carry. The valid
// set is fixed by workspace-server's recognizer (platform_managed | byok
// | disabled); "disabled" is never auto-selected by a provider choice —
// it's an explicit operator action via the LLM Billing section.
export type LLMBillingMode = "platform_managed" | "byok";
export function billingModeForProvider(provider: string): LLMBillingMode {
const v = provider.trim().toLowerCase();
if (v === "" || v === "platform") return "platform_managed";
return "byok";
}
// Fallback used when /templates can't be fetched (offline, older backend).
// Keep in sync with manifest.json workspace_templates as a defensive default.
// Model + env suggestions only flow when the backend is reachable.
@@ -736,36 +701,6 @@ export function ConfigTab({ workspaceId }: Props) {
}
}
// Provider → billing_mode linkage (internal#703 Gap 2). When the
// provider actually changed AND its implied billing_mode differs
// from the previously-selected provider's, push the new mode to
// the per-tenant llm-billing-mode endpoint (same path the LLM
// Billing section uses). Without this, selecting a non-Platform
// provider leaves billing_mode=platform_managed → CP keeps
// injecting the platform proxy → BYOK never takes.
//
// Gated on (a) the provider PUT having succeeded — no point setting
// byok if the credential write failed — and (b) the mode actually
// changing, so an unrelated provider tweak between two BYOK vendors
// (e.g. minimax → openrouter) doesn't re-issue a redundant
// platform_managed→byok PUT and trigger a needless restart.
let billingModeSaveError: string | null = null;
if (providerChanged && !providerSaveError) {
const nextMode = billingModeForProvider(provider);
const prevMode = billingModeForProvider(originalProvider);
if (nextMode !== prevMode) {
try {
await api.put(
`/admin/workspaces/${workspaceId}/llm-billing-mode`,
{ mode: nextMode },
);
} catch (e) {
billingModeSaveError =
e instanceof Error ? e.message : "Billing mode update was rejected";
}
}
}
setOriginalYaml(content);
if (rawMode) {
const parsed = parseYaml(content);
@@ -785,22 +720,16 @@ export function ConfigTab({ workspaceId }: Props) {
} else if (!restart) {
useCanvasStore.getState().updateNodeData(workspaceId, { needsRestart: !providerWillAutoRestart });
}
// Aggregate partial-save errors. modelSaveError, providerSaveError,
// and billingModeSaveError describe rejected updates from
// independent endpoints — show whichever fired so the user knows
// which field reverts on next reload (otherwise they'd see "Saved"
// and be confused why Provider snapped back). The billing-mode case
// is the most important to surface: the provider credential saved
// but BYOK won't actually take until billing_mode flips, so a
// silent failure here is exactly the #703 "selecting a provider has
// no effect" symptom.
// Aggregate partial-save errors. Both modelSaveError and
// providerSaveError describe rejected updates from independent
// endpoints — show whichever fired so the user knows which
// field reverts on next reload (otherwise they'd see "Saved" and
// be confused why Provider snapped back).
const partialError = providerSaveError
? `Other fields saved, but provider update failed: ${providerSaveError}`
: billingModeSaveError
? `Provider saved, but switching billing mode failed — your own provider key/OAuth may not take effect until billing mode is set: ${billingModeSaveError}`
: modelSaveError
? `Other fields saved, but model update failed: ${modelSaveError}`
: null;
: modelSaveError
? `Other fields saved, but model update failed: ${modelSaveError}`
: null;
if (partialError) {
setError(partialError);
} else {
@@ -1179,8 +1108,6 @@ export function ConfigTab({ workspaceId }: Props) {
</div>
</Section>
<LLMBillingSection workspaceId={workspaceId} />
<SecretsSection
workspaceId={workspaceId}
requiredEnv={config.runtime_config?.required_env}
@@ -1,255 +0,0 @@
// @vitest-environment jsdom
//
// Tests for the provider → llm_billing_mode linkage (internal#703 Gap 2).
//
// What this pins: when the operator changes the PROVIDER in the Config
// tab, the workspace's llm_billing_mode must follow — a non-Platform
// provider sets billing_mode=byok; Platform sets platform_managed. Before
// this wiring, selecting "Claude Code subscription (OAuth)" or any vendor
// key wrote the credential env but left billing_mode=platform_managed, so
// CP kept injecting the platform proxy base URL and the OAuth token /
// vendor key was never used — BYOK silently no-op'd (the live jrs-auto
// SEO-Agent symptom in #703).
//
// The billing-mode PUT targets the same per-tenant endpoint the LLM
// Billing section uses: PUT /admin/workspaces/:id/llm-billing-mode with
// body {mode: "byok" | "platform_managed"}.
import { describe, it, expect, vi, afterEach, beforeEach } from "vitest";
import { render, screen, cleanup, waitFor, fireEvent } from "@testing-library/react";
import React from "react";
afterEach(cleanup);
const apiGet = vi.fn();
const apiPatch = vi.fn();
const apiPut = vi.fn();
vi.mock("@/lib/api", () => ({
api: {
get: (path: string) => apiGet(path),
patch: (path: string, body: unknown) => apiPatch(path, body),
put: (path: string, body: unknown) => apiPut(path, body),
post: vi.fn(),
del: vi.fn(),
},
}));
const storeUpdateNodeData = vi.fn();
const storeRestartWorkspace = vi.fn();
vi.mock("@/store/canvas", () => ({
useCanvasStore: Object.assign(
(selector: (s: unknown) => unknown) =>
selector({ restartWorkspace: storeRestartWorkspace, updateNodeData: storeUpdateNodeData }),
{
getState: () => ({
restartWorkspace: storeRestartWorkspace,
updateNodeData: storeUpdateNodeData,
}),
},
),
}));
vi.mock("../AgentCardSection", () => ({
AgentCardSection: () => <div data-testid="agent-card-stub" />,
}));
import { ConfigTab, billingModeForProvider } from "../ConfigTab";
function wireApi(opts: { providerValue?: string | "missing" }) {
apiGet.mockImplementation((path: string) => {
if (path === `/workspaces/ws-test`) {
return Promise.resolve({ runtime: "hermes" });
}
if (path === `/workspaces/ws-test/model`) {
return Promise.resolve({ model: "nousresearch/hermes-4-70b" });
}
if (path === `/workspaces/ws-test/provider`) {
if (opts.providerValue === "missing") return Promise.reject(new Error("404"));
return Promise.resolve({
provider: opts.providerValue ?? "",
source: opts.providerValue ? "workspace_secrets" : "default",
});
}
if (path === `/workspaces/ws-test/files/config.yaml`) {
return Promise.resolve({ content: "name: ws\nruntime: hermes\n" });
}
if (path === "/templates") return Promise.resolve([]);
return Promise.reject(new Error(`unmocked api.get: ${path}`));
});
}
function billingModeCalls() {
return apiPut.mock.calls.filter(
([path]) => path === "/admin/workspaces/ws-test/llm-billing-mode",
);
}
beforeEach(() => {
apiGet.mockReset();
apiPatch.mockReset();
apiPut.mockReset();
storeUpdateNodeData.mockReset();
storeRestartWorkspace.mockReset();
});
describe("billingModeForProvider — pure mapping (internal#703 Gap 2)", () => {
// Platform / empty → platform_managed. Empty means "no explicit
// override → inherit", which resolves to platform on the backend, so
// it must NOT flip the workspace into byok.
it("maps Platform and empty to platform_managed", () => {
expect(billingModeForProvider("platform")).toBe("platform_managed");
expect(billingModeForProvider("")).toBe("platform_managed");
expect(billingModeForProvider(" ")).toBe("platform_managed");
expect(billingModeForProvider("PLATFORM")).toBe("platform_managed");
});
// Every non-Platform provider → byok. If this regresses to returning
// platform_managed for a vendor, BYOK silently no-ops again (#703).
it("maps non-Platform providers to byok", () => {
expect(billingModeForProvider("anthropic-oauth")).toBe("byok"); // Claude Code subscription
expect(billingModeForProvider("anthropic")).toBe("byok"); // Anthropic API key
expect(billingModeForProvider("minimax")).toBe("byok");
expect(billingModeForProvider("openrouter")).toBe("byok");
expect(billingModeForProvider("openai")).toBe("byok");
});
});
describe("ConfigTab — provider change drives billing_mode (internal#703 Gap 2)", () => {
// The core fix: picking a non-Platform provider (here "anthropic-oauth"
// = Claude Code subscription OAuth) from a fresh/empty provider must
// PUT mode=byok to the per-tenant llm-billing-mode endpoint. This is
// the exact path that was missing — the credential env saved but the
// billing mode never followed, so the proxy stayed engaged.
it("PUTs mode=byok when switching to a non-Platform provider", async () => {
wireApi({ providerValue: "" });
apiPut.mockResolvedValue({ status: "saved" });
render(<ConfigTab workspaceId="ws-test" />);
const input = await screen.findByTestId("provider-input");
fireEvent.change(input, { target: { value: "anthropic-oauth" } });
fireEvent.click(screen.getByRole("button", { name: /^save$/i }));
await waitFor(() => {
const calls = billingModeCalls();
expect(calls.length).toBe(1);
expect(calls[0][1]).toEqual({ mode: "byok" });
});
// Provider credential PUT still happens too (independent endpoint).
expect(
apiPut.mock.calls.some(([path]) => path === "/workspaces/ws-test/provider"),
).toBe(true);
});
// Switching FROM a byok provider back TO Platform must PUT
// mode=platform_managed so the workspace re-engages the proxy and stops
// expecting a (now-absent) vendor key.
it("PUTs mode=platform_managed when switching back to Platform", async () => {
wireApi({ providerValue: "anthropic-oauth" });
apiPut.mockResolvedValue({ status: "saved" });
render(<ConfigTab workspaceId="ws-test" />);
const input = await screen.findByTestId("provider-input");
await waitFor(() => expect((input as HTMLInputElement).value).toBe("anthropic-oauth"));
fireEvent.change(input, { target: { value: "platform" } });
fireEvent.click(screen.getByRole("button", { name: /^save$/i }));
await waitFor(() => {
const calls = billingModeCalls();
expect(calls.length).toBe(1);
expect(calls[0][1]).toEqual({ mode: "platform_managed" });
});
});
// Changing between two BYOK vendors (minimax → openrouter) keeps
// billing_mode=byok — the implied mode is unchanged, so re-PUTing it
// would be a wasteful no-op that risks an extra restart. Must NOT fire.
it("does NOT PUT billing-mode when the implied mode is unchanged", async () => {
wireApi({ providerValue: "minimax" });
apiPut.mockResolvedValue({ status: "saved" });
render(<ConfigTab workspaceId="ws-test" />);
const input = await screen.findByTestId("provider-input");
await waitFor(() => expect((input as HTMLInputElement).value).toBe("minimax"));
fireEvent.change(input, { target: { value: "openrouter" } });
fireEvent.click(screen.getByRole("button", { name: /^save$/i }));
await waitFor(() => {
// Provider PUT fires (vendor changed)...
expect(
apiPut.mock.calls.some(([path]) => path === "/workspaces/ws-test/provider"),
).toBe(true);
});
// ...but billing-mode does NOT (byok → byok is a no-op).
expect(billingModeCalls().length).toBe(0);
});
// A Save that doesn't touch the provider must not PUT billing-mode —
// editing tier/name shouldn't disturb the workspace's billing mode.
it("does NOT PUT billing-mode on a Save that leaves provider unchanged", async () => {
wireApi({ providerValue: "anthropic-oauth" });
apiPut.mockResolvedValue({ status: "saved" });
render(<ConfigTab workspaceId="ws-test" />);
await screen.findByTestId("provider-input");
// Dirty an unrelated field so Save is enabled.
const tierSelect = screen.getByLabelText(/tier/i) as HTMLSelectElement;
fireEvent.change(tierSelect, { target: { value: "3" } });
fireEvent.click(screen.getByRole("button", { name: /^save$/i }));
await waitFor(() => {
// Some PUT may fire (e.g. /model); just assert billing-mode did not.
expect(billingModeCalls().length).toBe(0);
});
});
// If the provider credential PUT itself fails, we must NOT set byok —
// flipping billing_mode while the credential write failed would leave
// the workspace expecting a key it doesn't have (worse than no-op).
it("does NOT PUT billing-mode when the provider PUT fails", async () => {
wireApi({ providerValue: "" });
apiPut.mockImplementation((path: string) => {
if (path === "/workspaces/ws-test/provider") return Promise.reject(new Error("boom"));
return Promise.resolve({ status: "saved" });
});
render(<ConfigTab workspaceId="ws-test" />);
const input = await screen.findByTestId("provider-input");
fireEvent.change(input, { target: { value: "anthropic-oauth" } });
fireEvent.click(screen.getByRole("button", { name: /^save$/i }));
await waitFor(() => {
// The provider-failure error is surfaced (getByText throws if absent).
expect(screen.getByText(/provider update failed/i)).toBeTruthy();
});
expect(billingModeCalls().length).toBe(0);
});
// If the credential saved but the billing-mode PUT is rejected, the
// user must be warned that BYOK may not take — a silent failure here
// is precisely the #703 symptom we're fixing.
it("surfaces an error when billing-mode PUT fails after a successful provider save", async () => {
wireApi({ providerValue: "" });
apiPut.mockImplementation((path: string) => {
if (path === "/admin/workspaces/ws-test/llm-billing-mode") {
return Promise.reject(new Error("403 forbidden"));
}
return Promise.resolve({ status: "saved" });
});
render(<ConfigTab workspaceId="ws-test" />);
const input = await screen.findByTestId("provider-input");
fireEvent.change(input, { target: { value: "anthropic-oauth" } });
fireEvent.click(screen.getByRole("button", { name: /^save$/i }));
await waitFor(() => {
expect(screen.getByText(/switching billing mode failed/i)).toBeTruthy();
});
});
});
@@ -1,176 +0,0 @@
// @vitest-environment jsdom
import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
import {
render,
screen,
waitFor,
cleanup,
fireEvent,
} from "@testing-library/react";
import { LLMBillingSection } from "../llm-billing-section";
// Tests for LLMBillingSection (internal#691). Locks in:
// - the section renders the resolved mode + source label
// - the dropdown maps "inherit" → PUT {mode: null}
// - the dropdown maps "byok" → PUT {mode: "byok"}
// - a garbled override surfaces the warning banner
// - the post-write resolution updates the UI without a refetch
const apiGet = vi.fn();
const apiPut = vi.fn();
vi.mock("@/lib/api", () => ({
api: {
get: (...args: unknown[]) => apiGet(...args),
put: (...args: unknown[]) => apiPut(...args),
post: vi.fn().mockResolvedValue({}),
del: vi.fn().mockResolvedValue({}),
patch: vi.fn().mockResolvedValue({}),
},
}));
// Collapsed-by-default Section wrapper would hide the content; replace
// it with a passthrough so the dropdown is reachable in the test DOM.
vi.mock("../form-inputs", async () => {
const actual = await vi.importActual<typeof import("../form-inputs")>(
"../form-inputs",
);
return {
...actual,
Section: ({ children }: { children: React.ReactNode }) => (
<div>{children}</div>
),
};
});
beforeEach(() => {
vi.clearAllMocks();
});
afterEach(() => {
cleanup();
});
describe("LLMBillingSection — internal#691", () => {
it("renders the resolved mode + source for an inherited workspace", async () => {
apiGet.mockResolvedValueOnce({
workspace_id: "ws-1",
resolved_mode: "platform_managed",
workspace_override: null,
org_default: "platform_managed",
source: "org_default",
});
render(<LLMBillingSection workspaceId="ws-1" />);
await waitFor(() => {
expect(apiGet).toHaveBeenCalledWith(
"/admin/workspaces/ws-1/llm-billing-mode",
);
});
// Resolved mode appears.
expect(screen.getByText(/Resolved mode:/i).textContent).toMatch(/platform_managed/);
// Source label appears.
expect(
screen.getByText(/inherited from org default/i),
).toBeTruthy();
});
it('PUTs {mode: "byok"} when user picks BYOK and reflects the new resolution', async () => {
apiGet.mockResolvedValueOnce({
workspace_id: "ws-2",
resolved_mode: "platform_managed",
workspace_override: null,
org_default: "platform_managed",
source: "org_default",
});
apiPut.mockResolvedValueOnce({
workspace_id: "ws-2",
resolved_mode: "byok",
workspace_override: "byok",
org_default: "platform_managed",
source: "workspace_override",
});
render(<LLMBillingSection workspaceId="ws-2" />);
await waitFor(() => expect(apiGet).toHaveBeenCalled());
const select = (await screen.findByLabelText(
/llm billing mode override/i,
)) as HTMLSelectElement;
fireEvent.change(select, { target: { value: "byok" } });
await waitFor(() => {
expect(apiPut).toHaveBeenCalledWith(
"/admin/workspaces/ws-2/llm-billing-mode",
{ mode: "byok" },
);
});
// Post-write resolution propagated to UI.
await waitFor(() => {
expect(
screen.getByText(/explicit override on this workspace/i),
).toBeTruthy();
});
});
it("PUTs {mode: null} when user picks Inherit (clears the override)", async () => {
apiGet.mockResolvedValueOnce({
workspace_id: "ws-3",
resolved_mode: "byok",
workspace_override: "byok",
org_default: "platform_managed",
source: "workspace_override",
});
apiPut.mockResolvedValueOnce({
workspace_id: "ws-3",
resolved_mode: "platform_managed",
workspace_override: null,
org_default: "platform_managed",
source: "org_default",
});
render(<LLMBillingSection workspaceId="ws-3" />);
await waitFor(() => expect(apiGet).toHaveBeenCalled());
const select = (await screen.findByLabelText(
/llm billing mode override/i,
)) as HTMLSelectElement;
fireEvent.change(select, { target: { value: "inherit" } });
await waitFor(() => {
expect(apiPut).toHaveBeenCalledWith(
"/admin/workspaces/ws-3/llm-billing-mode",
{ mode: null },
);
});
});
it("surfaces a warning banner when the override value is garbled", async () => {
apiGet.mockResolvedValueOnce({
workspace_id: "ws-4",
resolved_mode: "platform_managed", // resolver fell through, default-closed
workspace_override: "byokk", // typo persisted somehow
org_default: "platform_managed",
source: "org_default",
});
render(<LLMBillingSection workspaceId="ws-4" />);
await waitFor(() => {
expect(
screen.getByText(/non-standard value/i),
).toBeTruthy();
});
});
it("renders an error banner when the GET fails", async () => {
apiGet.mockRejectedValueOnce(new Error("network down"));
render(<LLMBillingSection workspaceId="ws-5" />);
await waitFor(() => {
expect(screen.getByText(/network down/i)).toBeTruthy();
});
});
});
@@ -1,4 +1,3 @@
export { type ConfigData, DEFAULT_CONFIG, TextInput, NumberInput, Toggle, TagList, Section } from "./form-inputs";
export { parseYaml, toYaml } from "./yaml-utils";
export { SecretsSection } from "./secrets-section";
export { LLMBillingSection } from "./llm-billing-section";
@@ -1,219 +0,0 @@
"use client";
// llm-billing-section.tsx — Config-tab section for the per-workspace
// llm_billing_mode override (internal#691).
//
// Surfaces:
// - The currently RESOLVED mode for this workspace (the mode the
// workspace-server's strip gate will use at next provision).
// - The org-level default (so the user sees what they're inheriting).
// - A dropdown to set / clear the workspace-level override.
// - A "source" line so operators can answer "is this inherited or
// explicit?" without DB archeology (RFC Observability hot-spot).
//
// Hits:
// GET /admin/workspaces/:id/llm-billing-mode — read resolution
// PUT /admin/workspaces/:id/llm-billing-mode — write {mode: "..."|null}
//
// Both routes are on the per-tenant workspace-server (same origin as the
// other canvas /admin calls). CP's proxy at /cp/admin/workspaces/:id/
// llm-billing-mode exists for ops use; the canvas uses the per-tenant
// path directly to keep the round-trip cheap.
import { useState, useEffect, useCallback } from "react";
import { api } from "@/lib/api";
import { Section } from "./form-inputs";
// Mirrors workspace-server/internal/handlers/llm_billing_mode.go::BillingModeResolution.
// Kept as a literal shape (not imported) because canvas has no Go-type bridge.
export interface BillingModeResolution {
workspace_id: string;
resolved_mode: "platform_managed" | "byok" | "disabled";
// Pointer-typed on the Go side: nil = inherit, non-nil = the raw
// workspace-level override (even if garbled and falling through).
workspace_override: string | null;
org_default: "platform_managed" | "byok" | "disabled";
source: "workspace_override" | "org_default" | "constant_fallback";
}
// The dropdown emits one of these values. "inherit" is the UX-only label
// that maps to a `null` body in the PUT request.
type DropdownChoice = "inherit" | "platform_managed" | "byok" | "disabled";
interface Props {
workspaceId: string;
}
const MODE_LABELS: Record<DropdownChoice, string> = {
inherit: "Inherit from org default",
platform_managed: "Platform-managed (uses Molecule credits)",
byok: "BYOK (your own OAuth / vendor keys)",
disabled: "Disabled (no LLM access)",
};
const MODE_DESCRIPTIONS: Record<DropdownChoice, string> = {
inherit:
"Use whichever mode is set at the organization level. Recommended unless this specific workspace needs a different billing source.",
platform_managed:
"Strip CLAUDE_CODE_OAUTH_TOKEN and vendor API keys from the workspace; route all LLM traffic through Molecule's proxy and bill your org credits.",
byok:
"Keep CLAUDE_CODE_OAUTH_TOKEN / vendor API keys in the workspace; LLM traffic goes directly to your provider and is billed to your OAuth subscription or API account.",
disabled:
"Block all LLM access for this workspace. Useful for sandbox workspaces that should not consume credits or hit external providers.",
};
const SOURCE_LABELS: Record<BillingModeResolution["source"], string> = {
workspace_override: "explicit override on this workspace",
org_default: "inherited from org default",
constant_fallback:
"fallback (workspace + org defaults missing or unrecognized — defaulted to platform_managed)",
};
export function LLMBillingSection({ workspaceId }: Props) {
const [resolution, setResolution] = useState<BillingModeResolution | null>(
null,
);
const [loading, setLoading] = useState(true);
const [saving, setSaving] = useState(false);
const [error, setError] = useState<string | null>(null);
const [success, setSuccess] = useState(false);
const load = useCallback(async () => {
setLoading(true);
setError(null);
try {
const res = await api.get<BillingModeResolution>(
`/admin/workspaces/${workspaceId}/llm-billing-mode`,
);
setResolution(res);
} catch (e) {
setError(e instanceof Error ? e.message : "Failed to load billing mode");
} finally {
setLoading(false);
}
}, [workspaceId]);
useEffect(() => {
void load();
}, [load]);
// Current dropdown selection is derived from the resolution. If the
// override is null, we show "inherit"; otherwise we mirror the raw
// workspace_override (NOT resolved_mode — that would conflate "explicit
// platform_managed override" with "inherit while org happens to be
// platform_managed", which has different semantics on the write side).
const currentChoice: DropdownChoice = (() => {
if (!resolution) return "inherit";
if (resolution.workspace_override == null) return "inherit";
const raw = resolution.workspace_override;
if (raw === "platform_managed" || raw === "byok" || raw === "disabled") {
return raw;
}
// Garbled value persisted via some external write. Show inherit so
// the user can pick a clean value; on save they'll either clear it
// (PUT null) or overwrite it with a valid one.
return "inherit";
})();
const handleChange = async (choice: DropdownChoice) => {
if (!resolution) return;
setSaving(true);
setError(null);
setSuccess(false);
try {
// "inherit" → PUT {mode: null}; otherwise → PUT {mode: choice}.
const body = choice === "inherit" ? { mode: null } : { mode: choice };
const updated = await api.put<BillingModeResolution>(
`/admin/workspaces/${workspaceId}/llm-billing-mode`,
body,
);
setResolution(updated);
setSuccess(true);
setTimeout(() => setSuccess(false), 2000);
} catch (e) {
setError(e instanceof Error ? e.message : "Failed to update billing mode");
} finally {
setSaving(false);
}
};
return (
<Section title="LLM Billing" defaultOpen={false}>
{loading && (
<div className="text-[10px] text-ink-mid">Loading billing mode</div>
)}
{error && (
<div
role="alert"
aria-live="assertive"
className="px-2 py-1 bg-red-900/30 border border-red-800 rounded text-[10px] text-bad mb-2"
>
{error}
</div>
)}
{resolution && (
<div className="space-y-2">
<div className="text-[10px] text-ink-mid">
Resolved mode: <strong className="text-ink">{resolution.resolved_mode}</strong>{" "}
<span className="text-ink-mid">
({SOURCE_LABELS[resolution.source]})
</span>
</div>
<div className="text-[10px] text-ink-mid">
Org default: <span className="text-ink">{resolution.org_default}</span>
</div>
<label
className="block text-[10px] text-ink-mid"
htmlFor={`llm-billing-mode-${workspaceId}`}
>
Override
</label>
<select
id={`llm-billing-mode-${workspaceId}`}
aria-label="LLM billing mode override"
value={currentChoice}
disabled={saving}
onChange={(e) => void handleChange(e.target.value as DropdownChoice)}
className="w-full bg-surface-card border border-line rounded p-1 text-[10px] text-ink focus:outline-none focus:border-accent disabled:opacity-50"
>
{(Object.keys(MODE_LABELS) as DropdownChoice[]).map((m) => (
<option key={m} value={m}>
{MODE_LABELS[m]}
</option>
))}
</select>
<div
className="text-[10px] text-ink-mid leading-snug"
aria-live="polite"
>
{MODE_DESCRIPTIONS[currentChoice]}
</div>
{success && (
<div className="mt-1 px-2 py-1 bg-green-900/30 border border-green-800 rounded text-[10px] text-good">
Updated. Restart the workspace to apply.
</div>
)}
{resolution.workspace_override != null &&
!["platform_managed", "byok", "disabled"].includes(
resolution.workspace_override,
) && (
<div
role="alert"
className="mt-1 px-2 py-1 bg-yellow-900/30 border border-yellow-800 rounded text-[10px] text-warning"
>
Workspace override has a non-standard value (
<code>{resolution.workspace_override}</code>) and is being
ignored. Pick a valid mode above to clear the corrupt value.
</div>
)}
</div>
)}
</Section>
);
}
@@ -0,0 +1,98 @@
// @vitest-environment jsdom
/**
* Tests for design-tokens.ts — STATUS_CONFIG, TIER_CONFIG, COMM_TYPE_LABELS
* plus the statusDotClass function exported from design-tokens.ts.
*
* Note: statusDotClass is also tested in statusDotClass.test.ts; this file
* covers the remaining exports and edge cases.
*/
import { describe, it, expect } from "vitest";
import {
STATUS_CONFIG,
statusDotClass,
TIER_CONFIG,
COMM_TYPE_LABELS,
} from "../design-tokens";
describe("STATUS_CONFIG", () => {
it("has entries for all known status values", () => {
const statuses = ["online", "offline", "paused", "degraded", "failed", "provisioning", "not_configured"];
for (const s of statuses) {
expect(STATUS_CONFIG[s]).toBeTruthy();
expect(typeof STATUS_CONFIG[s].dot).toBe("string");
expect(typeof STATUS_CONFIG[s].label).toBe("string");
expect(typeof STATUS_CONFIG[s].bar).toBe("string");
}
});
it("provisioning has motion-safe:animate-pulse in dot class", () => {
expect(STATUS_CONFIG.provisioning.dot).toContain("animate-pulse");
});
it("failed and degraded have glow classes", () => {
expect(STATUS_CONFIG.failed.glow).toBeTruthy();
expect(STATUS_CONFIG.degraded.glow).toBeTruthy();
});
});
describe("statusDotClass", () => {
it("returns dot class for known status", () => {
expect(statusDotClass("online")).toBe("bg-emerald-400");
});
it("returns fallback bg-zinc-500 for unknown status", () => {
expect(statusDotClass("nonsense")).toBe("bg-zinc-500");
});
it("returns fallback bg-zinc-500 for empty string", () => {
expect(statusDotClass("")).toBe("bg-zinc-500");
});
});
describe("TIER_CONFIG", () => {
it("has entries for tiers 1-4", () => {
for (let tier = 1; tier <= 4; tier++) {
expect(TIER_CONFIG[tier]).toBeTruthy();
expect(typeof TIER_CONFIG[tier].label).toBe("string");
expect(typeof TIER_CONFIG[tier].color).toBe("string");
expect(typeof TIER_CONFIG[tier].border).toBe("string");
}
});
it("tier labels are T{num}", () => {
expect(TIER_CONFIG[1].label).toBe("T1");
expect(TIER_CONFIG[2].label).toBe("T2");
expect(TIER_CONFIG[3].label).toBe("T3");
expect(TIER_CONFIG[4].label).toBe("T4");
});
it("tier 1 uses ink-mid (safe/read-only)", () => {
expect(TIER_CONFIG[1].color).toContain("text-ink-mid");
});
it("tier 2 uses accent (full agents, read+write)", () => {
expect(TIER_CONFIG[2].color).toContain("bg-accent");
});
it("tier 3 uses violet (privileged)", () => {
expect(TIER_CONFIG[3].color).toContain("bg-violet-600");
});
it("tier 4 uses warm (full-host)", () => {
expect(TIER_CONFIG[4].color).toContain("bg-warm");
});
});
describe("COMM_TYPE_LABELS", () => {
it("maps a2a_send to 'sent'", () => {
expect(COMM_TYPE_LABELS.a2a_send).toBe("sent");
});
it("maps a2a_receive to 'received'", () => {
expect(COMM_TYPE_LABELS.a2a_receive).toBe("received");
});
it("maps task_update to 'task update'", () => {
expect(COMM_TYPE_LABELS.task_update).toBe("task update");
});
});
+54 -151
View File
@@ -1,205 +1,108 @@
// @vitest-environment jsdom
"use client";
/**
* Tests for palette-context.tsx — MobileAccentProvider context + usePalette hook.
* Tests for palette-context.tsx — normalizeStatus, tierCode, getPalette.
*
* Test coverage (9 cases):
* 1. MobileAccentProvider renders children
* 2. usePalette(false) without provider → MOL_LIGHT
* 3. usePalette(true) without provider → MOL_DARK
* 4. accent=null returns base palette unchanged
* 5. accent=base.accent returns base palette unchanged (identity guard)
* 6. accent="#custom" overrides both accent and online
* 7. MOL_LIGHT singleton never mutated
* 8. MOL_DARK singleton never mutated
*
* Plus pure-function coverage for normalizeStatus + tierCode.
* Pure functions that don't require the React context to test.
*/
import { describe, expect, it, vi, beforeEach, afterEach } from "vitest";
import React from "react";
import { render, screen, cleanup } from "@testing-library/react";
import { describe, it, expect } from "vitest";
import {
MOL_LIGHT,
MOL_DARK,
getPalette,
normalizeStatus,
tierCode,
MobileAccentProvider,
usePalette,
getPalette,
MOL_LIGHT,
MOL_DARK,
} from "../palette-context";
// ─── usePalette test helper ───────────────────────────────────────────────────
// usePalette reads document.documentElement.dataset.theme internally.
// We set this before rendering so the hook sees the right value.
function setDataTheme(theme: "light" | "dark") {
if (typeof document !== "undefined") {
document.documentElement.dataset.theme = theme;
}
}
// ─── Pure function tests ──────────────────────────────────────────────────────
describe("normalizeStatus", () => {
it("returns emerald-400 for online status", () => {
it("online → bg-emerald-400", () => {
expect(normalizeStatus("online", false)).toBe("bg-emerald-400");
expect(normalizeStatus("online", true)).toBe("bg-emerald-400");
});
it("returns emerald-400 for degraded status", () => {
it("degraded → bg-emerald-400", () => {
expect(normalizeStatus("degraded", false)).toBe("bg-emerald-400");
expect(normalizeStatus("degraded", true)).toBe("bg-emerald-400");
});
it("returns red-400 for failed status", () => {
it("failed → bg-red-400", () => {
expect(normalizeStatus("failed", false)).toBe("bg-red-400");
expect(normalizeStatus("failed", true)).toBe("bg-red-400");
});
it("returns amber-400 for paused status", () => {
it("paused → bg-amber-400", () => {
expect(normalizeStatus("paused", false)).toBe("bg-amber-400");
expect(normalizeStatus("paused", true)).toBe("bg-amber-400");
});
it("returns amber-400 for not_configured status", () => {
it("not_configured → bg-amber-400", () => {
expect(normalizeStatus("not_configured", false)).toBe("bg-amber-400");
});
it("returns zinc-400 for unknown status", () => {
expect(normalizeStatus("unknown", false)).toBe("bg-zinc-400");
it("unknown status → bg-zinc-400", () => {
expect(normalizeStatus("offline", false)).toBe("bg-zinc-400");
expect(normalizeStatus("provisioning", false)).toBe("bg-zinc-400");
expect(normalizeStatus("nonsense", false)).toBe("bg-zinc-400");
expect(normalizeStatus("", false)).toBe("bg-zinc-400");
});
});
describe("tierCode", () => {
it("returns T1 for tier 1", () => {
it("maps tier 1-4 to T1-T4", () => {
expect(tierCode(1)).toBe("T1");
});
it("returns T2 for tier 2", () => {
expect(tierCode(2)).toBe("T2");
});
it("returns T4 for tier 4", () => {
expect(tierCode(3)).toBe("T3");
expect(tierCode(4)).toBe("T4");
});
it("returns generic T{n} for non-standard tiers", () => {
expect(tierCode(99)).toBe("T99");
it("negative tier", () => {
expect(tierCode(0)).toBe("T0");
expect(tierCode(-1)).toBe("T-1");
});
});
// ─── getPalette tests ─────────────────────────────────────────────────────────
describe("getPalette — accent override", () => {
it("accent=null returns base palette unchanged (light)", () => {
const result = getPalette(null, false);
expect(result).toEqual({ ...MOL_LIGHT });
expect(result).not.toBe(MOL_LIGHT); // returned object is a copy
describe("getPalette", () => {
it("null accent with light → MOL_LIGHT", () => {
const p = getPalette(null, false);
expect(p.accent).toBe(MOL_LIGHT.accent);
expect(p.online).toBe(MOL_LIGHT.online);
});
it("accent=null returns base palette unchanged (dark)", () => {
const result = getPalette(null, true);
expect(result).toEqual({ ...MOL_DARK });
expect(result).not.toBe(MOL_DARK);
it("null accent with dark → MOL_DARK", () => {
const p = getPalette(null, true);
expect(p.accent).toBe(MOL_DARK.accent);
expect(p.online).toBe(MOL_DARK.online);
});
it("accent=base.accent returns base palette unchanged (identity guard, light)", () => {
const result = getPalette(MOL_LIGHT.accent, false);
expect(result).toEqual({ ...MOL_LIGHT });
expect(result).not.toBe(MOL_LIGHT);
it("returns a new object, not the singleton", () => {
const p = getPalette(null, false);
expect(p).not.toBe(MOL_LIGHT);
expect(p).not.toBe(MOL_DARK);
});
it("accent=base.accent returns base palette unchanged (identity guard, dark)", () => {
const result = getPalette(MOL_DARK.accent, true);
expect(result).toEqual({ ...MOL_DARK });
expect(result).not.toBe(MOL_DARK);
it("identity guard: same accent as base → returns copy of base", () => {
const p = getPalette(MOL_LIGHT.accent, false);
expect(p.accent).toBe(MOL_LIGHT.accent);
expect(p).not.toBe(MOL_LIGHT);
});
it("accent='#custom' overrides accent and online (light)", () => {
const result = getPalette("#ff0000", false);
expect(result.accent).toBe("#ff0000");
expect(result.online).toBe("bg-emerald-400"); // normalizeStatus("online", false)
it("custom accent → overrides accent and online", () => {
const p = getPalette("#ff0000", false);
expect(p.accent).toBe("#ff0000");
// online should be normalizeStatus("online", false) = bg-emerald-400
expect(p.online).toBe("bg-emerald-400");
// other fields unchanged
expect(p.ink).toBe(MOL_LIGHT.ink);
expect(p.surface).toBe(MOL_LIGHT.surface);
});
it("accent='#custom' overrides accent and online (dark)", () => {
const result = getPalette("#00ff00", true);
expect(result.accent).toBe("#00ff00");
expect(result.online).toBe("bg-emerald-400"); // normalizeStatus("online", true)
it("custom accent in dark mode", () => {
const p = getPalette("#00ff00", true);
expect(p.accent).toBe("#00ff00");
expect(p.online).toBe("bg-emerald-400"); // normalizeStatus is dark-agnostic for online
});
it("MOL_LIGHT singleton is never mutated", () => {
getPalette("#mutate", false);
// All fields must still match the original freeze definition
expect(MOL_LIGHT.accent).toBe("bg-blue-500");
expect(MOL_LIGHT.online).toBe("bg-emerald-400");
expect(MOL_LIGHT.surface).toBe("bg-zinc-900");
expect(MOL_LIGHT.ink).toBe("text-zinc-100");
expect(MOL_LIGHT.line).toBe("border-zinc-700");
expect(MOL_LIGHT.bg).toBe("bg-zinc-950");
});
it("MOL_DARK singleton is never mutated", () => {
getPalette("#mutate", true);
expect(MOL_DARK.accent).toBe("bg-sky-400");
expect(MOL_DARK.online).toBe("bg-emerald-400");
expect(MOL_DARK.surface).toBe("bg-zinc-800");
expect(MOL_DARK.ink).toBe("text-zinc-100");
expect(MOL_DARK.line).toBe("border-zinc-700");
expect(MOL_DARK.bg).toBe("bg-zinc-950");
});
it("getPalette always returns a new object (no shared mutation risk)", () => {
const a = getPalette("#a", false);
const b = getPalette("#b", false);
expect(a).not.toBe(b);
expect(a.accent).not.toBe(b.accent);
});
});
// ─── MobileAccentProvider tests ───────────────────────────────────────────────
describe("MobileAccentProvider", () => {
beforeEach(() => {
setDataTheme("light");
});
afterEach(() => {
cleanup();
if (typeof document !== "undefined") {
document.documentElement.dataset.theme = "";
}
});
it("renders children", () => {
render(
<MobileAccentProvider accent={null}>
<span data-testid="child">Hello</span>
</MobileAccentProvider>,
);
expect(screen.getByTestId("child")).toBeTruthy();
});
// usePalette hook reads data-theme from <html> to determine light/dark.
// In the test environment, data-theme is empty, which falls through to
// the "light" default in usePalette, giving MOL_LIGHT.
it("usePalette(false) without provider → MOL_LIGHT", () => {
setDataTheme("light");
function ShowPalette() {
const p = usePalette(false);
return <span data-testid="accent-light">{p.accent}</span>;
}
render(<ShowPalette />);
expect(screen.getByTestId("accent-light").textContent).toBe(MOL_LIGHT.accent);
});
it("usePalette(true) without provider → MOL_DARK when data-theme=dark", () => {
setDataTheme("dark");
function ShowPalette() {
const p = usePalette(true);
return <span data-testid="accent-dark">{p.accent}</span>;
}
render(<ShowPalette />);
expect(screen.getByTestId("accent-dark").textContent).toBe(MOL_DARK.accent);
it("custom accent does not mutate MOL_LIGHT or MOL_DARK", () => {
getPalette("#custom", false);
expect(MOL_LIGHT.accent).toBe("bg-blue-500"); // unchanged
getPalette("#custom2", true);
expect(MOL_DARK.accent).toBe("bg-sky-400"); // unchanged
});
});
@@ -0,0 +1,46 @@
// @vitest-environment jsdom
/**
* Tests for theme-provider.tsx.
*
* Re-export contract:
* - THEME_COOKIE value (string "mol_theme") from theme-cookie
* - themeBootScript value from theme-cookie
* - ThemePreference + ResolvedTheme types (runtime value = undefined)
*
* The ThemeProvider component itself requires full React context rendering;
* prop contract is enforced by TypeScript.
*/
import { describe, it, expect, beforeEach } from "vitest";
describe("applyResolvedTheme", () => {
beforeEach(() => {
document.documentElement.removeAttribute("data-theme");
});
it("sets data-theme on html element", () => {
document.documentElement.dataset.theme = "dark";
expect(document.documentElement.dataset.theme).toBe("dark");
document.documentElement.dataset.theme = "light";
expect(document.documentElement.dataset.theme).toBe("light");
});
});
describe("ThemeProvider component", () => {
it("is a function (React component)", async () => {
const { ThemeProvider } = await import("../theme-provider");
expect(typeof ThemeProvider).toBe("function");
});
});
describe("re-exports from theme-cookie", () => {
it("re-exports THEME_COOKIE = 'mol_theme'", async () => {
const { THEME_COOKIE } = await import("../theme-provider");
expect(THEME_COOKIE).toBe("mol_theme");
});
it("re-exports themeBootScript as a string value", async () => {
const { themeBootScript } = await import("../theme-provider");
expect(typeof themeBootScript).toBe("string");
expect(themeBootScript.length).toBeGreaterThan(0);
});
});
-10
View File
@@ -658,11 +658,6 @@
outline-offset: var(--focus-ring-offset);
}
.delete-dialog__cancel-btn:focus-visible {
outline: var(--focus-ring);
outline-offset: var(--focus-ring-offset);
}
.delete-dialog__confirm-btn {
background: var(--status-invalid);
color: #ffffff;
@@ -676,11 +671,6 @@
outline-offset: var(--focus-ring-offset);
}
.delete-dialog__confirm-btn:focus-visible {
outline: var(--focus-ring);
outline-offset: var(--focus-ring-offset);
}
.delete-dialog__confirm-btn:disabled { opacity: 0.4; cursor: not-allowed; }
/* ── Unsaved changes guard ─────────────────────────── */
+29 -47
View File
@@ -73,15 +73,7 @@ else
fi
# Test 4: Create workspace B (needs bearer — tokens now exist in DB)
# #1953 cross-tenant isolation: Summarizer is created as a CHILD of Echo so the
# two live in the SAME org (Echo is the org root; Summarizer hangs off it via
# parent_id). The peer-discovery tests below assert same-org peer enumeration
# (Echo sees its child, the child sees its parent). Previously both were created
# parent_id=NULL — two DISTINCT org roots — and "peers" only listed each other
# via the `WHERE parent_id IS NULL` branch that returned every tenant's org root.
# That branch WAS the cross-tenant leak (#1953) and is now removed, so two org
# roots no longer see each other; the assertions must run inside one org.
R=$(acurl -X POST "$BASE/workspaces" -H "Content-Type: application/json" -d "{\"name\":\"Summarizer Agent\",\"tier\":1,\"runtime\":\"external\",\"external\":true,\"parent_id\":\"$ECHO_ID\"}")
R=$(acurl -X POST "$BASE/workspaces" -H "Content-Type: application/json" -d '{"name":"Summarizer Agent","tier":1,"runtime":"external","external":true}')
check "POST /workspaces (create summarizer)" '"status":"awaiting_agent"' "$R"
SUM_ID=$(echo "$R" | python3 -c "import sys,json; print(json.load(sys.stdin)['id'])")
@@ -141,23 +133,21 @@ check "Heartbeat updated uptime" '"uptime_seconds":120' "$R"
R=$(curl -s "$BASE/registry/discover/$ECHO_ID")
check "GET /registry/discover/:id (missing caller rejected)" 'X-Workspace-ID header is required' "$R"
# Test 12: Discover (from same-org child — allowed)
# Test 12: Discover (from sibling — allowed)
R=$(curl -s "$BASE/registry/discover/$ECHO_ID" -H "X-Workspace-ID: $SUM_ID" -H "Authorization: Bearer $SUM_TOKEN")
check "GET /registry/discover/:id (same-org)" '"url"' "$R"
check "GET /registry/discover/:id (sibling)" '"url"' "$R"
# Test 13: Peers — same-org parent/child see each other (#1953). Echo is the org
# root and lists its child Summarizer; Summarizer lists its parent Echo. A
# cross-org workspace would NOT appear here (see cross_tenant_isolation_test.go).
# Test 13: Peers (root siblings see each other)
R=$(curl -s "$BASE/registry/$ECHO_ID/peers" -H "Authorization: Bearer $ECHO_TOKEN")
check "GET /registry/:id/peers (has summarizer)" '"Summarizer' "$R"
R=$(curl -s "$BASE/registry/$SUM_ID/peers" -H "Authorization: Bearer $SUM_TOKEN")
check "GET /registry/:id/peers (has echo)" '"Echo Agent"' "$R"
# Test 14: Check access (same-org parent↔child — allowed)
# Test 14: Check access (root siblings)
R=$(curl -s -X POST "$BASE/registry/check-access" -H "Content-Type: application/json" \
-d "{\"caller_id\":\"$ECHO_ID\",\"target_id\":\"$SUM_ID\"}")
check "POST /registry/check-access (same-org allowed)" '"allowed":true' "$R"
check "POST /registry/check-access (siblings allowed)" '"allowed":true' "$R"
# Test 15: PATCH workspace (update position)
R=$(acurl -X PATCH "$BASE/workspaces/$ECHO_ID" -H "Content-Type: application/json" -d '{"x":100,"y":200}')
@@ -299,40 +289,32 @@ R=$(curl -s "$BASE/workspaces" -H "Authorization: Bearer $ECHO_TOKEN")
check "current_task in list response" '"current_task"' "$R"
# Test 21: Delete
# #1953: Summarizer is now a CHILD of Echo (same-org, for the peer-discovery
# tests above). DELETE on the *parent* (Echo) cascade-removes its descendants
# (CascadeDelete walks the recursive `parent_id` CTE), so deleting Echo first
# would also remove Summarizer and the "one survives" assertion would see 0.
# Delete the CHILD (Summarizer) here instead: a child delete does NOT cascade
# upward, so the parent Echo survives and count=1 holds. The bundle round-trip
# below needs Summarizer's exported config, so capture it BEFORE this delete.
BUNDLE=$(curl -s "$BASE/bundles/export/$SUM_ID" -H "Authorization: Bearer $SUM_TOKEN")
check "GET /bundles/export/:id" '"name":"Summarizer Agent"' "$BUNDLE"
ORIG_NAME=$(echo "$BUNDLE" | python3 -c "import sys,json; print(json.load(sys.stdin)['name'])")
ORIG_TIER=$(echo "$BUNDLE" | python3 -c "import sys,json; print(json.load(sys.stdin)['tier'])")
R=$(acurl -X DELETE "$BASE/workspaces/$SUM_ID?confirm=true" \
-H "Authorization: Bearer $SUM_TOKEN" \
-H "X-Confirm-Name: Summarizer Agent")
check "DELETE /workspaces/:id" '"status":"removed"' "$R"
# Parent Echo must survive a child delete — list as Echo and expect count=1.
R=$(curl -s "$BASE/workspaces" -H "Authorization: Bearer $ECHO_TOKEN")
COUNT=$(echo "$R" | python3 -c "import sys,json; print(len(json.load(sys.stdin)))")
check "List after delete (count=1)" "1" "$COUNT"
# Test 22: Bundle round-trip — export → delete → import → verify same config.
# Summarizer's bundle was captured above; now delete the parent Echo (the only
# remaining workspace) so the import lands in a clean org, then re-import the
# Summarizer bundle.
echo ""
echo "--- Bundle Round-Trip Test ---"
# Delete the remaining parent Echo — use ECHO_TOKEN (per-workspace) for
# WorkspaceAuth and ADMIN_TOKEN for the AdminAuth layer.
R=$(acurl -X DELETE "$BASE/workspaces/$ECHO_ID?confirm=true" \
-H "Authorization: Bearer $ECHO_TOKEN" \
-H "X-Confirm-Name: Echo Agent v2")
check "DELETE /workspaces/:id" '"status":"removed"' "$R"
R=$(curl -s "$BASE/workspaces" -H "Authorization: Bearer $SUM_TOKEN")
COUNT=$(echo "$R" | python3 -c "import sys,json; print(len(json.load(sys.stdin)))")
check "List after delete (count=1)" "1" "$COUNT"
# Test 22: Bundle round-trip — export → delete → import → verify same config
echo ""
echo "--- Bundle Round-Trip Test ---"
# Export the summarizer workspace (#165 / PR #167 — admin-gated)
BUNDLE=$(curl -s "$BASE/bundles/export/$SUM_ID" -H "Authorization: Bearer $SUM_TOKEN")
check "GET /bundles/export/:id" '"name":"Summarizer Agent"' "$BUNDLE"
# Capture original config for comparison
ORIG_NAME=$(echo "$BUNDLE" | python3 -c "import sys,json; print(json.load(sys.stdin)['name'])")
ORIG_TIER=$(echo "$BUNDLE" | python3 -c "import sys,json; print(json.load(sys.stdin)['tier'])")
# Delete the workspace — use SUM_TOKEN (per-workspace) for WorkspaceAuth
# and ADMIN_TOKEN for the AdminAuth layer.
R=$(curl -s -X DELETE "$BASE/workspaces/$SUM_ID?confirm=true" \
-H "Authorization: Bearer $SUM_TOKEN" \
-H "X-Confirm-Name: Summarizer Agent")
check "Delete before re-import" '"status":"removed"' "$R"
# After deleting both workspaces, all per-workspace tokens are revoked.
+2 -7
View File
@@ -149,13 +149,8 @@ func main() {
result, err := db.DB.ExecContext(ctx, `DELETE FROM activity_logs WHERE created_at < now() - ($1 || ' days')::interval`, retentionDays)
if err != nil {
log.Printf("Activity log cleanup error: %v", err)
} else {
n, err := result.RowsAffected()
if err != nil {
log.Printf("Activity log cleanup RowsAffected error: %v", err)
} else if n > 0 {
log.Printf("Activity log cleanup: purged %d old entries", n)
}
} else if n, _ := result.RowsAffected(); n > 0 {
log.Printf("Activity log cleanup: purged %d old entries", n)
}
}
}
+4 -11
View File
@@ -3,7 +3,6 @@ package bundle
import (
"context"
"fmt"
"log"
"strings"
"git.moleculesai.app/molecule-ai/molecule-core/workspace-server/internal/db"
@@ -73,9 +72,7 @@ func Import(
}
}
// Store runtime in DB
if _, err := db.DB.ExecContext(ctx, `UPDATE workspaces SET runtime = $1 WHERE id = $2`, bundleRuntime, wsID); err != nil {
log.Printf("bundle import: failed to store runtime for workspace %s: %v", wsID, err)
}
_, _ = db.DB.ExecContext(ctx, `UPDATE workspaces SET runtime = $1 WHERE id = $2`, bundleRuntime, wsID)
// Provision the container if provisioner is available
if prov != nil {
@@ -95,9 +92,7 @@ func Import(
if err != nil {
markFailed(provCtx, wsID, broadcaster, err)
} else if url != "" {
if _, err := db.DB.ExecContext(provCtx, `UPDATE workspaces SET url = $1 WHERE id = $2`, url, wsID); err != nil {
log.Printf("bundle import: failed to store URL for workspace %s: %v", wsID, err)
}
db.DB.ExecContext(provCtx, `UPDATE workspaces SET url = $1 WHERE id = $2`, url, wsID)
}
}()
}
@@ -144,11 +139,9 @@ func markFailed(ctx context.Context, wsID string, broadcaster *events.Broadcaste
// markProvisionFailed in workspace-server/internal/handlers/
// workspace_provision_shared.go.
msg := err.Error()
if _, dbErr := db.DB.ExecContext(ctx,
db.DB.ExecContext(ctx,
`UPDATE workspaces SET status = $1, last_sample_error = $2, updated_at = now() WHERE id = $3`,
models.StatusFailed, msg, wsID); dbErr != nil {
log.Printf("bundle import: failed to mark workspace %s as failed: %v", wsID, dbErr)
}
models.StatusFailed, msg, wsID)
broadcaster.RecordAndBroadcast(ctx, string(events.EventWorkspaceProvisionFailed), wsID, map[string]interface{}{
"error": msg,
})
+9 -30
View File
@@ -82,10 +82,7 @@ func NewManager(proxy A2AProxy, broadcaster Broadcaster) *Manager {
log.Printf("Channels: failed to disable telegram chat_id=%s: %v", chatID, err)
return
}
rows, err := res.RowsAffected()
if err != nil {
log.Printf("Channels: disable telegram RowsAffected error chat_id=%s: %v", chatID, err)
} else if rows > 0 {
if rows, _ := res.RowsAffected(); rows > 0 {
log.Printf("Channels: disabled %d telegram channel(s) for chat_id=%s (bot removed)", rows, chatID)
// Reload so the in-memory poller map drops the now-disabled row.
m.Reload(ctx)
@@ -313,7 +310,7 @@ func (m *Manager) HandleInbound(ctx context.Context, ch ChannelRow, msg *Inbound
history := m.loadHistory(ctx, historyKey)
// Build A2A JSON-RPC payload
a2aBody, marshalErr := json.Marshal(map[string]interface{}{
a2aBody, _ := json.Marshal(map[string]interface{}{
"method": "message/send",
"params": map[string]interface{}{
"message": map[string]interface{}{
@@ -333,10 +330,6 @@ func (m *Manager) HandleInbound(ctx context.Context, ch ChannelRow, msg *Inbound
},
},
})
if marshalErr != nil {
log.Printf("Channels %s: json.Marshal a2aBody failed: %v", ch.ChannelType, marshalErr)
return fmt.Errorf("marshal a2a body: %w", marshalErr)
}
callerID := "channel:" + ch.ChannelType
@@ -396,13 +389,11 @@ func (m *Manager) HandleInbound(ctx context.Context, ch ChannelRow, msg *Inbound
// Update stats in DB
if db.DB != nil {
if _, err := db.DB.ExecContext(ctx, `
db.DB.ExecContext(ctx, `
UPDATE workspace_channels
SET last_message_at = now(), message_count = message_count + 1, updated_at = now()
WHERE id = $1
`, ch.ID); err != nil {
log.Printf("Channels: inbound stats update failed for channel %s: %v", ch.ID, err)
}
`, ch.ID)
}
// Broadcast event
@@ -443,13 +434,11 @@ func (m *Manager) SendOutbound(ctx context.Context, channelID string, text strin
}
if db.DB != nil {
if _, err := db.DB.ExecContext(ctx, `
db.DB.ExecContext(ctx, `
UPDATE workspace_channels
SET last_message_at = now(), message_count = message_count + 1, updated_at = now()
WHERE id = $1
`, channelID); err != nil {
log.Printf("Channels: outbound stats update failed for channel %s: %v", channelID, err)
}
`, channelID)
}
if m.broadcaster != nil {
@@ -519,20 +508,14 @@ func (m *Manager) FetchWorkspaceChannelContext(ctx context.Context, workspaceID
}
defer rows.Close()
if !rows.Next() {
if err := rows.Err(); err != nil {
log.Printf("ChannelManager: FetchWorkspaceChannelContext rows error for %s: %v", workspaceID, err)
}
return ""
}
var configJSON []byte
if err := rows.Scan(&configJSON); err != nil {
log.Printf("ChannelManager: FetchWorkspaceChannelContext scan error for %s: %v", workspaceID, err)
if rows.Scan(&configJSON) != nil {
return ""
}
var config map[string]interface{}
if err := json.Unmarshal(configJSON, &config); err != nil {
log.Printf("ChannelManager: unmarshal config: %v", err)
}
json.Unmarshal(configJSON, &config)
if err := DecryptSensitiveFields(config); err != nil {
return ""
}
@@ -669,16 +652,12 @@ func (m *Manager) appendHistory(ctx context.Context, key string, username, userM
if db.RDB == nil {
return
}
entry, marshalErr := json.Marshal(map[string]string{
entry, _ := json.Marshal(map[string]string{
"user": username,
"message": userMsg,
"reply": agentReply,
"time": time.Now().UTC().Format(time.RFC3339),
})
if marshalErr != nil {
log.Printf("appendHistory %s: json.Marshal entry failed: %v", key, marshalErr)
return
}
db.RDB.LPush(ctx, key, string(entry))
db.RDB.LTrim(ctx, key, 0, int64(maxHistoryEntries-1))
db.RDB.Expire(ctx, key, historyTTL)
+1 -6
View File
@@ -6,7 +6,6 @@ import (
"encoding/json"
"fmt"
"io"
"log"
"net/http"
"strings"
"time"
@@ -160,11 +159,7 @@ func (s *SlackAdapter) sendBotMessage(ctx context.Context, config map[string]int
payload["icon_emoji"] = iconEmoji
}
body, marshalErr := json.Marshal(payload)
if marshalErr != nil {
log.Printf("slack SendMessage: json.Marshal payload failed: %v", marshalErr)
return fmt.Errorf("slack: marshal payload: %w", marshalErr)
}
body, _ := json.Marshal(payload)
req, err := http.NewRequestWithContext(ctx, http.MethodPost, "https://slack.com/api/chat.postMessage", bytes.NewReader(body))
if err != nil {
return fmt.Errorf("slack: build request: %w", err)
@@ -482,14 +482,12 @@ func (t *TelegramAdapter) StartPolling(ctx context.Context, config map[string]in
if apiErr.Code == 429 {
retryAfter := time.Duration(apiErr.RetryAfter) * time.Second
log.Printf("Channels: Telegram poll rate-limited, sleeping %s", retryAfter)
timer := time.NewTimer(retryAfter)
select {
case <-ctx.Done():
timer.Stop()
return nil
case <-timer.C:
case <-time.After(retryAfter):
continue
}
continue
}
if apiErr.Code == 401 {
invalidateBot(token)
@@ -497,14 +495,12 @@ func (t *TelegramAdapter) StartPolling(ctx context.Context, config map[string]in
}
}
log.Printf("Channels: Telegram poll error: %v", err)
timer := time.NewTimer(telegramPollInterval)
select {
case <-ctx.Done():
timer.Stop()
return nil
case <-timer.C:
case <-time.After(telegramPollInterval):
continue
}
continue
}
for _, update := range updates {
@@ -375,30 +375,6 @@ func (h *WorkspaceHandler) proxyA2ARequest(ctx context.Context, workspaceID stri
Response: gin.H{"error": "access denied: workspaces cannot communicate per hierarchy rules"},
}
}
// #1953 cross-tenant isolation. CanCommunicate alone does NOT enforce
// org boundaries: its "root-level siblings — both have no parent" rule
// treats every tenant's org root as a sibling, so a caller that is an
// org root could resolve and route a2a to another tenant's org root
// (and resolveAgentURL accepts ANY workspace id with no org check).
// Gate on the SAME parent_id-chain org scoping the OFFSEC-015 broadcast
// fix uses: reject before resolveAgentURL when caller and target are in
// different orgs. Fail-closed — a DB error denies cross-org routing.
ok, err := sameOrg(ctx, db.DB, callerID, workspaceID)
if err != nil {
log.Printf("ProxyA2A: org-scope check failed %s → %s: %v — denying", callerID, workspaceID, err)
return 0, nil, &proxyA2AError{
Status: http.StatusForbidden,
Response: gin.H{"error": "access denied: org isolation check failed"},
}
}
if !ok {
log.Printf("ProxyA2A: cross-org routing denied %s → %s (#1953)", callerID, workspaceID)
return 0, nil, &proxyA2AError{
Status: http.StatusForbidden,
Response: gin.H{"error": "access denied: target workspace is in a different org"},
}
}
}
// Budget enforcement: reject A2A calls when the workspace has exceeded its
@@ -115,15 +115,12 @@ func (h *WorkspaceHandler) handleA2ADispatchError(ctx context.Context, workspace
if logActivity {
h.logA2ABusyQueued(ctx, workspaceID, callerID, body, a2aMethod, durationMs)
}
respBody, marshalErr := json.Marshal(gin.H{
respBody, _ := json.Marshal(gin.H{
"queued": true,
"queue_id": qid,
"queue_depth": depth,
"message": "workspace agent busy — request queued, will dispatch when capacity available",
})
if marshalErr != nil {
log.Printf("ProxyA2A %s: json.Marshal respBody failed: %v", workspaceID, marshalErr)
}
return http.StatusAccepted, respBody, nil
} else {
// Queue insert failed — fall through to legacy 503 behavior
@@ -426,34 +423,16 @@ func nilIfEmpty(s string) *string {
// (their next /registry/register will mint their first token, after
// which this branch never fires again for them).
//
// Post-RFC#637 addition: a request may instead be carrying a HUMAN's
// canvas-user identity (e.g. the 344a2623-… identity workspace from the
// RFC#637 rollout). That human sits OUTSIDE the workspace org hierarchy, so
// the returned isCanvasUser flag lets the A2A proxy bypass CanCommunicate for
// it. Canvas-user classification is decided by isGenuineCanvasUser using
// NON-FORGEABLE credentials only (see that function) — never by the caller's
// X-Workspace-ID alone, and never by a bare same-origin Host/Referer in a
// SaaS image (those are forgeable; see middleware.IsSameOriginCanvas).
//
// #1673: this canvas-user check is now evaluated BEFORE the HasAnyLiveToken
// peer-token contract. Previously it lived only in the !hasLive branch, so a
// canvas-user identity workspace that had acquired live tokens fell into the
// hasLive=true branch, which demands a bearer the canvas frontend never sends
// → silent 401 → the message was dropped before logA2AReceiveQueued wrote the
// activity_logs row, breaking canvas chat for poll-mode workspaces. A genuine
// canvas user is identified by the human's session/admin/org credential, which
// is independent of whether the identity workspace happens to hold peer tokens.
// Post-RFC#637 addition: when the tokenless workspace is accompanied by
// canvas or admin auth (same-origin request, admin bearer, or org-level
// token), the caller is identified as a canvas-user identity rather than
// a legacy peer agent. The returned isCanvasUser flag lets the A2A proxy
// bypass CanCommunicate for human users, who sit outside the workspace
// hierarchy.
//
// On auth failure this writes the 401 via c and returns an error so the
// handler aborts without running the proxy.
func validateCallerToken(ctx context.Context, c *gin.Context, callerID string) (isCanvasUser bool, err error) {
// Genuine canvas-user identity? Decided independently of the caller
// workspace's token state (the #1673 fix) and using only non-forgeable
// signals (the #1944 escalation guard).
if isGenuineCanvasUser(ctx, c) {
return true, nil
}
hasLive, dbErr := wsauth.HasAnyLiveToken(ctx, db.DB, callerID)
if dbErr != nil {
// Fail-open here matches the heartbeat path — A2A caller auth is
@@ -464,10 +443,22 @@ func validateCallerToken(ctx context.Context, c *gin.Context, callerID string) (
return false, nil
}
if !hasLive {
// Tokenless, non-canvas-user workspace — legacy / pre-upgrade peer.
// Grandfather it through (its next /registry/register mints its
// first token, after which it lands in the hasLive=true branch).
return false, nil
// Tokenless workspace — could be legacy/pre-upgrade caller or
// canvas-user identity. Distinguish by request auth signals.
if middleware.IsSameOriginCanvas(c) {
return true, nil
}
tok := wsauth.BearerTokenFromHeader(c.GetHeader("Authorization"))
if tok != "" {
adminSecret := os.Getenv("ADMIN_TOKEN")
if adminSecret != "" && subtle.ConstantTimeCompare([]byte(tok), []byte(adminSecret)) == 1 {
return true, nil
}
if _, _, _, err := orgtoken.Validate(ctx, db.DB, tok); err == nil {
return true, nil
}
}
return false, nil // legacy / pre-upgrade caller
}
tok := wsauth.BearerTokenFromHeader(c.GetHeader("Authorization"))
if tok == "" {
@@ -481,61 +472,6 @@ func validateCallerToken(ctx context.Context, c *gin.Context, callerID string) (
return false, nil
}
// isGenuineCanvasUser reports whether the request is a real human acting
// through the canvas UI (RFC#637 canvas-user identity), as opposed to a peer
// workspace agent. A true result lets the A2A proxy bypass CanCommunicate, so
// it MUST only accept signals an attacker on the platform network cannot forge:
//
// - A control-plane-verified canvas session: the WorkOS session cookie is
// confirmed upstream to belong to a MEMBER of THIS tenant's org
// (middleware.IsVerifiedCanvasSession → /cp/auth/tenant-member). This is
// the production SaaS canvas path.
// - An Authorization: Bearer matching ADMIN_TOKEN (break-glass / molecli).
// - An Authorization: Bearer matching a live org_api_tokens row (user-minted
// org-scoped API token).
//
// Deliberately NOT accepted as a canvas-user signal in a SaaS image:
//
// - A bare same-origin Host/Referer/Origin (middleware.IsSameOriginCanvas).
// Those headers are trivially forgeable by any container on the Docker
// network, and the combined-tenant image (CANVAS_PROXY_URL set) is exactly
// where a forged Referer + an arbitrary X-Workspace-ID could otherwise
// bypass CanCommunicate and reach cross-workspace A2A — the PR #1944
// privilege escalation. Same-origin is only honored as a fallback when CP
// session verification is NOT configured (self-hosted / dev), a
// single-tenant topology with no cross-tenant boundary to escalate across;
// even there the org hierarchy still owns intra-org routing.
//
// Note this classification is about the human's credential, not the caller
// workspace's X-Workspace-ID — so it never trusts an attacker-supplied caller
// ID, and it is independent of whether that workspace holds peer tokens.
func isGenuineCanvasUser(ctx context.Context, c *gin.Context) bool {
// Production SaaS: control-plane-verified org-member session cookie.
if middleware.IsVerifiedCanvasSession(c) {
return true
}
if tok := wsauth.BearerTokenFromHeader(c.GetHeader("Authorization")); tok != "" {
adminSecret := os.Getenv("ADMIN_TOKEN")
if adminSecret != "" && subtle.ConstantTimeCompare([]byte(tok), []byte(adminSecret)) == 1 {
return true
}
if _, _, _, err := orgtoken.Validate(ctx, db.DB, tok); err == nil {
return true
}
}
// Self-hosted / dev fallback ONLY: when upstream session verification is
// not configured there is no verified-cookie signal to use, and the
// deployment is single-tenant, so the forgeable same-origin check is an
// acceptable canvas signal. In SaaS (CP session configured) this branch is
// skipped, closing the forged-same-origin escalation.
if !middleware.CPSessionConfigured() && middleware.IsSameOriginCanvas(c) {
return true
}
return false
}
// errInvalidCallerToken is a sentinel for validateCallerToken's "missing
// token" branch so the handler-level guard can detect it without string
// matching (the wsauth errors are typed for the invalid case).
@@ -11,7 +11,6 @@ import (
"net/http"
"net/http/httptest"
"os"
"os/exec"
"strings"
"testing"
"time"
@@ -437,10 +436,6 @@ func TestProxyA2A_CallerIDPropagated(t *testing.T) {
WithArgs("ws-target").
WillReturnRows(sqlmock.NewRows([]string{"id", "parent_id"}).AddRow("ws-target", "ws-parent"))
// #1953 cross-tenant guard: same-org check after CanCommunicate. Both
// workspaces resolve to the same org root → routing allowed.
mockSameOrg(mock, "ws-caller", "ws-target", true)
expectBudgetCheck(mock, "ws-target")
// Expect activity log with source_id set
@@ -469,24 +464,6 @@ func TestProxyA2A_CallerIDPropagated(t *testing.T) {
}
}
// mockSameOrg sets up the two org-root recursive-CTE expectations that the
// #1953 cross-tenant guard in proxyA2ARequest runs after CanCommunicate passes.
// sameOrg=true returns the SAME root_id for both caller and target (same tenant);
// sameOrg=false returns different root_ids (cross-tenant → routing must be denied).
func mockSameOrg(mock sqlmock.Sqlmock, caller, target string, sameOrg bool) {
callerRoot := "org-root-shared"
targetRoot := "org-root-shared"
if !sameOrg {
targetRoot = "org-root-other-tenant"
}
mock.ExpectQuery("WITH RECURSIVE org_chain AS").
WithArgs(caller).
WillReturnRows(sqlmock.NewRows([]string{"root_id"}).AddRow(callerRoot))
mock.ExpectQuery("WITH RECURSIVE org_chain AS").
WithArgs(target).
WillReturnRows(sqlmock.NewRows([]string{"root_id"}).AddRow(targetRoot))
}
// mockCanCommunicate sets up sqlmock expectations for CanCommunicate(caller, target).
// allowed=true sets up rows that satisfy the access policy (siblings under same parent).
// allowed=false sets up rows that don't (different parents).
@@ -681,9 +658,6 @@ func TestProxyA2A_CallerIDDerivedFromBearer(t *testing.T) {
WithArgs("ws-target").
WillReturnRows(sqlmock.NewRows([]string{"id", "parent_id"}).AddRow("ws-target", "ws-parent"))
// 3b. #1953 cross-tenant guard — same org root → routing allowed.
mockSameOrg(mock, "ws-caller", "ws-target", true)
expectBudgetCheck(mock, "ws-target")
// 4. activity_logs INSERT — verify source_id arg is the derived ws-caller
@@ -1270,12 +1244,13 @@ func TestValidateCallerToken_WrongWorkspaceBindingRejected(t *testing.T) {
}
func TestValidateCallerToken_CanvasUser_AdminToken(t *testing.T) {
setupTestDB(t)
mock := setupTestDB(t)
setupTestRedis(t)
// #1673/#1944: the genuine-canvas-user check (admin bearer here) now runs
// BEFORE HasAnyLiveToken, so no SELECT COUNT(*) is issued — the human's
// credential, not the caller workspace's token state, decides canvas-user.
// Tokenless workspace
mock.ExpectQuery(`SELECT COUNT\(\*\) FROM workspace_auth_tokens`).
WithArgs("ws-canvas-admin").
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
t.Setenv("ADMIN_TOKEN", "admin-secret-42")
@@ -1301,9 +1276,10 @@ func TestValidateCallerToken_CanvasUser_OrgToken(t *testing.T) {
mock := setupTestDB(t)
setupTestRedis(t)
// #1673/#1944: the genuine-canvas-user check (org token here) now runs
// BEFORE HasAnyLiveToken, so the first DB query is orgtoken.Validate's
// lookup — there is no SELECT COUNT(*) expectation anymore.
// Tokenless workspace
mock.ExpectQuery(`SELECT COUNT\(\*\) FROM workspace_auth_tokens`).
WithArgs("ws-canvas-org").
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
// orgtoken.Validate lookup
mock.ExpectQuery(`SELECT id, prefix, org_id FROM org_api_tokens WHERE token_hash = .* AND revoked_at IS NULL`).
@@ -2365,197 +2341,6 @@ func TestProxyA2A_PollMode_ShortCircuits_NoSSRF_NoDispatch(t *testing.T) {
}
}
// stubVerifiedCPSession points VerifiedCPSession at a stub control-plane that
// confirms the given cookie belongs to a tenant-member, so tests can exercise
// the genuine (non-forgeable) canvas-session path end-to-end without a live CP.
// It sets CP_UPSTREAM_URL + MOLECULE_ORG_SLUG for the test's lifetime; the
// real middleware.VerifiedCPSession HTTP+cache code path runs unchanged.
func stubVerifiedCPSession(t *testing.T, member bool) {
t.Helper()
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
if member {
fmt.Fprint(w, `{"member":true,"user_id":"user-canvas-1"}`)
} else {
w.WriteHeader(http.StatusForbidden)
fmt.Fprint(w, `{"member":false}`)
}
}))
t.Cleanup(srv.Close)
t.Setenv("CP_UPSTREAM_URL", srv.URL)
t.Setenv("MOLECULE_ORG_SLUG", "test-tenant")
}
// TestProxyA2A_PollMode_CanvasUserWithVerifiedSession is the #1673 regression
// guard. A poll-mode canvas-user identity workspace that HAS acquired live
// tokens (the exact condition that made #1673 fire) sends a canvas message
// carrying a control-plane-verified session cookie but no bearer token. The
// fix must classify it as a canvas user BEFORE the HasAnyLiveToken peer-token
// contract, so the request is queued (200) and logA2AReceiveQueued writes the
// activity_logs row — instead of the pre-fix silent 401 that dropped the
// message before any row landed (breaking canvas chat + chat-history).
//
// Runs in a subprocess with CANVAS_PROXY_URL set so middleware.canvasProxyActive
// is true at package-init time (matching the combined-tenant image), proving the
// fix does not depend on disabling same-origin detection.
func TestProxyA2A_PollMode_CanvasUserWithVerifiedSession(t *testing.T) {
if os.Getenv("CANVAS_PROXY_URL") == "" {
cmd := exec.Command(os.Args[0], "-test.run=^TestProxyA2A_PollMode_CanvasUserWithVerifiedSession$", "-test.v")
cmd.Env = append(os.Environ(), "CANVAS_PROXY_URL=http://localhost")
out, err := cmd.CombinedOutput()
if err != nil {
t.Fatalf("subprocess test failed: %v\n%s", err, out)
}
return
}
stubVerifiedCPSession(t, true)
mock := setupTestDB(t)
setupTestRedis(t)
broadcaster := newTestBroadcaster()
handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", t.TempDir())
const wsTarget = "ws-poll-canvas-target"
const wsCanvasUser = "ws-canvas-user-344a"
// CRUCIAL: no SELECT COUNT(*) FROM workspace_auth_tokens expectation. The
// genuine-canvas-user check (verified session) must short-circuit BEFORE
// HasAnyLiveToken — that is the #1673 regression path. An identity
// workspace that already holds live tokens must NOT fall into the
// hasLive=true bearer-required branch.
// isCanvasUser=true → CanCommunicate is skipped (no parent_id lookups).
expectBudgetCheck(mock, wsTarget)
mock.ExpectQuery("SELECT delivery_mode FROM workspaces WHERE id").
WithArgs(wsTarget).
WillReturnRows(sqlmock.NewRows([]string{"delivery_mode"}).AddRow("poll"))
// logA2AReceiveQueued must fire synchronously and write the row.
mock.ExpectExec("INSERT INTO activity_logs").
WillReturnResult(sqlmock.NewResult(0, 1))
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Params = gin.Params{{Key: "id", Value: wsTarget}}
body := `{"jsonrpc":"2.0","id":"canvas-1","method":"message/send","params":{"message":{"role":"user","parts":[{"text":"hello from canvas"}]}}}`
req := httptest.NewRequest("POST", "/workspaces/"+wsTarget+"/a2a", bytes.NewBufferString(body))
req.Header.Set("Content-Type", "application/json")
req.Header.Set("X-Workspace-ID", wsCanvasUser)
// Verified canvas session cookie (the genuine, non-forgeable signal).
req.Header.Set("Cookie", "wos-session=valid-canvas-session-cookie")
// Same-origin headers, present as a real canvas request would send them —
// but they are NOT what authorizes the bypass here (the verified session is).
req.Host = "localhost"
req.Header.Set("Referer", "https://localhost/")
c.Request = req
handler.ProxyA2A(c)
time.Sleep(50 * time.Millisecond)
if w.Code != http.StatusOK {
t.Fatalf("expected 200 (queued) for canvas-user with verified session, got %d: %s", w.Code, w.Body.String())
}
var resp map[string]interface{}
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
t.Fatalf("response is not valid JSON: %v", err)
}
if resp["status"] != "queued" {
t.Errorf("response.status = %v, want %q", resp["status"], "queued")
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Errorf("unmet sqlmock expectations (activity_logs row must be written): %v", err)
}
}
// TestProxyA2A_ForgedSameOrigin_CannotBypassCanCommunicate is the security
// crux of the #1673 fix and the reason PR #1944 was held. In the combined-
// tenant SaaS image (CANVAS_PROXY_URL set, CP session verification configured),
// an attacker forges a same-origin request — correct Host + a matching
// `Referer: https://<host>/` — and supplies an arbitrary X-Workspace-ID naming
// a workspace it does not control, targeting a workspace it is NOT authorized
// to reach. It presents NO verified session cookie, NO admin token, NO org
// token.
//
// PR #1944's same-origin bypass would have classified this as a canvas user and
// skipped CanCommunicate, granting cross-workspace A2A — a privilege
// escalation. The safe fix must instead fall through to the standard
// peer-token contract and CanCommunicate, which rejects the cross-hierarchy
// call with 403. This test proves the escalation is closed.
func TestProxyA2A_ForgedSameOrigin_CannotBypassCanCommunicate(t *testing.T) {
if os.Getenv("CANVAS_PROXY_URL") == "" {
cmd := exec.Command(os.Args[0], "-test.run=^TestProxyA2A_ForgedSameOrigin_CannotBypassCanCommunicate$", "-test.v")
cmd.Env = append(os.Environ(), "CANVAS_PROXY_URL=http://localhost")
out, err := cmd.CombinedOutput()
if err != nil {
t.Fatalf("subprocess test failed: %v\n%s", err, out)
}
return
}
// SaaS image with CP session verification configured. The stub CP rejects
// any cookie as a non-member; the attacker sends none anyway. This asserts
// that with verification configured, same-origin alone is NOT a canvas
// signal (CPSessionConfigured()==true disables the dev fallback).
stubVerifiedCPSession(t, false)
mock := setupTestDB(t)
setupTestRedis(t)
broadcaster := newTestBroadcaster()
handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", t.TempDir())
const wsTarget = "ws-victim-target"
const wsForgedCaller = "ws-attacker-caller"
// validateCallerToken: not a genuine canvas user (no verified session, no
// admin/org token, and the dev same-origin fallback is disabled in SaaS).
// So it consults the peer-token contract: HasAnyLiveToken for the forged
// caller. Return 0 → tokenless legacy peer → grandfathered through token
// validation (isCanvasUser stays false). The request must then still be
// gated by CanCommunicate.
mock.ExpectQuery(`SELECT COUNT\(\*\) FROM workspace_auth_tokens`).
WithArgs(wsForgedCaller).
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
// CanCommunicate MUST run (the escalation guard) and DENY: caller and
// target sit under different parents.
mockCanCommunicate(mock, wsForgedCaller, wsTarget, false)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Params = gin.Params{{Key: "id", Value: wsTarget}}
body := `{"jsonrpc":"2.0","id":"exploit-1","method":"message/send","params":{"message":{"role":"user","parts":[{"text":"cross-workspace exploit"}]}}}`
req := httptest.NewRequest("POST", "/workspaces/"+wsTarget+"/a2a", bytes.NewBufferString(body))
req.Header.Set("Content-Type", "application/json")
// Arbitrary caller workspace the attacker does not own.
req.Header.Set("X-Workspace-ID", wsForgedCaller)
// Forged same-origin signals (the #1944 bypass vector).
req.Host = "localhost"
req.Header.Set("Referer", "https://localhost/")
req.Header.Set("Origin", "https://localhost")
// No Cookie / Authorization — no genuine canvas credential.
c.Request = req
handler.ProxyA2A(c)
if w.Code != http.StatusForbidden {
t.Fatalf("ESCALATION NOT CLOSED: forged same-origin + arbitrary X-Workspace-ID "+
"reached an unauthorized target with status %d (want 403): %s", w.Code, w.Body.String())
}
var resp map[string]interface{}
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
t.Fatalf("body not JSON: %v", err)
}
if !strings.Contains(fmt.Sprint(resp["error"]), "access denied") {
t.Errorf("expected an access-denied error from CanCommunicate, got %v", resp["error"])
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Errorf("unmet sqlmock expectations — CanCommunicate must have been consulted: %v", err)
}
}
// TestProxyA2A_PushMode_NoShortCircuit verifies the symmetric contract:
// a push-mode workspace (default) is NOT affected by the new short-circuit.
// It still proceeds to resolveAgentURL + dispatch. Without this guard, a
@@ -160,12 +160,10 @@ func EnqueueA2A(
}
// Return current queue depth for the caller's visibility.
if err := db.DB.QueryRowContext(ctx, `
_ = db.DB.QueryRowContext(ctx, `
SELECT COUNT(*) FROM a2a_queue
WHERE workspace_id = $1 AND status = 'queued'
`, workspaceID).Scan(&depth); err != nil {
log.Printf("A2AQueue: depth query failed for workspace %s: %v", workspaceID, err)
}
`, workspaceID).Scan(&depth)
log.Printf("A2AQueue: enqueued %s for workspace %s (priority=%d, depth=%d)", id, workspaceID, priority, depth)
return id, depth, nil
@@ -251,12 +249,10 @@ func MarkQueueItemFailed(ctx context.Context, id, errMsg string) {
// can see how many ahead of them.
func QueueDepth(ctx context.Context, workspaceID string) int {
var n int
if err := db.DB.QueryRowContext(ctx,
_ = db.DB.QueryRowContext(ctx,
`SELECT COUNT(*) FROM a2a_queue WHERE workspace_id = $1 AND status = 'queued'`,
workspaceID,
).Scan(&n); err != nil {
log.Printf("A2AQueue: QueueDepth query failed for workspace %s: %v", workspaceID, err)
}
).Scan(&n)
return n
}
@@ -419,14 +415,10 @@ func (h *WorkspaceHandler) stitchDrainResponseToDelegation(ctx context.Context,
return
}
responseText := extractResponseText(respBody)
respJSON, marshalErr := json.Marshal(map[string]interface{}{
respJSON, _ := json.Marshal(map[string]interface{}{
"text": responseText,
"delegation_id": delegationID,
})
if marshalErr != nil {
log.Printf("a2aQueue stitch %s: json.Marshal respJSON failed: %v", delegationID, marshalErr)
return
}
res, err := db.DB.ExecContext(ctx, `
UPDATE activity_logs
SET status = 'completed',
@@ -442,12 +434,7 @@ func (h *WorkspaceHandler) stitchDrainResponseToDelegation(ctx context.Context,
log.Printf("A2AQueue drain stitch: update failed for delegation %s: %v", delegationID, err)
return
}
rows, err := res.RowsAffected()
if err != nil {
log.Printf("A2AQueue drain stitch: RowsAffected error for delegation %s: %v", delegationID, err)
return
}
if rows == 0 {
if rows, _ := res.RowsAffected(); rows == 0 {
log.Printf("A2AQueue drain stitch: no delegate_result row for delegation %s (queued-row may not exist yet)", delegationID)
return
}
@@ -153,15 +153,7 @@ func queueRowAuthFields(ctx context.Context, queueID string) (callerID, workspac
if err != nil {
return "", "", err
}
callerID = ""
if callerNS.Valid {
callerID = callerNS.String
}
workspaceID = ""
if workspaceNS.Valid {
workspaceID = workspaceNS.String
}
return callerID, workspaceID, nil
return callerNS.String, workspaceNS.String, nil
}
// GetA2AQueueStatus handles GET /workspaces/:id/a2a/queue/:queue_id.
@@ -1,62 +1,9 @@
package handlers
import (
"context"
"testing"
"github.com/DATA-DOG/go-sqlmock"
)
// TestQueueRowAuthFields_NilSafeScan proves queueRowAuthFields returns empty
// strings (not a panic / garbage) when the a2a_queue row has NULL caller_id
// or workspace_id. Before the fix it dereferenced NullString.String directly,
// which is only the zero value when Valid is false but masked the NULL-vs-""
// distinction; the guard makes the intent explicit and safe.
func TestQueueRowAuthFields_NilSafeScan(t *testing.T) {
mock := setupTestDB(t)
queueID := "queue-123"
mock.ExpectQuery(`SELECT caller_id, workspace_id FROM a2a_queue WHERE id = \$1`).
WithArgs(queueID).
WillReturnRows(sqlmock.NewRows([]string{"caller_id", "workspace_id"}).AddRow(nil, nil))
caller, workspace, err := queueRowAuthFields(context.Background(), queueID)
if err != nil {
t.Fatalf("queueRowAuthFields returned error: %v", err)
}
if caller != "" {
t.Errorf("callerID = %q, want empty string for NULL caller_id", caller)
}
if workspace != "" {
t.Errorf("workspaceID = %q, want empty string for NULL workspace_id", workspace)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatalf("unmet expectations: %v", err)
}
}
// TestQueueRowAuthFields_PopulatedRow confirms the non-NULL path still returns
// the scanned values unchanged.
func TestQueueRowAuthFields_PopulatedRow(t *testing.T) {
mock := setupTestDB(t)
queueID := "queue-456"
mock.ExpectQuery(`SELECT caller_id, workspace_id FROM a2a_queue WHERE id = \$1`).
WithArgs(queueID).
WillReturnRows(sqlmock.NewRows([]string{"caller_id", "workspace_id"}).AddRow("caller-x", "ws-y"))
caller, workspace, err := queueRowAuthFields(context.Background(), queueID)
if err != nil {
t.Fatalf("queueRowAuthFields returned error: %v", err)
}
if caller != "caller-x" || workspace != "ws-y" {
t.Fatalf("got caller=%q workspace=%q, want caller-x / ws-y", caller, workspace)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatalf("unmet expectations: %v", err)
}
}
// TestExtractExpiresInSeconds covers the JSON parser used at enqueue time
// to honor a caller-specified TTL. Zero return = "no TTL" — caller leaves
// expires_at NULL on the queue row.
@@ -164,11 +164,7 @@ func (w *AgentMessageWriter) Send(
}
respPayload["parts"] = fileParts
}
respJSON, marshalErr := json.Marshal(respPayload)
if marshalErr != nil {
log.Printf("AgentMessageWriter %s: json.Marshal respPayload failed: %v", workspaceID, marshalErr)
return nil
}
respJSON, _ := json.Marshal(respPayload)
preview := textutil.TruncateRunes(message, 80)
if _, err := w.db.ExecContext(ctx, `
INSERT INTO activity_logs (workspace_id, activity_type, method, summary, response_body, status)
@@ -34,10 +34,7 @@ func (h *ApprovalsHandler) Create(c *gin.Context) {
return
}
ctxJSON, marshalErr := json.Marshal(body.Context)
if marshalErr != nil {
log.Printf("Approvals create %s: json.Marshal context failed: %v", workspaceID, marshalErr)
}
ctxJSON, _ := json.Marshal(body.Context)
if ctxJSON == nil {
ctxJSON = []byte("{}")
}
@@ -83,12 +80,10 @@ func (h *ApprovalsHandler) ListAll(c *gin.Context) {
ctx := c.Request.Context()
// Auto-expire stale approvals (older than 10 min)
if _, err := db.DB.ExecContext(ctx, `
db.DB.ExecContext(ctx, `
UPDATE approval_requests SET status = 'denied', decided_by = 'auto-expired', decided_at = now()
WHERE status = 'pending' AND created_at < now() - interval '10 minutes'
`); err != nil {
log.Printf("approvals: auto-expire failed: %v", err)
}
`)
rows, err := db.DB.QueryContext(ctx, `
SELECT a.id, a.workspace_id, w.name, a.action, a.reason, a.status, a.created_at
@@ -205,12 +200,7 @@ func (h *ApprovalsHandler) Decide(c *gin.Context) {
return
}
rows, err := result.RowsAffected()
if err != nil {
log.Printf("Approval decision RowsAffected error approval=%s workspace=%s: %v", approvalID, workspaceID, err)
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to update"})
return
}
rows, _ := result.RowsAffected()
if rows == 0 {
c.JSON(http.StatusNotFound, gin.H{"error": "approval not found or already decided"})
return
+1 -5
View File
@@ -344,11 +344,7 @@ func computeAuditHMAC(key []byte, ev *auditEventRow) string {
"timestamp": ev.Timestamp.UTC().Format("2006-01-02T15:04:05Z"),
}
payload, marshalErr := json.Marshal(canonical) // compact, sorted keys
if marshalErr != nil {
log.Printf("auditChainHash: json.Marshal canonical failed: %v", marshalErr)
return ""
}
payload, _ := json.Marshal(canonical) // compact, sorted keys
mac := hmac.New(sha256.New, key)
mac.Write(payload)
return hex.EncodeToString(mac.Sum(nil))
+12 -55
View File
@@ -26,10 +26,6 @@ type ChannelHandler struct {
manager *channels.Manager
}
// channelSlugRe matches valid agent slugs used in [slug] routing.
// Compiled once at init to avoid recompilation on every webhook call.
var channelSlugRe = regexp.MustCompile(`^[a-zA-Z0-9 _-]+$`)
// NewChannelHandler creates a channel handler with the given manager.
func NewChannelHandler(manager *channels.Manager) *ChannelHandler {
return &ChannelHandler{manager: manager}
@@ -71,9 +67,7 @@ func (h *ChannelHandler) List(c *gin.Context) {
}
var config map[string]interface{}
if err := json.Unmarshal(configJSON, &config); err != nil {
log.Printf("Channels: unmarshal config for channel %s: %v", id, err)
}
json.Unmarshal(configJSON, &config)
// #319: decrypt sensitive fields first so the mask operates on
// plaintext (first-4 / last-4 of the real token, not the ciphertext
// prefix). Decrypt errors are logged but non-fatal — List must keep
@@ -92,9 +86,7 @@ func (h *ChannelHandler) List(c *gin.Context) {
}
var allowed []string
if err := json.Unmarshal(allowedJSON, &allowed); err != nil {
log.Printf("Channels: unmarshal allowed_users for channel %s: %v", id, err)
}
json.Unmarshal(allowedJSON, &allowed)
entry := map[string]interface{}{
"id": id,
@@ -169,18 +161,8 @@ func (h *ChannelHandler) Create(c *gin.Context) {
return
}
configJSON, marshalErr := json.Marshal(body.Config)
if marshalErr != nil {
log.Printf("Channels create %s: json.Marshal config failed: %v", workspaceID, marshalErr)
c.JSON(http.StatusInternalServerError, gin.H{"error": "marshal config failed"})
return
}
allowedJSON, marshalErr := json.Marshal(body.AllowedUsers)
if marshalErr != nil {
log.Printf("Channels create %s: json.Marshal allowed_users failed: %v", workspaceID, marshalErr)
c.JSON(http.StatusInternalServerError, gin.H{"error": "marshal allowed_users failed"})
return
}
configJSON, _ := json.Marshal(body.Config)
allowedJSON, _ := json.Marshal(body.AllowedUsers)
enabled := true
if body.Enabled != nil {
enabled = *body.Enabled
@@ -235,21 +217,11 @@ func (h *ChannelHandler) Update(c *gin.Context) {
c.JSON(http.StatusInternalServerError, gin.H{"error": "encrypt failed"})
return
}
j, marshalErr := json.Marshal(body.Config)
if marshalErr != nil {
log.Printf("Channels update %s: json.Marshal config failed: %v", workspaceID, marshalErr)
c.JSON(http.StatusInternalServerError, gin.H{"error": "marshal config failed"})
return
}
j, _ := json.Marshal(body.Config)
configArg = string(j)
}
if body.AllowedUsers != nil {
j, marshalErr := json.Marshal(body.AllowedUsers)
if marshalErr != nil {
log.Printf("Channels update %s: json.Marshal allowed_users failed: %v", workspaceID, marshalErr)
c.JSON(http.StatusInternalServerError, gin.H{"error": "marshal allowed_users failed"})
return
}
j, _ := json.Marshal(body.AllowedUsers)
allowedArg = string(j)
}
@@ -266,13 +238,7 @@ func (h *ChannelHandler) Update(c *gin.Context) {
return
}
n, err := result.RowsAffected()
if err != nil {
log.Printf("Channel update RowsAffected error channel=%s workspace=%s: %v", channelID, workspaceID, err)
c.JSON(http.StatusInternalServerError, gin.H{"error": "update failed"})
return
}
if n == 0 {
if n, _ := result.RowsAffected(); n == 0 {
c.JSON(http.StatusNotFound, gin.H{"error": "channel not found"})
return
}
@@ -297,13 +263,7 @@ func (h *ChannelHandler) Delete(c *gin.Context) {
return
}
n, err := result.RowsAffected()
if err != nil {
log.Printf("Channel delete RowsAffected error channel=%s workspace=%s: %v", channelID, workspaceID, err)
c.JSON(http.StatusInternalServerError, gin.H{"error": "delete failed"})
return
}
if n == 0 {
if n, _ := result.RowsAffected(); n == 0 {
c.JSON(http.StatusNotFound, gin.H{"error": "channel not found"})
return
}
@@ -504,10 +464,11 @@ func (h *ChannelHandler) Webhook(c *gin.Context) {
// in a shared channel and route to a specific agent.
targetSlug := ""
routedText := msg.Text
validSlugRe := regexp.MustCompile(`^[a-zA-Z0-9 _-]+$`)
if len(msg.Text) > 2 && msg.Text[0] == '[' {
if idx := strings.Index(msg.Text, "]"); idx > 1 && idx < 40 {
candidate := strings.ToLower(strings.TrimSpace(msg.Text[1:idx]))
if channelSlugRe.MatchString(candidate) {
if validSlugRe.MatchString(candidate) {
targetSlug = candidate
routedText = strings.TrimSpace(msg.Text[idx+1:])
if routedText == "" {
@@ -538,12 +499,8 @@ func (h *ChannelHandler) Webhook(c *gin.Context) {
if err := rows.Scan(&row.ID, &row.WorkspaceID, &row.ChannelType, &configJSON, &row.Enabled, &allowedJSON); err != nil {
continue
}
if err := json.Unmarshal(configJSON, &row.Config); err != nil {
log.Printf("Channels: unmarshal config for webhook row %s: %v", row.ID, err)
}
if err := json.Unmarshal(allowedJSON, &row.AllowedUsers); err != nil {
log.Printf("Channels: unmarshal allowed_users for webhook row %s: %v", row.ID, err)
}
json.Unmarshal(configJSON, &row.Config)
json.Unmarshal(allowedJSON, &row.AllowedUsers)
if err := channels.DecryptSensitiveFields(row.Config); err != nil {
log.Printf("Channels: decrypt webhook row %s: %v", row.ID, err)
continue
@@ -229,12 +229,7 @@ func (h *CheckpointsHandler) Delete(c *gin.Context) {
return
}
n, err := result.RowsAffected()
if err != nil {
log.Printf("Delete checkpoints RowsAffected error workspace=%s wf=%s: %v", workspaceID, workflowID, err)
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to delete checkpoints"})
return
}
n, _ := result.RowsAffected()
if n == 0 {
c.JSON(http.StatusNotFound, gin.H{"error": "no checkpoints found for workflow"})
return
@@ -1,427 +0,0 @@
package handlers
// cross_tenant_isolation_test.go — #1953 regression tests.
//
// Three workspace-server paths historically derived an "org-root sibling set"
// as `WHERE parent_id IS NULL`, which matches EVERY tenant's org root (the
// workspaces table has no org_id column) → cross-tenant data exposure:
//
// 1. GET /registry/:id/peers (discovery.Peers)
// 2. MCP toolListPeers (mcp_tools.toolListPeers)
// 3. a2a routing (a2a_proxy.proxyA2ARequest → resolveAgentURL)
//
// These tests assert that a workspace in a DIFFERENT org is never returned as a
// peer and that a2a refuses to resolve/route to a workspace outside the caller's
// org, while same-org peers/targets still work. They reuse the SAME parent_id-
// chain org scoping the OFFSEC-015 broadcast fix introduced (org_scope.go).
import (
"bytes"
"context"
"database/sql"
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"git.moleculesai.app/molecule-ai/molecule-core/workspace-server/internal/db"
"github.com/DATA-DOG/go-sqlmock"
"github.com/gin-gonic/gin"
)
// dbHandleForTest returns the global sqlmock-backed *sql.DB that setupTestDB
// installs, for tests that need to hand a *sql.DB to a component (e.g.
// MCPHandler.database, sameOrg) rather than relying on the package-global.
func dbHandleForTest() *sql.DB { return db.DB }
// peerColsForIsolation matches queryPeerMaps' SELECT column set.
var peerColsForIsolation = []string{
"id", "name", "role", "tier", "status", "agent_card", "url", "parent_id", "active_tasks",
}
// -------------------------------------------------------------------------
// Path 1: GET /registry/:id/peers — discovery.Peers
// -------------------------------------------------------------------------
// TestPeers_CrossTenant_OrgRootNotLeaked is the core #1953 regression for the
// discovery path. The caller is an org root (parent_id IS NULL). Pre-fix the
// handler ran `SELECT ... WHERE w.parent_id IS NULL AND w.id != $1`, returning
// every OTHER tenant's org root as a "sibling" peer. Post-fix an org-root caller
// issues NO sibling query — its only peers are its own children. If the handler
// regressed and issued the cross-tenant sibling query, sqlmock would report an
// unexpected query (the expectation below is intentionally NOT registered) and
// the test fails.
func TestPeers_CrossTenant_OrgRootNotLeaked(t *testing.T) {
mock := setupTestDB(t)
setupTestRedis(t)
handler := NewDiscoveryHandler()
// Behavioural leak test: register the OLD leaky `parent_id IS NULL` sibling
// query so that IF the handler still issues it, it returns another tenant's
// org root (org-b-root). The fix removes that query for an org-root caller,
// so org-b-root must never appear in the output. Unordered matching makes
// the leaky-sibling expectation optional — the fix simply never consumes it.
mock.MatchExpectationsInOrder(false)
caller := "org-a-root" // parent_id IS NULL — an org root for tenant A
// parent_id lookup → NULL (caller is an org root)
mock.ExpectQuery("SELECT parent_id FROM workspaces WHERE id =").
WithArgs(caller).
WillReturnRows(sqlmock.NewRows([]string{"parent_id"}).AddRow(nil))
// LEAKY sibling query (pre-fix). Returns a DIFFERENT tenant's org root.
// The fix must NOT issue this query; if it does, org-b-root leaks into the
// peer list and the output assertion below fails.
mock.ExpectQuery("SELECT w.id, w.name.*WHERE w.parent_id IS NULL AND w.id != \\$1").
WithArgs(caller).
WillReturnRows(sqlmock.NewRows(peerColsForIsolation).
AddRow("org-b-root", "Org B Root", "lead", 0, "online", []byte("null"), "http://b-root", nil, 0))
// Children query — caller's own org-A children only. Return one child.
mock.ExpectQuery("SELECT w.id, w.name.*WHERE w.parent_id = \\$1 AND w.id != \\$2").
WithArgs(caller, caller).
WillReturnRows(sqlmock.NewRows(peerColsForIsolation).
AddRow("org-a-child", "Org A Child", "worker", 1, "online", []byte("null"), "http://a-child", caller, 0))
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Params = gin.Params{{Key: "id", Value: caller}}
c.Request = httptest.NewRequest("GET", "/registry/"+caller+"/peers", nil)
handler.Peers(c)
if w.Code != http.StatusOK {
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
}
var peers []map[string]interface{}
if err := json.Unmarshal(w.Body.Bytes(), &peers); err != nil {
t.Fatalf("failed to parse response: %v", err)
}
// The other-tenant org root must NEVER appear; only the same-org child.
for _, p := range peers {
if id, _ := p["id"].(string); id == "org-b-root" {
t.Fatalf("cross-tenant leak (#1953): org-b-root appeared in org-a-root's peer list: %v", peers)
}
}
if len(peers) != 1 {
t.Fatalf("expected exactly 1 peer (same-org child), got %d: %v", len(peers), peers)
}
// NOTE: ExpectationsWereMet is intentionally NOT asserted — the leaky
// sibling expectation is deliberately left unconsumed by the fixed path.
}
// TestPeers_SameOrg_SiblingsStillWork is the positive companion: a non-root
// child caller still sees its same-org siblings, children, and parent. This
// guards against the fix over-scoping and breaking legitimate intra-org
// discovery.
func TestPeers_SameOrg_SiblingsStillWork(t *testing.T) {
mock := setupTestDB(t)
setupTestRedis(t)
handler := NewDiscoveryHandler()
caller := "org-a-child-1"
parent := "org-a-root"
mock.ExpectQuery("SELECT parent_id FROM workspaces WHERE id =").
WithArgs(caller).
WillReturnRows(sqlmock.NewRows([]string{"parent_id"}).AddRow(parent))
// Siblings — scoped to the shared parent (one tenant).
mock.ExpectQuery("SELECT w.id, w.name.*WHERE w.parent_id = \\$1 AND w.id != \\$2").
WithArgs(parent, caller).
WillReturnRows(sqlmock.NewRows(peerColsForIsolation).
AddRow("org-a-child-2", "Org A Sibling", "worker", 1, "online", []byte("null"), "http://a-sib", parent, 0))
// Children — none.
mock.ExpectQuery("SELECT w.id, w.name.*WHERE w.parent_id = \\$1 AND w.id != \\$2 AND w.status").
WithArgs(caller, caller).
WillReturnRows(sqlmock.NewRows(peerColsForIsolation))
// Parent.
mock.ExpectQuery("SELECT w.id, w.name.*WHERE w.id = \\$1 AND w.id != \\$2 AND w.status").
WithArgs(parent, caller).
WillReturnRows(sqlmock.NewRows(peerColsForIsolation).
AddRow(parent, "Org A Root", "lead", 0, "online", []byte("null"), "http://a-root", nil, 0))
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Params = gin.Params{{Key: "id", Value: caller}}
c.Request = httptest.NewRequest("GET", "/registry/"+caller+"/peers", nil)
handler.Peers(c)
if w.Code != http.StatusOK {
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
}
var peers []map[string]interface{}
if err := json.Unmarshal(w.Body.Bytes(), &peers); err != nil {
t.Fatalf("failed to parse response: %v", err)
}
// Sibling + parent = 2 same-org peers.
if len(peers) != 2 {
t.Fatalf("expected 2 same-org peers (sibling + parent), got %d: %v", len(peers), peers)
}
names := map[string]bool{}
for _, p := range peers {
names[fmt.Sprint(p["name"])] = true
}
if !names["Org A Sibling"] || !names["Org A Root"] {
t.Errorf("expected same-org sibling + parent in peer list, got %v", names)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Errorf("unmet sqlmock expectations: %v", err)
}
}
// -------------------------------------------------------------------------
// Path 2: MCP toolListPeers — mcp_tools.toolListPeers
// -------------------------------------------------------------------------
// mcpPeerCols matches toolListPeers' SELECT column set.
var mcpPeerCols = []string{"id", "name", "role", "status", "tier"}
// TestToolListPeers_CrossTenant_OrgRootNotLeaked is the #1953 regression for
// the MCP path. Same shape as the discovery test: an org-root caller must NOT
// enumerate other tenants' org roots. The cross-tenant `parent_id IS NULL`
// sibling query is intentionally not registered, so if it runs sqlmock fails.
func TestToolListPeers_CrossTenant_OrgRootNotLeaked(t *testing.T) {
mock := setupTestDB(t)
mock.MatchExpectationsInOrder(false)
h := &MCPHandler{database: dbHandleForTest()}
caller := "org-a-root"
mock.ExpectQuery("SELECT parent_id FROM workspaces WHERE id =").
WithArgs(caller).
WillReturnRows(sqlmock.NewRows([]string{"parent_id"}).AddRow(nil))
// LEAKY sibling query (pre-fix). Returns another tenant's org root. The fix
// must NOT issue this for an org-root caller; if it does, org-b-root leaks
// into the output and the assertion below fails. Left optional via
// unordered matching, so the fixed path simply never consumes it.
mock.ExpectQuery("WHERE w.parent_id IS NULL AND w.id != \\$1").
WithArgs(caller).
WillReturnRows(sqlmock.NewRows(mcpPeerCols).
AddRow("org-b-root", "Org B Root", "lead", "online", 0))
// Children — caller's own org-A children only.
mock.ExpectQuery("WHERE w.parent_id = \\$1 AND w.status").
WithArgs(caller).
WillReturnRows(sqlmock.NewRows(mcpPeerCols).
AddRow("org-a-child", "Org A Child", "worker", "online", 1))
out, err := h.toolListPeers(context.Background(), caller)
if err != nil {
t.Fatalf("toolListPeers returned error: %v", err)
}
if strings.Contains(out, "org-b-root") || strings.Contains(out, "Org B Root") {
t.Fatalf("cross-tenant leak (#1953): another tenant's org root appeared in toolListPeers output:\n%s", out)
}
if !strings.Contains(out, "org-a-child") {
t.Errorf("same-org child missing from toolListPeers output:\n%s", out)
}
// ExpectationsWereMet intentionally NOT asserted — leaky sibling expectation
// is deliberately left unconsumed by the fixed path.
}
// TestToolListPeers_SameOrg_SiblingsStillWork — positive companion for the MCP
// path: a non-root child still enumerates its same-org siblings + children + parent.
func TestToolListPeers_SameOrg_SiblingsStillWork(t *testing.T) {
mock := setupTestDB(t)
h := &MCPHandler{database: dbHandleForTest()}
caller := "org-a-child-1"
parent := "org-a-root"
mock.ExpectQuery("SELECT parent_id FROM workspaces WHERE id =").
WithArgs(caller).
WillReturnRows(sqlmock.NewRows([]string{"parent_id"}).AddRow(parent))
// Siblings — scoped to shared parent.
mock.ExpectQuery("WHERE w.parent_id = \\$1 AND w.id != \\$2 AND w.status").
WithArgs(parent, caller).
WillReturnRows(sqlmock.NewRows(mcpPeerCols).
AddRow("org-a-child-2", "Org A Sibling", "worker", "online", 1))
// Children — none.
mock.ExpectQuery("WHERE w.parent_id = \\$1 AND w.status").
WithArgs(caller).
WillReturnRows(sqlmock.NewRows(mcpPeerCols))
// Parent.
mock.ExpectQuery("WHERE w.id = \\$1 AND w.status").
WithArgs(parent).
WillReturnRows(sqlmock.NewRows(mcpPeerCols).
AddRow(parent, "Org A Root", "lead", "online", 0))
out, err := h.toolListPeers(context.Background(), caller)
if err != nil {
t.Fatalf("toolListPeers returned error: %v", err)
}
if !strings.Contains(out, "Org A Sibling") || !strings.Contains(out, "Org A Root") {
t.Errorf("expected same-org sibling + parent in toolListPeers output:\n%s", out)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Errorf("unmet sqlmock expectations: %v", err)
}
}
// -------------------------------------------------------------------------
// Path 3: a2a routing — a2a_proxy.proxyA2ARequest / resolveAgentURL
// -------------------------------------------------------------------------
// TestProxyA2A_CrossTenant_RoutingDenied is the #1953 regression for a2a
// routing. Caller and target are both org roots (parent_id IS NULL) belonging
// to DIFFERENT tenants. Pre-fix, CanCommunicate's "root-level siblings" rule
// waved this through and resolveAgentURL routed to the foreign tenant. Post-fix
// the org-scope guard resolves each to a different org root and returns 403
// BEFORE resolveAgentURL/dispatch.
func TestProxyA2A_CrossTenant_RoutingDenied(t *testing.T) {
mock := setupTestDB(t)
mr := setupTestRedis(t)
broadcaster := newTestBroadcaster()
handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", t.TempDir())
caller := "org-a-root"
target := "org-b-root" // different tenant
// A URL exists for the target; the guard must deny BEFORE it is used.
mr.Set(fmt.Sprintf("ws:%s:url", target), "http://localhost:1")
// CanCommunicate: both root-level (parent_id NULL) → its weak "root-level
// siblings" rule ALLOWS this. The org guard must catch it afterward.
mock.ExpectQuery("SELECT id, parent_id FROM workspaces WHERE id = ").
WithArgs(caller).
WillReturnRows(sqlmock.NewRows([]string{"id", "parent_id"}).AddRow(caller, nil))
mock.ExpectQuery("SELECT id, parent_id FROM workspaces WHERE id = ").
WithArgs(target).
WillReturnRows(sqlmock.NewRows([]string{"id", "parent_id"}).AddRow(target, nil))
// #1953 org-scope guard: caller resolves to org-a-root, target to org-b-root
// → different orgs → 403. (Each org root resolves to itself.)
mock.ExpectQuery("WITH RECURSIVE org_chain AS").
WithArgs(caller).
WillReturnRows(sqlmock.NewRows([]string{"root_id"}).AddRow(caller))
mock.ExpectQuery("WITH RECURSIVE org_chain AS").
WithArgs(target).
WillReturnRows(sqlmock.NewRows([]string{"root_id"}).AddRow(target))
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Params = gin.Params{{Key: "id", Value: target}}
body := `{"method":"message/send","params":{"message":{"role":"user","parts":[{"text":"cross-tenant"}]}}}`
c.Request = httptest.NewRequest("POST", "/workspaces/"+target+"/a2a", bytes.NewBufferString(body))
c.Request.Header.Set("Content-Type", "application/json")
c.Request.Header.Set("X-Workspace-ID", caller)
handler.ProxyA2A(c)
if w.Code != http.StatusForbidden {
t.Fatalf("expected 403 for cross-tenant a2a routing, got %d: %s", w.Code, w.Body.String())
}
var resp map[string]interface{}
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
t.Fatalf("body not JSON: %v", err)
}
if msg, _ := resp["error"].(string); !strings.Contains(msg, "different org") {
t.Errorf("expected cross-org denial message, got %v", resp["error"])
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Errorf("unmet sqlmock expectations: %v", err)
}
}
// TestResolveAgentURL_CrossTenant_RejectedViaSameOrg is a direct unit test of
// the sameOrg primitive that gates resolveAgentURL: a target in a different org
// must be reported as NOT same-org, so the a2a guard rejects it before
// resolveAgentURL is ever called.
func TestResolveAgentURL_CrossTenant_RejectedViaSameOrg(t *testing.T) {
mock := setupTestDB(t)
caller := "org-a-root"
target := "org-b-root"
mock.ExpectQuery("WITH RECURSIVE org_chain AS").
WithArgs(caller).
WillReturnRows(sqlmock.NewRows([]string{"root_id"}).AddRow(caller))
mock.ExpectQuery("WITH RECURSIVE org_chain AS").
WithArgs(target).
WillReturnRows(sqlmock.NewRows([]string{"root_id"}).AddRow(target))
ok, err := sameOrg(context.Background(), dbHandleForTest(), caller, target)
if err != nil {
t.Fatalf("sameOrg returned unexpected error: %v", err)
}
if ok {
t.Errorf("expected cross-tenant workspaces to be reported as DIFFERENT orgs, got sameOrg=true")
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Errorf("unmet sqlmock expectations: %v", err)
}
}
// TestProxyA2A_SameOrg_RoutingAllowed — positive companion for a2a: two
// same-org siblings route successfully (mirrors TestProxyA2A_CallerIDPropagated
// but named to document the #1953 same-org allow path).
func TestProxyA2A_SameOrg_RoutingAllowed(t *testing.T) {
mock := setupTestDB(t)
mr := setupTestRedis(t)
allowLoopbackForTest(t)
broadcaster := newTestBroadcaster()
handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", t.TempDir())
waitForHandlerAsyncBeforeDBCleanup(t, handler)
caller := "org-a-child-1"
target := "org-a-child-2"
parent := "org-a-root"
agentServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
fmt.Fprint(w, `{"jsonrpc":"2.0","id":"1","result":{}}`)
}))
defer agentServer.Close()
mr.Set(fmt.Sprintf("ws:%s:url", target), agentServer.URL)
// CanCommunicate — siblings under shared parent.
mock.ExpectQuery("SELECT id, parent_id FROM workspaces WHERE id = ").
WithArgs(caller).
WillReturnRows(sqlmock.NewRows([]string{"id", "parent_id"}).AddRow(caller, parent))
mock.ExpectQuery("SELECT id, parent_id FROM workspaces WHERE id = ").
WithArgs(target).
WillReturnRows(sqlmock.NewRows([]string{"id", "parent_id"}).AddRow(target, parent))
// #1953 org guard — both resolve to the same org root → allowed.
mock.ExpectQuery("WITH RECURSIVE org_chain AS").
WithArgs(caller).
WillReturnRows(sqlmock.NewRows([]string{"root_id"}).AddRow(parent))
mock.ExpectQuery("WITH RECURSIVE org_chain AS").
WithArgs(target).
WillReturnRows(sqlmock.NewRows([]string{"root_id"}).AddRow(parent))
expectBudgetCheck(mock, target)
mock.ExpectExec("INSERT INTO activity_logs").WillReturnResult(sqlmock.NewResult(0, 1))
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Params = gin.Params{{Key: "id", Value: target}}
body := `{"method":"message/send","params":{"message":{"role":"user","parts":[{"text":"same-org"}]}}}`
c.Request = httptest.NewRequest("POST", "/workspaces/"+target+"/a2a", bytes.NewBufferString(body))
c.Request.Header.Set("Content-Type", "application/json")
c.Request.Header.Set("X-Workspace-ID", caller)
handler.ProxyA2A(c)
time.Sleep(50 * time.Millisecond) // allow the async logA2ASuccess INSERT to flush
if w.Code != http.StatusOK {
t.Fatalf("expected 200 for same-org a2a routing, got %d: %s", w.Code, w.Body.String())
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Errorf("unmet sqlmock expectations: %v", err)
}
}
@@ -57,18 +57,10 @@ func pushDelegationResultToInbox(ctx context.Context, sourceID, delegationID, st
"text": responsePreview,
"delegation_id": delegationID,
}
respJSON, marshalErr := json.Marshal(respPayload)
if marshalErr != nil {
log.Printf("Delegation %s: json.Marshal respPayload failed: %v", delegationID, marshalErr)
return
}
reqJSON, marshalErr := json.Marshal(map[string]interface{}{
respJSON, _ := json.Marshal(respPayload)
reqJSON, _ := json.Marshal(map[string]interface{}{
"delegation_id": delegationID,
})
if marshalErr != nil {
log.Printf("Delegation %s: json.Marshal reqPayload failed: %v", delegationID, marshalErr)
return
}
logStatus := "ok"
if status == "failed" {
logStatus = "error"
@@ -173,7 +165,7 @@ func (h *DelegationHandler) Delegate(c *gin.Context) {
// check_task_status returned status='queued' forever even after a
// real reply landed). messageId mirrors delegation_id so the
// platform's idempotency-key extraction also keys off the same id.
a2aBody, marshalErr := json.Marshal(map[string]interface{}{
a2aBody, _ := json.Marshal(map[string]interface{}{
"method": "message/send",
"params": map[string]interface{}{
"message": map[string]interface{}{
@@ -184,9 +176,6 @@ func (h *DelegationHandler) Delegate(c *gin.Context) {
},
},
})
if marshalErr != nil {
log.Printf("Delegation %s: json.Marshal a2aBody failed: %v", delegationID, marshalErr)
}
// Fire-and-forget: send A2A in a background goroutine.
//
@@ -272,12 +261,10 @@ func lookupIdempotentDelegation(ctx context.Context, c *gin.Context, sourceID, i
return false
}
if existingStatus == "failed" {
if _, err := db.DB.ExecContext(ctx, `
_, _ = db.DB.ExecContext(ctx, `
DELETE FROM activity_logs
WHERE workspace_id = $1 AND idempotency_key = $2 AND status = 'failed'
`, sourceID, idempotencyKey); err != nil {
log.Printf("delegation: failed to clean up failed idempotency row for %s/%s: %v", sourceID, idempotencyKey, err)
}
`, sourceID, idempotencyKey)
return false
}
c.JSON(http.StatusOK, gin.H{
@@ -315,24 +302,16 @@ const (
// insertDelegationRow stores the pending delegation row. See
// insertDelegationOutcome for the three possible return values.
func insertDelegationRow(ctx context.Context, c *gin.Context, sourceID string, body delegateRequest, delegationID string) insertDelegationOutcome {
taskJSON, marshalErr := json.Marshal(map[string]interface{}{
taskJSON, _ := json.Marshal(map[string]interface{}{
"task": body.Task,
"delegation_id": delegationID,
})
if marshalErr != nil {
log.Printf("Delegation %s: json.Marshal taskJSON failed: %v", delegationID, marshalErr)
return insertTrackingUnavailable
}
// Store delegation_id in response_body so agent check_delegation_status
// (which reads response_body->>delegation_id) can locate this row even
// when request_body hasn't propagated yet. Fixes mc#984.
respJSON, marshalErr := json.Marshal(map[string]interface{}{
respJSON, _ := json.Marshal(map[string]interface{}{
"delegation_id": delegationID,
})
if marshalErr != nil {
log.Printf("Delegation %s: json.Marshal respJSON failed: %v", delegationID, marshalErr)
return insertTrackingUnavailable
}
var idemArg interface{}
if body.IdempotencyKey != "" {
idemArg = body.IdempotencyKey
@@ -435,12 +414,10 @@ func (h *DelegationHandler) executeDelegation(ctx context.Context, sourceID, tar
if proxyErr != nil && isTransientProxyError(proxyErr) && len(respBody) == 0 {
log.Printf("Delegation %s: first attempt failed (%s) — retrying in %s after reactive URL refresh",
delegationID, proxyErr.Error(), delegationRetryDelay)
timer := time.NewTimer(delegationRetryDelay)
select {
case <-ctx.Done():
timer.Stop()
// outer timeout hit before retry window elapsed
case <-timer.C:
case <-time.After(delegationRetryDelay):
status, respBody, proxyErr = h.workspace.proxyA2ARequest(ctx, targetID, a2aBody, sourceID, true, false)
}
}
@@ -505,19 +482,15 @@ handleSuccess:
// dispatch eventually succeeds. Without the key, the drain finds
// the row by (workspace_id, target_id, method) but can't tell
// multiple-queued-delegations-to-same-target apart.
queuedJSON, marshalErr := json.Marshal(map[string]interface{}{
queuedJSON, _ := json.Marshal(map[string]interface{}{
"delegation_id": delegationID,
"queued": true,
})
if marshalErr != nil {
log.Printf("Delegation %s: json.Marshal queuedJSON failed: %v", delegationID, marshalErr)
} else {
if _, err := db.DB.ExecContext(ctx, `
INSERT INTO activity_logs (workspace_id, activity_type, method, source_id, target_id, summary, response_body, status)
VALUES ($1, 'delegation', 'delegate_result', $2, $3, $4, $5::jsonb, 'queued')
`, sourceID, sourceID, targetID, "Delegation queued — target at capacity", string(queuedJSON)); err != nil {
log.Printf("Delegation %s: failed to insert queued log: %v", delegationID, err)
}
if _, err := db.DB.ExecContext(ctx, `
INSERT INTO activity_logs (workspace_id, activity_type, method, source_id, target_id, summary, response_body, status)
VALUES ($1, 'delegation', 'delegate_result', $2, $3, $4, $5::jsonb, 'queued')
`, sourceID, sourceID, targetID, "Delegation queued — target at capacity", string(queuedJSON)); err != nil {
log.Printf("Delegation %s: failed to insert queued log: %v", delegationID, err)
}
h.broadcaster.RecordAndBroadcast(ctx, string(events.EventDelegationStatus), sourceID, map[string]interface{}{
"delegation_id": delegationID, "target_id": targetID, "status": "queued",
@@ -532,19 +505,15 @@ handleSuccess:
log.Printf("Delegation %s: step=inserting_success_log", delegationID)
// Store success (response_body must be JSONB, include delegation_id)
respJSON, marshalErr := json.Marshal(map[string]interface{}{
respJSON, _ := json.Marshal(map[string]interface{}{
"text": responseText,
"delegation_id": delegationID,
})
if marshalErr != nil {
log.Printf("Delegation %s: json.Marshal respJSON failed: %v", delegationID, marshalErr)
} else {
if _, err := db.DB.ExecContext(ctx, `
INSERT INTO activity_logs (workspace_id, activity_type, method, source_id, target_id, summary, response_body, status)
VALUES ($1, 'delegation', 'delegate_result', $2, $3, $4, $5::jsonb, 'completed')
`, sourceID, sourceID, targetID, "Delegation completed ("+textutil.TruncateBytes(responseText, 80)+")", string(respJSON)); err != nil {
log.Printf("Delegation %s: failed to insert success log: %v", delegationID, err)
}
if _, err := db.DB.ExecContext(ctx, `
INSERT INTO activity_logs (workspace_id, activity_type, method, source_id, target_id, summary, response_body, status)
VALUES ($1, 'delegation', 'delegate_result', $2, $3, $4, $5::jsonb, 'completed')
`, sourceID, sourceID, targetID, "Delegation completed ("+textutil.TruncateBytes(responseText, 80)+")", string(respJSON)); err != nil {
log.Printf("Delegation %s: failed to insert success log: %v", delegationID, err)
}
log.Printf("Delegation %s: step=recording_ledger_completed", delegationID)
@@ -621,25 +590,15 @@ func (h *DelegationHandler) Record(c *gin.Context) {
return
}
taskJSON, marshalErr := json.Marshal(map[string]interface{}{
taskJSON, _ := json.Marshal(map[string]interface{}{
"task": body.Task,
"delegation_id": body.DelegationID,
})
if marshalErr != nil {
log.Printf("Delegation %s: json.Marshal taskJSON failed: %v", body.DelegationID, marshalErr)
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to marshal task"})
return
}
// Store delegation_id in response_body so agent check_delegation_status
// can locate this row. Fixes mc#984.
respJSON, marshalErr := json.Marshal(map[string]interface{}{
respJSON, _ := json.Marshal(map[string]interface{}{
"delegation_id": body.DelegationID,
})
if marshalErr != nil {
log.Printf("Delegation %s: json.Marshal respJSON failed: %v", body.DelegationID, marshalErr)
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to marshal response"})
return
}
if _, err := db.DB.ExecContext(ctx, `
INSERT INTO activity_logs (workspace_id, activity_type, method, source_id, target_id, summary, request_body, response_body, status)
VALUES ($1, 'delegation', 'delegate', $2, $3, $4, $5::jsonb, $6::jsonb, 'dispatched')
@@ -703,19 +662,15 @@ func (h *DelegationHandler) UpdateStatus(c *gin.Context) {
h.updateDelegationStatus(ctx, sourceID, delegationID, body.Status, body.Error)
if body.Status == "completed" {
respJSON, marshalErr := json.Marshal(map[string]interface{}{
respJSON, _ := json.Marshal(map[string]interface{}{
"text": body.ResponsePreview,
"delegation_id": delegationID,
})
if marshalErr != nil {
log.Printf("Delegation UpdateStatus %s: json.Marshal respJSON failed: %v", delegationID, marshalErr)
} else {
if _, err := db.DB.ExecContext(ctx, `
INSERT INTO activity_logs (workspace_id, activity_type, method, source_id, summary, response_body, status)
VALUES ($1, 'delegation', 'delegate_result', $2, $3, $4::jsonb, 'completed')
`, sourceID, sourceID, "Delegation completed ("+textutil.TruncateBytes(body.ResponsePreview, 80)+")", string(respJSON)); err != nil {
log.Printf("Delegation UpdateStatus: result insert failed for %s: %v", delegationID, err)
}
if _, err := db.DB.ExecContext(ctx, `
INSERT INTO activity_logs (workspace_id, activity_type, method, source_id, summary, response_body, status)
VALUES ($1, 'delegation', 'delegate_result', $2, $3, $4::jsonb, 'completed')
`, sourceID, sourceID, "Delegation completed ("+textutil.TruncateBytes(body.ResponsePreview, 80)+")", string(respJSON)); err != nil {
log.Printf("Delegation UpdateStatus: result insert failed for %s: %v", delegationID, err)
}
h.broadcaster.RecordAndBroadcast(ctx, string(events.EventDelegationComplete), sourceID, map[string]interface{}{
"delegation_id": delegationID,
@@ -140,14 +140,7 @@ func buildHTTPResponse(statusCode int, body string) []byte {
}
// setupIntegrationFixtures inserts the rows executeDelegation requires:
// - workspaces: source (org root) + target as its CHILD, so both live in the
// SAME org. CanCommunicate=true (parent↔child) AND the #1953 sameOrg() guard
// in proxyA2ARequest passes (both resolve to the same org root). A real
// delegation happens INSIDE one org. (Previously both were parent_id=NULL —
// two DISTINCT org roots — which only "communicated" via CanCommunicate's
// root-sibling rule; #1953 added a sameOrg() guard that now denies routing
// between two org roots as cross-tenant, so the success-path tests below
// must use a same-org source/target pair.)
// - workspaces: source and target (siblings, parent_id=NULL so CanCommunicate=true)
// - activity_logs: the 'delegate' row that updateDelegationStatus UPDATE will find
// - delegations: the ledger row that recordLedgerStatus will UPDATE
//
@@ -155,14 +148,13 @@ func buildHTTPResponse(statusCode int, body string) []byte {
func setupIntegrationFixtures(t *testing.T, conn *sql.DB) func() {
t.Helper()
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
sourceID := integrationTestSourceID // org root (parent_id NULL); target hangs off it
for _, ws := range []struct {
id string
name string
parentID *string
}{
{integrationTestSourceID, "test-source", nil},
{integrationTestTargetID, "test-target", &sourceID}, // child of source → same org
{integrationTestTargetID, "test-target", nil},
} {
if _, err := conn.ExecContext(ctx,
`INSERT INTO workspaces (id, name, parent_id) VALUES ($1::uuid, $2, $3) ON CONFLICT (id) DO NOTHING`,
@@ -518,94 +510,6 @@ func TestIntegration_ExecuteDelegation_RedisDown_FallsBackToDB(t *testing.T) {
}
}
// TestIntegration_SameOrg_RealCTE_ResolvesAncestorChain is the regression gate
// for the org_scope.go recursive-CTE bug (#1953 follow-up). The sqlmock unit
// tests feed sameOrg() a pre-computed root_id row, so they CANNOT catch a wrong
// CTE — they assume it already returns the right value. Only a real Postgres
// run exercises orgRootSubtreeCTE itself.
//
// The bug: the CTE carried `id AS root_id` from the recursive SEED, so a
// non-root workspace resolved to ITSELF instead of its topmost ancestor. That
// made sameOrg() return false for two genuinely same-org workspaces and 403 a
// legitimate same-org a2a route (over-block). This test seeds a real
// root → child → grandchild chain plus a separate org root, and asserts:
// - every node in the chain resolves to the SAME org root (root, child, grandchild)
// - two workspaces in the same chain are sameOrg (incl. grandchild ↔ root)
// - a workspace in a DIFFERENT chain is NOT sameOrg (cross-tenant stays closed)
func TestIntegration_SameOrg_RealCTE_ResolvesAncestorChain(t *testing.T) {
conn := integrationDB(t)
const (
rootA = "11111111-1111-1111-1111-111111111111"
childA = "22222222-2222-2222-2222-222222222222"
grandchildA = "33333333-3333-3333-3333-333333333333"
rootB = "44444444-4444-4444-4444-444444444444"
)
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
t.Cleanup(func() {
c2, cancel2 := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel2()
// Delete leaf-first to respect the parent_id self-FK.
for _, id := range []string{grandchildA, childA, rootA, rootB} {
conn.ExecContext(c2, `DELETE FROM workspaces WHERE id = $1`, id)
}
})
// Insert parent-before-child to satisfy the self-referential FK.
seed := []struct {
id, name string
parent *string
}{
{rootA, "org-a-root", nil},
{childA, "org-a-child", strPtr(rootA)},
{grandchildA, "org-a-grandchild", strPtr(childA)},
{rootB, "org-b-root", nil},
}
for _, s := range seed {
if _, err := conn.ExecContext(ctx,
`INSERT INTO workspaces (id, name, parent_id) VALUES ($1::uuid, $2, $3) ON CONFLICT (id) DO NOTHING`,
s.id, s.name, s.parent); err != nil {
t.Fatalf("seed %s: %v", s.name, err)
}
}
// Every node in chain A must resolve to rootA via the REAL CTE.
for _, id := range []string{rootA, childA, grandchildA} {
got, err := orgRootID(ctx, conn, id)
if err != nil {
t.Fatalf("orgRootID(%s): %v", id, err)
}
if got != rootA {
t.Errorf("orgRootID(%s) = %q, want rootA %q (CTE must walk to topmost ancestor)", id, got, rootA)
}
}
// Same-org positives — including the grandchild↔root pair that the buggy
// CTE got wrong.
for _, pair := range [][2]string{{childA, grandchildA}, {rootA, grandchildA}, {rootA, childA}} {
ok, err := sameOrg(ctx, conn, pair[0], pair[1])
if err != nil {
t.Fatalf("sameOrg(%s,%s): %v", pair[0], pair[1], err)
}
if !ok {
t.Errorf("sameOrg(%s,%s) = false, want true (same org chain)", pair[0], pair[1])
}
}
// Cross-org negative — isolation must stay closed.
for _, pair := range [][2]string{{rootA, rootB}, {grandchildA, rootB}, {childA, rootB}} {
ok, err := sameOrg(ctx, conn, pair[0], pair[1])
if err != nil {
t.Fatalf("sameOrg(%s,%s): %v", pair[0], pair[1], err)
}
if ok {
t.Errorf("sameOrg(%s,%s) = true, want false (different orgs — cross-tenant must stay denied)", pair[0], pair[1])
}
}
}
// extractHostPort parses "http://127.0.0.1:PORT/" and returns "127.0.0.1:PORT".
func extractHostPort(rawURL string) string {
// Simple parse: strip "http://" prefix and trailing slash.
@@ -1059,25 +1059,13 @@ func expectExecuteDelegationBase(mock sqlmock.Sqlmock) {
WillReturnResult(sqlmock.NewResult(0, 1))
// CanCommunicate: getWorkspaceRef(source) + getWorkspaceRef(target).
// Source and target are siblings under one shared parent (one tenant) →
// CanCommunicate allowed. (#1953: they must NOT both be parent_id=NULL —
// two distinct org roots are now treated as DIFFERENT orgs and routing
// between them is denied. A real delegation happens inside one org.)
// Both are root-level workspaces (parent_id=NULL) → root-level siblings → allowed.
mock.ExpectQuery("SELECT id, parent_id FROM workspaces WHERE id = ").
WithArgs(testDeliverySourceID).
WillReturnRows(sqlmock.NewRows([]string{"id", "parent_id"}).AddRow(testDeliverySourceID, "ws-org-root-159"))
WillReturnRows(sqlmock.NewRows([]string{"id", "parent_id"}).AddRow(testDeliverySourceID, nil))
mock.ExpectQuery("SELECT id, parent_id FROM workspaces WHERE id = ").
WithArgs(testDeliveryTargetID).
WillReturnRows(sqlmock.NewRows([]string{"id", "parent_id"}).AddRow(testDeliveryTargetID, "ws-org-root-159"))
// #1953 cross-tenant guard: same-org check after CanCommunicate. Both
// resolve to the same org root → routing allowed.
mock.ExpectQuery("WITH RECURSIVE org_chain AS").
WithArgs(testDeliverySourceID).
WillReturnRows(sqlmock.NewRows([]string{"root_id"}).AddRow("ws-org-root-159"))
mock.ExpectQuery("WITH RECURSIVE org_chain AS").
WithArgs(testDeliveryTargetID).
WillReturnRows(sqlmock.NewRows([]string{"root_id"}).AddRow("ws-org-root-159"))
WillReturnRows(sqlmock.NewRows([]string{"id", "parent_id"}).AddRow(testDeliveryTargetID, nil))
// resolveAgentURL: test callers always set the URL in Redis (mr.Set ws:{id}:url),
// so resolveAgentURL gets a cache hit and never falls back to DB.
@@ -237,17 +237,7 @@ func (h *DiscoveryHandler) Peers(c *gin.Context) {
var peers []map[string]interface{}
// Siblings — workspaces sharing the caller's parent.
//
// #1953 cross-tenant isolation: the OLD code's else-branch handled the
// org-root caller (parent_id IS NULL) by returning EVERY workspace with
// parent_id IS NULL — i.e. every other tenant's org root, since the
// workspaces table has no org_id column. That leaked peer identities/URLs
// across tenants. An org root has no siblings inside its own org (each
// tenant is a distinct org root), so the org-root caller now gets an empty
// sibling set; its real peers are its children, returned below. Only the
// parent_id-bound branch enumerates siblings, and that is already scoped to
// one parent (one tenant).
// Siblings
if parentID.Valid {
siblings, _ := queryPeerMaps(`
SELECT w.id, w.name, COALESCE(w.role, ''), w.tier, w.status,
@@ -256,6 +246,14 @@ func (h *DiscoveryHandler) Peers(c *gin.Context) {
FROM workspaces w WHERE w.parent_id = $1 AND w.id != $2 AND w.status != 'removed'`,
parentID.String, workspaceID)
peers = append(peers, siblings...)
} else {
siblings, _ := queryPeerMaps(`
SELECT w.id, w.name, COALESCE(w.role, ''), w.tier, w.status,
COALESCE(w.agent_card, 'null'::jsonb), COALESCE(w.url, ''),
w.parent_id, w.active_tasks
FROM workspaces w WHERE w.parent_id IS NULL AND w.id != $1 AND w.status != 'removed'`,
workspaceID)
peers = append(peers, siblings...)
}
// Children — exclude self defensively. A child row whose parent_id
@@ -223,10 +223,10 @@ func TestPeers_RootWorkspace_NoPeers(t *testing.T) {
peerCols := []string{"id", "name", "role", "tier", "status", "agent_card", "url", "parent_id", "active_tasks"}
// #1953: an org-root caller (parent_id IS NULL) now issues NO sibling
// query at all. The old `WHERE w.parent_id IS NULL` sibling read returned
// EVERY tenant's org root (cross-tenant leak); an org root has no siblings
// inside its own org, so the handler skips the sibling read entirely.
// Siblings (other root-level workspaces) — none
mock.ExpectQuery("SELECT w.id, w.name.*WHERE w.parent_id IS NULL AND w.id != \\$1").
WithArgs("ws-root-alone").
WillReturnRows(sqlmock.NewRows(peerCols))
// Children — none. #383 added explicit `w.id != $2` self-filter.
mock.ExpectQuery("SELECT w.id, w.name.*WHERE w.parent_id = \\$1 AND w.id != \\$2").
@@ -155,10 +155,7 @@ func generateAppInstallationToken() (string, time.Time, error) {
if err != nil {
return "", time.Time{}, fmt.Errorf("sign JWT: %w", err)
}
req, err := http.NewRequest("POST", fmt.Sprintf("https://api.github.com/app/installations/%d/access_tokens", installID), nil)
if err != nil {
return "", time.Time{}, fmt.Errorf("build request: %w", err)
}
req, _ := http.NewRequest("POST", fmt.Sprintf("https://api.github.com/app/installations/%d/access_tokens", installID), nil)
req.Header.Set("Authorization", "Bearer "+signed)
req.Header.Set("Accept", "application/vnd.github+json")
client := &http.Client{Timeout: 30 * time.Second}
@@ -167,9 +164,6 @@ func generateAppInstallationToken() (string, time.Time, error) {
return "", time.Time{}, err
}
defer func() { _ = resp.Body.Close() }()
if resp.StatusCode != http.StatusCreated {
return "", time.Time{}, fmt.Errorf("github token endpoint returned status %d", resp.StatusCode)
}
var result struct {
Token string `json:"token"`
ExpiresAt time.Time `json:"expires_at"`
@@ -255,23 +255,9 @@ func TestExtended_SecretsListEmpty(t *testing.T) {
// ---------- TestSecretsSet (Extended) ----------
func TestExtended_SecretsSet(t *testing.T) {
// internal#691: the per-workspace strip gate now defaults to platform_managed
// on empty MOLECULE_LLM_BILLING_MODE (closed default). This test's intent is
// the happy path of persisting a vendor key, so put the org into byok which
// matches the pre-#691 implicit behavior of an unset env.
t.Setenv("MOLECULE_LLM_BILLING_MODE", "byok")
mock := setupTestDB(t)
handler := NewSecretsHandler(nil)
// internal#691: secrets.Set now consults ResolveLLMBillingMode before the
// strip gate. Mock returns no row → resolver falls through to the org
// default (byok, set via t.Setenv above) → bypass-list check is skipped
// and the write proceeds. This pattern is the test-side mirror of the
// real-prod fall-through behavior for a fresh workspace with no override.
mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs("22222222-2222-2222-2222-222222222222").
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}))
// Expect INSERT (encrypted value is dynamic, use AnyArg)
mock.ExpectExec("INSERT INTO workspace_secrets").
WithArgs("22222222-2222-2222-2222-222222222222", "OPENAI_API_KEY", sqlmock.AnyArg(), sqlmock.AnyArg()).
@@ -316,7 +302,7 @@ func TestExtended_SecretsSetRejectsHermesCustomProviderInPlatformManagedMode(t *
c, _ := gin.CreateTestContext(w)
c.Params = gin.Params{{Key: "id", Value: "22222222-2222-2222-2222-222222222222"}}
body := `{"key":"KIMI_API_KEY","value":"sk-test-moonshot"}`
body := `{"key":"HERMES_CUSTOM_BASE_URL","value":"https://api.moonshot.ai/v1"}`
c.Request = httptest.NewRequest("POST", "/workspaces/22222222-2222-2222-2222-222222222222/secrets", bytes.NewBufferString(body))
c.Request.Header.Set("Content-Type", "application/json")
@@ -453,14 +439,6 @@ func TestExtended_DiscoverMissingHeader(t *testing.T) {
// ---------- TestPeers (Extended) ----------
// TestExtended_Peers verifies a root-level (org-root) workspace's peer view.
//
// #1953: previously a root-level caller issued `WHERE w.parent_id IS NULL`
// for siblings, which returned EVERY other tenant's org root as a "peer"
// (cross-tenant leak, since the workspaces table has no org_id column). After
// the fix an org root has no cross-tenant siblings; its only peers are its own
// children. This test asserts the child is returned and that NO sibling query
// is issued (no `parent_id IS NULL` read).
func TestExtended_Peers(t *testing.T) {
mock := setupTestDB(t)
setupTestRedis(t)
@@ -471,14 +449,17 @@ func TestExtended_Peers(t *testing.T) {
WithArgs("ws-peer").
WillReturnRows(sqlmock.NewRows([]string{"parent_id"}).AddRow(nil))
// NO root-level sibling query is issued for an org-root caller anymore.
// Expect root-level siblings query (parent IS NULL, excluding self)
mock.ExpectQuery("SELECT w.id, w.name").
WithArgs("ws-peer").
WillReturnRows(sqlmock.NewRows([]string{"id", "name", "role", "tier", "status", "agent_card", "url", "parent_id", "active_tasks"}).
AddRow("ws-sibling", "Sibling Agent", "worker", 1, "online", []byte("null"), "http://localhost:9001", nil, 0))
// Children query (workspaces with parent_id = ws-peer, excluding self).
// Query binds (parent_id, self_id) for the self-filter guard added in #383.
// Expect children query (workspaces with parent_id = ws-peer, excluding self)
// Query now binds (parent_id, self_id) for the self-filter guard added in #383.
mock.ExpectQuery("SELECT w.id, w.name").
WithArgs("ws-peer", "ws-peer").
WillReturnRows(sqlmock.NewRows([]string{"id", "name", "role", "tier", "status", "agent_card", "url", "parent_id", "active_tasks"}).
AddRow("ws-child", "Child Agent", "worker", 1, "online", []byte("null"), "http://localhost:9001", "ws-peer", 0))
WillReturnRows(sqlmock.NewRows([]string{"id", "name", "role", "tier", "status", "agent_card", "url", "parent_id", "active_tasks"}))
// No parent query since workspace is root-level
@@ -498,10 +479,10 @@ func TestExtended_Peers(t *testing.T) {
t.Fatalf("failed to parse response: %v", err)
}
if len(resp) != 1 {
t.Fatalf("expected 1 peer (the child), got %d", len(resp))
t.Fatalf("expected 1 peer, got %d", len(resp))
}
if resp[0]["name"] != "Child Agent" {
t.Errorf("expected peer name 'Child Agent', got %v", resp[0]["name"])
if resp[0]["name"] != "Sibling Agent" {
t.Errorf("expected peer name 'Sibling Agent', got %v", resp[0]["name"])
}
if err := mock.ExpectationsWereMet(); err != nil {
@@ -169,13 +169,7 @@ func (h *InstructionsHandler) Update(c *gin.Context) {
c.JSON(http.StatusInternalServerError, gin.H{"error": "update failed"})
return
}
n, err := result.RowsAffected()
if err != nil {
log.Printf("Instructions update RowsAffected error: %v", err)
c.JSON(http.StatusInternalServerError, gin.H{"error": "update failed"})
return
}
if n == 0 {
if n, _ := result.RowsAffected(); n == 0 {
c.JSON(http.StatusNotFound, gin.H{"error": "instruction not found"})
return
}
@@ -192,13 +186,7 @@ func (h *InstructionsHandler) Delete(c *gin.Context) {
c.JSON(http.StatusInternalServerError, gin.H{"error": "delete failed"})
return
}
n, err := result.RowsAffected()
if err != nil {
log.Printf("Instructions delete RowsAffected error: %v", err)
c.JSON(http.StatusInternalServerError, gin.H{"error": "delete failed"})
return
}
if n == 0 {
if n, _ := result.RowsAffected(); n == 0 {
c.JSON(http.StatusNotFound, gin.H{"error": "instruction not found"})
return
}
@@ -1,234 +0,0 @@
package handlers
// llm_billing_mode.go — per-workspace LLM billing mode resolution (internal#691).
//
// The resolver answers a single question at provision time:
// "Should we strip CLAUDE_CODE_OAUTH_TOKEN + every vendor key from this
// workspace's env, force-route to the CP proxy, and bill org credits?"
//
// That question used to be a single env-var read inside applyPlatformManagedLLMEnv:
//
// os.Getenv("MOLECULE_LLM_BILLING_MODE") == "platform_managed" → strip
//
// where MOLECULE_LLM_BILLING_MODE was an ORG-level value, fetched from CP's
// tenant_config and exported into the workspace-server process at boot. That
// shape made it impossible to mix billing modes across workspaces in the same
// org: turning the org dial to `byok` so one workspace could keep its OAuth
// stops the strip for EVERY workspace in the org. Turning it to `platform_managed`
// blocks every workspace's own OAuth/vendor keys.
//
// The resolver replaces the env-var read with a per-workspace lookup:
//
// workspaces.llm_billing_mode (per-workspace override, NULLABLE)
// ?? organizations.llm_billing_mode (org default, fetched via tenant_config)
// ?? "platform_managed" (closed default — the existing implicit default)
//
// Default-closed contract — non-negotiable per the RFC Safety axis:
//
// - workspace row missing (sql.ErrNoRows) → fall through to org default
// - DB error on the lookup → "platform_managed" + propagated error
// - workspace override = NULL → fall through to org default
// - workspace override = unknown string → "platform_managed" (default-closed)
// - org default = NULL / empty / unknown string → "platform_managed" (closed default)
// - org default = recognized non-pm string + ws null → org default (byok/disabled honored)
//
// The ONLY way to resolve to "byok" or "disabled" is an explicit, recognized
// string in the workspace override OR the org default. A NULL JOIN, transient
// resolver error, or garbled enum value MUST NOT silently flip a workspace
// off of platform_managed — that would shadow the org's billing policy and
// is the exact failure mode the RFC's Safety hot-spot calls out.
import (
"context"
"database/sql"
"errors"
"fmt"
"git.moleculesai.app/molecule-ai/molecule-core/workspace-server/internal/db"
)
// Constants mirror molecule-controlplane/internal/credits/llm_billing.go.
// Kept as string literals (not imports) because workspace-server has no
// build-time dependency on the CP module; the values are stable wire
// strings used in the tenant_config response, the workspaces.llm_billing_mode
// column check constraint, and the CP route bodies.
const (
LLMBillingModePlatformManaged = "platform_managed"
LLMBillingModeBYOK = "byok"
LLMBillingModeDisabled = "disabled"
)
// BillingModeSource describes which layer of the resolution stack supplied
// the final mode. Surfaced via the admin route for operator debug
// ("why is this workspace being stripped?") per the RFC Observability axis.
type BillingModeSource string
const (
BillingModeSourceWorkspaceOverride BillingModeSource = "workspace_override"
BillingModeSourceOrgDefault BillingModeSource = "org_default"
BillingModeSourceConstantFallback BillingModeSource = "constant_fallback"
)
// BillingModeResolution is the structured answer the admin GET route returns
// and the strip gate logs at INFO. The same struct is the unit-test fixture
// shape, so the resolver test asserts both the mode AND the source per case
// (catches a bug where the right mode is returned via the wrong layer).
type BillingModeResolution struct {
WorkspaceID string `json:"workspace_id"`
ResolvedMode string `json:"resolved_mode"`
WorkspaceOverride *string `json:"workspace_override"` // nil = inherit
OrgDefault string `json:"org_default"` // already default-closed by CP
Source BillingModeSource `json:"source"`
}
// isKnownBillingMode is the enum-recognizer for the resolver's default-closed
// branch. Returning false for an unknown string forces the resolver to fall
// through to the next layer (or the constant fallback) — NEVER to honor a
// garbled value as if it were valid. This is what makes a row with mode='byokk'
// (typo) resolve to platform_managed instead of accidentally to byok.
func isKnownBillingMode(s string) bool {
switch s {
case LLMBillingModePlatformManaged, LLMBillingModeBYOK, LLMBillingModeDisabled:
return true
default:
return false
}
}
// normalizeOrgDefault applies the same default-closed contract to the
// org-level input as the workspace override gets. The org_default arrives
// from tenant_config which already COALESCEs NULL → platform_managed at the
// CP SQL layer, but we DO NOT trust that contract here — if CP regresses or
// the tenant_config env wasn't populated (race on boot), we still default-
// close. Same principle: never honor a garbled value.
func normalizeOrgDefault(orgMode string) string {
if isKnownBillingMode(orgMode) {
return orgMode
}
return LLMBillingModePlatformManaged
}
// ResolveLLMBillingMode is the canonical resolver. Every code path that
// previously gated on `os.Getenv("MOLECULE_LLM_BILLING_MODE") == "platform_managed"`
// must call this instead and gate on the returned mode. The architectural
// test (resolver_ast_test.go) asserts there is no remaining call site of
// the old shape outside the resolver-input wiring.
//
// Returning an error does NOT prevent the caller from making a decision —
// the returned mode is always a valid enum value (default-closed to
// platform_managed) so the caller can proceed without a separate fail-closed
// branch. The error is informational: log it, surface it to operators, but
// the strip-gate decision is already safe.
func ResolveLLMBillingMode(ctx context.Context, workspaceID, orgMode string) (BillingModeResolution, error) {
res := BillingModeResolution{
WorkspaceID: workspaceID,
OrgDefault: normalizeOrgDefault(orgMode),
}
if workspaceID == "" {
// No workspace ID = pre-provision context (templating, validation).
// Resolve against the org default only, no DB read.
res.ResolvedMode = res.OrgDefault
res.Source = BillingModeSourceOrgDefault
if !isKnownBillingMode(orgMode) {
// Org default was garbled/NULL and we clamped to platform_managed.
// Mark the source as constant_fallback so the operator can see
// the clamp happened, not that the org "really" said platform_managed.
res.Source = BillingModeSourceConstantFallback
}
return res, nil
}
var wsOverride sql.NullString
err := db.DB.QueryRowContext(ctx,
`SELECT llm_billing_mode FROM workspaces WHERE id = $1`,
workspaceID,
).Scan(&wsOverride)
switch {
case errors.Is(err, sql.ErrNoRows):
// Workspace row missing — concurrent delete, or pre-create call. Don't
// silently flip; fall through to org default. Source stays org_default
// so operators can see the row-missing case is being handled as a
// fallback, not a workspace-explicit decision.
res.ResolvedMode = res.OrgDefault
res.Source = BillingModeSourceOrgDefault
if !isKnownBillingMode(orgMode) {
res.Source = BillingModeSourceConstantFallback
}
return res, nil
case err != nil:
// DB error — default-closed to platform_managed AND propagate the
// error so operators get a structured log line. The caller is
// expected to log and continue with the safe default.
res.ResolvedMode = LLMBillingModePlatformManaged
res.Source = BillingModeSourceConstantFallback
return res, fmt.Errorf("resolve workspace llm_billing_mode for %s: %w", workspaceID, err)
}
if wsOverride.Valid && isKnownBillingMode(wsOverride.String) {
mode := wsOverride.String
res.WorkspaceOverride = &mode
res.ResolvedMode = mode
res.Source = BillingModeSourceWorkspaceOverride
return res, nil
}
// Override row present but the value is NULL or garbled. Fall through.
// If the value was non-NULL but garbled (CHECK constraint should prevent
// this, but defense in depth — a future migration could relax the check
// or another path could write the column directly), surface the raw
// override value so operators can spot the corrupt row.
if wsOverride.Valid {
raw := wsOverride.String
res.WorkspaceOverride = &raw
}
res.ResolvedMode = res.OrgDefault
res.Source = BillingModeSourceOrgDefault
if !isKnownBillingMode(orgMode) {
res.Source = BillingModeSourceConstantFallback
}
return res, nil
}
// SetWorkspaceLLMBillingMode writes the override column. Pass mode=="" to
// clear (set to NULL = inherit). Validates the mode against the enum set
// so the route handler doesn't have to duplicate validation; a garbled
// mode round-trips as an explicit 400 from the caller, not a CHECK-
// constraint error from the DB driver.
func SetWorkspaceLLMBillingMode(ctx context.Context, workspaceID, mode string) error {
if workspaceID == "" {
return errors.New("SetWorkspaceLLMBillingMode: workspace id required")
}
if mode == "" {
// NULL = inherit. Caller asked to clear the override.
res, err := db.DB.ExecContext(ctx,
`UPDATE workspaces SET llm_billing_mode = NULL WHERE id = $1`,
workspaceID,
)
if err != nil {
return fmt.Errorf("clear workspace llm_billing_mode for %s: %w", workspaceID, err)
}
n, _ := res.RowsAffected()
if n == 0 {
return sql.ErrNoRows
}
return nil
}
if !isKnownBillingMode(mode) {
return fmt.Errorf("unknown billing mode %q (allowed: %s, %s, %s)",
mode, LLMBillingModePlatformManaged, LLMBillingModeBYOK, LLMBillingModeDisabled)
}
res, err := db.DB.ExecContext(ctx,
`UPDATE workspaces SET llm_billing_mode = $1 WHERE id = $2`,
mode, workspaceID,
)
if err != nil {
return fmt.Errorf("set workspace llm_billing_mode for %s: %w", workspaceID, err)
}
n, _ := res.RowsAffected()
if n == 0 {
return sql.ErrNoRows
}
return nil
}
@@ -1,154 +0,0 @@
package handlers
// llm_billing_mode_handler.go — workspace-server admin routes that read /
// write the per-workspace billing mode override (internal#691). These are
// the per-tenant routes that CP's new /cp/admin/workspaces/:id/llm-billing-mode
// proxies to; the canvas hits them via the CP route, not directly.
//
// Route shape:
//
// GET /admin/workspaces/:id/llm-billing-mode
// -> 200 BillingModeResolution
// -> 400 on malformed UUID
// -> 500 on DB error (response still includes a safe_default the caller
// can fall through to — the resolver always returns a valid mode
// even on error, per the default-closed contract)
//
// PUT /admin/workspaces/:id/llm-billing-mode
// body: {"mode": "byok" | "platform_managed" | "disabled" | null}
// -> 200 BillingModeResolution (post-write)
// -> 400 on bad UUID / unknown mode / malformed body / missing "mode" key
// -> 404 when the workspace row doesn't exist
//
// Auth: mounted under wsAdmin (middleware.AdminAuth) — admin_token required.
import (
"database/sql"
"encoding/json"
"errors"
"io"
"net/http"
"os"
"strings"
"github.com/gin-gonic/gin"
)
// GetWorkspaceLLMBillingMode handles GET /admin/workspaces/:id/llm-billing-mode.
//
// Reads the workspace override + the org-level default (from the same
// MOLECULE_LLM_BILLING_MODE env var the provisioner reads at strip-gate time —
// keeps the two paths consistent so the GET result matches what the strip
// gate would compute) and returns the structured resolution.
func GetWorkspaceLLMBillingMode(c *gin.Context) {
workspaceID := strings.TrimSpace(c.Param("id"))
if !uuidRegex.MatchString(workspaceID) {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid workspace id"})
return
}
orgMode := strings.ToLower(strings.TrimSpace(os.Getenv("MOLECULE_LLM_BILLING_MODE")))
res, err := ResolveLLMBillingMode(c.Request.Context(), workspaceID, orgMode)
if err != nil {
// Resolver returns a safe default-closed mode alongside the error;
// surface the error so the operator sees the DB issue, but the
// response still has a usable mode field for the caller to fall
// through to without a separate fail-closed branch.
c.JSON(http.StatusInternalServerError, gin.H{
"error": "resolve workspace billing mode failed",
"detail": err.Error(),
"safe_default": res.ResolvedMode,
"workspace_id": res.WorkspaceID,
})
return
}
c.JSON(http.StatusOK, res)
}
// PutWorkspaceLLMBillingMode handles PUT /admin/workspaces/:id/llm-billing-mode.
//
// Body shape: {"mode": "byok" | "platform_managed" | "disabled" | null}
// where null clears the override (workspace inherits the org default again).
// Omitting "mode" entirely is a 400 — callers must be explicit about whether
// they want to set or clear, so a typo'd field name can't silently no-op.
//
// On success returns the post-write resolution so the canvas can re-render
// without a follow-up GET.
func PutWorkspaceLLMBillingMode(c *gin.Context) {
workspaceID := strings.TrimSpace(c.Param("id"))
if !uuidRegex.MatchString(workspaceID) {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid workspace id"})
return
}
// Read raw body so we can distinguish three cases:
// {"mode": "byok"} → set override
// {"mode": null} → clear override
// {} → 400 (caller must be explicit)
// json.RawMessage zero length ⇔ key absent; raw "null" ⇔ explicit clear;
// raw quoted string ⇔ set.
raw, readErr := io.ReadAll(c.Request.Body)
if readErr != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "read body", "detail": readErr.Error()})
return
}
var body struct {
Mode json.RawMessage `json:"mode"`
}
if err := json.Unmarshal(raw, &body); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid json", "detail": err.Error()})
return
}
if len(body.Mode) == 0 {
c.JSON(http.StatusBadRequest, gin.H{"error": "missing required field 'mode' (use null to clear override)"})
return
}
var writeErr error
if string(body.Mode) == "null" {
writeErr = SetWorkspaceLLMBillingMode(c.Request.Context(), workspaceID, "")
} else {
var modeStr string
if err := json.Unmarshal(body.Mode, &modeStr); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "mode must be a string or null", "detail": err.Error()})
return
}
modeStr = strings.TrimSpace(modeStr)
if modeStr == "" {
// Empty string is ambiguous (could be "clear" or "user error");
// reject as 400 so the caller picks null explicitly.
c.JSON(http.StatusBadRequest, gin.H{"error": "mode must be one of platform_managed, byok, disabled, or null to clear"})
return
}
writeErr = SetWorkspaceLLMBillingMode(c.Request.Context(), workspaceID, modeStr)
}
if errors.Is(writeErr, sql.ErrNoRows) {
c.JSON(http.StatusNotFound, gin.H{"error": "workspace not found"})
return
}
if writeErr != nil {
// Validation errors from SetWorkspaceLLMBillingMode (unknown mode
// string) come back as a plain error; map to 400.
if strings.HasPrefix(writeErr.Error(), "unknown billing mode") {
c.JSON(http.StatusBadRequest, gin.H{"error": writeErr.Error()})
return
}
c.JSON(http.StatusInternalServerError, gin.H{"error": "set workspace billing mode failed", "detail": writeErr.Error()})
return
}
// Read back the resolution so the response reflects post-write state.
orgMode := strings.ToLower(strings.TrimSpace(os.Getenv("MOLECULE_LLM_BILLING_MODE")))
res, resolveErr := ResolveLLMBillingMode(c.Request.Context(), workspaceID, orgMode)
if resolveErr != nil {
// Write succeeded but readback failed — still return 200 with the
// best-effort resolution; the safe default is set even on error.
c.JSON(http.StatusOK, gin.H{
"workspace_id": workspaceID,
"resolved_mode": res.ResolvedMode,
"readback_error": resolveErr.Error(),
})
return
}
c.JSON(http.StatusOK, res)
}
@@ -1,205 +0,0 @@
package handlers
// llm_billing_mode_handler_test.go — admin route coverage for the per-workspace
// LLM billing mode endpoint (internal#691).
//
// What this guards:
// - GET path validates UUID + returns the BillingModeResolution shape
// - PUT distinguishes "omitted mode" (400) from "explicit null" (clear)
// from "string value" (set), so a typo'd field name can't silently no-op
// - Unknown mode strings 400 from the validator, not from a PG CHECK
// constraint round-trip (matters because the error message must be
// actionable to a canvas user)
// - 404 propagates when the workspace row is missing on a set/clear
import (
"bytes"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"github.com/DATA-DOG/go-sqlmock"
"github.com/gin-gonic/gin"
)
func init() {
gin.SetMode(gin.TestMode)
}
const testWSID = "44444444-4444-4444-4444-444444444444"
func TestGetWorkspaceLLMBillingMode_HappyPath_InheritsOrgDefault(t *testing.T) {
t.Setenv("MOLECULE_LLM_BILLING_MODE", LLMBillingModeBYOK)
mock := setupTestDB(t)
// Workspace has no override → resolver returns org_default = byok.
mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(testWSID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(nil))
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Params = gin.Params{{Key: "id", Value: testWSID}}
c.Request = httptest.NewRequest("GET", "/admin/workspaces/"+testWSID+"/llm-billing-mode", nil)
GetWorkspaceLLMBillingMode(c)
if w.Code != http.StatusOK {
t.Fatalf("status: got %d want 200, body=%s", w.Code, w.Body.String())
}
var res BillingModeResolution
if err := json.Unmarshal(w.Body.Bytes(), &res); err != nil {
t.Fatalf("decode: %v", err)
}
if res.ResolvedMode != LLMBillingModeBYOK {
t.Errorf("resolved mode: got %q want %q", res.ResolvedMode, LLMBillingModeBYOK)
}
if res.Source != BillingModeSourceOrgDefault {
t.Errorf("source: got %q want %q", res.Source, BillingModeSourceOrgDefault)
}
if res.WorkspaceOverride != nil {
t.Errorf("expected nil override, got %v", *res.WorkspaceOverride)
}
}
func TestGetWorkspaceLLMBillingMode_BadUUID_400(t *testing.T) {
setupTestDB(t)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Params = gin.Params{{Key: "id", Value: "not-a-uuid"}}
c.Request = httptest.NewRequest("GET", "/admin/workspaces/not-a-uuid/llm-billing-mode", nil)
GetWorkspaceLLMBillingMode(c)
if w.Code != http.StatusBadRequest {
t.Fatalf("status: got %d want 400", w.Code)
}
}
func TestPutWorkspaceLLMBillingMode_SetByok(t *testing.T) {
t.Setenv("MOLECULE_LLM_BILLING_MODE", LLMBillingModePlatformManaged)
mock := setupTestDB(t)
mock.ExpectExec(`UPDATE workspaces SET llm_billing_mode = \$1 WHERE id = \$2`).
WithArgs(LLMBillingModeBYOK, testWSID).
WillReturnResult(sqlmock.NewResult(0, 1))
// Readback after write.
mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(testWSID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(LLMBillingModeBYOK))
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Params = gin.Params{{Key: "id", Value: testWSID}}
body := `{"mode":"byok"}`
c.Request = httptest.NewRequest("PUT",
"/admin/workspaces/"+testWSID+"/llm-billing-mode",
bytes.NewBufferString(body))
c.Request.Header.Set("Content-Type", "application/json")
PutWorkspaceLLMBillingMode(c)
if w.Code != http.StatusOK {
t.Fatalf("status: got %d want 200, body=%s", w.Code, w.Body.String())
}
var res BillingModeResolution
if err := json.Unmarshal(w.Body.Bytes(), &res); err != nil {
t.Fatalf("decode: %v", err)
}
if res.ResolvedMode != LLMBillingModeBYOK {
t.Errorf("post-write resolved: got %q want %q", res.ResolvedMode, LLMBillingModeBYOK)
}
if res.Source != BillingModeSourceWorkspaceOverride {
t.Errorf("post-write source: got %q want %q", res.Source, BillingModeSourceWorkspaceOverride)
}
}
func TestPutWorkspaceLLMBillingMode_ExplicitNullClearsOverride(t *testing.T) {
t.Setenv("MOLECULE_LLM_BILLING_MODE", LLMBillingModePlatformManaged)
mock := setupTestDB(t)
mock.ExpectExec(`UPDATE workspaces SET llm_billing_mode = NULL WHERE id = \$1`).
WithArgs(testWSID).
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(testWSID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(nil))
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Params = gin.Params{{Key: "id", Value: testWSID}}
body := `{"mode":null}`
c.Request = httptest.NewRequest("PUT",
"/admin/workspaces/"+testWSID+"/llm-billing-mode",
bytes.NewBufferString(body))
c.Request.Header.Set("Content-Type", "application/json")
PutWorkspaceLLMBillingMode(c)
if w.Code != http.StatusOK {
t.Fatalf("status: got %d want 200, body=%s", w.Code, w.Body.String())
}
var res BillingModeResolution
if err := json.Unmarshal(w.Body.Bytes(), &res); err != nil {
t.Fatalf("decode: %v", err)
}
if res.ResolvedMode != LLMBillingModePlatformManaged {
t.Errorf("post-clear resolved: got %q want %q", res.ResolvedMode, LLMBillingModePlatformManaged)
}
if res.Source != BillingModeSourceOrgDefault {
t.Errorf("post-clear source: got %q want %q", res.Source, BillingModeSourceOrgDefault)
}
if res.WorkspaceOverride != nil {
t.Errorf("post-clear override should be nil, got %v", *res.WorkspaceOverride)
}
}
func TestPutWorkspaceLLMBillingMode_MissingModeField_400(t *testing.T) {
setupTestDB(t)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Params = gin.Params{{Key: "id", Value: testWSID}}
body := `{}`
c.Request = httptest.NewRequest("PUT",
"/admin/workspaces/"+testWSID+"/llm-billing-mode",
bytes.NewBufferString(body))
c.Request.Header.Set("Content-Type", "application/json")
PutWorkspaceLLMBillingMode(c)
if w.Code != http.StatusBadRequest {
t.Fatalf("status: got %d want 400, body=%s", w.Code, w.Body.String())
}
}
func TestPutWorkspaceLLMBillingMode_UnknownMode_400(t *testing.T) {
setupTestDB(t)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Params = gin.Params{{Key: "id", Value: testWSID}}
body := `{"mode":"totally-bogus"}`
c.Request = httptest.NewRequest("PUT",
"/admin/workspaces/"+testWSID+"/llm-billing-mode",
bytes.NewBufferString(body))
c.Request.Header.Set("Content-Type", "application/json")
PutWorkspaceLLMBillingMode(c)
if w.Code != http.StatusBadRequest {
t.Fatalf("status: got %d want 400, body=%s", w.Code, w.Body.String())
}
}
func TestPutWorkspaceLLMBillingMode_NoSuchWorkspace_404(t *testing.T) {
mock := setupTestDB(t)
// SET path: rows affected = 0 → SetWorkspaceLLMBillingMode returns sql.ErrNoRows
// → handler maps to 404.
mock.ExpectExec(`UPDATE workspaces SET llm_billing_mode = \$1 WHERE id = \$2`).
WithArgs(LLMBillingModeBYOK, testWSID).
WillReturnResult(sqlmock.NewResult(0, 0))
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Params = gin.Params{{Key: "id", Value: testWSID}}
body := `{"mode":"byok"}`
c.Request = httptest.NewRequest("PUT",
"/admin/workspaces/"+testWSID+"/llm-billing-mode",
bytes.NewBufferString(body))
c.Request.Header.Set("Content-Type", "application/json")
PutWorkspaceLLMBillingMode(c)
if w.Code != http.StatusNotFound {
t.Fatalf("status: got %d want 404, body=%s", w.Code, w.Body.String())
}
}
@@ -1,261 +0,0 @@
package handlers
// llm_billing_mode_test.go — table-driven tests for the per-workspace
// resolver (internal#691). The cases below enumerate every documented
// branch in the default-closed contract; if one of them flips behavior
// later the test names will tell the reviewer exactly which RFC clause
// regressed.
import (
"context"
"errors"
"testing"
"github.com/DATA-DOG/go-sqlmock"
)
func TestResolveLLMBillingMode_TableDriven(t *testing.T) {
ctx := context.Background()
const wsID = "11111111-1111-1111-1111-111111111111"
type want struct {
mode string
source BillingModeSource
// hasOverride asserts whether the resolver surfaced the override
// value in the result (nil pointer = clean inherit, non-nil = the
// row was present even if it ultimately fell through because it
// was garbled). Lets us distinguish "row missing, fell through"
// from "row present but garbled, fell through" — both resolve to
// the same mode but the resolver tells operators which case it was.
hasOverride bool
}
type tc struct {
name string
workspaceID string
orgMode string
setupMock func(m sqlmock.Sqlmock)
want want
wantErr bool
}
cases := []tc{
{
name: "workspace_override_byok_overrides_pm_org",
workspaceID: wsID,
orgMode: LLMBillingModePlatformManaged,
setupMock: func(m sqlmock.Sqlmock) {
m.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(wsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(LLMBillingModeBYOK))
},
want: want{mode: LLMBillingModeBYOK, source: BillingModeSourceWorkspaceOverride, hasOverride: true},
},
{
name: "workspace_override_disabled_overrides_pm_org",
workspaceID: wsID,
orgMode: LLMBillingModePlatformManaged,
setupMock: func(m sqlmock.Sqlmock) {
m.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(wsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(LLMBillingModeDisabled))
},
want: want{mode: LLMBillingModeDisabled, source: BillingModeSourceWorkspaceOverride, hasOverride: true},
},
{
name: "workspace_override_null_inherits_byok_org",
workspaceID: wsID,
orgMode: LLMBillingModeBYOK,
setupMock: func(m sqlmock.Sqlmock) {
m.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(wsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(nil))
},
want: want{mode: LLMBillingModeBYOK, source: BillingModeSourceOrgDefault, hasOverride: false},
},
{
name: "workspace_override_null_inherits_pm_org",
workspaceID: wsID,
orgMode: LLMBillingModePlatformManaged,
setupMock: func(m sqlmock.Sqlmock) {
m.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(wsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(nil))
},
want: want{mode: LLMBillingModePlatformManaged, source: BillingModeSourceOrgDefault, hasOverride: false},
},
{
name: "workspace_override_garbled_falls_through_to_pm_org_DEFAULT_CLOSED",
workspaceID: wsID,
orgMode: LLMBillingModePlatformManaged,
setupMock: func(m sqlmock.Sqlmock) {
// CHECK constraint would normally prevent this but if a future
// migration loosens it (or a direct UPDATE bypasses it on a
// non-PG driver in a test stub), a garbled value MUST NOT
// be honored as if it were valid. This is the default-closed
// safety axis the RFC calls out.
m.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(wsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow("byokk"))
},
want: want{mode: LLMBillingModePlatformManaged, source: BillingModeSourceOrgDefault, hasOverride: true},
},
{
name: "workspace_override_garbled_org_garbled_constant_fallback",
workspaceID: wsID,
orgMode: "garbled-or-empty",
setupMock: func(m sqlmock.Sqlmock) {
m.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(wsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow("nonsense"))
},
// Both layers garbled → constant fallback. Source is constant_fallback
// so operators can see the org-default-was-also-bad case explicitly.
want: want{mode: LLMBillingModePlatformManaged, source: BillingModeSourceConstantFallback, hasOverride: true},
},
{
name: "workspace_row_missing_falls_through_to_org_byok",
workspaceID: wsID,
orgMode: LLMBillingModeBYOK,
setupMock: func(m sqlmock.Sqlmock) {
m.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(wsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}))
},
want: want{mode: LLMBillingModeBYOK, source: BillingModeSourceOrgDefault, hasOverride: false},
},
{
name: "workspace_id_empty_pre_provision_org_only",
workspaceID: "",
orgMode: LLMBillingModeBYOK,
setupMock: func(m sqlmock.Sqlmock) { /* no DB read expected — empty ws id short-circuits */ },
want: want{mode: LLMBillingModeBYOK, source: BillingModeSourceOrgDefault, hasOverride: false},
},
{
name: "workspace_id_empty_org_garbled_constant_fallback",
workspaceID: "",
orgMode: "",
setupMock: func(m sqlmock.Sqlmock) { /* no DB read */ },
want: want{mode: LLMBillingModePlatformManaged, source: BillingModeSourceConstantFallback, hasOverride: false},
},
{
name: "db_error_default_closed_to_pm_with_error",
workspaceID: wsID,
orgMode: LLMBillingModeBYOK, // org says byok but DB errored — DO NOT honor org
setupMock: func(m sqlmock.Sqlmock) {
m.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(wsID).
WillReturnError(errors.New("connection refused"))
},
// Critical: even though orgMode=byok, a DB error means we can't
// confirm the workspace doesn't have an override, so we default
// to the closed mode. This is the safer of the two failures —
// silently flipping to org-byok on a DB error would leak the
// OAuth-keeping behavior to workspaces whose row says NULL.
want: want{mode: LLMBillingModePlatformManaged, source: BillingModeSourceConstantFallback, hasOverride: false},
wantErr: true,
},
}
for _, c := range cases {
t.Run(c.name, func(t *testing.T) {
mock := setupTestDB(t)
c.setupMock(mock)
res, err := ResolveLLMBillingMode(ctx, c.workspaceID, c.orgMode)
if (err != nil) != c.wantErr {
t.Fatalf("err: got %v wantErr=%v", err, c.wantErr)
}
if res.ResolvedMode != c.want.mode {
t.Errorf("mode: got %q want %q", res.ResolvedMode, c.want.mode)
}
if res.Source != c.want.source {
t.Errorf("source: got %q want %q", res.Source, c.want.source)
}
if (res.WorkspaceOverride != nil) != c.want.hasOverride {
t.Errorf("hasOverride: got %v want %v (override=%v)",
res.WorkspaceOverride != nil, c.want.hasOverride, res.WorkspaceOverride)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Errorf("sqlmock expectations: %v", err)
}
})
}
}
// TestResolveLLMBillingMode_ResolvedModeIsAlwaysValid asserts the resolver's
// post-condition: the returned mode is ALWAYS one of the three known enum
// values, never an empty string and never a garbled passthrough. The strip
// gate downstream relies on this so it can switch on res.ResolvedMode
// without a separate is-valid check on every call site.
func TestResolveLLMBillingMode_ResolvedModeIsAlwaysValid(t *testing.T) {
ctx := context.Background()
const wsID = "22222222-2222-2222-2222-222222222222"
// Throw a pathological row at the resolver: garbled override + garbled
// org default. Resolved mode must still be a recognized enum.
mock := setupTestDB(t)
mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(wsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow("totally-bogus"))
res, err := ResolveLLMBillingMode(ctx, wsID, "also-bogus")
if err != nil {
t.Fatalf("unexpected err: %v", err)
}
if !isKnownBillingMode(res.ResolvedMode) {
t.Errorf("post-condition violated: resolved mode %q is not a known enum value", res.ResolvedMode)
}
if res.ResolvedMode != LLMBillingModePlatformManaged {
t.Errorf("default-closed contract: garbled-x-garbled must resolve to platform_managed, got %q", res.ResolvedMode)
}
}
// TestSetWorkspaceLLMBillingMode_Validation guards the SET path. The CHECK
// constraint at the DB layer is the second line of defense; the route
// handler relies on this function rejecting unknown modes with a clean
// error (so it can map to 400) instead of letting them hit Postgres and
// surfacing as a sql-driver error string.
func TestSetWorkspaceLLMBillingMode_Validation(t *testing.T) {
ctx := context.Background()
const wsID = "33333333-3333-3333-3333-333333333333"
t.Run("rejects_unknown_mode_without_db_call", func(t *testing.T) {
setupTestDB(t) // mock expects nothing — the function must short-circuit
if err := SetWorkspaceLLMBillingMode(ctx, wsID, "totally-bogus"); err == nil {
t.Fatal("expected error for unknown mode, got nil")
}
})
t.Run("rejects_empty_workspace_id", func(t *testing.T) {
setupTestDB(t)
if err := SetWorkspaceLLMBillingMode(ctx, "", LLMBillingModeBYOK); err == nil {
t.Fatal("expected error for empty workspace id, got nil")
}
})
t.Run("clear_uses_NULL_update", func(t *testing.T) {
mock := setupTestDB(t)
mock.ExpectExec(`UPDATE workspaces SET llm_billing_mode = NULL WHERE id = \$1`).
WithArgs(wsID).
WillReturnResult(sqlmock.NewResult(0, 1))
if err := SetWorkspaceLLMBillingMode(ctx, wsID, ""); err != nil {
t.Fatalf("unexpected err: %v", err)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
})
t.Run("set_byok_uses_value_update", func(t *testing.T) {
mock := setupTestDB(t)
mock.ExpectExec(`UPDATE workspaces SET llm_billing_mode = \$1 WHERE id = \$2`).
WithArgs(LLMBillingModeBYOK, wsID).
WillReturnResult(sqlmock.NewResult(0, 1))
if err := SetWorkspaceLLMBillingMode(ctx, wsID, LLMBillingModeBYOK); err != nil {
t.Fatalf("unexpected err: %v", err)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
})
}
+3 -89
View File
@@ -14,9 +14,9 @@ import (
"errors"
"github.com/DATA-DOG/go-sqlmock"
"git.moleculesai.app/molecule-ai/molecule-core/workspace-server/internal/db"
"git.moleculesai.app/molecule-ai/molecule-core/workspace-server/internal/memory/contract"
"github.com/DATA-DOG/go-sqlmock"
"github.com/gin-gonic/gin"
)
@@ -196,7 +196,7 @@ func TestMCPHandler_DelegateTask_RoutesThroughPlatformA2AProxy(t *testing.T) {
expectCanCommunicateSiblings(mock, callerID, targetID, parentID)
mock.ExpectExec(`(?s)INSERT INTO activity_logs.*'delegation'.*'delegate'`).
WithArgs(callerID, callerID, targetID, "Delegating to "+targetID, sqlmock.AnyArg(), "pending").
WithArgs(callerID, callerID, targetID, "Delegating to "+targetID, sqlmock.AnyArg()).
WillReturnResult(sqlmock.NewResult(1, 1))
mock.ExpectExec(`UPDATE activity_logs`).
WithArgs("dispatched", "", callerID, sqlmock.AnyArg()).
@@ -241,7 +241,7 @@ func TestMCPHandler_DelegateTaskAsync_RoutesThroughPlatformA2AProxy(t *testing.T
expectCanCommunicateSiblings(mock, callerID, targetID, parentID)
mock.ExpectExec(`(?s)INSERT INTO activity_logs.*'delegation'.*'delegate'`).
WithArgs(callerID, callerID, targetID, "Delegating to "+targetID, sqlmock.AnyArg(), "pending").
WithArgs(callerID, callerID, targetID, "Delegating to "+targetID, sqlmock.AnyArg()).
WillReturnResult(sqlmock.NewResult(1, 1))
mock.ExpectExec(`UPDATE activity_logs`).
WithArgs("dispatched", "", callerID, sqlmock.AnyArg()).
@@ -280,92 +280,6 @@ func TestMCPHandler_DelegateTaskAsync_RoutesThroughPlatformA2AProxy(t *testing.T
}
}
// TestMCPHandler_DelegateTaskAsync_MarshalFailureDoesNotCallProxy proves the
// extracted #1933 fix: when the A2A body fails to marshal, the detached
// goroutine returns early and never calls proxyA2ARequest with a nil/empty
// body. Before the fix the goroutine logged the error and fell through,
// dispatching a malformed A2A request.
func TestMCPHandler_DelegateTaskAsync_MarshalFailureDoesNotCallProxy(t *testing.T) {
h, mock := newMCPHandler(t)
callerID := "11111111-1111-1111-1111-111111111111"
targetID := "22222222-2222-2222-2222-222222222222"
parentID := "33333333-3333-3333-3333-333333333333"
expectCanCommunicateSiblings(mock, callerID, targetID, parentID)
mock.ExpectExec(`(?s)INSERT INTO activity_logs.*'delegation'.*'delegate'`).
WithArgs(callerID, callerID, targetID, "Delegating to "+targetID, sqlmock.AnyArg(), "pending").
WillReturnResult(sqlmock.NewResult(1, 1))
mock.ExpectExec(`UPDATE activity_logs`).
WithArgs("dispatched", "", callerID, sqlmock.AnyArg()).
WillReturnResult(sqlmock.NewResult(0, 1))
// Force the (otherwise near-impossible) marshal failure for the A2A body.
origMarshal := marshalA2ABody
marshalA2ABody = func(any) ([]byte, error) {
return nil, errors.New("forced marshal failure")
}
t.Cleanup(func() { marshalA2ABody = origMarshal })
proxyCalled := make(chan struct{}, 1)
h.a2aProxy = func(ctx context.Context, workspaceID string, body []byte, proxyCallerID string, logActivity bool) (int, []byte, error) {
proxyCalled <- struct{}{}
return 200, []byte(`{}`), nil
}
out, err := h.toolDelegateTaskAsync(context.Background(), callerID, map[string]interface{}{
"workspace_id": targetID,
"task": "async work",
})
if err != nil {
t.Fatalf("delegate_task_async returned error: %v", err)
}
if !strings.Contains(out, `"status":"dispatched"`) {
t.Fatalf("delegate_task_async response = %s", out)
}
// Wait for the detached goroutine to finish, then assert the proxy was
// never reached because of the early return on marshal failure.
waitGlobalAsyncForTest()
select {
case <-proxyCalled:
t.Fatal("proxyA2ARequest was called after marshal failure; expected early return")
default:
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatalf("unmet expectations: %v", err)
}
}
// TestMCPHandler_CheckTaskStatus_NullStatusDefaultsToUnknown proves the
// extracted #1933 hardening: when the activity_logs row has a NULL status,
// check_task_status reports "unknown" instead of an empty string (the old
// status.String zero value).
func TestMCPHandler_CheckTaskStatus_NullStatusDefaultsToUnknown(t *testing.T) {
h, mock := newMCPHandler(t)
callerID := "11111111-1111-1111-1111-111111111111"
targetID := "22222222-2222-2222-2222-222222222222"
taskID := "task-abc"
mock.ExpectQuery(`(?s)SELECT status, error_detail, response_body.*FROM activity_logs`).
WithArgs(callerID, targetID, taskID).
WillReturnRows(sqlmock.NewRows([]string{"status", "error_detail", "response_body"}).
AddRow(nil, nil, nil))
out, err := h.toolCheckTaskStatus(context.Background(), callerID, map[string]interface{}{
"workspace_id": targetID,
"task_id": taskID,
})
if err != nil {
t.Fatalf("check_task_status returned error: %v", err)
}
if !strings.Contains(out, `"status": "unknown"`) {
t.Fatalf("expected status \"unknown\" for NULL status row, got: %s", out)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatalf("unmet expectations: %v", err)
}
}
// ─────────────────────────────────────────────────────────────────────────────
// notifications/initialized
// ─────────────────────────────────────────────────────────────────────────────
+19 -51
View File
@@ -20,26 +20,18 @@ import (
"github.com/google/uuid"
)
// marshalA2ABody marshals the JSON-RPC body for an async A2A dispatch.
// Indirected through a package var so tests can force the (otherwise
// near-impossible) marshal-failure path and assert the early return.
var marshalA2ABody = json.Marshal
// insertMCPDelegationRow writes a delegation activity row so the canvas
// Agent Comms tab can show the task text for MCP-initiated delegations.
// Mirrors insertDelegationRow (delegation.go) for the MCP tool path.
func insertMCPDelegationRow(ctx context.Context, db *sql.DB, workspaceID, targetID, delegationID, task string) error {
taskJSON, marshalErr := json.Marshal(map[string]interface{}{
taskJSON, _ := json.Marshal(map[string]interface{}{
"task": task,
"delegation_id": delegationID,
})
if marshalErr != nil {
log.Printf("insertMCPDelegationRow %s: json.Marshal taskJSON failed: %v", delegationID, marshalErr)
}
_, err := db.ExecContext(ctx, `
INSERT INTO activity_logs (workspace_id, activity_type, method, source_id, target_id, summary, request_body, status)
VALUES ($1, 'delegation', 'delegate', $2, $3, $4, $5::jsonb, $6)
`, workspaceID, workspaceID, targetID, "Delegating to "+targetID, string(taskJSON), "pending")
VALUES ($1, 'delegation', 'delegate', $2, $3, $4, $5::jsonb, 'pending')
`, workspaceID, workspaceID, targetID, "Delegating to "+targetID, string(taskJSON))
return err
}
@@ -97,15 +89,7 @@ func (h *MCPHandler) toolListPeers(ctx context.Context, workspaceID string) (str
const cols = `SELECT w.id, w.name, COALESCE(w.role,''), w.status, w.tier`
// Siblings — workspaces sharing the caller's parent.
//
// #1953 cross-tenant isolation: the OLD else-branch returned every
// workspace with parent_id IS NULL when the caller was itself an org root,
// i.e. every other tenant's org root (the workspaces table has no org_id
// column). That leaked peer identities across tenants via MCP list_peers.
// An org root has no siblings inside its own org, so the org-root caller
// now gets no siblings; its peers are its children, enumerated below. Only
// the parent_id-bound branch enumerates siblings, scoped to one tenant.
// Siblings
if parentID.Valid {
rows, err := h.database.QueryContext(ctx,
cols+` FROM workspaces w WHERE w.parent_id = $1 AND w.id != $2 AND w.status != 'removed'`,
@@ -115,6 +99,15 @@ func (h *MCPHandler) toolListPeers(ctx context.Context, workspaceID string) (str
log.Printf("MCP toolListPeers: sibling scan error: %v", scanErr)
}
}
} else {
rows, err := h.database.QueryContext(ctx,
cols+` FROM workspaces w WHERE w.parent_id IS NULL AND w.id != $1 AND w.status != 'removed'`,
workspaceID)
if err == nil {
if scanErr := scanPeers(rows); scanErr != nil {
log.Printf("MCP toolListPeers: sibling scan error: %v", scanErr)
}
}
}
// Children
@@ -145,11 +138,7 @@ func (h *MCPHandler) toolListPeers(ctx context.Context, workspaceID string) (str
return "No peers found.", nil
}
b, marshalErr := json.MarshalIndent(peers, "", " ")
if marshalErr != nil {
log.Printf("toolListPeers: json.MarshalIndent peers failed: %v", marshalErr)
return "", fmt.Errorf("marshal response: %w", marshalErr)
}
b, _ := json.MarshalIndent(peers, "", " ")
return string(b), nil
}
@@ -179,11 +168,7 @@ func (h *MCPHandler) toolGetWorkspaceInfo(ctx context.Context, workspaceID strin
if parentID.Valid {
info["parent_id"] = parentID.String
}
b, marshalErr := json.MarshalIndent(info, "", " ")
if marshalErr != nil {
log.Printf("toolGetWorkspaceInfo %s: json.MarshalIndent info failed: %v", workspaceID, marshalErr)
return "", fmt.Errorf("marshal response: %w", marshalErr)
}
b, _ := json.MarshalIndent(info, "", " ")
return string(b), nil
}
@@ -275,7 +260,7 @@ func (h *MCPHandler) toolDelegateTaskAsync(ctx context.Context, callerID string,
bgCtx, cancel := context.WithTimeout(context.Background(), mcpAsyncCallTimeout)
defer cancel()
a2aBody, marshalErr := marshalA2ABody(map[string]interface{}{
a2aBody, _ := json.Marshal(map[string]interface{}{
"jsonrpc": "2.0",
"id": delegationID,
"method": "message/send",
@@ -287,12 +272,6 @@ func (h *MCPHandler) toolDelegateTaskAsync(ctx context.Context, callerID string,
},
},
})
if marshalErr != nil {
log.Printf("toolDelegateTask %s: json.Marshal a2aBody failed: %v", delegationID, marshalErr)
// Bail out: proceeding would call proxyA2ARequest with a
// nil/empty body, dispatching a malformed A2A request.
return
}
status, _, err := h.proxyA2ARequest(bgCtx, targetID, a2aBody, callerID, true)
if err != nil || status < 200 || status >= 300 {
@@ -339,24 +318,16 @@ func (h *MCPHandler) toolCheckTaskStatus(ctx context.Context, callerID string, a
result := map[string]interface{}{
"task_id": taskID,
"status": status.String,
"target_id": targetID,
}
if status.Valid {
result["status"] = status.String
} else {
result["status"] = "unknown"
}
if errorDetail.Valid && errorDetail.String != "" {
result["error"] = errorDetail.String
}
if len(responseBody) > 0 {
result["result"] = extractA2AText(responseBody)
}
b, marshalErr := json.MarshalIndent(result, "", " ")
if marshalErr != nil {
log.Printf("toolCheckTaskStatus: json.MarshalIndent result failed: %v", marshalErr)
return "", fmt.Errorf("marshal response: %w", marshalErr)
}
b, _ := json.MarshalIndent(result, "", " ")
return string(b), nil
}
@@ -511,9 +482,6 @@ func extractA2AText(body []byte) string {
}
// Fallback: marshal result as JSON.
b, marshalErr := json.Marshal(result)
if marshalErr != nil {
log.Printf("extractA2AText: json.Marshal result failed: %v", marshalErr)
}
b, _ := json.Marshal(result)
return string(b)
}
@@ -25,7 +25,6 @@ import (
"context"
"encoding/json"
"fmt"
"log"
"strings"
"git.moleculesai.app/molecule-ai/molecule-core/workspace-server/internal/memory/contract"
@@ -191,11 +190,7 @@ func (h *MCPHandler) recallMemoryLegacyShim(ctx context.Context, workspaceID str
if len(out) == 0 {
return "No memories found.", nil
}
b, marshalErr := json.MarshalIndent(out, "", " ")
if marshalErr != nil {
log.Printf("toolRecallMemory: json.MarshalIndent out failed: %v", marshalErr)
return "", fmt.Errorf("marshal response: %w", marshalErr)
}
b, _ := json.MarshalIndent(out, "", " ")
return string(b), nil
}
@@ -48,7 +48,6 @@ type memoryV2Deps struct {
// call. Defining an interface here lets handler tests stub the plugin
// without spinning up an HTTP server.
type memoryPluginAPI interface {
UpsertNamespace(ctx context.Context, name string, body contract.NamespaceUpsert) (*contract.Namespace, error)
CommitMemory(ctx context.Context, namespace string, body contract.MemoryWrite) (*contract.MemoryWriteResponse, error)
Search(ctx context.Context, body contract.SearchRequest) (*contract.SearchResponse, error)
ForgetMemory(ctx context.Context, id string, body contract.ForgetRequest) error
@@ -118,9 +117,6 @@ func (h *MCPHandler) toolCommitMemoryV2(ctx context.Context, workspaceID string,
if !ok {
return "", fmt.Errorf("workspace %s cannot write to namespace %s", workspaceID, ns)
}
if _, err := h.memv2.plugin.UpsertNamespace(ctx, ns, contract.NamespaceUpsert{Kind: kindFromNamespace(ns)}); err != nil {
return "", fmt.Errorf("plugin upsert namespace: %w", err)
}
// SAFE-T1201: scrub credential-shaped strings BEFORE the plugin sees
// them. Non-negotiable; see memories.go:180.
@@ -167,27 +163,10 @@ func (h *MCPHandler) toolCommitMemoryV2(ctx context.Context, workspaceID string,
summary := "commit_memory to " + ns
logMemoryMCPActivity(ctx, h.broadcaster, workspaceID, "memory_write", resp.ID, ns, &summary)
out, marshalErr := json.Marshal(resp)
if marshalErr != nil {
log.Printf("toolCommitMemoryV2 %s: json.Marshal resp failed: %v", workspaceID, marshalErr)
return "", fmt.Errorf("marshal response: %w", marshalErr)
}
out, _ := json.Marshal(resp)
return string(out), nil
}
func kindFromNamespace(ns string) contract.NamespaceKind {
switch {
case strings.HasPrefix(ns, "workspace:"):
return contract.NamespaceKindWorkspace
case strings.HasPrefix(ns, "team:"):
return contract.NamespaceKindTeam
case strings.HasPrefix(ns, "org:"):
return contract.NamespaceKindOrg
default:
return contract.NamespaceKindCustom
}
}
// ─────────────────────────────────────────────────────────────────────────────
// search_memory
// ─────────────────────────────────────────────────────────────────────────────
@@ -238,11 +217,7 @@ func (h *MCPHandler) toolSearchMemory(ctx context.Context, workspaceID string, a
}
}
out, marshalErr := json.Marshal(resp)
if marshalErr != nil {
log.Printf("toolSearchMemory %s: json.Marshal resp failed: %v", workspaceID, marshalErr)
return "", fmt.Errorf("marshal response: %w", marshalErr)
}
out, _ := json.Marshal(resp)
return string(out), nil
}
@@ -297,11 +272,7 @@ func (h *MCPHandler) toolCommitSummary(ctx context.Context, workspaceID string,
summary := "commit_summary to " + ns
logMemoryMCPActivity(ctx, h.broadcaster, workspaceID, "memory_summary_write", resp.ID, ns, &summary)
out, marshalErr := json.Marshal(resp)
if marshalErr != nil {
log.Printf("toolCommitSummary %s: json.Marshal resp failed: %v", workspaceID, marshalErr)
return "", fmt.Errorf("marshal response: %w", marshalErr)
}
out, _ := json.Marshal(resp)
return string(out), nil
}
@@ -317,11 +288,7 @@ func (h *MCPHandler) toolListWritableNamespaces(ctx context.Context, workspaceID
if err != nil {
return "", fmt.Errorf("resolve writable: %w", err)
}
b, marshalErr := json.MarshalIndent(ns, "", " ")
if marshalErr != nil {
log.Printf("toolListWritableNamespaces %s: json.MarshalIndent ns failed: %v", workspaceID, marshalErr)
return "", fmt.Errorf("marshal response: %w", marshalErr)
}
b, _ := json.MarshalIndent(ns, "", " ")
return string(b), nil
}
@@ -333,11 +300,7 @@ func (h *MCPHandler) toolListReadableNamespaces(ctx context.Context, workspaceID
if err != nil {
return "", fmt.Errorf("resolve readable: %w", err)
}
b, marshalErr := json.MarshalIndent(ns, "", " ")
if marshalErr != nil {
log.Printf("toolListReadableNamespaces %s: json.MarshalIndent ns failed: %v", workspaceID, marshalErr)
return "", fmt.Errorf("marshal response: %w", marshalErr)
}
b, _ := json.MarshalIndent(ns, "", " ")
return string(b), nil
}
@@ -20,18 +20,11 @@ import (
// --- stubs ---
type stubMemoryPlugin struct {
upsertFn func(ctx context.Context, name string, body contract.NamespaceUpsert) (*contract.Namespace, error)
commitFn func(ctx context.Context, ns string, body contract.MemoryWrite) (*contract.MemoryWriteResponse, error)
searchFn func(ctx context.Context, body contract.SearchRequest) (*contract.SearchResponse, error)
forgetFn func(ctx context.Context, id string, body contract.ForgetRequest) error
}
func (s *stubMemoryPlugin) UpsertNamespace(ctx context.Context, name string, body contract.NamespaceUpsert) (*contract.Namespace, error) {
if s.upsertFn != nil {
return s.upsertFn(ctx, name, body)
}
return &contract.Namespace{Name: name, Kind: body.Kind}, nil
}
func (s *stubMemoryPlugin) CommitMemory(ctx context.Context, ns string, body contract.MemoryWrite) (*contract.MemoryWriteResponse, error) {
if s.commitFn != nil {
return s.commitFn(ctx, ns, body)
@@ -166,15 +159,7 @@ func TestMemoryV2Available(t *testing.T) {
func TestCommitMemoryV2_HappyPathDefaultNamespace(t *testing.T) {
db, _, _ := sqlmock.New()
defer db.Close()
gotUpsertNS := ""
h := newV2Handler(t, db, &stubMemoryPlugin{
upsertFn: func(_ context.Context, name string, body contract.NamespaceUpsert) (*contract.Namespace, error) {
gotUpsertNS = name
if body.Kind != contract.NamespaceKindWorkspace {
t.Errorf("upsert kind = %q, want workspace", body.Kind)
}
return &contract.Namespace{Name: name, Kind: body.Kind}, nil
},
commitFn: func(_ context.Context, ns string, body contract.MemoryWrite) (*contract.MemoryWriteResponse, error) {
if ns != "workspace:root-1" {
t.Errorf("ns = %q, want default workspace:root-1", ns)
@@ -195,9 +180,6 @@ func TestCommitMemoryV2_HappyPathDefaultNamespace(t *testing.T) {
if !strings.Contains(got, `"id":"mem-1"`) {
t.Errorf("got = %s", got)
}
if gotUpsertNS != "workspace:root-1" {
t.Errorf("upsert namespace = %q, want workspace:root-1", gotUpsertNS)
}
}
func TestCommitMemoryV2_NamespaceParamUsed(t *testing.T) {
@@ -1,12 +1,8 @@
package handlers
import (
"context"
"encoding/json"
"testing"
"git.moleculesai.app/molecule-ai/molecule-core/workspace-server/internal/db"
"github.com/DATA-DOG/go-sqlmock"
)
// ─────────────────────────────────────────────────────────────────────────────
@@ -195,115 +191,3 @@ func TestExtractA2AText_PriorityArtifactsOverMessage(t *testing.T) {
t.Errorf("artifacts should take priority: got %q, want %q", got, want)
}
}
// ─────────────────────────────────────────────────────────────────────────────
// insertMCPDelegationRow tests
// ─────────────────────────────────────────────────────────────────────────────
func TestInsertMCPDelegationRow_Success(t *testing.T) {
mockDB, mock, err := sqlmock.New()
if err != nil {
t.Fatalf("failed to create sqlmock: %v", err)
}
prevDB := db.DB
db.DB = mockDB
t.Cleanup(func() { db.DB = prevDB; mockDB.Close() })
mock.ExpectExec(`INSERT INTO activity_logs`).
WithArgs("ws-src", "ws-src", "ws-tgt", "Delegating to ws-tgt", sqlmock.AnyArg(), "pending").
WillReturnResult(sqlmock.NewResult(0, 1))
err = insertMCPDelegationRow(context.Background(), mockDB, "ws-src", "ws-tgt", "del-123", "summarise the report")
if err != nil {
t.Errorf("unexpected error: %v", err)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Errorf("sqlmock expectations: %v", err)
}
}
func TestInsertMCPDelegationRow_DBError(t *testing.T) {
mockDB, mock, err := sqlmock.New()
if err != nil {
t.Fatalf("failed to create sqlmock: %v", err)
}
prevDB := db.DB
db.DB = mockDB
t.Cleanup(func() { db.DB = prevDB; mockDB.Close() })
mock.ExpectExec(`INSERT INTO activity_logs`).
WithArgs("ws-src", "ws-src", "ws-tgt", sqlmock.AnyArg(), sqlmock.AnyArg(), "pending").
WillReturnError(context.DeadlineExceeded)
err = insertMCPDelegationRow(context.Background(), mockDB, "ws-src", "ws-tgt", "del-456", "check the logs")
if err == nil {
t.Error("expected error, got nil")
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Errorf("sqlmock expectations: %v", err)
}
}
// ─────────────────────────────────────────────────────────────────────────────
// updateMCPDelegationStatus tests
// ─────────────────────────────────────────────────────────────────────────────
func TestUpdateMCPDelegationStatus_Success(t *testing.T) {
mockDB, mock, err := sqlmock.New()
if err != nil {
t.Fatalf("failed to create sqlmock: %v", err)
}
prevDB := db.DB
db.DB = mockDB
t.Cleanup(func() { db.DB = prevDB; mockDB.Close() })
mock.ExpectExec(`UPDATE activity_logs`).
WithArgs("completed", "", "ws-src", "del-789").
WillReturnResult(sqlmock.NewResult(0, 1))
// Should not panic, should not error
updateMCPDelegationStatus(context.Background(), mockDB, "ws-src", "del-789", "completed", "")
if err := mock.ExpectationsWereMet(); err != nil {
t.Errorf("sqlmock expectations: %v", err)
}
}
func TestUpdateMCPDelegationStatus_WithErrorDetail(t *testing.T) {
mockDB, mock, err := sqlmock.New()
if err != nil {
t.Fatalf("failed to create sqlmock: %v", err)
}
prevDB := db.DB
db.DB = mockDB
t.Cleanup(func() { db.DB = prevDB; mockDB.Close() })
mock.ExpectExec(`UPDATE activity_logs`).
WithArgs("failed", "timeout", "ws-src", "del-000").
WillReturnResult(sqlmock.NewResult(0, 1))
updateMCPDelegationStatus(context.Background(), mockDB, "ws-src", "del-000", "failed", "timeout")
if err := mock.ExpectationsWereMet(); err != nil {
t.Errorf("sqlmock expectations: %v", err)
}
}
func TestUpdateMCPDelegationStatus_DBError_LoggedNotReturned(t *testing.T) {
mockDB, mock, err := sqlmock.New()
if err != nil {
t.Fatalf("failed to create sqlmock: %v", err)
}
prevDB := db.DB
db.DB = mockDB
t.Cleanup(func() { db.DB = prevDB; mockDB.Close() })
mock.ExpectExec(`UPDATE activity_logs`).
WithArgs("failed", sqlmock.AnyArg(), "ws-src", "del-abc").
WillReturnError(context.DeadlineExceeded)
// Function returns no value — error is logged, not propagated.
// Verify it does not panic.
updateMCPDelegationStatus(context.Background(), mockDB, "ws-src", "del-abc", "failed", "connection refused")
if err := mock.ExpectationsWereMet(); err != nil {
t.Errorf("sqlmock expectations: %v", err)
}
}
+7 -11
View File
@@ -240,21 +240,17 @@ func (h *MemoriesHandler) Commit(c *gin.Context) {
// Hash the sanitised content so the audit trail reflects what was
// actually persisted (not the raw, potentially secret-bearing input).
sum := sha256.Sum256([]byte(content))
auditBody, marshalErr := json.Marshal(map[string]string{
auditBody, _ := json.Marshal(map[string]string{
"memory_id": memoryID,
"namespace": nsName,
"content_sha256": hex.EncodeToString(sum[:]),
})
if marshalErr != nil {
log.Printf("Commit %s: json.Marshal auditBody failed: %v", workspaceID, marshalErr)
} else {
summary := "GLOBAL memory written: id=" + memoryID + " namespace=" + nsName
if _, auditErr := db.DB.ExecContext(ctx, `
INSERT INTO activity_logs (workspace_id, activity_type, source_id, summary, request_body, status)
VALUES ($1, $2, $3, $4, $5::jsonb, $6)
`, workspaceID, "memory_write_global", workspaceID, summary, string(auditBody), "ok"); auditErr != nil {
log.Printf("Commit: GLOBAL memory audit log failed for %s/%s: %v", workspaceID, memoryID, auditErr)
}
summary := "GLOBAL memory written: id=" + memoryID + " namespace=" + nsName
if _, auditErr := db.DB.ExecContext(ctx, `
INSERT INTO activity_logs (workspace_id, activity_type, source_id, summary, request_body, status)
VALUES ($1, $2, $3, $4, $5::jsonb, $6)
`, workspaceID, "memory_write_global", workspaceID, summary, string(auditBody), "ok"); auditErr != nil {
log.Printf("Commit: GLOBAL memory audit log failed for %s/%s: %v", workspaceID, memoryID, auditErr)
}
}
@@ -45,9 +45,6 @@ type fakePlugin struct {
forgetReq contract.ForgetRequest
}
func (f *fakePlugin) UpsertNamespace(ctx context.Context, name string, body contract.NamespaceUpsert) (*contract.Namespace, error) {
return &contract.Namespace{Name: name, Kind: body.Kind}, nil
}
func (f *fakePlugin) CommitMemory(ctx context.Context, ns string, body contract.MemoryWrite) (*contract.MemoryWriteResponse, error) {
return nil, errors.New("not implemented in fake")
}
@@ -514,11 +511,11 @@ func TestMemoriesV2_Forget_MissingMemoryID_400(t *testing.T) {
// DisplayName over UUID-prefix fallback (issue #2988).
func TestNamespaceLabelWithName_PrefersDisplayNameWhenSet(t *testing.T) {
cases := []struct {
name string
raw string
kind contract.NamespaceKind
display string
want string
name string
raw string
kind contract.NamespaceKind
display string
want string
}{
{"workspace with name", "workspace:abc-1234", contract.NamespaceKindWorkspace, "mac laptop", "Workspace (mac laptop)"},
{"team with name", "team:abc-1234", contract.NamespaceKindTeam, "Engineering", "Team (Engineering)"},
@@ -628,12 +625,12 @@ func TestParseLimit(t *testing.T) {
}{
{"", memoriesV2DefaultLimit},
{"10", 10},
{"0", memoriesV2DefaultLimit}, // ≤0 → default, not error
{"-5", memoriesV2DefaultLimit}, // negative → default
{"abc", memoriesV2DefaultLimit}, // non-numeric → default
{"99999", memoriesV2MaxLimit}, // over cap → clamped
{"100", memoriesV2MaxLimit}, // exactly cap → kept
{"99", 99}, // just under cap → kept
{"0", memoriesV2DefaultLimit}, // ≤0 → default, not error
{"-5", memoriesV2DefaultLimit}, // negative → default
{"abc", memoriesV2DefaultLimit}, // non-numeric → default
{"99999", memoriesV2MaxLimit}, // over cap → clamped
{"100", memoriesV2MaxLimit}, // exactly cap → kept
{"99", 99}, // just under cap → kept
}
for _, tc := range cases {
t.Run("raw="+tc.raw, func(t *testing.T) {
@@ -744,11 +741,11 @@ func TestWithMemoryV2_FluentReturnsReceiver(t *testing.T) {
func TestShortID(t *testing.T) {
cases := map[string]string{
"": "",
"short": "short",
"exactly8": "exactly8",
"longer-than-eight": "longer-t",
"abc-1234-5678-90ab": "abc-1234",
"": "",
"short": "short",
"exactly8": "exactly8",
"longer-than-eight": "longer-t",
"abc-1234-5678-90ab": "abc-1234",
}
for in, want := range cases {
if got := shortID(in); got != want {
@@ -1,104 +0,0 @@
package handlers
// org_scope.go — cross-tenant isolation helpers (#1953).
//
// The `workspaces` table has no `org_id` column; an "org" is the subtree of
// workspaces reachable through the `parent_id` chain from a single org root
// (a row with parent_id IS NULL). Several code paths historically computed an
// org-root sibling set as `WHERE parent_id IS NULL`, which matches EVERY
// tenant's org root and therefore leaks peer metadata / routing across tenants.
//
// This file centralises the org-scoping primitive so peer discovery, the MCP
// list_peers tool, and a2a routing all derive "the caller's org" the SAME way
// the OFFSEC-015 broadcast fix (commit 5a05302c, workspace_broadcast.go) does:
// a recursive CTE that walks the parent_id chain up to the org root. Keeping
// the CTE in one place means there is a single, testable source of truth for
// tenant isolation rather than four hand-copied queries that can drift.
//
// NOTE: this is the parent_id-chain scoping that the broadcast fix already
// ships. It is deliberately NOT an `org_id` column — adding that column is a
// separate architecture decision pending CTO sign-off. See #1953.
import (
"context"
"database/sql"
"errors"
)
// errNoOrgRoot is returned by orgRootID when the workspace id has no row (and
// therefore no resolvable org root). Callers translate this into a 404/not-found
// at their own layer; it is distinct from a transient DB error so a missing
// workspace never gets treated as "belongs to every org".
var errNoOrgRoot = errors.New("org root not found for workspace")
// orgRootSubtreeCTE is the recursive CTE — identical in shape to the OFFSEC-015
// broadcast fix — that walks UP the parent_id chain from a single workspace to
// its org root. The org root is the row on the chain whose parent_id IS NULL.
//
// $1 = workspace id to resolve
//
// The recursive member walks UP the parent_id chain: each step joins to the row
// whose id is the current row's parent_id. The topmost ancestor is the single
// chain row with parent_id IS NULL — and THAT row's own `id` is the org root.
//
// We select that parentless row's `id` (aliased root_id). We must NOT carry a
// fixed `id AS root_id` from the recursive seed: that value is just the input
// workspace id, so a non-root caller (e.g. a child delegating to a sibling)
// would resolve to ITSELF instead of its org root, and sameOrg() would wrongly
// report two genuinely same-org workspaces as different orgs and 403 a
// legitimate a2a route. A workspace that already IS an org root has a one-row
// chain whose id == itself, so it correctly resolves to itself.
const orgRootSubtreeCTE = `
WITH RECURSIVE org_chain AS (
SELECT id, parent_id
FROM workspaces
WHERE id = $1
UNION ALL
SELECT w.id, w.parent_id
FROM workspaces w
JOIN org_chain c ON w.id = c.parent_id
)
SELECT id AS root_id FROM org_chain WHERE parent_id IS NULL LIMIT 1
`
// orgRootID resolves the org root of `workspaceID` by walking the parent_id
// chain via orgRootSubtreeCTE. Returns errNoOrgRoot when the workspace (or its
// chain) yields no org root row, and the underlying error on any DB failure.
//
// This is the SAME lookup the broadcast handler performs inline; the three
// leak paths in #1953 call this instead of re-deriving "the org" from
// `parent_id IS NULL` (which spans all tenants).
func orgRootID(ctx context.Context, database *sql.DB, workspaceID string) (string, error) {
var root string
err := database.QueryRowContext(ctx, orgRootSubtreeCTE, workspaceID).Scan(&root)
if errors.Is(err, sql.ErrNoRows) {
return "", errNoOrgRoot
}
if err != nil {
return "", err
}
if root == "" {
return "", errNoOrgRoot
}
return root, nil
}
// sameOrg reports whether workspaces `a` and `b` share an org root, i.e. they
// belong to the same tenant. Used by a2a routing to reject resolving/dispatching
// to a workspace id outside the caller's org. Fail-CLOSED: any lookup error or
// missing org root yields (false, err) so a DB hiccup denies cross-tenant
// routing rather than allowing it.
func sameOrg(ctx context.Context, database *sql.DB, a, b string) (bool, error) {
if a == b {
return true, nil
}
rootA, err := orgRootID(ctx, database, a)
if err != nil {
return false, err
}
rootB, err := orgRootID(ctx, database, b)
if err != nil {
return false, err
}
return rootA == rootB, nil
}
@@ -1,7 +1,6 @@
package handlers
import (
"io"
"log"
"net/http"
@@ -69,10 +68,7 @@ type createOrgTokenResponse struct {
func (h *OrgTokenHandler) Create(c *gin.Context) {
var req createOrgTokenRequest
// Optional body — an empty POST should still work (unnamed token).
if err := c.ShouldBindJSON(&req); err != nil && err != io.EOF {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid JSON body"})
return
}
_ = c.ShouldBindJSON(&req)
if len(req.Name) > 100 {
c.JSON(http.StatusBadRequest, gin.H{"error": "name too long (max 100 chars)"})
return
+4 -16
View File
@@ -345,16 +345,8 @@ func (h *RegistryHandler) Register(c *gin.Context) {
if qErr := db.DB.QueryRowContext(ctx,
`SELECT name, role FROM workspaces WHERE id = $1`, payload.ID,
).Scan(&dbName, &dbRole); qErr == nil {
name := ""
if dbName.Valid {
name = dbName.String
}
role := ""
if dbRole.Valid {
role = dbRole.String
}
if rc, did := reconcileAgentCardIdentity(
payload.AgentCard, payload.ID, name, role,
payload.AgentCard, payload.ID, dbName.String, dbRole.String,
); did {
reconciledCard = rc
log.Printf("Registry register: reconciled agent_card identity for %s from workspaces row", payload.ID)
@@ -538,9 +530,7 @@ func (h *RegistryHandler) Heartbeat(c *gin.Context) {
// Read previous current_task to detect changes (before the UPDATE)
var prevTask string
if err := db.DB.QueryRowContext(ctx, `SELECT COALESCE(current_task, '') FROM workspaces WHERE id = $1`, payload.WorkspaceID).Scan(&prevTask); err != nil {
log.Printf("registry heartbeat: prev_task query failed for workspace %s: %v", payload.WorkspaceID, err)
}
_ = db.DB.QueryRowContext(ctx, `SELECT COALESCE(current_task, '') FROM workspaces WHERE id = $1`, payload.WorkspaceID).Scan(&prevTask)
// #615: Clamp monthly_spend to a safe range before any DB write.
// A malicious or buggy agent could report math.MaxInt64, causing
@@ -822,12 +812,10 @@ func (h *RegistryHandler) evaluateStatus(c *gin.Context, payload models.Heartbea
// timeouts, retry logic, and activity_logs wiring.
if h.drainQueue != nil {
var maxConcurrent int
if err := db.DB.QueryRowContext(ctx,
_ = db.DB.QueryRowContext(ctx,
`SELECT COALESCE(max_concurrent_tasks, 1) FROM workspaces WHERE id = $1`,
payload.WorkspaceID,
).Scan(&maxConcurrent); err != nil {
log.Printf("registry heartbeat: max_concurrent query failed for workspace %s: %v", payload.WorkspaceID, err)
}
).Scan(&maxConcurrent)
if payload.ActiveTasks < maxConcurrent {
// context.WithoutCancel: heartbeat handler's ctx is about to
// expire as soon as we return. The drain needs to outlive it.
@@ -177,12 +177,10 @@ func waitForWorkspaceOnline(ctx context.Context, workspaceID string, timeout tim
).Scan(&status); err == nil && status == "online" {
return true
}
timer := time.NewTimer(restartContextOnlinePollInterval)
select {
case <-ctx.Done():
timer.Stop()
return false
case <-timer.C:
case <-time.After(restartContextOnlinePollInterval):
}
}
return false
@@ -215,12 +213,10 @@ func waitForFreshHeartbeat(ctx context.Context, workspaceID string, restartStart
lastHB.Valid && lastHB.Time.After(restartStartTs) {
return true
}
timer := time.NewTimer(restartContextOnlinePollInterval)
select {
case <-ctx.Done():
timer.Stop()
return false
case <-timer.C:
case <-time.After(restartContextOnlinePollInterval):
}
}
return false
@@ -80,10 +80,7 @@ func (h *WorkspaceHandler) gracefulPreRestart(ctx context.Context, workspaceID s
},
"id": nil,
}
body, marshalErr := json.Marshal(payload)
if marshalErr != nil {
log.Printf("A2AGracefulRestart %s: json.Marshal payload failed: %v", workspaceID, marshalErr)
}
body, _ := json.Marshal(payload)
req, reqErr := http.NewRequestWithContext(signalCtx, http.MethodPost, url, bytes.NewReader(body))
if reqErr != nil {
@@ -160,14 +160,13 @@ func (h *ScheduleHandler) Create(c *gin.Context) {
}
// Validate timezone
loc, err := time.LoadLocation(body.Timezone)
if err != nil {
if _, err := time.LoadLocation(body.Timezone); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid timezone: " + body.Timezone})
return
}
// Validate and compute next run
nextRun, err := scheduler.ComputeNextRun(body.CronExpr, body.Timezone, time.Now().In(loc))
nextRun, err := scheduler.ComputeNextRun(body.CronExpr, body.Timezone, time.Now())
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request body"})
return
@@ -261,12 +260,11 @@ func (h *ScheduleHandler) Update(c *gin.Context) {
if body.Timezone != nil {
tz = *body.Timezone
}
loc, err := time.LoadLocation(tz)
if err != nil {
if _, err := time.LoadLocation(tz); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid timezone: " + tz})
return
}
nextRun, err := scheduler.ComputeNextRun(cronExpr, tz, time.Now().In(loc))
nextRun, err := scheduler.ComputeNextRun(cronExpr, tz, time.Now())
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request body"})
return
@@ -290,12 +288,7 @@ func (h *ScheduleHandler) Update(c *gin.Context) {
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to update schedule"})
return
}
n, err := result.RowsAffected()
if err != nil {
log.Printf("Schedules.Update: RowsAffected error: %v", err)
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to update schedule"})
return
}
n, _ := result.RowsAffected()
if n == 0 {
c.JSON(http.StatusNotFound, gin.H{"error": "schedule not found"})
return
@@ -328,12 +321,7 @@ func (h *ScheduleHandler) Delete(c *gin.Context) {
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to delete schedule"})
return
}
n, err := result.RowsAffected()
if err != nil {
log.Printf("Schedules.Delete: RowsAffected error: %v", err)
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to delete schedule"})
return
}
n, _ := result.RowsAffected()
if n == 0 {
c.JSON(http.StatusNotFound, gin.H{"error": "schedule not found"})
return
+3 -107
View File
@@ -19,28 +19,8 @@ import (
var uuidRegex = regexp.MustCompile(`^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$`)
var platformManagedDirectLLMBypassKeys = map[string]struct{}{
"AI_GATEWAY_API_KEY": {},
"ANTHROPIC_API_KEY": {},
"ANTHROPIC_AUTH_TOKEN": {},
"ARCEEAI_API_KEY": {},
"CLAUDE_CODE_OAUTH_TOKEN": {},
"DASHSCOPE_API_KEY": {},
"DEEPSEEK_API_KEY": {},
"GEMINI_API_KEY": {},
"GLM_API_KEY": {},
"HERMES_CUSTOM_API_KEY": {},
"HERMES_CUSTOM_BASE_URL": {},
"HF_TOKEN": {},
"KIMI_API_KEY": {},
"KIMI_CN_API_KEY": {},
"MINIMAX_API_KEY": {},
"MINIMAX_CN_API_KEY": {},
"NOUS_API_KEY": {},
"OPENAI_API_KEY": {},
"OPENAI_BASE_URL": {},
"OPENROUTER_API_KEY": {},
"XAI_API_KEY": {},
"ZAI_API_KEY": {},
"HERMES_CUSTOM_API_KEY": {},
"HERMES_CUSTOM_BASE_URL": {},
}
func isPlatformManagedDirectLLMBypassKey(key string) bool {
@@ -48,54 +28,10 @@ func isPlatformManagedDirectLLMBypassKey(key string) bool {
return ok
}
// platformManagedLLMModeForWorkspace replaces the org-level platformManagedLLMMode
// gate with a per-workspace resolved-mode check (internal#691). The strip-list
// is enforced ONLY when this specific workspace's resolved mode is
// platform_managed — a workspace with a byok override is allowed to write its
// own CLAUDE_CODE_OAUTH_TOKEN / vendor key via the canvas Secrets tab.
//
// Default-closed: if the resolver hits a DB error, falls back to
// platform_managed (the safe-default behavior), so a transient DB failure
// during a secret write still rejects the bypass-list keys — fail safer not
// freer. This matches the resolver's documented contract.
func platformManagedLLMModeForWorkspace(c *gin.Context, workspaceID string) bool {
orgMode := strings.ToLower(strings.TrimSpace(os.Getenv("MOLECULE_LLM_BILLING_MODE")))
res, err := ResolveLLMBillingMode(c.Request.Context(), workspaceID, orgMode)
if err != nil {
log.Printf("secrets: resolve billing mode for workspace=%s failed: %v (defaulting to platform_managed for safety)", workspaceID, err)
}
return strings.EqualFold(res.ResolvedMode, LLMBillingModePlatformManaged)
}
// platformManagedLLMMode is the legacy org-level gate retained for any test
// harness still asserting the env-var-only behavior. Production code paths
// must call platformManagedLLMModeForWorkspace instead so a workspace-level
// byok override actually takes effect on the secrets-write path.
func platformManagedLLMMode() bool {
return strings.EqualFold(strings.TrimSpace(os.Getenv("MOLECULE_LLM_BILLING_MODE")), "platform_managed")
}
// rejectPlatformManagedDirectLLMBypassForWorkspace is the per-workspace
// successor to rejectPlatformManagedDirectLLMBypass (internal#691). The
// strip-list ONLY applies when this specific workspace resolves to
// platform_managed; byok/disabled workspaces can write their own vendor keys.
func rejectPlatformManagedDirectLLMBypassForWorkspace(c *gin.Context, workspaceID, key string) bool {
if !platformManagedLLMModeForWorkspace(c, workspaceID) || !isPlatformManagedDirectLLMBypassKey(key) {
return false
}
c.JSON(http.StatusBadRequest, gin.H{
"error": "direct vendor key writes are blocked for platform-managed workspaces; use MODEL/LLM_PROVIDER or the platform LLM proxy env instead, or set this workspace's billing mode to 'byok' via /admin/workspaces/:id/llm-billing-mode",
"key": key,
"workspace_id": workspaceID,
})
return true
}
// rejectPlatformManagedDirectLLMBypass is the legacy org-level shim. Retained
// only for backwards compatibility with any external/test caller still on the
// old shape; new code MUST use the per-workspace variant above. Production
// code paths (the secrets.go handlers + workspace.go create-secret path) all
// switched in internal#691.
func rejectPlatformManagedDirectLLMBypass(c *gin.Context, key string) bool {
if !platformManagedLLMMode() || !isPlatformManagedDirectLLMBypassKey(key) {
return false
@@ -245,11 +181,6 @@ func (h *SecretsHandler) Values(c *gin.Context) {
// provisioner path in workspace_provision.go so env-vars look identical
// whether the workspace was bootstrapped locally or remotely).
out := map[string]string{}
// Provenance side-channel (internal#711): which keys in `out` originated
// from global_secrets and were NOT overridden by a workspace_secrets row.
// Used by the provider-aware gate below so a non-platform workspace's
// remote pull never receives the platform's scope:global LLM credential.
globalKeys := map[string]struct{}{}
// Track decrypt failures so we can refuse the response with a list
// instead of returning a partial bundle that boots a broken agent.
var failedKeys []string
@@ -275,7 +206,6 @@ func (h *SecretsHandler) Values(c *gin.Context) {
continue
}
out[k] = string(decrypted)
globalKeys[k] = struct{}{}
}
}
if err := globalRows.Err(); err != nil {
@@ -300,10 +230,6 @@ func (h *SecretsHandler) Values(c *gin.Context) {
continue
}
out[k] = string(decrypted) // workspace override wins over global
// User explicitly re-set this via the canvas Secrets tab — it is
// no longer "the operator-store version", so drop the global
// provenance flag (mirrors loadWorkspaceSecrets).
delete(globalKeys, k)
}
}
if err := wsRows.Err(); err != nil {
@@ -319,32 +245,6 @@ func (h *SecretsHandler) Values(c *gin.Context) {
return
}
// internal#711: provider-aware gate on the remote-pull path. A workspace
// whose resolved billing mode is NOT platform_managed (byok / subscription)
// must NOT receive the platform's scope:global LLM credentials
// (CLAUDE_CODE_OAUTH_TOKEN + the rest of the bypass-key set). Those keys
// were merged from global_secrets above; here we drop any that are still
// of global provenance (a workspace override survives, since its flag was
// cleared). Symmetric with applyPlatformManagedLLMEnv's strip on the
// provision/restart env path — both injection vectors are now gated.
//
// Default-closed: ResolveLLMBillingMode collapses any DB error / NULL /
// garbled value to platform_managed, so a transient failure leaves the
// existing (global-inheriting) behavior in place rather than stripping a
// platform_managed workspace's creds.
orgMode := strings.ToLower(strings.TrimSpace(os.Getenv("MOLECULE_LLM_BILLING_MODE")))
res, resolveErr := ResolveLLMBillingMode(ctx, workspaceID, orgMode)
if resolveErr != nil {
log.Printf("secrets.Values: resolve billing mode workspace=%s err=%v (defaulting to platform_managed)", workspaceID, resolveErr)
}
if res.ResolvedMode != LLMBillingModePlatformManaged {
for k := range globalKeys {
if isPlatformManagedDirectLLMBypassKey(k) {
delete(out, k)
}
}
}
c.JSON(http.StatusOK, out)
}
@@ -365,7 +265,7 @@ func (h *SecretsHandler) Set(c *gin.Context) {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request body"})
return
}
if rejectPlatformManagedDirectLLMBypassForWorkspace(c, workspaceID, body.Key) {
if rejectPlatformManagedDirectLLMBypass(c, body.Key) {
return
}
@@ -437,8 +337,6 @@ func (h *SecretsHandler) Delete(c *gin.Context) {
rows, err := result.RowsAffected()
if err != nil {
log.Printf("DeleteWorkspace: RowsAffected error: %v", err)
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to delete secret"})
return
}
if rows == 0 {
c.JSON(http.StatusNotFound, gin.H{"error": "secret not found"})
@@ -623,8 +521,6 @@ func (h *SecretsHandler) DeleteGlobal(c *gin.Context) {
rows, err := result.RowsAffected()
if err != nil {
log.Printf("DeleteGlobal: RowsAffected error: %v", err)
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to delete"})
return
}
if rows == 0 {
c.JSON(http.StatusNotFound, gin.H{"error": "secret not found"})
@@ -865,12 +865,6 @@ func TestSecretsValues_LegacyWorkspaceGrandfathered(t *testing.T) {
WithArgs(testWsID).
WillReturnRows(sqlmock.NewRows([]string{"key", "encrypted_value", "encryption_version"}).
AddRow("WS_KEY", []byte("ws_plainvalue"), 0))
// internal#711: Values now resolves billing mode to gate the global LLM-cred
// merge. Neither key here is a platform-managed LLM bypass key, so the mode
// is immaterial to the assertions — but the resolver query must be mocked.
mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(testWsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(LLMBillingModePlatformManaged))
w := httptest.NewRecorder()
c := secretsValuesRequest(w, "") // no auth — grandfathered
@@ -948,12 +942,6 @@ func TestSecretsValues_ValidTokenReturnsDecryptedMerge(t *testing.T) {
WillReturnRows(sqlmock.NewRows([]string{"key", "encrypted_value", "encryption_version"}).
AddRow("ONLY_WS", []byte("ws_val"), 0).
AddRow("SHARED_KEY", []byte("ws_wins"), 0))
// internal#711: billing-mode resolver query. None of these keys is a
// platform-managed LLM bypass key, so the resolved mode does not affect the
// merge assertions; platform_managed keeps the existing pass-through.
mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(testWsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(LLMBillingModePlatformManaged))
w := httptest.NewRecorder()
c := secretsValuesRequest(w, "Bearer good-token")
@@ -975,68 +963,6 @@ func TestSecretsValues_ValidTokenReturnsDecryptedMerge(t *testing.T) {
}
}
// TestSecretsValues_ByokStripsGlobalLLMCred is the internal#711 regression
// guard for the remote-pull injection vector. A non-platform (byok) workspace
// that pulls its secrets via GET /workspaces/:id/secrets/values must NOT
// receive the platform's scope:global CLAUDE_CODE_OAUTH_TOKEN — that key is
// of global_secrets provenance and is dropped by the provider-aware gate.
// Its OWN ANTHROPIC_API_KEY (a workspace_secrets row) survives, and unrelated
// non-LLM global secrets are untouched.
func TestSecretsValues_ByokStripsGlobalLLMCred(t *testing.T) {
mock := setupTestDB(t)
handler := NewSecretsHandler(nil)
mock.ExpectQuery(`SELECT COUNT\(\*\) FROM workspace_auth_tokens`).
WithArgs(testWsID).
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(1))
mock.ExpectQuery(`SELECT t\.id, t\.workspace_id.*FROM workspace_auth_tokens t.*JOIN workspaces`).
WithArgs(sqlmock.AnyArg()).
WillReturnRows(sqlmock.NewRows([]string{"id", "workspace_id"}).AddRow("tok-1", testWsID))
mock.ExpectExec(`UPDATE workspace_auth_tokens SET last_used_at`).
WithArgs("tok-1").
WillReturnResult(sqlmock.NewResult(0, 1))
// global_secrets holds the platform's scope:global OAuth token + a
// non-LLM operator global (should be untouched).
mock.ExpectQuery(`SELECT key, encrypted_value, encryption_version FROM global_secrets`).
WillReturnRows(sqlmock.NewRows([]string{"key", "encrypted_value", "encryption_version"}).
AddRow("CLAUDE_CODE_OAUTH_TOKEN", []byte("PLATFORM-GLOBAL-OAUTH"), 0).
AddRow("SENTRY_DSN", []byte("https://sentry.example/123"), 0))
// The workspace brought its OWN Anthropic API key via the Secrets tab.
mock.ExpectQuery(`SELECT key, encrypted_value, encryption_version FROM workspace_secrets WHERE workspace_id`).
WithArgs(testWsID).
WillReturnRows(sqlmock.NewRows([]string{"key", "encrypted_value", "encryption_version"}).
AddRow("ANTHROPIC_API_KEY", []byte("CUSTOMER-OWN-ANTHROPIC-KEY"), 0))
// Resolver: this workspace is byok.
mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(testWsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(LLMBillingModeBYOK))
w := httptest.NewRecorder()
c := secretsValuesRequest(w, "Bearer good-token")
handler.Values(c)
if w.Code != http.StatusOK {
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
}
var body map[string]string
_ = json.Unmarshal(w.Body.Bytes(), &body)
// 1. Platform global OAuth token stripped — the leak is closed on the pull path.
if got, ok := body["CLAUDE_CODE_OAUTH_TOKEN"]; ok {
t.Fatalf("CLAUDE_CODE_OAUTH_TOKEN = %q present — platform scope:global token must be stripped for byok pull", got)
}
// 2. The workspace's own LLM key survives.
if body["ANTHROPIC_API_KEY"] != "CUSTOMER-OWN-ANTHROPIC-KEY" {
t.Fatalf("ANTHROPIC_API_KEY = %q, want the workspace's own key preserved", body["ANTHROPIC_API_KEY"])
}
// 3. Unrelated non-LLM global secrets are untouched.
if body["SENTRY_DSN"] != "https://sentry.example/123" {
t.Fatalf("SENTRY_DSN = %q, want non-LLM globals untouched", body["SENTRY_DSN"])
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Errorf("unmet sqlmock expectations: %v", err)
}
}
func TestSecretsValues_InvalidWorkspaceID(t *testing.T) {
setupTestDB(t)
handler := NewSecretsHandler(nil)
@@ -1,180 +0,0 @@
package handlers
// template_schedules.go — read a workspace template's `schedules:`
// block and seed workspace_schedules with source='template'. Mirrors
// the org/import flow (org_import.go) so a workspace created directly
// from a workspace template (e.g. via WorkspaceHandler.Create) lands
// with the same schedule grid the org/import path would have produced.
//
// Issue #24 contract (also enforced by org_import + schedules.go):
// - INSERT new rows with source='template'
// - On (workspace_id, name) collision, only refresh template-source
// rows; runtime-added rows survive re-provisioning untouched
// - Never DELETE (additive only)
//
// The actual INSERT statement is the canonical orgImportScheduleSQL
// defined in org.go — reused here verbatim so the four guarantees
// stay in one place.
//
// Hostile-template defenses (a tenant can upload a config.yaml via
// POST /templates/import or webhook-sync a repo they control):
// - config.yaml is loaded through a 1 MiB LimitReader so a YAML
// anchor-bomb / billion-laughs cannot pre-explode memory before
// unmarshal returns.
// - len(schedules), per-schedule cron length, and resolved prompt
// body length are all bounded; over-sized entries are skipped
// rather than committed.
// - Per-row insert errors and ctx cancellation surface to the
// caller via the returned counts so partial-seed states are
// observable (workspace.go Create logs the (seeded, skipped)
// pair when skipped > 0).
import (
"context"
"errors"
"fmt"
"io"
"log"
"os"
"path/filepath"
"time"
"gopkg.in/yaml.v3"
"git.moleculesai.app/molecule-ai/molecule-core/workspace-server/internal/db"
"git.moleculesai.app/molecule-ai/molecule-core/workspace-server/internal/scheduler"
)
// Bounds protecting the seeder against hostile or buggy templates.
// All chosen with generous headroom relative to legitimate use
// (reno-stars org template — the largest production schedule grid —
// runs ~10 entries per workspace, each prompt body well under 1 KiB).
const (
maxTemplateConfigYAMLBytes int64 = 1 << 20 // 1 MiB — hard cap on config.yaml size
maxTemplateSchedules = 100 // 10x current largest grid
maxScheduleCronExprLen = 128 // cron-spec syntax is short by construction
maxSchedulePromptBytes = 16 << 10 // 16 KiB after prompt_file resolution
)
// templateConfigSchedules is the minimal shape parsed from a workspace
// template's config.yaml. Only the `schedules:` block is modelled;
// the rest of the file (providers, runtime_config, …) is opaque to
// this loader and continues to flow through the existing pass-through
// in workspace_provision.go.
type templateConfigSchedules struct {
Schedules []OrgSchedule `yaml:"schedules"`
}
// parseTemplateSchedules reads `<templatePath>/config.yaml` and
// returns its `schedules:` block (nil + nil error when the file is
// absent or the block is empty).
//
// The file is read through a 1 MiB LimitReader so a billion-laughs
// or anchor-explosion YAML cannot pre-explode memory before
// Unmarshal returns. Returns an error only when a present
// config.yaml fails to read or parse — callers should treat that as
// a template-author bug rather than a runtime fault. The Create
// handler logs the error and continues so a broken schedules block
// can never block workspace provisioning.
func parseTemplateSchedules(templatePath string) ([]OrgSchedule, error) {
if templatePath == "" {
return nil, nil
}
f, err := os.Open(filepath.Join(templatePath, "config.yaml"))
if err != nil {
if errors.Is(err, os.ErrNotExist) {
return nil, nil
}
return nil, fmt.Errorf("open template config.yaml: %w", err)
}
defer f.Close()
// Read maxTemplateConfigYAMLBytes+1 — if we filled the buffer the
// underlying file exceeded the cap and we refuse to unmarshal.
data, err := io.ReadAll(io.LimitReader(f, maxTemplateConfigYAMLBytes+1))
if err != nil {
return nil, fmt.Errorf("read template config.yaml: %w", err)
}
if int64(len(data)) > maxTemplateConfigYAMLBytes {
return nil, fmt.Errorf("template config.yaml exceeds %d-byte cap", maxTemplateConfigYAMLBytes)
}
var cfg templateConfigSchedules
if err := yaml.Unmarshal(data, &cfg); err != nil {
return nil, fmt.Errorf("parse template config.yaml schedules: %w", err)
}
if len(cfg.Schedules) > maxTemplateSchedules {
return nil, fmt.Errorf("template declares %d schedules; cap is %d", len(cfg.Schedules), maxTemplateSchedules)
}
return cfg.Schedules, nil
}
// seedTemplateSchedules INSERTs (or refreshes) each schedule into
// workspace_schedules with source='template'. Returns (seeded,
// skipped) counts so the caller can observe partial-seed states.
//
// Prompt body resolution mirrors org_import.go: inline `prompt:` wins,
// else `prompt_file:` is resolved relative to templatePath via
// resolvePromptRef. Per-schedule failures (bad cron, missing prompt
// file, DB error, oversize input) are logged with the schedule name
// quoted via %q (CRLF-safe) and skipped so one bad row never breaks
// the rest of the grid. A cancelled ctx breaks the loop early.
//
// Timezone defaults to "UTC" when unset. Env-var expansion in the
// timezone field is intentionally not performed — that mirrors the
// org/import behavior; template authors should pick a literal IANA
// zone (or rely on UTC + operator overrides per-tenant).
func seedTemplateSchedules(ctx context.Context, workspaceID, templatePath string, schedules []OrgSchedule) (seeded, skipped int) {
for _, sched := range schedules {
// Honour caller cancellation — protects against long seed
// loops on a request whose client already gave up.
if err := ctx.Err(); err != nil {
log.Printf("Template schedule seed: ctx cancelled after %d/%d on %s: %v", seeded, len(schedules), workspaceID, err)
skipped += len(schedules) - seeded - skipped
return
}
if len(sched.CronExpr) > maxScheduleCronExprLen {
log.Printf("Template schedule seed: cron_expr too long (%d > %d) for %q on %s — skipping", len(sched.CronExpr), maxScheduleCronExprLen, sched.Name, workspaceID)
skipped++
continue
}
tz := sched.Timezone
if tz == "" {
tz = "UTC"
}
enabled := true
if sched.Enabled != nil {
enabled = *sched.Enabled
}
prompt, promptErr := resolvePromptRef(sched.Prompt, sched.PromptFile, templatePath, "")
if promptErr != nil {
log.Printf("Template schedule seed: failed to resolve prompt for %q on %s: %v — skipping", sched.Name, workspaceID, promptErr)
skipped++
continue
}
if prompt == "" {
log.Printf("Template schedule seed: schedule %q on %s has empty prompt — skipping", sched.Name, workspaceID)
skipped++
continue
}
if len(prompt) > maxSchedulePromptBytes {
log.Printf("Template schedule seed: prompt too long (%d > %d bytes) for %q on %s — skipping", len(prompt), maxSchedulePromptBytes, sched.Name, workspaceID)
skipped++
continue
}
nextRun, nextRunErr := scheduler.ComputeNextRun(sched.CronExpr, tz, time.Now())
if nextRunErr != nil {
log.Printf("Template schedule seed: invalid cron for %q on %s: %v — skipping", sched.Name, workspaceID, nextRunErr)
skipped++
continue
}
if _, err := db.DB.ExecContext(ctx, orgImportScheduleSQL,
workspaceID, sched.Name, sched.CronExpr, tz, prompt, enabled, nextRun); err != nil {
log.Printf("Template schedule seed: failed to upsert %q on %s: %v", sched.Name, workspaceID, err)
skipped++
continue
}
seeded++
log.Printf("Template schedule seed: %q (%s, %d chars) upserted on %s (source=template)", sched.Name, sched.CronExpr, len(prompt), workspaceID)
}
return
}
@@ -1,141 +0,0 @@
package handlers
// template_schedules_test.go — unit tests for parseTemplateSchedules.
//
// seedTemplateSchedules' DB INSERT path is already covered indirectly
// by TestImport_OrgScheduleSQLShape (schedules_test.go) since both
// code paths share the canonical orgImportScheduleSQL constant; the
// loop logic (default tz, default enabled, prompt resolution, cron
// validation) is exercised at the parser level here and at the
// orgImportScheduleSQL level there.
import (
"path/filepath"
"testing"
)
func TestParseTemplateSchedules_AbsentFile(t *testing.T) {
dir := t.TempDir()
// No config.yaml in dir.
got, err := parseTemplateSchedules(dir)
if err != nil {
t.Fatalf("expected nil error for absent config.yaml, got %v", err)
}
if got != nil {
t.Fatalf("expected nil slice, got %#v", got)
}
}
func TestParseTemplateSchedules_EmptyTemplatePath(t *testing.T) {
got, err := parseTemplateSchedules("")
if err != nil {
t.Fatalf("expected nil error for empty path, got %v", err)
}
if got != nil {
t.Fatalf("expected nil slice for empty path, got %#v", got)
}
}
func TestParseTemplateSchedules_NoSchedulesBlock(t *testing.T) {
dir := t.TempDir()
mustWriteFile(t, filepath.Join(dir, "config.yaml"), `
name: Some Template
runtime: claude-code
model: foo/bar
`)
got, err := parseTemplateSchedules(dir)
if err != nil {
t.Fatalf("expected nil error when schedules: absent, got %v", err)
}
if len(got) != 0 {
t.Fatalf("expected zero schedules, got %d", len(got))
}
}
func TestParseTemplateSchedules_HappyPath(t *testing.T) {
dir := t.TempDir()
mustWriteFile(t, filepath.Join(dir, "config.yaml"), `
name: SEO Agent
schedules:
- name: Continuous tick
cron_expr: "*/30 * * * *"
timezone: America/Vancouver
prompt: |
Run one SEO tick.
- name: Monday GSC
cron_expr: "0 8 * * 1"
timezone: America/Vancouver
prompt: /seo google
enabled: true
- name: Disabled placeholder
cron_expr: "0 0 1 1 *"
prompt: noop
enabled: false
`)
got, err := parseTemplateSchedules(dir)
if err != nil {
t.Fatalf("parse: %v", err)
}
if len(got) != 3 {
t.Fatalf("expected 3 schedules, got %d", len(got))
}
if got[0].Name != "Continuous tick" || got[0].CronExpr != "*/30 * * * *" {
t.Errorf("schedule[0] mismatch: %+v", got[0])
}
if got[1].Timezone != "America/Vancouver" {
t.Errorf("schedule[1].Timezone = %q, want America/Vancouver", got[1].Timezone)
}
// Enabled is *bool: nil means "default true" at seed time, false is
// explicit opt-out and must survive the YAML round-trip.
if got[2].Enabled == nil {
t.Errorf("schedule[2].Enabled = nil, want *false")
} else if *got[2].Enabled {
t.Errorf("schedule[2].Enabled = true, want false")
}
}
func TestParseTemplateSchedules_MalformedYAML(t *testing.T) {
dir := t.TempDir()
mustWriteFile(t, filepath.Join(dir, "config.yaml"), `
name: Broken
schedules:
- this is: [not, a, valid
`)
_, err := parseTemplateSchedules(dir)
if err == nil {
t.Fatal("expected parse error on malformed YAML, got nil")
}
}
// TestParseTemplateSchedules_RejectsOversizeFile gates against the
// billion-laughs / anchor-bomb DoS class: a hostile config.yaml over
// the 1 MiB cap must be refused before yaml.Unmarshal runs.
func TestParseTemplateSchedules_RejectsOversizeFile(t *testing.T) {
dir := t.TempDir()
// One byte over the cap — fastest path to the gate.
pad := make([]byte, maxTemplateConfigYAMLBytes+1)
for i := range pad {
pad[i] = '#'
}
mustWriteFile(t, filepath.Join(dir, "config.yaml"), string(pad))
if _, err := parseTemplateSchedules(dir); err == nil {
t.Fatal("expected oversize-file error, got nil")
}
}
// TestParseTemplateSchedules_RejectsTooManySchedules gates against a
// hostile config.yaml that flips one row into a 10k-row insert storm.
func TestParseTemplateSchedules_RejectsTooManySchedules(t *testing.T) {
dir := t.TempDir()
var b []byte
b = append(b, []byte("schedules:\n")...)
// maxTemplateSchedules+1 minimal entries — they don't have to be
// valid as schedules because the gate trips before resolution.
for i := 0; i <= maxTemplateSchedules; i++ {
b = append(b, []byte(" - name: s\n cron_expr: \"* * * * *\"\n prompt: x\n")...)
}
mustWriteFile(t, filepath.Join(dir, "config.yaml"), string(b))
if _, err := parseTemplateSchedules(dir); err == nil {
t.Fatal("expected schedule-count error, got nil")
}
}
@@ -4,7 +4,6 @@ import (
"bytes"
"context"
"fmt"
"log"
"net/http"
"os"
"os/exec"
@@ -120,11 +119,9 @@ func (h *TerminalHandler) HandleDiagnose(c *gin.Context) {
}
var instanceID string
if err := db.DB.QueryRowContext(ctx,
_ = db.DB.QueryRowContext(ctx,
`SELECT COALESCE(instance_id, '') FROM workspaces WHERE id = $1`,
workspaceID).Scan(&instanceID); err != nil {
log.Printf("terminal diagnose: instance_id query failed for workspace %s: %v", workspaceID, err)
}
workspaceID).Scan(&instanceID)
var res diagnoseResult
if instanceID != "" {
+1 -6
View File
@@ -153,12 +153,7 @@ func (h *TokenHandler) Revoke(c *gin.Context) {
return
}
rows, err := result.RowsAffected()
if err != nil {
log.Printf("tokens: revoke RowsAffected error token=%s workspace=%s: %v", tokenID, workspaceID, err)
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to revoke token"})
return
}
rows, _ := result.RowsAffected()
if rows == 0 {
c.JSON(http.StatusNotFound, gin.H{"error": "token not found or already revoked"})
return
+6 -14
View File
@@ -394,13 +394,9 @@ func (h *WebhookHandler) handleCronTriggerEvent(c *gin.Context, eventType string
log.Printf("Webhook: cron trigger (issues/opened) DB error: %v", err)
return true, fmt.Errorf("failed to trigger schedules: %w", err)
}
affected, err := result.RowsAffected()
if err != nil {
log.Printf("Webhook: issues/opened RowsAffected error: %v", err)
} else {
log.Printf("Webhook: issues/opened in %s #%d by %s — triggered %d pick-up-work schedule(s)",
payload.Repository.FullName, payload.Issue.Number, payload.Sender.Login, affected)
}
affected, _ := result.RowsAffected()
log.Printf("Webhook: issues/opened in %s #%d by %s — triggered %d pick-up-work schedule(s)",
payload.Repository.FullName, payload.Issue.Number, payload.Sender.Login, affected)
c.JSON(http.StatusOK, gin.H{
"status": "triggered",
@@ -433,13 +429,9 @@ func (h *WebhookHandler) handleCronTriggerEvent(c *gin.Context, eventType string
log.Printf("Webhook: cron trigger (pull_request_review/submitted) DB error: %v", err)
return true, fmt.Errorf("failed to trigger schedules: %w", err)
}
affected, err := result.RowsAffected()
if err != nil {
log.Printf("Webhook: pull_request_review/submitted RowsAffected error: %v", err)
} else {
log.Printf("Webhook: pull_request_review/submitted in %s PR #%d by %s (state=%s) — triggered %d review schedule(s)",
payload.Repository.FullName, payload.PullRequest.Number, payload.Sender.Login, payload.Review.State, affected)
}
affected, _ := result.RowsAffected()
log.Printf("Webhook: pull_request_review/submitted in %s PR #%d by %s (state=%s) — triggered %d review schedule(s)",
payload.Repository.FullName, payload.PullRequest.Number, payload.Sender.Login, payload.Review.State, affected)
c.JSON(http.StatusOK, gin.H{
"status": "triggered",
@@ -568,7 +568,7 @@ func (h *WorkspaceHandler) Create(c *gin.Context) {
// nil/empty map is a no-op. Any failure rolls back the workspace insert
// so we never have a workspace row without its intended secrets.
for k, v := range payload.Secrets {
if rejectPlatformManagedDirectLLMBypassForWorkspace(c, id, k) {
if rejectPlatformManagedDirectLLMBypass(c, k) {
tx.Rollback() //nolint:errcheck
return
}
@@ -677,9 +677,7 @@ func (h *WorkspaceHandler) Create(c *gin.Context) {
// Preserve BYO-compute runtime label (kimi, kimi-cli, external) —
// don't coerce to generic "external" so the canvas can show the
// correct runtime name in the node card.
if _, err := db.DB.ExecContext(ctx, `UPDATE workspaces SET url = $1, status = $2, runtime = $3, updated_at = now() WHERE id = $4`, payload.URL, models.StatusOnline, normalizeExternalRuntime(payload.Runtime), id); err != nil {
log.Printf("External workspace: failed to update URL/status for %s: %v", id, err)
}
db.DB.ExecContext(ctx, `UPDATE workspaces SET url = $1, status = $2, runtime = $3, updated_at = now() WHERE id = $4`, payload.URL, models.StatusOnline, normalizeExternalRuntime(payload.Runtime), id)
if err := db.CacheURL(ctx, id, payload.URL); err != nil {
log.Printf("External workspace: failed to cache URL for %s: %v", id, err)
}
@@ -692,9 +690,7 @@ func (h *WorkspaceHandler) Create(c *gin.Context) {
// from the external agent (with this token + its URL)
// flips the row to online.
// Preserve BYO-compute runtime label (kimi, kimi-cli, external).
if _, err := db.DB.ExecContext(ctx, `UPDATE workspaces SET status = $1, runtime = $2, updated_at = now() WHERE id = $3`, models.StatusAwaitingAgent, normalizeExternalRuntime(payload.Runtime), id); err != nil {
log.Printf("External workspace: failed to update status for %s: %v", id, err)
}
db.DB.ExecContext(ctx, `UPDATE workspaces SET status = $1, runtime = $2, updated_at = now() WHERE id = $3`, models.StatusAwaitingAgent, normalizeExternalRuntime(payload.Runtime), id)
tok, tokErr := wsauth.IssueToken(ctx, db.DB, id)
if tokErr != nil {
log.Printf("External workspace %s: token issuance failed: %v", id, tokErr)
@@ -769,8 +765,7 @@ func (h *WorkspaceHandler) Create(c *gin.Context) {
// runtime/model/tier as JSON — the Config tab needs that to render
// even on failed workspaces, so Create owns this Create-only side
// effect rather than coupling Auto to a UI concern.
provisionOK := h.provisionWorkspaceAuto(id, templatePath, configFiles, payload)
if !provisionOK {
if !h.provisionWorkspaceAuto(id, templatePath, configFiles, payload) {
cfgJSON := fmt.Sprintf(`{"name":%q,"runtime":%q,"tier":%d,"template":%q}`,
payload.Name, payload.Runtime, payload.Tier, payload.Template)
if _, err := db.DB.ExecContext(ctx, `
@@ -781,32 +776,6 @@ func (h *WorkspaceHandler) Create(c *gin.Context) {
}
}
// Seed schedules declared in the workspace template's config.yaml
// AFTER provisionWorkspaceAuto succeeds so the scheduler never
// fires cron rows against a workspace whose backend never wired
// (review feedback PR #1929#1). Async EC2 provisioning may still
// fail downstream; scheduler.go is expected to handle non-online
// status as a no-op tick. Idempotent across re-creates via
// orgImportScheduleSQL's ON CONFLICT clause; runtime-added rows
// are preserved (Issue #24 contract). Restart does not re-seed
// (so user-deleted template rows stay deleted).
//
// Non-fatal: a broken schedules: block must never block workspace
// provisioning — the workspace row is already live and the grid
// is recoverable via POST /workspaces/{id}/schedules.
if provisionOK && templatePath != "" {
if templateScheds, parseErr := parseTemplateSchedules(templatePath); parseErr != nil {
log.Printf("Create %s: parsing template schedules: %v (continuing)", id, parseErr)
} else if len(templateScheds) > 0 {
seeded, skipped := seedTemplateSchedules(ctx, id, templatePath, templateScheds)
if skipped > 0 {
log.Printf("Create %s: template schedule partial-seed: seeded=%d skipped=%d total=%d", id, seeded, skipped, len(templateScheds))
} else {
log.Printf("Create %s: seeded %d/%d template schedules", id, seeded, len(templateScheds))
}
}
}
c.JSON(http.StatusCreated, gin.H{
"id": id,
"status": "provisioning",
@@ -1023,11 +992,9 @@ func (h *WorkspaceHandler) Get(c *gin.Context) {
// the client would otherwise see — the actionable signal
// is the 410 + hint, not the timestamp.
var removedAt time.Time
if err := db.DB.QueryRowContext(c.Request.Context(),
_ = db.DB.QueryRowContext(c.Request.Context(),
`SELECT updated_at FROM workspaces WHERE id = $1`, id,
).Scan(&removedAt); err != nil {
log.Printf("workspace GET: removed_at query failed for %s: %v", id, err)
}
).Scan(&removedAt)
body := gin.H{
"error": "workspace removed",
"id": id,
@@ -450,12 +450,8 @@ func (h *WorkspaceHandler) Delete(c *gin.Context) {
}
}
// Null out parent_id / forwarded_to references
if _, err := db.DB.ExecContext(ctx, "UPDATE workspaces SET parent_id = NULL WHERE parent_id = ANY($1::uuid[])", purgeIDs); err != nil {
log.Printf("Purge parent_id null error for %v: %v", allIDs, err)
}
if _, err := db.DB.ExecContext(ctx, "UPDATE workspaces SET forwarded_to = NULL WHERE forwarded_to = ANY($1::uuid[])", purgeIDs); err != nil {
log.Printf("Purge forwarded_to null error for %v: %v", allIDs, err)
}
db.DB.ExecContext(ctx, "UPDATE workspaces SET parent_id = NULL WHERE parent_id = ANY($1::uuid[])", purgeIDs)
db.DB.ExecContext(ctx, "UPDATE workspaces SET forwarded_to = NULL WHERE forwarded_to = ANY($1::uuid[])", purgeIDs)
// Hard delete the workspace row
if _, err := db.DB.ExecContext(ctx, "DELETE FROM workspaces WHERE id = ANY($1::uuid[])", purgeIDs); err != nil {
log.Printf("Purge workspace row error for %v: %v", allIDs, err)
@@ -574,12 +570,7 @@ func (h *WorkspaceHandler) CascadeDelete(ctx context.Context, id string) ([]stri
var stopErrs []error
stopAndRemove := func(wsID string) {
// Delete-path stop uses bounded retry (matches the restart path) and
// records a durable structure_events row on exhaustion so a leaked /
// pending EC2 is queryable and handed off to the CP-orphan-sweeper —
// rather than the bare one-shot StopWorkspaceAuto that produced the
// silent-leak class (task #15 / workspace-ec2-leak).
if err := h.stopWorkspaceForDelete(cleanupCtx, wsID); err != nil {
if err := h.StopWorkspaceAuto(cleanupCtx, wsID); err != nil {
log.Printf("CascadeDelete %s stop failed: %v — leaving cleanup for orphan sweeper", wsID, err)
stopErrs = append(stopErrs, fmt.Errorf("stop %s: %w", wsID, err))
return
@@ -165,43 +165,3 @@ func TestValidateWorkspaceFields_YAMLCharsAllowedInEmptyName(t *testing.T) {
t.Errorf("empty name with valid role: expected nil, got %v", err)
}
}
// ─── validateWorkspaceID ───────────────────────────────────────────────────────
func TestValidateWorkspaceID_ValidUUIDv4(t *testing.T) {
if err := validateWorkspaceID("550e8400-e29b-41d4-a716-446655440000"); err != nil {
t.Errorf("valid v4 UUID: expected nil, got %v", err)
}
}
func TestValidateWorkspaceID_ValidUUIDv1(t *testing.T) {
// UUIDv1 format is also accepted by uuid.Parse.
if err := validateWorkspaceID("6ba7b810-9dad-11d1-80b4-00c04fd430c8"); err != nil {
t.Errorf("valid v1 UUID: expected nil, got %v", err)
}
}
func TestValidateWorkspaceID_EmptyString(t *testing.T) {
if err := validateWorkspaceID(""); err == nil {
t.Error("empty string: expected error, got nil")
}
}
func TestValidateWorkspaceID_NotAUuid(t *testing.T) {
if err := validateWorkspaceID("not-a-uuid"); err == nil {
t.Error("not-a-uuid: expected error, got nil")
}
}
func TestValidateWorkspaceID_WrongLength(t *testing.T) {
if err := validateWorkspaceID("550e8400-e29b-41d4-a716"); err == nil {
t.Error("short UUID: expected error, got nil")
}
}
func TestValidateWorkspaceID_InvalidCharacters(t *testing.T) {
// 'g' is not a valid hex character.
if err := validateWorkspaceID("550e8400-e29b-41d4-a716-44665544000g"); err == nil {
t.Error("invalid hex char: expected error, got nil")
}
}
@@ -1,102 +0,0 @@
package handlers
// workspace_delete_stop_retry_test.go — pins the contract of the
// delete-path EC2 stop retry (task #15 / workspace-ec2-leak).
//
// Background (Phase 1 evidence): the DELETE path's StopWorkspaceAuto →
// cpProv.Stop had NO retry, while the restart path used cpStopWithRetry
// (bounded exponential backoff). A transient CP/AWS hiccup on delete left
// the workspace row at status='removed' with instance_id still populated,
// returned a 500, and relied entirely on the 60s CP-orphan-sweeper to
// re-drive the terminate. For a cascade *descendant* whose own row is
// already 'removed', the inline retry-via-client-replay is defeated by
// CascadeDelete's `status != 'removed'` CTE filter — so the only inline
// recovery is this bounded retry.
//
// Contract of stopWorkspaceForDelete:
// - CP path: bounded retry (cpStopRetryAttempts, exp backoff) on
// cpProv.Stop; returns nil on eventual success.
// - On retry exhaustion: returns the terminal error AND emits a
// `workspace.delete.terminate_retry_exhausted` structure_events row so
// the leak decision is queryable (structured-logging gate), not just a
// log.Printf. The row is the durable pending-terminate signal: the row
// stays status='removed' with instance_id populated, which is exactly
// what the CP-orphan-sweeper (registry/cp_orphan_sweeper.go) re-drives.
// - Docker path: single Stop, no retry (local daemon failure won't heal
// on retry — matches RestartWorkspaceAuto's Docker rationale).
// - No backend wired: nil (nothing to stop).
import (
"context"
"errors"
"strings"
"testing"
"github.com/DATA-DOG/go-sqlmock"
)
func TestStopWorkspaceForDelete_CPRetriesTransientThenSucceeds(t *testing.T) {
shrinkRetryBackoff(t)
buf := captureLog(t)
// 2 transient failures then success — within the 3-attempt budget.
stub := &scriptedCPStop{errs: []error{
errors.New("cp 503 attempt 1"),
errors.New("cp 503 attempt 2"),
}}
h := &WorkspaceHandler{cpProv: stub}
err := h.stopWorkspaceForDelete(context.Background(), "ws-del-1")
if err != nil {
t.Fatalf("expected nil error on eventual success, got %v", err)
}
if stub.calls != 3 {
t.Errorf("expected 3 Stop calls (2 fails + 1 success), got %d", stub.calls)
}
if strings.Contains(buf.String(), "terminate_retry_exhausted") {
t.Errorf("eventual success must NOT log retry-exhausted; got %q", buf.String())
}
}
func TestStopWorkspaceForDelete_CPExhaustsEmitsDurableEventAndReturnsError(t *testing.T) {
shrinkRetryBackoff(t)
mock := setupTestDB(t)
buf := captureLog(t)
stub := &scriptedCPStop{errs: []error{
errors.New("cp 502 attempt 1"),
errors.New("cp 502 attempt 2"),
errors.New("cp 502 final"),
}}
h := &WorkspaceHandler{cpProv: stub}
// On exhaustion the helper persists a durable pending-terminate row so
// the leak decision is queryable. structure_events is the audit-of-record.
mock.ExpectExec("INSERT INTO structure_events").
WillReturnResult(sqlmock.NewResult(0, 1))
err := h.stopWorkspaceForDelete(context.Background(), "ws-doomed")
if err == nil {
t.Fatal("expected terminal error on retry exhaustion, got nil")
}
if stub.calls != cpStopRetryAttempts {
t.Errorf("expected %d Stop calls when all fail, got %d", cpStopRetryAttempts, stub.calls)
}
if !strings.Contains(err.Error(), "cp 502 final") {
t.Errorf("returned error should wrap the LAST attempt's error, got %v", err)
}
if e := mock.ExpectationsWereMet(); e != nil {
t.Fatalf("expected structure_events INSERT on exhaustion: %v", e)
}
// The LEAK-SUSPECT line stays the operator-facing prose bridge to the
// orphan reconciler; assert it carries the delete source so triage can
// distinguish delete-leaks from restart-leaks.
if !strings.Contains(buf.String(), "LEAK-SUSPECT") {
t.Errorf("expected LEAK-SUSPECT log on exhaustion, got %q", buf.String())
}
}
func TestStopWorkspaceForDelete_NoBackendIsNoOp(t *testing.T) {
h := &WorkspaceHandler{} // cpProv nil, provisioner nil
if err := h.stopWorkspaceForDelete(context.Background(), "ws-x"); err != nil {
t.Errorf("expected nil no-op with no backend, got %v", err)
}
}
@@ -31,11 +31,9 @@ package handlers
import (
"context"
"encoding/json"
"log"
"time"
"git.moleculesai.app/molecule-ai/molecule-core/workspace-server/internal/db"
"git.moleculesai.app/molecule-ai/molecule-core/workspace-server/internal/models"
"git.moleculesai.app/molecule-ai/molecule-core/workspace-server/internal/provlog"
)
@@ -209,86 +207,6 @@ func (h *WorkspaceHandler) StopWorkspaceAuto(ctx context.Context, workspaceID st
return nil
}
// stopWorkspaceForDelete is the DELETE-path stop dispatcher. It differs
// from StopWorkspaceAuto in exactly one way: the CP (EC2) path gets the
// same bounded retry the restart path uses (cpStopWithRetryErr), and on
// retry exhaustion it persists a durable `workspace.delete.terminate_retry_exhausted`
// event to structure_events (the structured-logging gate) so the leak
// decision is queryable, not just stdout prose.
//
// Why retry here (task #15 / workspace-ec2-leak): the bare cpProv.Stop on
// delete left a transient CP/AWS hiccup as an immediate 500 with no inline
// recovery. For a cascade *descendant* the "client retries → replays
// terminate" recovery is defeated by CascadeDelete's `status != 'removed'`
// CTE filter (the descendant's row is already 'removed', so a retry walks
// zero descendant rows). Bounded retry absorbs the transient class inline;
// the durable event + the row staying status='removed'+instance_id is the
// hand-off to the 60s CP-orphan-sweeper (registry/cp_orphan_sweeper.go) for
// the (rarer) sustained-outage case.
//
// We deliberately do NOT clear status='removed' on exhaustion — the
// CP-orphan-sweeper's recovery query keys on exactly that state, so
// reverting it would break the existing backstop. The error is still
// returned so the HTTP Delete handler surfaces the retryable 500.
//
// Docker path: single Stop, no retry — a local daemon that fails to stop a
// container won't heal on retry (matches RestartWorkspaceAuto's Docker
// rationale); the orphan-container sweeper (registry/orphan_sweeper.go) is
// the Docker-side backstop.
func (h *WorkspaceHandler) stopWorkspaceForDelete(ctx context.Context, workspaceID string) error {
if h.cpProv != nil {
if err := h.cpStopWithRetryErr(ctx, workspaceID, "Delete"); err != nil {
h.emitDeleteTerminateRetryExhausted(ctx, workspaceID, err)
return err
}
return nil
}
if h.provisioner != nil {
return h.provisioner.Stop(ctx, workspaceID)
}
return nil
}
// emitDeleteTerminateRetryExhausted persists a durable record that the
// delete-path EC2 terminate could not be completed inline after the full
// retry budget. Per the §Persistent structured logging gate: a
// state-mutating decision (we are leaving a known-leaked-or-pending EC2 for
// the orphan sweeper) must land in structure_events, not just log.Printf.
//
// Event-type taxonomy (append-only; never rename):
//
// workspace.delete.terminate_retry_exhausted — delete-path cpProv.Stop
// exhausted its retry budget; row stays status='removed' with
// instance_id populated for the CP-orphan-sweeper to re-drive.
//
// Telemetry never blocks the request path: marshal / INSERT failures are
// logged and swallowed.
func (h *WorkspaceHandler) emitDeleteTerminateRetryExhausted(ctx context.Context, workspaceID string, cause error) {
payload := map[string]any{
"workspace_id": workspaceID,
"attempts": cpStopRetryAttempts,
"last_error": cause.Error(),
// recovery_path documents WHO is expected to finish the terminate,
// so a reader of the audit row doesn't have to grep the code to
// know the EC2 isn't simply abandoned.
"recovery_path": "cp_orphan_sweeper",
}
payloadJSON, err := json.Marshal(payload)
if err != nil {
log.Printf("emitDeleteTerminateRetryExhausted: marshal payload failed for %s: %v", workspaceID, err)
return
}
if db.DB == nil {
return
}
if _, err := db.DB.ExecContext(ctx, `
INSERT INTO structure_events (event_type, workspace_id, payload, created_at)
VALUES ($1, $2, $3, now())
`, "workspace.delete.terminate_retry_exhausted", workspaceID, payloadJSON); err != nil {
log.Printf("emitDeleteTerminateRetryExhausted: insert failed for %s: %v", workspaceID, err)
}
}
// RestartWorkspaceAuto stops the running workload (with retry semantics
// tuned for the restart hot path) then starts provisioning again, in a
// detached goroutine. Returns true when a backend was kicked off, false
@@ -75,21 +75,3 @@ func formatMissingEnvError(missing []string) string {
strings.Join(missing, ", "),
)
}
// formatMissingBYOKCredentialError builds the user-facing message for a
// provision failure caused by a non-platform (byok/subscription) workspace
// that has no usable LLM credential of its own (internal#711). The platform's
// scope:global LLM credentials are NOT a valid fallback for a non-platform
// workspace — resolving to them would bill the platform's Anthropic credits —
// so the provision fails closed here rather than starting the workspace on
// stripped/absent creds. Rendered verbatim in the canvas Events tab.
func formatMissingBYOKCredentialError(mode string) string {
return fmt.Sprintf(
"this workspace's LLM billing mode is %q (not platform-managed) but it has no LLM credential of its own. "+
"Add a workspace-scoped credential (e.g. CLAUDE_CODE_OAUTH_TOKEN or your provider's API key) under "+
"Config → Secrets, or switch the workspace to platform-managed billing via "+
"/admin/workspaces/:id/llm-billing-mode, then retry. The platform's shared LLM credentials are not "+
"used for non-platform workspaces.",
mode,
)
}
@@ -922,124 +922,21 @@ func applyRuntimeModelEnv(envVars map[string]string, runtime, model string) {
}
// applyPlatformManagedLLMEnv wires the control-plane LLM proxy into a
// workspace only when the RESOLVED billing mode for this workspace is
// platform_managed. "Resolved" means: the workspace-level override (if any)
// wins over the org default (delivered via tenant_config in MOLECULE_LLM_BILLING_MODE).
//
// Pre-internal#691 this gate read the org-level env var directly, which made
// it impossible to mix billing modes across workspaces in the same org. The
// resolver (ResolveLLMBillingMode) is the single source of truth now; the
// architectural test asserts no remaining code path gates on os.Getenv
// ("MOLECULE_LLM_BILLING_MODE") for strip-decision purposes — that env value
// is still read INTO the resolver as the org-default input, but it is never
// the final decision.
//
// Default-closed: any resolver error / NULL JOIN / garbled enum value
// collapses to platform_managed (see llm_billing_mode.go for the contract).
// This preserves the existing implicit default exactly while making the
// per-workspace opt-out path safe.
//
// The resolved mode is exported into the workspace container as
// MOLECULE_LLM_BILLING_MODE_RESOLVED so an in-container debug check can
// answer "what mode is this workspace running under" without DB queries
// (RFC Observability hot-spot).
//
// internal#711 — PROVIDER-AWARE GLOBAL-LLM-CRED GATE. The platform's
// LLM credentials (CLAUDE_CODE_OAUTH_TOKEN + the rest of the
// platformManagedDirectLLMBypassKeys set) live in `global_secrets` and
// are merged into EVERY workspace's env by loadWorkspaceSecrets — that
// merge is provenance-blind. Pre-fix, the non-platform (byok/disabled)
// early-return left envVars untouched, so a BYOK / subscription
// workspace that brought NO LLM credential of its own still inherited
// the platform's scope:global CLAUDE_CODE_OAUTH_TOKEN and ran Opus on
// the platform's (Molecule's) Anthropic credits (Reno Stars SEO +
// Marketing agents, confirmed live 2026-05-27).
//
// The gate: on the non-platform path we strip every platform-managed
// LLM key whose PROVENANCE is `global_secrets` (the globalKeys set).
// A workspace's OWN LLM credential — set via the canvas Secrets tab,
// i.e. a `workspace_secrets` row — has had its global provenance flag
// dropped by loadWorkspaceSecrets, so it is NOT in globalKeys and
// survives. Net effect: platform global LLM creds reach a workspace
// ONLY when its resolved mode is platform_managed; a non-platform
// workspace resolves to its own (workspace-scoped) credential or none.
//
// The boolean return reports whether, after the gate, the workspace
// still has at least one usable LLM credential. The caller
// (prepareProvisionContext) uses it to FAIL CLOSED — a non-platform
// workspace with no usable LLM credential is aborted with a clear
// MISSING_BYOK_CREDENTIAL error at provision time rather than being
// started on (now-stripped) platform creds.
// platformLLMEnvResult is the structured outcome of applyPlatformManagedLLMEnv.
// ResolvedMode is the per-workspace billing/provider mode the resolver
// landed on. HasUsableLLMCred reports whether — AFTER the provider-aware
// global-cred gate — the workspace still has at least one platform-managed
// LLM credential key in its env (its own, workspace-scoped one). Only the
// non-platform path consults HasUsableLLMCred for the fail-closed decision;
// the platform_managed path always returns true (it forces the CP proxy
// usage token, which IS the usable credential).
type platformLLMEnvResult struct {
ResolvedMode string
HasUsableLLMCred bool
}
func applyPlatformManagedLLMEnv(ctx context.Context, envVars map[string]string, globalKeys map[string]struct{}, workspaceID, runtime, model string) platformLLMEnvResult {
orgMode := strings.ToLower(strings.TrimSpace(os.Getenv("MOLECULE_LLM_BILLING_MODE")))
res, resolveErr := ResolveLLMBillingMode(ctx, workspaceID, orgMode)
if resolveErr != nil {
// resolveErr != nil ⇒ resolver hit a DB error AND already defaulted
// res.ResolvedMode to platform_managed. Log + proceed; the safe default
// is already in place, no early return needed.
log.Printf("workspace_provision: resolve billing mode workspace=%s err=%v (defaulting to platform_managed)", workspaceID, resolveErr)
}
log.Printf("workspace_provision: billing mode workspace=%s resolved=%s source=%s org_default=%s", workspaceID, res.ResolvedMode, res.Source, res.OrgDefault)
// internal#703: MOLECULE_LLM_BILLING_MODE in the container must reflect the
// RESOLVED per-workspace mode, not a hardcoded literal. Pre-fix this var was
// only emitted (hardcoded "platform_managed") on the strip path below, so a
// byok/disabled container never carried a truthful billing-mode value — only
// MOLECULE_LLM_BILLING_MODE_RESOLVED. Emit both here, resolver-driven, for
// every mode so the value is correct on the byok/disabled early-return path
// too (and downstream consumers / debug shells see byok, not platform_managed).
envVars["MOLECULE_LLM_BILLING_MODE"] = res.ResolvedMode
// Observability: surface the resolved mode in the container env so the
// agent / debug shell can answer "why is my key being stripped" without
// pulling logs or hitting the admin route.
envVars["MOLECULE_LLM_BILLING_MODE_RESOLVED"] = res.ResolvedMode
if res.ResolvedMode != LLMBillingModePlatformManaged {
// byok or disabled — DO NOT force-route to CP, DO NOT override the
// workspace's own ANTHROPIC_BASE_URL / OAuth token.
//
// internal#711: but DO strip platform-origin LLM credentials. The
// platform's scope:global CLAUDE_CODE_OAUTH_TOKEN (+ the rest of the
// bypass-key set) was merged into envVars by loadWorkspaceSecrets
// from global_secrets; without this strip a BYOK workspace that
// brought no LLM credential of its own would inherit the platform's
// global token and bill the platform's Anthropic credits. The strip
// is PROVENANCE-AWARE: only keys still flagged as global_secrets
// origin are removed; a workspace's own LLM cred (a workspace_secrets
// row — provenance flag already dropped by loadWorkspaceSecrets)
// survives so the workspace talks to its own provider directly.
stripGlobalOriginLLMCreds(envVars, globalKeys)
return platformLLMEnvResult{
ResolvedMode: res.ResolvedMode,
HasUsableLLMCred: hasAnyPlatformManagedLLMKey(envVars),
}
// workspace only when the org is in platform-managed mode. Provider keys
// never enter the tenant; provider SDK API-key envs receive the tenant token
// for the CP proxy only when the workspace has not supplied BYOK/OAuth auth.
func applyPlatformManagedLLMEnv(envVars map[string]string, runtime string, model string) {
if strings.ToLower(strings.TrimSpace(os.Getenv("MOLECULE_LLM_BILLING_MODE"))) != "platform_managed" {
return
}
baseURL := firstNonEmptyEnv("MOLECULE_LLM_BASE_URL", "OPENAI_BASE_URL")
anthropicBaseURL := firstNonEmptyEnv("MOLECULE_LLM_ANTHROPIC_BASE_URL", "ANTHROPIC_BASE_URL")
token := firstNonEmptyEnv("MOLECULE_LLM_USAGE_TOKEN", "OPENAI_API_KEY")
if baseURL == "" || token == "" {
// Proxy not configured (boot race / misconfig). On the platform_managed
// path the workspace IS entitled to platform creds, so we do NOT strip
// here — but we report HasUsableLLMCred from whatever survived so the
// caller's fail-closed branch (non-platform only) is never reached on
// this path.
return platformLLMEnvResult{ResolvedMode: res.ResolvedMode, HasUsableLLMCred: true}
return
}
stripPlatformManagedLLMBypassEnv(envVars)
// MOLECULE_LLM_BILLING_MODE is already set to res.ResolvedMode (==
// platform_managed on this path) above (internal#703); no hardcode here.
envVars["MOLECULE_LLM_BILLING_MODE"] = "platform_managed"
envVars["MOLECULE_LLM_BASE_URL"] = baseURL
envVars["MOLECULE_LLM_USAGE_TOKEN"] = token
if anthropicBaseURL != "" {
@@ -1049,11 +946,11 @@ func applyPlatformManagedLLMEnv(ctx context.Context, envVars map[string]string,
envVars["MOLECULE_LLM_USAGE_URL"] = usageURL
}
if !runtimeUsesAnthropicNativeProxy(runtime) {
if strings.TrimSpace(envVars["OPENAI_API_KEY"]) == "" && !runtimeUsesAnthropicNativeProxy(runtime) {
envVars["OPENAI_API_KEY"] = token
envVars["OPENAI_BASE_URL"] = baseURL
}
if runtimeUsesAnthropicNativeProxy(runtime) && anthropicBaseURL != "" {
if runtimeUsesAnthropicNativeProxy(runtime) && anthropicBaseURL != "" && workspaceHasNoAnthropicAuth(envVars) {
envVars["ANTHROPIC_API_KEY"] = token
envVars["ANTHROPIC_BASE_URL"] = anthropicBaseURL
}
@@ -1063,55 +960,27 @@ func applyPlatformManagedLLMEnv(ctx context.Context, envVars map[string]string,
envVars["MOLECULE_MODEL"] = defaultModel
}
}
// platform_managed: the CP proxy usage token (injected as ANTHROPIC_API_KEY
// / OPENAI_API_KEY above) IS the usable credential, so the workspace is
// never fail-closed on this path.
return platformLLMEnvResult{ResolvedMode: res.ResolvedMode, HasUsableLLMCred: true}
}
func stripPlatformManagedLLMBypassEnv(envVars map[string]string) {
for key := range platformManagedDirectLLMBypassKeys {
delete(envVars, key)
}
}
// stripGlobalOriginLLMCreds removes platform-managed LLM credential keys
// (CLAUDE_CODE_OAUTH_TOKEN + the rest of platformManagedDirectLLMBypassKeys)
// from envVars ONLY when they originated from the operator-controlled
// `global_secrets` table (i.e. their key is present in globalKeys).
//
// internal#711 provider-aware gate. A platform global LLM credential is the
// platform's own credential and must never be the credential a non-platform
// (byok / subscription) workspace runs on. loadWorkspaceSecrets drops the
// global-provenance flag for any key the workspace re-set via the canvas
// Secrets tab (a workspace_secrets row), so a workspace's OWN LLM credential
// is NOT in globalKeys and survives this strip — only the inherited platform
// global creds are removed.
func stripGlobalOriginLLMCreds(envVars map[string]string, globalKeys map[string]struct{}) {
for key := range platformManagedDirectLLMBypassKeys {
if _, fromGlobal := globalKeys[key]; fromGlobal {
delete(envVars, key)
}
}
}
// hasAnyPlatformManagedLLMKey reports whether envVars still carries at least
// one non-empty platform-managed LLM credential key after the provider-aware
// gate. Used by the non-platform fail-closed branch: a byok/subscription
// workspace with no surviving (workspace-scoped) LLM credential must be
// aborted with MISSING_BYOK_CREDENTIAL rather than started credential-less or
// on stripped platform creds.
func hasAnyPlatformManagedLLMKey(envVars map[string]string) bool {
for key := range platformManagedDirectLLMBypassKeys {
if strings.TrimSpace(envVars[key]) != "" {
return true
}
}
return false
}
func runtimeUsesAnthropicNativeProxy(runtime string) bool {
return strings.EqualFold(strings.TrimSpace(runtime), "claude-code")
return strings.TrimSpace(strings.ToLower(runtime)) == "claude-code"
}
func workspaceHasNoAnthropicAuth(envVars map[string]string) bool {
for _, key := range []string{
"CLAUDE_CODE_OAUTH_TOKEN",
"ANTHROPIC_API_KEY",
"ANTHROPIC_AUTH_TOKEN",
"MINIMAX_API_KEY",
"KIMI_API_KEY",
"GLM_API_KEY",
"DEEPSEEK_API_KEY",
} {
if strings.TrimSpace(envVars[key]) != "" {
return false
}
}
return true
}
func firstNonEmptyEnv(names ...string) string {
@@ -193,35 +193,7 @@ func (h *WorkspaceHandler) prepareProvisionContext(
// continue to rely on workspace_secrets / org-import persona-env
// merge for their git auth.
applyAgentGitHTTPCreds(envVars, payload.Role)
// internal#711: provider-aware LLM-credential resolution. On a non-platform
// (byok/subscription) workspace this strips the platform's scope:global LLM
// creds inherited from global_secrets and reports whether the workspace
// still has a usable (workspace-scoped) LLM credential of its own.
llmRes := applyPlatformManagedLLMEnv(ctx, envVars, globalSecretKeys, workspaceID, payload.Runtime, payload.Model)
// Fail closed for a BYOK workspace with no usable LLM credential: do NOT
// start it on the platform's (now-stripped) global creds. Mirror the
// "model+provider+credential REQUIRED at create" spirit (internal#711)
// with an actionable error surfaced at provision time.
//
// Scoped to byok specifically (NOT disabled): "byok" means "the user
// intends to run an LLM on their own credential" — a missing one is a
// misconfiguration worth surfacing loudly. "disabled" means "this
// workspace runs no platform-billed LLM at all" (terminal / file work, or
// a runtime that talks to a non-bypass-key endpoint); stripping the
// inherited platform globals is sufficient there and aborting would
// regress a legitimate no-LLM workspace. The strip above already ran for
// both non-platform modes.
//
// The bypass-key check is intentionally broad — any surviving bypass key
// (the workspace's own, of workspace_secrets provenance) clears it.
if llmRes.ResolvedMode == LLMBillingModeBYOK && !llmRes.HasUsableLLMCred {
msg := formatMissingBYOKCredentialError(llmRes.ResolvedMode)
log.Printf("Provisioner: ABORT workspace=%s — byok billing mode has no usable LLM credential (MISSING_BYOK_CREDENTIAL, internal#711)", workspaceID)
return nil, &provisionAbort{
Msg: msg,
Extra: map[string]interface{}{"error": msg, "code": "MISSING_BYOK_CREDENTIAL", "billing_mode": llmRes.ResolvedMode, "issue": "711"},
}
}
applyPlatformManagedLLMEnv(envVars, payload.Runtime, payload.Model)
applyRuntimeModelEnv(envVars, payload.Runtime, payload.Model)
if payload.Role != "" {
envVars["MOLECULE_AGENT_ROLE"] = payload.Role
@@ -494,57 +494,6 @@ func TestPrepareProvisionContext_WorkspaceSecretWinsOverPersonaToken(t *testing.
}
}
// TestPrepareProvisionContext_ByokWithOnlyGlobalOAuthFailsClosed is the
// internal#711 end-to-end guard for the live Reno Stars leak. A byok
// workspace whose ONLY LLM credential is the platform's scope:global
// CLAUDE_CODE_OAUTH_TOKEN (inherited from global_secrets, no workspace
// override) must:
//
// 1. have that platform token STRIPPED from the prepared env (no leak), and
// 2. ABORT the provision with the MISSING_BYOK_CREDENTIAL code rather than
// start the workspace on the platform's credits.
//
// This is the discriminating end-to-end test: pre-fix prepared.EnvVars would
// carry CLAUDE_CODE_OAUTH_TOKEN=<platform token> and the provision would
// succeed, running Opus on Molecule's Anthropic credits.
func TestPrepareProvisionContext_ByokWithOnlyGlobalOAuthFailsClosed(t *testing.T) {
const wsID = "352e3c2b-0546-4e9c-b487-1e2ff1cf29fc" // Reno Stars SEO agent
t.Setenv("MOLECULE_LLM_BILLING_MODE", LLMBillingModePlatformManaged)
mock := setupTestDB(t)
// global_secrets carries the platform's scope:global OAuth token.
mock.ExpectQuery(`SELECT key, encrypted_value, encryption_version FROM global_secrets`).
WillReturnRows(sqlmock.NewRows([]string{"key", "encrypted_value", "encryption_version"}).
AddRow("CLAUDE_CODE_OAUTH_TOKEN", []byte("PLATFORM-GLOBAL-OAUTH"), 0))
// Workspace set NO secrets of its own.
mock.ExpectQuery(`SELECT key, encrypted_value, encryption_version FROM workspace_secrets`).
WithArgs(wsID).
WillReturnRows(sqlmock.NewRows([]string{"key", "encrypted_value", "encryption_version"}))
// Resolver: workspace override = byok.
mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(wsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(LLMBillingModeBYOK))
handler := NewWorkspaceHandler(&captureBroadcaster{}, nil, "http://localhost:8080", t.TempDir())
payload := models.CreateWorkspacePayload{
Name: "Reno Stars SEO",
Runtime: "claude-code",
Tier: 1,
}
prepared, abort := handler.prepareProvisionContext(
context.Background(), wsID, "/nonexistent", nil, payload, false)
if abort == nil {
t.Fatalf("expected MISSING_BYOK_CREDENTIAL abort, got success (prepared=%v) — the leak would still ship", prepared)
}
if code, _ := abort.Extra["code"].(string); code != "MISSING_BYOK_CREDENTIAL" {
t.Fatalf("abort.Extra[code] = %v, want MISSING_BYOK_CREDENTIAL", abort.Extra["code"])
}
if mode, _ := abort.Extra["billing_mode"].(string); mode != LLMBillingModeBYOK {
t.Fatalf("abort.Extra[billing_mode] = %v, want %q", abort.Extra["billing_mode"], LLMBillingModeBYOK)
}
}
// TestReadOrLazyHealInboundSecret pins the four branches of the
// shared lazy-heal helper directly. Each call site (chat_files,
// registry) has its own integration test, but those go through the
@@ -1023,7 +972,7 @@ func TestApplyPlatformManagedLLMEnv_NonClaudeRuntimeDefaultsOpenAIProxyWhenNoWor
t.Setenv("MOLECULE_LLM_DEFAULT_MODEL", "moonshot/kimi-k2.6")
envVars := map[string]string{}
applyPlatformManagedLLMEnv(context.Background(), envVars, nil, "", "codex", "")
applyPlatformManagedLLMEnv(envVars, "codex", "")
applyRuntimeModelEnv(envVars, "codex", "")
if got := envVars["OPENAI_BASE_URL"]; got != "https://api.example.test/api/v1/internal/llm/openai/v1" {
@@ -1043,7 +992,7 @@ func TestApplyPlatformManagedLLMEnv_NonClaudeRuntimeDefaultsOpenAIProxyWhenNoWor
}
}
func TestApplyPlatformManagedLLMEnv_StripsWorkspaceOpenAIKeyForClaudeCode(t *testing.T) {
func TestApplyPlatformManagedLLMEnv_DoesNotOverrideWorkspaceOpenAIKey(t *testing.T) {
t.Setenv("MOLECULE_LLM_BILLING_MODE", "platform_managed")
t.Setenv("MOLECULE_LLM_BASE_URL", "https://api.example.test/api/v1/internal/llm/openai/v1")
t.Setenv("MOLECULE_LLM_USAGE_TOKEN", "tenant-admin-token")
@@ -1053,13 +1002,13 @@ func TestApplyPlatformManagedLLMEnv_StripsWorkspaceOpenAIKeyForClaudeCode(t *tes
"OPENAI_BASE_URL": "https://api.openai.com/v1",
"MODEL": "openai/gpt-5.5",
}
applyPlatformManagedLLMEnv(context.Background(), envVars, nil, "", "claude-code", "")
applyPlatformManagedLLMEnv(envVars, "claude-code", "")
if _, ok := envVars["OPENAI_API_KEY"]; ok {
t.Fatalf("OPENAI_API_KEY should be stripped for claude-code platform-managed mode")
if got := envVars["OPENAI_API_KEY"]; got != "user-openai-key" {
t.Fatalf("OPENAI_API_KEY was overwritten: %q", got)
}
if _, ok := envVars["OPENAI_BASE_URL"]; ok {
t.Fatalf("OPENAI_BASE_URL should be stripped for claude-code platform-managed mode")
if got := envVars["OPENAI_BASE_URL"]; got != "https://api.openai.com/v1" {
t.Fatalf("OPENAI_BASE_URL was overwritten: %q", got)
}
if got := envVars["MOLECULE_LLM_USAGE_TOKEN"]; got != "tenant-admin-token" {
t.Fatalf("MOLECULE_LLM_USAGE_TOKEN = %q", got)
@@ -1069,7 +1018,7 @@ func TestApplyPlatformManagedLLMEnv_StripsWorkspaceOpenAIKeyForClaudeCode(t *tes
}
}
func TestApplyPlatformManagedLLMEnv_ClaudeCodeUsesAnthropicProxyOverOAuth(t *testing.T) {
func TestApplyPlatformManagedLLMEnv_ClaudeCodeUsesAnthropicProxyWithoutOverwritingOAuth(t *testing.T) {
t.Setenv("MOLECULE_LLM_BILLING_MODE", "platform_managed")
t.Setenv("MOLECULE_LLM_BASE_URL", "https://api.example.test/api/v1/internal/llm/openai/v1")
t.Setenv("MOLECULE_LLM_ANTHROPIC_BASE_URL", "https://api.example.test/api/v1/internal/llm/anthropic/v1")
@@ -1079,16 +1028,13 @@ func TestApplyPlatformManagedLLMEnv_ClaudeCodeUsesAnthropicProxyOverOAuth(t *tes
"CLAUDE_CODE_OAUTH_TOKEN": "user-oauth-token",
"MODEL": "sonnet",
}
applyPlatformManagedLLMEnv(context.Background(), envVars, nil, "", "claude-code", "")
applyPlatformManagedLLMEnv(envVars, "claude-code", "")
if _, ok := envVars["CLAUDE_CODE_OAUTH_TOKEN"]; ok {
t.Fatalf("CLAUDE_CODE_OAUTH_TOKEN should be stripped in platform-managed mode")
if got := envVars["CLAUDE_CODE_OAUTH_TOKEN"]; got != "user-oauth-token" {
t.Fatalf("CLAUDE_CODE_OAUTH_TOKEN was overwritten: %q", got)
}
if got := envVars["ANTHROPIC_API_KEY"]; got != "tenant-admin-token" {
t.Fatalf("ANTHROPIC_API_KEY = %q", got)
}
if got := envVars["ANTHROPIC_BASE_URL"]; got != "https://api.example.test/api/v1/internal/llm/anthropic/v1" {
t.Fatalf("ANTHROPIC_BASE_URL = %q", got)
if _, ok := envVars["ANTHROPIC_API_KEY"]; ok {
t.Fatalf("ANTHROPIC_API_KEY should not be set when Claude OAuth is present")
}
if got := envVars["MOLECULE_LLM_ANTHROPIC_BASE_URL"]; got != "https://api.example.test/api/v1/internal/llm/anthropic/v1" {
t.Fatalf("MOLECULE_LLM_ANTHROPIC_BASE_URL = %q", got)
@@ -1102,7 +1048,7 @@ func TestApplyPlatformManagedLLMEnv_ClaudeCodeInjectsAnthropicProxyWhenNoWorkspa
t.Setenv("MOLECULE_LLM_USAGE_TOKEN", "tenant-admin-token")
envVars := map[string]string{}
applyPlatformManagedLLMEnv(context.Background(), envVars, nil, "", "claude-code", "minimax/MiniMax-M2.7")
applyPlatformManagedLLMEnv(envVars, "claude-code", "minimax/MiniMax-M2.7")
if got := envVars["ANTHROPIC_BASE_URL"]; got != "https://api.example.test/api/v1/internal/llm/anthropic/v1" {
t.Fatalf("ANTHROPIC_BASE_URL = %q", got)
@@ -1115,7 +1061,7 @@ func TestApplyPlatformManagedLLMEnv_ClaudeCodeInjectsAnthropicProxyWhenNoWorkspa
}
}
func TestApplyPlatformManagedLLMEnv_ClaudeCodeStripsVendorBYOK(t *testing.T) {
func TestApplyPlatformManagedLLMEnv_ClaudeCodeDoesNotOverrideVendorBYOK(t *testing.T) {
t.Setenv("MOLECULE_LLM_BILLING_MODE", "platform_managed")
t.Setenv("MOLECULE_LLM_BASE_URL", "https://api.example.test/api/v1/internal/llm/openai/v1")
t.Setenv("MOLECULE_LLM_ANTHROPIC_BASE_URL", "https://api.example.test/api/v1/internal/llm/anthropic/v1")
@@ -1125,16 +1071,16 @@ func TestApplyPlatformManagedLLMEnv_ClaudeCodeStripsVendorBYOK(t *testing.T) {
"MINIMAX_API_KEY": "user-minimax-key",
"MODEL": "MiniMax-M2.7",
}
applyPlatformManagedLLMEnv(context.Background(), envVars, nil, "", "claude-code", "")
applyPlatformManagedLLMEnv(envVars, "claude-code", "")
if _, ok := envVars["MINIMAX_API_KEY"]; ok {
t.Fatalf("MINIMAX_API_KEY should be stripped in platform-managed mode")
if got := envVars["MINIMAX_API_KEY"]; got != "user-minimax-key" {
t.Fatalf("MINIMAX_API_KEY was overwritten: %q", got)
}
if got := envVars["ANTHROPIC_API_KEY"]; got != "tenant-admin-token" {
t.Fatalf("ANTHROPIC_API_KEY = %q", got)
if _, ok := envVars["ANTHROPIC_API_KEY"]; ok {
t.Fatalf("ANTHROPIC_API_KEY should not be set when vendor BYOK is present")
}
if got := envVars["ANTHROPIC_BASE_URL"]; got != "https://api.example.test/api/v1/internal/llm/anthropic/v1" {
t.Fatalf("ANTHROPIC_BASE_URL = %q", got)
if _, ok := envVars["ANTHROPIC_BASE_URL"]; ok {
t.Fatalf("ANTHROPIC_BASE_URL should not be set when vendor BYOK is present")
}
if got := envVars["MOLECULE_LLM_USAGE_TOKEN"]; got != "tenant-admin-token" {
t.Fatalf("MOLECULE_LLM_USAGE_TOKEN = %q", got)
@@ -1147,7 +1093,7 @@ func TestApplyPlatformManagedLLMEnv_NoopsOutsidePlatformManaged(t *testing.T) {
t.Setenv("MOLECULE_LLM_USAGE_TOKEN", "tenant-admin-token")
envVars := map[string]string{}
applyPlatformManagedLLMEnv(context.Background(), envVars, nil, "", "claude-code", "")
applyPlatformManagedLLMEnv(envVars, "claude-code", "")
if _, ok := envVars["OPENAI_API_KEY"]; ok {
t.Fatalf("OPENAI_API_KEY should not be set outside platform-managed mode")
@@ -1157,288 +1103,6 @@ func TestApplyPlatformManagedLLMEnv_NoopsOutsidePlatformManaged(t *testing.T) {
}
}
// TestApplyPlatformManagedLLMEnv_ClaudeCodeByokKeepsOwnProviderEnv is the
// internal#703 regression guard: a per-workspace byok override (org-level
// MOLECULE_LLM_BILLING_MODE left at the platform_managed bootstrap floor)
// must resolve to byok and leave the workspace own provider env intact —
// the CP-injected proxy ANTHROPIC_BASE_URL / usage token must NOT be forced,
// the OAuth token must NOT be stripped, and MOLECULE_LLM_BILLING_MODE in the
// container must read the RESOLVED mode (byok), not the hardcoded literal.
//
// This is the discriminating test for the byok end-to-end fix: pre-fix the
// strip path was the only emitter of MOLECULE_LLM_BILLING_MODE (hardcoded
// "platform_managed"), so a byok container carried no truthful billing mode.
func TestApplyPlatformManagedLLMEnv_ClaudeCodeByokKeepsOwnProviderEnv(t *testing.T) {
const wsID = "77777777-7777-7777-7777-777777777777"
mock := setupTestDB(t)
mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(wsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(LLMBillingModeBYOK))
// Org-level env left at the bootstrap floor — the per-workspace override
// is what must flip this workspace to byok (the realistic prod shape).
t.Setenv("MOLECULE_LLM_BILLING_MODE", LLMBillingModePlatformManaged)
t.Setenv("MOLECULE_LLM_BASE_URL", "https://api.example.test/api/v1/internal/llm/openai/v1")
t.Setenv("MOLECULE_LLM_ANTHROPIC_BASE_URL", "https://api.example.test/api/v1/internal/llm/anthropic")
t.Setenv("MOLECULE_LLM_USAGE_TOKEN", "tenant-admin-token")
// The workspace brought its own Claude Code OAuth token (BYOK via the
// subscription provider). It must survive untouched.
envVars := map[string]string{
"CLAUDE_CODE_OAUTH_TOKEN": "user-oauth-token",
"MODEL": "sonnet",
}
applyPlatformManagedLLMEnv(context.Background(), envVars, nil, wsID, "claude-code", "")
// 1. OAuth token intact — not stripped.
if got := envVars["CLAUDE_CODE_OAUTH_TOKEN"]; got != "user-oauth-token" {
t.Fatalf("CLAUDE_CODE_OAUTH_TOKEN = %q, want it left intact for byok", got)
}
// 2. No CP proxy base URL / usage token forced onto the workspace.
if got, ok := envVars["ANTHROPIC_BASE_URL"]; ok {
t.Fatalf("ANTHROPIC_BASE_URL must NOT be injected for byok, got %q", got)
}
if got, ok := envVars["ANTHROPIC_API_KEY"]; ok {
t.Fatalf("ANTHROPIC_API_KEY must NOT be injected for byok, got %q", got)
}
if got, ok := envVars["MOLECULE_LLM_ANTHROPIC_BASE_URL"]; ok {
t.Fatalf("MOLECULE_LLM_ANTHROPIC_BASE_URL must NOT be injected for byok, got %q", got)
}
if got, ok := envVars["MOLECULE_LLM_USAGE_TOKEN"]; ok {
t.Fatalf("MOLECULE_LLM_USAGE_TOKEN must NOT be injected for byok, got %q", got)
}
// 3. Billing mode in the container reflects the RESOLVED mode (byok).
if got := envVars["MOLECULE_LLM_BILLING_MODE"]; got != LLMBillingModeBYOK {
t.Fatalf("MOLECULE_LLM_BILLING_MODE = %q, want %q (resolver-driven, not hardcoded)", got, LLMBillingModeBYOK)
}
if got := envVars["MOLECULE_LLM_BILLING_MODE_RESOLVED"]; got != LLMBillingModeBYOK {
t.Fatalf("MOLECULE_LLM_BILLING_MODE_RESOLVED = %q, want %q", got, LLMBillingModeBYOK)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Errorf("unmet sqlmock expectations: %v", err)
}
}
// TestApplyPlatformManagedLLMEnv_ByokStripsGlobalOriginOAuthToken is the
// internal#711 regression guard for the live 2026-05-27 leak (Reno Stars SEO
// + Marketing claude-code agents). A non-platform (byok) workspace that
// brought NO LLM credential of its own, but which inherited the platform's
// scope:global CLAUDE_CODE_OAUTH_TOKEN from global_secrets (provenance =
// globalKeys), must have that platform token STRIPPED — not run on it.
//
// Pre-fix the byok early-return left envVars untouched, so the platform's
// global OAuth token survived into the container and the agent ran Opus on
// the platform's Anthropic credits. The fix gates the global-cred merge on
// provider==platform: a non-platform workspace keeps only its own
// (workspace_secrets) creds, of which there are none here.
func TestApplyPlatformManagedLLMEnv_ByokStripsGlobalOriginOAuthToken(t *testing.T) {
const wsID = "352e3c2b-0546-4e9c-b487-1e2ff1cf29fc" // Reno Stars SEO agent
mock := setupTestDB(t)
mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(wsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(LLMBillingModeBYOK))
t.Setenv("MOLECULE_LLM_BILLING_MODE", LLMBillingModePlatformManaged)
t.Setenv("MOLECULE_LLM_BASE_URL", "https://api.example.test/api/v1/internal/llm/openai/v1")
t.Setenv("MOLECULE_LLM_ANTHROPIC_BASE_URL", "https://api.example.test/api/v1/internal/llm/anthropic")
t.Setenv("MOLECULE_LLM_USAGE_TOKEN", "tenant-admin-token")
// The ONLY LLM credential in env is the platform's scope:global OAuth
// token, merged from global_secrets (so its key is in globalKeys). The
// workspace set none of its own.
envVars := map[string]string{
"CLAUDE_CODE_OAUTH_TOKEN": "PLATFORM-GLOBAL-OAUTH-TOKEN",
"MODEL": "opus",
}
globalKeys := map[string]struct{}{"CLAUDE_CODE_OAUTH_TOKEN": {}}
res := applyPlatformManagedLLMEnv(context.Background(), envVars, globalKeys, wsID, "claude-code", "")
// 1. The platform global OAuth token must be STRIPPED — the leak is closed.
if got, ok := envVars["CLAUDE_CODE_OAUTH_TOKEN"]; ok {
t.Fatalf("CLAUDE_CODE_OAUTH_TOKEN = %q present — platform scope:global token must be stripped for a byok workspace", got)
}
// 2. No CP proxy creds forced (byok = workspace talks to its own provider).
if got, ok := envVars["ANTHROPIC_API_KEY"]; ok {
t.Fatalf("ANTHROPIC_API_KEY must NOT be injected for byok, got %q", got)
}
// 3. Resolver reports byok with NO usable LLM credential → caller fails closed.
if res.ResolvedMode != LLMBillingModeBYOK {
t.Fatalf("ResolvedMode = %q, want %q", res.ResolvedMode, LLMBillingModeBYOK)
}
if res.HasUsableLLMCred {
t.Fatalf("HasUsableLLMCred = true, want false (only the stripped platform global token was present)")
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Errorf("unmet sqlmock expectations: %v", err)
}
}
// TestApplyPlatformManagedLLMEnv_ByokKeepsWorkspaceOwnOAuthEvenWithGlobal is
// the discriminating companion to the strip test: a byok workspace that DID
// set its own CLAUDE_CODE_OAUTH_TOKEN via the canvas Secrets tab (a
// workspace_secrets row) keeps it. loadWorkspaceSecrets drops the global
// provenance flag on a workspace override, so the key is NOT in globalKeys
// and the provenance-aware strip leaves it alone. Proves the fix strips only
// platform-origin creds, never the customer's own.
func TestApplyPlatformManagedLLMEnv_ByokKeepsWorkspaceOwnOAuthEvenWithGlobal(t *testing.T) {
const wsID = "6b66de8d-9337-4fb4-be8d-6d49dca0d809" // Reno Stars Marketing agent
mock := setupTestDB(t)
mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(wsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(LLMBillingModeBYOK))
t.Setenv("MOLECULE_LLM_BILLING_MODE", LLMBillingModePlatformManaged)
t.Setenv("MOLECULE_LLM_BASE_URL", "https://api.example.test/api/v1/internal/llm/openai/v1")
t.Setenv("MOLECULE_LLM_USAGE_TOKEN", "tenant-admin-token")
// Workspace set its OWN OAuth token — loadWorkspaceSecrets would have
// dropped its global provenance flag, so globalKeys does NOT contain it.
envVars := map[string]string{
"CLAUDE_CODE_OAUTH_TOKEN": "CUSTOMER-OWN-OAUTH-TOKEN",
"MODEL": "opus",
}
globalKeys := map[string]struct{}{} // not from global_secrets
res := applyPlatformManagedLLMEnv(context.Background(), envVars, globalKeys, wsID, "claude-code", "")
if got := envVars["CLAUDE_CODE_OAUTH_TOKEN"]; got != "CUSTOMER-OWN-OAUTH-TOKEN" {
t.Fatalf("CLAUDE_CODE_OAUTH_TOKEN = %q, want the workspace's own token left intact", got)
}
if !res.HasUsableLLMCred {
t.Fatalf("HasUsableLLMCred = false, want true (workspace brought its own credential)")
}
if res.ResolvedMode != LLMBillingModeBYOK {
t.Fatalf("ResolvedMode = %q, want %q", res.ResolvedMode, LLMBillingModeBYOK)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Errorf("unmet sqlmock expectations: %v", err)
}
}
// TestApplyPlatformManagedLLMEnv_DisabledStripsGlobalButReportsNoCred proves
// that "disabled" mode also strips the platform's global LLM creds (the leak
// is closed for disabled too), and reports HasUsableLLMCred=false. The
// caller's fail-closed abort is scoped to byok only, so a disabled workspace
// with no LLM cred still boots (for terminal / non-LLM work); here we pin the
// function-level strip + report.
func TestApplyPlatformManagedLLMEnv_DisabledStripsGlobalButReportsNoCred(t *testing.T) {
const wsID = "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"
mock := setupTestDB(t)
mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(wsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(LLMBillingModeDisabled))
t.Setenv("MOLECULE_LLM_BILLING_MODE", LLMBillingModePlatformManaged)
envVars := map[string]string{
"CLAUDE_CODE_OAUTH_TOKEN": "PLATFORM-GLOBAL-OAUTH-TOKEN",
}
globalKeys := map[string]struct{}{"CLAUDE_CODE_OAUTH_TOKEN": {}}
res := applyPlatformManagedLLMEnv(context.Background(), envVars, globalKeys, wsID, "claude-code", "")
if _, ok := envVars["CLAUDE_CODE_OAUTH_TOKEN"]; ok {
t.Fatalf("CLAUDE_CODE_OAUTH_TOKEN must be stripped for disabled mode too")
}
if res.ResolvedMode != LLMBillingModeDisabled {
t.Fatalf("ResolvedMode = %q, want %q", res.ResolvedMode, LLMBillingModeDisabled)
}
if res.HasUsableLLMCred {
t.Fatalf("HasUsableLLMCred = true, want false")
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Errorf("unmet sqlmock expectations: %v", err)
}
}
// TestApplyPlatformManagedLLMEnv_PlatformManagedStillReceivesGlobalCreds is
// the no-regression guard for the OTHER side of the gate (internal#711): a
// platform-managed workspace MUST still receive the platform's creds. Here
// the proxy IS configured, so the contract is the existing one — the global
// OAuth token is replaced by the proxy usage token (HasUsableLLMCred=true).
func TestApplyPlatformManagedLLMEnv_PlatformManagedStillReceivesGlobalCreds(t *testing.T) {
const wsID = "99999999-9999-9999-9999-999999999999"
mock := setupTestDB(t)
mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(wsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(LLMBillingModePlatformManaged))
t.Setenv("MOLECULE_LLM_BILLING_MODE", LLMBillingModePlatformManaged)
t.Setenv("MOLECULE_LLM_BASE_URL", "https://api.example.test/api/v1/internal/llm/openai/v1")
t.Setenv("MOLECULE_LLM_ANTHROPIC_BASE_URL", "https://api.example.test/api/v1/internal/llm/anthropic")
t.Setenv("MOLECULE_LLM_USAGE_TOKEN", "tenant-admin-token")
envVars := map[string]string{
"CLAUDE_CODE_OAUTH_TOKEN": "PLATFORM-GLOBAL-OAUTH-TOKEN",
"MODEL": "opus",
}
globalKeys := map[string]struct{}{"CLAUDE_CODE_OAUTH_TOKEN": {}}
res := applyPlatformManagedLLMEnv(context.Background(), envVars, globalKeys, wsID, "claude-code", "")
// Platform-managed routes through the CP proxy: OAuth stripped, proxy creds forced.
if _, ok := envVars["CLAUDE_CODE_OAUTH_TOKEN"]; ok {
t.Fatalf("CLAUDE_CODE_OAUTH_TOKEN should be stripped + replaced by the proxy token for platform_managed")
}
if got := envVars["ANTHROPIC_API_KEY"]; got != "tenant-admin-token" {
t.Fatalf("ANTHROPIC_API_KEY = %q, want proxy usage token for platform_managed", got)
}
if !res.HasUsableLLMCred {
t.Fatalf("HasUsableLLMCred = false, want true for platform_managed (proxy token is the credential)")
}
if res.ResolvedMode != LLMBillingModePlatformManaged {
t.Fatalf("ResolvedMode = %q, want %q", res.ResolvedMode, LLMBillingModePlatformManaged)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Errorf("unmet sqlmock expectations: %v", err)
}
}
// TestApplyPlatformManagedLLMEnv_PlatformManagedStillEmitsResolvedMode is the
// no-regression companion: a workspace that resolves to platform_managed must
// still strip + force the proxy AND emit MOLECULE_LLM_BILLING_MODE=
// platform_managed (now resolver-driven, internal#703). Proves the byok fix
// did not alter the platform_managed contract.
func TestApplyPlatformManagedLLMEnv_PlatformManagedStillEmitsResolvedMode(t *testing.T) {
const wsID = "88888888-8888-8888-8888-888888888888"
mock := setupTestDB(t)
mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WithArgs(wsID).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}).AddRow(LLMBillingModePlatformManaged))
t.Setenv("MOLECULE_LLM_BILLING_MODE", LLMBillingModePlatformManaged)
t.Setenv("MOLECULE_LLM_BASE_URL", "https://api.example.test/api/v1/internal/llm/openai/v1")
t.Setenv("MOLECULE_LLM_ANTHROPIC_BASE_URL", "https://api.example.test/api/v1/internal/llm/anthropic")
t.Setenv("MOLECULE_LLM_USAGE_TOKEN", "tenant-admin-token")
envVars := map[string]string{
"CLAUDE_CODE_OAUTH_TOKEN": "user-oauth-token",
"MODEL": "sonnet",
}
applyPlatformManagedLLMEnv(context.Background(), envVars, nil, wsID, "claude-code", "")
// OAuth stripped, proxy forced — unchanged platform_managed contract.
if _, ok := envVars["CLAUDE_CODE_OAUTH_TOKEN"]; ok {
t.Fatalf("CLAUDE_CODE_OAUTH_TOKEN should be stripped for platform_managed")
}
if got := envVars["ANTHROPIC_BASE_URL"]; got != "https://api.example.test/api/v1/internal/llm/anthropic" {
t.Fatalf("ANTHROPIC_BASE_URL = %q, want proxy forced for platform_managed", got)
}
if got := envVars["ANTHROPIC_API_KEY"]; got != "tenant-admin-token" {
t.Fatalf("ANTHROPIC_API_KEY = %q, want usage token for platform_managed", got)
}
if got := envVars["MOLECULE_LLM_BILLING_MODE"]; got != LLMBillingModePlatformManaged {
t.Fatalf("MOLECULE_LLM_BILLING_MODE = %q, want %q", got, LLMBillingModePlatformManaged)
}
if got := envVars["MOLECULE_LLM_BILLING_MODE_RESOLVED"]; got != LLMBillingModePlatformManaged {
t.Fatalf("MOLECULE_LLM_BILLING_MODE_RESOLVED = %q, want %q", got, LLMBillingModePlatformManaged)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Errorf("unmet sqlmock expectations: %v", err)
}
}
// TestApplyRuntimeModelEnv_PersonaEnvMODELSecretPreserved locks in the
// 2026-05-08 fix that prevents the MODEL_PROVIDER-as-slug fallback from
// silently overwriting a per-persona MODEL workspace_secret on restart,
@@ -1616,28 +1616,3 @@ func (*mockResolver) Scheme() string { return "" }
func (m *mockResolver) Fetch(_ context.Context, _, _ string) (string, error) {
return m.fetchName, m.fetchErr
}
// TestRuntimeUsesAnthropicNativeProxy_CaseAndWhitespace proves the
// strings.EqualFold hardening: the runtime check now matches "claude-code"
// case-insensitively (and after trimming whitespace) instead of relying on
// a lowercased exact compare.
func TestRuntimeUsesAnthropicNativeProxy_CaseAndWhitespace(t *testing.T) {
cases := []struct {
runtime string
want bool
}{
{"claude-code", true},
{"Claude-Code", true},
{"CLAUDE-CODE", true},
{" claude-code ", true},
{"\tClaude-Code\n", true},
{"claude-code-x", false},
{"codex", false},
{"", false},
}
for _, c := range cases {
if got := runtimeUsesAnthropicNativeProxy(c.runtime); got != c.want {
t.Errorf("runtimeUsesAnthropicNativeProxy(%q) = %v, want %v", c.runtime, got, c.want)
}
}
}
@@ -3,7 +3,6 @@ package handlers
import (
"context"
"database/sql"
"io"
"log"
"net/http"
"runtime/debug"
@@ -284,10 +283,7 @@ func (h *WorkspaceHandler) Restart(c *gin.Context) {
Reset bool `json:"reset"` // #12: discard claude-sessions volume before restart
RebuildConfig bool `json:"rebuild_config"` // #239: re-render config volume from org-template source (recovery path when volume was destroyed)
}
if err := c.ShouldBindJSON(&body); err != nil && err != io.EOF {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid JSON body"})
return
}
c.ShouldBindJSON(&body)
// Read runtime from container's config.yaml before stopping. Docker-
// only: in SaaS mode the workspace runs on a remote EC2 and we can't
@@ -296,10 +292,8 @@ func (h *WorkspaceHandler) Restart(c *gin.Context) {
containerRuntime := h.restartRuntimeFromConfig(ctx, id, wsName, dbRuntime, body.ApplyTemplate)
// Reset to provisioning
if _, err := db.DB.ExecContext(ctx,
`UPDATE workspaces SET status = $1, url = '', updated_at = now() WHERE id = $2`, models.StatusProvisioning, id); err != nil {
log.Printf("Restart: failed to set provisioning status for %s: %v", id, err)
}
db.DB.ExecContext(ctx,
`UPDATE workspaces SET status = $1, url = '', updated_at = now() WHERE id = $2`, models.StatusProvisioning, id)
h.broadcaster.RecordAndBroadcast(ctx, string(events.EventWorkspaceProvisioning), id, map[string]interface{}{
"name": wsName,
"tier": tier,
@@ -389,9 +383,7 @@ func (h *WorkspaceHandler) restartRuntimeFromConfig(ctx context.Context, id, wsN
if parsed != "" && parsed != containerRuntime {
log.Printf("Restart: runtime changed in config.yaml %q→%q for %s", containerRuntime, parsed, wsName)
containerRuntime = parsed
if _, err := db.DB.ExecContext(ctx, `UPDATE workspaces SET runtime = $1 WHERE id = $2`, containerRuntime, id); err != nil {
log.Printf("Restart: failed to persist runtime %q for %s: %v", containerRuntime, id, err)
}
db.DB.ExecContext(ctx, `UPDATE workspaces SET runtime = $1 WHERE id = $2`, containerRuntime, id)
}
break
}
@@ -474,11 +466,7 @@ func (h *WorkspaceHandler) HibernateWorkspace(ctx context.Context, workspaceID s
log.Printf("Hibernate: atomic claim failed for %s: %v", workspaceID, err)
return
}
rowsAffected, err := result.RowsAffected()
if err != nil {
log.Printf("Hibernate: RowsAffected error for %s: %v", workspaceID, err)
return
}
rowsAffected, _ := result.RowsAffected()
if rowsAffected == 0 {
// Either already hibernating/hibernated/paused/removed, or active_tasks > 0 —
// safe to abort without side-effects.
@@ -721,31 +709,8 @@ var cpStopRetryBaseDelay = 1 * time.Second
//
// Returns nothing — caller's contract is unchanged.
func (h *WorkspaceHandler) cpStopWithRetry(ctx context.Context, workspaceID, source string) {
// Restart's contract is "make the workspace alive again": it proceeds
// with reprovision regardless of the Stop outcome, so it discards the
// terminal error. The delete path needs the error (it must keep the
// row recoverable for the orphan-sweeper + emit a durable event), so
// the actual retry loop lives in cpStopWithRetryErr below.
_ = h.cpStopWithRetryErr(ctx, workspaceID, source)
}
// cpStopWithRetryErr is the shared bounded-retry core for cpProv.Stop.
// It returns the terminal error so callers that need to react to a leak
// (the DELETE path's stopWorkspaceForDelete) can do so, while
// cpStopWithRetry keeps its void contract for the restart paths.
//
// Behaviour (unchanged from the original cpStopWithRetry loop):
// - cpProv nil → nil (no-op; nothing to stop).
// - success on attempt N → nil; logs a retry-success line when N > 1.
// - ctx cancelled mid-retry → returns ctx.Err(); logs an "abandoned"
// line and deliberately does NOT emit LEAK-SUSPECT (operator-initiated
// drain is a different signal than "we tried hard and failed").
// - all attempts fail → returns the LAST attempt's error and emits the
// stable `LEAK-SUSPECT cpProv.Stop ...` log line so the CP-side orphan
// reconciler can correlate by workspace_id.
func (h *WorkspaceHandler) cpStopWithRetryErr(ctx context.Context, workspaceID, source string) error {
if h.cpProv == nil {
return nil
return
}
var lastErr error
delay := cpStopRetryBaseDelay
@@ -755,7 +720,7 @@ func (h *WorkspaceHandler) cpStopWithRetryErr(ctx context.Context, workspaceID,
if attempt > 1 {
log.Printf("%s: cpProv.Stop(%s) succeeded on attempt %d", source, workspaceID, attempt)
}
return nil
return
}
lastErr = err
if attempt == cpStopRetryAttempts {
@@ -763,14 +728,12 @@ func (h *WorkspaceHandler) cpStopWithRetryErr(ctx context.Context, workspaceID,
}
// Sleep with ctx awareness so a cancelled ctx exits early instead
// of stalling the goroutine through the remaining backoff.
timer := time.NewTimer(delay)
select {
case <-ctx.Done():
timer.Stop()
log.Printf("%s: cpProv.Stop(%s) abandoned mid-retry: ctx cancelled (last_err=%v)",
source, workspaceID, lastErr)
return ctx.Err()
case <-timer.C:
return
case <-time.After(delay):
}
delay *= 2
}
@@ -778,7 +741,6 @@ func (h *WorkspaceHandler) cpStopWithRetryErr(ctx context.Context, workspaceID,
// so logs are greppable / parseable for the CP-side orphan reconciler.
log.Printf("LEAK-SUSPECT cpProv.Stop workspace_id=%s source=%s attempts=%d last_err=%q",
workspaceID, source, cpStopRetryAttempts, lastErr.Error())
return lastErr
}
// runRestartCycle does the actual stop+provision work for one restart
@@ -832,10 +794,8 @@ func (h *WorkspaceHandler) runRestartCycle(workspaceID string) {
h.stopForRestart(ctx, workspaceID)
if _, err := db.DB.ExecContext(ctx,
`UPDATE workspaces SET status = $1, url = '', updated_at = now() WHERE id = $2`, models.StatusProvisioning, workspaceID); err != nil {
log.Printf("Auto-restart: failed to set provisioning status for %s: %v", workspaceID, err)
}
db.DB.ExecContext(ctx,
`UPDATE workspaces SET status = $1, url = '', updated_at = now() WHERE id = $2`, models.StatusProvisioning, workspaceID)
h.broadcaster.RecordAndBroadcast(ctx, string(events.EventWorkspaceProvisioning), workspaceID, map[string]interface{}{
"name": wsName, "tier": tier, "runtime": dbRuntime,
})
@@ -894,15 +854,12 @@ func (h *WorkspaceHandler) Pause(c *gin.Context) {
// Collect this workspace + all descendants to pause
toPause := []struct{ id, name string }{{id, wsName}}
rows, err := db.DB.QueryContext(ctx,
rows, _ := db.DB.QueryContext(ctx,
`WITH RECURSIVE descendants AS (
SELECT id, name FROM workspaces WHERE parent_id = $1 AND status NOT IN ('removed', 'paused')
UNION ALL
SELECT w.id, w.name FROM workspaces w JOIN descendants d ON w.parent_id = d.id WHERE w.status NOT IN ('removed', 'paused')
) SELECT id, name FROM descendants`, id)
if err != nil {
log.Printf("Pause: descendant query failed for %s: %v", id, err)
}
if rows != nil {
defer rows.Close()
for rows.Next() {
@@ -929,10 +886,8 @@ func (h *WorkspaceHandler) Pause(c *gin.Context) {
if err := h.StopWorkspaceAuto(ctx, ws.id); err != nil {
log.Printf("Pause: stop %s failed: %v — orphan sweeper will reconcile", ws.id, err)
}
if _, err := db.DB.ExecContext(ctx,
`UPDATE workspaces SET status = $1, url = '', updated_at = now() WHERE id = $2`, models.StatusPaused, ws.id); err != nil {
log.Printf("Pause: failed to set paused status for %s: %v", ws.id, err)
}
db.DB.ExecContext(ctx,
`UPDATE workspaces SET status = $1, url = '', updated_at = now() WHERE id = $2`, models.StatusPaused, ws.id)
db.ClearWorkspaceKeys(ctx, ws.id)
h.broadcaster.RecordAndBroadcast(ctx, string(events.EventWorkspacePaused), ws.id, map[string]interface{}{
"name": ws.name,
@@ -983,15 +938,12 @@ func (h *WorkspaceHandler) Resume(c *gin.Context) {
tier int
}
toResume := []wsInfo{{id, wsName, dbRuntime, tier}}
rows, err := db.DB.QueryContext(ctx,
rows, _ := db.DB.QueryContext(ctx,
`WITH RECURSIVE descendants AS (
SELECT id, name, tier, COALESCE(runtime, 'claude-code') AS runtime FROM workspaces WHERE parent_id = $1 AND status = 'paused'
UNION ALL
SELECT w.id, w.name, w.tier, COALESCE(w.runtime, 'claude-code') FROM workspaces w JOIN descendants d ON w.parent_id = d.id WHERE w.status = 'paused'
) SELECT id, name, tier, runtime FROM descendants`, id)
if err != nil {
log.Printf("Resume: descendant query failed for %s: %v", id, err)
}
if rows != nil {
defer rows.Close()
for rows.Next() {
@@ -1007,10 +959,8 @@ func (h *WorkspaceHandler) Resume(c *gin.Context) {
// Re-provision all
for _, ws := range toResume {
if _, err := db.DB.ExecContext(ctx,
`UPDATE workspaces SET status = $1, updated_at = now() WHERE id = $2`, models.StatusProvisioning, ws.id); err != nil {
log.Printf("Resume: failed to set provisioning status for %s: %v", ws.id, err)
}
db.DB.ExecContext(ctx,
`UPDATE workspaces SET status = $1, updated_at = now() WHERE id = $2`, models.StatusProvisioning, ws.id)
h.broadcaster.RecordAndBroadcast(ctx, string(events.EventWorkspaceProvisioning), ws.id, map[string]interface{}{
"name": ws.name, "tier": ws.tier, "runtime": ws.runtime,
})
@@ -248,13 +248,8 @@ func TestRestart_CPStopOnlyInsideRetryHelper(t *testing.T) {
if !ok || fn.Body == nil || fn.Recv == nil {
continue
}
// cpStopWithRetryErr is the ONE allowed home for h.cpProv.Stop
// the bounded-retry loop. cpStopWithRetry is the void-returning
// wrapper (restart path) that delegates to it; the delete path uses
// cpStopWithRetryErr directly via stopWorkspaceForDelete to capture
// the terminal error (task #15). Both wrappers are exempt from this
// gate; any OTHER direct cpProv.Stop is the silent-leak regression.
if fn.Name.Name == "cpStopWithRetry" || fn.Name.Name == "cpStopWithRetryErr" {
// cpStopWithRetry is the ONE allowed home for h.cpProv.Stop.
if fn.Name.Name == "cpStopWithRetry" {
continue
}
ast.Inspect(fn.Body, func(n ast.Node) bool {
@@ -501,10 +501,6 @@ func TestWorkspaceCreate_WithSecrets_Persists(t *testing.T) {
// while persisting a secret causes the entire transaction to roll back and
// the handler to return 500. The workspace row must NOT be committed.
func TestWorkspaceCreate_SecretPersistFails_RollsBack(t *testing.T) {
// internal#691: see TestExtended_SecretsSet — same default-closed reasoning.
// This test is asserting the rollback path on DB failure, not the strip gate;
// keep the org in byok so the OPENAI_API_KEY write reaches the INSERT.
t.Setenv("MOLECULE_LLM_BILLING_MODE", "byok")
mock := setupTestDB(t)
setupTestRedis(t)
broadcaster := newTestBroadcaster()
@@ -513,14 +509,6 @@ func TestWorkspaceCreate_SecretPersistFails_RollsBack(t *testing.T) {
mock.ExpectBegin()
mock.ExpectExec("INSERT INTO workspaces").
WillReturnResult(sqlmock.NewResult(0, 1))
// internal#691: Create() now resolves billing mode per-workspace before
// the secret-strip gate. The workspace row was just inserted in the same
// transaction so it isn't readable from a separate query yet; the
// resolver expects the SELECT and the mock returns no row → falls back
// to the org default (byok, set above) so the OPENAI_API_KEY write
// reaches the INSERT-and-fail path this test exercises.
mock.ExpectQuery(`SELECT llm_billing_mode FROM workspaces WHERE id = \$1`).
WillReturnRows(sqlmock.NewRows([]string{"llm_billing_mode"}))
mock.ExpectExec("INSERT INTO workspace_secrets").
WillReturnError(sql.ErrConnDone) // DB failure while writing secret
mock.ExpectRollback() // workspace insert must be rolled back
@@ -15,7 +15,6 @@ import (
"context"
"database/sql"
"encoding/json"
"log"
"path"
"strings"
"time"
@@ -392,9 +391,7 @@ func extractFilesFromResponse(body json.RawMessage) []ChatAttachment {
var probe struct {
Result json.RawMessage `json:"result"`
}
if err := json.Unmarshal(body, &probe); err != nil {
log.Printf("messagestore: unmarshal probe body: %v", err)
}
_ = json.Unmarshal(body, &probe)
feed := body
if len(probe.Result) > 0 {
trimmed := bytesTrimSpace(probe.Result)
@@ -110,15 +110,10 @@ func WorkspaceAuth(database *sql.DB) gin.HandlerFunc {
c.Next()
return
}
// SaaS-canvas path: a browser cookie is acceptable only after the
// control plane confirms membership in this tenant. Referer/Origin
// are forgeable and must never authenticate workspace data routes.
if cookieHeader := c.GetHeader("Cookie"); cookieHeader != "" {
if ok, _ := VerifiedCPSession(cookieHeader); ok {
c.Set("cp_session_actor", cpSessionActor(cookieHeader))
c.Next()
return
}
// Same-origin canvas on tenant image — Referer matches Host.
if isSameOriginCanvas(c) {
c.Next()
return
}
// Local-dev escape hatch — see devmode.go. Unreachable on SaaS
// (hosted tenants always have ADMIN_TOKEN + MOLECULE_ENV=production).
@@ -412,40 +407,3 @@ func isSameOriginCanvas(c *gin.Context) bool {
origin := c.GetHeader("Origin")
return origin == "https://"+host || origin == "http://"+host
}
// cpSessionConfigured reports whether this platform is wired for upstream
// session-cookie verification — i.e. it runs as a SaaS tenant image with
// both CP_UPSTREAM_URL and MOLECULE_ORG_SLUG set. When false (self-hosted /
// dev), VerifiedCPSession can never succeed, so callers that want a
// non-forgeable canvas signal in SaaS while still working in dev can use
// this to decide whether the forgeable same-origin fallback is acceptable.
func cpSessionConfigured() bool {
return os.Getenv("CP_UPSTREAM_URL") != "" && tenantSlug() != ""
}
// CPSessionConfigured is the exported form of cpSessionConfigured for callers
// outside this package (e.g. the A2A proxy's canvas-user classification).
func CPSessionConfigured() bool {
return cpSessionConfigured()
}
// IsVerifiedCanvasSession returns true ONLY when the request carries a WorkOS
// session cookie that the control plane confirms belongs to a member of THIS
// tenant's org (via /cp/auth/tenant-member). Unlike IsSameOriginCanvas — whose
// Host/Referer/Origin inputs are trivially forgeable by any container on the
// Docker network and which is therefore documented as cosmetic-only (see
// AdminAuth / CanvasOrBearer comments above, #623/#194) — this is a real,
// upstream-verified authentication boundary. It is the correct gate for
// non-cosmetic actions such as A2A dispatch on behalf of a canvas user.
//
// Returns false (no network call) in self-hosted / dev deployments where
// CP_UPSTREAM_URL / MOLECULE_ORG_SLUG are unset; callers should treat that as
// "no verified canvas session available" and fall back accordingly.
func IsVerifiedCanvasSession(c *gin.Context) bool {
cookie := c.GetHeader("Cookie")
if cookie == "" {
return false
}
valid, _ := VerifiedCPSession(cookie)
return valid
}
@@ -75,90 +75,6 @@ func TestWorkspaceAuth_351_NoBearer_Returns401_NoDBCalls(t *testing.T) {
}
}
// TestWorkspaceAuth_ForgedSameOriginHeaders_Returns401 pins the production
// boundary for combined tenant images: Referer/Origin are forgeable request
// headers and must not authenticate workspace data routes.
func TestWorkspaceAuth_ForgedSameOriginHeaders_Returns401(t *testing.T) {
t.Setenv("MOLECULE_ENV", "production")
t.Setenv("ADMIN_TOKEN", "admin-secret")
prev := canvasProxyActive
canvasProxyActive = true
defer func() { canvasProxyActive = prev }()
mockDB, mock, err := sqlmock.New()
if err != nil {
t.Fatalf("sqlmock.New: %v", err)
}
defer mockDB.Close()
r := gin.New()
r.GET("/workspaces/:id/secrets", WorkspaceAuth(mockDB), func(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"ok": true})
})
r.DELETE("/workspaces/:id/secrets/:key", WorkspaceAuth(mockDB), func(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"ok": true})
})
for _, tt := range []struct {
name string
method string
path string
}{
{"list secrets", http.MethodGet, "/workspaces/ws-forged/secrets"},
{"delete secret", http.MethodDelete, "/workspaces/ws-forged/secrets/HERMES_CUSTOM_API_KEY"},
} {
t.Run(tt.name, func(t *testing.T) {
w := httptest.NewRecorder()
req, _ := http.NewRequest(tt.method, tt.path, nil)
req.Host = "tenant.example.test"
req.Header.Set("Referer", "https://tenant.example.test/")
req.Header.Set("Origin", "https://tenant.example.test")
r.ServeHTTP(w, req)
if w.Code != http.StatusUnauthorized {
t.Fatalf("forged same-origin headers: expected 401, got %d: %s", w.Code, w.Body.String())
}
})
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Errorf("unmet sqlmock expectations: %v", err)
}
}
func TestWorkspaceAuth_VerifiedTenantSessionCookie_AllowsCanvas(t *testing.T) {
resetSessionCache()
t.Setenv("MOLECULE_ENV", "production")
t.Setenv("ADMIN_TOKEN", "admin-secret")
t.Setenv("MOLECULE_ORG_SLUG", "tenant-a")
srv, _ := mockCPServer(t, http.StatusOK, `{"member":true,"user_id":"u_1","role":"owner","org_id":"org_1"}`)
t.Setenv("CP_UPSTREAM_URL", srv.URL)
mockDB, mock, err := sqlmock.New()
if err != nil {
t.Fatalf("sqlmock.New: %v", err)
}
defer mockDB.Close()
r := gin.New()
r.GET("/workspaces/:id/secrets", WorkspaceAuth(mockDB), func(c *gin.Context) {
if _, ok := c.Get("cp_session_actor"); !ok {
t.Errorf("cp_session_actor was not set")
}
c.JSON(http.StatusOK, gin.H{"ok": true})
})
w := httptest.NewRecorder()
req, _ := http.NewRequest(http.MethodGet, "/workspaces/ws-session/secrets", nil)
req.Header.Set("Cookie", "molecule_session=valid")
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("verified tenant session: expected 200, got %d: %s", w.Code, w.Body.String())
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Errorf("unmet sqlmock expectations: %v", err)
}
}
// TestWorkspaceAuth_C4_C8_NoBearer_Returns401 — C4/C8 critical path:
// when a workspace has live tokens and the caller sends NO bearer token,
// the middleware must return 401. This was the confirmed attack vector —
+3 -9
View File
@@ -24,7 +24,6 @@ import (
"encoding/base64"
"errors"
"fmt"
"log"
"time"
)
@@ -131,10 +130,8 @@ func Validate(ctx context.Context, db *sql.DB, plaintext string) (id, prefix, or
// Best-effort last_used_at bump. Failure here is acceptable — the
// request is already authenticated; we don't want a transient DB
// blip to flip a 200 into a 500.
if _, err := db.ExecContext(ctx,
`UPDATE org_api_tokens SET last_used_at = now() WHERE id = $1`, id); err != nil {
log.Printf("orgtoken: last_used_at bump failed for %s: %v", id, err)
}
_, _ = db.ExecContext(ctx,
`UPDATE org_api_tokens SET last_used_at = now() WHERE id = $1`, id)
return id, prefix, orgID, nil
}
@@ -195,10 +192,7 @@ func Revoke(ctx context.Context, db *sql.DB, id string) (bool, error) {
if err != nil {
return false, fmt.Errorf("orgtoken: revoke: %w", err)
}
n, err := res.RowsAffected()
if err != nil {
return false, fmt.Errorf("orgtoken: revoke RowsAffected: %w", err)
}
n, _ := res.RowsAffected()
return n > 0, nil
}
@@ -202,9 +202,7 @@ func (p *CPProvisioner) Start(ctx context.Context, cfg WorkspaceConfig) (string,
// - Rejects symlinks at the template root (prevents bypass via symlink traversal)
// - Skips symlinks during WalkDir (prevents /etc/passwd etc. inclusion)
// - Validates all paths are relative and non-escaping
// - Caps total size at cpConfigFilesMaxBytes (a transport-DoS guard,
// not the retired 12 KiB user-data ceiling — config now ships off
// user-data via the CP's Secrets-Manager seeding path)
// - Caps total size at 12 KiB to prevent payload bloat
configFiles, err := collectCPConfigFiles(cfg)
if err != nil {
return "", fmt.Errorf("cp provisioner: collect config files: %w", err)
@@ -279,27 +277,7 @@ func (p *CPProvisioner) Start(ctx context.Context, cfg WorkspaceConfig) (string,
return result.InstanceID, nil
}
// cpConfigFilesMaxBytes bounds the aggregate config bundle this tenant
// ships to the control plane. It is a transport-DoS guard, NOT the old
// EC2-user-data ceiling.
//
// History: this was 12 KiB (12<<10) because the CP embedded the bundle in
// EC2 user-data, which AWS caps at 16 KiB (the cap left ~4 KiB for bootstrap
// overhead). That ceiling failed real customers — the jrs-auto SEO Agent's
// config (long SEO system prompt + SERVICES_REPO_WEBSITE + a 12-schedule
// block baked into config.yaml) exceeds 12 KiB, so Start() rejected it
// client-side with "config files exceed 12288 bytes" and the workspace
// could never provision.
//
// Config delivery now goes OFF user-data: the CP stages the bundle to AWS
// Secrets Manager (molecule/workspace/<id>/config) at provision time and the
// workspace fetches it into /configs at boot (mirrors the proven tenant
// bootstrap-secrets pattern). The bundle travels here only inside the JSON
// HTTP request body to the CP, which has no 16 KiB limit. The remaining
// bound exists purely so a buggy/hostile tenant can't stream an unbounded
// body and OOM the CP provision path — set generous (256 KiB) so legitimate
// growth (more schedules, longer prompts, more skills) never re-hits a wall.
const cpConfigFilesMaxBytes = 256 << 10
const cpConfigFilesMaxBytes = 12 << 10
// isCPTemplateConfigFile restricts which files from a template directory are
// eligible for transport to the control plane. Only config.yaml (the runtime
@@ -420,10 +398,7 @@ func (p *CPProvisioner) Stop(ctx context.Context, workspaceID string) error {
return ErrNoBackend
}
url := fmt.Sprintf("%s/cp/workspaces/%s?instance_id=%s", p.baseURL, workspaceID, instanceID)
req, err := http.NewRequestWithContext(ctx, "DELETE", url, nil)
if err != nil {
return fmt.Errorf("cp provisioner: stop: build request: %w", err)
}
req, _ := http.NewRequestWithContext(ctx, "DELETE", url, nil)
p.provisionAuthHeaders(req)
resp, err := p.httpClient.Do(req)
if err != nil {
@@ -538,10 +513,7 @@ func (p *CPProvisioner) IsRunning(ctx context.Context, workspaceID string) (bool
return false, ErrNoBackend
}
url := fmt.Sprintf("%s/cp/workspaces/%s/status?instance_id=%s", p.baseURL, workspaceID, instanceID)
req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
if err != nil {
return true, fmt.Errorf("cp provisioner: status: build request: %w", err)
}
req, _ := http.NewRequestWithContext(ctx, "GET", url, nil)
p.provisionAuthHeaders(req)
resp, err := p.httpClient.Do(req)
if err != nil {
@@ -575,10 +547,7 @@ func (p *CPProvisioner) IsRunning(ctx context.Context, workspaceID string) (bool
// to render to the user.
func (p *CPProvisioner) GetConsoleOutput(ctx context.Context, workspaceID string) (string, error) {
url := fmt.Sprintf("%s/cp/admin/workspaces/%s/console", p.baseURL, workspaceID)
req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
if err != nil {
return "", fmt.Errorf("cp provisioner: console: build request: %w", err)
}
req, _ := http.NewRequestWithContext(ctx, "GET", url, nil)
p.adminAuthHeaders(req)
resp, err := p.httpClient.Do(req)
if err != nil {
@@ -1,151 +0,0 @@
package provisioner
import (
"context"
"encoding/base64"
"encoding/json"
"io"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"strings"
"testing"
)
// TestStart_OversizedConfigBundleProvisions is the Prove-It reproduction for
// the jrs-auto SEO Agent provisioning failure:
//
// CPProvisioner: workspace start failed: cp provisioner: collect config
// files: config files exceed 12288 bytes
//
// Root cause: collectCPConfigFiles hard-capped the *eligible* config bundle
// (config.yaml + prompts/*) at 12 KiB because the controlplane embedded it in
// EC2 user-data (16 KiB AWS ceiling bootstrap overhead). The SEO agent's
// config (long SEO system prompt + SERVICES_REPO_WEBSITE + the 12-schedule
// block baked into config.yaml) exceeds 12 KiB, so Start() failed before it
// ever reached the wire — blocking a paying customer from provisioning.
//
// After moving config delivery OFF user-data and onto the persistent
// secondary volume (CP stages the bundle to Secrets Manager; the workspace
// fetches it at boot into /configs), the 12 KiB ceiling is obsolete: the
// bundle travels in the JSON HTTP body to CP, which has no 16 KiB limit. This
// test pins that a realistically-oversized (>12288 B) config bundle now
// reaches the CP request body intact instead of being rejected client-side.
func TestStart_OversizedConfigBundleProvisions(t *testing.T) {
// SEO-sized config.yaml: a 12-schedule block + SERVICES_REPO_WEBSITE +
// a long system prompt, comfortably over the retired 12 KiB cap.
var sb strings.Builder
sb.WriteString("name: jrs-auto-seo\nruntime: claude-code\n")
sb.WriteString("env:\n SERVICES_REPO_WEBSITE: https://example.com/jrs-auto/website-repo\n")
sb.WriteString("schedules:\n")
for i := 0; i < 12; i++ {
sb.WriteString(" - id: seo-task-")
sb.WriteString(strings.Repeat("x", 8))
sb.WriteString("\n cron: \"0 */2 * * *\"\n prompt: |\n")
sb.WriteString(" Run the SEO audit pass, refresh keyword rankings, regenerate the\n")
sb.WriteString(" sitemap, and publish the digest to the marketing channel.\n")
}
configYAML := sb.String()
seoPrompt := strings.Repeat(
"You are an expert SEO agent. Audit pages, find ranking gaps, and act. ", 200)
cfg := map[string][]byte{
"config.yaml": []byte(configYAML),
"prompts/system.md": []byte(seoPrompt),
}
total := len(configYAML) + len(seoPrompt)
if total <= 12<<10 {
t.Fatalf("fixture not representative: bundle is %d bytes, must exceed 12288 to reproduce the failure", total)
}
t.Logf("oversized config bundle: %d bytes (> old 12288 cap)", total)
var body cpProvisionRequest
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
t.Errorf("decode request: %v", err)
}
w.WriteHeader(http.StatusCreated)
_, _ = io.WriteString(w, `{"instance_id":"i-seo","state":"pending"}`)
}))
defer srv.Close()
p := &CPProvisioner{baseURL: srv.URL, orgID: "org-seo", httpClient: srv.Client()}
_, err := p.Start(context.Background(), WorkspaceConfig{
WorkspaceID: "ws-seo",
Runtime: "claude-code",
Tier: 4,
PlatformURL: "http://tenant",
ConfigFiles: cfg,
})
if err != nil {
t.Fatalf("Start with oversized config bundle failed: %v — the 12288-byte cap must be gone now config delivery is off user-data", err)
}
// The full bundle must have reached the CP request body intact.
wantCfg := base64.StdEncoding.EncodeToString([]byte(configYAML))
if got := body.ConfigFiles["config.yaml"]; got != wantCfg {
t.Errorf("config.yaml not delivered intact to CP (len got=%d want=%d)", len(got), len(wantCfg))
}
wantPrompt := base64.StdEncoding.EncodeToString([]byte(seoPrompt))
if got := body.ConfigFiles["prompts/system.md"]; got != wantPrompt {
t.Errorf("prompts/system.md not delivered intact to CP (len got=%d want=%d)", len(got), len(wantPrompt))
}
}
// TestCollectCPConfigFiles_DoSGuardStillBounds pins that retiring the 12 KiB
// cap did NOT remove the bound entirely — an absurdly large bundle (a buggy
// or hostile tenant) is still rejected so a compromised workspace-server
// can't OOM the CP request path. The guard just moved from a 12 KiB
// user-data ceiling to a generous transport-DoS ceiling.
func TestCollectCPConfigFiles_DoSGuardStillBounds(t *testing.T) {
huge := make([]byte, cpConfigFilesMaxBytes+1)
for i := range huge {
huge[i] = 'a'
}
_, err := collectCPConfigFiles(WorkspaceConfig{
ConfigFiles: map[string][]byte{"config.yaml": huge},
})
if err == nil {
t.Fatalf("expected the DoS guard to reject a %d-byte bundle, got nil", len(huge))
}
if !strings.Contains(err.Error(), "config files exceed") {
t.Errorf("unexpected error %q, want the size-guard message", err.Error())
}
}
// TestCollectCPConfigFiles_AcceptsSEOSizedBundle is the unit-level companion:
// collectCPConfigFiles itself (not just Start) must accept the SEO-sized
// bundle. Guards the exact constant that caused the outage.
func TestCollectCPConfigFiles_AcceptsSEOSizedBundle(t *testing.T) {
// 30 KiB of eligible config — far over the retired 12288 cap, far under
// the new DoS guard.
cfgBlob := make([]byte, 18<<10)
for i := range cfgBlob {
cfgBlob[i] = 'c'
}
promptBlob := make([]byte, 12<<10)
for i := range promptBlob {
promptBlob[i] = 'p'
}
files, err := collectCPConfigFiles(WorkspaceConfig{
ConfigFiles: map[string][]byte{
"config.yaml": cfgBlob,
"prompts/system.md": promptBlob,
},
})
if err != nil {
t.Fatalf("collectCPConfigFiles rejected a %d-byte SEO-sized bundle: %v", len(cfgBlob)+len(promptBlob), err)
}
if len(files) != 2 {
t.Fatalf("expected 2 files collected, got %d", len(files))
}
// Also confirm a template-dir path stays size-bounded the same way.
tmpl := t.TempDir()
if err := os.WriteFile(filepath.Join(tmpl, "config.yaml"), cfgBlob, 0o600); err != nil {
t.Fatal(err)
}
if _, err := collectCPConfigFiles(WorkspaceConfig{TemplatePath: tmpl}); err != nil {
t.Fatalf("collectCPConfigFiles rejected an SEO-sized template config.yaml: %v", err)
}
}
@@ -190,11 +190,7 @@ func sweepStuckProvisioning(ctx context.Context, emitter ProvisionTimeoutEmitter
log.Printf("Provision-timeout sweep: failed to flip %s to failed: %v", c.id, err)
continue
}
affected, err := res.RowsAffected()
if err != nil {
log.Printf("Provision-timeout sweep: RowsAffected error for %s: %v", c.id, err)
continue
}
affected, _ := res.RowsAffected()
if affected == 0 {
// Raced with restart / register — no harm, just skip.
continue
+3 -7
View File
@@ -173,12 +173,6 @@ func Setup(hub *ws.Hub, broadcaster *events.Broadcaster, prov *provisioner.Provi
// so the canvas flips to failed in seconds instead of waiting
// for the 10-minute provision-timeout sweeper.
wsAdmin.POST("/admin/workspaces/:id/bootstrap-failed", wh.BootstrapFailed)
// Per-workspace LLM billing mode override (internal#691). Used by
// CP's /cp/admin/workspaces/:id/llm-billing-mode proxy + (via that
// proxy) by the canvas Config-tab "LLM Billing" section. Default-
// closed resolver lives in handlers/llm_billing_mode.go.
wsAdmin.GET("/admin/workspaces/:id/llm-billing-mode", handlers.GetWorkspaceLLMBillingMode)
wsAdmin.PUT("/admin/workspaces/:id/llm-billing-mode", handlers.PutWorkspaceLLMBillingMode)
// Proxy to CP's serial-console endpoint so the canvas's "View
// Logs" button can render the actual boot trace without handing
// the tenant AWS credentials. Admin-gated because console output
@@ -223,7 +217,9 @@ func Setup(hub *ws.Hub, broadcaster *events.Broadcaster, prov *provisioner.Provi
{
// #680: PATCH /workspaces/:id moved under WorkspaceAuth (#680 IDOR fix).
// WorkspaceAuth enforces that the caller holds a valid bearer token for
// this specific workspace, or a control-plane-verified tenant session.
// this specific workspace — both auth AND ownership in one check. Cosmetic
// updates (x/y drag-reposition, inline rename) from the combined tenant
// image canvas still pass via the isSameOriginCanvas bypass in WorkspaceAuth.
wsAuth.PATCH("", wh.Update)
// Lifecycle
@@ -406,7 +406,7 @@ func (s *Scheduler) fireSchedule(ctx context.Context, sched scheduleRow) {
msgID := fmt.Sprintf("cron-%s-%s", short(sched.ID, 8), uuid.New().String()[:8])
a2aBody, marshalErr := json.Marshal(map[string]interface{}{
a2aBody, _ := json.Marshal(map[string]interface{}{
"method": "message/send",
"params": map[string]interface{}{
"message": map[string]interface{}{
@@ -416,10 +416,6 @@ func (s *Scheduler) fireSchedule(ctx context.Context, sched scheduleRow) {
},
},
})
if marshalErr != nil {
log.Printf("Scheduler '%s': json.Marshal a2aBody failed: %v", sched.Name, marshalErr)
return
}
log.Printf("Scheduler: firing '%s' → workspace %s", sched.Name, short(sched.WorkspaceID, 12))
@@ -494,13 +490,11 @@ func (s *Scheduler) fireSchedule(ctx context.Context, sched scheduleRow) {
} else if lastStatus == "ok" {
// Non-empty success — reset the counter
resetCtx, resetCancel := context.WithTimeout(context.Background(), dbQueryTimeout)
if _, err := db.DB.ExecContext(resetCtx, `
_, _ = db.DB.ExecContext(resetCtx, `
UPDATE workspace_schedules
SET consecutive_empty_runs = 0,
updated_at = now()
WHERE id = $1`, sched.ID); err != nil {
log.Printf("Scheduler: '%s' empty-run reset failed: %v", sched.Name, err)
}
WHERE id = $1`, sched.ID)
resetCancel()
}
@@ -531,11 +525,9 @@ func (s *Scheduler) fireSchedule(ctx context.Context, sched scheduleRow) {
log.Printf("Scheduler: '%s' AUTO-DISABLING after %d consecutive SDK errors (workspace %s)",
sched.Name, consecSDK, short(sched.WorkspaceID, 12))
autoDisableCtx, autoDisableCancel := context.WithTimeout(context.Background(), dbQueryTimeout)
if _, err := db.DB.ExecContext(autoDisableCtx, `
_, _ = db.DB.ExecContext(autoDisableCtx, `
UPDATE workspace_schedules SET enabled = false, updated_at = now() WHERE id = $1 AND enabled = true`,
sched.ID); err != nil {
log.Printf("Scheduler: '%s' auto-disable failed: %v", sched.Name, err)
}
sched.ID)
autoDisableCancel()
}
} else {
@@ -545,13 +537,11 @@ func (s *Scheduler) fireSchedule(ctx context.Context, sched scheduleRow) {
// and we should clear the streak.
if lastStatus == "ok" {
resetCtx, resetCancel := context.WithTimeout(context.Background(), dbQueryTimeout)
if _, err := db.DB.ExecContext(resetCtx, `
_, _ = db.DB.ExecContext(resetCtx, `
UPDATE workspace_schedules
SET consecutive_sdk_errors = 0,
updated_at = now()
WHERE id = $1`, sched.ID); err != nil {
log.Printf("Scheduler: '%s' SDK-error reset failed: %v", sched.Name, err)
}
WHERE id = $1`, sched.ID)
resetCancel()
}
}
@@ -596,32 +586,28 @@ func (s *Scheduler) fireSchedule(ctx context.Context, sched scheduleRow) {
// #2026: sanitize the truncated prompt — even UTF-8-safe truncate() can
// carry pre-existing invalid bytes from an agent-edited template. jsonb
// columns reject invalid UTF-8 and hold the transaction open.
cronMeta, marshalErr := json.Marshal(map[string]interface{}{
cronMeta, _ := json.Marshal(map[string]interface{}{
"schedule_id": sched.ID,
"schedule_name": sched.Name,
"cron_expr": sched.CronExpr,
"prompt": sanitizeUTF8(textutil.TruncateBytes(sched.Prompt, 200)),
})
if marshalErr != nil {
log.Printf("Scheduler '%s': json.Marshal cronMeta failed: %v", sched.Name, marshalErr)
} else {
// #152: persist lastError into error_detail on the activity_logs row
// so GET /workspaces/:id/schedules/:id/history can surface why a run
// failed (previously dropped — history returned status without any
// error context, making root-cause debugging impossible).
// #2026: bounded Background() context — this INSERT was observed wedging
// indefinitely on invalid-UTF-8 jsonb payloads, blocking wg.Wait() in
// tick() and stalling the whole scheduler. Now: 10s deadline, survives
// outer ctx cancellation, and every string is UTF-8 sanitized.
insertCtx, insertCancel := context.WithTimeout(context.Background(), dbQueryTimeout)
if _, insErr := db.DB.ExecContext(insertCtx, `
INSERT INTO activity_logs (workspace_id, activity_type, source_id, method, summary, request_body, status, error_detail, created_at)
VALUES ($1, 'cron_run', NULL, 'cron', $2, $3::jsonb, $4, $5, now())
`, sched.WorkspaceID, sanitizeUTF8("Cron: "+sched.Name), string(cronMeta), lastStatus, sanitizeUTF8(lastError)); insErr != nil {
log.Printf("Scheduler: activity_logs insert failed for '%s' (%s): %v", sched.Name, sched.ID, insErr)
}
insertCancel()
// #152: persist lastError into error_detail on the activity_logs row
// so GET /workspaces/:id/schedules/:id/history can surface why a run
// failed (previously dropped — history returned status without any
// error context, making root-cause debugging impossible).
// #2026: bounded Background() context — this INSERT was observed wedging
// indefinitely on invalid-UTF-8 jsonb payloads, blocking wg.Wait() in
// tick() and stalling the whole scheduler. Now: 10s deadline, survives
// outer ctx cancellation, and every string is UTF-8 sanitized.
insertCtx, insertCancel := context.WithTimeout(context.Background(), dbQueryTimeout)
if _, insErr := db.DB.ExecContext(insertCtx, `
INSERT INTO activity_logs (workspace_id, activity_type, source_id, method, summary, request_body, status, error_detail, created_at)
VALUES ($1, 'cron_run', NULL, 'cron', $2, $3::jsonb, $4, $5, now())
`, sched.WorkspaceID, sanitizeUTF8("Cron: "+sched.Name), string(cronMeta), lastStatus, sanitizeUTF8(lastError)); insErr != nil {
log.Printf("Scheduler: activity_logs insert failed for '%s' (%s): %v", sched.Name, sched.ID, insErr)
}
insertCancel()
if s.broadcaster != nil {
s.broadcaster.RecordAndBroadcast(ctx, string(events.EventCronExecuted), sched.WorkspaceID, map[string]interface{}{
@@ -672,7 +658,7 @@ func (s *Scheduler) recordSkipped(ctx context.Context, sched scheduleRow, active
// #2026: bounded Background() context so the bookkeeping can't block
// on a stuck DB and stall the scheduler.
skipUpdCtx, skipUpdCancel := context.WithTimeout(context.Background(), dbQueryTimeout)
if _, err := db.DB.ExecContext(skipUpdCtx, `
_, _ = db.DB.ExecContext(skipUpdCtx, `
UPDATE workspace_schedules
SET last_run_at = now(),
next_run_at = COALESCE($2, next_run_at),
@@ -681,32 +667,24 @@ func (s *Scheduler) recordSkipped(ctx context.Context, sched scheduleRow, active
last_error = $3,
updated_at = now()
WHERE id = $1
`, sched.ID, nextRunPtr, sanitizeUTF8(reason)); err != nil {
log.Printf("Scheduler: '%s' skip update failed: %v", sched.Name, err)
}
`, sched.ID, nextRunPtr, sanitizeUTF8(reason))
skipUpdCancel()
cronMeta, marshalErr := json.Marshal(map[string]interface{}{
cronMeta, _ := json.Marshal(map[string]interface{}{
"schedule_id": sched.ID,
"schedule_name": sched.Name,
"cron_expr": sched.CronExpr,
"skipped": true,
"active_tasks": activeTasks,
})
if marshalErr != nil {
log.Printf("Scheduler '%s': json.Marshal cronMeta failed: %v", sched.Name, marshalErr)
} else {
// #2026: bounded Background() context on the skipped activity log INSERT
// for the same reason as the fireSchedule activity_logs INSERT above.
skipInsCtx, skipInsCancel := context.WithTimeout(context.Background(), dbQueryTimeout)
if _, err := db.DB.ExecContext(skipInsCtx, `
INSERT INTO activity_logs (workspace_id, activity_type, source_id, method, summary, request_body, status, error_detail, created_at)
VALUES ($1, 'cron_run', NULL, 'cron', $2, $3::jsonb, 'skipped', $4, now())
`, sched.WorkspaceID, sanitizeUTF8("Cron skipped: "+sched.Name), string(cronMeta), sanitizeUTF8(reason)); err != nil {
log.Printf("Scheduler: '%s' skip activity log failed: %v", sched.Name, err)
}
skipInsCancel()
}
// #2026: bounded Background() context on the skipped activity log INSERT
// for the same reason as the fireSchedule activity_logs INSERT above.
skipInsCtx, skipInsCancel := context.WithTimeout(context.Background(), dbQueryTimeout)
_, _ = db.DB.ExecContext(skipInsCtx, `
INSERT INTO activity_logs (workspace_id, activity_type, source_id, method, summary, request_body, status, error_detail, created_at)
VALUES ($1, 'cron_run', NULL, 'cron', $2, $3::jsonb, 'skipped', $4, now())
`, sched.WorkspaceID, sanitizeUTF8("Cron skipped: "+sched.Name), string(cronMeta), sanitizeUTF8(reason))
skipInsCancel()
if s.broadcaster != nil {
_ = s.broadcaster.RecordAndBroadcast(ctx, string(events.EventCronSkipped), sched.WorkspaceID, map[string]interface{}{
@@ -60,12 +60,10 @@ func RunWithRecover(ctx context.Context, name string, fn func(context.Context))
}
// Panic → back off and restart.
timer := time.NewTimer(backoff)
select {
case <-ctx.Done():
timer.Stop()
return
case <-timer.C:
case <-time.After(backoff):
}
if backoff < maxBackoff {
backoff *= 2
+4 -9
View File
@@ -19,7 +19,6 @@ import (
"encoding/base64"
"errors"
"fmt"
"log"
"strings"
)
@@ -125,10 +124,8 @@ func ValidateToken(ctx context.Context, db *sql.DB, expectedWorkspaceID, plainte
// Best-effort last_used_at update. A failure here (DB hiccup, etc.)
// must not cause an otherwise-valid request to 401.
if _, err := db.ExecContext(ctx,
`UPDATE workspace_auth_tokens SET last_used_at = now() WHERE id = $1`, tokenID); err != nil {
log.Printf("wsauth: last_used_at bump failed for %s: %v", tokenID, err)
}
_, _ = db.ExecContext(ctx,
`UPDATE workspace_auth_tokens SET last_used_at = now() WHERE id = $1`, tokenID)
return nil
}
@@ -253,9 +250,7 @@ func ValidateAnyToken(ctx context.Context, db *sql.DB, plaintext string) error {
}
// Best-effort last_used_at update.
if _, err := db.ExecContext(ctx,
`UPDATE workspace_auth_tokens SET last_used_at = now() WHERE id = $1`, tokenID); err != nil {
log.Printf("wsauth: last_used_at bump failed for %s: %v", tokenID, err)
}
_, _ = db.ExecContext(ctx,
`UPDATE workspace_auth_tokens SET last_used_at = now() WHERE id = $1`, tokenID)
return nil
}
@@ -1,4 +0,0 @@
-- Reverse internal#691 per-workspace billing mode column.
-- The column is nullable + check-constrained; dropping it is non-destructive
-- to org-level behavior (workspaces fall back to the org default again).
ALTER TABLE workspaces DROP COLUMN IF EXISTS llm_billing_mode;
@@ -1,17 +0,0 @@
-- Per-workspace llm_billing_mode override (internal#691).
--
-- NULL = inherit the org-level default (organizations.llm_billing_mode on CP,
-- propagated to workspace-server via tenant_config as MOLECULE_LLM_BILLING_MODE).
-- A non-NULL value overrides the org default for this workspace only.
--
-- Resolver contract: workspaces.llm_billing_mode ?? org_default ?? 'platform_managed'.
-- Default-closed: any NULL, error, unknown enum, or JOIN miss resolves to
-- 'platform_managed' (the existing implicit default — see internal#691
-- spec sketch + Phase 1 design comment).
--
-- The check constraint mirrors the CP-side credits.LLMBillingMode* constants
-- (molecule-controlplane/internal/credits/llm_billing.go). Keep in sync if
-- a new mode is ever added; the resolver also enumerates them explicitly.
ALTER TABLE workspaces
ADD COLUMN IF NOT EXISTS llm_billing_mode TEXT
CHECK (llm_billing_mode IS NULL OR llm_billing_mode IN ('platform_managed', 'byok', 'disabled'));