From dc812742619788b5e6abab42013b53ace2c70442 Mon Sep 17 00:00:00 2001 From: Gage Krumbach Date: Fri, 19 Dec 2025 08:38:43 -0600 Subject: [PATCH 1/4] Enhance Google OAuth secret handling in agentic session events - Create a placeholder Google OAuth secret if it doesn't exist, ensuring the volume mount is always present for K8s to sync credentials post-OAuth. - Log creation of the placeholder secret and its mounting to the ambient-code-runner container. - Maintain existing functionality to mount the Google OAuth secret, allowing for updates once the backend populates the credentials. --- .../operator/internal/handlers/sessions.go | 82 +++++++++++++------ 1 file changed, 58 insertions(+), 24 deletions(-) diff --git a/components/operator/internal/handlers/sessions.go b/components/operator/internal/handlers/sessions.go index 862d8202..c67fae7f 100644 --- a/components/operator/internal/handlers/sessions.go +++ b/components/operator/internal/handlers/sessions.go @@ -1376,35 +1376,69 @@ func handleAgenticSessionEvent(obj *unstructured.Unstructured) error { } } - // Check for Google OAuth secret and mount it if present (for MCP Google Drive integration) + // Create placeholder Google OAuth secret if it doesn't exist (for MCP Google Workspace integration) + // This ensures the volume mount is always present so K8s can sync credentials after OAuth completion googleOAuthSecretName := fmt.Sprintf("%s-google-oauth", name) - if _, err := config.K8sClient.CoreV1().Secrets(sessionNamespace).Get(context.TODO(), googleOAuthSecretName, v1.GetOptions{}); err == nil { - log.Printf("Found Google OAuth secret %s, mounting to runner container", googleOAuthSecretName) - job.Spec.Template.Spec.Volumes = append(job.Spec.Template.Spec.Volumes, corev1.Volume{ - Name: "google-oauth", - VolumeSource: corev1.VolumeSource{ - Secret: &corev1.SecretVolumeSource{ - SecretName: googleOAuthSecretName, - Optional: boolPtr(true), // Don't fail if secret disappears before pod starts + if _, err := config.K8sClient.CoreV1().Secrets(sessionNamespace).Get(context.TODO(), googleOAuthSecretName, v1.GetOptions{}); errors.IsNotFound(err) { + // Create empty placeholder secret - backend will populate it after user completes OAuth + placeholderSecret := &corev1.Secret{ + ObjectMeta: v1.ObjectMeta{ + Name: googleOAuthSecretName, + Namespace: sessionNamespace, + Labels: map[string]string{ + "app": "ambient-code", + "ambient-code.io/session": name, + "ambient-code.io/provider": "google", + "ambient-code.io/oauth": "placeholder", + }, + OwnerReferences: []v1.OwnerReference{ + { + APIVersion: "vteam.ambient-code/v1", + Kind: "AgenticSession", + Name: currentObj.GetName(), + UID: currentObj.GetUID(), + Controller: boolPtr(true), + }, }, }, - }) - // Mount to the ambient-code-runner container - for i := range job.Spec.Template.Spec.Containers { - if job.Spec.Template.Spec.Containers[i].Name == "ambient-code-runner" { - job.Spec.Template.Spec.Containers[i].VolumeMounts = append(job.Spec.Template.Spec.Containers[i].VolumeMounts, corev1.VolumeMount{ - Name: "google-oauth", - MountPath: "/app/.google_workspace_mcp/credentials", - ReadOnly: true, - }) - log.Printf("Mounted Google OAuth secret to /app/.google_workspace_mcp/credentials in runner container for session %s", name) - break - } + Type: corev1.SecretTypeOpaque, + Data: map[string][]byte{ + "credentials.json": []byte(""), // Empty placeholder, runner checks for content + }, } - } else if !errors.IsNotFound(err) { - log.Printf("Error checking for Google OAuth secret %s: %v (continuing without MCP)", googleOAuthSecretName, err) + if _, createErr := config.K8sClient.CoreV1().Secrets(sessionNamespace).Create(context.TODO(), placeholderSecret, v1.CreateOptions{}); createErr != nil { + log.Printf("Warning: Failed to create placeholder Google OAuth secret %s: %v", googleOAuthSecretName, createErr) + } else { + log.Printf("Created placeholder Google OAuth secret %s (will be populated after user OAuth)", googleOAuthSecretName) + } + } else if err != nil { + log.Printf("Error checking for Google OAuth secret %s: %v", googleOAuthSecretName, err) } else { - log.Printf("No Google OAuth secret found (session %s), MCP Google Drive integration will not be available", name) + log.Printf("Found existing Google OAuth secret %s", googleOAuthSecretName) + } + + // Always mount Google OAuth secret (with Optional: true so pod starts even if empty) + // K8s will sync updates when backend populates credentials after OAuth completion (~60s) + job.Spec.Template.Spec.Volumes = append(job.Spec.Template.Spec.Volumes, corev1.Volume{ + Name: "google-oauth", + VolumeSource: corev1.VolumeSource{ + Secret: &corev1.SecretVolumeSource{ + SecretName: googleOAuthSecretName, + Optional: boolPtr(true), // Don't fail if secret is empty or missing + }, + }, + }) + // Mount to the ambient-code-runner container + for i := range job.Spec.Template.Spec.Containers { + if job.Spec.Template.Spec.Containers[i].Name == "ambient-code-runner" { + job.Spec.Template.Spec.Containers[i].VolumeMounts = append(job.Spec.Template.Spec.Containers[i].VolumeMounts, corev1.VolumeMount{ + Name: "google-oauth", + MountPath: "/app/.google_workspace_mcp/credentials", + ReadOnly: true, + }) + log.Printf("Mounted Google OAuth secret to /app/.google_workspace_mcp/credentials in runner container for session %s", name) + break + } } // Do not mount runner Secret volume; runner fetches tokens on demand From 0ad326329e6fc6b3fd2becb38d31f2419c2af9f0 Mon Sep 17 00:00:00 2001 From: Gage Krumbach Date: Fri, 19 Dec 2025 08:52:44 -0600 Subject: [PATCH 2/4] fix: Skip INITIAL_PROMPT on session resume using marker file - Use .initial_prompt_sent marker file instead of .claude/state - Create marker file after successfully sending initial prompt - More reliable detection of session resume vs fresh start --- components/runners/claude-code-runner/main.py | 24 +++++++++++++------ 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/components/runners/claude-code-runner/main.py b/components/runners/claude-code-runner/main.py index dee4e6c9..f7cbe576 100644 --- a/components/runners/claude-code-runner/main.py +++ b/components/runners/claude-code-runner/main.py @@ -7,6 +7,7 @@ import json import logging from contextlib import asynccontextmanager +from pathlib import Path from typing import Optional, List, Dict, Any, Union from fastapi import FastAPI, Request, HTTPException @@ -80,7 +81,6 @@ async def lifespan(app: FastAPI): # Import adapter here to avoid circular imports from adapter import ClaudeCodeAdapter - from pathlib import Path # Initialize context from environment session_id = os.getenv("SESSION_ID", "unknown") @@ -98,16 +98,17 @@ async def lifespan(app: FastAPI): logger.info("Adapter initialized - fresh client will be created for each run") - # Check if this is a restart (conversation history exists) - history_marker = Path(workspace_path) / ".claude" / "state" + # Check if this is a restart/resume (initial prompt was already sent) + # We use our own marker file since Claude SDK's internal state location varies + initial_prompt_marker = Path(workspace_path) / ".initial_prompt_sent" # Check for INITIAL_PROMPT and auto-execute (only if this is first run) initial_prompt = os.getenv("INITIAL_PROMPT", "").strip() - if initial_prompt and not history_marker.exists(): + if initial_prompt and not initial_prompt_marker.exists(): logger.info(f"INITIAL_PROMPT detected ({len(initial_prompt)} chars), will auto-execute after 3s delay") - asyncio.create_task(auto_execute_initial_prompt(initial_prompt, session_id)) + asyncio.create_task(auto_execute_initial_prompt(initial_prompt, session_id, initial_prompt_marker)) elif initial_prompt: - logger.info(f"INITIAL_PROMPT detected but conversation history exists - skipping auto-execution (session restart)") + logger.info(f"INITIAL_PROMPT detected but marker exists - skipping auto-execution (session resume)") logger.info(f"AG-UI server ready for session {session_id}") @@ -117,11 +118,13 @@ async def lifespan(app: FastAPI): logger.info("Shutting down AG-UI server...") -async def auto_execute_initial_prompt(prompt: str, session_id: str): +async def auto_execute_initial_prompt(prompt: str, session_id: str, marker_path: Path): """Auto-execute INITIAL_PROMPT by POSTing to backend after short delay. The 3-second delay gives the runner time to fully start. Backend has retry logic to handle if Service DNS isn't ready yet. + + Creates a marker file on success to prevent re-sending on session resume. """ import uuid import aiohttp @@ -171,6 +174,13 @@ async def auto_execute_initial_prompt(prompt: str, session_id: str): if resp.status == 200: result = await resp.json() logger.info(f"INITIAL_PROMPT auto-execution started: {result}") + # Create marker file to prevent re-sending on resume + try: + marker_path.parent.mkdir(parents=True, exist_ok=True) + marker_path.write_text(f"sent_at={asyncio.get_event_loop().time()}\n") + logger.info(f"Created initial prompt marker: {marker_path}") + except Exception as marker_err: + logger.warning(f"Failed to create initial prompt marker: {marker_err}") else: error_text = await resp.text() logger.warning(f"INITIAL_PROMPT failed with status {resp.status}: {error_text[:200]}") From 81159b4025de6c4329bc15681596a98eba491bb3 Mon Sep 17 00:00:00 2001 From: Gage Krumbach Date: Fri, 19 Dec 2025 09:03:56 -0600 Subject: [PATCH 3/4] fix: Also skip INITIAL_PROMPT for continuation sessions When PARENT_SESSION_ID is set, this is a child session continuing from a parent - skip the initial prompt since the conversation already has context from the parent session. --- components/runners/claude-code-runner/main.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/components/runners/claude-code-runner/main.py b/components/runners/claude-code-runner/main.py index f7cbe576..f52b4eb9 100644 --- a/components/runners/claude-code-runner/main.py +++ b/components/runners/claude-code-runner/main.py @@ -98,15 +98,19 @@ async def lifespan(app: FastAPI): logger.info("Adapter initialized - fresh client will be created for each run") - # Check if this is a restart/resume (initial prompt was already sent) - # We use our own marker file since Claude SDK's internal state location varies + # Check if this is a restart/resume or continuation + # - Marker file: session was already initialized (resume) + # - Parent session ID: this is a child session continuing from parent initial_prompt_marker = Path(workspace_path) / ".initial_prompt_sent" + parent_session_id = os.getenv("PARENT_SESSION_ID", "").strip() - # Check for INITIAL_PROMPT and auto-execute (only if this is first run) + # Check for INITIAL_PROMPT and auto-execute (only if this is a fresh start) initial_prompt = os.getenv("INITIAL_PROMPT", "").strip() - if initial_prompt and not initial_prompt_marker.exists(): + if initial_prompt and not initial_prompt_marker.exists() and not parent_session_id: logger.info(f"INITIAL_PROMPT detected ({len(initial_prompt)} chars), will auto-execute after 3s delay") asyncio.create_task(auto_execute_initial_prompt(initial_prompt, session_id, initial_prompt_marker)) + elif initial_prompt and parent_session_id: + logger.info(f"INITIAL_PROMPT detected but this is a continuation (parent={parent_session_id[:12]}...) - skipping") elif initial_prompt: logger.info(f"INITIAL_PROMPT detected but marker exists - skipping auto-execution (session resume)") From 5ad096226877c1b28e018d0b7ec9d29dc218afa9 Mon Sep 17 00:00:00 2001 From: Gage Krumbach Date: Fri, 19 Dec 2025 09:06:51 -0600 Subject: [PATCH 4/4] refactor: Simplify INITIAL_PROMPT skip logic - only use PARENT_SESSION_ID Remove marker file approach. Now only check PARENT_SESSION_ID to determine if we should skip the initial prompt. If a parent session exists, this is a continuation and we skip the initial prompt. --- components/runners/claude-code-runner/main.py | 28 ++++++------------- 1 file changed, 8 insertions(+), 20 deletions(-) diff --git a/components/runners/claude-code-runner/main.py b/components/runners/claude-code-runner/main.py index f52b4eb9..7f14b166 100644 --- a/components/runners/claude-code-runner/main.py +++ b/components/runners/claude-code-runner/main.py @@ -7,7 +7,6 @@ import json import logging from contextlib import asynccontextmanager -from pathlib import Path from typing import Optional, List, Dict, Any, Union from fastapi import FastAPI, Request, HTTPException @@ -98,21 +97,17 @@ async def lifespan(app: FastAPI): logger.info("Adapter initialized - fresh client will be created for each run") - # Check if this is a restart/resume or continuation - # - Marker file: session was already initialized (resume) - # - Parent session ID: this is a child session continuing from parent - initial_prompt_marker = Path(workspace_path) / ".initial_prompt_sent" + # Check if this is a continuation (has parent session) + # PARENT_SESSION_ID is set when continuing from another session parent_session_id = os.getenv("PARENT_SESSION_ID", "").strip() - # Check for INITIAL_PROMPT and auto-execute (only if this is a fresh start) + # Check for INITIAL_PROMPT and auto-execute (only if no parent session) initial_prompt = os.getenv("INITIAL_PROMPT", "").strip() - if initial_prompt and not initial_prompt_marker.exists() and not parent_session_id: + if initial_prompt and not parent_session_id: logger.info(f"INITIAL_PROMPT detected ({len(initial_prompt)} chars), will auto-execute after 3s delay") - asyncio.create_task(auto_execute_initial_prompt(initial_prompt, session_id, initial_prompt_marker)) - elif initial_prompt and parent_session_id: - logger.info(f"INITIAL_PROMPT detected but this is a continuation (parent={parent_session_id[:12]}...) - skipping") + asyncio.create_task(auto_execute_initial_prompt(initial_prompt, session_id)) elif initial_prompt: - logger.info(f"INITIAL_PROMPT detected but marker exists - skipping auto-execution (session resume)") + logger.info(f"INITIAL_PROMPT detected but has parent session ({parent_session_id[:12]}...) - skipping") logger.info(f"AG-UI server ready for session {session_id}") @@ -122,13 +117,13 @@ async def lifespan(app: FastAPI): logger.info("Shutting down AG-UI server...") -async def auto_execute_initial_prompt(prompt: str, session_id: str, marker_path: Path): +async def auto_execute_initial_prompt(prompt: str, session_id: str): """Auto-execute INITIAL_PROMPT by POSTing to backend after short delay. The 3-second delay gives the runner time to fully start. Backend has retry logic to handle if Service DNS isn't ready yet. - Creates a marker file on success to prevent re-sending on session resume. + Only called for fresh sessions (no PARENT_SESSION_ID set). """ import uuid import aiohttp @@ -178,13 +173,6 @@ async def auto_execute_initial_prompt(prompt: str, session_id: str, marker_path: if resp.status == 200: result = await resp.json() logger.info(f"INITIAL_PROMPT auto-execution started: {result}") - # Create marker file to prevent re-sending on resume - try: - marker_path.parent.mkdir(parents=True, exist_ok=True) - marker_path.write_text(f"sent_at={asyncio.get_event_loop().time()}\n") - logger.info(f"Created initial prompt marker: {marker_path}") - except Exception as marker_err: - logger.warning(f"Failed to create initial prompt marker: {marker_err}") else: error_text = await resp.text() logger.warning(f"INITIAL_PROMPT failed with status {resp.status}: {error_text[:200]}")