addressed code review

Signed-off-by: Nir Rozenbaum <[email protected]>

Author: nirrozenbaum

pkg/epp/plugins/plugin_state_test.go

@@ -66,8 +66,14 @@ func TestPluginState_ReadWrite(t *testing.T) {
 	assert.True(t, ok, "should be able to cast to pluginTestData")
 	assert.Equal(t, data1, td.value)
 
-	// Delete the req2 data and verify it's removed
+	// Delete the req2 data and verify content that was read before is still valid
+	readData, err = state.Read(req2, key)
+	assert.NoError(t, err)
 	state.Delete(req2)
+	td, ok = readData.(*pluginTestData)
+	assert.True(t, ok, "should be able to cast to pluginTestData")
+	assert.Equal(t, data2, td.value)
+	// try to read again aftet deletion, verify error
 	readData, err = state.Read(req2, key)
 	assert.Equal(t, ErrNotFound, err)
 	assert.Nil(t, readData, "expected no data after delete")

pkg/epp/scheduling/framework/plugins/multi/prefix/plugin.go

@@ -209,6 +209,7 @@ func (p *Plugin) PreRequest(ctx context.Context, request *types.LLMRequest, sche
 	targetPod := primaryProfileResult.TargetPods[0].GetPod() // get the first pod of the primary profile
 
 	state, err := plugins.ReadPluginStateKey[*SchedulingContextState](p.pluginState, request.RequestId, PrefixCachePluginType)
+	p.pluginState.Delete(request.RequestId) // delete the state explicitly after completing using it
 	if err != nil {
 		log.FromContext(ctx).Error(err, "failed to read prefix plugin state", "requestID", request.RequestId)
 		return