Mid-session USE over TDS no longer hangs SqlClient: the database-change ENVCHANGE + INFO 5701 are now written before the response's final DONE, matching real SQL Server's token order.

RyanLamansky · RyanLamansky · commit 0ff26f30538d · 2026-07-14T10:42:44.000-05:00
diff --git a/SqlServerSimulator.Tests.SqlClient/UseDatabaseTests.cs b/SqlServerSimulator.Tests.SqlClient/UseDatabaseTests.cs
@@ -0,0 +1,126 @@
+using Microsoft.Data.SqlClient;
+using static Microsoft.VisualStudio.TestTools.UnitTesting.Assert;
+
+namespace SqlServerSimulator;
+
+/// <summary>
+/// Mid-session <c>USE</c> over the wire. The database-change ENVCHANGE (+
+/// INFO 5701) must precede the response's final DONE — SqlClient's token
+/// reader stalls until command timeout on an ENVCHANGE that arrives after
+/// the last DONE (probe-confirmed 2026-07-15; go-mssqldb tolerates the late
+/// position, so only a real-SqlClient oracle catches the ordering).
+/// </summary>
+[TestClass]
+public sealed class UseDatabaseTests
+{
+    public TestContext TestContext { get; set; } = null!;
+
+    [TestMethod]
+    public async Task UseMaster_ExecuteNonQuery_SwitchesDatabase()
+    {
+        var simulation = new Simulation();
+        await using var listener = await simulation.ListenAsync(0, TestContext.CancellationToken);
+        await using var connection = await Wire.OpenAsync(listener, TestContext.CancellationToken);
+
+        await using (var use = new SqlCommand("use [master]", connection))
+            _ = await use.ExecuteNonQueryAsync(TestContext.CancellationToken);
+
+        AreEqual("master", connection.Database);
+        await using var query = new SqlCommand("select db_name()", connection);
+        AreEqual("master", await query.ExecuteScalarAsync(TestContext.CancellationToken));
+    }
+
+    [TestMethod]
+    public async Task UseMaster_RaisesChangedDatabaseContextInfoMessage()
+    {
+        var simulation = new Simulation();
+        await using var listener = await simulation.ListenAsync(0, TestContext.CancellationToken);
+        await using var connection = await Wire.OpenAsync(listener, TestContext.CancellationToken);
+
+        var messages = new List<string>();
+        connection.InfoMessage += (_, e) => messages.Add(e.Message);
+        await using var use = new SqlCommand("use [master]", connection);
+        _ = await use.ExecuteNonQueryAsync(TestContext.CancellationToken);
+
+        Contains("Changed database context to 'master'.", messages);
+    }
+
+    [TestMethod]
+    public async Task UseInMultiStatementBatch_WithResultSet()
+    {
+        var simulation = new Simulation();
+        await using var listener = await simulation.ListenAsync(0, TestContext.CancellationToken);
+        await using var connection = await Wire.OpenAsync(listener, TestContext.CancellationToken);
+
+        await using var command = new SqlCommand("use [master] select db_name()", connection);
+        AreEqual("master", await command.ExecuteScalarAsync(TestContext.CancellationToken));
+        AreEqual("master", connection.Database);
+    }
+
+    [TestMethod]
+    public async Task UseViaRpc_ParameterizedCommand_SwitchesDatabase()
+    {
+        var simulation = new Simulation();
+        await using var listener = await simulation.ListenAsync(0, TestContext.CancellationToken);
+        await using var connection = await Wire.OpenAsync(listener, TestContext.CancellationToken);
+
+        // A parameter forces the sp_executesql RPC path instead of SQLBatch.
+        await using var command = new SqlCommand("use [master] select @x", connection);
+        _ = command.Parameters.AddWithValue("@x", 1);
+        AreEqual(1, await command.ExecuteScalarAsync(TestContext.CancellationToken));
+        AreEqual("master", connection.Database);
+    }
+
+    [TestMethod]
+    public async Task UseThenError_StillSwitchesDatabase()
+    {
+        var simulation = new Simulation();
+        await using var listener = await simulation.ListenAsync(0, TestContext.CancellationToken);
+        await using var connection = await Wire.OpenAsync(listener, TestContext.CancellationToken);
+
+        await using (var command = new SqlCommand("use [master] select * from nosuchtable", connection))
+        {
+            var ex = await ThrowsExactlyAsync<SqlException>(() => command.ExecuteScalarAsync(TestContext.CancellationToken));
+            AreEqual(208, ex.Number);
+        }
+
+        AreEqual("master", connection.Database);
+        await using var query = new SqlCommand("select db_name()", connection);
+        AreEqual("master", await query.ExecuteScalarAsync(TestContext.CancellationToken));
+    }
+
+    [TestMethod]
+    public async Task UseRoundTrip_MasterAndBack()
+    {
+        var simulation = new Simulation();
+        await using var listener = await simulation.ListenAsync(0, TestContext.CancellationToken);
+        await using var connection = await Wire.OpenAsync(listener, TestContext.CancellationToken);
+
+        await using (var toMaster = new SqlCommand("use [master]", connection))
+            _ = await toMaster.ExecuteNonQueryAsync(TestContext.CancellationToken);
+        await using (var back = new SqlCommand("use [simulated]", connection))
+            _ = await back.ExecuteNonQueryAsync(TestContext.CancellationToken);
+
+        AreEqual("simulated", connection.Database);
+        await using var query = new SqlCommand("select db_name()", connection);
+        AreEqual("simulated", await query.ExecuteScalarAsync(TestContext.CancellationToken));
+    }
+
+    [TestMethod]
+    public async Task UseMissingDatabase_RaisesMsg911_SessionSurvives()
+    {
+        var simulation = new Simulation();
+        await using var listener = await simulation.ListenAsync(0, TestContext.CancellationToken);
+        await using var connection = await Wire.OpenAsync(listener, TestContext.CancellationToken);
+
+        await using (var command = new SqlCommand("use [nosuchdb]", connection))
+        {
+            var ex = await ThrowsExactlyAsync<SqlException>(() => command.ExecuteNonQueryAsync(TestContext.CancellationToken));
+            AreEqual(911, ex.Number);
+        }
+
+        AreEqual("simulated", connection.Database);
+        await using var query = new SqlCommand("select db_name()", connection);
+        AreEqual("simulated", await query.ExecuteScalarAsync(TestContext.CancellationToken));
+    }
+}
diff --git a/SqlServerSimulator/Network/TdsSession.Rpc.cs b/SqlServerSimulator/Network/TdsSession.Rpc.cs
@@ -46,7 +46,7 @@ private async ValueTask ExecuteRpcMessageAsync(TdsMessage message, TdsTokenWrite
             return;
         }
 
-        var databaseBefore = this.connection!.Database;
+        this.databaseAtMessageStart = this.connection!.Database;
         for (var i = 0; i < requests.Count; i++)
         {
             var moreRequests = i < requests.Count - 1;
@@ -58,19 +58,19 @@ private async ValueTask ExecuteRpcMessageAsync(TdsMessage message, TdsTokenWrite
             {
                 _ = this.FlushInfoMessages(writer);
                 WriteErrors(writer, ex);
+                if (!moreRequests)
+                    this.WriteDatabaseChangeIfAny(writer);
                 writer.WriteDoneToken(Tds.TokenDoneProc, (ushort)(Tds.DoneError | (moreRequests ? Tds.DoneMore : Tds.DoneFinal)), 0);
             }
             catch (NotSupportedException ex)
             {
                 _ = this.FlushInfoMessages(writer);
                 writer.WriteErrorOrInfo(Tds.TokenError, 50000, 1, 16, $"SqlServerSimulator: {ex.Message}", "SIMULATED", "", 1);
+                if (!moreRequests)
+                    this.WriteDatabaseChangeIfAny(writer);
                 writer.WriteDoneToken(Tds.TokenDoneProc, (ushort)(Tds.DoneError | (moreRequests ? Tds.DoneMore : Tds.DoneFinal)), 0);
             }
         }
-
-        var databaseAfter = this.connection.Database;
-        if (!string.Equals(databaseAfter, databaseBefore, StringComparison.Ordinal))
-            writer.WriteEnvChange(Tds.EnvDatabase, databaseAfter, databaseBefore);
     }
 
     private async ValueTask DispatchRpcAsync(TdsRpcRequest request, TdsTokenWriter writer, bool moreRequests, CancellationToken cancellationToken)
@@ -180,6 +180,8 @@ private async ValueTask ExecuteStatementRpcAsync(
         foreach (var (ordinal, wire, bound) in outputs)
             TdsTypeCodec.WriteReturnValue(writer, checked((ushort)ordinal), wire.Name, wire.DbType, bound.Value);
 
+        if (!moreRequests)
+            this.WriteDatabaseChangeIfAny(writer);
         writer.WriteDoneToken(Tds.TokenDoneProc, moreRequests ? Tds.DoneMore : Tds.DoneFinal, 0);
     }
 
@@ -212,6 +214,8 @@ private async ValueTask ExecuteProcedureRpcAsync(TdsRpcRequest request, TdsToken
         foreach (var (ordinal, wire, bound) in outputs)
             TdsTypeCodec.WriteReturnValue(writer, checked((ushort)ordinal), wire.Name, wire.DbType, bound.Value);
 
+        if (!moreRequests)
+            this.WriteDatabaseChangeIfAny(writer);
         writer.WriteDoneToken(Tds.TokenDoneProc, moreRequests ? Tds.DoneMore : Tds.DoneFinal, 0);
     }
 
diff --git a/SqlServerSimulator/Network/TdsSession.cs b/SqlServerSimulator/Network/TdsSession.cs
@@ -206,7 +206,7 @@ private async ValueTask ExecuteBatchAsync(TdsMessage message, TdsTokenWriter wri
             writer.WriteResetConnectionAck();
         }
 
-        var databaseBefore = this.connection!.Database;
+        this.databaseAtMessageStart = this.connection!.Database;
         try
         {
             using var command = this.connection.CreateCommand();
@@ -219,19 +219,46 @@ private async ValueTask ExecuteBatchAsync(TdsMessage message, TdsTokenWriter wri
         {
             _ = this.FlushInfoMessages(writer);
             WriteErrors(writer, ex);
+            this.WriteDatabaseChangeIfAny(writer);
             writer.WriteDone(Tds.DoneError, 0);
         }
         catch (NotSupportedException ex)
         {
             _ = this.FlushInfoMessages(writer);
             writer.WriteErrorOrInfo(Tds.TokenError, 50000, 1, 16, $"SqlServerSimulator: {ex.Message}", "SIMULATED", "", 1);
+            this.WriteDatabaseChangeIfAny(writer);
             writer.WriteDone(Tds.DoneError, 0);
         }
+    }
 
-        _ = this.FlushInfoMessages(writer);
-        var databaseAfter = this.connection.Database;
-        if (!string.Equals(databaseAfter, databaseBefore, StringComparison.Ordinal))
-            writer.WriteEnvChange(Tds.EnvDatabase, databaseAfter, databaseBefore);
+    /// <summary>
+    /// The session database when the current batch / RPC message began, for
+    /// detecting a mid-message <c>USE</c>. Emitted as ENVCHANGE type 1 +
+    /// INFO 5701 via <see cref="WriteDatabaseChangeIfAny"/>, which must run
+    /// BEFORE the response's final DONE: SqlClient's token reader stalls
+    /// until command timeout on an ENVCHANGE that arrives after the last
+    /// DONE (probe-confirmed 2026-07-15 — the SSMS freeze on
+    /// <c>use [master]</c>; go-mssqldb tolerates the late position, which is
+    /// how the ordering shipped unnoticed).
+    /// </summary>
+    private string? databaseAtMessageStart;
+
+    /// <summary>
+    /// Writes the database-change ENVCHANGE + INFO 5701 when the session
+    /// database differs from <see cref="databaseAtMessageStart"/>, matching
+    /// real SQL Server's token order for <c>USE</c> (ENVCHANGE, then INFO,
+    /// then the statement's DONE). Idempotent — the first call records the
+    /// new baseline, so the multiple call sites (per-final-DONE seams and
+    /// error paths) emit at most once per change.
+    /// </summary>
+    private void WriteDatabaseChangeIfAny(TdsTokenWriter writer)
+    {
+        var current = this.connection!.Database;
+        if (this.databaseAtMessageStart is null || string.Equals(current, this.databaseAtMessageStart, StringComparison.Ordinal))
+            return;
+        writer.WriteEnvChange(Tds.EnvDatabase, current, this.databaseAtMessageStart);
+        writer.WriteErrorOrInfo(Tds.TokenInfo, 5701, 2, 0, $"Changed database context to '{current}'.", "SIMULATED", "", 1);
+        this.databaseAtMessageStart = current;
     }
 
     /// <summary>
@@ -276,7 +303,10 @@ private async ValueTask StreamOutcomesAsync(SimulatedDbCommand command, TdsToken
                 }
 
                 hasOutcome = outcomes.MoveNext();
-                writer.WriteDoneToken(doneToken, (ushort)(this.OutcomeDoneStatus(hasOutcome, trailingTokensFollow) | Tds.DoneCount), rows);
+                var queryStatus = (ushort)(this.OutcomeDoneStatus(hasOutcome, trailingTokensFollow) | Tds.DoneCount);
+                if ((queryStatus & Tds.DoneMore) == 0)
+                    this.WriteDatabaseChangeIfAny(writer);
+                writer.WriteDoneToken(doneToken, queryStatus, rows);
             }
             else
             {
@@ -286,16 +316,22 @@ private async ValueTask StreamOutcomesAsync(SimulatedDbCommand command, TdsToken
                 if (affected >= 0)
                     status |= Tds.DoneCount;
 
+                if ((status & Tds.DoneMore) == 0)
+                    this.WriteDatabaseChangeIfAny(writer);
                 writer.WriteDoneToken(doneToken, status, Math.Max(affected, 0));
             }
         }
 
         // Trailing messages (batch ends in PRINT): INFO may never follow the
         // final DONE, so the last outcome's DONE stayed DONE_MORE (see
         // OutcomeDoneStatus) and the batch closes with its own final DONE.
+        // A mid-batch USE's ENVCHANGE must likewise precede the final DONE.
         var flushedTrailing = this.FlushInfoMessages(writer);
         if (!trailingTokensFollow && (flushedTrailing || !anyOutcome))
+        {
+            this.WriteDatabaseChangeIfAny(writer);
             writer.WriteDoneToken(doneToken, Tds.DoneFinal, 0);
+        }
     }
 
     /// <summary>
diff --git a/docs/claude/tds-endpoint.md b/docs/claude/tds-endpoint.md
@@ -28,7 +28,7 @@ Everything lives in `Network/` (internal) except the public `SimulatedNetworkLis
 - **SQLBatch** (type 1): ALL_HEADERS skipped via its leading length DWORD; UCS-2 text executed on the session connection. Per result set: COLMETADATA + ROW stream + DONE (`DONE_COUNT` + `DONE_MORE` when more outcomes follow); per non-query statement: DONE with `DONE_COUNT` only when `RecordsAffected >= 0`. Zero outcomes → single final DONE. All tokens are `0xFD DONE` (DONEPROC/DONEINPROC are proc-scoped and RPC isn't shipped).
 - **Errors**: `SimulatedSqlException.Errors` map field-for-field onto ERROR tokens (number/state/class/message/server/procedure/line) + DONE with `DONE_ERROR`; the session survives and keeps serving. `NotSupportedException` becomes a synthetic ERROR number 50000 class 16 prefixed `SqlServerSimulator:`.
 - **PRINT / low-severity RAISERROR**: the session subscribes to `SimulatedDbConnection.InfoMessage` and drains the queue as INFO tokens between statements and at batch end. **Each INFO flush in a SQLBatch response gets its own DONE**, mirroring real per-statement DONEs: info preceding an outcome is followed by `DONE_MORE` count 0 before the outcome's tokens, and trailing info (batch ends in PRINT) forces the last outcome's DONE to `DONE_MORE` with a closing `DONE_FINAL` count 0 after the INFO — an INFO token must never follow the final DONE. Without both, SqlClient's token reader stalls until command timeout on any batch mixing PRINT with a result set, and go-mssqldb silently drops the message (go-sqlcmd shakedown, 2026-07-14; the pre-fix oracle only covered PRINT-without-result-set). RPC responses are unaffected — every DONEINPROC already carries `DONE_MORE`.
-- **`USE`**: database change detected by before/after comparison and emitted as ENVCHANGE type 1 (after the DONEs, before flush — SqlClient processes it anywhere pre-EOM).
+- **`USE`**: database change detected by comparing the session database against its value at message start (`databaseAtMessageStart`) and emitted as ENVCHANGE type 1 + INFO 5701 (`Changed database context to '<db>'.`) **before the response's final DONE** — the seam fires at every final-DONE site (batch outcome DONEs, the closing DONE, error-path DONEs, and the RPC handlers' final DONEPROC) and is idempotent, so it emits at most once per message. Ordering is load-bearing: SqlClient's token reader stalls until command timeout on an ENVCHANGE that arrives after the last DONE (probe-confirmed 2026-07-15 — this froze SSMS on its first `use [master]` once master existed; go-mssqldb tolerates the late position, which is how the original after-the-DONEs ordering shipped unnoticed). The INFO 5701 is wire-layer-only — the in-process engine raises no InfoMessage for `USE`, a minor in-process/wire asymmetry matching the login response's synthesized 5701.
 - **Reset-connection status bit** (pooled-connection recycle): backing connection disposed and recreated on the same database, acked with the empty ENVCHANGE type 18 before the batch's tokens.
 - **Attention** (type 6): acked with DONE `DONE_ATTN`. Execution is synchronous per message, so attention is only observed between messages — a cancel never interrupts a running statement server-side, it just gets acked when the stream drains. In-process execution is fast enough that this matches observable SqlClient behavior.
 - **Bulk-load (7)**: ERROR 50000 naming the unsupported request type + DONE error (`SqlBulkCopy` is a planned follow-up).

Original file line number	Diff line number	Diff line change
`@@ -46,7 +46,7 @@ private async ValueTask ExecuteRpcMessageAsync(TdsMessage message, TdsTokenWrite`
`46`	`46`	`return;`
`47`	`47`	`}`
`48`	`48`
`49`		`- var databaseBefore = this.connection!.Database;`
	`49`	`+ this.databaseAtMessageStart = this.connection!.Database;`
`50`	`50`	`for (var i = 0; i < requests.Count; i++)`
`51`	`51`	`{`
`52`	`52`	`var moreRequests = i < requests.Count - 1;`
`@@ -58,19 +58,19 @@ private async ValueTask ExecuteRpcMessageAsync(TdsMessage message, TdsTokenWrite`
`58`	`58`	`{`
`59`	`59`	`_ = this.FlushInfoMessages(writer);`
`60`	`60`	`WriteErrors(writer, ex);`
	`61`	`+ if (!moreRequests)`
	`62`	`+ this.WriteDatabaseChangeIfAny(writer);`
`61`	`63`	`writer.WriteDoneToken(Tds.TokenDoneProc, (ushort)(Tds.DoneError \| (moreRequests ? Tds.DoneMore : Tds.DoneFinal)), 0);`
`62`	`64`	`}`
`63`	`65`	`catch (NotSupportedException ex)`
`64`	`66`	`{`
`65`	`67`	`_ = this.FlushInfoMessages(writer);`
`66`	`68`	`writer.WriteErrorOrInfo(Tds.TokenError, 50000, 1, 16, $"SqlServerSimulator: {ex.Message}", "SIMULATED", "", 1);`
	`69`	`+ if (!moreRequests)`
	`70`	`+ this.WriteDatabaseChangeIfAny(writer);`
`67`	`71`	`writer.WriteDoneToken(Tds.TokenDoneProc, (ushort)(Tds.DoneError \| (moreRequests ? Tds.DoneMore : Tds.DoneFinal)), 0);`
`68`	`72`	`}`
`69`	`73`	`}`
`70`		`-`
`71`		`- var databaseAfter = this.connection.Database;`
`72`		`- if (!string.Equals(databaseAfter, databaseBefore, StringComparison.Ordinal))`
`73`		`- writer.WriteEnvChange(Tds.EnvDatabase, databaseAfter, databaseBefore);`
`74`	`74`	`}`
`75`	`75`
`76`	`76`	`private async ValueTask DispatchRpcAsync(TdsRpcRequest request, TdsTokenWriter writer, bool moreRequests, CancellationToken cancellationToken)`
`@@ -180,6 +180,8 @@ private async ValueTask ExecuteStatementRpcAsync(`
`180`	`180`	`foreach (var (ordinal, wire, bound) in outputs)`
`181`	`181`	`TdsTypeCodec.WriteReturnValue(writer, checked((ushort)ordinal), wire.Name, wire.DbType, bound.Value);`
`182`	`182`
	`183`	`+ if (!moreRequests)`
	`184`	`+ this.WriteDatabaseChangeIfAny(writer);`
`183`	`185`	`writer.WriteDoneToken(Tds.TokenDoneProc, moreRequests ? Tds.DoneMore : Tds.DoneFinal, 0);`
`184`	`186`	`}`
`185`	`187`
`@@ -212,6 +214,8 @@ private async ValueTask ExecuteProcedureRpcAsync(TdsRpcRequest request, TdsToken`
`212`	`214`	`foreach (var (ordinal, wire, bound) in outputs)`
`213`	`215`	`TdsTypeCodec.WriteReturnValue(writer, checked((ushort)ordinal), wire.Name, wire.DbType, bound.Value);`
`214`	`216`
	`217`	`+ if (!moreRequests)`
	`218`	`+ this.WriteDatabaseChangeIfAny(writer);`
`215`	`219`	`writer.WriteDoneToken(Tds.TokenDoneProc, moreRequests ? Tds.DoneMore : Tds.DoneFinal, 0);`
`216`	`220`	`}`
`217`	`221`