X-Git-Url: http://git.giacomolli.uk/onechat/blobdiff_plain/6b21fcf78372995c6d8130c63f2f8aa2609bf7fe..a6a82d96b4906e9e1f09f45f69f48124f11ae1bb:/model_test.go?ds=sidebyside

diff --git a/model_test.go b/model_test.go
index 84b767e..11a7d00 100644
--- a/model_test.go
+++ b/model_test.go
@@ -85,6 +85,19 @@ func TestClientSendsTheSessionHeaderToChatGPT(t *testing.T) {
 			call := server.Call(0)
 			require.Equal(t, tc.want, call.header.Get("session_id"))
 			require.Equal(t, "agent-key", call.body["prompt_cache_key"])
+
+			// A subagent has its own key.
+			request := modelTestRequest()
+			request.Agent = "GreenOtter"
+			_, err = client.Create(t.Context(), request)
+			require.NoError(t, err)
+
+			call = server.Call(1)
+			if tc.want != "" {
+				require.Equal(t, "agent-key-GreenOtter", call.header.Get("session_id"))
+			}
+
+			require.Equal(t, "agent-key-GreenOtter", call.body["prompt_cache_key"])
 		})
 	}
 }
@@ -145,10 +158,12 @@ func TestClientRecordsEachCall(t *testing.T) {
 	client := newTestModelClient(t, server)
 	request := modelTestRequest()
 	request.Node = &NodeKey{2, 5}
+	request.Agent = "GreenOtter"
 	_, err := client.Create(t.Context(), request)
 	require.NoError(t, err)
 	require.NotContains(t, server.Call(0).body, "node")
 	require.NotContains(t, string(server.Call(0).raw), "[2,5]")
+	require.NotContains(t, server.Call(0).body, "agent")
 
 	_, err = client.Create(t.Context(), modelTestRequest())
 	require.ErrorContains(t, err, "HTTP 429")
@@ -168,12 +183,14 @@ func TestClientRecordsEachCall(t *testing.T) {
 	require.Equal(t, 7, records[0].OutputTokens)
 	require.Equal(t, quota, records[0].Quota)
 	require.Equal(t, &[2]int{2, 5}, records[0].Node)
+	require.Equal(t, "GreenOtter", records[0].Agent)
 	require.Equal(t, len("ok"), records[0].TextBytes)
 	require.Greater(t, records[0].Seconds, 0.0)
 	require.Empty(t, records[0].Error)
 	require.NotEmpty(t, records[0].Date)
 
 	require.Nil(t, records[1].Node)
+	require.Empty(t, records[1].Agent)
 	require.Zero(t, records[1].TextBytes)
 
 	require.Equal(t, http.StatusTooManyRequests, records[1].Status)