diff --git a/.github/workflows/mysql-session-service-integration.yml b/.github/workflows/mysql-session-service-integration.yml new file mode 100644 index 000000000..59e1f8fa8 --- /dev/null +++ b/.github/workflows/mysql-session-service-integration.yml @@ -0,0 +1,48 @@ +name: MySQL session service integration + +on: + pull_request: + branches: + - main + paths: + - "contrib/mysql-session-service/**" + - "core/**" + - "pom.xml" + - ".mvn/**" + - ".github/workflows/mysql-session-service-integration.yml" + - "mvnw" + - "mvnw.cmd" + +concurrency: + group: ${{ github.workflow }}-${{ github.head_ref || github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + mysql-integration: + runs-on: ubuntu-24.04 + timeout-minutes: 20 + strategy: + matrix: + mysql-image: ["mysql:8.0", "mysql:8.4"] + + steps: + - name: Checkout code + uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 + with: + persist-credentials: false + + - name: Set up Java 17 + uses: actions/setup-java@b6effb05e454b25005698d916606bdc6ffcbf961 # v5 + with: + distribution: temurin + java-version: "17" + cache: maven + + - name: Install ADK core without tests + run: ./mvnw -pl core -am -DskipTests install + + - name: Run MySQL unit and container integration tests + run: ./mvnw -pl contrib/mysql-session-service -Pintegration-tests -Dmysql.image=${{ matrix.mysql-image }} verify diff --git a/contrib/mysql-session-service/README.md b/contrib/mysql-session-service/README.md new file mode 100644 index 000000000..b33dc764f --- /dev/null +++ b/contrib/mysql-session-service/README.md @@ -0,0 +1,117 @@ +# MySQL Session Service for ADK + +This module provides a MySQL-backed implementation of ADK's `BaseSessionService`. It persists +sessions, events, and app- and user-scoped state while keeping the MySQL JDBC driver and connection +pool under the application's control. + +## Dependency + +Use the same version for ADK core and this module. The MySQL driver and a production connection pool +are intentionally not forced on consumers; add the implementations used by your application. + + + +```xml + + 1.12.0 + +``` + + + +```xml + + + com.google.adk + google-adk + ${adk.version} + + + com.google.adk + google-adk-mysql-session-service + ${adk.version} + + + com.mysql + mysql-connector-j + 9.6.0 + + + com.zaxxer + HikariCP + 7.0.2 + + +``` + +The version above tracks ADK releases; until this contribution is released, build this module and +ADK core together from the PR branch and use its snapshot version. Connector/J 9.6 retains support +for MySQL 8.0. Applications using newer drivers should check their minimum supported server version. + +## Database schema + +MySQL 8.0 or newer is required. Apply the packaged +[`mysql-session-service-schema.sql`](src/main/resources/mysql-session-service-schema.sql) before +constructing the service. + +The session identity is the ADK `SessionKey` tuple `(app_name, user_id, session_id)`, so the same +client-generated session ID can safely be reused in a different app or user scope. Events use an +internal monotonic sequence to provide deterministic ordering when multiple events have the same +millisecond timestamp. + +The schema in earlier preview revisions of this contribution used globally unique `session_id` and +`event_id` primary keys. Recreate or migrate those tables before upgrading to this schema. + +## Usage + +Configure a `DataSource`, construct the service, and pass it to the current `Runner` builder. + +```java +import com.google.adk.runner.Runner; +import com.google.adk.sessions.MySqlSessionService; +import com.zaxxer.hikari.HikariConfig; +import com.zaxxer.hikari.HikariDataSource; + +HikariConfig config = new HikariConfig(); +config.setJdbcUrl("jdbc:mysql://localhost:3306/my_db"); +config.setUsername("user"); +config.setPassword("password"); +HikariDataSource dataSource = new HikariDataSource(config); + +MySqlSessionService sessionService = new MySqlSessionService(dataSource); +Runner runner = + Runner.builder() + .agent(myAgent) + .appName("my-app") + .sessionService(sessionService) + .build(); +``` + +App-prefixed (`app:`) and user-prefixed (`user:`) values supplied when creating a session are +persisted in their respective shared state scopes. Temporary (`temp:`) values are returned on the +newly created session but are intentionally not persisted and therefore do not appear after a +reload. + +`appendEvent` is lazy: subscribe to the returned `Single` (or use `blockingGet`) to persist the event. +It updates the caller's session state and event list only after the database transaction commits. +A stale session fails with `ConcurrentModificationException`; reload it before retrying. + +## Testing + +Unit and contract tests use H2 in MySQL compatibility mode and require no external services: + +```bash +./mvnw -pl contrib/mysql-session-service -am test +``` + +The integration profile runs the `*IT` tests against a real MySQL container and therefore requires +Docker: + +```bash +./mvnw -pl contrib/mysql-session-service -am -Pintegration-tests verify +``` + +The same contract tests run against H2 and MySQL, including event JSON round trips, scoped state, +timestamp precision, and rollback without mutating the caller's session. The container defaults to +MySQL 8.4; select MySQL 8.0 with `-Dmysql.image=mysql:8.0`. A running Docker Engine compatible with +Testcontainers 2.x is required. diff --git a/contrib/mysql-session-service/pom.xml b/contrib/mysql-session-service/pom.xml new file mode 100644 index 000000000..5d872c58d --- /dev/null +++ b/contrib/mysql-session-service/pom.xml @@ -0,0 +1,143 @@ + + + + 4.0.0 + + + com.google.adk + google-adk-parent + 1.12.1-SNAPSHOT + ../../pom.xml + + + google-adk-mysql-session-service + Agent Development Kit - MySQL Session Management + MySQL integration with Agent Development Kit for User Session Management + + + 2.0.5 + + 9.6.0 + mysql:8.4 + "-javaagent:${org.mockito:mockito-core:jar}" + + + + + + org.testcontainers + testcontainers-bom + ${testcontainers.version} + pom + import + + + + + + + com.google.adk + google-adk + ${project.version} + + + com.mysql + mysql-connector-j + ${mysql-connector.version} + provided + + + + + com.google.truth + truth + test + + + org.mockito + mockito-core + test + + + org.junit.jupiter + junit-jupiter-api + test + + + org.junit.jupiter + junit-jupiter-engine + test + + + + org.testcontainers + testcontainers-mysql + test + + + org.testcontainers + testcontainers-junit-jupiter + test + + + com.zaxxer + HikariCP + test + + + org.slf4j + slf4j-simple + test + + + com.h2database + h2 + test + + + + + + integration-tests + + + + org.apache.maven.plugins + maven-failsafe-plugin + 3.5.2 + + ${surefire.argLine} + + ${mysql.image} + + + + + + integration-test + verify + + + + + + + + + diff --git a/contrib/mysql-session-service/src/main/java/com/google/adk/sessions/MySqlDBHelper.java b/contrib/mysql-session-service/src/main/java/com/google/adk/sessions/MySqlDBHelper.java new file mode 100644 index 000000000..4217aa8bb --- /dev/null +++ b/contrib/mysql-session-service/src/main/java/com/google/adk/sessions/MySqlDBHelper.java @@ -0,0 +1,557 @@ +/* + * Copyright 2025 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.adk.sessions; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import com.google.adk.JsonBaseModel; +import com.google.adk.events.Event; +import io.reactivex.rxjava3.core.Completable; +import io.reactivex.rxjava3.core.Maybe; +import io.reactivex.rxjava3.core.Single; +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Timestamp; +import java.time.Instant; +import java.time.temporal.ChronoUnit; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.ConcurrentHashMap; +import javax.sql.DataSource; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** Helper class for MySQL database operations. */ +public class MySqlDBHelper { + + private static final Logger logger = LoggerFactory.getLogger(MySqlDBHelper.class); + private static final String LEGACY_REMOVED_SENTINEL = "__ADK_SENTINEL_REMOVED__"; + private static final int EVENT_FORMAT_VERSION = 2; + + private static final String UPSERT_APP_STATE_SQL = + "INSERT INTO adk_app_state (app_name, state_key, state_value) VALUES (?, ?, ?)" + + " ON DUPLICATE KEY UPDATE state_value = ?"; + private static final String UPSERT_USER_STATE_SQL = + "INSERT INTO adk_user_state (app_name, user_id, state_key, state_value) VALUES" + + " (?, ?, ?, ?) ON DUPLICATE KEY UPDATE state_value = ?"; + private static final String DELETE_APP_STATE_SQL = + "DELETE FROM adk_app_state WHERE app_name = ? AND state_key = ?"; + private static final String DELETE_USER_STATE_SQL = + "DELETE FROM adk_user_state WHERE app_name = ? AND user_id = ? AND state_key = ?"; + + private final DataSource dataSource; + private final ObjectMapper objectMapper; + + public MySqlDBHelper(DataSource dataSource) { + this.dataSource = dataSource; + this.objectMapper = JsonBaseModel.getMapper(); + } + + /** Saves a session and any app- or user-scoped values supplied in its initial state. */ + public Completable saveSession(Session session) { + return Completable.fromAction( + () -> { + String insertSessionSql = + "INSERT INTO adk_sessions (app_name, user_id, session_id, state, created_at," + + " updated_at) VALUES (?, ?, ?, ?, ?, ?)"; + Connection conn = null; + try { + conn = dataSource.getConnection(); + conn.setAutoCommit(false); + + try (PreparedStatement pstmt = conn.prepareStatement(insertSessionSql)) { + String stateJson = objectMapper.writeValueAsString(localState(session.state())); + pstmt.setString(1, session.appName()); + pstmt.setString(2, session.userId()); + pstmt.setString(3, session.id()); + pstmt.setString(4, stateJson); + pstmt.setTimestamp(5, Timestamp.from(session.lastUpdateTime())); + pstmt.setTimestamp(6, Timestamp.from(session.lastUpdateTime())); + pstmt.executeUpdate(); + } + + applyGlobalStateChanges(conn, session.appName(), session.userId(), session.state()); + conn.commit(); + } catch (Exception e) { + rollback(conn, e); + if (e instanceof SQLException sqlException && isDuplicateKey(sqlException)) { + throw new SessionException(SessionException.SESSION_ALREADY_EXISTS, sqlException); + } + if (e instanceof SessionException sessionException) { + throw sessionException; + } + throw new SessionException("Failed to save session", e); + } finally { + closeTransactionConnection(conn); + } + }); + } + + /** Persists an event and all resulting state changes in one transaction. */ + public Completable appendEventAndUpdateState(Session session, Event event) { + return Completable.fromAction( + () -> { + String insertEventSql = + "INSERT INTO adk_events (app_name, user_id, session_id, event_id, event_data," + + " created_at) VALUES (?, ?, ?, ?, ?, ?)"; + String updateSessionSql = + "UPDATE adk_sessions SET updated_at = ?, state = ? WHERE app_name = ? AND user_id =" + + " ? AND session_id = ? AND updated_at = ?"; + + Connection conn = null; + try { + conn = dataSource.getConnection(); + conn.setAutoCommit(false); + + try (PreparedStatement pstmt = conn.prepareStatement(insertEventSql)) { + pstmt.setString(1, session.appName()); + pstmt.setString(2, session.userId()); + pstmt.setString(3, session.id()); + pstmt.setString(4, event.id()); + pstmt.setString(5, serializeEvent(event)); + pstmt.setTimestamp(6, Timestamp.from(Instant.ofEpochMilli(event.timestamp()))); + pstmt.executeUpdate(); + } + + Map stateDelta = + event.actions() == null ? Map.of() : event.actions().stateDelta(); + applyGlobalStateChanges(conn, session.appName(), session.userId(), stateDelta); + + Map localState = localState(session.state()); + applyLocalStateChanges(localState, stateDelta); + Instant newLastUpdateTime = nextUpdateTime(session.lastUpdateTime()); + try (PreparedStatement pstmt = conn.prepareStatement(updateSessionSql)) { + pstmt.setTimestamp(1, Timestamp.from(newLastUpdateTime)); + pstmt.setString(2, objectMapper.writeValueAsString(localState)); + pstmt.setString(3, session.appName()); + pstmt.setString(4, session.userId()); + pstmt.setString(5, session.id()); + pstmt.setTimestamp(6, Timestamp.from(session.lastUpdateTime())); + if (pstmt.executeUpdate() == 0) { + throw new java.util.ConcurrentModificationException( + "Session has been modified by another transaction"); + } + } + + conn.commit(); + session.lastUpdateTime(newLastUpdateTime); + } catch (Exception e) { + rollback(conn, e); + if (e instanceof java.util.ConcurrentModificationException concurrentModification) { + throw concurrentModification; + } + if (e instanceof SessionException sessionException) { + throw sessionException; + } + throw new SessionException("Failed to append event and update state", e); + } finally { + closeTransactionConnection(conn); + } + }); + } + + public Maybe getSession( + String appName, String userId, String sessionId, Optional configOpt) { + return Maybe.fromCallable( + () -> { + String sessionSql = + "SELECT session_id, app_name, user_id, updated_at, state FROM adk_sessions WHERE" + + " app_name = ? AND user_id = ? AND session_id = ?"; + GetSessionConfig config = configOpt.orElseGet(() -> GetSessionConfig.builder().build()); + + try (Connection conn = dataSource.getConnection(); + PreparedStatement sessionPstmt = conn.prepareStatement(sessionSql)) { + sessionPstmt.setString(1, appName); + sessionPstmt.setString(2, userId); + sessionPstmt.setString(3, sessionId); + + Session session; + try (ResultSet rs = sessionPstmt.executeQuery()) { + if (!rs.next()) { + return null; + } + Map mergedState = loadGlobalState(conn, appName, userId); + mergeNonNullValues(mergedState, readState(rs.getString("state"))); + session = sessionFromRow(rs, mergedState, new ArrayList<>()); + } + + List events = new ArrayList<>(); + loadEvents(conn, appName, userId, sessionId, config, events); + session.addEvents(events); + return session; + } catch (SQLException | JsonProcessingException e) { + throw new SessionException("Failed to get session", e); + } + }); + } + + public Single> listEvents(String appName, String userId, String sessionId) { + return Single.fromCallable( + () -> { + String sql = + "SELECT event_data FROM adk_events WHERE app_name = ? AND user_id = ? AND" + + " session_id = ? ORDER BY created_at ASC, event_sequence ASC"; + List events = new ArrayList<>(); + try (Connection conn = dataSource.getConnection(); + PreparedStatement pstmt = conn.prepareStatement(sql)) { + pstmt.setString(1, appName); + pstmt.setString(2, userId); + pstmt.setString(3, sessionId); + try (ResultSet rs = pstmt.executeQuery()) { + while (rs.next()) { + events.add(deserializeEvent(rs.getString("event_data"))); + } + } + } catch (SQLException | JsonProcessingException e) { + throw new SessionException("Failed to list events", e); + } + return events; + }); + } + + public Single> listSessions(String appName, String userId) { + return Single.fromCallable( + () -> { + String sql = + "SELECT session_id, app_name, user_id, updated_at, state FROM adk_sessions WHERE" + + " app_name = ? AND user_id = ? ORDER BY created_at DESC, session_id ASC"; + List sessions = new ArrayList<>(); + try (Connection conn = dataSource.getConnection()) { + Map globalState = loadGlobalState(conn, appName, userId); + try (PreparedStatement pstmt = conn.prepareStatement(sql)) { + pstmt.setString(1, appName); + pstmt.setString(2, userId); + try (ResultSet rs = pstmt.executeQuery()) { + while (rs.next()) { + Map mergedState = new HashMap<>(globalState); + mergeNonNullValues(mergedState, readState(rs.getString("state"))); + sessions.add(sessionFromRow(rs, mergedState, new ArrayList<>())); + } + } + } + } catch (SQLException | JsonProcessingException e) { + throw new SessionException("Failed to list sessions", e); + } + return sessions; + }); + } + + public Single> getInitialState(String appName, String userId) { + return Single.fromCallable( + () -> { + try (Connection conn = dataSource.getConnection()) { + return new ConcurrentHashMap<>(loadGlobalState(conn, appName, userId)); + } catch (SQLException | JsonProcessingException e) { + throw new SessionException("Failed to get initial state", e); + } + }); + } + + public Completable deleteSession(String appName, String userId, String sessionId) { + return Completable.fromAction( + () -> { + String sql = + "DELETE FROM adk_sessions WHERE app_name = ? AND user_id = ? AND session_id = ?"; + try (Connection conn = dataSource.getConnection(); + PreparedStatement pstmt = conn.prepareStatement(sql)) { + pstmt.setString(1, appName); + pstmt.setString(2, userId); + pstmt.setString(3, sessionId); + pstmt.executeUpdate(); + } catch (SQLException e) { + throw new SessionException("Failed to delete session", e); + } + }); + } + + private void loadEvents( + Connection conn, + String appName, + String userId, + String sessionId, + GetSessionConfig config, + List destination) + throws SQLException, JsonProcessingException { + StringBuilder sql = + new StringBuilder( + "SELECT event_data FROM adk_events WHERE app_name = ? AND user_id = ? AND" + + " session_id = ?"); + config.afterTimestamp().ifPresent(unused -> sql.append(" AND created_at >= ?")); + if (config.numRecentEvents().isPresent()) { + sql.append(" ORDER BY created_at DESC, event_sequence DESC LIMIT ?"); + } else { + sql.append(" ORDER BY created_at ASC, event_sequence ASC"); + } + + try (PreparedStatement pstmt = conn.prepareStatement(sql.toString())) { + int index = 1; + pstmt.setString(index++, appName); + pstmt.setString(index++, userId); + pstmt.setString(index++, sessionId); + if (config.afterTimestamp().isPresent()) { + // Stored events have millisecond precision. Round the inclusive cursor up before JDBC + // can round or truncate its finer precision and accidentally include an earlier event. + Instant cursor = config.afterTimestamp().get(); + Instant roundedCursor = cursor.truncatedTo(ChronoUnit.MILLIS); + if (cursor.isAfter(roundedCursor)) { + roundedCursor = roundedCursor.plusMillis(1); + } + pstmt.setTimestamp(index++, Timestamp.from(roundedCursor)); + } + if (config.numRecentEvents().isPresent()) { + pstmt.setInt(index, config.numRecentEvents().get()); + } + try (ResultSet rs = pstmt.executeQuery()) { + while (rs.next()) { + destination.add(deserializeEvent(rs.getString("event_data"))); + } + } + } + if (config.numRecentEvents().isPresent()) { + Collections.reverse(destination); + } + } + + private Map loadGlobalState(Connection conn, String appName, String userId) + throws SQLException, JsonProcessingException { + Map state = new HashMap<>(); + try (PreparedStatement pstmt = + conn.prepareStatement( + "SELECT state_key, state_value FROM adk_app_state WHERE app_name = ?")) { + pstmt.setString(1, appName); + loadStateRows(pstmt, State.APP_PREFIX, state); + } + try (PreparedStatement pstmt = + conn.prepareStatement( + "SELECT state_key, state_value FROM adk_user_state WHERE app_name = ? AND user_id =" + + " ?")) { + pstmt.setString(1, appName); + pstmt.setString(2, userId); + loadStateRows(pstmt, State.USER_PREFIX, state); + } + return state; + } + + private void loadStateRows( + PreparedStatement pstmt, String prefix, Map destination) + throws SQLException, JsonProcessingException { + try (ResultSet rs = pstmt.executeQuery()) { + while (rs.next()) { + Object value = objectMapper.readValue(rs.getString("state_value"), Object.class); + if (value != null) { + destination.put(prefix + rs.getString("state_key"), value); + } + } + } + } + + private void applyGlobalStateChanges( + Connection conn, String appName, String userId, Map stateDelta) + throws SQLException, JsonProcessingException { + if (stateDelta == null || stateDelta.isEmpty()) { + return; + } + try (PreparedStatement upsertApp = conn.prepareStatement(UPSERT_APP_STATE_SQL); + PreparedStatement upsertUser = conn.prepareStatement(UPSERT_USER_STATE_SQL); + PreparedStatement deleteApp = conn.prepareStatement(DELETE_APP_STATE_SQL); + PreparedStatement deleteUser = conn.prepareStatement(DELETE_USER_STATE_SQL)) { + for (Map.Entry entry : stateDelta.entrySet()) { + String key = entry.getKey(); + Object value = entry.getValue(); + if (key.startsWith(State.APP_PREFIX)) { + String dbKey = key.substring(State.APP_PREFIX.length()); + if (isRemoved(value)) { + deleteApp.setString(1, appName); + deleteApp.setString(2, dbKey); + deleteApp.addBatch(); + } else { + String valueJson = objectMapper.writeValueAsString(value); + upsertApp.setString(1, appName); + upsertApp.setString(2, dbKey); + upsertApp.setString(3, valueJson); + upsertApp.setString(4, valueJson); + upsertApp.addBatch(); + } + } else if (key.startsWith(State.USER_PREFIX)) { + String dbKey = key.substring(State.USER_PREFIX.length()); + if (isRemoved(value)) { + deleteUser.setString(1, appName); + deleteUser.setString(2, userId); + deleteUser.setString(3, dbKey); + deleteUser.addBatch(); + } else { + String valueJson = objectMapper.writeValueAsString(value); + upsertUser.setString(1, appName); + upsertUser.setString(2, userId); + upsertUser.setString(3, dbKey); + upsertUser.setString(4, valueJson); + upsertUser.setString(5, valueJson); + upsertUser.addBatch(); + } + } + } + upsertApp.executeBatch(); + upsertUser.executeBatch(); + deleteApp.executeBatch(); + deleteUser.executeBatch(); + } + } + + private String serializeEvent(Event event) throws JsonProcessingException { + JsonNode eventJson = objectMapper.valueToTree(event); + JsonNode stateDelta = eventJson.path("actions").path("stateDelta"); + if (stateDelta instanceof ObjectNode stateDeltaObject && event.actions() != null) { + event + .actions() + .stateDelta() + .forEach( + (key, value) -> { + if (value == State.REMOVED) { + stateDeltaObject.putNull(key); + } + }); + } + ObjectNode storedEvent = objectMapper.createObjectNode(); + storedEvent.put("formatVersion", EVENT_FORMAT_VERSION); + storedEvent.set("event", eventJson); + return objectMapper.writeValueAsString(storedEvent); + } + + private Event deserializeEvent(String eventJson) throws JsonProcessingException { + JsonNode storedEvent = objectMapper.readTree(eventJson); + boolean currentFormat = + storedEvent.path("formatVersion").asInt() == EVENT_FORMAT_VERSION + && storedEvent.has("event"); + JsonNode serializedEvent = currentFormat ? storedEvent.get("event") : storedEvent; + Event event = objectMapper.treeToValue(serializedEvent, Event.class); + if (!currentFormat && event.actions() != null) { + event + .actions() + .stateDelta() + .replaceAll( + (unused, value) -> LEGACY_REMOVED_SENTINEL.equals(value) ? State.REMOVED : value); + } + return event; + } + + @SuppressWarnings("unchecked") + private Map readState(String stateJson) throws JsonProcessingException { + return stateJson == null ? Map.of() : objectMapper.readValue(stateJson, Map.class); + } + + private static Session sessionFromRow(ResultSet rs, Map state, List events) + throws SQLException { + return Session.builder(rs.getString("session_id")) + .appName(rs.getString("app_name")) + .userId(rs.getString("user_id")) + .state(new State(state)) + .events(events) + .lastUpdateTime(rs.getTimestamp("updated_at").toInstant()) + .build(); + } + + private static Map localState(Map state) { + Map localState = new HashMap<>(); + state.forEach( + (key, value) -> { + if (!isGlobalOrTemporary(key) && !isRemoved(value)) { + localState.put(key, value); + } + }); + return localState; + } + + private static void applyLocalStateChanges( + Map localState, Map stateDelta) { + if (stateDelta == null) { + return; + } + stateDelta.forEach( + (key, value) -> { + if (!isGlobalOrTemporary(key)) { + if (isRemoved(value)) { + localState.remove(key); + } else { + localState.put(key, value); + } + } + }); + } + + private static boolean isGlobalOrTemporary(String key) { + return key.startsWith(State.APP_PREFIX) + || key.startsWith(State.USER_PREFIX) + || key.startsWith(State.TEMP_PREFIX); + } + + private static boolean isRemoved(Object value) { + return value == null || value == State.REMOVED; + } + + private static void mergeNonNullValues( + Map destination, Map values) { + values.forEach( + (key, value) -> { + if (value != null && value != State.REMOVED) { + destination.put(key, value); + } + }); + } + + private static Instant nextUpdateTime(Instant previousUpdateTime) { + Instant now = Instant.now().truncatedTo(ChronoUnit.MILLIS); + return now.isAfter(previousUpdateTime) ? now : previousUpdateTime.plusMillis(1); + } + + private static boolean isDuplicateKey(SQLException e) { + return e.getErrorCode() == 1062 || "23505".equals(e.getSQLState()); + } + + private static void rollback(Connection conn, Exception original) { + if (conn == null) { + return; + } + try { + conn.rollback(); + } catch (SQLException rollbackFailure) { + original.addSuppressed(rollbackFailure); + } + } + + private static void closeTransactionConnection(Connection conn) { + if (conn == null) { + return; + } + try { + conn.setAutoCommit(true); + } catch (SQLException e) { + logger.warn("Failed to reset auto-commit", e); + } + try { + conn.close(); + } catch (SQLException e) { + logger.warn("Failed to close connection", e); + } + } +} diff --git a/contrib/mysql-session-service/src/main/java/com/google/adk/sessions/MySqlSessionService.java b/contrib/mysql-session-service/src/main/java/com/google/adk/sessions/MySqlSessionService.java new file mode 100644 index 000000000..79a0fcf47 --- /dev/null +++ b/contrib/mysql-session-service/src/main/java/com/google/adk/sessions/MySqlSessionService.java @@ -0,0 +1,174 @@ +/* + * Copyright 2025 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.adk.sessions; + +import com.google.adk.events.Event; +import io.reactivex.rxjava3.core.Completable; +import io.reactivex.rxjava3.core.Maybe; +import io.reactivex.rxjava3.core.Single; +import java.time.Instant; +import java.time.temporal.ChronoUnit; +import java.util.ArrayList; +import java.util.Map; +import java.util.Optional; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import javax.sql.DataSource; +import org.jspecify.annotations.Nullable; + +/** A MySQL-backed implementation of {@link BaseSessionService}. */ +public class MySqlSessionService implements BaseSessionService { + + private final MySqlDBHelper dbHelper; + + public MySqlSessionService(DataSource dataSource) { + this.dbHelper = new MySqlDBHelper(dataSource); + } + + // For testing purposes + MySqlSessionService(MySqlDBHelper dbHelper) { + this.dbHelper = dbHelper; + } + + @Deprecated + @Override + public Single createSession( + String appName, + String userId, + @Nullable ConcurrentMap state, + @Nullable String sessionId) { + return createSession(appName, userId, (Map) state, sessionId); + } + + @Override + public Single createSession( + String appName, + String userId, + @Nullable Map state, + @Nullable String sessionId) { + return validateAppAndUser(appName, userId) + .andThen(createNewSession(appName, userId, state, sessionId)); + } + + private Single createNewSession( + String appName, + String userId, + @Nullable Map state, + @Nullable String sessionId) { + ConcurrentMap sessionState = + (state != null) ? new ConcurrentHashMap<>(state) : new ConcurrentHashMap<>(); + String newSessionId = + (sessionId == null || sessionId.trim().isEmpty()) + ? UUID.randomUUID().toString() + : sessionId; + + return dbHelper + .getInitialState(appName, userId) + .flatMap( + globalState -> { + // Start with global state (defaults) + ConcurrentHashMap mergedState = new ConcurrentHashMap<>(globalState); + // Overlay initial session state (overrides) + mergedState.putAll(sessionState); + + Session session = + Session.builder(newSessionId) + .appName(appName) + .userId(userId) + .state(new State(mergedState)) + .events(new ArrayList<>()) + .lastUpdateTime(Instant.now().truncatedTo(ChronoUnit.MILLIS)) + .build(); + + return dbHelper.saveSession(session).toSingleDefault(session); + }); + } + + @Override + public Maybe getSession( + String appName, String userId, String sessionId, Optional config) { + if (config == null) { + return Maybe.error(new IllegalArgumentException("config cannot be null.")); + } + return validateAppAndUser(appName, userId) + .andThen(validateSessionId(sessionId)) + .andThen(dbHelper.getSession(appName, userId, sessionId, config)); + } + + @Override + public Single listSessions(String appName, String userId) { + return validateAppAndUser(appName, userId) + .andThen( + dbHelper + .listSessions(appName, userId) + .map(sessions -> ListSessionsResponse.builder().sessions(sessions).build())); + } + + @Override + public Completable deleteSession(String appName, String userId, String sessionId) { + return validateAppAndUser(appName, userId) + .andThen(validateSessionId(sessionId)) + .andThen(dbHelper.deleteSession(appName, userId, sessionId)); + } + + @Override + public Single listEvents(String appName, String userId, String sessionId) { + return validateAppAndUser(appName, userId) + .andThen(validateSessionId(sessionId)) + .andThen( + dbHelper + .listEvents(appName, userId, sessionId) + .map(events -> ListEventsResponse.builder().events(events).build())); + } + + @Override + public Single appendEvent(Session session, Event event) { + if (session == null) { + return Single.error(new IllegalArgumentException("session cannot be null.")); + } + if (event == null) { + return Single.error(new IllegalArgumentException("event cannot be null.")); + } + + return validateAppAndUser(session.appName(), session.userId()) + .andThen(validateSessionId(session.id())) + .andThen( + event.partial().orElse(false) + ? Single.just(event) + : dbHelper + .appendEventAndUpdateState(session, event) + .andThen( + Single.defer(() -> BaseSessionService.super.appendEvent(session, event)))); + } + + private Completable validateAppAndUser(String appName, String userId) { + if (appName == null || appName.trim().isEmpty()) { + return Completable.error(new IllegalArgumentException("appName cannot be null or empty.")); + } + if (userId == null || userId.trim().isEmpty()) { + return Completable.error(new IllegalArgumentException("userId cannot be null or empty.")); + } + return Completable.complete(); + } + + private Completable validateSessionId(String sessionId) { + if (sessionId == null || sessionId.trim().isEmpty()) { + return Completable.error(new IllegalArgumentException("sessionId cannot be null or empty.")); + } + return Completable.complete(); + } +} diff --git a/contrib/mysql-session-service/src/main/resources/mysql-session-service-schema.sql b/contrib/mysql-session-service/src/main/resources/mysql-session-service-schema.sql new file mode 100644 index 000000000..57bb581b3 --- /dev/null +++ b/contrib/mysql-session-service/src/main/resources/mysql-session-service-schema.sql @@ -0,0 +1,40 @@ +CREATE TABLE IF NOT EXISTS adk_sessions ( + app_name VARCHAR(255) NOT NULL, + user_id VARCHAR(255) NOT NULL, + session_id VARCHAR(255) NOT NULL, + state JSON NOT NULL, + created_at TIMESTAMP(3) NOT NULL, + updated_at TIMESTAMP(3) NOT NULL, + PRIMARY KEY (app_name, user_id, session_id), + INDEX idx_adk_sessions_updated (app_name, user_id, updated_at) +); + +CREATE TABLE IF NOT EXISTS adk_events ( + event_sequence BIGINT NOT NULL AUTO_INCREMENT PRIMARY KEY, + app_name VARCHAR(255) NOT NULL, + user_id VARCHAR(255) NOT NULL, + session_id VARCHAR(255) NOT NULL, + event_id VARCHAR(255) NOT NULL, + event_data JSON NOT NULL, + created_at TIMESTAMP(3) NOT NULL, + CONSTRAINT fk_adk_events_session + FOREIGN KEY (app_name, user_id, session_id) + REFERENCES adk_sessions (app_name, user_id, session_id) + ON DELETE CASCADE, + INDEX idx_adk_events_session (app_name, user_id, session_id) +); + +CREATE TABLE IF NOT EXISTS adk_app_state ( + app_name VARCHAR(255) NOT NULL, + state_key VARCHAR(255) NOT NULL, + state_value JSON NOT NULL, + PRIMARY KEY (app_name, state_key) +); + +CREATE TABLE IF NOT EXISTS adk_user_state ( + app_name VARCHAR(255) NOT NULL, + user_id VARCHAR(255) NOT NULL, + state_key VARCHAR(255) NOT NULL, + state_value JSON NOT NULL, + PRIMARY KEY (app_name, user_id, state_key) +); diff --git a/contrib/mysql-session-service/src/test/java/com/google/adk/sessions/MySqlDBHelperTest.java b/contrib/mysql-session-service/src/test/java/com/google/adk/sessions/MySqlDBHelperTest.java new file mode 100644 index 000000000..16cb83bf3 --- /dev/null +++ b/contrib/mysql-session-service/src/test/java/com/google/adk/sessions/MySqlDBHelperTest.java @@ -0,0 +1,394 @@ +/* + * Copyright 2025 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.adk.sessions; + +import static com.google.common.truth.Truth.assertThat; + +import com.google.adk.events.Event; +import com.google.adk.events.EventActions; +import com.zaxxer.hikari.HikariConfig; +import com.zaxxer.hikari.HikariDataSource; +import io.reactivex.rxjava3.core.Single; +import java.sql.Connection; +import java.sql.Statement; +import java.time.Instant; +import java.time.temporal.ChronoUnit; +import java.util.ArrayList; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import javax.sql.DataSource; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** + * Unit tests for {@link MySqlDBHelper} using H2 in MySQL compatibility mode. This ensures SQL + * syntax is valid without requiring a real MySQL container. + */ +class MySqlDBHelperTest { + + private DataSource dataSource; + private MySqlDBHelper dbHelper; + + @BeforeEach + void setUp() throws Exception { + HikariConfig config = new HikariConfig(); + config.setJdbcUrl("jdbc:h2:mem:testdb;MODE=MySQL;DB_CLOSE_DELAY=-1"); + config.setUsername("sa"); + config.setPassword(""); + dataSource = new HikariDataSource(config); + + MySqlTestSchema.reset(dataSource); + + dbHelper = new MySqlDBHelper(dataSource); + } + + @AfterEach + void tearDown() { + if (dataSource instanceof HikariDataSource) { + ((HikariDataSource) dataSource).close(); + } + } + + @Test + void appendEvent_withAppState_executesSqlSuccessfully() { + // Create a session first + Session session = + Session.builder(UUID.randomUUID().toString()) + .appName("testApp") + .userId("testUser") + .state(new ConcurrentHashMap<>()) + .events(new ArrayList<>()) + .lastUpdateTime(Instant.now().truncatedTo(ChronoUnit.MILLIS)) + .build(); + + dbHelper.saveSession(session).blockingAwait(); + + // Create event with App State delta + ConcurrentMap stateDelta = new ConcurrentHashMap<>(); + stateDelta.put(State.APP_PREFIX + "someKey", "someValue"); + + Event event = + Event.builder() + .id(UUID.randomUUID().toString()) + .timestamp(System.currentTimeMillis()) + .actions(EventActions.builder().stateDelta(stateDelta).build()) + .build(); + + // This triggers the upsertAppStateSql. + // If there is a parameter mismatch (3 columns vs 4 placeholders), this will fail. + dbHelper.appendEventAndUpdateState(session, event).blockingAwait(); + + // Verify it was inserted + Single> stateSingle = + dbHelper.getInitialState("testApp", "testUser"); + ConcurrentHashMap state = stateSingle.blockingGet(); + assertThat(state).containsEntry(State.APP_PREFIX + "someKey", "someValue"); + } + + @Test + void appendEvent_withUserState_executesSqlSuccessfully() { + // Create a session + Session session = + Session.builder(UUID.randomUUID().toString()) + .appName("testApp") + .userId("testUser") + .state(new ConcurrentHashMap<>()) + .events(new ArrayList<>()) + .lastUpdateTime(Instant.now().truncatedTo(ChronoUnit.MILLIS)) + .build(); + + dbHelper.saveSession(session).blockingAwait(); + + // Create event with User State delta + ConcurrentMap stateDelta = new ConcurrentHashMap<>(); + stateDelta.put(State.USER_PREFIX + "someUserKey", "someUserValue"); + + Event event = + Event.builder() + .id(UUID.randomUUID().toString()) + .timestamp(System.currentTimeMillis()) + .actions(EventActions.builder().stateDelta(stateDelta).build()) + .build(); + + dbHelper.appendEventAndUpdateState(session, event).blockingAwait(); + + // Verify it was inserted + ConcurrentHashMap state = + dbHelper.getInitialState("testApp", "testUser").blockingGet(); + assertThat(state).containsEntry(State.USER_PREFIX + "someUserKey", "someUserValue"); + } + + @Test + void listSessions_retrievesSessionState() { + ConcurrentMap state = new ConcurrentHashMap<>(); + state.put("key1", "value1"); + + Session session = + Session.builder(UUID.randomUUID().toString()) + .appName("testApp") + .userId("testUser") + .state(state) + .events(new ArrayList<>()) + .lastUpdateTime(Instant.now().truncatedTo(ChronoUnit.MILLIS)) + .build(); + + dbHelper.saveSession(session).blockingAwait(); + + java.util.List sessions = dbHelper.listSessions("testApp", "testUser").blockingGet(); + + assertThat(sessions).hasSize(1); + assertThat(sessions.get(0).state()).containsEntry("key1", "value1"); + } + + @Test + void listEvents_enforcesOwnership() { + String sessionId = UUID.randomUUID().toString(); + Session session = + Session.builder(sessionId) + .appName("correctApp") + .userId("correctUser") + .state(new ConcurrentHashMap<>()) + .events(new ArrayList<>()) + .lastUpdateTime(Instant.now().truncatedTo(ChronoUnit.MILLIS)) + .build(); + + dbHelper.saveSession(session).blockingAwait(); + + Event event = + Event.builder() + .id(UUID.randomUUID().toString()) + .timestamp(System.currentTimeMillis()) + .build(); + + dbHelper.appendEventAndUpdateState(session, event).blockingAwait(); + + // 1. Correct credentials should return the event + java.util.List events = + dbHelper.listEvents("correctApp", "correctUser", sessionId).blockingGet(); + assertThat(events).hasSize(1); + + // 2. Wrong App Name should return empty + java.util.List eventsWrongApp = + dbHelper.listEvents("wrongApp", "correctUser", sessionId).blockingGet(); + assertThat(eventsWrongApp).isEmpty(); + + // 3. Wrong User ID should return empty + java.util.List eventsWrongUser = + dbHelper.listEvents("correctApp", "wrongUser", sessionId).blockingGet(); + assertThat(eventsWrongUser).isEmpty(); + } + + @Test + void getSession_mergesStateCorrectly() throws Exception { + // 1. Setup Global App State directly in DB + try (Connection conn = dataSource.getConnection(); + Statement stmt = conn.createStatement()) { + stmt.execute( + "INSERT INTO adk_app_state (app_name, state_key, state_value) VALUES ('testApp', 'sharedKey', '\"globalValue\"')"); + stmt.execute( + "INSERT INTO adk_app_state (app_name, state_key, state_value) VALUES ('testApp', 'globalOnly', '\"globalOnlyValue\"')"); + } + + // 2. Create Session with conflicting key + ConcurrentMap sessionState = new ConcurrentHashMap<>(); + sessionState.put("sharedKey", "sessionValue"); + + Session session = + Session.builder(UUID.randomUUID().toString()) + .appName("testApp") + .userId("testUser") + .state(sessionState) + .events(new ArrayList<>()) + .lastUpdateTime(Instant.now().truncatedTo(ChronoUnit.MILLIS)) + .build(); + + dbHelper.saveSession(session).blockingAwait(); + + // 3. Retrieve Session + Session retrievedSession = + dbHelper + .getSession("testApp", "testUser", session.id(), java.util.Optional.empty()) + .blockingGet(); + + // 4. Verify Merge Logic + // Session-specific value should override Global value + assertThat(retrievedSession.state()).containsEntry("sharedKey", "sessionValue"); + // Non-conflicting Global value should be present with prefix + assertThat(retrievedSession.state()) + .containsEntry(State.APP_PREFIX + "globalOnly", "globalOnlyValue"); + // The shared key from global state should also be present with prefix + assertThat(retrievedSession.state()) + .containsEntry(State.APP_PREFIX + "sharedKey", "globalValue"); + } + + @Test + void deleteSession_executesSqlSuccessfully() { + Session session = + Session.builder(UUID.randomUUID().toString()) + .appName("testApp") + .userId("testUser") + .state(new ConcurrentHashMap<>()) + .events(new ArrayList<>()) + .lastUpdateTime(Instant.now().truncatedTo(ChronoUnit.MILLIS)) + .build(); + + dbHelper.saveSession(session).blockingAwait(); + + dbHelper.deleteSession("testApp", "testUser", session.id()).blockingAwait(); + + Session deletedSession = + dbHelper + .getSession("testApp", "testUser", session.id(), java.util.Optional.empty()) + .blockingGet(); + assertThat(deletedSession).isNull(); + } + + @Test + public void listSessions_mergesGlobalState() throws Exception { + // 1. Setup Global App State directly in DB + try (Connection conn = dataSource.getConnection(); + Statement stmt = conn.createStatement()) { + stmt.execute( + "INSERT INTO adk_app_state (app_name, state_key, state_value) VALUES ('testApp', 'globalKey', '\"globalValue\"')"); + stmt.execute( + "INSERT INTO adk_user_state (app_name, user_id, state_key, state_value) VALUES ('testApp', 'testUser', 'userKey', '\"userValue\"')"); + } + + Session session = + Session.builder(UUID.randomUUID().toString()) + .appName("testApp") + .userId("testUser") + .state(new ConcurrentHashMap<>()) + .events(new ArrayList<>()) + .lastUpdateTime(Instant.now().truncatedTo(ChronoUnit.MILLIS)) + .build(); + + dbHelper.saveSession(session).blockingAwait(); + + java.util.List sessions = dbHelper.listSessions("testApp", "testUser").blockingGet(); + + assertThat(sessions).hasSize(1); + assertThat(sessions.get(0).state()) + .containsEntry(State.APP_PREFIX + "globalKey", "globalValue"); + assertThat(sessions.get(0).state()).containsEntry(State.USER_PREFIX + "userKey", "userValue"); + } + + @Test + void prefixHandling_storesCleanKeysInDb() throws Exception { + Session session = + Session.builder(UUID.randomUUID().toString()) + .appName("testApp") + .userId("testUser") + .state(new ConcurrentHashMap<>()) + .events(new ArrayList<>()) + .lastUpdateTime(Instant.now().truncatedTo(ChronoUnit.MILLIS)) + .build(); + + dbHelper.saveSession(session).blockingAwait(); + + ConcurrentMap stateDelta = new ConcurrentHashMap<>(); + stateDelta.put(State.APP_PREFIX + "cleanAppKey", "val1"); + stateDelta.put(State.USER_PREFIX + "cleanUserKey", "val2"); + + Event event = + Event.builder() + .id(UUID.randomUUID().toString()) + .timestamp(System.currentTimeMillis()) + .actions(EventActions.builder().stateDelta(stateDelta).build()) + .build(); + + dbHelper.appendEventAndUpdateState(session, event).blockingAwait(); + + // Verify DB storage directly + try (Connection conn = dataSource.getConnection(); + Statement stmt = conn.createStatement()) { + + // Check App State + try (java.sql.ResultSet rs = + stmt.executeQuery( + "SELECT state_key FROM adk_app_state WHERE app_name = 'testApp' AND state_key = 'cleanAppKey'")) { + assertThat(rs.next()).isTrue(); + assertThat(rs.getString("state_key")).isEqualTo("cleanAppKey"); // Should NOT have prefix + } + + // Check User State + try (java.sql.ResultSet rs = + stmt.executeQuery( + "SELECT state_key FROM adk_user_state WHERE app_name = 'testApp' AND user_id = 'testUser' AND state_key = 'cleanUserKey'")) { + assertThat(rs.next()).isTrue(); + assertThat(rs.getString("state_key")).isEqualTo("cleanUserKey"); // Should NOT have prefix + } + } + } + + @Test + void prefixHandling_removesKeysFromDb() throws Exception { + // 1. Setup existing state in DB + try (Connection conn = dataSource.getConnection(); + Statement stmt = conn.createStatement()) { + stmt.execute( + "INSERT INTO adk_app_state (app_name, state_key, state_value) VALUES ('testApp', 'keyToRemove', '\"val\"')"); + stmt.execute( + "INSERT INTO adk_user_state (app_name, user_id, state_key, state_value) VALUES ('testApp', 'testUser', 'keyToRemove', '\"val\"')"); + } + + Session session = + Session.builder(UUID.randomUUID().toString()) + .appName("testApp") + .userId("testUser") + .state(new ConcurrentHashMap<>()) + .events(new ArrayList<>()) + .lastUpdateTime(Instant.now().truncatedTo(ChronoUnit.MILLIS)) + .build(); + + dbHelper.saveSession(session).blockingAwait(); + + // 2. Create event with REMOVED delta + ConcurrentMap stateDelta = new ConcurrentHashMap<>(); + stateDelta.put(State.APP_PREFIX + "keyToRemove", State.REMOVED); + stateDelta.put(State.USER_PREFIX + "keyToRemove", State.REMOVED); + + Event event = + Event.builder() + .id(UUID.randomUUID().toString()) + .timestamp(System.currentTimeMillis()) + .actions(EventActions.builder().stateDelta(stateDelta).build()) + .build(); + + dbHelper.appendEventAndUpdateState(session, event).blockingAwait(); + + // 3. Verify deletion in DB + try (Connection conn = dataSource.getConnection(); + Statement stmt = conn.createStatement()) { + + // Check App State + try (java.sql.ResultSet rs = + stmt.executeQuery( + "SELECT state_key FROM adk_app_state WHERE app_name = 'testApp' AND state_key = 'keyToRemove'")) { + assertThat(rs.next()).isFalse(); // Should be gone + } + + // Check User State + try (java.sql.ResultSet rs = + stmt.executeQuery( + "SELECT state_key FROM adk_user_state WHERE app_name = 'testApp' AND user_id = 'testUser' AND state_key = 'keyToRemove'")) { + assertThat(rs.next()).isFalse(); // Should be gone + } + } + } +} diff --git a/contrib/mysql-session-service/src/test/java/com/google/adk/sessions/MySqlSessionServiceContract.java b/contrib/mysql-session-service/src/test/java/com/google/adk/sessions/MySqlSessionServiceContract.java new file mode 100644 index 000000000..0f4cc92f4 --- /dev/null +++ b/contrib/mysql-session-service/src/test/java/com/google/adk/sessions/MySqlSessionServiceContract.java @@ -0,0 +1,583 @@ +/* + * Copyright 2025 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.adk.sessions; + +import static com.google.common.truth.Truth.assertThat; + +import com.google.adk.JsonBaseModel; +import com.google.adk.events.Event; +import com.google.adk.events.EventActions; +import com.google.adk.events.EventCompaction; +import com.google.adk.events.ToolConfirmation; +import com.google.genai.types.Blob; +import com.google.genai.types.Content; +import com.google.genai.types.FunctionCall; +import com.google.genai.types.FunctionResponse; +import com.google.genai.types.Part; +import com.zaxxer.hikari.HikariDataSource; +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.time.Instant; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import javax.sql.DataSource; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** Contract tests aligned with the behavior of {@link InMemorySessionService}. */ +abstract class MySqlSessionServiceContract { + + // MySQL TIMESTAMP cannot represent the first second after the Unix epoch. + private static final long EVENT_EPOCH = Instant.parse("2026-01-01T00:00:00Z").toEpochMilli(); + + protected DataSource dataSource; + protected MySqlSessionService sessionService; + + protected abstract DataSource createDataSource(); + + @BeforeEach + void setUp() throws Exception { + dataSource = createDataSource(); + MySqlTestSchema.reset(dataSource); + sessionService = new MySqlSessionService(dataSource); + } + + @AfterEach + void tearDown() { + closeDataSource(); + } + + protected void closeDataSource() { + if (dataSource instanceof HikariDataSource hikariDataSource) { + hikariDataSource.close(); + } + } + + @Test + void sameSessionIdCanExistInDifferentScopes() { + Session first = + sessionService.createSession("app-1", "user-1", Map.of(), "shared").blockingGet(); + Session second = + sessionService.createSession("app-2", "user-1", Map.of(), "shared").blockingGet(); + Session third = + sessionService.createSession("app-1", "user-2", Map.of(), "shared").blockingGet(); + + sessionService.appendEvent(first, event("first", 100, Map.of("value", 1))).blockingGet(); + sessionService.appendEvent(second, event("second", 100, Map.of("value", 2))).blockingGet(); + sessionService.appendEvent(third, event("third", 100, Map.of("value", 3))).blockingGet(); + + assertThat(get(first).state()).containsEntry("value", 1); + assertThat(get(second).state()).containsEntry("value", 2); + assertThat(get(third).state()).containsEntry("value", 3); + + sessionService.deleteSession("app-1", "user-1", "shared").blockingAwait(); + assertThat(getOrNull("app-1", "user-1", "shared")).isNull(); + assertThat(get(second)).isNotNull(); + assertThat(get(third)).isNotNull(); + } + + @Test + void initialScopedStateIsPersistedAndTemporaryStateIsEphemeral() { + Map initialState = new HashMap<>(); + initialState.put("local", "local-value"); + initialState.put(State.APP_PREFIX + "setting", "app-value"); + initialState.put(State.USER_PREFIX + "setting", "user-value"); + initialState.put(State.TEMP_PREFIX + "request", "temporary-value"); + + Session created = + sessionService.createSession("app", "user", initialState, "session-1").blockingGet(); + assertThat(created.state()).containsEntry(State.TEMP_PREFIX + "request", "temporary-value"); + + Session reloaded = get(created); + assertThat(reloaded.state()).containsEntry("local", "local-value"); + assertThat(reloaded.state()).containsEntry(State.APP_PREFIX + "setting", "app-value"); + assertThat(reloaded.state()).containsEntry(State.USER_PREFIX + "setting", "user-value"); + assertThat(reloaded.state()).doesNotContainKey(State.TEMP_PREFIX + "request"); + + Session another = + sessionService.createSession("app", "user", Map.of(), "session-2").blockingGet(); + assertThat(another.state()).containsEntry(State.APP_PREFIX + "setting", "app-value"); + assertThat(another.state()).containsEntry(State.USER_PREFIX + "setting", "user-value"); + assertThat(another.state()).doesNotContainKey("local"); + } + + @Test + void appStateIsSharedAcrossUsersButUserStateIsIsolated() { + sessionService + .createSession( + "app", + "user-1", + Map.of( + State.APP_PREFIX + "shared", "app-value", + State.USER_PREFIX + "private", "user-value"), + "first") + .blockingGet(); + + Session anotherUser = + sessionService.createSession("app", "user-2", Map.of(), "second").blockingGet(); + assertThat(anotherUser.state()).containsEntry(State.APP_PREFIX + "shared", "app-value"); + assertThat(anotherUser.state()).doesNotContainKey(State.USER_PREFIX + "private"); + + Session anotherApp = + sessionService.createSession("other-app", "user-1", Map.of(), "third").blockingGet(); + assertThat(anotherApp.state()).doesNotContainKey(State.APP_PREFIX + "shared"); + assertThat(anotherApp.state()).doesNotContainKey(State.USER_PREFIX + "private"); + } + + @Test + void getSessionComposesRecentAndInclusiveTimestampFilters() { + Session session = sessionService.createSession("app", "user").blockingGet(); + for (long timestamp : new long[] {100, 200, 300, 400, 500}) { + sessionService + .appendEvent(session, event("event-" + timestamp, timestamp, Map.of())) + .blockingGet(); + } + GetSessionConfig config = + GetSessionConfig.builder() + .numRecentEvents(4) + .afterTimestamp(Instant.ofEpochMilli(EVENT_EPOCH + 300)) + .build(); + + Session retrieved = + sessionService.getSession("app", "user", session.id(), Optional.of(config)).blockingGet(); + + assertThat(retrieved.immutableEvents().stream().map(Event::timestamp)) + .containsExactly(EVENT_EPOCH + 300, EVENT_EPOCH + 400, EVENT_EPOCH + 500) + .inOrder(); + } + + @Test + void equalTimestampsRetainInsertionOrder() { + Session session = sessionService.createSession("app", "user").blockingGet(); + for (String id : List.of("first", "second", "third")) { + sessionService.appendEvent(session, event(id, 1000, Map.of())).blockingGet(); + } + + assertThat( + sessionService.listEvents("app", "user", session.id()).blockingGet().events().stream() + .map(Event::id)) + .containsExactly("first", "second", "third") + .inOrder(); + } + + @Test + void removedStateRoundTripsAsSentinelAndUsesJsonNull() throws Exception { + Session session = + sessionService + .createSession( + "app", + "user", + Map.of( + "local", + "value", + State.APP_PREFIX + "app-key", + "value", + State.USER_PREFIX + "user-key", + "value"), + "session") + .blockingGet(); + ConcurrentMap removals = new ConcurrentHashMap<>(); + removals.put("local", State.REMOVED); + removals.put(State.APP_PREFIX + "app-key", State.REMOVED); + removals.put(State.USER_PREFIX + "user-key", State.REMOVED); + sessionService.appendEvent(session, event("removed", 1000, removals)).blockingGet(); + + Event persisted = + sessionService.listEvents("app", "user", "session").blockingGet().events().get(0); + assertThat(persisted.actions().stateDelta().get("local")).isSameInstanceAs(State.REMOVED); + assertThat(persisted.actions().stateDelta().get(State.APP_PREFIX + "app-key")) + .isSameInstanceAs(State.REMOVED); + assertThat(get(session).state()).doesNotContainKey("local"); + + String eventJson = readOnlyEventJson(); + assertThat(eventJson).contains("\"local\":null"); + assertThat(eventJson).doesNotContain("__ADK_SENTINEL_REMOVED__"); + } + + @Test + void legacyRemovedSentinelStillDeserializes() throws Exception { + Session session = sessionService.createSession("app", "user").blockingGet(); + ConcurrentMap removals = new ConcurrentHashMap<>(); + removals.put("legacy", State.REMOVED); + sessionService.appendEvent(session, event("legacy", 1000, removals)).blockingGet(); + + String currentJson = readOnlyEventJson(); + String directEventJson = + JsonBaseModel.getMapper() + .writeValueAsString(JsonBaseModel.getMapper().readTree(currentJson).get("event")); + String legacyJson = + directEventJson.replace("\"legacy\":null", "\"legacy\":\"__ADK_SENTINEL_REMOVED__\""); + assertThat(legacyJson).isNotEqualTo(directEventJson); + try (Connection conn = dataSource.getConnection(); + PreparedStatement pstmt = + conn.prepareStatement("UPDATE adk_events SET event_data = ? WHERE event_id = ?")) { + pstmt.setString(1, legacyJson); + pstmt.setString(2, "legacy"); + pstmt.executeUpdate(); + } + + Event persisted = + sessionService.listEvents("app", "user", session.id()).blockingGet().events().get(0); + assertThat(persisted.actions().stateDelta().get("legacy")).isSameInstanceAs(State.REMOVED); + } + + @Test + void literalLegacySentinelStringIsPreservedInCurrentFormat() { + Session session = sessionService.createSession("app", "user").blockingGet(); + sessionService + .appendEvent(session, event("literal", 1000, Map.of("literal", "__ADK_SENTINEL_REMOVED__"))) + .blockingGet(); + + Event persisted = + sessionService.listEvents("app", "user", session.id()).blockingGet().events().get(0); + assertThat(persisted.actions().stateDelta()) + .containsEntry("literal", "__ADK_SENTINEL_REMOVED__"); + } + + @Test + void partialEventIsNotPersistedOrApplied() { + Session session = sessionService.createSession("app", "user").blockingGet(); + Event partial = + Event.builder() + .id("partial") + .partial(true) + .actions(EventActions.builder().stateDelta(Map.of("key", "value")).build()) + .build(); + + sessionService.appendEvent(session, partial).blockingGet(); + + assertThat(session.immutableEvents()).isEmpty(); + assertThat(session.state()).doesNotContainKey("key"); + assertThat(sessionService.listEvents("app", "user", session.id()).blockingGet().events()) + .isEmpty(); + } + + @Test + void staleSessionFailsOptimisticLockAndRollsBackEvent() { + Session created = sessionService.createSession("app", "user").blockingGet(); + Session firstWriter = get(created); + Session staleWriter = get(created); + Instant previousUpdateTime = staleWriter.lastUpdateTime(); + + sessionService + .appendEvent(firstWriter, event("committed", 1000, Map.of("key", "committed"))) + .blockingGet(); + sessionService + .appendEvent( + staleWriter, + event("rolled-back", 2000, Map.of("key", "stale", "user:rolled-back", "stale"))) + .test() + .assertError(java.util.ConcurrentModificationException.class); + + Session reloaded = get(created); + assertThat(reloaded.state()).containsEntry("key", "committed"); + assertThat(reloaded.immutableEvents().stream().map(Event::id)).containsExactly("committed"); + assertThat(reloaded.state()).doesNotContainKey("user:rolled-back"); + assertThat(staleWriter.state()).doesNotContainKey("key"); + assertThat(staleWriter.state()).doesNotContainKey("user:rolled-back"); + assertThat(staleWriter.immutableEvents()).isEmpty(); + assertThat(staleWriter.lastUpdateTime()).isEqualTo(previousUpdateTime); + } + + @Test + void appendEventDoesNotMutateOrPersistUntilSubscribed() { + Session session = sessionService.createSession("app", "user").blockingGet(); + Instant previousUpdateTime = session.lastUpdateTime(); + Event event = event("deferred", 1000, Map.of("key", "value")); + + var pending = sessionService.appendEvent(session, event); + + assertThat(session.immutableEvents()).isEmpty(); + assertThat(session.state()).doesNotContainKey("key"); + assertThat(session.lastUpdateTime()).isEqualTo(previousUpdateTime); + assertThat(get(session).immutableEvents()).isEmpty(); + + assertThat(pending.blockingGet()).isSameInstanceAs(event); + assertThat(session.immutableEvents()).containsExactly(event); + assertThat(session.state()).containsEntry("key", "value"); + assertThat(get(session).immutableEvents()).hasSize(1); + } + + @Test + void duplicateSessionUsesCoreErrorContractWithoutChangingScopedState() { + Session session = + sessionService + .createSession("app", "user", Map.of("user:name", "original"), "same") + .blockingGet(); + + sessionService + .createSession("app", "user", Map.of("user:name", "replacement"), "same") + .test() + .assertError( + error -> + error instanceof SessionException + && SessionException.SESSION_ALREADY_EXISTS.equals(error.getMessage())); + + assertThat(get(session).state()).containsEntry("user:name", "original"); + } + + @Test + void afterTimestampKeepsExactBoundaryAndExcludesEarlierMillisecondForFinerCursors() { + Session session = sessionService.createSession("app", "user").blockingGet(); + sessionService.appendEvent(session, event("boundary", 1400, Map.of())).blockingGet(); + sessionService.appendEvent(session, event("later", 1401, Map.of())).blockingGet(); + + for (long extraNanos : new long[] {0, 1, 999, 1000, 300_000, 999_999}) { + GetSessionConfig config = + GetSessionConfig.builder() + .afterTimestamp(Instant.ofEpochMilli(EVENT_EPOCH + 1400).plusNanos(extraNanos)) + .build(); + Session filtered = + sessionService.getSession("app", "user", session.id(), Optional.of(config)).blockingGet(); + assertThat(filtered.immutableEvents().stream().map(Event::id)) + .containsExactlyElementsIn( + extraNanos == 0 ? List.of("boundary", "later") : List.of("later")) + .inOrder(); + } + } + + @Test + void wholeSecondAndFractionalEventsKeepTimeOrderWithStableTies() { + Session session = sessionService.createSession("app", "user").blockingGet(); + sessionService.appendEvent(session, event("whole", 1000, Map.of())).blockingGet(); + sessionService.appendEvent(session, event("fraction", 1400, Map.of())).blockingGet(); + sessionService.appendEvent(session, event("tie", 1400, Map.of())).blockingGet(); + sessionService.appendEvent(session, event("next", 2000, Map.of())).blockingGet(); + + assertThat(get(session).immutableEvents().stream().map(Event::id)) + .containsExactly("whole", "fraction", "tie", "next") + .inOrder(); + GetSessionConfig config = GetSessionConfig.builder().numRecentEvents(2).build(); + Session recent = + sessionService.getSession("app", "user", session.id(), Optional.of(config)).blockingGet(); + assertThat(recent.immutableEvents().stream().map(Event::id)) + .containsExactly("tie", "next") + .inOrder(); + } + + @Test + void listEventsDoesNotExposeAnotherUserOrAppsHistory() { + Session session = sessionService.createSession("app", "user").blockingGet(); + sessionService.appendEvent(session, event("private", 1000, Map.of())).blockingGet(); + + assertThat(sessionService.listEvents("app", "other-user", session.id()).blockingGet().events()) + .isEmpty(); + assertThat(sessionService.listEvents("other-app", "user", session.id()).blockingGet().events()) + .isEmpty(); + } + + @Test + void sharedMapValuesAreReplacedAndDottedKeysRemainLiteral() { + Session first = + sessionService + .createSession( + "app", "user", Map.of("user:prefs", Map.of("lang", "ko", "theme", "dark")), "first") + .blockingGet(); + Session second = sessionService.createSession("app", "user", Map.of(), "second").blockingGet(); + sessionService + .appendEvent( + first, + event( + "replace", + 1000, + Map.of( + "user:prefs", + Map.of("lang", "en"), + "user:settings.v2", + Map.of("enabled", true)))) + .blockingGet(); + + assertThat(get(second).state()).containsEntry("user:prefs", Map.of("lang", "en")); + assertThat(get(second).state()).containsEntry("user:settings.v2", Map.of("enabled", true)); + } + + @Test + void deletedSharedValuesAndTemporaryContextStateDoNotReappearAfterLocalWrite() { + Session first = + sessionService + .createSession("app", "user", Map.of("app:theme", "dark", "user:name", "Jae"), "first") + .blockingGet(); + Session second = sessionService.createSession("app", "user", Map.of(), "second").blockingGet(); + second.state().put("temp:context", "ephemeral"); + sessionService.appendEvent(second, event("local", 1000, Map.of("counter", 1))).blockingGet(); + sessionService + .appendEvent( + first, + event("remove", 2000, Map.of("app:theme", State.REMOVED, "user:name", State.REMOVED))) + .blockingGet(); + + Session reloaded = get(second); + assertThat(reloaded.state()).containsEntry("counter", 1); + assertThat(reloaded.state()).doesNotContainKey("app:theme"); + assertThat(reloaded.state()).doesNotContainKey("user:name"); + assertThat(reloaded.state()).doesNotContainKey("temp:context"); + assertThat( + sessionService.listSessions("app", "user").blockingGet().sessions().stream() + .allMatch(s -> !s.state().containsKey("user:name"))) + .isTrue(); + } + + @Test + void fullEventPreservesIdsAuthorSignaturesAndInlineDataOnReload() throws Exception { + Session session = sessionService.createSession("app", "account-42").blockingGet(); + Event original = + Event.builder() + .id("stable-id") + .invocationId("invocation") + .author("user") + .branch("root.child") + .timestamp(EVENT_EPOCH + 1000) + .longRunningToolIds(Set.of("call-1")) + .actions( + EventActions.builder() + .requestedToolConfirmations( + Map.of("call-1", ToolConfirmation.builder().hint("Approve?").build())) + .build()) + .content( + Content.fromParts( + Part.builder() + .text("thought") + .thought(true) + .thoughtSignature(new byte[] {1, 2, 3}) + .build(), + Part.builder() + .functionCall( + FunctionCall.builder() + .id("call-1") + .name("tool") + .args(Map.of("input", "value")) + .build()) + .build(), + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("call-1") + .name("tool") + .response(Map.of("result", "value")) + .build()) + .build(), + Part.builder() + .inlineData( + Blob.builder().data(new byte[] {4, 5, 6}).mimeType("image/png").build()) + .build())) + .build(); + sessionService.appendEvent(session, original).blockingGet(); + + assertSameEventJson(original, get(session).immutableEvents().get(0)); + assertSameEventJson( + original, + sessionService.listEvents("app", "account-42", session.id()).blockingGet().events().get(0)); + } + + @Test + void toolConfirmationWithStructuredArgsAndCompactionSurviveReload() throws Exception { + Session session = sessionService.createSession("app", "user").blockingGet(); + Event confirmation = + event("confirmation", 1000, Map.of()).toBuilder() + .content( + Content.fromParts( + Part.builder() + .functionCall( + FunctionCall.builder() + .id("confirm-1") + .name("adk_request_confirmation") + .args( + Map.of( + "originalFunctionCall", + FunctionCall.builder().id("call-1").name("tool").build(), + "toolConfirmation", + ToolConfirmation.builder().hint("Approve?").build())) + .build()) + .build())) + .build(); + Event compaction = + event("compaction", 2000, Map.of()).toBuilder() + .actions( + EventActions.builder() + .compaction( + EventCompaction.builder() + .startTimestamp(EVENT_EPOCH) + .endTimestamp(EVENT_EPOCH + 1000) + .compactedContent(Content.fromParts(Part.fromText("summary"))) + .build()) + .build()) + .build(); + sessionService.appendEvent(session, confirmation).blockingGet(); + sessionService.appendEvent(session, compaction).blockingGet(); + + List reloaded = get(session).immutableEvents(); + assertSameEventJson(confirmation, reloaded.get(0)); + assertSameEventJson(compaction, reloaded.get(1)); + } + + private static void assertSameEventJson(Event expected, Event actual) throws Exception { + assertThat(JsonBaseModel.getMapper().readTree(actual.toJson())) + .isEqualTo(JsonBaseModel.getMapper().readTree(expected.toJson())); + } + + @Test + void duplicateEventIdsAreAllowedLikeInMemoryService() { + Session session = sessionService.createSession("app", "user").blockingGet(); + sessionService.appendEvent(session, event("duplicate", 1000, Map.of())).blockingGet(); + sessionService.appendEvent(session, event("duplicate", 2000, Map.of())).blockingGet(); + + assertThat(sessionService.listEvents("app", "user", session.id()).blockingGet().events()) + .hasSize(2); + } + + @Test + void missingSessionOperationsAreNoOpsOrEmpty() { + assertThat(getOrNull("app", "user", "missing")).isNull(); + assertThat(sessionService.listEvents("app", "user", "missing").blockingGet().events()) + .isEmpty(); + sessionService.deleteSession("app", "user", "missing").blockingAwait(); + } + + private Session get(Session session) { + return sessionService + .getSession(session.appName(), session.userId(), session.id(), Optional.empty()) + .blockingGet(); + } + + private Session getOrNull(String appName, String userId, String sessionId) { + return sessionService.getSession(appName, userId, sessionId, Optional.empty()).blockingGet(); + } + + private String readOnlyEventJson() throws Exception { + try (Connection conn = dataSource.getConnection(); + PreparedStatement pstmt = + conn.prepareStatement("SELECT event_data FROM adk_events ORDER BY event_sequence"); + ResultSet rs = pstmt.executeQuery()) { + assertThat(rs.next()).isTrue(); + return rs.getString(1); + } + } + + private static Event event(String id, long timestamp, Map stateDelta) { + return Event.builder() + .id(id) + .timestamp(EVENT_EPOCH + timestamp) + .actions(EventActions.builder().stateDelta(stateDelta).build()) + .build(); + } +} diff --git a/contrib/mysql-session-service/src/test/java/com/google/adk/sessions/MySqlSessionServiceContractTest.java b/contrib/mysql-session-service/src/test/java/com/google/adk/sessions/MySqlSessionServiceContractTest.java new file mode 100644 index 000000000..2505acff6 --- /dev/null +++ b/contrib/mysql-session-service/src/test/java/com/google/adk/sessions/MySqlSessionServiceContractTest.java @@ -0,0 +1,33 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.adk.sessions; + +import com.zaxxer.hikari.HikariConfig; +import com.zaxxer.hikari.HikariDataSource; +import javax.sql.DataSource; + +/** Runs the shared service contract against H2 without an external database. */ +class MySqlSessionServiceContractTest extends MySqlSessionServiceContract { + + @Override + protected DataSource createDataSource() { + HikariConfig config = new HikariConfig(); + config.setJdbcUrl("jdbc:h2:mem:contractdb;MODE=MySQL;DB_CLOSE_DELAY=-1"); + config.setUsername("sa"); + config.setPassword(""); + return new HikariDataSource(config); + } +} diff --git a/contrib/mysql-session-service/src/test/java/com/google/adk/sessions/MySqlSessionServiceIT.java b/contrib/mysql-session-service/src/test/java/com/google/adk/sessions/MySqlSessionServiceIT.java new file mode 100644 index 000000000..01072147b --- /dev/null +++ b/contrib/mysql-session-service/src/test/java/com/google/adk/sessions/MySqlSessionServiceIT.java @@ -0,0 +1,315 @@ +/* + * Copyright 2025 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.adk.sessions; + +import static com.google.common.truth.Truth.assertThat; + +import com.google.adk.events.Event; +import com.google.adk.events.EventActions; +import com.zaxxer.hikari.HikariConfig; +import com.zaxxer.hikari.HikariDataSource; +import io.reactivex.rxjava3.core.Single; +import java.sql.Connection; +import java.sql.Statement; +import java.util.Optional; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import javax.sql.DataSource; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.testcontainers.junit.jupiter.Container; +import org.testcontainers.junit.jupiter.Testcontainers; +import org.testcontainers.mysql.MySQLContainer; + +/** Unit tests for {@link MySqlSessionService} using Testcontainers. */ +@Testcontainers +public final class MySqlSessionServiceIT extends MySqlSessionServiceContract { + + @Container + private static final MySQLContainer mysql = + new MySQLContainer(System.getProperty("mysql.image", "mysql:8.4")) + .withDatabaseName("adk_test"); + + private static HikariDataSource mysqlDataSource; + + @BeforeAll + public static void setUpClass() { + HikariConfig config = new HikariConfig(); + config.setJdbcUrl(mysql.getJdbcUrl()); + config.setUsername(mysql.getUsername()); + config.setPassword(mysql.getPassword()); + mysqlDataSource = new HikariDataSource(config); + } + + @AfterAll + public static void tearDownClass() { + if (mysqlDataSource != null) { + mysqlDataSource.close(); + } + } + + @Override + protected DataSource createDataSource() { + return mysqlDataSource; + } + + @Override + protected void closeDataSource() { + // The shared container pool is closed once, after all integration tests. + } + + @Test + public void lifecycle_noSession() { + assertThat( + sessionService + .getSession("app-name", "user-id", "session-id", Optional.empty()) + .blockingGet()) + .isNull(); + + assertThat(sessionService.listSessions("app-name", "user-id").blockingGet().sessions()) + .isEmpty(); + + assertThat( + sessionService.listEvents("app-name", "user-id", "session-id").blockingGet().events()) + .isEmpty(); + } + + @Test + public void lifecycle_createSession() { + Single sessionSingle = sessionService.createSession("app-name", "user-id"); + + Session session = sessionSingle.blockingGet(); + + assertThat(session.id()).isNotNull(); + assertThat(session.appName()).isEqualTo("app-name"); + assertThat(session.userId()).isEqualTo("user-id"); + assertThat(session.state()).isEmpty(); + } + + @Test + public void lifecycle_getSession() { + Session session = sessionService.createSession("app-name", "user-id").blockingGet(); + + Session retrievedSession = + sessionService + .getSession(session.appName(), session.userId(), session.id(), Optional.empty()) + .blockingGet(); + + assertThat(retrievedSession).isNotNull(); + assertThat(retrievedSession.id()).isEqualTo(session.id()); + } + + @Test + public void lifecycle_listSessions() { + Session session = sessionService.createSession("app-name", "user-id").blockingGet(); + + ListSessionsResponse response = + sessionService.listSessions(session.appName(), session.userId()).blockingGet(); + + assertThat(response.sessions()).hasSize(1); + assertThat(response.sessions().get(0).id()).isEqualTo(session.id()); + } + + @Test + public void lifecycle_deleteSession() { + Session session = sessionService.createSession("app-name", "user-id").blockingGet(); + + sessionService.deleteSession(session.appName(), session.userId(), session.id()).blockingAwait(); + + assertThat( + sessionService + .getSession(session.appName(), session.userId(), session.id(), Optional.empty()) + .blockingGet()) + .isNull(); + } + + @Test + public void appendEvent_updatesSessionState() { + Session session = + sessionService + .createSession("app", "user", new ConcurrentHashMap<>(), "session1") + .blockingGet(); + + ConcurrentMap stateDelta = new ConcurrentHashMap<>(); + stateDelta.put("sessionKey", "sessionValue"); + stateDelta.put(State.APP_PREFIX + "appKey", "appValue"); + stateDelta.put(State.USER_PREFIX + "userKey", "userValue"); + + Event event = + Event.builder() + .id(UUID.randomUUID().toString()) // Assign a unique ID + .actions(EventActions.builder().stateDelta(stateDelta).build()) + .build(); + + sessionService.appendEvent(session, event).blockingGet(); + + assertThat(session.state()).containsEntry("sessionKey", "sessionValue"); + + Session retrievedSession = + sessionService + .getSession(session.appName(), session.userId(), session.id(), Optional.empty()) + .blockingGet(); + + assertThat(retrievedSession.state()).containsEntry("sessionKey", "sessionValue"); + assertThat(retrievedSession.state()).containsEntry(State.APP_PREFIX + "appKey", "appValue"); + assertThat(retrievedSession.state()).containsEntry(State.USER_PREFIX + "userKey", "userValue"); + } + + @Test + public void listSessions_includesGlobalState() { + Session session = + sessionService + .createSession("app", "user", new ConcurrentHashMap<>(), "session1") + .blockingGet(); + + ConcurrentMap stateDelta = new ConcurrentHashMap<>(); + stateDelta.put(State.APP_PREFIX + "appKey", "appValue"); + stateDelta.put(State.USER_PREFIX + "userKey", "userValue"); + + Event event = + Event.builder() + .id(UUID.randomUUID().toString()) + .actions(EventActions.builder().stateDelta(stateDelta).build()) + .build(); + + sessionService.appendEvent(session, event).blockingGet(); + + ListSessionsResponse response = sessionService.listSessions("app", "user").blockingGet(); + + assertThat(response.sessions()).hasSize(1); + Session listedSession = response.sessions().get(0); + + // This asserts that the list operation merged the global state correctly + assertThat(listedSession.state()).containsEntry(State.APP_PREFIX + "appKey", "appValue"); + assertThat(listedSession.state()).containsEntry(State.USER_PREFIX + "userKey", "userValue"); + } + + @Test + public void storesCleanKeysInDb_integration() throws Exception { + Session session = + sessionService + .createSession("app", "user", new ConcurrentHashMap<>(), "session1") + .blockingGet(); + + ConcurrentMap stateDelta = new ConcurrentHashMap<>(); + stateDelta.put(State.APP_PREFIX + "cleanAppKey", "val1"); + stateDelta.put(State.USER_PREFIX + "cleanUserKey", "val2"); + + Event event = + Event.builder() + .id(UUID.randomUUID().toString()) + .actions(EventActions.builder().stateDelta(stateDelta).build()) + .build(); + + sessionService.appendEvent(session, event).blockingGet(); + + // Verify DB storage directly + try (Connection conn = dataSource.getConnection(); + Statement stmt = conn.createStatement()) { + + // Check App State + try (java.sql.ResultSet rs = + stmt.executeQuery( + "SELECT state_key FROM adk_app_state WHERE app_name = 'app' AND state_key =" + + " 'cleanAppKey'")) { + assertThat(rs.next()).isTrue(); + assertThat(rs.getString("state_key")).isEqualTo("cleanAppKey"); // Should NOT have prefix + } + + // Check User State + try (java.sql.ResultSet rs = + stmt.executeQuery( + "SELECT state_key FROM adk_user_state WHERE app_name = 'app' AND user_id = 'user'" + + " AND state_key = 'cleanUserKey'")) { + assertThat(rs.next()).isTrue(); + assertThat(rs.getString("state_key")).isEqualTo("cleanUserKey"); // Should NOT have prefix + } + } + + // Verify retrieval adds prefix back + Session retrievedSession = + sessionService + .getSession(session.appName(), session.userId(), session.id(), Optional.empty()) + .blockingGet(); + + assertThat(retrievedSession.state()).containsEntry(State.APP_PREFIX + "cleanAppKey", "val1"); + assertThat(retrievedSession.state()).containsEntry(State.USER_PREFIX + "cleanUserKey", "val2"); + } + + @Test + public void removesKeysFromDb_integration() throws Exception { + Session session = + sessionService + .createSession("app", "user", new ConcurrentHashMap<>(), "session1") + .blockingGet(); + + // 1. Setup existing state in DB directly + try (Connection conn = dataSource.getConnection(); + Statement stmt = conn.createStatement()) { + stmt.execute( + "INSERT INTO adk_app_state (app_name, state_key, state_value) VALUES ('app'," + + " 'keyToRemove', '\"val\"')"); + stmt.execute( + "INSERT INTO adk_user_state (app_name, user_id, state_key, state_value) VALUES ('app'," + + " 'user', 'keyToRemove', '\"val\"')"); + } + + // 2. Create event with REMOVED delta + ConcurrentMap stateDelta = new ConcurrentHashMap<>(); + stateDelta.put(State.APP_PREFIX + "keyToRemove", State.REMOVED); + stateDelta.put(State.USER_PREFIX + "keyToRemove", State.REMOVED); + + Event event = + Event.builder() + .id(UUID.randomUUID().toString()) + .actions(EventActions.builder().stateDelta(stateDelta).build()) + .build(); + + sessionService.appendEvent(session, event).blockingGet(); + + // 3. Verify deletion in DB + try (Connection conn = dataSource.getConnection(); + Statement stmt = conn.createStatement()) { + + // Check App State + try (java.sql.ResultSet rs = + stmt.executeQuery( + "SELECT state_key FROM adk_app_state WHERE app_name = 'app' AND state_key =" + + " 'keyToRemove'")) { + assertThat(rs.next()).isFalse(); + } + + // Check User State + try (java.sql.ResultSet rs = + stmt.executeQuery( + "SELECT state_key FROM adk_user_state WHERE app_name = 'app' AND user_id = 'user'" + + " AND state_key = 'keyToRemove'")) { + assertThat(rs.next()).isFalse(); + } + } + + // 4. Verify retrieval does not have keys + Session retrievedSession = + sessionService + .getSession(session.appName(), session.userId(), session.id(), Optional.empty()) + .blockingGet(); + + assertThat(retrievedSession.state()).doesNotContainKey(State.APP_PREFIX + "keyToRemove"); + assertThat(retrievedSession.state()).doesNotContainKey(State.USER_PREFIX + "keyToRemove"); + } +} diff --git a/contrib/mysql-session-service/src/test/java/com/google/adk/sessions/MySqlSessionServiceUnitTest.java b/contrib/mysql-session-service/src/test/java/com/google/adk/sessions/MySqlSessionServiceUnitTest.java new file mode 100644 index 000000000..f145d20eb --- /dev/null +++ b/contrib/mysql-session-service/src/test/java/com/google/adk/sessions/MySqlSessionServiceUnitTest.java @@ -0,0 +1,105 @@ +/* + * Copyright 2025 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.adk.sessions; + +import static com.google.common.truth.Truth.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import io.reactivex.rxjava3.core.Completable; +import io.reactivex.rxjava3.core.Single; +import java.util.concurrent.ConcurrentHashMap; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; + +/** + * Unit tests for {@link MySqlSessionService}. These tests use a mock {@link MySqlDBHelper} and do + * not require a running database. + */ +class MySqlSessionServiceUnitTest { + + private MySqlDBHelper mockDbHelper; + private MySqlSessionService sessionService; + + @BeforeEach + void setUp() { + mockDbHelper = mock(MySqlDBHelper.class); + sessionService = new MySqlSessionService(mockDbHelper); + } + + @Test + void createSession_validInput_delegatesToDbHelper() { + String appName = "testApp"; + String userId = "testUser"; + + when(mockDbHelper.saveSession(any(Session.class))).thenReturn(Completable.complete()); + when(mockDbHelper.getInitialState(anyString(), anyString())) + .thenReturn(Single.just(new ConcurrentHashMap<>())); + + sessionService.createSession(appName, userId, null, null).blockingGet(); + + verify(mockDbHelper).saveSession(any(Session.class)); + verify(mockDbHelper).getInitialState(eq(appName), eq(userId)); + } + + @Test + void createSession_mergesGlobalState() { + String appName = "testApp"; + String userId = "testUser"; + + // Global state + ConcurrentHashMap globalState = new ConcurrentHashMap<>(); + globalState.put("_user_pref", "dark_mode"); + globalState.put("shared_key", "global_value"); + + // Initial Session state + ConcurrentHashMap initialSessionState = new ConcurrentHashMap<>(); + initialSessionState.put("shared_key", "session_value"); + + when(mockDbHelper.saveSession(any(Session.class))).thenReturn(Completable.complete()); + when(mockDbHelper.getInitialState(anyString(), anyString())) + .thenReturn(Single.just(globalState)); + + Session session = + sessionService.createSession(appName, userId, initialSessionState, null).blockingGet(); + + // Verify global only key is present + assertThat(session.state()).containsEntry("_user_pref", "dark_mode"); + // Verify session key overrides global key + assertThat(session.state()).containsEntry("shared_key", "session_value"); + + // Verify that saveSession was called WITH the merged state + ArgumentCaptor sessionCaptor = ArgumentCaptor.forClass(Session.class); + verify(mockDbHelper).saveSession(sessionCaptor.capture()); + Session savedSession = sessionCaptor.getValue(); + assertThat(savedSession.state()).containsEntry("_user_pref", "dark_mode"); + assertThat(savedSession.state()).containsEntry("shared_key", "session_value"); + } + + @Test + void getSession_validInput_delegatesToDbHelper() { + when(mockDbHelper.getSession(anyString(), anyString(), anyString(), any())) + .thenReturn(io.reactivex.rxjava3.core.Maybe.empty()); + + sessionService.getSession("app", "user", "session1", java.util.Optional.empty()); + verify(mockDbHelper).getSession("app", "user", "session1", java.util.Optional.empty()); + } +} diff --git a/contrib/mysql-session-service/src/test/java/com/google/adk/sessions/MySqlTestSchema.java b/contrib/mysql-session-service/src/test/java/com/google/adk/sessions/MySqlTestSchema.java new file mode 100644 index 000000000..fe63855a8 --- /dev/null +++ b/contrib/mysql-session-service/src/test/java/com/google/adk/sessions/MySqlTestSchema.java @@ -0,0 +1,58 @@ +/* + * Copyright 2025 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.adk.sessions; + +import java.io.IOException; +import java.io.InputStream; +import java.nio.charset.StandardCharsets; +import java.sql.Connection; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.Objects; +import javax.sql.DataSource; + +final class MySqlTestSchema { + + private MySqlTestSchema() {} + + static void reset(DataSource dataSource) throws SQLException, IOException { + try (Connection conn = dataSource.getConnection(); + Statement stmt = conn.createStatement()) { + stmt.execute("DROP TABLE IF EXISTS adk_events"); + stmt.execute("DROP TABLE IF EXISTS adk_sessions"); + stmt.execute("DROP TABLE IF EXISTS adk_app_state"); + stmt.execute("DROP TABLE IF EXISTS adk_user_state"); + + try (InputStream schemaStream = + Objects.requireNonNull( + MySqlTestSchema.class + .getClassLoader() + .getResourceAsStream("mysql-session-service-schema.sql"))) { + String schema = new String(schemaStream.readAllBytes(), StandardCharsets.UTF_8); + if (conn.getMetaData().getDatabaseProductName().equals("H2")) { + // H2's JSON type JSON-encodes strings passed through PreparedStatement#setString, + // unlike MySQL. TEXT preserves the JDBC behavior exercised by these DAO tests. + schema = schema.replace(" JSON NOT NULL", " TEXT NOT NULL"); + } + for (String sql : schema.split(";")) { + if (!sql.isBlank()) { + stmt.execute(sql); + } + } + } + } + } +} diff --git a/pom.xml b/pom.xml index 16d4d08f2..7f6e1875f 100644 --- a/pom.xml +++ b/pom.xml @@ -32,6 +32,7 @@ contrib/spring-ai contrib/planners contrib/firestore-session-service + contrib/mysql-session-service examples tutorials/city-time-weather tutorials/live-audio-single-agent diff --git a/release-please-config.json b/release-please-config.json index 6b3a1302c..9972a1826 100644 --- a/release-please-config.json +++ b/release-please-config.json @@ -6,6 +6,7 @@ "release-type": "maven", "extra-files": [ "core/src/main/java/com/google/adk/Version.java", - "README.md" + "README.md", + "contrib/mysql-session-service/README.md" ] }