From 72bc3c0319580b4b937eaebcba6d72f904a3ddcd Mon Sep 17 00:00:00 2001
From: Haitao Pan <haitao.pan@xworkmate.ai>
Date: Sun, 7 Jun 2026 21:57:53 +0800
Subject: [PATCH] fix: delay desktop capture until WebRTC is ready

---
 internal/acp/rpc_handler.go |  4 ++++
 internal/desktop/service.go | 29 ++++++++++++++++++++---------
 2 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/internal/acp/rpc_handler.go b/internal/acp/rpc_handler.go
index d53f3c2..06f4b9b 100644
--- a/internal/acp/rpc_handler.go
+++ b/internal/acp/rpc_handler.go
@@ -491,6 +491,10 @@ func (s *Server) handleDesktopMethod(ctx context.Context, method string, params
 			srv.StopSession(sessionID)
 			return nil, &shared.RPCError{Code: -32002, Message: fmt.Sprintf("failed to process SDP offer: %v", err)}
 		}
+		if err := srv.StartCapture(sessionID); err != nil {
+			srv.StopSession(sessionID)
+			return nil, &shared.RPCError{Code: -32004, Message: fmt.Sprintf("failed to start desktop capture: %v", err)}
+		}
 
 		return map[string]any{
 			"sessionId": sessionID,
diff --git a/internal/desktop/service.go b/internal/desktop/service.go
index 52042ca..12b95cb 100644
--- a/internal/desktop/service.go
+++ b/internal/desktop/service.go
@@ -11,6 +11,7 @@ import (
 type DesktopSession struct {
 	SessionID string
 	Port      int
+	Config    PipelineConfig
 	Pipeline  *PipelineManager
 	Injector  *XdotoolInjector
 	WebRTC    *WebRTCServer
@@ -72,18 +73,11 @@ func (s *Service) StartSession(sessionID string, cfg PipelineConfig, iceServers
 		return nil, fmt.Errorf("failed to start RTP receiver: %w", err)
 	}
 
-	// 3. Initialize screen capture pipeline
-	pipeline := NewPipelineManager()
-	if err := pipeline.Start(cfg); err != nil {
-		webrtcSrv.Close()
-		_ = injector.Close()
-		return nil, fmt.Errorf("failed to start capture pipeline: %w", err)
-	}
-
 	sess := &DesktopSession{
 		SessionID: sessionID,
 		Port:      cfg.Port,
-		Pipeline:  pipeline,
+		Config:    cfg,
+		Pipeline:  NewPipelineManager(),
 		Injector:  injector,
 		WebRTC:    webrtcSrv,
 	}
@@ -92,6 +86,23 @@ func (s *Service) StartSession(sessionID string, cfg PipelineConfig, iceServers
 	return sess, nil
 }
 
+func (s *Service) StartCapture(sessionID string) error {
+	sess, err := s.GetSession(sessionID)
+	if err != nil {
+		return err
+	}
+	if sess.Pipeline == nil {
+		sess.Pipeline = NewPipelineManager()
+	}
+	if sess.Pipeline.IsRunning() {
+		return nil
+	}
+	if err := sess.Pipeline.Start(sess.Config); err != nil {
+		return fmt.Errorf("failed to start capture pipeline: %w", err)
+	}
+	return nil
+}
+
 func (s *Service) GetSession(sessionID string) (*DesktopSession, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()