Merge pull request #88 from AgentWorkforce/architecture-adjustments

Architecture adjustments

Author: khaliqgant

deploy/workspace/entrypoint.sh

@@ -6,8 +6,13 @@ log() {
   echo "[workspace] $*"
 }
 
-# Drop to workspace user if running as root
+# Fix volume permissions and drop to workspace user if running as root
 if [[ "$(id -u)" == "0" ]]; then
+  # Fix /data and /workspace permissions before dropping privileges
+  # Fresh Fly.io volumes are root-owned, need to chown for workspace user
+  log "Fixing volume permissions..."
+  chown -R workspace:workspace /data /workspace 2>/dev/null || true
+
   log "Dropping privileges to workspace user..."
   exec gosu workspace "$0" "$@"
 fi
@@ -16,12 +21,54 @@ PORT="${AGENT_RELAY_DASHBOARD_PORT:-${PORT:-3888}}"
 export AGENT_RELAY_DASHBOARD_PORT="${PORT}"
 export PORT="${PORT}"
 
+# ============================================================================
+# Per-user credential storage setup
+# Create user-specific HOME on persistent volume (/data)
+# This enables multi-user workspaces where each user has their own credentials
+# ============================================================================
+DATA_DIR="${AGENT_RELAY_DATA_DIR:-/data}"
+if [[ -n "${WORKSPACE_OWNER_USER_ID:-}" ]]; then
+  USER_HOME="${DATA_DIR}/users/${WORKSPACE_OWNER_USER_ID}"
+  log "Setting up per-user HOME at ${USER_HOME}"
+  mkdir -p "${USER_HOME}"
+  mkdir -p "${USER_HOME}/.claude"
+  mkdir -p "${USER_HOME}/.codex"
+  mkdir -p "${USER_HOME}/.config/gcloud"
+  mkdir -p "${USER_HOME}/.config/gh"
+  export HOME="${USER_HOME}"
+  export XDG_CONFIG_HOME="${USER_HOME}/.config"
+  export AGENT_RELAY_USER_ID="${WORKSPACE_OWNER_USER_ID}"
+  log "HOME set to ${HOME} (user: ${WORKSPACE_OWNER_USER_ID})"
+else
+  log "No WORKSPACE_OWNER_USER_ID set, using default HOME: ${HOME}"
+fi
+
 WORKSPACE_DIR="${WORKSPACE_DIR:-/workspace}"
 REPO_LIST="${REPOSITORIES:-}"
 
 mkdir -p "${WORKSPACE_DIR}"
 cd "${WORKSPACE_DIR}"
 
+# ============================================================================
+# Persist workspace environment for SSH sessions
+# SSH sessions don't inherit the container's runtime environment, so we write
+# critical variables to /etc/profile.d/ which gets sourced by login shells
+# ============================================================================
+if [[ -n "${CLOUD_API_URL:-}" || -n "${WORKSPACE_ID:-}" ]]; then
+  log "Persisting workspace environment for SSH sessions"
+  cat > /etc/profile.d/workspace-env.sh <<ENVEOF
+# Workspace environment variables (auto-generated by entrypoint.sh)
+export WORKSPACE_ID="${WORKSPACE_ID:-}"
+export CLOUD_API_URL="${CLOUD_API_URL:-}"
+export WORKSPACE_TOKEN="${WORKSPACE_TOKEN:-}"
+export WORKSPACE_DIR="${WORKSPACE_DIR}"
+export HOME="${HOME}"
+export AGENT_RELAY_USER_ID="${AGENT_RELAY_USER_ID:-}"
+export AGENT_RELAY_DATA_DIR="${AGENT_RELAY_DATA_DIR:-/data}"
+ENVEOF
+  chmod 644 /etc/profile.d/workspace-env.sh
+fi
+
 # Configure Git credentials via the gateway (tokens auto-refresh via Nango)
 # The credential helper fetches fresh tokens from the cloud API on each git operation
 if [[ -n "${CLOUD_API_URL:-}" && -n "${WORKSPACE_ID:-}" && -n "${WORKSPACE_TOKEN:-}" ]]; then

docs/auth-revocation-handling.md

@@ -0,0 +1,146 @@
+# Auth Revocation Handling Design
+
+## Problem
+
+Claude (and other AI CLIs) have limited active OAuth sessions. When a user authenticates:
+1. **Via relay** → Can revoke their local Claude instance's auth
+2. **Locally** → Can revoke the relay workspace agent's auth
+
+Both scenarios need graceful handling.
+
+## Detection Patterns
+
+### Claude CLI Auth Revocation Indicators
+```
+- "Your session has expired"
+- "Please log in again"
+- "Authentication required"
+- "Unauthorized"
+- "session expired"
+- "invalid credentials"
+- API responses with 401/403
+- CLI exit with auth-related error message
+```
+
+## Implementation Plan
+
+### 1. Add Auth Error Detection to Parser/Wrapper
+
+**File: `src/wrapper/auth-detection.ts` (new)**
+```typescript
+export const AUTH_REVOCATION_PATTERNS = [
+  /session\s+(has\s+)?expired/i,
+  /please\s+log\s*in\s+again/i,
+  /authentication\s+required/i,
+  /unauthorized/i,
+  /invalid\s+credentials/i,
+  /not\s+authenticated/i,
+  /login\s+required/i,
+];
+
+export function detectAuthRevocation(output: string): boolean {
+  return AUTH_REVOCATION_PATTERNS.some(pattern => pattern.test(output));
+}
+```
+
+**File: `src/wrapper/tmux-wrapper.ts` (modify)**
+- Import auth detection
+- In output processing, check for auth revocation patterns
+- When detected, emit 'auth_revoked' event and update status
+
+### 2. Add Agent Status: `auth_revoked`
+
+**File: `src/cloud/db/schema.ts`**
+- Document that `status` can be: `active`, `idle`, `ended`, `auth_revoked`
+
+**File: `src/cloud/api/workspaces.ts`**
+- Add endpoint to update agent auth status
+- Add endpoint to trigger re-authentication
+
+### 3. Relay Protocol Extension
+
+**File: `src/protocol/types.ts`**
+Add new envelope type for auth status:
+```typescript
+export interface AuthStatusPayload {
+  agentName: string;
+  status: 'revoked' | 'valid';
+  provider: string; // 'claude', 'codex', etc.
+  message?: string;
+}
+```
+
+### 4. Dashboard UI
+
+**File: `src/dashboard/react-components/AgentCard.tsx`**
+- Show "Auth Required" badge when agent status is `auth_revoked`
+- Show "Re-authenticate" button
+
+**File: `src/dashboard/react-components/AuthRevocationNotification.tsx` (new)**
+- Toast/banner notification when auth is revoked
+- Explains what happened and how to fix
+
+**File: `src/dashboard/react-components/AuthWarningModal.tsx` (new)**
+- Warning before authenticating: "This may revoke other active sessions"
+- Checkbox: "Don't show again"
+- Continue / Cancel buttons
+
+### 5. Re-authentication Flow
+
+When user clicks "Re-authenticate":
+1. Opens the existing CLI auth flow (`/api/cli/:provider/start`)
+2. On success, agent status updated back to `active`
+3. Agent resumes operation (may need to restart or reconnect)
+
+## Files to Create/Modify
+
+### New Files
+- `src/wrapper/auth-detection.ts` - Auth error patterns and detection
+- `src/dashboard/react-components/AuthRevocationNotification.tsx`
+- `src/dashboard/react-components/AuthWarningModal.tsx`
+
+### Modified Files
+- `src/wrapper/tmux-wrapper.ts` - Add auth detection in output processing
+- `src/wrapper/base-wrapper.ts` - Add auth status events
+- `src/protocol/types.ts` - Add AUTH_STATUS envelope type
+- `src/daemon/router.ts` - Handle auth status messages
+- `src/cloud/api/workspaces.ts` - Add auth status endpoints
+- `src/dashboard/react-components/AgentCard.tsx` - Show auth status
+- `src/dashboard/react-components/App.tsx` - Handle auth notifications
+- `src/cloud/api/onboarding.ts` - Add pre-auth warning flag
+
+## API Endpoints
+
+### POST /api/workspaces/:id/agents/:agentName/reauth
+Triggers re-authentication for an agent with revoked auth.
+
+### GET /api/workspaces/:id/agents/:agentName/auth-status
+Returns current auth status for an agent.
+
+### POST /api/cli/:provider/start (existing, modify)
+Add `skipWarning` parameter to bypass the session limit warning.
+
+## Event Flow
+
+```
+1. Agent running in workspace
+2. User authenticates Claude locally
+3. Cloud auth is revoked
+4. Agent CLI outputs "session expired" or similar
+5. Wrapper detects pattern → emits AUTH_REVOKED event
+6. Daemon receives event → updates agent status
+7. Cloud DB updated → status = 'auth_revoked'
+8. Dashboard polls/receives status → shows notification
+9. User clicks "Re-authenticate"
+10. CLI auth flow starts
+11. On success → agent status = 'active'
+12. Agent resumes (or user restarts agent)
+```
+
+## Pre-Auth Warning
+
+Before starting any CLI auth:
+1. Check if this is a provider with session limits (Claude, etc.)
+2. Show modal: "Authenticating Claude here may sign out other sessions"
+3. User confirms → proceed with auth
+4. User can check "Don't warn me again" (stored in localStorage)

src/cloud/api/billing.ts

@@ -7,13 +7,66 @@
 import { Router, Request } from 'express';
 import { getBillingService, getAllPlans, getPlan, comparePlans } from '../billing/index.js';
 import type { SubscriptionTier } from '../billing/types.js';
-import { getConfig } from '../config.js';
-import { db } from '../db/index.js';
+import { getConfig, isAdminUser } from '../config.js';
+import { db, type PlanType } from '../db/index.js';
 import { requireAuth } from './auth.js';
+import { getProvisioner, RESOURCE_TIERS } from '../provisioner/index.js';
+import { getResourceTierForPlan } from '../services/planLimits.js';
 import type Stripe from 'stripe';
 
 export const billingRouter = Router();
 
+/**
+ * Resize user's workspaces to match their new plan tier
+ * Called after plan upgrade/downgrade to adjust compute resources
+ *
+ * Strategy:
+ * - Stopped workspaces: Resize immediately (no disruption)
+ * - Running workspaces: Save config for next restart (no agent disruption)
+ *
+ * User can manually restart to get new resources immediately, or wait for
+ * natural restart (auto-stop idle, manual restart, etc.)
+ */
+async function resizeWorkspacesForPlan(userId: string, newPlan: PlanType): Promise<void> {
+  try {
+    const workspaces = await db.workspaces.findByUserId(userId);
+    if (workspaces.length === 0) return;
+
+    const provisioner = getProvisioner();
+    const targetTierName = getResourceTierForPlan(newPlan);
+    const targetTier = RESOURCE_TIERS[targetTierName];
+
+    console.log(`[billing] Upgrading ${workspaces.length} workspace(s) for user ${userId.substring(0, 8)} to ${targetTierName}`);
+
+    for (const workspace of workspaces) {
+      if (workspace.status !== 'running' && workspace.status !== 'stopped') {
+        console.log(`[billing] Skipping workspace ${workspace.id.substring(0, 8)} (status: ${workspace.status})`);
+        continue;
+      }
+
+      try {
+        // For running workspaces: don't restart, apply on next restart
+        // This prevents disrupting active agents
+        const skipRestart = workspace.status === 'running';
+
+        await provisioner.resize(workspace.id, targetTier, skipRestart);
+
+        if (skipRestart) {
+          console.log(`[billing] Queued resize for workspace ${workspace.id.substring(0, 8)} to ${targetTierName} (will apply on next restart)`);
+          // TODO: Store pending upgrade in workspace metadata so we can show in UI
+        } else {
+          console.log(`[billing] Resized workspace ${workspace.id.substring(0, 8)} to ${targetTierName}`);
+        }
+      } catch (error) {
+        console.error(`[billing] Failed to resize workspace ${workspace.id}:`, error);
+        // Continue with other workspaces even if one fails
+      }
+    }
+  } catch (error) {
+    console.error('[billing] Failed to resize workspaces:', error);
+  }
+}
+
 /**
  * GET /api/billing/plans
  * Get all available billing plans
@@ -85,7 +138,6 @@ billingRouter.get('/compare', (req, res) => {
  */
 billingRouter.get('/subscription', requireAuth, async (req, res) => {
   const userId = req.session.userId!;
-  const billing = getBillingService();
 
   try {
     // Fetch user from database
@@ -94,6 +146,28 @@ billingRouter.get('/subscription', requireAuth, async (req, res) => {
       return res.status(404).json({ error: 'User not found' });
     }
 
+    // Admin users have special status - show their current plan without Stripe
+    if (isAdminUser(user.githubUsername)) {
+      return res.json({
+        tier: user.plan || 'enterprise',
+        subscription: null,
+        customer: null,
+        isAdmin: true,
+      });
+    }
+
+    // If user doesn't have a Stripe customer ID and is on free tier, skip Stripe calls entirely
+    // This prevents hanging on Stripe API calls for users who have never paid
+    if (!user.stripeCustomerId && user.plan === 'free') {
+      return res.json({
+        tier: 'free',
+        subscription: null,
+        customer: null,
+      });
+    }
+
+    const billing = getBillingService();
+
     // Get or create Stripe customer
     const customerId = user.stripeCustomerId ||
       await billing.getOrCreateCustomer(user.id, user.email || '', user.githubUsername);
@@ -150,7 +224,6 @@ billingRouter.post('/checkout', requireAuth, async (req, res) => {
     return;
   }
 
-  const billing = getBillingService();
   const config = getConfig();
 
   try {
@@ -160,6 +233,26 @@ billingRouter.post('/checkout', requireAuth, async (req, res) => {
       return res.status(404).json({ error: 'User not found' });
     }
 
+    // Admin users get free upgrades - skip Stripe entirely
+    if (isAdminUser(user.githubUsername)) {
+      // Update user plan directly
+      await db.users.update(userId, { plan: tier });
+      console.log(`[billing] Admin user ${user.githubUsername} upgraded to ${tier} (free)`);
+
+      // Resize workspaces to match new plan (async)
+      resizeWorkspacesForPlan(userId, tier as PlanType).catch((err) => {
+        console.error(`[billing] Failed to resize workspaces for admin ${user.githubUsername}:`, err);
+      });
+
+      // Return a fake session that redirects to success
+      return res.json({
+        sessionId: 'admin-upgrade',
+        checkoutUrl: `${config.publicUrl}/billing/success?admin=true`,
+      });
+    }
+
+    const billing = getBillingService();
+
     // Get or create customer
     const customerId = user.stripeCustomerId ||
       await billing.getOrCreateCustomer(user.id, user.email || '', user.githubUsername);
@@ -370,9 +463,9 @@ billingRouter.get('/invoices', requireAuth, async (req, res) => {
       return res.status(404).json({ error: 'User not found' });
     }
 
+    // No Stripe customer = no invoices, skip Stripe call entirely
     if (!user.stripeCustomerId) {
-      res.json({ invoices: [] });
-      return;
+      return res.json({ invoices: [] });
     }
 
     const billing = getBillingService();
@@ -466,11 +559,16 @@ billingRouter.post(
           // Extract subscription tier and update user's plan
           if (billingEvent.userId) {
             const subscription = billingEvent.data as unknown as Stripe.Subscription;
-            const tier = billing.getTierFromSubscription(subscription);
+            const tier = billing.getTierFromSubscription(subscription) as PlanType;
 
             // Update user's plan in database
             await db.users.update(billingEvent.userId, { plan: tier });
             console.log(`Updated user ${billingEvent.userId} plan to: ${tier}`);
+
+            // Resize workspaces to match new plan (async, don't block webhook)
+            resizeWorkspacesForPlan(billingEvent.userId, tier).catch((err) => {
+              console.error(`Failed to resize workspaces for user ${billingEvent.userId}:`, err);
+            });
           } else {
             console.warn('Subscription event received without userId:', billingEvent.id);
           }
@@ -482,6 +580,11 @@ billingRouter.post(
           if (billingEvent.userId) {
             await db.users.update(billingEvent.userId, { plan: 'free' });
             console.log(`User ${billingEvent.userId} subscription canceled, reset to free plan`);
+
+            // Resize workspaces down to free tier (async)
+            resizeWorkspacesForPlan(billingEvent.userId, 'free').catch((err) => {
+              console.error(`Failed to resize workspaces for user ${billingEvent.userId}:`, err);
+            });
           }
           break;
         }