Merge pull request #830 from rumpl/remote-oauth

Add unmanaged mode for oauth flow

Author: dgageot

pkg/agent/agent_test.go

@@ -37,6 +37,7 @@ func (s *stubToolSet) Tools(context.Context) ([]tools.Tool, error) {
 func (s *stubToolSet) Instructions() string                           { return s.instructions }
 func (s *stubToolSet) SetElicitationHandler(tools.ElicitationHandler) {}
 func (s *stubToolSet) SetOAuthSuccessHandler(func())                  {}
+func (s *stubToolSet) SetManagedOAuth(bool)                           {}
 
 func TestAgentTools(t *testing.T) {
 	tests := []struct {

pkg/api/types.go

@@ -155,17 +155,6 @@ type DesktopTokenResponse struct {
 	Token string `json:"token"`
 }
 
-// ResumeStartOauthRequest represents the user approval to start the OAuth flow
-type ResumeStartOauthRequest struct {
-	Confirmation bool `json:"confirmation"`
-}
-
-// ResumeCodeReceivedOauthRequest represents the response from getting the OAuth URL with code and state
-type ResumeCodeReceivedOauthRequest struct {
-	Code  string `json:"code"`
-	State string `json:"state"`
-}
-
 // ResumeElicitationRequest represents a request to resume with an elicitation response
 type ResumeElicitationRequest struct {
 	Action  string         `json:"action"`  // "accept", "decline", or "cancel"

pkg/app/app.go

@@ -9,6 +9,7 @@ import (
 
 	"github.com/docker/cagent/pkg/runtime"
 	"github.com/docker/cagent/pkg/session"
+	"github.com/docker/cagent/pkg/tools"
 )
 
 type App struct {
@@ -87,9 +88,12 @@ func (a *App) Subscribe(ctx context.Context, program *tea.Program) {
 
 // Resume resumes the runtime with the given confirmation type
 func (a *App) Resume(resumeType runtime.ResumeType) {
-	if a.runtime != nil {
-		a.runtime.Resume(context.Background(), resumeType)
-	}
+	a.runtime.Resume(context.Background(), resumeType)
+}
+
+// ResumeElicitation resumes an elicitation request with the given action and content
+func (a *App) ResumeElicitation(ctx context.Context, action tools.ElicitationAction, content map[string]any) error {
+	return a.runtime.ResumeElicitation(ctx, action, content)
 }
 
 func (a *App) NewSession() {
@@ -105,7 +109,7 @@ func (a *App) Session() *session.Session {
 }
 
 func (a *App) CompactSession() {
-	if a.runtime != nil && a.session != nil {
+	if a.session != nil {
 		events := make(chan runtime.Event, 100)
 		a.runtime.Summarize(context.Background(), a.session, events)
 		close(events)
@@ -115,14 +119,6 @@ func (a *App) CompactSession() {
 	}
 }
 
-// ResumeStartOAuth resumes the runtime with OAuth authorization confirmation
-func (a *App) ResumeStartOAuth(bool) {
-	if a.runtime != nil {
-		// TODO(rumpl): handle the error
-		_ = a.runtime.ResumeElicitation(context.Background(), "accept", nil)
-	}
-}
-
 func (a *App) PlainTextTranscript() string {
 	return transcript(a.session)
 }

pkg/runtime/client.go

@@ -16,6 +16,7 @@ import (
 	"github.com/docker/cagent/pkg/api"
 	v2 "github.com/docker/cagent/pkg/config/v2"
 	"github.com/docker/cagent/pkg/session"
+	"github.com/docker/cagent/pkg/tools"
 )
 
 // Client is an HTTP client for the cagent server API
@@ -373,17 +374,7 @@ func (c *Client) runAgentWithAgentName(ctx context.Context, sessionID, agent, ag
 	return eventChan, nil
 }
 
-func (c *Client) ResumeStartAuthorizationFlow(ctx context.Context, id string, confirmation bool) error {
-	req := api.ResumeStartOauthRequest{Confirmation: confirmation}
-	return c.doRequest(ctx, http.MethodPost, "/api/"+id+"/resumeStartOauth", req, nil)
-}
-
-func (c *Client) ResumeCodeReceived(ctx context.Context, code, state string) error {
-	req := api.ResumeCodeReceivedOauthRequest{Code: code, State: state}
-	return c.doRequest(ctx, http.MethodPost, "/api/resumeCodeReceivedOauth", req, nil)
-}
-
-func (c *Client) ResumeElicitation(ctx context.Context, action string, content map[string]any) error {
-	req := api.ResumeElicitationRequest{Action: action, Content: content}
-	return c.doRequest(ctx, http.MethodPost, "/api/resumeElicitation", req, nil)
+func (c *Client) ResumeElicitation(ctx context.Context, sessionID string, action tools.ElicitationAction, content map[string]any) error {
+	req := api.ResumeElicitationRequest{Action: string(action), Content: content}
+	return c.doRequest(ctx, http.MethodPost, "/api/sessions/"+sessionID+"/elicitation", req, nil)
 }

pkg/runtime/event.go

@@ -293,14 +293,14 @@ func ElicitationRequest(message string, schema any, meta map[string]any, agentNa
 func (e *ElicitationRequestEvent) GetAgentName() string { return e.AgentName }
 
 type AuthorizationEvent struct {
-	Type         string `json:"type"`
-	Confirmation string `json:"confirmation"` // only "confirmed"
+	Type         string                  `json:"type"`
+	Confirmation tools.ElicitationAction `json:"confirmation"`
 	AgentContext
 }
 
 func (e *AuthorizationEvent) GetAgentName() string { return "" }
 
-func Authorization(confirmation, agentName string) Event {
+func Authorization(confirmation tools.ElicitationAction, agentName string) Event {
 	return &AuthorizationEvent{
 		Type:         "authorization_event",
 		Confirmation: confirmation,

pkg/runtime/remote_runtime.go

@@ -2,23 +2,30 @@ package runtime
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"log/slog"
+	"time"
+
+	"golang.org/x/oauth2"
 
 	"github.com/docker/cagent/pkg/api"
 	"github.com/docker/cagent/pkg/chat"
 	latest "github.com/docker/cagent/pkg/config/v2"
 	"github.com/docker/cagent/pkg/session"
 	"github.com/docker/cagent/pkg/team"
+	"github.com/docker/cagent/pkg/tools"
+	"github.com/docker/cagent/pkg/tools/mcp"
 )
 
 // RemoteRuntime implements the Interface using a remote client
 type RemoteRuntime struct {
-	client        *Client
-	currentAgent  string
-	agentFilename string
-	sessionID     string
-	team          *team.Team
+	client                  *Client
+	currentAgent            string
+	agentFilename           string
+	sessionID               string
+	team                    *team.Team
+	pendingOAuthElicitation *ElicitationRequestEvent
 }
 
 // RemoteRuntimeOption is a function for configuring the RemoteRuntime
@@ -115,6 +122,10 @@ func (r *RemoteRuntime) RunStream(ctx context.Context, sess *session.Session) <-
 		}
 
 		for streamEvent := range streamChan {
+			if elicitationRequest, ok := streamEvent.(*ElicitationRequestEvent); ok {
+				// Store pending OAuth elicitation request
+				r.pendingOAuthElicitation = elicitationRequest
+			}
 			events <- streamEvent
 		}
 	}()
@@ -176,51 +187,186 @@ func (r *RemoteRuntime) convertSessionMessages(sess *session.Session) []api.Mess
 	return messages
 }
 
-// ResumeStartAuthorizationFlow allows resuming execution after user confirmation
-func (r *RemoteRuntime) ResumeStartAuthorizationFlow(ctx context.Context, confirmationType bool) {
-	slog.Debug("Resuming remote runtime", "agent", r.currentAgent, "confirmation_type", confirmationType, "session_id", r.sessionID)
+// ResumeElicitation sends an elicitation response back to a waiting elicitation request
+func (r *RemoteRuntime) ResumeElicitation(ctx context.Context, action tools.ElicitationAction, content map[string]any) error {
+	slog.Debug("Resuming remote runtime with elicitation response", "agent", r.currentAgent, "action", action, "session_id", r.sessionID)
 
-	if r.sessionID == "" {
-		slog.Error("Cannot resume: no session ID available")
-		return
+	err := r.handleOAuthElicitation(ctx, r.pendingOAuthElicitation)
+	if err != nil {
+		return err
 	}
+	// TODO: once we get here and the elicitation is the OAuth type, we need to start the managed OAuth flow
 
-	if err := r.client.ResumeStartAuthorizationFlow(ctx, r.sessionID, confirmationType); err != nil {
-		slog.Error("Failed to resume remote session", "error", err, "session_id", r.sessionID)
+	if err := r.client.ResumeElicitation(ctx, r.sessionID, action, content); err != nil {
+		return err
 	}
+
+	return nil
 }
 
-// ResumeCodeReceived allows resuming execution after user confirmation
-func (r *RemoteRuntime) ResumeCodeReceived(ctx context.Context, code, state string) error {
-	slog.Debug("Resuming remote runtime", "agent", r.currentAgent, "code", code, "state", state, "session_id", r.sessionID)
+// HandleOAuthElicitation handles OAuth elicitation requests from remote MCP servers
+func (r *RemoteRuntime) handleOAuthElicitation(ctx context.Context, req *ElicitationRequestEvent) error {
+	slog.Debug("Handling OAuth elicitation request", "server_url", req.Meta["cagent/server_url"])
 
-	if r.sessionID == "" {
-		slog.Error("Cannot resume: no session ID available")
-		return fmt.Errorf("session ID cannot be empty")
+	// Extract OAuth parameters from metadata
+	serverURL, ok := req.Meta["cagent/server_url"].(string)
+	if !ok {
+		err := fmt.Errorf("server_url missing from elicitation metadata")
+		slog.Error("Failed to extract server_url", "error", err)
+		_ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil)
+		return err
 	}
 
-	if err := r.client.ResumeCodeReceived(ctx, code, state); err != nil {
-		slog.Error("Failed to resume remote session", "error", err, "session_id", r.sessionID)
+	// Extract authorization server metadata
+	authServerMetadata, ok := req.Meta["auth_server_metadata"].(map[string]any)
+	if !ok {
+		err := fmt.Errorf("auth_server_metadata missing from elicitation metadata")
+		slog.Error("Failed to extract auth_server_metadata", "error", err)
+		_ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil)
 		return err
 	}
 
-	return nil
-}
+	// Unmarshal authorization server metadata
+	var authMetadata mcp.AuthorizationServerMetadata
+	metadataBytes, err := json.Marshal(authServerMetadata)
+	if err != nil {
+		slog.Error("Failed to marshal auth_server_metadata", "error", err)
+		_ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil)
+		return fmt.Errorf("failed to marshal auth_server_metadata: %w", err)
+	}
+	if err := json.Unmarshal(metadataBytes, &authMetadata); err != nil {
+		slog.Error("Failed to unmarshal auth_server_metadata", "error", err)
+		_ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil)
+		return fmt.Errorf("failed to unmarshal auth_server_metadata: %w", err)
+	}
 
-// ResumeElicitation sends an elicitation response back to a waiting elicitation request
-func (r *RemoteRuntime) ResumeElicitation(ctx context.Context, action string, content map[string]any) error {
-	slog.Debug("Resuming remote runtime with elicitation response", "agent", r.currentAgent, "action", action, "session_id", r.sessionID)
+	slog.Debug("Authorization server metadata extracted", "issuer", authMetadata.Issuer)
 
-	if r.sessionID == "" {
-		slog.Error("Cannot resume: no session ID available")
-		return fmt.Errorf("session ID cannot be empty")
+	// Create timeout context for OAuth flow (5 minutes)
+	oauthCtx, cancel := context.WithTimeout(ctx, 5*time.Minute)
+	defer cancel()
+
+	// Create and start callback server
+	slog.Debug("Creating OAuth callback server")
+	callbackServer, err := mcp.NewCallbackServer()
+	if err != nil {
+		slog.Error("Failed to create callback server", "error", err)
+		_ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil)
+		return fmt.Errorf("failed to create callback server: %w", err)
 	}
+	defer func() {
+		shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer shutdownCancel()
+		if err := callbackServer.Shutdown(shutdownCtx); err != nil {
+			slog.Error("Failed to shutdown callback server", "error", err)
+		}
+	}()
+
+	if err := callbackServer.Start(); err != nil {
+		slog.Error("Failed to start callback server", "error", err)
+		_ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil)
+		return fmt.Errorf("failed to start callback server: %w", err)
+	}
+
+	redirectURI := callbackServer.GetRedirectURI()
+	slog.Debug("Callback server started", "redirect_uri", redirectURI)
+
+	// Register client
+	var clientID, clientSecret string
+	if authMetadata.RegistrationEndpoint != "" {
+		slog.Debug("Attempting dynamic client registration")
+		clientID, clientSecret, err = mcp.RegisterClient(oauthCtx, &authMetadata, redirectURI, nil)
+		if err != nil {
+			slog.Error("Dynamic client registration failed", "error", err)
+			_ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil)
+			return fmt.Errorf("failed to register client: %w", err)
+		}
+		slog.Debug("Client registered successfully", "client_id", clientID)
+	} else {
+		err := fmt.Errorf("authorization server does not support dynamic client registration")
+		slog.Error("Client registration not supported", "error", err)
+		_ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil)
+		return err
+	}
+
+	// Generate state and PKCE verifier
+	state, err := mcp.GenerateState()
+	if err != nil {
+		slog.Error("Failed to generate state", "error", err)
+		_ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil)
+		return fmt.Errorf("failed to generate state: %w", err)
+	}
+
+	callbackServer.SetExpectedState(state)
+	verifier := mcp.GeneratePKCEVerifier()
+
+	// Build authorization URL
+	authURL := mcp.BuildAuthorizationURL(
+		authMetadata.AuthorizationEndpoint,
+		clientID,
+		redirectURI,
+		state,
+		oauth2.S256ChallengeFromVerifier(verifier),
+		serverURL,
+	)
 
-	if err := r.client.ResumeElicitation(ctx, action, content); err != nil {
-		slog.Error("Failed to resume remote session with elicitation", "error", err, "session_id", r.sessionID)
+	slog.Debug("Authorization URL built", "url", authURL)
+
+	// Request authorization code (this opens the browser)
+	slog.Debug("Requesting authorization code")
+	code, receivedState, err := mcp.RequestAuthorizationCode(oauthCtx, authURL, callbackServer, state)
+	if err != nil {
+		slog.Error("Failed to get authorization code", "error", err)
+		_ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil)
+		return fmt.Errorf("failed to get authorization code: %w", err)
+	}
+
+	if receivedState != state {
+		err := fmt.Errorf("state mismatch: expected %s, got %s", state, receivedState)
+		slog.Error("State mismatch in authorization response", "error", err)
+		_ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil)
 		return err
 	}
 
+	slog.Debug("Authorization code received, exchanging for token")
+
+	// Exchange code for token
+	token, err := mcp.ExchangeCodeForToken(
+		oauthCtx,
+		authMetadata.TokenEndpoint,
+		code,
+		verifier,
+		clientID,
+		clientSecret,
+		redirectURI,
+	)
+	if err != nil {
+		slog.Error("Failed to exchange code for token", "error", err)
+		_ = r.client.ResumeElicitation(ctx, r.sessionID, "decline", nil)
+		return fmt.Errorf("failed to exchange code for token: %w", err)
+	}
+
+	slog.Debug("Token obtained successfully", "token_type", token.TokenType)
+
+	// Send token back to server via ResumeElicitation
+	tokenData := map[string]any{
+		"access_token": token.AccessToken,
+		"token_type":   token.TokenType,
+	}
+	if token.ExpiresIn > 0 {
+		tokenData["expires_in"] = token.ExpiresIn
+	}
+	if token.RefreshToken != "" {
+		tokenData["refresh_token"] = token.RefreshToken
+	}
+
+	slog.Debug("Sending token to server")
+	if err := r.client.ResumeElicitation(ctx, r.sessionID, tools.ElicitationActionAccept, tokenData); err != nil {
+		slog.Error("Failed to send token to server", "error", err)
+		return fmt.Errorf("failed to send token to server: %w", err)
+	}
+
+	slog.Debug("OAuth flow completed successfully")
 	return nil
 }