File: screen_ai_service_handler_base.h

package info (click to toggle)
chromium 139.0.7258.127-1
  • links: PTS, VCS
  • area: main
  • in suites:
  • size: 6,122,068 kB
  • sloc: cpp: 35,100,771; ansic: 7,163,530; javascript: 4,103,002; python: 1,436,920; asm: 946,517; xml: 746,709; pascal: 187,653; perl: 88,691; sh: 88,436; objc: 79,953; sql: 51,488; cs: 44,583; fortran: 24,137; makefile: 22,147; tcl: 15,277; php: 13,980; yacc: 8,984; ruby: 7,485; awk: 3,720; lisp: 3,096; lex: 1,327; ada: 727; jsp: 228; sed: 36
file content (144 lines) | stat: -rw-r--r-- 4,907 bytes parent folder | download | duplicates (5)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
// Copyright 2025 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#ifndef CHROME_BROWSER_SCREEN_AI_SCREEN_AI_SERVICE_HANDLER_BASE_H_
#define CHROME_BROWSER_SCREEN_AI_SCREEN_AI_SERVICE_HANDLER_BASE_H_

#include <optional>
#include <set>
#include <string_view>

#include "base/memory/weak_ptr.h"
#include "base/time/time.h"
#include "chrome/browser/screen_ai/resource_monitor.h"
#include "mojo/public/cpp/bindings/pending_receiver.h"
#include "mojo/public/cpp/bindings/receiver.h"
#include "mojo/public/cpp/bindings/remote.h"
#include "services/screen_ai/public/mojom/screen_ai_factory.mojom.h"

namespace base {
class Process;
}

namespace screen_ai {

using ServiceStateCallback = base::OnceCallback<void(bool)>;

// Base class for ScreenAI service management.
class ScreenAIServiceHandlerBase
    : screen_ai::mojom::ScreenAIServiceShutdownHandler {
 public:
  ScreenAIServiceHandlerBase();
  ScreenAIServiceHandlerBase(const ScreenAIServiceHandlerBase&) = delete;
  ScreenAIServiceHandlerBase& operator=(const ScreenAIServiceHandlerBase&) =
      delete;
  ~ScreenAIServiceHandlerBase() override;

  // Schedules library download and initializes the service if needed, and
  // calls `callback` with initialization result when service is ready or
  // failed to initialize.
  // Returns the already known state of the service (see `GetServiceState`).
  std::optional<bool> GetServiceStateAsync(ServiceStateCallback callback);

  // Called when library availability state is changed.
  void OnLibraryAvailablityChanged(bool available);

  // screen_ai::mojom::ScreenAIServiceShutdownHandler::
  void ShuttingDownOnIdle() override;

  // Returns true if the connection for the service is bound.
  bool IsConnectionBoundForTesting();

  // Returns true if sandboxed process is running.
  bool IsProcessRunningForTesting();

 protected:
  std::string GetMetricFullName(std::string_view metric_name) const;

  // Initialzies the service if it's not already done.
  void InitializeServiceIfNeeded();

  // Callback from Screen AI service with library load result.
  void SetLibraryLoadState(base::TimeTicks request_start_time, bool successful);

  mojo::Remote<mojom::ScreenAIServiceFactory>& screen_ai_service_factory() {
    return screen_ai_service_factory_;
  }

 private:
  friend class ScreenAIServiceRouterFactory;
  friend class ScreenAIServiceShutdownHandlerTest;

  bool GetAndRecordSuspendedState();
  void ResetSuspend() { shutdown_handler_data_.suspended = false; }

  // Launches the service if it's not already launched.
  void LaunchIfNotRunning();

  // Called after service is launched.
  void OnServiceLaunched(const std::string& process_name,
                         const base::Process& process);

  // Creates resource monitor to track CPU and memory load.
  void CreateResourceMonitor(const std::string& process_name);

  // True if service is already initialized, false if it is disabled, and
  // nullopt if not known.
  std::optional<bool> GetServiceState();

  // Calls back all pendnding service state requests.
  void CallPendingStatusRequests(bool successful);

  // Called when ScreenAI service factory is disconnected.
  void OnScreenAIServiceDisconnected();

  // Returns the service name.
  virtual std::string GetServiceName() const = 0;

  // Loads the library model files and initializes the service.
  virtual void LoadModelFilesAndInitialize(
      base::TimeTicks request_start_time) = 0;

  virtual bool IsConnectionBound() const = 0;
  virtual bool IsServiceEnabled() const = 0;
  virtual void ResetConnection() = 0;

  // Pending requests to receive service state for each service type.
  std::vector<ServiceStateCallback> pending_state_requests_;

  std::unique_ptr<ResourceMonitor> resource_monitor_;
  struct ShutdownHandlerData {
    bool shutdown_message_received = false;
    bool suspended = false;
    int crash_count = 0;
  } shutdown_handler_data_;

  base::TimeTicks service_start_time_;

  mojo::Receiver<screen_ai::mojom::ScreenAIServiceShutdownHandler>
      screen_ai_service_shutdown_handler_;

  mojo::Remote<mojom::ScreenAIServiceFactory> screen_ai_service_factory_;

  base::WeakPtrFactory<ScreenAIServiceHandlerBase> weak_ptr_factory_{this};
};

class ComponentFiles {
 public:
  explicit ComponentFiles(const base::FilePath& library_binary_path,
                          const base::FilePath::CharType* files_list_file_name);
  ComponentFiles(const ComponentFiles&) = delete;
  ComponentFiles& operator=(const ComponentFiles&) = delete;
  ~ComponentFiles();

  static std::unique_ptr<ComponentFiles> Load(
      const base::FilePath::CharType* files_list_file_name);

  base::flat_map<base::FilePath, base::File> model_files_;
  base::FilePath library_binary_path_;
};

}  // namespace screen_ai

#endif  // CHROME_BROWSER_SCREEN_AI_SCREEN_AI_SERVICE_HANDLER_BASE_H_