Google Workspace MCP を Managed Agents で動かす

import static com.google.common.base.Preconditions.checkState;
import static java.nio.charset.StandardCharsets.UTF_8;

import com.anthropic.client.AnthropicClient;
import com.anthropic.client.okhttp.AnthropicOkHttpClient;
import com.anthropic.core.JsonValue;
import com.anthropic.core.http.StreamResponse;
import com.anthropic.models.beta.agents.AgentCreateParams;
import com.anthropic.models.beta.agents.BetaManagedAgentsAlwaysAllowPolicy;
import com.anthropic.models.beta.agents.BetaManagedAgentsMcpToolConfigParams;
import com.anthropic.models.beta.agents.BetaManagedAgentsMcpToolsetDefaultConfigParams;
import com.anthropic.models.beta.agents.BetaManagedAgentsMcpToolsetParams;
import com.anthropic.models.beta.agents.BetaManagedAgentsModel;
import com.anthropic.models.beta.agents.BetaManagedAgentsModelConfigParams;
import com.anthropic.models.beta.agents.BetaManagedAgentsUrlMcpServerParams;
import com.anthropic.models.beta.environments.EnvironmentCreateParams;
import com.anthropic.models.beta.sessions.SessionCreateParams;
import com.anthropic.models.beta.sessions.events.BetaManagedAgentsStreamSessionEvents;
import com.anthropic.models.beta.sessions.events.BetaManagedAgentsTextBlock;
import com.anthropic.models.beta.sessions.events.BetaManagedAgentsUserMessageEventParams;
import com.anthropic.models.beta.sessions.events.EventSendParams;
import com.anthropic.models.beta.vaults.VaultCreateParams;
import com.anthropic.models.beta.vaults.credentials.BetaManagedAgentsMcpOAuthCreateParams;
import com.anthropic.models.beta.vaults.credentials.BetaManagedAgentsMcpOAuthRefreshParams;
import com.anthropic.models.beta.vaults.credentials.CredentialCreateParams;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.SerializationFeature;
import com.google.common.base.Joiner;
import com.google.common.base.Strings;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Maps;
import java.io.BufferedWriter;
import java.io.IOException;
import java.net.URI;
import java.net.URLEncoder;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.nio.file.Files;
import java.nio.file.Path;
import java.time.OffsetDateTime;
import java.time.ZoneOffset;
import java.util.Iterator;
import java.util.List;
import java.util.Map;


/**
 * Google Workspace MCP に Managed Agents から接続するデモ。
 *
 * <p>Calendar は読み取りだけ、Gmail は下書き作成だけを許可する。 OAuth トークンは Vault に保存し、実行時は Vault ID
 * だけを Session に渡す。
 * 
 * <p>セットアップ: 自前の認可フローで取得済みのトークンを環境変数に入れて、setupを行う
 * 実行前に設定: ANTHROPIC_API_KEY / GOOGLE_CLIENT_ID / GOOGLE_CLIENT_SECRET / GOOGLE_REFRESH_TOKEN
 * mvn -q compile exec:java -Dexec.args="setup"
 *
 * <p>実行: ANTHROPIC_ENVIRONMENT_ID / ANTHROPIC_AGENT_ID / ANTHROPIC_VAULT_ID を設定して
 * mvn -q compile exec:java
*/
public final class WorkspaceMcpAgent {

  private static final String CALENDAR_MCP_URL = "https://calendarmcp.googleapis.com/mcp/v1";
  private static final String GMAIL_MCP_URL = "https://gmailmcp.googleapis.com/mcp/v1";
  private static final String GOOGLE_TOKEN_ENDPOINT = "https://oauth2.googleapis.com/token";

  private static final String CALENDAR_SCOPE =
      "https://www.googleapis.com/auth/calendar.events.readonly";
  private static final String GMAIL_SCOPE = "https://www.googleapis.com/auth/gmail.compose";

  private static final Path DEFAULT_OUTPUT_PATH = Path.of("output.md");
  private static BufferedWriter outputWriter;

  private static final String TASK =
      """
      今日から3日以内に予定されている打ち合わせをすべて確認してください。

      該当した打ち合わせ1件につき1通、その参加者宛に次の内容のメールの下書きを作成してください。
      複数該当する場合は、すべてについて作成してください。

      - 体調不良のため、対面ではなくリモート会議に切り替えてほしいこと
      - どの打ち合わせのことか分かるよう、会議名と日時に触れること
      - 内容は簡潔に、事務的になりすぎない範囲で丁寧に

      3日以内に該当する打ち合わせが無ければ、下書きは作らず、その旨だけを答えてください。
      予定の内容を推測で補わないでください。下書きの作成までで構いません。送信はしないでください。
      """;

  /** SDK が返した MCP input/result と Google OAuth レスポンスを扱う。 */
  private static final ObjectMapper JSON =
      new ObjectMapper().findAndRegisterModules().enable(SerializationFeature.INDENT_OUTPUT);

  private static final BetaManagedAgentsAlwaysAllowPolicy ALWAYS_ALLOW =
      BetaManagedAgentsAlwaysAllowPolicy.builder()
          .type(BetaManagedAgentsAlwaysAllowPolicy.Type.ALWAYS_ALLOW)
          .build();

  private final AnthropicClient client;
  private final HttpClient httpClient = HttpClient.newHttpClient();

  /** agent.mcp_tool_use の event id -> 呼び出し情報。tool_result と対応付ける。 */
  private final Map<String, PendingCall> pendingCalls = Maps.newHashMap();

  private record PendingCall(String server, String tool, String inputJson) {}

  private record RunResult(int successfulDrafts, String answer) {}

  private record GoogleAccessToken(String accessToken, OffsetDateTime expiresAt) {}

  WorkspaceMcpAgent(AnthropicClient client) {
    this.client = client;
  }

  public static void main(String[] args) throws IOException {
    Path outputPath = resolveOutputPath();
    createParentDirectories(outputPath);

    try (BufferedWriter writer = Files.newBufferedWriter(outputPath, UTF_8)) {
      outputWriter = writer;

      try {
        WorkspaceMcpAgent app = new WorkspaceMcpAgent(AnthropicOkHttpClient.fromEnv());

        if (args.length > 0 && "setup".equalsIgnoreCase(args[0])) {
          app.setup();
          return;
        }

        app.run(
            requiredEnvironment("ANTHROPIC_ENVIRONMENT_ID"),
            requiredEnvironment("ANTHROPIC_AGENT_ID"),
            requiredEnvironment("ANTHROPIC_VAULT_ID"));
      } catch (RuntimeException exception) {
        out();
        out("## エラー");
        out();
        out("- 種類: `" + exception.getClass().getName() + "`");
        out("- メッセージ: " + String.valueOf(exception.getMessage()));
        throw exception;
      }
    } finally {
      outputWriter = null;
    }
  }

  /** Environment / Agent / Vault を作成し、Google の OAuth credential を Vault に登録する。 */
  void setup() {
    String clientId = requiredEnvironment("GOOGLE_CLIENT_ID");
    String clientSecret = requiredEnvironment("GOOGLE_CLIENT_SECRET");

    var environment =
        client
            .beta()
            .environments()
            .create(EnvironmentCreateParams.builder().name("workspace-mcp-environment").build());

    var agent =
        client
            .beta()
            .agents()
            .create(
                AgentCreateParams.builder()
                    .name("meeting-assistant")
                    .model(
                        BetaManagedAgentsModelConfigParams.builder()
                            .id(BetaManagedAgentsModel.CLAUDE_SONNET_5)
                            .build())
                    .system(
                        "カレンダーとメールを扱うアシスタントです。取得した情報だけを根拠に行動し、"
                            + "推測で予定や宛先を補わないでください。下書きの作成までで止め、送信はしません。")
                    .addMcpServer(mcpServer("calendar", CALENDAR_MCP_URL))
                    .addMcpServer(mcpServer("gmail", GMAIL_MCP_URL))
                    .addTool(mcpToolset("calendar", ImmutableList.of("list_events", "get_event")))
                    .addTool(mcpToolset("gmail", ImmutableList.of("create_draft")))
                    .build());

    var vault =
        client
            .beta()
            .vaults()
            .create(
                VaultCreateParams.builder()
                    .displayName("Masahiro")
                    .metadata(
                        VaultCreateParams.Metadata.builder()
                            .putAdditionalProperty(
                                "external_user_id", JsonValue.from("usr_masahiro"))
                            .build())
                    .build());

    registerCredential(
        vault.id(),
        "Google Calendar",
        CALENDAR_MCP_URL,
        CALENDAR_SCOPE,
        clientId,
        clientSecret,
        requiredEnvironment("GOOGLE_REFRESH_TOKEN"));

    registerCredential(
        vault.id(),
        "Gmail",
        GMAIL_MCP_URL,
        GMAIL_SCOPE,
        clientId,
        clientSecret,
        requiredEnvironment("GOOGLE_REFRESH_TOKEN"));

    out("# setup completed");
    out();
    out("export ANTHROPIC_ENVIRONMENT_ID=\"" + environment.id() + "\"");
    out("export ANTHROPIC_AGENT_ID=\"" + agent.id() + "\"");
    out("export ANTHROPIC_VAULT_ID=\"" + vault.id() + "\"");
  }

  /** MCP server 定義。 */
  private static BetaManagedAgentsUrlMcpServerParams mcpServer(String name, String url) {
    return BetaManagedAgentsUrlMcpServerParams.builder()
        .type(BetaManagedAgentsUrlMcpServerParams.Type.URL)
        .name(name)
        .url(url)
        .build();
  }

  /**
   * MCP toolset の許可リストを作る。
   *
   * <p>default は enabled=false にして、明示した tool だけ enabled=true にする。 permission は always_allow
   * に固定するため、実行途中で確認待ちにならない。
   */
  private static BetaManagedAgentsMcpToolsetParams mcpToolset(
      String serverName, List<String> enabledTools) {
    var builder =
        BetaManagedAgentsMcpToolsetParams.builder()
            .type(BetaManagedAgentsMcpToolsetParams.Type.MCP_TOOLSET)
            .mcpServerName(serverName)
            .defaultConfig(
                BetaManagedAgentsMcpToolsetDefaultConfigParams.builder()
                    .enabled(false)
                    .permissionPolicy(ALWAYS_ALLOW)
                    .build());

    enabledTools.forEach(
        toolName ->
            builder.addConfig(
                BetaManagedAgentsMcpToolConfigParams.builder()
                    .name(toolName)
                    .enabled(true)
                    .build()));

    return builder.build();
  }

  /** refresh token から最新 access token を取得し、expires_in から expiresAt を自動計算して Vault に登録する。 */
  private void registerCredential(
      String vaultId,
      String displayName,
      String serverUrl,
      String scope,
      String clientId,
      String clientSecret,
      String refreshToken) {

    GoogleAccessToken token = refreshGoogleAccessToken(clientId, clientSecret, refreshToken);

    var refresh =
        BetaManagedAgentsMcpOAuthRefreshParams.builder()
            .tokenEndpoint(GOOGLE_TOKEN_ENDPOINT)
            .clientId(clientId)
            .clientSecretPostTokenEndpointAuth(clientSecret)
            .scope(scope)
            .refreshToken(refreshToken)
            .build();

    var oauth =
        BetaManagedAgentsMcpOAuthCreateParams.builder()
            .type(BetaManagedAgentsMcpOAuthCreateParams.Type.MCP_OAUTH)
            .mcpServerUrl(serverUrl)
            .accessToken(token.accessToken())
            .expiresAt(token.expiresAt())
            .refresh(refresh)
            .build();

    client
        .beta()
        .vaults()
        .credentials()
        .create(
            CredentialCreateParams.builder()
                .vaultId(vaultId)
                .displayName(displayName)
                .auth(oauth)
                .build());

    out("- " + displayName + " credential registered (expiresAt: " + token.expiresAt() + ")");
  }

  /** Google OAuth token endpoint に refresh_token grant を送り、新しい access token を取得する。 */
  private GoogleAccessToken refreshGoogleAccessToken(
      String clientId, String clientSecret, String refreshToken) {

    String form =
        "client_id="
            + formEncode(clientId)
            + "&client_secret="
            + formEncode(clientSecret)
            + "&refresh_token="
            + formEncode(refreshToken)
            + "&grant_type=refresh_token";

    HttpRequest request =
        HttpRequest.newBuilder(URI.create(GOOGLE_TOKEN_ENDPOINT))
            .header("Content-Type", "application/x-www-form-urlencoded")
            .POST(HttpRequest.BodyPublishers.ofString(form, UTF_8))
            .build();

    try {
      HttpResponse<String> response =
          httpClient.send(request, HttpResponse.BodyHandlers.ofString(UTF_8));

      if (response.statusCode() < 200 || response.statusCode() >= 300) {
        throw new IllegalStateException(
            "Google OAuth token refresh failed: HTTP "
                + response.statusCode()
                + " / "
                + sanitizeGoogleError(response.body()));
      }

      JsonNode body = JSON.readTree(response.body());
      String accessToken = body.path("access_token").asText("");
      long expiresIn = body.path("expires_in").asLong(-1);

      checkState(!accessToken.isBlank(), "Google OAuth response に access_token がありません");
      checkState(expiresIn > 0, "Google OAuth response の expires_in が不正です: %s", expiresIn);

      OffsetDateTime expiresAt = OffsetDateTime.now(ZoneOffset.UTC).plusSeconds(expiresIn);
      return new GoogleAccessToken(accessToken, expiresAt);

    } catch (InterruptedException exception) {
      Thread.currentThread().interrupt();
      throw new IllegalStateException("Google OAuth token refresh が中断されました", exception);
    } catch (IOException exception) {
      throw new IllegalStateException("Google OAuth token refresh に失敗しました", exception);
    }
  }

  private static String formEncode(String value) {
    return URLEncoder.encode(value, UTF_8);
  }

  /** エラー本文に token 等が含まれないよう、Google が返した代表的なエラー項目だけを取り出す。 */
  private static String sanitizeGoogleError(String responseBody) {
    try {
      JsonNode body = JSON.readTree(responseBody);
      String error = body.path("error").asText("unknown_error");
      String description = body.path("error_description").asText("");
      return description.isBlank() ? error : error + ": " + description;
    } catch (JsonProcessingException exception) {
      return "response body could not be parsed";
    }
  }

  /** Session を作成して TASK を送り、完了までストリームを読む。 */
  void run(String environmentId, String agentId, String vaultId) {
    pendingCalls.clear();

    var session =
        client
            .beta()
            .sessions()
            .create(
                SessionCreateParams.builder()
                    .environmentId(environmentId)
                    .agent(agentId)
                    .addVaultId(vaultId)
                    .title("打ち合わせのリモート切り替え依頼")
                    .build());

    sendTask(session.id());
    RunResult result = consumeEvents(session.id());

    out();
    out("作成に成功した下書き: " + result.successfulDrafts() + " 件");
    out("A. " + result.answer());
  }

  private void sendTask(String sessionId) {
    client
        .beta()
        .sessions()
        .events()
        .send(
            EventSendParams.builder()
                .sessionId(sessionId)
                .addUserMessageEvent(
                    ImmutableList.of(
                        BetaManagedAgentsUserMessageEventParams.Content.ofText(
                            BetaManagedAgentsTextBlock.builder()
                                .type(BetaManagedAgentsTextBlock.Type.TEXT)
                                .text(TASK)
                                .build())))
                .build());
  }

  private RunResult consumeEvents(String sessionId) {
    StringBuilder answer = new StringBuilder();
    int successfulDrafts = 0;

    try (StreamResponse<BetaManagedAgentsStreamSessionEvents> response =
        client.beta().sessions().events().streamStreaming(sessionId)) {

      Iterator<BetaManagedAgentsStreamSessionEvents> events = response.stream().iterator();

      while (events.hasNext()) {
        BetaManagedAgentsStreamSessionEvents event = events.next();

        if (event.isAgentMcpToolUse()) {
          handleMcpToolUse(event);
          continue;
        }

        if (event.isAgentMcpToolResult()) {
          if (isSuccessfulDraftResult(event)) {
            successfulDrafts++;
          }
          continue;
        }

        if (event.isAgentMessage()) {
          event.asAgentMessage().content().forEach(block -> answer.append(block.text()));
          continue;
        }

        if (event.isSessionError()) {
          throw new IllegalStateException("session error: " + event.asSessionError());
        }

        if (event.isSessionStatusIdle()) {
          var stopReason = event.asSessionStatusIdle().stopReason();

          if (stopReason.isEndTurn()) {
            break;
          }

          if (stopReason.isRequiresAction()) {
            String calls =
                Joiner.on(", ")
                    .join(
                        stopReason.asRequiresAction().eventIds().stream()
                            .map(id -> String.valueOf(pendingCalls.get(id)))
                            .toList());
            throw new IllegalStateException("承認待ちで停止しました。permission policy を確認してください: " + calls);
          }
        }
      }
    }

    return new RunResult(successfulDrafts, answer.toString().strip());
  }

  private void handleMcpToolUse(BetaManagedAgentsStreamSessionEvents event) {
    var toolUse = event.asAgentMcpToolUse();
    String inputJson = toPrettyJson(toolUse.input());

    pendingCalls.put(
        toolUse.id(), new PendingCall(toolUse.mcpServerName(), toolUse.name(), inputJson));

    out("[mcp_tool_use] " + toolUse.mcpServerName() + " / " + toolUse.name());
    out(inputJson);
  }

  /** create_draft の MCP 実行結果が成功した場合だけ true。 */
  private boolean isSuccessfulDraftResult(BetaManagedAgentsStreamSessionEvents event) {
    var result = event.asAgentMcpToolResult();
    PendingCall call = pendingCalls.remove(result.mcpToolUseId());
    boolean failed = result.isError().orElse(false);

    out(
        "[mcp_tool_result] "
            + (call == null ? "unknown" : call.server() + " / " + call.tool())
            + " / success="
            + !failed);
    result.content().ifPresent(content -> out(toPrettyJson(content)));

    return !failed
        && call != null
        && "gmail".equals(call.server())
        && "create_draft".equals(call.tool());
  }

  private static String toPrettyJson(Object value) {
    try {
      return JSON.writeValueAsString(value);
    } catch (JsonProcessingException | IllegalArgumentException exception) {
      return String.valueOf(value);
    }
  }

  private static void out() {
    out("");
  }

  private static void out(String text) {
    System.out.println(text);
    writeText(text + System.lineSeparator());
  }

  private static void writeText(String text) {
    if (outputWriter == null) {
      return;
    }

    try {
      outputWriter.write(text);
      outputWriter.flush();
    } catch (IOException exception) {
      throw new IllegalStateException(resolveOutputPath() + " への書き込みに失敗しました", exception);
    }
  }

  private static Path resolveOutputPath() {
    String configured = System.getenv("OUTPUT_FILE");
    return Strings.isNullOrEmpty(configured) ? DEFAULT_OUTPUT_PATH : Path.of(configured);
  }

  private static void createParentDirectories(Path outputPath) throws IOException {
    Path parent = outputPath.toAbsolutePath().getParent();
    if (parent != null) {
      Files.createDirectories(parent);
    }
  }

  private static String requiredEnvironment(String name) {
    String value = System.getenv(name);
    checkState(!Strings.isNullOrEmpty(value) && !value.isBlank(), "%s を設定してください", name);
    return value;
  }
}